1 files changed, 122 insertions, 0 deletions
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
new file mode 100644
index 0000000..eaf0190
--- /dev/null
+++ b/internal/server/server_test.go
@@ -0,0 +1,122 @@
+package server
+
+import (
+	"testing"
+
+	"github.com/miekg/dns"
+)
+
+func TestBuildResponse(t *testing.T) {
+	tests := []struct {
+		name        string
+		req         *dns.Msg
+		wantRcode   int
+		wantAnswers int
+		wantEdns0   bool
+	}{
+		{
+			name: "example.com A returns 127.0.0.1",
+			req: func() *dns.Msg {
+				m := new(dns.Msg)
+				m.SetQuestion("example.com.", dns.TypeA)
+				return m
+			}(),
+			wantRcode:   dns.RcodeSuccess,
+			wantAnswers: 1,
+			wantEdns0:   false,
+		},
+		{
+			name: "google.com A returns NXDOMAIN",
+			req: func() *dns.Msg {
+				m := new(dns.Msg)
+				m.SetQuestion("google.com.", dns.TypeA)
+				return m
+			}(),
+			wantRcode:   dns.RcodeNameError,
+			wantAnswers: 0,
+			wantEdns0:   false,
+		},
+		{
+			name: "other.com A returns NXDOMAIN",
+			req: func() *dns.Msg {
+				m := new(dns.Msg)
+				m.SetQuestion("other.com.", dns.TypeA)
+				return m
+			}(),
+			wantRcode:   dns.RcodeNameError,
+			wantAnswers: 0,
+			wantEdns0:   false,
+		},
+		{
+			name: "no questions returns FORMERR",
+			req: func() *dns.Msg {
+				return new(dns.Msg)
+			}(),
+			wantRcode:   dns.RcodeFormatError,
+			wantAnswers: 0,
+			wantEdns0:   false,
+		},
+		{
+			name: "EDNS0 query preserved with 4096 buffer",
+			req: func() *dns.Msg {
+				m := new(dns.Msg)
+				m.SetQuestion("example.com.", dns.TypeA)
+				m.SetEdns0(1232, true)
+				return m
+			}(),
+			wantRcode:   dns.RcodeSuccess,
+			wantAnswers: 1,
+			wantEdns0:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			resp := buildResponse(tt.req)
+			if resp.Rcode != tt.wantRcode {
+				t.Errorf("rcode: got %d, want %d", resp.Rcode, tt.wantRcode)
+			}
+			if len(resp.Answer) != tt.wantAnswers {
+				t.Errorf("answers: got %d, want %d", len(resp.Answer), tt.wantAnswers)
+			}
+			if tt.wantEdns0 {
+				if opt := resp.IsEdns0(); opt == nil {
+					t.Error("expected EDNS0 in response, got none")
+				} else if opt.UDPSize() != 4096 {
+					t.Errorf("edns0 udp size: got %d, want 4096", opt.UDPSize())
+				}
+			}
+		})
+	}
+}
+
+func FuzzBuildResponse(f *testing.F) {
+	seed := []byte{
+		0x00, 0x00, // ID
+		0x01, 0x00, // flags: RD
+		0x00, 0x01, // QDCOUNT: 1
+		0x00, 0x00, // ANCOUNT
+		0x00, 0x00, // NSCOUNT
+		0x00, 0x00, // ARCOUNT
+		// Question: example.com A
+		0x07, 0x65, 0x78, 0x61, 0x6d, 0x70, 0x6c, 0x65,
+		0x03, 0x63, 0x6f, 0x6d,
+		0x00,
+		0x00, 0x01, // QTYPE: A
+		0x00, 0x01, // QCLASS: IN
+	}
+	f.Add(seed)
+	f.Fuzz(func(t *testing.T, data []byte) {
+		msg := new(dns.Msg)
+		if err := msg.Unpack(data); err != nil {
+			return
+		}
+		resp := buildResponse(msg)
+		if resp == nil {
+			t.Fatal("buildResponse returned nil")
+		}
+		if _, err := resp.Pack(); err != nil {
+			t.Errorf("pack failed: %v", err)
+		}
+	})
+}