From 1c9c9bf4c93ee029810272d3e5b8a126aee5bf1f Mon Sep 17 00:00:00 2001
From: chantra <chantra@users.noreply.github.com>
Date: Fri, 16 Nov 2018 16:00:14 -0800
Subject: [PATCH] properly set extended rcode when packing (#791)

* properly set extended rcode when packing

When calling `SetExtendedRcode`, we expect to get the full extended
rcode, not the rcode after we shift 4 bytes right.

* fix extended rcode

* fix TestOPTTtl test

* set error messages in TestPackExtendedBadCookie

* Set Rcode with extended rcode

* |=

* Set extended RCODE field to 0 when RCODE is not an extended one.
+ unittests

* Force setting extended rcode if we have an OPT available.

* go fmt + @tmthrgd comments

* comments and nits

* reformat comment
---
 edns.go      | 11 ++++----
 edns_test.go | 11 ++++++--
 msg.go       | 20 +++++++++------
 msg_test.go  | 71 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 98 insertions(+), 15 deletions(-)

diff --git a/edns.go b/edns.go
index 18d05413..0f6c4995 100644
--- a/edns.go
+++ b/edns.go
@@ -102,15 +102,14 @@ func (rr *OPT) SetVersion(v uint8) {
 
 // ExtendedRcode returns the EDNS extended RCODE field (the upper 8 bits of the TTL).
 func (rr *OPT) ExtendedRcode() int {
-	return int(rr.Hdr.Ttl&0xFF000000>>24) + 15
+	return int(rr.Hdr.Ttl&0xFF000000>>24) << 4
 }
 
 // SetExtendedRcode sets the EDNS extended RCODE field.
-func (rr *OPT) SetExtendedRcode(v uint8) {
-	if v < RcodeBadVers { // Smaller than 16.. Use the 4 bits you have!
-		return
-	}
-	rr.Hdr.Ttl = rr.Hdr.Ttl&0x00FFFFFF | uint32(v-15)<<24
+//
+// If the RCODE is not an extended RCODE, will reset the extended RCODE field to 0.
+func (rr *OPT) SetExtendedRcode(v uint16) {
+	rr.Hdr.Ttl = rr.Hdr.Ttl&0x00FFFFFF | uint32(v>>4)<<24
 }
 
 // UDPSize returns the UDP buffer size.
diff --git a/edns_test.go b/edns_test.go
index f7cf1575..6d27a4aa 100644
--- a/edns_test.go
+++ b/edns_test.go
@@ -62,7 +62,14 @@ func TestOPTTtl(t *testing.T) {
 	}
 
 	e.SetExtendedRcode(42)
-	if e.ExtendedRcode() != 42 {
-		t.Errorf("set 42, expected %d, got %d", 42, e.ExtendedRcode())
+	// ExtendedRcode has the last 4 bits set to 0.
+	if e.ExtendedRcode() != 42 & 0xFFFFFFF0 {
+		t.Errorf("set 42, expected %d, got %d", 42 & 0xFFFFFFF0, e.ExtendedRcode())
+	}
+
+	// This will reset the 8 upper bits of the extended rcode
+	e.SetExtendedRcode(RcodeNotAuth)
+	if e.ExtendedRcode() != 0 {
+		t.Errorf("Setting a non-extended rcode is expected to set extended rcode to 0, got: %d", e.ExtendedRcode())
 	}
 }
diff --git a/msg.go b/msg.go
index 47ac6cf2..b1367c1f 100644
--- a/msg.go
+++ b/msg.go
@@ -684,13 +684,14 @@ func (dns *Msg) packBufferWithCompressionMap(buf []byte, compression map[string]
 	if dns.Rcode < 0 || dns.Rcode > 0xFFF {
 		return nil, ErrRcode
 	}
-	if dns.Rcode > 0xF {
-		// Regular RCODE field is 4 bits
-		opt := dns.IsEdns0()
-		if opt == nil {
-			return nil, ErrExtendedRcode
-		}
-		opt.SetExtendedRcode(uint8(dns.Rcode >> 4))
+
+	// Set extended rcode unconditionally if we have an opt, this will allow
+	// reseting the extended rcode bits if they need to.
+	if opt := dns.IsEdns0(); opt != nil {
+		opt.SetExtendedRcode(uint16(dns.Rcode))
+	} else if dns.Rcode > 0xF {
+		// If Rcode is an extended one and opt is nil, error out.
+		return nil, ErrExtendedRcode
 	}
 
 	// Convert convenient Msg into wire-like Header.
@@ -837,6 +838,11 @@ func (dns *Msg) Unpack(msg []byte) (err error) {
 	// The header counts might have been wrong so we need to update it
 	dh.Arcount = uint16(len(dns.Extra))
 
+	// Set extended Rcode
+	if opt := dns.IsEdns0(); opt != nil {
+		dns.Rcode |= opt.ExtendedRcode()
+	}
+
 	if off != len(msg) {
 		// TODO(miek) make this an error?
 		// use PackOpt to let people tell how detailed the error reporting should be?
diff --git a/msg_test.go b/msg_test.go
index e9dce7ec..8a180018 100644
--- a/msg_test.go
+++ b/msg_test.go
@@ -45,6 +45,77 @@ func TestPackNoSideEffect(t *testing.T) {
 	}
 }
 
+func TestPackExtendedBadCookie(t *testing.T) {
+	m := new(Msg)
+	m.SetQuestion(Fqdn("example.com."), TypeNS)
+
+	a := new(Msg)
+	a.SetReply(m)
+	o := &OPT{
+		Hdr: RR_Header{
+			Name:   ".",
+			Rrtype: TypeOPT,
+		},
+	}
+	o.SetUDPSize(DefaultMsgSize)
+	a.Extra = append(a.Extra, o)
+
+	a.SetRcode(m, RcodeBadCookie)
+
+	edns0 := a.IsEdns0()
+	if edns0 == nil {
+		t.Fatal("Expected OPT RR")
+	}
+	// SetExtendedRcode is only called as part of `Pack()`, hence at this stage,
+	// the OPT RR is not set yet.
+	if edns0.ExtendedRcode() == RcodeBadCookie&0xFFFFFFF0 {
+		t.Errorf("ExtendedRcode is expected to not be BADCOOKIE before Pack")
+	}
+
+	a.Pack()
+
+	edns0 = a.IsEdns0()
+	if edns0 == nil {
+		t.Fatal("Expected OPT RR")
+	}
+
+	if edns0.ExtendedRcode() != RcodeBadCookie&0xFFFFFFF0 {
+		t.Errorf("ExtendedRcode is expected to be BADCOOKIE after Pack")
+	}
+}
+
+func TestUnPackExtendedRcode(t *testing.T) {
+	m := new(Msg)
+	m.SetQuestion(Fqdn("example.com."), TypeNS)
+
+	a := new(Msg)
+	a.SetReply(m)
+	o := &OPT{
+		Hdr: RR_Header{
+			Name:   ".",
+			Rrtype: TypeOPT,
+		},
+	}
+	o.SetUDPSize(DefaultMsgSize)
+	a.Extra = append(a.Extra, o)
+
+	a.SetRcode(m, RcodeBadCookie)
+
+	packed, err := a.Pack()
+	if err != nil {
+		t.Fatalf("Could not unpack %v", a)
+	}
+
+	unpacked := new(Msg)
+	if err := unpacked.Unpack(packed); err != nil {
+		t.Fatalf("Failed to unpack message")
+	}
+
+	if unpacked.Rcode != RcodeBadCookie {
+		t.Fatalf("Rcode should be matching RcodeBadCookie (%d), got (%d)", RcodeBadCookie, unpacked.Rcode)
+	}
+}
+
 func TestUnpackDomainName(t *testing.T) {
 	var cases = []struct {
 		label          string