From 1bf0864632738787fcbae20c2f24fa8fb6dbeec5 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Thu, 31 May 2012 21:40:07 +0200
Subject: [PATCH] First pass for edns subnet support

---
 edns.go   | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++-----
 ex/q/q.go | 28 +++++++++++++++-----
 msg.go    |  7 +++--
 3 files changed, 98 insertions(+), 15 deletions(-)

diff --git a/edns.go b/edns.go
index 32592b50..ed516506 100644
--- a/edns.go
+++ b/edns.go
@@ -2,18 +2,19 @@ package dns
 
 import (
 	"encoding/hex"
+	"errors"
 	"net"
 	"strconv"
 )
 
 // EDNS0 Option codes.
 const (
-	_                = iota
+	_           = iota
 	EDNS0LLQ             // not used
 	EDNS0UL              // not used
 	EDNS0NSID            // NSID, RFC5001
 	EDNS0SUBNET = 0x50fa // client-subnet draft
-	_DO              = 1 << 7 // dnssec ok
+	_DO         = 1 << 7 // dnssec ok
 )
 
 /* 
@@ -58,6 +59,8 @@ func (rr *RR_OPT) String() string {
 				}
 				s += "  " + r
 			}
+		case *EDNS0_SUBNET:
+			s += "\n; SUBNET: " + o.String()
 		}
 	}
 	return s
@@ -130,8 +133,8 @@ type EDNS0 interface {
 }
 
 type EDNS0_NSID struct {
-	Code uint16
-	Nsid string		// This string needs to be hex encoded
+	Code uint16 // Always EDNS0NSID
+	Nsid string // This string needs to be hex encoded
 }
 
 func (e *EDNS0_NSID) Option() uint16 {
@@ -155,9 +158,70 @@ func (e *EDNS0_NSID) String() string {
 }
 
 type EDNS0_SUBNET struct {
-	Code          uint16
-	Family        uint16
+	Code          uint16 // Always EDNS0SUBNET
+	Family        uint16 // 1 for IP, 2 for IP6
 	SourceNetmask uint8
 	SourceScope   uint8
-	Address       []net.IP
+	Address       net.IP
+}
+
+func (e *EDNS0_SUBNET) Option() uint16 {
+	return e.Code
+}
+
+func (e *EDNS0_SUBNET) Pack() ([]byte, error) {
+	b := make([]byte, 4)
+	b[0], b[1] = packUint16(e.Family)
+	b[2] = e.SourceNetmask
+	b[3] = e.SourceScope
+	switch e.Family {
+	case 1:
+		// just copy? TODO (also in msg.go...)
+		ip := make([]byte, net.IPv4len)
+		for i := 0; i < net.IPv4len; i++ {
+			if i+1 > len(e.Address) {
+				break
+			}
+			ip[i] = e.Address[i]
+		}
+		b = append(b, ip...)
+	case 2:
+		ip := make([]byte, net.IPv6len)
+		for i := 0; i < net.IPv6len; i++ {
+			if i+1 > len(e.Address) {
+				break
+			}
+			ip[i] = e.Address[i]
+		}
+		b = append(b, ip...)
+	default:
+		return nil, errors.New("bad address family")
+	}
+	return b, nil
+}
+
+func (e *EDNS0_SUBNET) Unpack(b []byte) {
+	// TODO: length of b
+	e.Family, _ = unpackUint16(b, 0)
+	e.SourceNetmask = b[2]
+	e.SourceScope = b[3]
+	switch e.Family {
+	case 1:
+		if len(b) == 8 {
+			e.Address = net.IPv4(b[4], b[5], b[6], b[7])
+		}
+	case 2:
+		if len(b) == 20 {
+			e.Address = net.IP{b[4], b[4+1], b[4+2], b[4+3], b[4+4],
+				b[4+5], b[4+6], b[4+7], b[4+8], b[4+9], b[4+10],
+				b[4+11], b[4+12], b[4+13], b[4+14], b[4+15]}
+		}
+	}
+	return
+}
+
+func (e *EDNS0_SUBNET) String() string {
+	return e.Address.String() + "/" +
+		strconv.Itoa(int(e.SourceNetmask)) + "/" +
+		strconv.Itoa(int(e.SourceScope))
 }
diff --git a/ex/q/q.go b/ex/q/q.go
index 01820392..a7f69e1c 100644
--- a/ex/q/q.go
+++ b/ex/q/q.go
@@ -4,6 +4,7 @@ import (
 	"dns"
 	"flag"
 	"fmt"
+	"net"
 	"os"
 	"strconv"
 	"strings"
@@ -47,7 +48,8 @@ func main() {
 	rd := flag.Bool("rd", true, "set RD flag in query")
 	fallback := flag.Bool("fallback", false, "fallback to 4096 bytes bufsize and after that TCP")
 	tcp := flag.Bool("tcp", false, "TCP mode")
-	nsid := flag.Bool("nsid", false, "ask for NSID")
+	nsid := flag.Bool("nsid", false, "set edns nsid option")
+	client := flag.String("client", "", "set edns client-subnet option")
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr, "Usage: %s [@server] [qtype] [qclass] [name ...]\n", os.Args[0])
 		flag.PrintDefaults()
@@ -162,7 +164,8 @@ Flags:
 	m.MsgHdr.CheckingDisabled = *cd
 	m.MsgHdr.RecursionDesired = *rd
 	m.Question = make([]dns.Question, 1)
-	if *dnssec || *nsid {
+
+	if *dnssec || *nsid || *client != "" {
 		o := new(dns.RR_OPT)
 		o.Hdr.Name = "."
 		o.Hdr.Rrtype = dns.TypeOPT
@@ -171,12 +174,26 @@ Flags:
 			o.SetUDPSize(dns.DefaultMsgSize)
 		}
 		if *nsid {
-			// Ask for it
 			e := new(dns.EDNS0_NSID)
-			e.Nsid = ""
 			e.Code = dns.EDNS0NSID
 			o.Option = append(o.Option, e)
 		}
+		if *client != "" {
+			e := new(dns.EDNS0_SUBNET)
+			e.Code = dns.EDNS0SUBNET
+			e.SourceNetmask = 0
+			e.SourceScope = 0
+			e.Address = net.ParseIP(*client)
+			if e.Address == nil {
+				fmt.Fprintf(os.Stderr, "Failure to parse IP address: %s\n", *client)
+				return
+			}
+			e.Family = 1	// IP4
+			if len(e.Address) > net.IPv4len {
+				e.Family = 2 // IP6
+			}
+			o.Option = append(o.Option, e)
+		}
 		m.Extra = append(m.Extra, o)
 	}
 
@@ -243,8 +260,7 @@ forever:
 				}
 
 				fmt.Printf("%v", r.Reply)
-				fmt.Printf("\n;; query time: %.3d µs, server: %s(%s), size: %dB\n", r.Rtt/1e3, r.RemoteAddr, r.RemoteAddr.Network(), r.Reply.Len())
-				// Server maybe
+				fmt.Printf("\n;; query time: %.3d µs, server: %s(%s), size: %d bytes\n", r.Rtt/1e3, r.RemoteAddr, r.RemoteAddr.Network(), r.Reply.Len())
 			}
 			i++
 			if i == len(qname) {
diff --git a/msg.go b/msg.go
index 02b484f5..89783260 100644
--- a/msg.go
+++ b/msg.go
@@ -404,7 +404,7 @@ func packStructValue(val reflect.Value, msg []byte, off int, compression map[str
 					element := val.Field(i).Index(j).Interface()
 					b, e := element.(EDNS0).Pack()
 					if e != nil {
-						println("dns: failure packing OTP")
+						println("dns: failure packing OPT")
 						return lenmsg, false
 					}
 					// Option code
@@ -707,7 +707,10 @@ func unpackStructValue(val reflect.Value, msg []byte, off int) (off1 int, ok boo
 					edns = append(edns, e)
 					off = off1 + int(optlen)
 				case EDNS0SUBNET:
-					// ..
+					e := new(EDNS0_SUBNET)
+					e.Unpack(msg[off1 : off1+int(optlen)])
+					edns = append(edns, e)
+					off = off1 + int(optlen)
 				}
 				fv.Set(reflect.ValueOf(edns))
 				// goto ??