git » dnss » commit 1416dd9

diff --git a/internal/dnsserver/server.go b/internal/dnsserver/server.go
index 7904582..bda16b9 100644
--- a/internal/dnsserver/server.go
+++ b/internal/dnsserver/server.go
@@ -135,20 +135,30 @@ func (s *Server) Handler(w dns.ResponseWriter, r *dns.Msg) {
 }
 
 func (s *Server) writeReply(tr *trace.Trace, w dns.ResponseWriter, r, reply *dns.Msg) {
-	if w.RemoteAddr().Network() == "udp" {
-		// We need to check if the response fits.
-		// UDP by default has a maximum of 512 bytes. This can be extended via
-		// the client in the EDNS0 record.
-		max := 512
+	// We need to check if the response fits.
+	// UDP by default has a maximum of 512 bytes. This can be extended via
+	// the client in the EDNS0 record. TCP has a maximum of 64 KiB.
+	network := w.RemoteAddr().Network()
+	max := dns.MaxMsgSize
+	if network == "udp" {
+		max = 512
 		ednsOPT := r.IsEdns0()
 		if ednsOPT != nil {
 			max = int(ednsOPT.UDPSize())
 		}
-		reply.Truncate(max)
-		tr.Printf("UDP max:%d truncated:%v", max, reply.Truncated)
 	}
+	reply.Truncate(max)
+	tr.Printf("%s max:%d truncated:%v", network, max, reply.Truncated)
 
-	w.WriteMsg(reply)
+	// Always compress the reply. Messages are unpacked without compression,
+	// and packing them back as-is can make them much larger, even beyond what
+	// TCP can carry.
+	// Note this must be done after Truncate, which resets it.
+	reply.Compress = true
+
+	if err := w.WriteMsg(reply); err != nil {
+		tr.Errorf("error writing reply: %v", err)
+	}
 }
 
 // ListenAndServe launches the DNS proxy.
diff --git a/internal/dnsserver/server_test.go b/internal/dnsserver/server_test.go
index 2a939c5..2fafb65 100644
--- a/internal/dnsserver/server_test.go
+++ b/internal/dnsserver/server_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/miekg/dns"
 
 	"blitiri.com.ar/go/dnss/internal/testutil"
+	"blitiri.com.ar/go/dnss/internal/trace"
 )
 
 // Tests for the DNS server.
@@ -94,3 +95,48 @@ func queryFailure(t *testing.T, srv, domain string) {
 		t.Errorf("query %q: expected SERVFAIL, got message: %v", domain, m)
 	}
 }
