dns/dnsmessage: compress all names while appending to a buffer

Change-Id: Iedccbf3e47a63b2239def189ab41bab18a64c398 GitHub-Last-Rev: eb23195 GitHub-Pull-Request: #189 Reviewed-on: https://go-review.googlesource.com/c/net/+/522575 TryBot-Result: Gopher Robot <[email protected]> Reviewed-by: Ian Lance Taylor <[email protected]> Run-TryBot: Mateusz Poliwczak <[email protected]> Reviewed-by: Joedian Reid <[email protected]> Auto-Submit: Ian Lance Taylor <[email protected]>
golang · Aug 29, 2023 · b4d09be · b4d09be
1 parent 8b010a5
commit b4d09be
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/dns/dnsmessage/message.go b/dns/dnsmessage/message.go
@@ -2000,14 +2000,15 @@ func (n *Name) pack(msg []byte, compression map[string]int, compressionOff int)
 			}
 
 			// Miss. Add the suffix to the compression table if the
-			// offset can be stored in the available 14 bytes.
-			if len(msg) <= int(^uint16(0)>>2) {
+			// offset can be stored in the available 14 bits.
+			newPtr := len(msg) - compressionOff
+			if newPtr <= int(^uint16(0)>>2) {
 				if nameAsStr == "" {
 					// allocate n.Data on the heap once, to avoid allocating it
 					// multiple times (for next labels).
 					nameAsStr = string(n.Data[:n.Length])
 				}
-				compression[nameAsStr[i:]] = len(msg) - compressionOff
+				compression[nameAsStr[i:]] = newPtr
 			}
 		}
 	}

diff --git a/dns/dnsmessage/message_test.go b/dns/dnsmessage/message_test.go
@@ -1847,3 +1847,30 @@ func TestBuilderNameCompressionWithNonZeroedName(t *testing.T) {
 		t.Fatalf("b.Finish() = %v, want: %v", msg, expect)
 	}
 }
+
+func TestBuilderCompressionInAppendMode(t *testing.T) {
+	maxPtr := int(^uint16(0) >> 2)
+	b := NewBuilder(make([]byte, maxPtr, maxPtr+512), Header{})
+	b.EnableCompression()
+	if err := b.StartQuestions(); err != nil {
+		t.Fatalf("b.StartQuestions() unexpected error: %v", err)
+	}
+	if err := b.Question(Question{Name: MustNewName("go.dev.")}); err != nil {
+		t.Fatalf("b.Question() unexpected error: %v", err)
+	}
+	if err := b.Question(Question{Name: MustNewName("go.dev.")}); err != nil {
+		t.Fatalf("b.Question() unexpected error: %v", err)
+	}
+	msg, err := b.Finish()
+	if err != nil {
+		t.Fatalf("b.Finish() unexpected error: %v", err)
+	}
+	expect := []byte{
+		0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, // header
+		2, 'g', 'o', 3, 'd', 'e', 'v', 0, 0, 0, 0, 0, // question 1
+		0xC0, 12, 0, 0, 0, 0, // question 2
+	}
+	if !bytes.Equal(msg[maxPtr:], expect) {
+		t.Fatalf("msg[maxPtr:] = %v, want: %v", msg[maxPtr:], expect)
+	}
+}