From b4d09be75101024ceed6b173b49a5630084174e6 Mon Sep 17 00:00:00 2001 From: Mateusz Poliwczak Date: Mon, 28 Aug 2023 16:40:48 +0000 Subject: [PATCH] dns/dnsmessage: compress all names while appending to a buffer Change-Id: Iedccbf3e47a63b2239def189ab41bab18a64c398 GitHub-Last-Rev: eb23195734794ab2b211677e5e3616de5f0eb7be GitHub-Pull-Request: golang/net#189 Reviewed-on: https://go-review.googlesource.com/c/net/+/522575 TryBot-Result: Gopher Robot Reviewed-by: Ian Lance Taylor Run-TryBot: Mateusz Poliwczak Reviewed-by: Joedian Reid Auto-Submit: Ian Lance Taylor --- dns/dnsmessage/message.go | 7 ++++--- dns/dnsmessage/message_test.go | 27 +++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/dns/dnsmessage/message.go b/dns/dnsmessage/message.go index 9ddf2c2292..0215a5dde3 100644 --- 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 index ee42febbc2..23fb3d5748 100644 --- 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) + } +}