+
+// largeResolver answers all queries with testutil.NewLargeAnswer.
+// We can't use testutil.TestResolver because it returns the same message on
+// every query, and the server modifies it when truncating.
+type largeResolver struct{}
+
+func (largeResolver) Init() error { return nil }
+func (largeResolver) Maintain()   {}
+func (largeResolver) Query(r *dns.Msg, tr *trace.Trace) (*dns.Msg, error) {
+	return &dns.Msg{
+		Question: r.Question,
+		Answer:   testutil.NewLargeAnswer(300),
+	}, nil
+}
+
+// Test that large answers, which only fit when compressed, are delivered.
+func TestLargeAnswer(t *testing.T) {
+	srv := New(testutil.GetFreePort(), largeResolver{}, "", nil)
+	go srv.ListenAndServe()
+	testutil.WaitForDNSServer(srv.Addr)
+
+	m := &dns.Msg{}
+	m.SetQuestion(testutil.LargeAnswerName, dns.TypeA)
+
+	// Over TCP, we should get the full answer.
+	c := &dns.Client{Net: "tcp"}
+	in, _, err := c.Exchange(m, srv.Addr)
+	if err != nil {
+		t.Fatalf("TCP query failed: %v", err)
+	}
+	if in.Truncated || len(in.Answer) != 300 {
+		t.Errorf("TCP: expected 300 RRs, got %d (truncated:%v)",
+			len(in.Answer), in.Truncated)
+	}
+
+	// Over UDP without EDNS0, we should get a truncated answer.
+	in, err = dns.Exchange(m, srv.Addr)
+	if err != nil {
+		t.Fatalf("UDP query failed: %v", err)
+	}
+	if !in.Truncated || len(in.Answer) >= 300 {
+		t.Errorf("UDP: expected truncated answer, got %d RRs (truncated:%v)",
+			len(in.Answer), in.Truncated)
+	}
+}
diff --git a/internal/httpserver/server.go b/internal/httpserver/server.go
index 380445f..52a58c3 100644
--- a/internal/httpserver/server.go
+++ b/internal/httpserver/server.go
@@ -130,6 +130,11 @@ func (s *Server) resolveDoH(tr *trace.Trace, w http.ResponseWriter, dnsQuery []b
 
 	tr.Answer(fromUp)
 
+	// Compress the reply. Messages are unpacked without compression, and
+	// packing them back as-is can make them much larger, even beyond the
+	// DoH limit of 64 KiB.
+	fromUp.Compress = true
+
 	packed, err := fromUp.Pack()
 	if err != nil {
 		err = tr.Errorf("cannot pack reply: %v", err)
diff --git a/internal/httpserver/server_test.go b/internal/httpserver/server_test.go
index 40e019d..7caf7da 100644
--- a/internal/httpserver/server_test.go
+++ b/internal/httpserver/server_test.go
@@ -2,12 +2,15 @@ package httpserver
 
 import (
 	"errors"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
 	"blitiri.com.ar/go/dnss/internal/testutil"
+
+	"github.com/miekg/dns"
 )
 
 func TestBasic(t *testing.T) {
@@ -61,6 +64,57 @@ func TestBasic(t *testing.T) {
 	}
 }
 
+// Test that large answers, which only fit when compressed, are delivered.
+func TestLargeAnswer(t *testing.T) {
+	upstreamAddr := testutil.GetFreePort()
+	go testutil.ServeTestDNSServer(upstreamAddr,
+		func(w dns.ResponseWriter, r *dns.Msg) {
+			m := &dns.Msg{}
+			m.SetReply(r)
+			if r.Question[0].Name == testutil.LargeAnswerName {
+				m.Answer = testutil.NewLargeAnswer(300)
+			}
+			m.Compress = true
+			w.WriteMsg(m)
+		})
+	testutil.WaitForDNSServer(upstreamAddr)
+
+	srv := &Server{
+		Upstream: upstreamAddr,
+	}
+
+	// Use EDNS0 so the upstream query can get the whole answer over UDP,
+	// as our test server doesn't support TCP.
+	m := &dns.Msg{}
+	m.SetQuestion(testutil.LargeAnswerName, dns.TypeA)
+	m.SetEdns0(dns.MaxMsgSize, false)
+	packed, err := m.Pack()
+	if err != nil {
+		t.Fatalf("error packing query: %v", err)
+	}
+
+	resp := query(t, srv, "POST", "/ignored", string(packed))
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("expected http status ok, got %v", resp.StatusCode)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("error reading body: %v", err)
+	}
+	if len(body) > dns.MaxMsgSize {
+		t.Errorf("reply is too large: %d bytes", len(body))
+	}
+
+	reply := &dns.Msg{}
+	if err := reply.Unpack(body); err != nil {
+		t.Fatalf("error unpacking reply: %v", err)
+	}
+	if len(reply.Answer) != 300 {
+		t.Errorf("expected 300 RRs, got %d", len(reply.Answer))
+	}
+}
+
 func query(t *testing.T, srv *Server, method, url, body string) *http.Response {
 	t.Helper()
 
diff --git a/internal/testutil/testutil.go b/internal/testutil/testutil.go
index 80fdb84..11c2ca6 100644
--- a/internal/testutil/testutil.go
+++ b/internal/testutil/testutil.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 
@@ -163,3 +164,27 @@ func NewRR(tb testing.TB, s string) dns.RR {
 	}
 	return rr
 }
+
+// LargeAnswerName is the name used in the records from NewLargeAnswer.
+// It is as long as a name can be, so the answers compress very well.
+var LargeAnswerName = strings.Repeat(strings.Repeat("x", 62)+".", 4)
+
+// NewLargeAnswer returns n A records for LargeAnswerName. Each takes ~270
+// bytes uncompressed, but only ~16 bytes compressed. This is useful to test
+// compression: for example, 300 of them take ~80 KiB uncompressed (more than
+// the 64 KiB DNS limit), but only ~5 KiB compressed.
+func NewLargeAnswer(n int) []dns.RR {
+	rrs := []dns.RR{}
+	for i := 0; i < n; i++ {
+		rrs = append(rrs, &dns.A{
+			Hdr: dns.RR_Header{
+				Name:   LargeAnswerName,
+				Rrtype: dns.TypeA,
+				Class:  dns.ClassINET,
+				Ttl:    3600,
+			},
+			A: net.IPv4(10, 0, byte(i>>8), byte(i)).To4(),
+		})
+	}
+	return rrs
+}