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]>

Author: mateusz834

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
 			}
 		}
 	}

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)
+	}
+}