fix: harden dns proxy routing

router/dns.go

@@ -13,6 +13,14 @@ import (
 
 var errNoLocalIPForQuestionType = errors.New("no local IP for question type")
 
+type retryableDNSError struct {
+	rcode int
+}
+
+func (e retryableDNSError) Error() string {
+	return fmt.Sprintf("retryable dns response code %s", dns.RcodeToString[e.rcode])
+}
+
 func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
 	// https://stackoverflow.com/questions/4082081/requesting-a-and-aaaa-records-in-single-dns-query/4083071#4083071
 	if len(req.Question) != 1 {
@@ -22,13 +30,19 @@ func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
 
 	domain := req.Question[0].Name
 	qtype := req.Question[0].Qtype
+	if req.Question[0].Qclass != dns.ClassINET {
+		_ = w.WriteMsg(r.dnsFail(req, dns.RcodeNotImplemented))
+		return
+	}
+
 	// 1. rule_based( block > direct > proxy )
-	if r.BlockRule.Match(domain) {
+	routeDomains := dnsRouteDomains(domain, qtype)
+	if dnsRuleMatch(r.BlockRule, routeDomains) {
 		_ = w.WriteMsg(r.dnsFail(req, dns.RcodeNameError))
 		return
 	}
 
-	if !r.DirectRule.Match(domain) && r.ProxyRule.Match(domain) {
+	if !dnsRuleMatch(r.DirectRule, routeDomains) && dnsRuleMatch(r.ProxyRule, routeDomains) {
 		switch {
 		case isAddressQuestion(qtype):
 			resp, err := r.dnsProxyReply(domain, w.LocalAddr(), req)
@@ -39,7 +53,7 @@ func (r *Router) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
 			}
 			_ = w.WriteMsg(resp)
 			return
-		case isProxySensitiveQuestion(qtype):
+		default:
 			_ = w.WriteMsg(r.dnsNoData(req))
 			return
 		}
@@ -85,28 +99,35 @@ func (r *Router) dnsProxyReply(domain string, localAddr net.Addr, req *dns.Msg)
 }
 
 func (r *Router) proxyReplyIP(localAddr net.Addr, qtype uint16) (net.IP, error) {
-	host, _, err := net.SplitHostPort(localAddr.String())
-	if err != nil {
-		return nil, fmt.Errorf("split local dns address %q: %w", localAddr.String(), err)
-	}
-
-	ip := net.ParseIP(strings.Trim(host, "[]"))
-	if ip != nil && !ip.IsUnspecified() && matchesQuestionTypeIP(ip, qtype) {
-		return ip, nil
-	}
-	if ip == nil {
-		return nil, fmt.Errorf("parse local dns IP %q", host)
-	}
 	if !isAddressQuestion(qtype) {
 		return nil, fmt.Errorf("unsupported question type %d", qtype)
 	}
 
+	var localAddrErr error
+	if localAddr != nil {
+		host, _, err := net.SplitHostPort(localAddr.String())
+		if err != nil {
+			localAddrErr = fmt.Errorf("split local dns address %q: %w", localAddr.String(), err)
+		} else {
+			ip := net.ParseIP(strings.Trim(host, "[]"))
+			switch {
+			case ip == nil:
+				localAddrErr = fmt.Errorf("parse local dns IP %q", host)
+			case !ip.IsUnspecified() && matchesQuestionTypeIP(ip, qtype):
+				return ip, nil
+			}
+		}
+	}
+
 	for _, serveIP := range r.dns.serveIPs {
 		if !serveIP.IsUnspecified() && matchesQuestionTypeIP(serveIP, qtype) {
 			return serveIP, nil
 		}
 	}
 
+	if localAddrErr != nil {
+		return nil, localAddrErr
+	}
 	return nil, fmt.Errorf("%w %d", errNoLocalIPForQuestionType, qtype)
 }
 
@@ -114,8 +135,43 @@ func isAddressQuestion(qtype uint16) bool {
 	return qtype == dns.TypeA || qtype == dns.TypeAAAA
 }
 
-func isProxySensitiveQuestion(qtype uint16) bool {
-	return qtype == dns.TypeHTTPS || qtype == dns.TypeSVCB
+func dnsRuleMatch(rule interface{ Match(string) bool }, domains []string) bool {
+	for _, domain := range domains {
+		if rule.Match(domain) {
+			return true
+		}
+	}
+	return false
+}
+
+func dnsRouteDomains(domain string, qtype uint16) []string {
+	if !isServiceDiscoveryQuestion(qtype) {
+		return []string{domain}
+	}
+
+	base, ok := stripDNSServicePrefix(domain)
+	if !ok {
+		return []string{domain}
+	}
+	return []string{domain, base}
+}
+
+func isServiceDiscoveryQuestion(qtype uint16) bool {
+	switch qtype {
+	case dns.TypeSRV, dns.TypeSVCB, dns.TypeHTTPS:
+		return true
+	default:
+		return false
+	}
+}
+
+func stripDNSServicePrefix(domain string) (string, bool) {
+	trimmed := strings.TrimSuffix(domain, ".")
+	labels := strings.Split(trimmed, ".")
+	if len(labels) <= 2 || !strings.HasPrefix(labels[0], "_") || !strings.HasPrefix(labels[1], "_") {
+		return "", false
+	}
+	return strings.Join(labels[2:], ".") + ".", true
 }
 
 func dnsReply(req *dns.Msg) *dns.Msg {
@@ -183,9 +239,17 @@ func (r *Router) Exchange(req *dns.Msg) (_ *dns.Msg, err error) {
 		}
 	}
 
+	if resp != nil && isRetryableDNSResponseErr(err) {
+		return resp, nil
+	}
 	return resp, err
 }
 
+func isRetryableDNSResponseErr(err error) bool {
+	var retryableErr retryableDNSError
+	return errors.As(err, &retryableErr)
+}
+
 func (r *Router) currentUpstreamState(now time.Time) ([]string, int, bool, error) {
 	for {
 		addrs, index, shouldProbe, needRefresh, err := r.prepareUpstreamState(now)
@@ -383,10 +447,26 @@ func (r *Router) exchangeWithRetry(req *dns.Msg, addr string) (*dns.Msg, error)
 		return nil, err
 	}
 	if !resp.Truncated {
-		return resp, nil
+		return resp, dnsResponseErr(resp)
 	}
 
-	return r.exchangeTCP(req, addr)
+	resp, err = r.exchangeTCP(req, addr)
+	if err != nil {
+		return nil, err
+	}
+	return resp, dnsResponseErr(resp)
+}
+
+func dnsResponseErr(resp *dns.Msg) error {
+	if resp == nil {
+		return fmt.Errorf("empty dns response")
+	}
+	switch resp.Rcode {
+	case dns.RcodeServerFailure, dns.RcodeRefused, dns.RcodeNotImplemented:
+		return retryableDNSError{rcode: resp.Rcode}
+	default:
+		return nil
+	}
 }
 
 func (r *Router) exchangeUDP(req *dns.Msg, addr string) (*dns.Msg, error) {

router/router.go

@@ -52,7 +52,10 @@ var ErrBlocked = errors.New("route blocked")
 
 func NewRouter(serveIPs []string, upstreamDNS, fallbackDNS, mmdbFile string, proxyDial ProxyDialFn) *Router {
 	r := Router{
-		ProxyDial: proxyDial,
+		BlockRule:  suffixtree.NewNodeFromRules(),
+		DirectRule: suffixtree.NewNodeFromRules(),
+		ProxyRule:  suffixtree.NewNodeFromRules(),
+		ProxyDial:  proxyDial,
 	}
 
 	r.dns.upstreamDNS = upstreamDNS

router/router_test.go

@@ -24,6 +24,15 @@ func (fn roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
 	return fn(req)
 }
 
+func TestNewRouterInitializesRules(t *testing.T) {
+	t.Parallel()
+
+	r := NewRouter(nil, "", "223.5.5.5", "", nil)
+	if r.BlockRule == nil || r.DirectRule == nil || r.ProxyRule == nil {
+		t.Fatal("expected router rules to be initialized")
+	}
+}
+
 func TestDNSProxyReplyMatchesQuestionType(t *testing.T) {
 	t.Parallel()
 
@@ -98,6 +107,27 @@ func TestDNSProxyReplyIgnoresUnspecifiedLocalAddr(t *testing.T) {
 	}
 }
 
+func TestDNSProxyReplyFallsBackWhenLocalAddrMissing(t *testing.T) {
+	t.Parallel()
+
+	r := NewRouter([]string{"127.0.0.2"}, "", "223.5.5.5", "", nil)
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+
+	resp, err := r.dnsProxyReply("example.com.", nil, req)
+	if err != nil {
+		t.Fatalf("dnsProxyReply failed: %v", err)
+	}
+
+	answer, ok := resp.Answer[0].(*dns.A)
+	if !ok {
+		t.Fatalf("expected A record, got %T", resp.Answer[0])
+	}
+	if !answer.A.Equal(net.ParseIP("127.0.0.2")) {
+		t.Fatalf("expected configured serve IP 127.0.0.2, got %s", answer.A)
+	}
+}
+
 func TestExchangeSkipsServeIPInUpstreamList(t *testing.T) {
 	t.Parallel()
 
@@ -130,6 +160,153 @@ func TestBuildUpstreamAddrsPrefersConfiguredDNSBeforeFallback(t *testing.T) {
 	}
 }
 
+func TestExchangeFallsBackOnRetryableDNSResponse(t *testing.T) {
+	t.Parallel()
+
+	primaryAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := new(dns.Msg)
+		resp.SetRcode(req, dns.RcodeServerFailure)
+		_ = w.WriteMsg(resp)
+	}))
+	fallbackAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := new(dns.Msg)
+		resp.SetReply(req)
+		resp.Answer = []dns.RR{&dns.A{
+			Hdr: dns.RR_Header{Name: req.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 20},
+			A:   net.ParseIP("203.0.113.10"),
+		}}
+		_ = w.WriteMsg(resp)
+	}))
+
+	r := NewRouter(nil, "", "", "", nil)
+	r.dns.upstreamAddrs = []string{primaryAddr, fallbackAddr}
+	r.dns.upstreamIndex = 0
+
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+	resp, err := r.Exchange(req)
+	if err != nil {
+		t.Fatalf("exchange failed: %v", err)
+	}
+	if r.dns.upstreamIndex != 1 {
+		t.Fatalf("expected degraded upstream index 1, got %d", r.dns.upstreamIndex)
+	}
+	answer, ok := resp.Answer[0].(*dns.A)
+	if !ok {
+		t.Fatalf("expected fallback A answer, got %T", resp.Answer[0])
+	}
+	if !answer.A.Equal(net.ParseIP("203.0.113.10")) {
+		t.Fatalf("unexpected fallback answer: %s", answer.A)
+	}
+}
+
+func TestExchangeDoesNotFallbackOnTerminalDNSResponse(t *testing.T) {
+	t.Parallel()
+
+	var fallbackQueries atomic.Int32
+	primaryAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := new(dns.Msg)
+		resp.SetRcode(req, dns.RcodeNameError)
+		_ = w.WriteMsg(resp)
+	}))
+	fallbackAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		fallbackQueries.Add(1)
+		resp := new(dns.Msg)
+		resp.SetReply(req)
+		_ = w.WriteMsg(resp)
+	}))
+
+	r := NewRouter(nil, "", "", "", nil)
+	r.dns.upstreamAddrs = []string{primaryAddr, fallbackAddr}
+	r.dns.upstreamIndex = 0
+
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+	resp, err := r.Exchange(req)
+	if err != nil {
+		t.Fatalf("exchange failed: %v", err)
+	}
+	if resp.Rcode != dns.RcodeNameError {
+		t.Fatalf("expected NXDOMAIN response, got rcode %d", resp.Rcode)
+	}
+	if r.dns.upstreamIndex != 0 {
+		t.Fatalf("expected preferred upstream to remain selected, got %d", r.dns.upstreamIndex)
+	}
+	if fallbackQueries.Load() != 0 {
+		t.Fatalf("expected fallback to remain unused, got %d queries", fallbackQueries.Load())
+	}
+}
+
+func TestServeDNSReturnsRetryableResponseWhenUpstreamsExhausted(t *testing.T) {
+	t.Parallel()
+
+	upstreamAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := new(dns.Msg)
+		resp.SetRcode(req, dns.RcodeRefused)
+		_ = w.WriteMsg(resp)
+	}))
+
+	r := NewRouter(nil, "", "", "", nil)
+	r.dns.upstreamAddrs = []string{upstreamAddr}
+	r.dns.upstreamIndex = 0
+
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+	writer := &mockDNSWriter{localAddr: &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 53}}
+
+	r.ServeDNS(writer, req)
+
+	if writer.msg == nil {
+		t.Fatal("expected response message")
+	}
+	if writer.msg.Rcode != dns.RcodeRefused {
+		t.Fatalf("expected upstream REFUSED response, got rcode %d", writer.msg.Rcode)
+	}
+}
+
+func TestExchangeDoesNotPromoteRetryableProbeResponse(t *testing.T) {
+	t.Parallel()
+
+	var primaryQueries atomic.Int32
+	var fallbackQueries atomic.Int32
+	primaryAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		primaryQueries.Add(1)
+		resp := new(dns.Msg)
+		resp.SetRcode(req, dns.RcodeRefused)
+		_ = w.WriteMsg(resp)
+	}))
+	fallbackAddr := startUDPTestDNSServer(t, dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		fallbackQueries.Add(1)
+		resp := new(dns.Msg)
+		resp.SetReply(req)
+		resp.Answer = []dns.RR{&dns.A{
+			Hdr: dns.RR_Header{Name: req.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 20},
+			A:   net.ParseIP("203.0.113.20"),
+		}}
+		_ = w.WriteMsg(resp)
+	}))
+
+	r := NewRouter(nil, "", "", "", nil)
+	r.dns.upstreamAddrs = []string{primaryAddr, fallbackAddr}
+	r.dns.upstreamIndex = 1
+	r.dns.retryAt = time.Now().Add(-time.Second)
+
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+	if _, err := r.Exchange(req); err != nil {
+		t.Fatalf("exchange failed: %v", err)
+	}
+	if r.dns.upstreamIndex != 1 {
+		t.Fatalf("expected retryable probe response to keep fallback index 1, got %d", r.dns.upstreamIndex)
+	}
+	if primaryQueries.Load() != 1 {
+		t.Fatalf("expected one preferred probe, got %d", primaryQueries.Load())
+	}
+	if fallbackQueries.Load() != 1 {
+		t.Fatalf("expected one fallback query, got %d", fallbackQueries.Load())
+	}
+}
+
 func TestAddCountryCIDRsSkipsInvalidEntries(t *testing.T) {
 	t.Parallel()
 
@@ -195,7 +372,7 @@ func TestIsAccessChecksRequestedPortOnly(t *testing.T) {
 	}
 }
 
