Add dns interceptor based domain route functionality (#3032)

client/internal/dns/handler_chain.go

@@ -0,0 +1,192 @@
+package dns
+
+import (
+	"strings"
+	"sync"
+
+	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	PriorityDNSRoute    = 100
+	PriorityMatchDomain = 50
+	PriorityDefault     = 0
+)
+
+type HandlerEntry struct {
+	Handler     dns.Handler
+	Priority    int
+	Pattern     string
+	OrigPattern string
+	IsWildcard  bool
+	StopHandler handlerWithStop
+}
+
+// HandlerChain represents a prioritized chain of DNS handlers
+type HandlerChain struct {
+	mu       sync.RWMutex
+	handlers []HandlerEntry
+}
+
+// ResponseWriterChain wraps a dns.ResponseWriter to track if handler wants to continue chain
+type ResponseWriterChain struct {
+	dns.ResponseWriter
+	shouldContinue bool
+}
+
+func (w *ResponseWriterChain) WriteMsg(m *dns.Msg) error {
+	// Check if this is a continue signal (NXDOMAIN with Zero bit set)
+	if m.Rcode == dns.RcodeNameError && m.MsgHdr.Zero {
+		w.shouldContinue = true
+		return nil
+	}
+	return w.ResponseWriter.WriteMsg(m)
+}
+
+func NewHandlerChain() *HandlerChain {
+	return &HandlerChain{
+		handlers: make([]HandlerEntry, 0),
+	}
+}
+
+// AddHandler adds a new handler to the chain, replacing any existing handler with the same pattern and priority
+func (c *HandlerChain) AddHandler(pattern string, handler dns.Handler, priority int, stopHandler handlerWithStop) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	origPattern := pattern
+	isWildcard := strings.HasPrefix(pattern, "*.")
+	if isWildcard {
+		pattern = pattern[2:]
+	}
+	pattern = dns.Fqdn(pattern)
+	origPattern = dns.Fqdn(origPattern)
+
+	// First remove any existing handler with same original pattern and priority
+	for i := len(c.handlers) - 1; i >= 0; i-- {
+		if c.handlers[i].OrigPattern == origPattern && c.handlers[i].Priority == priority {
+			if c.handlers[i].StopHandler != nil {
+				c.handlers[i].StopHandler.stop()
+			}
+			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
+			break
+		}
+	}
+
+	log.Debugf("adding handler for pattern: %s (original: %s, wildcard: %v) with priority %d",
+		pattern, origPattern, isWildcard, priority)
+
+	entry := HandlerEntry{
+		Handler:     handler,
+		Priority:    priority,
+		Pattern:     pattern,
+		OrigPattern: origPattern,
+		IsWildcard:  isWildcard,
+		StopHandler: stopHandler,
+	}
+
+	// Insert handler in priority order
+	pos := 0
+	for i, h := range c.handlers {
+		if h.Priority < priority {
+			pos = i
+			break
+		}
+		pos = i + 1
+	}
+
+	c.handlers = append(c.handlers[:pos], append([]HandlerEntry{entry}, c.handlers[pos:]...)...)
+}
+
+// RemoveHandler removes a handler for the given pattern and priority
+func (c *HandlerChain) RemoveHandler(pattern string, priority int) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	pattern = dns.Fqdn(pattern)
+
+	// Find and remove handlers matching both original pattern and priority
+	for i := len(c.handlers) - 1; i >= 0; i-- {
+		entry := c.handlers[i]
+		if entry.OrigPattern == pattern && entry.Priority == priority {
+			if entry.StopHandler != nil {
+				entry.StopHandler.stop()
+			}
+			c.handlers = append(c.handlers[:i], c.handlers[i+1:]...)
+			return
+		}
+	}
+}
+
+// HasHandlers returns true if there are any handlers remaining for the given pattern
+func (c *HandlerChain) HasHandlers(pattern string) bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	pattern = dns.Fqdn(pattern)
+	for _, entry := range c.handlers {
+		if entry.Pattern == pattern {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *HandlerChain) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	if len(r.Question) == 0 {
+		return
+	}
+
+	qname := r.Question[0].Name
+	log.Debugf("handling DNS request for %s", qname)
+
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	log.Debugf("current handlers (%d):", len(c.handlers))
+	for _, h := range c.handlers {
+		log.Debugf("  - pattern: %s, original: %s, wildcard: %v, priority: %d",
+			h.Pattern, h.OrigPattern, h.IsWildcard, h.Priority)
+	}
+
+	// Try handlers in priority order
+	for _, entry := range c.handlers {
+		var matched bool
+		switch {
+		case entry.Pattern == ".":
+			matched = true
+		case entry.IsWildcard:
+			parts := strings.Split(strings.TrimSuffix(qname, entry.Pattern), ".")
+			matched = len(parts) >= 2 && strings.HasSuffix(qname, entry.Pattern)
+		default:
+			matched = qname == entry.Pattern || strings.HasSuffix(qname, "."+entry.Pattern)
+		}
+
+		if !matched {
+			log.Debugf("trying domain match: pattern=%s qname=%s wildcard=%v matched=false",
+				entry.OrigPattern, qname, entry.IsWildcard)
+			continue
+		}
+
+		log.Debugf("handler matched: pattern=%s qname=%s wildcard=%v",
+			entry.OrigPattern, qname, entry.IsWildcard)
+		chainWriter := &ResponseWriterChain{ResponseWriter: w}
+		entry.Handler.ServeDNS(chainWriter, r)
+
+		// If handler wants to continue, try next handler
+		if chainWriter.shouldContinue {
+			log.Debugf("handler requested continue to next handler")
+			continue
+		}
+		return
+	}
+
+	// No handler matched or all handlers passed
+	log.Debugf("no handler found for %s", qname)
+	resp := &dns.Msg{}
+	resp.SetRcode(r, dns.RcodeNameError)
+	if err := w.WriteMsg(resp); err != nil {
+		log.Errorf("failed to write DNS response: %v", err)
+	}
+}

