Merge pull request #842 from tmthrgd/compression-map-escaped

Put escaped names into compression map

length_test.go

@@ -382,3 +382,26 @@ func TestMsgCompressLengthLargeRecordsAllValues(t *testing.T) {
 		}
 	}
 }
+
+func TestMsgCompressLengthEscapingMatch(t *testing.T) {
+	// Although slightly non-optimal, "example.org." and "ex\\097mple.org."
+	// are not considered equal in the compression map, even though \097 is
+	// a valid escaping of a. This test ensures that the Len code and the
+	// Pack code don't disagree on this.
+
+	msg := new(Msg)
+	msg.Compress = true
+	msg.SetQuestion("www.example.org.", TypeA)
+	msg.Answer = append(msg.Answer, &NS{Hdr: RR_Header{Name: "ex\\097mple.org.", Rrtype: TypeNS, Class: ClassINET}, Ns: "ns.example.org."})
+
+	predicted := msg.Len()
+	buf, err := msg.Pack()
+	if err != nil {
+		t.Error(err)
+	}
+	// Len doesn't account for escaping when calculating the length *yet* so
+	// we're off by three here. This will be fixed in a follow up change.
+	if predicted != len(buf)+3 {
+		t.Fatalf("predicted compressed length is wrong: predicted %d, actual %d", predicted, len(buf))
+	}
+}

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
 		}

msg_test.go

@@ -214,8 +214,15 @@ func TestUnpackDomainName(t *testing.T) {
 }
 
 func TestPackDomainNameCompressionMap(t *testing.T) {
-	msg := make([]byte, 256)
+	expected := map[string]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)
 	for _, compress := range []bool{true, false} {
 		compression := make(map[string]int)
 
@@ -224,16 +231,8 @@ func TestPackDomainNameCompressionMap(t *testing.T) {
 			t.Fatalf("PackDomainName failed: %v", err)
 		}
 
-		for _, dname := range []string{
-			`www.this.is.\131an.example.org.`,
-			`is.\131an.example.org.`,
-			"\x83an.example.org.",
-			`example.org.`,
-			`org.`,
-		} {
-			if _, ok := compression[dname]; !ok {
-				t.Errorf("expected to find %q in compression map", dname)
-			}
+		if !compressionMapsEqual(expected, compression) {
+			t.Errorf("expected compression maps to be equal; expected %v, got %v", expected, compression)
 		}
 	}
 }