diff --git a/client/internal/dns/server.go b/client/internal/dns/server.go index dcd4cb9d0..c8ffd3412 100644 --- a/client/internal/dns/server.go +++ b/client/internal/dns/server.go @@ -135,7 +135,7 @@ type DefaultServer struct { disableSys bool mux sync.Mutex service service - dnsMuxMap registeredHandlerMap + dnsMuxHandlers []handlerWrapper localResolver *local.Resolver wgInterface WGIface hostManager hostManager @@ -199,8 +199,6 @@ type handlerWrapper struct { priority int } -type registeredHandlerMap map[types.HandlerID]handlerWrapper - // DefaultServerConfig holds configuration parameters for NewDefaultServer type DefaultServerConfig struct { WgInterface WGIface @@ -289,7 +287,6 @@ func newDefaultServer( service: dnsService, handlerChain: handlerChain, extraDomains: make(map[domain.Domain]int), - dnsMuxMap: make(registeredHandlerMap), localResolver: local.NewResolver(), wgInterface: wgInterface, statusRecorder: statusRecorder, @@ -328,7 +325,7 @@ func (s *DefaultServer) SetRouteSources(selected, active func() route.HAMap) { type routeSettable interface { setSelectedRoutes(func() route.HAMap) } - for _, entry := range s.dnsMuxMap { + for _, entry := range s.dnsMuxHandlers { if h, ok := entry.handler.(routeSettable); ok { h.setSelectedRoutes(selected) } @@ -978,19 +975,16 @@ func (s *DefaultServer) usableNameServers(nameServers []nbdns.NameServer) []neti func (s *DefaultServer) updateMux(muxUpdates []handlerWrapper) { // this will introduce a short period of time when the server is not able to handle DNS requests - for _, existing := range s.dnsMuxMap { + for _, existing := range s.dnsMuxHandlers { s.deregisterHandler([]string{existing.domain}, existing.priority) existing.handler.Stop() } - muxUpdateMap := make(registeredHandlerMap) - for _, update := range muxUpdates { s.registerHandler([]string{update.domain}, update.handler, update.priority) - muxUpdateMap[update.handler.ID()] = update } - s.dnsMuxMap = muxUpdateMap + s.dnsMuxHandlers = muxUpdates } // updateNSGroupStates records the new group set and pokes the refresher. @@ -1204,7 +1198,7 @@ func (s *DefaultServer) groupHasImmediateUpstream(servers []netip.AddrPort, snap // in more than one handler. func (s *DefaultServer) collectUpstreamHealth() map[netip.AddrPort]UpstreamHealth { merged := make(map[netip.AddrPort]UpstreamHealth) - for _, entry := range s.dnsMuxMap { + for _, entry := range s.dnsMuxHandlers { reporter, ok := entry.handler.(upstreamHealthReporter) if !ok { continue diff --git a/client/internal/dns/server_test.go b/client/internal/dns/server_test.go index 722c2abd7..44199150a 100644 --- a/client/internal/dns/server_test.go +++ b/client/internal/dns/server_test.go @@ -104,19 +104,6 @@ func init() { formatter.SetTextFormatter(log.StandardLogger()) } -func generateDummyHandler(d string, servers []nbdns.NameServer) *upstreamResolverBase { - var srvs []netip.AddrPort - for _, srv := range servers { - srvs = append(srvs, srv.AddrPort()) - } - u := &upstreamResolverBase{ - domain: domain.Domain(d), - cancel: func() {}, - } - u.addRace(srvs) - return u -} - func TestUpdateDNSServer(t *testing.T) { nameServers := []nbdns.NameServer{ @@ -132,22 +119,20 @@ func TestUpdateDNSServer(t *testing.T) { }, } - dummyHandler := local.NewResolver() - testCases := []struct { name string - initUpstreamMap registeredHandlerMap + initUpstreamMap []handlerWrapper initLocalZones []nbdns.CustomZone initSerial uint64 inputSerial uint64 inputUpdate nbdns.Config shouldFail bool - expectedUpstreamMap registeredHandlerMap + expectedUpstreamMap []handlerWrapper expectedLocalQs []dns.Question }{ { name: "Initial Config Should Succeed", - initUpstreamMap: make(registeredHandlerMap), + initUpstreamMap: nil, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ @@ -169,20 +154,17 @@ func TestUpdateDNSServer(t *testing.T) { }, }, }, - expectedUpstreamMap: registeredHandlerMap{ - generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{ + expectedUpstreamMap: []handlerWrapper{ + { domain: "netbird.io", - handler: dummyHandler, priority: PriorityUpstream, }, - dummyHandler.ID(): handlerWrapper{ + { domain: "netbird.cloud", - handler: dummyHandler, priority: PriorityLocal, }, - generateDummyHandler(".", nameServers).ID(): handlerWrapper{ + { domain: nbdns.RootZone, - handler: dummyHandler, priority: PriorityDefault, }, }, @@ -191,10 +173,10 @@ func TestUpdateDNSServer(t *testing.T) { { name: "New Config Should Succeed", initLocalZones: []nbdns.CustomZone{{Domain: "netbird.cloud", Records: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}}}}, - initUpstreamMap: registeredHandlerMap{ - generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{ + initUpstreamMap: []handlerWrapper{ + { domain: "netbird.cloud", - handler: dummyHandler, + handler: &mockHandler{}, priority: PriorityUpstream, }, }, @@ -215,15 +197,13 @@ func TestUpdateDNSServer(t *testing.T) { }, }, }, - expectedUpstreamMap: registeredHandlerMap{ - generateDummyHandler("netbird.io", nameServers).ID(): handlerWrapper{ + expectedUpstreamMap: []handlerWrapper{ + { domain: "netbird.io", - handler: dummyHandler, priority: PriorityUpstream, }, - "local-resolver": handlerWrapper{ + { domain: "netbird.cloud", - handler: dummyHandler, priority: PriorityLocal, }, }, @@ -232,7 +212,7 @@ func TestUpdateDNSServer(t *testing.T) { { name: "Smaller Config Serial Should Be Skipped", initLocalZones: []nbdns.CustomZone{}, - initUpstreamMap: make(registeredHandlerMap), + initUpstreamMap: nil, initSerial: 2, inputSerial: 1, shouldFail: true, @@ -240,7 +220,7 @@ func TestUpdateDNSServer(t *testing.T) { { name: "Empty NS Group Domain Or Not Primary Element Should Fail", initLocalZones: []nbdns.CustomZone{}, - initUpstreamMap: make(registeredHandlerMap), + initUpstreamMap: nil, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ @@ -262,7 +242,7 @@ func TestUpdateDNSServer(t *testing.T) { { name: "Invalid NS Group Nameservers list Should Fail", initLocalZones: []nbdns.CustomZone{}, - initUpstreamMap: make(registeredHandlerMap), + initUpstreamMap: nil, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ @@ -284,7 +264,7 @@ func TestUpdateDNSServer(t *testing.T) { { name: "Invalid Custom Zone Records list Should Skip", initLocalZones: []nbdns.CustomZone{}, - initUpstreamMap: make(registeredHandlerMap), + initUpstreamMap: nil, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ @@ -301,42 +281,41 @@ func TestUpdateDNSServer(t *testing.T) { }, }, }, - expectedUpstreamMap: registeredHandlerMap{generateDummyHandler(".", nameServers).ID(): handlerWrapper{ + expectedUpstreamMap: []handlerWrapper{{ domain: ".", - handler: dummyHandler, priority: PriorityDefault, }}, }, { name: "Empty Config Should Succeed and Clean Maps", initLocalZones: []nbdns.CustomZone{{Domain: "netbird.cloud", Records: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}}}}, - initUpstreamMap: registeredHandlerMap{ - generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{ + initUpstreamMap: []handlerWrapper{ + { domain: zoneRecords[0].Name, - handler: dummyHandler, + handler: &mockHandler{}, priority: PriorityUpstream, }, }, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ServiceEnable: true}, - expectedUpstreamMap: make(registeredHandlerMap), + expectedUpstreamMap: nil, expectedLocalQs: []dns.Question{}, }, { name: "Disabled Service Should clean map", initLocalZones: []nbdns.CustomZone{{Domain: "netbird.cloud", Records: []nbdns.SimpleRecord{{Name: "netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"}}}}, - initUpstreamMap: registeredHandlerMap{ - generateDummyHandler(zoneRecords[0].Name, nameServers).ID(): handlerWrapper{ + initUpstreamMap: []handlerWrapper{ + { domain: zoneRecords[0].Name, - handler: dummyHandler, + handler: &mockHandler{}, priority: PriorityUpstream, }, }, initSerial: 0, inputSerial: 1, inputUpdate: nbdns.Config{ServiceEnable: false}, - expectedUpstreamMap: make(registeredHandlerMap), + expectedUpstreamMap: nil, expectedLocalQs: []dns.Question{}, }, } @@ -393,7 +372,7 @@ func TestUpdateDNSServer(t *testing.T) { } }() - dnsServer.dnsMuxMap = testCase.initUpstreamMap + dnsServer.dnsMuxHandlers = testCase.initUpstreamMap dnsServer.localResolver.Update(testCase.initLocalZones) dnsServer.updateSerial = testCase.initSerial @@ -405,14 +384,20 @@ func TestUpdateDNSServer(t *testing.T) { t.Fatalf("update dns server should not fail, got error: %v", err) } - if len(dnsServer.dnsMuxMap) != len(testCase.expectedUpstreamMap) { - t.Fatalf("update upstream failed, map size is different than expected, want %d, got %d", len(testCase.expectedUpstreamMap), len(dnsServer.dnsMuxMap)) + if len(dnsServer.dnsMuxHandlers) != len(testCase.expectedUpstreamMap) { + t.Fatalf("update upstream failed, map size is different than expected, want %d, got %d", len(testCase.expectedUpstreamMap), len(dnsServer.dnsMuxHandlers)) } - for key := range testCase.expectedUpstreamMap { - _, found := dnsServer.dnsMuxMap[key] + for _, expected := range testCase.expectedUpstreamMap { + found := false + for _, got := range dnsServer.dnsMuxHandlers { + if got.domain == expected.domain && got.priority == expected.priority { + found = true + break + } + } if !found { - t.Fatalf("update upstream failed, key %s was not found in the dnsMuxMap: %#v", key, dnsServer.dnsMuxMap) + t.Fatalf("update upstream failed, handler for domain=%s priority=%d not found in dnsMuxHandlers: %#v", expected.domain, expected.priority, dnsServer.dnsMuxHandlers) } } @@ -512,8 +497,8 @@ func TestDNSFakeResolverHandleUpdates(t *testing.T) { } }() - dnsServer.dnsMuxMap = registeredHandlerMap{ - "id1": handlerWrapper{ + dnsServer.dnsMuxHandlers = []handlerWrapper{ + { domain: zoneRecords[0].Name, handler: &local.Resolver{}, priority: PriorityUpstream, @@ -1029,15 +1014,15 @@ func (m *mockService) RegisterMux(string, dns.Handler) {} func (m *mockService) DeregisterMux(string) {} func TestDefaultServer_UpdateMux(t *testing.T) { - baseMatchHandlers := registeredHandlerMap{ - "upstream-group1": { + baseMatchHandlers := []handlerWrapper{ + { domain: "example.com", handler: &mockHandler{ Id: "upstream-group1", }, priority: PriorityUpstream, }, - "upstream-group2": { + { domain: "example.com", handler: &mockHandler{ Id: "upstream-group2", @@ -1046,15 +1031,15 @@ func TestDefaultServer_UpdateMux(t *testing.T) { }, } - baseRootHandlers := registeredHandlerMap{ - "upstream-root1": { + baseRootHandlers := []handlerWrapper{ + { domain: ".", handler: &mockHandler{ Id: "upstream-root1", }, priority: PriorityDefault, }, - "upstream-root2": { + { domain: ".", handler: &mockHandler{ Id: "upstream-root2", @@ -1063,22 +1048,22 @@ func TestDefaultServer_UpdateMux(t *testing.T) { }, } - baseMixedHandlers := registeredHandlerMap{ - "upstream-group1": { + baseMixedHandlers := []handlerWrapper{ + { domain: "example.com", handler: &mockHandler{ Id: "upstream-group1", }, priority: PriorityUpstream, }, - "upstream-group2": { + { domain: "example.com", handler: &mockHandler{ Id: "upstream-group2", }, priority: PriorityUpstream - 1, }, - "upstream-other": { + { domain: "other.com", handler: &mockHandler{ Id: "upstream-other", @@ -1089,7 +1074,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) { tests := []struct { name string - initialHandlers registeredHandlerMap + initialHandlers []handlerWrapper updates []handlerWrapper expectedHandlers map[string]string // map[HandlerID]domain description string @@ -1373,32 +1358,38 @@ func TestDefaultServer_UpdateMux(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { server := &DefaultServer{ - dnsMuxMap: tt.initialHandlers, - handlerChain: NewHandlerChain(), - service: &mockService{}, + dnsMuxHandlers: tt.initialHandlers, + handlerChain: NewHandlerChain(), + service: &mockService{}, } // Perform the update server.updateMux(tt.updates) // Verify the results - assert.Equal(t, len(tt.expectedHandlers), len(server.dnsMuxMap), + assert.Equal(t, len(tt.expectedHandlers), len(server.dnsMuxHandlers), "Number of handlers after update doesn't match expected") // Check each expected handler for id, expectedDomain := range tt.expectedHandlers { - handler, exists := server.dnsMuxMap[types.HandlerID(id)] - assert.True(t, exists, "Expected handler %s not found", id) - if exists { - assert.Equal(t, expectedDomain, handler.domain, + var found *handlerWrapper + for i := range server.dnsMuxHandlers { + if server.dnsMuxHandlers[i].handler.ID() == types.HandlerID(id) { + found = &server.dnsMuxHandlers[i] + break + } + } + assert.NotNil(t, found, "Expected handler %s not found", id) + if found != nil { + assert.Equal(t, expectedDomain, found.domain, "Domain mismatch for handler %s", id) } } // Verify no unexpected handlers exist - for HandlerID := range server.dnsMuxMap { - _, expected := tt.expectedHandlers[string(HandlerID)] - assert.True(t, expected, "Unexpected handler found: %s", HandlerID) + for _, entry := range server.dnsMuxHandlers { + _, expected := tt.expectedHandlers[string(entry.handler.ID())] + assert.True(t, expected, "Unexpected handler found: %s", entry.handler.ID()) } // Verify the handlerChain state and order @@ -1413,7 +1404,7 @@ func TestDefaultServer_UpdateMux(t *testing.T) { // Verify handler exists in mux foundInMux := false - for _, muxEntry := range server.dnsMuxMap { + for _, muxEntry := range server.dnsMuxHandlers { if chainEntry.Handler == muxEntry.handler && chainEntry.Priority == muxEntry.priority && chainEntry.Pattern == dns.Fqdn(muxEntry.domain) { @@ -1422,12 +1413,62 @@ func TestDefaultServer_UpdateMux(t *testing.T) { } } assert.True(t, foundInMux, - "Handler in chain not found in dnsMuxMap") + "Handler in chain not found in dnsMuxHandlers") } }) } } +// chainHasPattern reports whether the handler chain holds an entry registered +// for the given fqdn pattern at the given priority. +func chainHasPattern(s *DefaultServer, pattern string, priority int) bool { + for _, h := range s.handlerChain.handlers { + if h.OrigPattern == pattern && h.Priority == priority { + return true + } + } + return false +} + +// TestDefaultServer_UpdateMux_SharedHandlerZoneRemoval guards against the custom +// zone teardown leak: every custom zone is served by the same handler instance +// (the local resolver, whose ID is the constant "local-resolver"). updateMux +// must track each (handler, domain) registration independently, so that removing +// one zone deregisters exactly that zone's chain entry. Keying the registration +// set by handler ID alone collapsed all zones onto one entry, leaving removed +// zones leaked in the chain where they answered authoritatively with no records. +func TestDefaultServer_UpdateMux_SharedHandlerZoneRemoval(t *testing.T) { + // One handler serves every custom zone, mirroring s.localResolver. + shared := &mockHandler{Id: "local-resolver"} + + server := &DefaultServer{ + handlerChain: NewHandlerChain(), + service: &mockService{}, + } + + // Two custom zones under the same handler. The surviving peer zone is + // registered last, mirroring the management emission order. + server.updateMux([]handlerWrapper{ + {domain: "userzone.test", handler: shared, priority: PriorityLocal}, + {domain: "peerzone.test", handler: shared, priority: PriorityLocal}, + }) + + require.True(t, chainHasPattern(server, "userzone.test.", PriorityLocal), + "userzone.test should be registered after the first update") + require.True(t, chainHasPattern(server, "peerzone.test.", PriorityLocal), + "peerzone.test should be registered after the first update") + + // Remove the user zone, keep the peer zone (the dist-group revert). + server.updateMux([]handlerWrapper{ + {domain: "peerzone.test", handler: shared, priority: PriorityLocal}, + }) + + assert.True(t, chainHasPattern(server, "peerzone.test.", PriorityLocal), + "peerzone.test should remain after removing userzone.test") + assert.False(t, chainHasPattern(server, "userzone.test.", PriorityLocal), + "userzone.test handler must be deregistered, not leaked in the chain") +} + func TestExtraDomains(t *testing.T) { tests := []struct { name string @@ -2049,7 +2090,6 @@ func TestBuildUpstreamHandler_MergesGroupsPerDomain(t *testing.T) { localResolver: local.NewResolver(), handlerChain: NewHandlerChain(), hostManager: &noopHostConfigurator{}, - dnsMuxMap: make(registeredHandlerMap), } groups := []*nbdns.NameServerGroup{ @@ -2207,7 +2247,7 @@ func TestEvaluateNSGroupHealth(t *testing.T) { } } -// healthStubHandler is a minimal dnsMuxMap entry that exposes a fixed +// healthStubHandler is a minimal dnsMuxHandlers entry that exposes a fixed // UpstreamHealth snapshot, letting tests drive recomputeNSGroupStates // without spinning up real handlers. type healthStubHandler struct { @@ -2283,12 +2323,11 @@ func newProjTestFixture(t *testing.T) *projTestFixture { ctx: context.Background(), wgInterface: &mocWGIface{}, statusRecorder: recorder, - dnsMuxMap: make(registeredHandlerMap), selectedRoutes: func() route.HAMap { return fx.selected }, activeRoutes: func() route.HAMap { return fx.active }, warningDelayBase: defaultWarningDelayBase, } - fx.server.dnsMuxMap["example.com"] = handlerWrapper{domain: "example.com", handler: fx.stub, priority: PriorityUpstream} + fx.server.dnsMuxHandlers = []handlerWrapper{{domain: "example.com", handler: fx.stub, priority: PriorityUpstream}} fx.server.mux.Lock() fx.server.updateNSGroupStates([]*nbdns.NameServerGroup{fx.group}) @@ -2395,7 +2434,6 @@ func TestProjection_OverlayAddrNoRouteDelaysWarning(t *testing.T) { ctx: context.Background(), wgInterface: &mocWGIface{}, statusRecorder: recorder, - dnsMuxMap: make(registeredHandlerMap), selectedRoutes: func() route.HAMap { return nil }, activeRoutes: func() route.HAMap { return nil }, warningDelayBase: 50 * time.Millisecond, @@ -2407,7 +2445,7 @@ func TestProjection_OverlayAddrNoRouteDelaysWarning(t *testing.T) { stub := &healthStubHandler{health: map[netip.AddrPort]UpstreamHealth{ overlayPeer: {LastFail: time.Now(), LastErr: "timeout"}, }} - server.dnsMuxMap["example.com"] = handlerWrapper{domain: "example.com", handler: stub, priority: PriorityUpstream} + server.dnsMuxHandlers = []handlerWrapper{{domain: "example.com", handler: stub, priority: PriorityUpstream}} server.mux.Lock() server.updateNSGroupStates([]*nbdns.NameServerGroup{group}) @@ -2444,7 +2482,6 @@ func TestProjection_StopClearsHealthState(t *testing.T) { service: NewServiceViaMemory(wgIface), hostManager: &noopHostConfigurator{}, extraDomains: map[domain.Domain]int{}, - dnsMuxMap: make(registeredHandlerMap), statusRecorder: peer.NewRecorder("mgm"), selectedRoutes: func() route.HAMap { return nil }, activeRoutes: func() route.HAMap { return nil }, @@ -2459,7 +2496,7 @@ func TestProjection_StopClearsHealthState(t *testing.T) { NameServers: []nbdns.NameServer{{IP: srv.Addr(), NSType: nbdns.UDPNameServerType, Port: int(srv.Port())}}, } stub := &healthStubHandler{health: map[netip.AddrPort]UpstreamHealth{srv: {LastOk: time.Now()}}} - server.dnsMuxMap["example.com"] = handlerWrapper{domain: "example.com", handler: stub, priority: PriorityUpstream} + server.dnsMuxHandlers = []handlerWrapper{{domain: "example.com", handler: stub, priority: PriorityUpstream}} server.mux.Lock() server.updateNSGroupStates([]*nbdns.NameServerGroup{group}) @@ -2595,7 +2632,6 @@ func TestProjection_MixedGroupEmitsImmediately(t *testing.T) { server := &DefaultServer{ ctx: context.Background(), statusRecorder: recorder, - dnsMuxMap: make(registeredHandlerMap), selectedRoutes: func() route.HAMap { return overlayMap }, activeRoutes: func() route.HAMap { return nil }, warningDelayBase: time.Hour, @@ -2613,7 +2649,7 @@ func TestProjection_MixedGroupEmitsImmediately(t *testing.T) { overlay: {LastFail: time.Now(), LastErr: "timeout"}, }, } - server.dnsMuxMap["example.com"] = handlerWrapper{domain: "example.com", handler: stub, priority: PriorityUpstream} + server.dnsMuxHandlers = []handlerWrapper{{domain: "example.com", handler: stub, priority: PriorityUpstream}} server.mux.Lock() server.updateNSGroupStates([]*nbdns.NameServerGroup{group}) @@ -2640,7 +2676,6 @@ func TestDNSLoopPrevention(t *testing.T) { localResolver: local.NewResolver(), handlerChain: NewHandlerChain(), hostManager: &noopHostConfigurator{}, - dnsMuxMap: make(registeredHandlerMap), } tests := []struct {