client/internal/dns/handler_chain_test.go

@@ -0,0 +1,490 @@
+package dns_test
+
+import (
+	"net"
+	"testing"
+
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	nbdns "github.com/netbirdio/netbird/client/internal/dns"
+)
+
+// MockHandler implements dns.Handler interface for testing
+type MockHandler struct {
+	mock.Mock
+}
+
+func (m *MockHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
+	m.Called(w, r)
+}
+
+// TestHandlerChain_ServeDNS_Priorities tests that handlers are executed in priority order
+func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	// Create mock handlers for different priorities
+	defaultHandler := &MockHandler{}
+	matchDomainHandler := &MockHandler{}
+	dnsRouteHandler := &MockHandler{}
+
+	// Setup handlers with different priorities
+	chain.AddHandler("example.com.", defaultHandler, nbdns.PriorityDefault, nil)
+	chain.AddHandler("example.com.", matchDomainHandler, nbdns.PriorityMatchDomain, nil)
+	chain.AddHandler("example.com.", dnsRouteHandler, nbdns.PriorityDNSRoute, nil)
+
+	// Create test request
+	r := new(dns.Msg)
+	r.SetQuestion("example.com.", dns.TypeA)
+
+	// Create test writer
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+	// Setup expectations - only highest priority handler should be called
+	dnsRouteHandler.On("ServeDNS", mock.Anything, r).Once()
+	matchDomainHandler.On("ServeDNS", mock.Anything, r).Maybe()
+	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe()
+
+	// Execute
+	chain.ServeDNS(w, r)
+
+	// Verify all expectations were met
+	dnsRouteHandler.AssertExpectations(t)
+	matchDomainHandler.AssertExpectations(t)
+	defaultHandler.AssertExpectations(t)
+}
+
+// TestHandlerChain_ServeDNS_DomainMatching tests various domain matching scenarios
+func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
+	tests := []struct {
+		name          string
+		handlerDomain string
+		queryDomain   string
+		isWildcard    bool
+		shouldMatch   bool
+	}{
+		{
+			name:          "exact match",
+			handlerDomain: "example.com.",
+			queryDomain:   "example.com.",
+			isWildcard:    false,
+			shouldMatch:   true,
+		},
+		{
+			name:          "subdomain with non-wildcard",
+			handlerDomain: "example.com.",
+			queryDomain:   "sub.example.com.",
+			isWildcard:    false,
+			shouldMatch:   true,
+		},
+		{
+			name:          "wildcard match",
+			handlerDomain: "*.example.com.",
+			queryDomain:   "sub.example.com.",
+			isWildcard:    true,
+			shouldMatch:   true,
+		},
+		{
+			name:          "wildcard no match on apex",
+			handlerDomain: "*.example.com.",
+			queryDomain:   "example.com.",
+			isWildcard:    true,
+			shouldMatch:   false,
+		},
+		{
+			name:          "root zone match",
+			handlerDomain: ".",
+			queryDomain:   "anything.com.",
+			isWildcard:    false,
+			shouldMatch:   true,
+		},
+		{
+			name:          "no match different domain",
+			handlerDomain: "example.com.",
+			queryDomain:   "example.org.",
+			isWildcard:    false,
+			shouldMatch:   false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			mockHandler := &MockHandler{}
+
+			pattern := tt.handlerDomain
+			if tt.isWildcard {
+				pattern = "*." + tt.handlerDomain[2:] // Remove the first two chars if it's a wildcard
+			}
+
+			chain.AddHandler(pattern, mockHandler, nbdns.PriorityDefault, nil)
+
+			r := new(dns.Msg)
+			r.SetQuestion(tt.queryDomain, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			if tt.shouldMatch {
+				mockHandler.On("ServeDNS", mock.Anything, r).Once()
+			}
+
+			chain.ServeDNS(w, r)
+			mockHandler.AssertExpectations(t)
+		})
+	}
+}
+
+// TestHandlerChain_ServeDNS_OverlappingDomains tests behavior with overlapping domain patterns
+func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
+	tests := []struct {
+		name     string
+		handlers []struct {
+			pattern  string
+			priority int
+		}
+		queryDomain     string
+		expectedCalls   int
+		expectedHandler int // index of the handler that should be called
+	}{
+		{
+			name: "wildcard and exact same priority - exact should win",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "example.com.", priority: nbdns.PriorityDefault},
+			},
+			queryDomain:     "example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // exact match handler should be called
+		},
+		{
+			name: "higher priority wildcard over lower priority exact",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // higher priority wildcard handler should be called
+		},
+		{
+			name: "multiple wildcards different priorities",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "*.example.com.", priority: nbdns.PriorityMatchDomain},
+				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 2, // highest priority handler should be called
+		},
+		{
+			name: "subdomain with mix of patterns",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
+				{pattern: "test.example.com.", priority: nbdns.PriorityMatchDomain},
+				{pattern: "*.test.example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "sub.test.example.com.",
+			expectedCalls:   1,
+			expectedHandler: 2, // highest priority matching handler should be called
+		},
+		{
+			name: "root zone with specific domain",
+			handlers: []struct {
+				pattern  string
+				priority int
+			}{
+				{pattern: ".", priority: nbdns.PriorityDefault},
+				{pattern: "example.com.", priority: nbdns.PriorityDNSRoute},
+			},
+			queryDomain:     "example.com.",
+			expectedCalls:   1,
+			expectedHandler: 1, // higher priority specific domain should win over root
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			var handlers []*MockHandler
+
+			// Setup handlers and expectations
+			for i := range tt.handlers {
+				handler := &MockHandler{}
+				handlers = append(handlers, handler)
+
+				// Set expectation based on whether this handler should be called
+				if i == tt.expectedHandler {
+					handler.On("ServeDNS", mock.Anything, mock.Anything).Once()
+				} else {
+					handler.On("ServeDNS", mock.Anything, mock.Anything).Maybe()
+				}
+
+				chain.AddHandler(tt.handlers[i].pattern, handler, tt.handlers[i].priority, nil)
+			}
+
+			// Create and execute request
+			r := new(dns.Msg)
+			r.SetQuestion(tt.queryDomain, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+			chain.ServeDNS(w, r)
+
+			// Verify expectations
+			for _, handler := range handlers {
+				handler.AssertExpectations(t)
+			}
+		})
+	}
+}
+
+// TestHandlerChain_ServeDNS_ChainContinuation tests the chain continuation functionality
+func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	// Create handlers
+	handler1 := &MockHandler{}
+	handler2 := &MockHandler{}
+	handler3 := &MockHandler{}
+
+	// Add handlers in priority order
+	chain.AddHandler("example.com.", handler1, nbdns.PriorityDNSRoute, nil)
+	chain.AddHandler("example.com.", handler2, nbdns.PriorityMatchDomain, nil)
+	chain.AddHandler("example.com.", handler3, nbdns.PriorityDefault, nil)
+
+	// Create test request
+	r := new(dns.Msg)
+	r.SetQuestion("example.com.", dns.TypeA)
+
+	// Setup mock responses to simulate chain continuation
+	handler1.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// First handler signals continue
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeNameError)
+		resp.MsgHdr.Zero = true // Signal to continue
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	handler2.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// Second handler signals continue
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeNameError)
+		resp.MsgHdr.Zero = true
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	handler3.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
+		// Last handler responds normally
+		w := args.Get(0).(*nbdns.ResponseWriterChain)
+		resp := new(dns.Msg)
+		resp.SetRcode(r, dns.RcodeSuccess)
+		assert.NoError(t, w.WriteMsg(resp))
+	}).Once()
+
+	// Execute
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	chain.ServeDNS(w, r)
+
+	// Verify all handlers were called in order
+	handler1.AssertExpectations(t)
+	handler2.AssertExpectations(t)
+	handler3.AssertExpectations(t)
+}
+
+// mockResponseWriter implements dns.ResponseWriter for testing
+type mockResponseWriter struct {
+	mock.Mock
+}
+
+func (m *mockResponseWriter) LocalAddr() net.Addr       { return nil }
+func (m *mockResponseWriter) RemoteAddr() net.Addr      { return nil }
+func (m *mockResponseWriter) WriteMsg(*dns.Msg) error   { return nil }
+func (m *mockResponseWriter) Write([]byte) (int, error) { return 0, nil }
+func (m *mockResponseWriter) Close() error              { return nil }
+func (m *mockResponseWriter) TsigStatus() error         { return nil }
+func (m *mockResponseWriter) TsigTimersOnly(bool)       {}
+func (m *mockResponseWriter) Hijack()                   {}
+
+func TestHandlerChain_PriorityDeregistration(t *testing.T) {
+	tests := []struct {
+		name string
+		ops  []struct {
+			action   string // "add" or "remove"
+			pattern  string
+			priority int
+		}
+		query         string
+		expectedCalls map[int]bool // map[priority]shouldBeCalled
+	}{
+		{
+			name: "remove high priority keeps lower priority handler",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"remove", "example.com.", nbdns.PriorityDNSRoute},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    false,
+				nbdns.PriorityMatchDomain: true,
+			},
+		},
+		{
+			name: "remove lower priority keeps high priority handler",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"remove", "example.com.", nbdns.PriorityMatchDomain},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    true,
+				nbdns.PriorityMatchDomain: false,
+			},
+		},
+		{
+			name: "remove all handlers in order",
+			ops: []struct {
+				action   string
+				pattern  string
+				priority int
+			}{
+				{"add", "example.com.", nbdns.PriorityDNSRoute},
+				{"add", "example.com.", nbdns.PriorityMatchDomain},
+				{"add", "example.com.", nbdns.PriorityDefault},
+				{"remove", "example.com.", nbdns.PriorityDNSRoute},
+				{"remove", "example.com.", nbdns.PriorityMatchDomain},
+			},
+			query: "example.com.",
+			expectedCalls: map[int]bool{
+				nbdns.PriorityDNSRoute:    false,
+				nbdns.PriorityMatchDomain: false,
+				nbdns.PriorityDefault:     true,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain := nbdns.NewHandlerChain()
+			handlers := make(map[int]*MockHandler)
+
+			// Execute operations
+			for _, op := range tt.ops {
+				if op.action == "add" {
+					handler := &MockHandler{}
+					handlers[op.priority] = handler
+					chain.AddHandler(op.pattern, handler, op.priority, nil)
+				} else {
+					chain.RemoveHandler(op.pattern, op.priority)
+				}
+			}
+
+			// Create test request
+			r := new(dns.Msg)
+			r.SetQuestion(tt.query, dns.TypeA)
+			w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+
+			// Setup expectations
+			for priority, handler := range handlers {
+				if shouldCall, exists := tt.expectedCalls[priority]; exists && shouldCall {
+					handler.On("ServeDNS", mock.Anything, r).Once()
+				} else {
+					handler.On("ServeDNS", mock.Anything, r).Maybe()
+				}
+			}
+
+			// Execute request
+			chain.ServeDNS(w, r)
+
+			// Verify expectations
+			for _, handler := range handlers {
+				handler.AssertExpectations(t)
+			}
+
+			// Verify handler exists check
+			for priority, shouldExist := range tt.expectedCalls {
+				if shouldExist {
+					assert.True(t, chain.HasHandlers(tt.ops[0].pattern),
+						"Handler chain should have handlers for pattern after removing priority %d", priority)
+				}
+			}
+		})
+	}
+}
+
+func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
+	chain := nbdns.NewHandlerChain()
+
+	testDomain := "example.com."
+	testQuery := "test.example.com."
+
+	// Create handlers for three priority levels
+	routeHandler := &MockHandler{}
+	matchHandler := &MockHandler{}
+	defaultHandler := &MockHandler{}
+
+	// Create test request that will be reused
+	r := new(dns.Msg)
+	r.SetQuestion(testQuery, dns.TypeA)
+
+	// Add handlers in mixed order
+	chain.AddHandler(testDomain, defaultHandler, nbdns.PriorityDefault, nil)
+	chain.AddHandler(testDomain, routeHandler, nbdns.PriorityDNSRoute, nil)
+	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityMatchDomain, nil)
+
+	// Test 1: Initial state with all three handlers
+	w := &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Highest priority handler (routeHandler) should be called
+	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	routeHandler.AssertExpectations(t)
+
+	// Test 2: Remove highest priority handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)
+	assert.True(t, chain.HasHandlers(testDomain))
+
+	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Now middle priority handler (matchHandler) should be called
+	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	matchHandler.AssertExpectations(t)
+
+	// Test 3: Remove middle priority handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityMatchDomain)
+	assert.True(t, chain.HasHandlers(testDomain))
+
+	w = &nbdns.ResponseWriterChain{ResponseWriter: &mockResponseWriter{}}
+	// Now lowest priority handler (defaultHandler) should be called
+	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()
+
+	chain.ServeDNS(w, r)
+	defaultHandler.AssertExpectations(t)
+
+	// Test 4: Remove last handler
+	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)
+	assert.False(t, chain.HasHandlers(testDomain))
+}

