From aea5c1da03e54b323a5cc27d5faefa0a3346f292 Mon Sep 17 00:00:00 2001
From: Omri Bahumi <omrib@everything.me>
Date: Tue, 8 Jul 2014 15:18:33 +0300
Subject: [PATCH] [ISSUE-95] Patch only relevant stacks (v4, v6, dual)

---
 udp.go       | 32 +++++++++++++++++++++++++++++++-
 udp_linux.go | 25 +++++++++++++++++--------
 udp_other.go |  4 ++--
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/udp.go b/udp.go
index c03f9653..7940d095 100644
--- a/udp.go
+++ b/udp.go
@@ -2,6 +2,7 @@ package dns
 
 import (
 	"net"
+	"syscall"
 )
 
 type UDPSession struct {
@@ -22,7 +23,7 @@ type UDPConn struct {
 // 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 = udpSocketOobData(conn)
+	err = udpPatchSocket(conn)
 
 	if err != nil {
 		return
@@ -51,3 +52,32 @@ func (conn *UDPConn) WriteToSessionUDP(b []byte, session *UDPSession) (n int, er
 	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)
+}
diff --git a/udp_linux.go b/udp_linux.go
index edf90283..00f948f5 100644
--- a/udp_linux.go
+++ b/udp_linux.go
@@ -10,20 +10,29 @@ import (
 // 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 udpSocketOobData(conn *net.UDPConn) (err error) {
+func udpPatchSocketTypes(conn *net.UDPConn, ipv4, ipv6 bool) (err error) {
 	file, err := conn.File()
 	if err != nil {
 		return
 	}
 
-	// IPv4 support
-	err = syscall.SetsockoptInt(int(file.Fd()), syscall.IPPROTO_IP, syscall.IP_PKTINFO, 1)
-	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
+		}
 	}
 
-	// IPv6 support
-	err = syscall.SetsockoptInt(int(file.Fd()), syscall.IPPROTO_IPV6, syscall.IPV6_RECVPKTINFO, 1)
+	if ipv6 {
+		// socket supports IPv6
 
-	return
+		err = syscall.SetsockoptInt(int(file.Fd()), syscall.IPPROTO_IPV6, syscall.IPV6_RECVPKTINFO, 1)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
diff --git a/udp_other.go b/udp_other.go
index 7fd2d163..37e77228 100644
--- a/udp_other.go
+++ b/udp_other.go
@@ -9,6 +9,6 @@ import (
 // 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 udpSocketOobData(conn *net.UDPConn) (err error) {
-	return
+func udpPatchSocketTypes(conn *net.UDPConn, ipv4, ipv6 bool) (err error) {
+	return nil
 }