dns/dnsmessage: add basic support for EDNS(0)

This change introduces OPTResourse type to support DNS messages
containing various extension options as defined in RFC 6891.
Parsing and building OPT pseudo records requires allocations like TXT
records.

Also adds DNSSECAllowed, ExtendedRCode and SetEDNS0 methods to
ResourceHeader for convenience.

Updates golang/go#6464.
Updates golang/go#16218.

Change-Id: Ib72cea277201e4122c6b5effa320084ff351c886
Reviewed-on: https://go-review.googlesource.com/47170
Reviewed-by: Ian Gudger <igudger@google.com>

dns/dnsmessage/message.go

@@ -5,6 +5,9 @@
 // Package dnsmessage provides a mostly RFC 1035 compliant implementation of
 // DNS message packing and unpacking.
 //
+// The package also supports messages with Extension Mechanisms for DNS
+// (EDNS(0)) as defined in RFC 6891.
+//
 // This implementation is designed to minimize heap allocations and avoid
 // unnecessary packing and unpacking as much as possible.
 package dnsmessage
@@ -39,6 +42,7 @@ const (
 	TypeTXT   Type = 16
 	TypeAAAA  Type = 28
 	TypeSRV   Type = 33
+	TypeOPT   Type = 41
 
 	// Question.Type
 	TypeWKS   Type = 11
@@ -802,6 +806,24 @@ func (p *Parser) AAAAResource() (AAAAResource, error) {
 	return r, nil
 }
 
+// OPTResource parses a single OPTResource.
+//
+// One of the XXXHeader methods must have been called before calling this
+// method.
+func (p *Parser) OPTResource() (OPTResource, error) {
+	if !p.resHeaderValid || p.resHeader.Type != TypeOPT {
+		return OPTResource{}, ErrNotStarted
+	}
+	r, err := unpackOPTResource(p.msg, p.off, p.resHeader.Length)
+	if err != nil {
+		return OPTResource{}, err
+	}
+	p.off += int(p.resHeader.Length)
+	p.resHeaderValid = false
+	p.index++
+	return r, nil
+}
+
 // Unpack parses a full Message.
 func (m *Message) Unpack(msg []byte) error {
 	var p Parser
@@ -1278,6 +1300,30 @@ func (b *Builder) AAAAResource(h ResourceHeader, r AAAAResource) error {
 	return nil
 }
 
+// OPTResource adds a single OPTResource.
+func (b *Builder) OPTResource(h ResourceHeader, r OPTResource) error {
+	if err := b.checkResourceSection(); err != nil {
+		return err
+	}
+	h.Type = r.realType()
+	msg, length, err := h.pack(b.msg, b.compression, b.start)
+	if err != nil {
+		return &nestedError{"ResourceHeader", err}
+	}
+	preLen := len(msg)
+	if msg, err = r.pack(msg, b.compression, b.start); err != nil {
+		return &nestedError{"OPTResource body", err}
+	}
+	if err := h.fixLen(msg, length, preLen); err != nil {
+		return err
+	}
+	if err := b.incrementSectionCount(); err != nil {
+		return err
+	}
+	b.msg = msg
+	return nil
+}
+
 // Finish ends message building and generates a binary message.
 func (b *Builder) Finish() ([]byte, error) {
 	if b.section < sectionHeader {
@@ -1366,6 +1412,44 @@ func (h *ResourceHeader) fixLen(msg []byte, length []byte, preLen int) error {
 	return nil
 }
 
+// EDNS(0) wire costants.
+const (
+	edns0Version = 0
+
+	edns0DNSSECOK     = 0x00008000
+	ednsVersionMask   = 0x00ff0000
+	edns0DNSSECOKMask = 0x00ff8000
+)
+
+// SetEDNS0 configures h for EDNS(0).
+//
+// The provided extRCode must be an extedned RCode.
+func (h *ResourceHeader) SetEDNS0(udpPayloadLen int, extRCode RCode, dnssecOK bool) error {
+	h.Name = Name{Data: [nameLen]byte{'.'}, Length: 1} // RFC 6891 section 6.1.2
+	h.Type = TypeOPT
+	h.Class = Class(udpPayloadLen)
+	h.TTL = uint32(extRCode) >> 4 << 24
+	if dnssecOK {
+		h.TTL |= edns0DNSSECOK
+	}
+	return nil
+}
+
+// DNSSECAllowed reports whether the DNSSEC OK bit is set.
+func (h *ResourceHeader) DNSSECAllowed() bool {
+	return h.TTL&edns0DNSSECOKMask == edns0DNSSECOK // RFC 6891 section 6.1.3
+}
+
+// ExtendedRCode returns an extended RCode.
+//
+// The provided rcode must be the RCode in DNS message header.
+func (h *ResourceHeader) ExtendedRCode(rcode RCode) RCode {
+	if h.TTL&ednsVersionMask == edns0Version { // RFC 6891 section 6.1.3
+		return RCode(h.TTL>>24<<4) | rcode
+	}
+	return rcode
+}
+
 func skipResource(msg []byte, off int) (int, error) {
 	newOff, err := skipName(msg, off)
 	if err != nil {
@@ -1794,6 +1878,11 @@ func unpackResourceBody(msg []byte, off int, hdr ResourceHeader) (ResourceBody,
 		rb, err = unpackSRVResource(msg, off)
 		r = &rb
 		name = "SRV"
+	case TypeOPT:
+		var rb OPTResource
+		rb, err = unpackOPTResource(msg, off, hdr.Length)
+		r = &rb
+		name = "OPT"
 	}
 	if err != nil {
 		return nil, off, &nestedError{name + " record", err}
@@ -2101,3 +2190,58 @@ func unpackAAAAResource(msg []byte, off int) (AAAAResource, error) {
 	}
 	return AAAAResource{aaaa}, nil
 }
+
+// An OPTResource is an OPT pseudo Resource record.
+//
+// The pseudo resource record is part of the extension mechanisms for DNS
+// as defined in RFC 6891.
+type OPTResource struct {
+	Options []Option
+}
+
+// An Option represents a DNS message option within OPTResource.
+//
+// The message option is part of the extension mechanisms for DNS as
+// defined in RFC 6891.
+type Option struct {
+	Code uint16 // option code
+	Data []byte
+}
+
+func (r *OPTResource) realType() Type {
+	return TypeOPT
+}
+
+func (r *OPTResource) pack(msg []byte, compression map[string]int, compressionOff int) ([]byte, error) {
+	for _, opt := range r.Options {
+		msg = packUint16(msg, opt.Code)
+		l := uint16(len(opt.Data))
+		msg = packUint16(msg, l)
+		msg = packBytes(msg, opt.Data)
+	}
+	return msg, nil
+}
+
+func unpackOPTResource(msg []byte, off int, length uint16) (OPTResource, error) {
+	var opts []Option
+	for oldOff := off; off < oldOff+int(length); {
+		var err error
+		var o Option
+		o.Code, off, err = unpackUint16(msg, off)
+		if err != nil {
+			return OPTResource{}, &nestedError{"Code", err}
+		}
+		var l uint16
+		l, off, err = unpackUint16(msg, off)
+		if err != nil {
+			return OPTResource{}, &nestedError{"Data", err}
+		}
+		o.Data = make([]byte, l)
+		if copy(o.Data, msg[off:]) != int(l) {
+			return OPTResource{}, &nestedError{"Data", errCalcLen}
+		}
+		off += int(l)
+		opts = append(opts, o)
+	}
+	return OPTResource{opts}, nil
+}

dns/dnsmessage/message_test.go

@@ -20,6 +20,14 @@ func mustNewName(name string) Name {
 	return n
 }
 
+func mustEDNS0ResourceHeader(l int, extrc RCode, do bool) ResourceHeader {
+	h := ResourceHeader{Class: ClassINET}
+	if err := h.SetEDNS0(l, extrc, do); err != nil {
+		panic(err)
+	}
+	return h
+}
+
 func (m *Message) String() string {
 	s := fmt.Sprintf("Message: %#v\n", &m.Header)
 	if len(m.Questions) > 0 {
@@ -569,6 +577,7 @@ func TestBuilderResourceError(t *testing.T) {
 		{"SRVResource", func(b *Builder) error { return b.SRVResource(ResourceHeader{}, SRVResource{}) }},
 		{"AResource", func(b *Builder) error { return b.AResource(ResourceHeader{}, AResource{}) }},
 		{"AAAAResource", func(b *Builder) error { return b.AAAAResource(ResourceHeader{}, AAAAResource{}) }},
+		{"OPTResource", func(b *Builder) error { return b.OPTResource(ResourceHeader{}, OPTResource{}) }},
 	}
 
 	envs := []struct {
@@ -675,8 +684,15 @@ func TestBuilder(t *testing.T) {
 		t.Fatal("b.StartAdditionals():", err)
 	}
 	for _, a := range msg.Additionals {
-		if err := b.TXTResource(a.Header, *a.Body.(*TXTResource)); err != nil {
-			t.Fatalf("b.TXTResource(%#v): %v", a, err)
+		switch a.Body.(type) {
+		case *TXTResource:
+			if err := b.TXTResource(a.Header, *a.Body.(*TXTResource)); err != nil {
+				t.Fatalf("Builder.TXTResource(%#v) = %v", a, err)
+			}
+		case *OPTResource:
+			if err := b.OPTResource(a.Header, *a.Body.(*OPTResource)); err != nil {
+				t.Fatalf("Builder.OPTResource(%#v) = %v", a, err)
+			}
 		}
 	}
 
@@ -745,6 +761,145 @@ func TestResourcePack(t *testing.T) {
 	}
 }
 
+func TestOptionPackUnpack(t *testing.T) {
+	for _, tt := range []struct {
+		name     string
+		w        []byte // wire format of m.Additionals
+		m        Message
+		dnssecOK bool
+		extRCode RCode
+	}{
+		{
+			name: "without EDNS(0) options",
+			w: []byte{
+				0x00, 0x00, 0x29, 0x10, 0x00, 0xfe, 0x00, 0x80,
+				0x00, 0x00, 0x00,
+			},
+			m: Message{
+				Header: Header{RCode: RCodeFormatError},
+				Questions: []Question{
+					{
+						Name:  mustNewName("."),
+						Type:  TypeA,
+						Class: ClassINET,
+					},
+				},
+				Additionals: []Resource{
+					{
+						mustEDNS0ResourceHeader(4096, 0xfe0|RCodeFormatError, true),
+						&OPTResource{},
+					},
+				},
+			},
+			dnssecOK: true,
+			extRCode: 0xfe0 | RCodeFormatError,
+		},
+		{
+			name: "with EDNS(0) options",
+			w: []byte{
+				0x00, 0x00, 0x29, 0x10, 0x00, 0xff, 0x00, 0x00,
+				0x00, 0x00, 0x0c, 0x00, 0x0c, 0x00, 0x02, 0x00,
+				0x00, 0x00, 0x0b, 0x00, 0x02, 0x12, 0x34,
+			},
+			m: Message{
+				Header: Header{RCode: RCodeServerFailure},
+				Questions: []Question{
+					{
+						Name:  mustNewName("."),
+						Type:  TypeAAAA,
+						Class: ClassINET,
+					},
+				},
+				Additionals: []Resource{
+					{
+						mustEDNS0ResourceHeader(4096, 0xff0|RCodeServerFailure, false),
+						&OPTResource{
+							Options: []Option{
+								{
+									Code: 12, // see RFC 7828
+									Data: []byte{0x00, 0x00},
+								},
+								{
+									Code: 11, // see RFC 7830
+									Data: []byte{0x12, 0x34},
+								},
+							},
+						},
+					},
+				},
+			},
+			dnssecOK: false,
+			extRCode: 0xff0 | RCodeServerFailure,
+		},
+		{
+			// Containing multiple OPT resources in a
+			// message is invalid, but it's necessary for
+			// protocol conformance testing.
+			name: "with multiple OPT resources",
+			w: []byte{
+				0x00, 0x00, 0x29, 0x10, 0x00, 0xff, 0x00, 0x00,
+				0x00, 0x00, 0x06, 0x00, 0x0b, 0x00, 0x02, 0x12,
+				0x34, 0x00, 0x00, 0x29, 0x10, 0x00, 0xff, 0x00,
+				0x00, 0x00, 0x00, 0x06, 0x00, 0x0c, 0x00, 0x02,
+				0x00, 0x00,
+			},
+			m: Message{
+				Header: Header{RCode: RCodeNameError},
+				Questions: []Question{
+					{
+						Name:  mustNewName("."),
+						Type:  TypeAAAA,
+						Class: ClassINET,
+					},
+				},
+				Additionals: []Resource{
+					{
+						mustEDNS0ResourceHeader(4096, 0xff0|RCodeNameError, false),
+						&OPTResource{
+							Options: []Option{
+								{
+									Code: 11, // see RFC 7830
+									Data: []byte{0x12, 0x34},
+								},
+							},
+						},
+					},
+					{
+						mustEDNS0ResourceHeader(4096, 0xff0|RCodeNameError, false),
+						&OPTResource{
+							Options: []Option{
+								{
+									Code: 12, // see RFC 7828
+									Data: []byte{0x00, 0x00},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	} {
+		w, err := tt.m.Pack()
+		if err != nil {
+			t.Errorf("Message.Pack() for %s = %v", tt.name, err)
+			continue
+		}
+		if !bytes.Equal(w[len(w)-len(tt.w):], tt.w) {
+			t.Errorf("got Message.Pack() for %s = %#v, want %#v", tt.name, w[len(w)-len(tt.w):], tt.w)
+			continue
+		}
+		var m Message
+		if err := m.Unpack(w); err != nil {
+			t.Errorf("Message.Unpack() for %s = %v", tt.name, err)
+			continue
+		}
+		if !reflect.DeepEqual(m.Additionals, tt.m.Additionals) {
+			t.Errorf("got Message.Pack/Unpack() roundtrip for %s = %+v, want %+v", tt.name, m, tt.m)
+			continue
+		}
+	}
+}
+
 func benchmarkParsingSetup() ([]byte, error) {
 	name := mustNewName("foo.bar.example.com.")
 	msg := Message{
@@ -837,6 +992,10 @@ func benchmarkParsing(tb testing.TB, buf []byte) {
 			if _, err := p.NSResource(); err != nil {
 				tb.Fatal("p.NSResource():", err)
 			}
+		case TypeOPT:
+			if _, err := p.OPTResource(); err != nil {
+				tb.Fatal("Parser.OPTResource() =", err)
+			}
 		default:
 			tb.Fatalf("unknown type: %T", h)
 		}
@@ -915,6 +1074,15 @@ func benchmarkBuilding(tb testing.TB, name Name, buf []byte) {
 		tb.Fatalf("bld.NSResource(%+v, %+v): %v", hdr, nsr, err)
 	}
 
+	extrc := 0xfe0 | RCodeNotImplemented
+	if err := (&hdr).SetEDNS0(4096, extrc, true); err != nil {
+		tb.Fatalf("ResourceHeader.SetEDNS0(4096, %#x, true) = %v", extrc, err)
+	}
+	optr := OPTResource{}
+	if err := bld.OPTResource(hdr, optr); err != nil {
+		tb.Fatalf("Builder.OPTResource(%+v, %+v) = %v", hdr, optr, err)
+	}
+
 	if _, err := bld.Finish(); err != nil {
 		tb.Fatal("bld.Finish():", err)
 	}
@@ -1132,6 +1300,17 @@ func largeTestMsg() Message {
 				},
 				&TXTResource{[]string{"Hamster Huey and the Gooey Kablooie"}},
 			},
+			{
+				mustEDNS0ResourceHeader(4096, 0xfe0|RCodeSuccess, false),
+				&OPTResource{
+					Options: []Option{
+						{
+							Code: 10, // see RFC 7873
+							Data: []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef},
+						},
+					},
+				},
+			},
 		},
 	}
 }