From 07ae768ab11faef931d3e1523023ee1a5287292f Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Thu, 29 Nov 2018 09:49:18 +1030
Subject: [PATCH] Put escaped names into compression map in PackDomainName

---
 msg.go      | 29 ++++++++++-------------------
 msg_test.go | 10 +++++-----
 2 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/msg.go b/msg.go
index b1cc4e43..94159d03 100644
--- a/msg.go
+++ b/msg.go
@@ -223,9 +223,11 @@ func packDomainName(s string, msg []byte, off int, compression map[string]int, c
 
 	// Emit sequence of counted strings, chopping at dots.
 	var (
-		begin  int
-		bs     []byte
-		wasDot bool
+		begin     int
+		compBegin int
+		compOff   int
+		bs        []byte
+		wasDot    bool
 	)
 loop:
 	for i := 0; i < ls; i++ {
@@ -251,9 +253,11 @@ loop:
 				bs[i] = dddToByte(bs[i+1:])
 				copy(bs[i+1:ls-3], bs[i+4:])
 				ls -= 3
+				compOff += 3
 			} else {
 				copy(bs[i:ls-1], bs[i+1:])
 				ls--
+				compOff++
 			}
 
 			wasDot = false
@@ -279,17 +283,7 @@ loop:
 			// We should only compress when compress is true, but we should also still pick
 			// up names that can be used for *future* compression(s).
 			if compression != nil && !isRootLabel(s, bs, begin, ls) {
-				var (
-					p  int
-					ok bool
-				)
-				if bs == nil {
-					p, ok = compression[s[begin:]]
-				} else {
-					p, ok = compression[string(bs[begin:ls])]
-				}
-
-				if ok {
+				if p, ok := compression[s[compBegin:]]; ok {
 					// The first hit is the longest matching dname
 					// keep the pointer offset we get back and store
 					// the offset of the current name, because that's
@@ -302,11 +296,7 @@ loop:
 					}
 				} else if off < maxCompressionOffset {
 					// Only offsets smaller than maxCompressionOffset can be used.
-					if bs == nil {
-						compression[s[begin:]] = off
-					} else {
-						compression[string(bs[begin:ls])] = off
-					}
+					compression[s[compBegin:]] = off
 				}
 			}
 
@@ -324,6 +314,7 @@ loop:
 
 			labels++
 			begin = i + 1
+			compBegin = begin + compOff
 		default:
 			wasDot = false
 		}
diff --git a/msg_test.go b/msg_test.go
index c8ff41d3..70dca3c3 100644
--- a/msg_test.go
+++ b/msg_test.go
@@ -215,11 +215,11 @@ func TestUnpackDomainName(t *testing.T) {
 
 func TestPackDomainNameCompressionMap(t *testing.T) {
 	expected := map[string]struct{}{
-		`www.this.is.\131an.example.org.`: struct{}{},
-		`is.\131an.example.org.`:          struct{}{},
-		"\x83an.example.org.":             struct{}{},
-		`example.org.`:                    struct{}{},
-		`org.`:                            struct{}{},
+		`www\.this.is.\131an.example.org.`: struct{}{},
+		`is.\131an.example.org.`:           struct{}{},
+		`\131an.example.org.`:              struct{}{},
+		`example.org.`:                     struct{}{},
+		`org.`:                             struct{}{},
 	}
 
 	msg := make([]byte, 256)