Put escaped names into compression map in PackDomainName

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
+		begin     int
-		bs     []byte
+		compBegin int
-		wasDot bool
+		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 (
+				if p, ok := compression[s[compBegin:]]; ok {
-					p  int
-					ok bool
-				)
-				if bs == nil {
-					p, ok = compression[s[begin:]]
-				} else {
-					p, ok = compression[string(bs[begin:ls])]
-				}
-
-				if 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[compBegin:]] = off
-						compression[s[begin:]] = off
-					} else {
-						compression[string(bs[begin:ls])] = off
-					}
 				}
 			}
 
@@ -324,6 +314,7 @@ loop:
 
 			labels++
 			begin = i + 1
+			compBegin = begin + compOff
 		default:
 			wasDot = false
 		}

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{}{},
+		`www\.this.is.\131an.example.org.`: struct{}{},
-		`is.\131an.example.org.`:          struct{}{},
+		`is.\131an.example.org.`:           struct{}{},
-		"\x83an.example.org.":             struct{}{},
+		`\131an.example.org.`:              struct{}{},
-		`example.org.`:                    struct{}{},
+		`example.org.`:                     struct{}{},
-		`org.`:                            struct{}{},
+		`org.`:                             struct{}{},
 	}
 
 	msg := make([]byte, 256)