client/internal/dns/mock_server.go

@@ -3,14 +3,30 @@ package dns
 import (
 	"fmt"
 
+	"github.com/miekg/dns"
+
 	nbdns "github.com/netbirdio/netbird/dns"
 )
 
 // MockServer is the mock instance of a dns server
 type MockServer struct {
-	InitializeFunc      func() error
-	StopFunc            func()
-	UpdateDNSServerFunc func(serial uint64, update nbdns.Config) error
+	InitializeFunc        func() error
+	StopFunc              func()
+	UpdateDNSServerFunc   func(serial uint64, update nbdns.Config) error
+	RegisterHandlerFunc   func([]string, dns.Handler, int)
+	DeregisterHandlerFunc func([]string, int)
+}
+
+func (m *MockServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+	if m.RegisterHandlerFunc != nil {
+		m.RegisterHandlerFunc(domains, handler, priority)
+	}
+}
+
+func (m *MockServer) DeregisterHandler(domains []string, priority int) {
+	if m.DeregisterHandlerFunc != nil {
+		m.DeregisterHandlerFunc(domains, priority)
+	}
 }
 
 // Initialize mock implementation of Initialize from Server interface

client/internal/dns/server.go

@@ -30,6 +30,8 @@ type IosDnsManager interface {
 
 // Server is a dns server interface
 type Server interface {
+	RegisterHandler(domains []string, handler dns.Handler, priority int)
+	DeregisterHandler(domains []string, priority int)
 	Initialize() error
 	Stop()
 	DnsIP() string
@@ -48,12 +50,14 @@ type DefaultServer struct {
 	mux                sync.Mutex
 	service            service
 	dnsMuxMap          registeredHandlerMap
+	handlerPriorities  map[string]int
 	localResolver      *localResolver
 	wgInterface        WGIface
 	hostManager        hostManager
 	updateSerial       uint64
 	previousConfigHash uint64
 	currentConfig      HostDNSConfig
+	handlerChain       *HandlerChain
 
 	// permanent related properties
 	permanent      bool
@@ -74,8 +78,9 @@ type handlerWithStop interface {
 }
 
 type muxUpdate struct {
-	domain  string
-	handler handlerWithStop
+	domain   string
+	handler  handlerWithStop
+	priority int
 }
 
 // NewDefaultServer returns a new dns server
@@ -135,10 +140,12 @@ func NewDefaultServerIos(
 func newDefaultServer(ctx context.Context, wgInterface WGIface, dnsService service, statusRecorder *peer.Status, stateManager *statemanager.Manager) *DefaultServer {
 	ctx, stop := context.WithCancel(ctx)
 	defaultServer := &DefaultServer{
-		ctx:       ctx,
-		ctxCancel: stop,
-		service:   dnsService,
-		dnsMuxMap: make(registeredHandlerMap),
+		ctx:               ctx,
+		ctxCancel:         stop,
+		service:           dnsService,
+		handlerChain:      NewHandlerChain(),
+		dnsMuxMap:         make(registeredHandlerMap),
+		handlerPriorities: make(map[string]int),
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
@@ -151,6 +158,41 @@ func newDefaultServer(ctx context.Context, wgInterface WGIface, dnsService servi
 	return defaultServer
 }
 
+func (s *DefaultServer) RegisterHandler(domains []string, handler dns.Handler, priority int) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	s.registerHandler(domains, handler, priority)
+}
+
+func (s *DefaultServer) registerHandler(domains []string, handler dns.Handler, priority int) {
+	log.Debugf("registering handler %s with priority %d", handler, priority)
+
+	for _, domain := range domains {
+		s.handlerChain.AddHandler(domain, handler, priority, nil)
+		s.handlerPriorities[domain] = priority
+		s.service.RegisterMux(nbdns.NormalizeZone(domain), s.handlerChain)
+	}
+}
+
+func (s *DefaultServer) DeregisterHandler(domains []string, priority int) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	s.deregisterHandler(domains, priority)
+}
+
+func (s *DefaultServer) deregisterHandler(domains []string, priority int) {
+	for _, domain := range domains {
+		s.handlerChain.RemoveHandler(domain, priority)
+
+		// Only deregister from service if no handlers remain
+		if !s.handlerChain.HasHandlers(domain) {
+			s.service.DeregisterMux(nbdns.NormalizeZone(domain))
+		}
+	}
+}
+
 // Initialize instantiate host manager and the dns service
 func (s *DefaultServer) Initialize() (err error) {
 	s.mux.Lock()
@@ -343,14 +385,14 @@ func (s *DefaultServer) buildLocalHandlerUpdate(customZones []nbdns.CustomZone)
 	localRecords := make(map[string]nbdns.SimpleRecord, 0)
 
 	for _, customZone := range customZones {
-
 		if len(customZone.Records) == 0 {
 			return nil, nil, fmt.Errorf("received an empty list of records")
 		}
 
 		muxUpdates = append(muxUpdates, muxUpdate{
-			domain:  customZone.Domain,
-			handler: s.localResolver,
+			domain:   customZone.Domain,
+			handler:  s.localResolver,
+			priority: PriorityMatchDomain,
 		})
 
 		for _, record := range customZone.Records {
@@ -412,8 +454,9 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 
 		if nsGroup.Primary {
 			muxUpdates = append(muxUpdates, muxUpdate{
-				domain:  nbdns.RootZone,
-				handler: handler,
+				domain:   nbdns.RootZone,
+				handler:  handler,
+				priority: PriorityDefault,
 			})
 			continue
 		}
@@ -429,8 +472,9 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 				return nil, fmt.Errorf("received a nameserver group with an empty domain element")
 			}
 			muxUpdates = append(muxUpdates, muxUpdate{
-				domain:  domain,
-				handler: handler,
+				domain:   domain,
+				handler:  handler,
+				priority: PriorityMatchDomain,
 			})
 		}
 	}
@@ -440,12 +484,16 @@ func (s *DefaultServer) buildUpstreamHandlerUpdate(nameServerGroups []*nbdns.Nam
 
 func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 	muxUpdateMap := make(registeredHandlerMap)
+	handlersByPriority := make(map[string]int)
 
 	var isContainRootUpdate bool
 
+	// First register new handlers
 	for _, update := range muxUpdates {
-		s.service.RegisterMux(update.domain, update.handler)
+		s.registerHandler([]string{update.domain}, update.handler, update.priority)
 		muxUpdateMap[update.domain] = update.handler
+		handlersByPriority[update.domain] = update.priority
+
 		if existingHandler, ok := s.dnsMuxMap[update.domain]; ok {
 			existingHandler.stop()
 		}
@@ -455,6 +503,7 @@ func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 		}
 	}
 
+	// Then deregister old handlers not in the update
 	for key, existingHandler := range s.dnsMuxMap {
 		_, found := muxUpdateMap[key]
 		if !found {
@@ -463,12 +512,16 @@ func (s *DefaultServer) updateMux(muxUpdates []muxUpdate) {
 				existingHandler.stop()
 			} else {
 				existingHandler.stop()
-				s.service.DeregisterMux(key)
+				// Deregister with the priority that was used to register
+				if oldPriority, ok := s.handlerPriorities[key]; ok {
+					s.deregisterHandler([]string{key}, oldPriority)
+				}
 			}
 		}
 	}
 
 	s.dnsMuxMap = muxUpdateMap
+	s.handlerPriorities = handlersByPriority
 }
 
 func (s *DefaultServer) updateLocalResolver(update map[string]nbdns.SimpleRecord) {
@@ -517,13 +570,13 @@ func (s *DefaultServer) upstreamCallbacks(
 		if nsGroup.Primary {
 			removeIndex[nbdns.RootZone] = -1
 			s.currentConfig.RouteAll = false
-			s.service.DeregisterMux(nbdns.RootZone)
+			s.deregisterHandler([]string{nbdns.RootZone}, PriorityDefault)
 		}
 
 		for i, item := range s.currentConfig.Domains {
 			if _, found := removeIndex[item.Domain]; found {
 				s.currentConfig.Domains[i].Disabled = true
-				s.service.DeregisterMux(item.Domain)
+				s.deregisterHandler([]string{item.Domain}, PriorityMatchDomain)
 				removeIndex[item.Domain] = i
 			}
 		}
@@ -554,7 +607,7 @@ func (s *DefaultServer) upstreamCallbacks(
 				continue
 			}
 			s.currentConfig.Domains[i].Disabled = false
-			s.service.RegisterMux(domain, handler)
+			s.registerHandler([]string{domain}, handler, PriorityMatchDomain)
 		}
 
 		l := log.WithField("nameservers", nsGroup.NameServers)
@@ -562,7 +615,7 @@ func (s *DefaultServer) upstreamCallbacks(
 
 		if nsGroup.Primary {
 			s.currentConfig.RouteAll = true
-			s.service.RegisterMux(nbdns.RootZone, handler)
+			s.registerHandler([]string{nbdns.RootZone}, handler, PriorityDefault)
 		}
 		if err := s.hostManager.applyDNSConfig(s.currentConfig, s.stateManager); err != nil {
 			l.WithError(err).Error("reactivate temporary disabled nameserver group, DNS update apply")
@@ -593,7 +646,8 @@ func (s *DefaultServer) addHostRootZone() {
 	}
 	handler.deactivate = func(error) {}
 	handler.reactivate = func() {}
-	s.service.RegisterMux(nbdns.RootZone, handler)
+
+	s.registerHandler([]string{nbdns.RootZone}, handler, PriorityDefault)
 }
 
 func (s *DefaultServer) updateNSGroupStates(groups []*nbdns.NameServerGroup) {

client/internal/dns/server_test.go

@@ -512,7 +512,7 @@ func TestDNSServerStartStop(t *testing.T) {
 				t.Error(err)
 			}
 
-			dnsServer.service.RegisterMux("netbird.cloud", dnsServer.localResolver)
+			dnsServer.registerHandler([]string{"netbird.cloud"}, dnsServer.localResolver, 1)
 
 			resolver := &net.Resolver{
 				PreferGo: true,
@@ -560,7 +560,9 @@ func TestDNSServerUpstreamDeactivateCallback(t *testing.T) {
 		localResolver: &localResolver{
 			registeredMap: make(registrationMap),
 		},
-		hostManager: hostManager,
+		handlerChain:      NewHandlerChain(),
+		handlerPriorities: make(map[string]int),
+		hostManager:       hostManager,
 		currentConfig: HostDNSConfig{
 			Domains: []DomainConfig{
 				{false, "domain0", false},

client/internal/dns/service_listener.go

@@ -105,6 +105,7 @@ func (s *serviceViaListener) Stop() {
 }
 
 func (s *serviceViaListener) RegisterMux(pattern string, handler dns.Handler) {
+	log.Debugf("registering dns handler for pattern: %s", pattern)
 	s.dnsMux.Handle(pattern, handler)
 }
 

client/internal/dns/upstream.go

@@ -66,6 +66,11 @@ func newUpstreamResolverBase(ctx context.Context, statusRecorder *peer.Status) *
 	}
 }
 
+// String returns a string representation of the upstream resolver
+func (u *upstreamResolverBase) String() string {
+	return fmt.Sprintf("%v", u.upstreamServers)
+}
+
 func (u *upstreamResolverBase) stop() {
 	log.Debugf("stopping serving DNS for upstreams %s", u.upstreamServers)
 	u.cancel()