-func TestServeDNSForProxyRuleForwardsNonAddressQuestion(t *testing.T) {
+func TestServeDNSForProxyRuleSuppressesNonAddressQuestion(t *testing.T) {
 	t.Parallel()
 
 	udpConn, err := net.ListenPacket("udp", "127.0.0.1:0")
@@ -204,7 +381,9 @@ func TestServeDNSForProxyRuleForwardsNonAddressQuestion(t *testing.T) {
 	}
 	defer udpConn.Close()
 
+	var upstreamQueries atomic.Int32
 	handler := dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		upstreamQueries.Add(1)
 		resp := new(dns.Msg)
 		resp.SetReply(req)
 		resp.Answer = []dns.RR{&dns.TXT{
@@ -239,15 +418,11 @@ func TestServeDNSForProxyRuleForwardsNonAddressQuestion(t *testing.T) {
 	if writer.msg.Rcode != dns.RcodeSuccess {
 		t.Fatalf("expected successful response, got rcode %d", writer.msg.Rcode)
 	}
-	if len(writer.msg.Answer) != 1 {
-		t.Fatalf("expected forwarded TXT answer, got %d answers", len(writer.msg.Answer))
+	if len(writer.msg.Answer) != 0 {
+		t.Fatalf("expected NODATA response, got %d answers", len(writer.msg.Answer))
 	}
-	answer, ok := writer.msg.Answer[0].(*dns.TXT)
-	if !ok {
-		t.Fatalf("expected TXT record, got %T", writer.msg.Answer[0])
-	}
-	if len(answer.Txt) != 1 || answer.Txt[0] != "forwarded" {
-		t.Fatalf("unexpected TXT payload: %v", answer.Txt)
+	if upstreamQueries.Load() != 0 {
+		t.Fatalf("expected TXT query to be suppressed locally, got %d upstream queries", upstreamQueries.Load())
 	}
 }
 
@@ -301,7 +476,7 @@ func TestServeDNSForProxyRuleSuppressesHTTPSQuestion(t *testing.T) {
 	}
 }
 
-func TestServeDNSForProxyRuleForwardsSRVQuestion(t *testing.T) {
+func TestServeDNSForProxyRuleSuppressesSRVQuestion(t *testing.T) {
 	t.Parallel()
 
 	udpConn, err := net.ListenPacket("udp", "127.0.0.1:0")
@@ -344,6 +519,60 @@ func TestServeDNSForProxyRuleForwardsSRVQuestion(t *testing.T) {
 
 	r.ServeDNS(writer, req)
 
+	if writer.msg == nil {
+		t.Fatal("expected response message")
+	}
+	if writer.msg.Rcode != dns.RcodeSuccess {
+		t.Fatalf("expected successful response, got rcode %d", writer.msg.Rcode)
+	}
+	if len(writer.msg.Answer) != 0 {
+		t.Fatalf("expected NODATA response, got %d answers", len(writer.msg.Answer))
+	}
+	if upstreamQueries.Load() != 0 {
+		t.Fatalf("expected SRV query to be suppressed locally, got %d upstream queries", upstreamQueries.Load())
+	}
+}
+
+func TestServeDNSDoesNotStripServicePrefixForNonServiceQuestion(t *testing.T) {
+	t.Parallel()
+
+	udpConn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+	defer udpConn.Close()
+
+	var upstreamQueries atomic.Int32
+	handler := dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		upstreamQueries.Add(1)
+		resp := new(dns.Msg)
+		resp.SetReply(req)
+		resp.Answer = []dns.RR{&dns.TXT{
+			Hdr: dns.RR_Header{Name: req.Question[0].Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 20},
+			Txt: []string{"forwarded"},
+		}}
+		_ = w.WriteMsg(resp)
+	})
+
+	server := &dns.Server{PacketConn: udpConn, Handler: handler}
+	go func() { _ = server.ActivateAndServe() }()
+	t.Cleanup(func() {
+		_ = server.Shutdown()
+	})
+
+	r := NewRouter([]string{"127.0.0.2"}, "", "223.5.5.5", "", nil)
+	r.BlockRule = suffixtree.NewNodeFromRules()
+	r.DirectRule = suffixtree.NewNodeFromRules()
+	r.ProxyRule = suffixtree.NewNodeFromRules("example.com.")
+	r.dns.upstreamAddrs = []string{udpConn.LocalAddr().String()}
+	r.dns.upstreamIndex = 0
+
+	req := new(dns.Msg)
+	req.SetQuestion("_xmpp-client._tcp.example.com.", dns.TypeTXT)
+	writer := &mockDNSWriter{localAddr: &net.UDPAddr{IP: net.ParseIP("127.0.0.2"), Port: 53}}
+
+	r.ServeDNS(writer, req)
+
 	if writer.msg == nil {
 		t.Fatal("expected response message")
 	}
@@ -351,17 +580,10 @@ func TestServeDNSForProxyRuleForwardsSRVQuestion(t *testing.T) {
 		t.Fatalf("expected successful response, got rcode %d", writer.msg.Rcode)
 	}
 	if len(writer.msg.Answer) != 1 {
-		t.Fatalf("expected forwarded SRV answer, got %d answers", len(writer.msg.Answer))
-	}
-	answer, ok := writer.msg.Answer[0].(*dns.SRV)
-	if !ok {
-		t.Fatalf("expected SRV record, got %T", writer.msg.Answer[0])
-	}
-	if answer.Target != "srv.example.com." || answer.Port != 5222 {
-		t.Fatalf("unexpected SRV answer: target=%s port=%d", answer.Target, answer.Port)
+		t.Fatalf("expected forwarded TXT answer, got %d answers", len(writer.msg.Answer))
 	}
 	if upstreamQueries.Load() != 1 {
-		t.Fatalf("expected SRV query to be forwarded once, got %d upstream queries", upstreamQueries.Load())
+		t.Fatalf("expected TXT query to be forwarded once, got %d upstream queries", upstreamQueries.Load())
 	}
 }
 
@@ -715,6 +937,25 @@ func TestServeDNSReturnsNXDomainForBlockedDomain(t *testing.T) {
 	}
 }
 
+func TestServeDNSRejectsUnsupportedQuestionClass(t *testing.T) {
+	t.Parallel()
+
+	r := NewRouter([]string{"127.0.0.1"}, "", "223.5.5.5", "", nil)
+	req := new(dns.Msg)
+	req.SetQuestion("example.com.", dns.TypeA)
+	req.Question[0].Qclass = dns.ClassCHAOS
+
+	writer := &mockDNSWriter{localAddr: &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 53}}
+	r.ServeDNS(writer, req)
+
+	if writer.msg == nil {
+		t.Fatal("expected response message")
+	}
+	if writer.msg.Rcode != dns.RcodeNotImplemented {
+		t.Fatalf("expected NOTIMP, got %d", writer.msg.Rcode)
+	}
+}
+
 func TestServeDNSRejectsMultipleQuestions(t *testing.T) {
 	t.Parallel()
 
@@ -803,3 +1044,20 @@ func (w *mockDNSWriter) Close() error                { return nil }
 func (w *mockDNSWriter) TsigStatus() error           { return nil }
 func (w *mockDNSWriter) TsigTimersOnly(bool)         {}
 func (w *mockDNSWriter) Hijack()                     {}
+
+func startUDPTestDNSServer(t *testing.T, handler dns.Handler) string {
+	t.Helper()
+
+	udpConn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen udp: %v", err)
+	}
+
+	server := &dns.Server{PacketConn: udpConn, Handler: handler}
+	go func() { _ = server.ActivateAndServe() }()
+	t.Cleanup(func() {
+		_ = server.Shutdown()
+		_ = udpConn.Close()
+	})
+	return udpConn.LocalAddr().String()
+}