From 21528a5346c75cdac6ff65726333106da17b99c9 Mon Sep 17 00:00:00 2001 From: Viktor Liu Date: Wed, 22 Jan 2025 12:11:21 +0100 Subject: [PATCH 1/2] Match more specific dns handler first --- client/internal/dns/handler_chain.go | 27 +++- client/internal/dns/handler_chain_test.go | 153 ++++++++++++++++++++++ 2 files changed, 173 insertions(+), 7 deletions(-) diff --git a/client/internal/dns/handler_chain.go b/client/internal/dns/handler_chain.go index 5f63d1ab3f8..673f410e295 100644 --- a/client/internal/dns/handler_chain.go +++ b/client/internal/dns/handler_chain.go @@ -105,17 +105,30 @@ func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority MatchSubdomains: matchSubdomains, } - // Insert handler in priority order - pos := 0 + pos := c.findHandlerPosition(entry) + c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...) +} + +// findHandlerPosition determines where to insert a new handler based on priority and specificity +func (c *HandlerChain) findHandlerPosition(newEntry HandlerEntry) int { for i, h := range c.handlers { - if h.Priority < priority { - pos = i - break + // prio first + if h.Priority < newEntry.Priority { + return i + } + + // domain specificity next + if h.Priority == newEntry.Priority { + newDots := strings.Count(newEntry.Pattern, ".") + existingDots := strings.Count(h.Pattern, ".") + if newDots > existingDots { + return i + } } - pos = i + 1 } - c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...) + // add at end + return len(c.handlers) } // RemoveHandler removes a handler for the given pattern and priority diff --git a/client/internal/dns/handler_chain_test.go b/client/internal/dns/handler_chain_test.go index eb40c907fb9..f744618a7d9 100644 --- a/client/internal/dns/handler_chain_test.go +++ b/client/internal/dns/handler_chain_test.go @@ -677,3 +677,156 @@ func TestHandlerChain_CaseSensitivity(t *testing.T) { }) } } + +func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) { + tests := []struct { + name string + scenario string + ops []struct { + action string + pattern string + priority int + subdomain bool + } + query string + expectedMatch string + }{ + { + name: "more specific domain matches first", + scenario: "sub.example.com should match before example.com", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "sub.example.com.", nbdns.PriorityMatchDomain, false}, + }, + query: "sub.example.com.", + expectedMatch: "sub.example.com.", + }, + { + name: "more specific domain matches first, both match subdomains", + scenario: "sub.example.com should match before example.com", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "sub.example.com.", nbdns.PriorityMatchDomain, true}, + }, + query: "sub.example.com.", + expectedMatch: "sub.example.com.", + }, + { + name: "maintain specificity order after removal", + scenario: "after removing most specific, should fall back to less specific", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "sub.example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "test.sub.example.com.", nbdns.PriorityMatchDomain, false}, + {"remove", "test.sub.example.com.", nbdns.PriorityMatchDomain, false}, + }, + query: "test.sub.example.com.", + expectedMatch: "sub.example.com.", + }, + { + name: "priority overrides specificity", + scenario: "less specific domain with higher priority should match first", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "sub.example.com.", nbdns.PriorityMatchDomain, false}, + {"add", "example.com.", nbdns.PriorityDNSRoute, true}, + }, + query: "sub.example.com.", + expectedMatch: "example.com.", + }, + { + name: "equal priority respects specificity", + scenario: "with equal priority, more specific domain should match", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "other.example.com.", nbdns.PriorityMatchDomain, true}, + {"add", "sub.example.com.", nbdns.PriorityMatchDomain, false}, + }, + query: "sub.example.com.", + expectedMatch: "sub.example.com.", + }, + { + name: "specific matches before wildcard", + scenario: "specific domain should match before wildcard at same priority", + ops: []struct { + action string + pattern string + priority int + subdomain bool + }{ + {"add", "*.example.com.", nbdns.PriorityDNSRoute, false}, + {"add", "sub.example.com.", nbdns.PriorityDNSRoute, false}, + }, + query: "sub.example.com.", + expectedMatch: "sub.example.com.", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + chain := nbdns.NewHandlerChain() + handlers := make(map[string]*nbdns.MockSubdomainHandler) + + for _, op := range tt.ops { + if op.action == "add" { + handler := &nbdns.MockSubdomainHandler{Subdomains: op.subdomain} + handlers[op.pattern] = handler + chain.AddHandler(op.pattern, handler, op.priority, nil) + } else { + chain.RemoveHandler(op.pattern, op.priority) + } + } + + r := new(dns.Msg) + r.SetQuestion(tt.query, dns.TypeA) + w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}} + + // Setup handler expectations + for pattern, handler := range handlers { + if pattern == tt.expectedMatch { + handler.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) { + w := args.Get(0).(dns.ResponseWriter) + r := args.Get(1).(*dns.Msg) + resp := new(dns.Msg) + resp.SetReply(r) + w.WriteMsg(resp) + }).Once() + } + } + + chain.ServeDNS(w, r) + + for pattern, handler := range handlers { + if pattern == tt.expectedMatch { + handler.AssertNumberOfCalls(t, "ServeDNS", 1) + } else { + handler.AssertNumberOfCalls(t, "ServeDNS", 0) + } + } + }) + } +} From 73e0b3fbf250cdf16b0ffe9757982ebda3f6f4fe Mon Sep 17 00:00:00 2001 From: Viktor Liu Date: Wed, 22 Jan 2025 12:20:32 +0100 Subject: [PATCH 2/2] Fix lint --- client/internal/dns/handler_chain_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/client/internal/dns/handler_chain_test.go b/client/internal/dns/handler_chain_test.go index f744618a7d9..d04bfbbb396 100644 --- a/client/internal/dns/handler_chain_test.go +++ b/client/internal/dns/handler_chain_test.go @@ -813,7 +813,7 @@ func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) { r := args.Get(1).(*dns.Msg) resp := new(dns.Msg) resp.SetReply(r) - w.WriteMsg(resp) + assert.NoError(t, w.WriteMsg(resp)) }).Once() } }