Merge branch 'ISSUE-95' of https://github.com/EverythingMe/dns into udp

CONTRIBUTORS

@@ -5,3 +5,4 @@ Dave Cheney
 Dusty Wilson
 Marek Majkowski
 Peter van Dijk
+Omri Bahumi

server.go

@@ -44,9 +44,10 @@ type response struct {
 	tsigTimersOnly bool
 	tsigRequestMAC string
 	tsigSecret     map[string]string // the tsig secrets
-	udp            *net.UDPConn      // i/o connection if UDP was used
-	tcp            *net.TCPConn      // i/o connection if TCP was used
-	remoteAddr     net.Addr          // address of the client
+	udp            *UDPConn          // i/o connection if UDP was used
+	udpSession     *UDPSession
+	tcp            *net.TCPConn // i/o connection if TCP was used
+	remoteAddr     net.Addr     // address of the client
 }
 
 // ServeMux is an DNS request multiplexer. It matches the
@@ -242,7 +243,13 @@ func (srv *Server) ListenAndServe() error {
 		if e != nil {
 			return e
 		}
-		return srv.serveUDP(l)
+
+		ll, e := NewUDPConn(l)
+		if e != nil {
+			return e
+		}
+
+		return srv.serveUDP(ll)
 	}
 	return &Error{err: "bad network"}
 }
@@ -268,14 +275,14 @@ func (srv *Server) serveTCP(l *net.TCPListener) error {
 		if e != nil {
 			continue
 		}
-		go srv.serve(rw.RemoteAddr(), handler, m, nil, rw)
+		go srv.serve(rw.RemoteAddr(), handler, m, nil, nil, rw)
 	}
 	panic("dns: not reached")
 }
 
 // serveUDP starts a UDP listener for the server.
 // Each request is handled in a seperate goroutine.
