Add more Conn wrappers

dns.go

@@ -97,6 +97,48 @@ type Conn struct {
 	Attempts int
 }
 
+// Dial connects to the remote address raddr on the network net.
+// If the string laddr is not empty, it is used as the local address
+// for the connection. Any errors are return in err otherwise err is nil.
+func Dial(n, laddr, raddr string) (*Conn, os.Error) {
+        d := new(Conn)
+        c, err := net.Dial(n, laddr, raddr)
+        if err != nil {
+                return nil, err
+        }
+        switch n {
+        case "tcp":
+                d.TCP = c.(*net.TCPConn)
+                d.Addr = d.TCP.RemoteAddr()
+                d.Port = d.TCP.RemoteAddr().(*net.TCPAddr).Port
+        case "udp":
+                d.UDP = c.(*net.UDPConn)
+                d.Addr = d.UDP.RemoteAddr()
+                d.Port = d.UDP.RemoteAddr().(*net.UDPAddr).Port
+        }
+        return d, nil
+}
+
+// Fill in a Conn from a TCPConn
+func (d *Conn) SetTCPConn(l *net.TCPConn, a net.Addr) {
+        d.TCP = l
+        d.UDP = nil
+        if a == nil {
+                d.Addr = l.RemoteAddr()
+        }
+        d.Port = d.Addr.(*net.TCPAddr).Port
+}
+
+// Fill in a Conn from a UDPConn
+func (d *Conn) SetUDPConn(l *net.UDPConn, a net.Addr) {
+        d.TCP = nil
+        d.UDP = l
+        if a == nil {
+                d.Addr = l.RemoteAddr()
+        }
+        d.Port = d.Addr.(*net.UDPAddr).Port
+}
+
 // Create a new buffer of the appropiate size. With
 // TCP the buffer is 64K, with UDP the returned buffer
 // has a length of 4K bytes.

resolver.go

@@ -32,7 +32,6 @@ func (res *Resolver) Query(q *Msg) (d *Msg, err os.Error) {
 
 // Send a query to the nameserver using res, but perform TSIG validation.
 func (res *Resolver) QueryTsig(q *Msg, tsig *Tsig) (d *Msg, err os.Error) {
-	var c net.Conn
 	var inb []byte
 	in := new(Msg)
 	port, err := check(res, q)
@@ -49,25 +48,20 @@ func (res *Resolver) QueryTsig(q *Msg, tsig *Tsig) (d *Msg, err os.Error) {
 	}
 
 	for i := 0; i < len(res.Servers); i++ {
-		d := new(Conn)
+                var d *Conn
 		server := res.Servers[i] + ":" + port
 		t := time.Nanoseconds()
 		if res.Tcp {
-			c, err = net.Dial("tcp", "", server)
+			d, err = Dial("tcp", "", server)
 			if err != nil {
 				continue
 			}
-			d.TCP = c.(*net.TCPConn)
-			d.Addr = d.TCP.RemoteAddr()
 		} else {
-			c, err = net.Dial("udp", "", server)
+			d, err = Dial("udp", "", server)
 			if err != nil {
 				continue
 			}
-			d.UDP = c.(*net.UDPConn)
-			d.Addr = d.UDP.RemoteAddr()
 		}
-
 		d.Tsig = tsig
 		inb, err = d.Exchange(sending, false)
 		if err != nil {
@@ -75,7 +69,7 @@ func (res *Resolver) QueryTsig(q *Msg, tsig *Tsig) (d *Msg, err os.Error) {
 		}
 		in.Unpack(inb) // Discard error.
 		res.Rtt[server] = time.Nanoseconds() - t
-		c.Close()
+                d.Close()
 		break
 	}
 	if err != nil {

server.go

@@ -25,9 +25,7 @@ func HandleUDP(l *net.UDPConn, f func(*Conn, *Msg)) os.Error {
 		m = m[:n]
 
 		d := new(Conn)
-		d.UDP = l
-		d.Addr = addr
-		d.Port = addr.Port
+                d.SetUDPConn(l, addr)
 
 		msg := new(Msg)
 		if !msg.Unpack(m) {
@@ -49,9 +47,7 @@ func HandleTCP(l *net.TCPListener, f func(*Conn, *Msg)) os.Error {
 			return e
 		}
 		d := new(Conn)
-		d.TCP = c
-		d.Addr = c.RemoteAddr()
-		d.Port = d.TCP.RemoteAddr().(*net.TCPAddr).Port
+                d.SetTCPConn(c, nil)
 
 		msg := new(Msg)
 		err := d.ReadMsg(msg)