Fix the reflect server

_examples/reflect/reflect.go

@@ -10,8 +10,8 @@
  *
  * Stephane Bortzmeyer <stephane+grong@bortzmeyer.org>
  *
- * Adapted to Go DNS (i.e. completely rewritten)
  * Miek Gieben <miek@miek.nl>
+ * Adapted to Go DNS (i.e. completely rewritten)
  */
 
 package main
@@ -20,46 +20,74 @@ import (
 	"net"
 	"dns"
 	"fmt"
-	//	"os"
+	"os/signal"
-	//	"os/signal"
+	"strconv"
-//	"strconv"
 )
 
 func handleReflect(w dns.ResponseWriter, r *dns.Msg) {
 	m := new(dns.Msg)
 	m.SetReply(r)
 	m.Extra = make([]dns.RR, 1)
-//	m.Answer = make([]dns.RR, 1)
+	m.Answer = make([]dns.RR, 1)
+        var (
+                v4 bool
+                rr dns.RR
+                str string
+                a net.IP
+        )
 
-        println(w.RemoteAddr())
+        if ip, ok := w.RemoteAddr().(*net.UDPAddr); ok {
-	ad := net.ParseIP(w.RemoteAddr())
+                str = "Port: " + strconv.Itoa(ip.Port) + " (udp)"
-        println(ad.String())
+                a = ip.IP
-//	if ad.To4() != nil {
+                v4 = a.To4() != nil
-//		rr := new(dns.RR_A)
+        }
-//		rr.Hdr = dns.RR_Header{Name: "whoami.miek.nl.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
+        if ip, ok := w.RemoteAddr().(*net.TCPAddr); ok {
-//		rr.A = ad
+                str = "Port: " + strconv.Itoa(ip.Port) + " (tcp)"
-//		m.Answer[0] = rr
+                a = ip.IP
-//	} else {
+                v4 = a.To4() != nil
-//		r := new(dns.RR_AAAA)
+        }
-//		r.Hdr = dns.RR_Header{Name: "whoami.miek.nl.", Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0}
+
-//		r.AAAA = ad
+	if v4 {
-//		m.Answer[0] = r
+		rr = new(dns.RR_A)
-//	}
+                rr.(*dns.RR_A).Hdr = dns.RR_Header{Name: "whoami.miek.nl.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
+		rr.(*dns.RR_A).A = a
+	} else {
+		rr = new(dns.RR_AAAA)
+                rr.(*dns.RR_AAAA).Hdr = dns.RR_Header{Name: "whoami.miek.nl.", Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0}
+		rr.(*dns.RR_AAAA).AAAA = a
+	}
 
 	t := new(dns.RR_TXT)
 	t.Hdr = dns.RR_Header{Name: "whoami.miek.nl.", Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
-//	t.Txt = "Port: " + strconv.Itoa(w.RemotePort()) + " (" + w.RemoteTransport() + ")"
+	t.Txt = str
-	t.Txt = "Port: " + " (" + w.RemoteTransport() + ")"
 	m.Extra[0] = t
+	m.Answer[0] = rr
 
-	b, _ := m.Pack()
+	b, ok := m.Pack()
+        if !ok { return }
 	w.Write(b)
 }
 
 func main() {
 	dns.HandleFunc(".", handleReflect)
-	err := dns.ListenAndServe(":8053", "udp", nil)
+        go func() {
-	if err != nil {
+                err := dns.ListenAndServe(":8053", "udp", nil)
-		fmt.Printf("Failed to setup the server")
+                if err != nil {
-	}
+                        fmt.Printf("Failed to setup the udp server")
+                }
+        }()
+        go func() {
+                err := dns.ListenAndServe(":8053", "tcp", nil)
+                if err != nil {
+                        fmt.Printf("Failed to setup the tcp server")
+                }
+        }()
+forever:
+        for {
+                select {
+                case <-signal.Incoming:
+                        fmt.Printf("Signal received, stopping")
+                        break forever
+                }
+        }
 }

server.go

@@ -21,10 +21,8 @@ type Handler interface {
 // A ResponseWriter interface is used by an DNS handler to
 // construct an DNS response.
 type ResponseWriter interface {
-	// RemoteAddr returns the address of the client that sent the current request
+	// RemoteAddr returns the net.Addr of the client that sent the current request
-	RemoteAddr() string
+	RemoteAddr() net.Addr
-        // RemoteTransport return "udp" or "tcp" depending on the transport used
-        RemoteTransport() string
 
         // Write a reply back
 	Write([]byte) (int, os.Error)
@@ -351,6 +349,7 @@ func (w *response) Write(data []byte) (n int, err os.Error) {
 	case w.conn._UDP != nil:
 		n, err = w.conn._UDP.WriteTo(data, w.conn.remoteAddr)
 		if err != nil {
+                        println(err.String())
 			return 0, err
 		}
 	case w.conn._TCP != nil:
@@ -382,12 +381,4 @@ func (w *response) Write(data []byte) (n int, err os.Error) {
 }
 
 // RemoteAddr implements the ResponseWriter.RemoteAddr method
-func (w *response) RemoteAddr() string { return w.conn.remoteAddr.String() }
+func (w *response) RemoteAddr() net.Addr { return w.conn.remoteAddr }
-
-// RemoteTransport implements the ResponseWriter.RemoteTransport method
-func (w *response) RemoteTransport() string {
-        if w.conn._UDP != nil {
-                return "udp"
-        }
-        return "tcp"
-}