-func (srv *Server) serveUDP(l *net.UDPConn) error {
+func (srv *Server) serveUDP(l *UDPConn) error {
 	defer l.Close()
 	handler := srv.Handler
 	if handler == nil {
@@ -286,19 +293,19 @@ func (srv *Server) serveUDP(l *net.UDPConn) error {
 		rtimeout = srv.ReadTimeout
 	}
 	for {
-		m, a, e := srv.readUDP(l, rtimeout)
+		m, s, e := srv.readUDP(l, rtimeout)
 		if e != nil {
 			// TODO(miek): logging?
 			continue
 		}
-		go srv.serve(a, handler, m, l, nil)
+		go srv.serve(s.RemoteAddr(), handler, m, l, s, nil)
 	}
 	panic("dns: not reached")
 }
 
 // Serve a new connection.
-func (srv *Server) serve(a net.Addr, h Handler, m []byte, u *net.UDPConn, t *net.TCPConn) {
-	w := &response{tsigSecret: srv.TsigSecret, udp: u, tcp: t, remoteAddr: a}
+func (srv *Server) serve(a net.Addr, h Handler, m []byte, u *UDPConn, s *UDPSession, t *net.TCPConn) {
+	w := &response{tsigSecret: srv.TsigSecret, udp: u, tcp: t, remoteAddr: a, udpSession: s}
 	q := 0
 Redo:
 	req := new(Msg)
@@ -385,9 +392,9 @@ func (srv *Server) readTCP(conn *net.TCPConn, timeout time.Duration) ([]byte, er
 	return m, nil
 }
 
-func (srv *Server) readUDP(conn *net.UDPConn, timeout time.Duration) ([]byte, net.Addr, error) {
+func (srv *Server) readUDP(conn *UDPConn, timeout time.Duration) ([]byte, *UDPSession, error) {
 	m := make([]byte, srv.UDPSize)
-	n, a, e := conn.ReadFromUDP(m)
+	n, s, e := conn.ReadFromSessionUDP(m)
 	if e != nil || n == 0 {
 		if e != nil {
 			return nil, nil, e
@@ -395,7 +402,7 @@ func (srv *Server) readUDP(conn *net.UDPConn, timeout time.Duration) ([]byte, ne
 		return nil, nil, ErrShortRead
 	}
 	m = m[:n]
-	return m, a, nil
+	return m, s, nil
 }
 
 // WriteMsg implements the ResponseWriter.WriteMsg method.
@@ -423,7 +430,7 @@ func (w *response) WriteMsg(m *Msg) (err error) {
 func (w *response) Write(m []byte) (int, error) {
 	switch {
 	case w.udp != nil:
-		n, err := w.udp.WriteTo(m, w.remoteAddr)
+		n, err := w.udp.WriteToSessionUDP(m, w.udpSession)
 		return n, err
 	case w.tcp != nil:
 		lm := len(m)

udp.go

@@ -0,0 +1,87 @@
+// Copyright 2011 Miek Gieben. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package dns
+
+import (
+	"net"
+	"syscall"
+)
+
+type UDPSession struct {
+	raddr   *net.UDPAddr
+	context []byte
+}
+
+func (session *UDPSession) RemoteAddr() net.Addr {
+	return session.raddr
+}
+
+type UDPConn struct {
+	*net.UDPConn
+}
+
+// Wrap a net.UDPConn with dns.UDPConn struct
+// Initialize the underlying net.UDPConn for supporting "sessions"
+// Sessions solve https://github.com/miekg/dns/issues/95
+func NewUDPConn(conn *net.UDPConn) (newconn *UDPConn, err error) {
+	// this function is implemented on a per platform basis. See udp_*.go for more details
+	err = udpPatchSocket(conn)
+
+	if err != nil {
+		return
+	}
+
+	return &UDPConn{conn}, nil
+}
+
+// Just like net.UDPConn.ReadFrom(), but returns a session object instead of net.UDPAddr
+// (RemoteAddr() is available from the UDPSession object)
+func (conn *UDPConn) ReadFromSessionUDP(b []byte) (n int, session *UDPSession, err error) {
+	oob := make([]byte, 40)
+
+	n, oobn, _, raddr, err := conn.ReadMsgUDP(b, oob)
+	if err != nil {
+		return
+	}
+
+	session = &UDPSession{raddr, oob[:oobn]}
+
+	return
+}
+
+// Just like net.UDPConn.WritetTo(), but uses a session object instead of net.Addr
+func (conn *UDPConn) WriteToSessionUDP(b []byte, session *UDPSession) (n int, err error) {
+	n, _, err = conn.WriteMsgUDP(b, session.context, session.raddr)
+	return
+}
+
+func udpPatchSocket(conn *net.UDPConn) (err error) {
+	file, err := conn.File()
+	if err != nil {
+		return
+	}
+
+	sa, err := syscall.Getsockname(int(file.Fd()))
+
+	ipv4, ipv6 := false, false
+	switch sa.(type) {
+	case *syscall.SockaddrInet6:
+		ipv6 = true
+
+		// dual stack. See http://stackoverflow.com/questions/1618240/how-to-support-both-ipv4-and-ipv6-connections
+		v6only, err := syscall.GetsockoptInt(int(file.Fd()), syscall.IPPROTO_IPV6, syscall.IPV6_V6ONLY)
+		if err != nil {
+			return err
+		}
+
+		if v6only == 0 {
+			ipv4 = true
+		}
+	case *syscall.SockaddrInet4:
+		ipv4 = true
+	}
+
+	return udpPatchSocketTypes(conn, ipv4, ipv6)
+}

udp_linux.go

@@ -0,0 +1,43 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+// Extensions of the original work are copyright (c) 2011 Miek Gieben
+
+// +build linux
+
+package dns
+
+import (
+	"net"
+	"syscall"
+)
+
+// Linux implementation for preparing the socket for sessions
+// Based on http://stackoverflow.com/questions/3062205/setting-the-source-ip-for-a-udp-socket
+// and http://blog.powerdns.com/2012/10/08/on-binding-datagram-udp-sockets-to-the-any-addresses/
+func udpPatchSocketTypes(conn *net.UDPConn, ipv4, ipv6 bool) (err error) {
+	file, err := conn.File()
+	if err != nil {
+		return
+	}
+
+	if ipv4 {
+		// socket supports IPv4
+
+		err = syscall.SetsockoptInt(int(file.Fd()), syscall.IPPROTO_IP, syscall.IP_PKTINFO, 1)
+		if err != nil {
+			return err
+		}
+	}
+
+	if ipv6 {
+		// socket supports IPv6
+
+		err = syscall.SetsockoptInt(int(file.Fd()), syscall.IPPROTO_IPV6, syscall.IPV6_RECVPKTINFO, 1)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

udp_other.go

@@ -0,0 +1,19 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+// Extensions of the original work are copyright (c) 2011 Miek Gieben
+
+// +build !linux
+
+package dns
+
+import (
+	"net"
+)
+
+// Default implementation for preparing the socket for sessions
+// This actually does nothing. See udp_linux.go for an example of how to implement this.
+// Make sure you edit the comment on the top of this file accordingly when adding implementations
+func udpPatchSocketTypes(conn *net.UDPConn, ipv4, ipv6 bool) (err error) {
+	return nil
+}