[client] Set up firewall rules for dns routes dynamically based on dns response (#3702)

client/internal/dnsfwd/forwarder.go

@@ -3,6 +3,7 @@ package dnsfwd
 import (
 	"context"
 	"errors"
+	"fmt"
 	"math"
 	"net"
 	"net/netip"
@@ -10,11 +11,16 @@ import (
 	"sync"
 	"time"
 
+	"github.com/hashicorp/go-multierror"
 	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/peer"
 	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/domain"
+	"github.com/netbirdio/netbird/route"
 )
 
 const errResolveFailed = "failed to resolve query for domain=%s: %v"
@@ -23,25 +29,27 @@ const upstreamTimeout = 15 * time.Second
 type DNSForwarder struct {
 	listenAddress  string
 	ttl            uint32
-	domains        []string
 	statusRecorder *peer.Status
 
 	dnsServer *dns.Server
 	mux       *dns.ServeMux
 
-	resId sync.Map
+	mutex      sync.RWMutex
+	fwdEntries []*ForwarderEntry
+	firewall   firewall.Manager
 }
 
-func NewDNSForwarder(listenAddress string, ttl uint32, statusRecorder *peer.Status) *DNSForwarder {
+func NewDNSForwarder(listenAddress string, ttl uint32, firewall firewall.Manager, statusRecorder *peer.Status) *DNSForwarder {
 	log.Debugf("creating DNS forwarder with listen_address=%s ttl=%d", listenAddress, ttl)
 	return &DNSForwarder{
 		listenAddress:  listenAddress,
 		ttl:            ttl,
+		firewall:       firewall,
 		statusRecorder: statusRecorder,
 	}
 }
 
-func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error {
+func (f *DNSForwarder) Listen(entries []*ForwarderEntry) error {
 	log.Infof("listen DNS forwarder on address=%s", f.listenAddress)
 	mux := dns.NewServeMux()
 
@@ -53,31 +61,35 @@ func (f *DNSForwarder) Listen(domains []string, resIds map[string]string) error
 	f.dnsServer = dnsServer
 	f.mux = mux
 
-	f.UpdateDomains(domains, resIds)
+	f.UpdateDomains(entries)
 
 	return dnsServer.ListenAndServe()
 }
 
-func (f *DNSForwarder) UpdateDomains(domains []string, resIds map[string]string) {
-	log.Debugf("Updating domains from %v to %v", f.domains, domains)
+func (f *DNSForwarder) UpdateDomains(entries []*ForwarderEntry) {
+	f.mutex.Lock()
+	defer f.mutex.Unlock()
 
-	for _, d := range f.domains {
-		f.mux.HandleRemove(d)
+	if f.mux == nil {
+		log.Debug("DNS mux is nil, skipping domain update")
+		f.fwdEntries = entries
+		return
 	}
-	f.resId.Clear()
 
-	newDomains := filterDomains(domains)
+	oldDomains := filterDomains(f.fwdEntries)
+
+	for _, d := range oldDomains {
+		f.mux.HandleRemove(d.PunycodeString())
+	}
+
+	newDomains := filterDomains(entries)
 	for _, d := range newDomains {
-		f.mux.HandleFunc(d, f.handleDNSQuery)
+		f.mux.HandleFunc(d.PunycodeString(), f.handleDNSQuery)
 	}
 
-	for domain, resId := range resIds {
-		if domain != "" {
-			f.resId.Store(domain, resId)
-		}
-	}
+	f.fwdEntries = entries
 
-	f.domains = newDomains
+	log.Debugf("Updated domains from %v to %v", oldDomains, newDomains)
 }
 
 func (f *DNSForwarder) Close(ctx context.Context) error {
@@ -91,11 +103,11 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 	if len(query.Question) == 0 {
 		return
 	}
-	log.Tracef("received DNS request for DNS forwarder: domain=%v type=%v class=%v",
-		query.Question[0].Name, query.Question[0].Qtype, query.Question[0].Qclass)
-
 	question := query.Question[0]
-	domain := question.Name
+	log.Tracef("received DNS request for DNS forwarder: domain=%v type=%v class=%v",
+		question.Name, question.Qtype, question.Qclass)
+
+	domain := strings.ToLower(question.Name)
 
 	resp := query.SetReply(query)
 	var network string
@@ -122,21 +134,7 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 		return
 	}
 
-	resId := f.getResIdForDomain(strings.TrimSuffix(domain, "."))
-	if resId != "" {
-		for _, ip := range ips {
-			var ipWithSuffix string
-			if ip.Is4() {
-				ipWithSuffix = ip.String() + "/32"
-				log.Tracef("resolved domain=%s to IPv4=%s", domain, ipWithSuffix)
-			} else {
-				ipWithSuffix = ip.String() + "/128"
-				log.Tracef("resolved domain=%s to IPv6=%s", domain, ipWithSuffix)
-			}
-			f.statusRecorder.AddResolvedIPLookupEntry(ipWithSuffix, resId)
-		}
-	}
-
+	f.updateInternalState(domain, ips)
 	f.addIPsToResponse(resp, domain, ips)
 
 	if err := w.WriteMsg(resp); err != nil {
@@ -144,6 +142,42 @@ func (f *DNSForwarder) handleDNSQuery(w dns.ResponseWriter, query *dns.Msg) {
 	}
 }
 
+func (f *DNSForwarder) updateInternalState(domain string, ips []netip.Addr) {
+	var prefixes []netip.Prefix
+	mostSpecificResId, matchingEntries := f.getMatchingEntries(strings.TrimSuffix(domain, "."))
+	if mostSpecificResId != "" {
+		for _, ip := range ips {
+			var prefix netip.Prefix
+			if ip.Is4() {
+				prefix = netip.PrefixFrom(ip, 32)
+			} else {
+				prefix = netip.PrefixFrom(ip, 128)
+			}
+			prefixes = append(prefixes, prefix)
+			f.statusRecorder.AddResolvedIPLookupEntry(prefix, mostSpecificResId)
+		}
+	}
+
+	if f.firewall != nil {
+		f.updateFirewall(matchingEntries, prefixes)
+	}
+}
+
+func (f *DNSForwarder) updateFirewall(matchingEntries []*ForwarderEntry, prefixes []netip.Prefix) {
+	var merr *multierror.Error
+	for _, entry := range matchingEntries {
+		if err := f.firewall.UpdateSet(entry.Set, prefixes); err != nil {
+			merr = multierror.Append(merr, fmt.Errorf("update set for domain=%s: %w", entry.Domain, err))
+		}
+	}
+	if merr != nil {
+		log.Errorf("failed to update firewall sets (%d/%d): %v",
+			len(merr.Errors),
+			len(matchingEntries),
+			nberrors.FormatErrorOrNil(merr))
+	}
+}
+
 // handleDNSError processes DNS lookup errors and sends an appropriate error response
 func (f *DNSForwarder) handleDNSError(w dns.ResponseWriter, resp *dns.Msg, domain string, err error) {
 	var dnsErr *net.DNSError
@@ -204,45 +238,53 @@ func (f *DNSForwarder) addIPsToResponse(resp *dns.Msg, domain string, ips []neti
 	}
 }
 
-func (f *DNSForwarder) getResIdForDomain(domain string) string {
-	var selectedResId string
+// getMatchingEntries retrieves the resource IDs for a given domain.
+// It returns the most specific match and all matching resource IDs.
+func (f *DNSForwarder) getMatchingEntries(domain string) (route.ResID, []*ForwarderEntry) {
+	var selectedResId route.ResID
 	var bestScore int
+	var matches []*ForwarderEntry
 
-	f.resId.Range(func(key, value interface{}) bool {
+	f.mutex.RLock()
+	defer f.mutex.RUnlock()
+
+	for _, entry := range f.fwdEntries {
 		var score int
-		pattern := key.(string)
+		pattern := entry.Domain.PunycodeString()
 
 		switch {
 		case strings.HasPrefix(pattern, "*."):
 			baseDomain := strings.TrimPrefix(pattern, "*.")
-			if domain == baseDomain || strings.HasSuffix(domain, "."+baseDomain) {
+
+			if strings.EqualFold(domain, baseDomain) || strings.HasSuffix(domain, "."+baseDomain) {
 				score = len(baseDomain)
+				matches = append(matches, entry)
 			}
 		case domain == pattern:
 			score = math.MaxInt
+			matches = append(matches, entry)
 		default:
-			return true
+			continue
 		}
 
 		if score > bestScore {
 			bestScore = score
-			selectedResId = value.(string)
+			selectedResId = entry.ResID
 		}
-		return true
-	})
+	}
 
-	return selectedResId
+	return selectedResId, matches
 }
 
 // filterDomains returns a list of normalized domains
-func filterDomains(domains []string) []string {
-	newDomains := make([]string, 0, len(domains))
-	for _, d := range domains {
-		if d == "" {
+func filterDomains(entries []*ForwarderEntry) domain.List {
+	newDomains := make(domain.List, 0, len(entries))
+	for _, d := range entries {
+		if d.Domain == "" {
 			log.Warn("empty domain in DNS forwarder")
 			continue
 		}
-		newDomains = append(newDomains, nbdns.NormalizeZone(d))
+		newDomains = append(newDomains, domain.Domain(nbdns.NormalizeZone(d.Domain.PunycodeString())))
 	}
 	return newDomains
 }

client/internal/dnsfwd/forwarder_test.go

@@ -1,56 +1,61 @@
 package dnsfwd
 
 import (
-	"sync"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/domain"
+	"github.com/netbirdio/netbird/route"
 )
 
-func TestGetResIdForDomain(t *testing.T) {
+func Test_getMatchingEntries(t *testing.T) {
 	testCases := []struct {
 		name           string
-		storedMappings map[string]string // key: domain pattern, value: resId
+		storedMappings map[string]route.ResID // key: domain pattern, value: resId
 		queryDomain    string
-		expectedResId  string
+		expectedResId  route.ResID
 	}{
 		{
 			name:           "Empty map returns empty string",
-			storedMappings: map[string]string{},
+			storedMappings: map[string]route.ResID{},
 			queryDomain:    "example.com",
 			expectedResId:  "",
 		},
 		{
 			name:           "Exact match returns stored resId",
-			storedMappings: map[string]string{"example.com": "res1"},
+			storedMappings: map[string]route.ResID{"example.com": "res1"},
 			queryDomain:    "example.com",
 			expectedResId:  "res1",
 		},
 		{
 			name:           "Wildcard pattern matches base domain",
-			storedMappings: map[string]string{"*.example.com": "res2"},
+			storedMappings: map[string]route.ResID{"*.example.com": "res2"},
 			queryDomain:    "example.com",
 			expectedResId:  "res2",
 		},
 		{
 			name:           "Wildcard pattern matches subdomain",
-			storedMappings: map[string]string{"*.example.com": "res3"},
+			storedMappings: map[string]route.ResID{"*.example.com": "res3"},
 			queryDomain:    "foo.example.com",
 			expectedResId:  "res3",
 		},
 		{
 			name:           "Wildcard pattern does not match different domain",
-			storedMappings: map[string]string{"*.example.com": "res4"},
+			storedMappings: map[string]route.ResID{"*.example.com": "res4"},
 			queryDomain:    "foo.notexample.com",
 			expectedResId:  "",
 		},
 		{
 			name:           "Non-wildcard pattern does not match subdomain",
-			storedMappings: map[string]string{"example.com": "res5"},
+			storedMappings: map[string]route.ResID{"example.com": "res5"},
 			queryDomain:    "foo.example.com",
 			expectedResId:  "",
 		},
 		{
 			name: "Exact match over overlapping wildcard",
-			storedMappings: map[string]string{
+			storedMappings: map[string]route.ResID{
 				"*.example.com":   "resWildcard",
 				"foo.example.com": "resExact",
 			},
@@ -59,7 +64,7 @@ func TestGetResIdForDomain(t *testing.T) {
 		},
 		{
 			name: "Overlapping wildcards: Select more specific wildcard",
-			storedMappings: map[string]string{
+			storedMappings: map[string]route.ResID{
 				"*.example.com":     "resA",
 				"*.sub.example.com": "resB",
 			},
@@ -68,7 +73,7 @@ func TestGetResIdForDomain(t *testing.T) {
 		},
 		{
 			name: "Wildcard multi-level subdomain match",
-			storedMappings: map[string]string{
+			storedMappings: map[string]route.ResID{
 				"*.example.com": "resMulti",
 			},
 			queryDomain:   "a.b.example.com",
@@ -78,18 +83,21 @@ func TestGetResIdForDomain(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			fwd := &DNSForwarder{
-				resId: sync.Map{},
-			}
+			fwd := &DNSForwarder{}
 
+			var entries []*ForwarderEntry
 			for domainPattern, resId := range tc.storedMappings {
-				fwd.resId.Store(domainPattern, resId)
+				d, err := domain.FromString(domainPattern)
+				require.NoError(t, err)
+				entries = append(entries, &ForwarderEntry{
+					Domain: d,
+					ResID:  resId,
+				})
 			}
+			fwd.UpdateDomains(entries)
 
-			got := fwd.getResIdForDomain(tc.queryDomain)
-			if got != tc.expectedResId {
-				t.Errorf("For query domain %q, expected resId %q, but got %q", tc.queryDomain, tc.expectedResId, got)
-			}
+			got, _ := fwd.getMatchingEntries(tc.queryDomain)
+			assert.Equal(t, got, tc.expectedResId)
 		})
 	}
 }

client/internal/dnsfwd/manager.go

@@ -11,6 +11,8 @@ import (
 	nberrors "github.com/netbirdio/netbird/client/errors"
 	firewall "github.com/netbirdio/netbird/client/firewall/manager"
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/management/domain"
+	"github.com/netbirdio/netbird/route"
 )
 
 const (
@@ -19,6 +21,13 @@ const (
 	dnsTTL     = 60 //seconds
 )
 
+// ForwarderEntry is a mapping from a domain to a resource ID and a hash of the parent domain list.
+type ForwarderEntry struct {
+	Domain domain.Domain
+	ResID  route.ResID
+	Set    firewall.Set
+}
+
 type Manager struct {
 	firewall       firewall.Manager
 	statusRecorder *peer.Status
@@ -34,7 +43,7 @@ func NewManager(fw firewall.Manager, statusRecorder *peer.Status) *Manager {
 	}
 }
 
-func (m *Manager) Start(domains []string, resIds map[string]string) error {
+func (m *Manager) Start(fwdEntries []*ForwarderEntry) error {
 	log.Infof("starting DNS forwarder")
 	if m.dnsForwarder != nil {
 		return nil
@@ -44,9 +53,9 @@ func (m *Manager) Start(domains []string, resIds map[string]string) error {
 		return err
 	}
 
-	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.statusRecorder)
+	m.dnsForwarder = NewDNSForwarder(fmt.Sprintf(":%d", ListenPort), dnsTTL, m.firewall, m.statusRecorder)
 	go func() {
-		if err := m.dnsForwarder.Listen(domains, resIds); err != nil {
+		if err := m.dnsForwarder.Listen(fwdEntries); err != nil {
 			// todo handle close error if it is exists
 			log.Errorf("failed to start DNS forwarder, err: %v", err)
 		}
@@ -55,12 +64,12 @@ func (m *Manager) Start(domains []string, resIds map[string]string) error {
 	return nil
 }
 
-func (m *Manager) UpdateDomains(domains []string, resIds map[string]string) {
+func (m *Manager) UpdateDomains(entries []*ForwarderEntry) {
 	if m.dnsForwarder == nil {
 		return
 	}
 
-	m.dnsForwarder.UpdateDomains(domains, resIds)
+	m.dnsForwarder.UpdateDomains(entries)
 }
 
 func (m *Manager) Stop(ctx context.Context) error {
@@ -81,34 +90,34 @@ func (m *Manager) Stop(ctx context.Context) error {
 	return nberrors.FormatErrorOrNil(mErr)
 }
 
-func (h *Manager) allowDNSFirewall() error {
+func (m *Manager) allowDNSFirewall() error {
 	dport := &firewall.Port{
 		IsRange: false,
 		Values:  []uint16{ListenPort},
 	}
 
-	if h.firewall == nil {
+	if m.firewall == nil {
 		return nil
 	}
 
-	dnsRules, err := h.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
+	dnsRules, err := m.firewall.AddPeerFiltering(nil, net.IP{0, 0, 0, 0}, firewall.ProtocolUDP, nil, dport, firewall.ActionAccept, "")
 	if err != nil {
 		log.Errorf("failed to add allow DNS router rules, err: %v", err)
 		return err
 	}
-	h.fwRules = dnsRules
+	m.fwRules = dnsRules
 
 	return nil
 }
 
-func (h *Manager) dropDNSFirewall() error {
+func (m *Manager) dropDNSFirewall() error {
 	var mErr *multierror.Error
-	for _, rule := range h.fwRules {
-		if err := h.firewall.DeletePeerRule(rule); err != nil {
+	for _, rule := range m.fwRules {
+		if err := m.firewall.DeletePeerRule(rule); err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf("failed to delete DNS router rules, err: %v", err))
 		}
 	}
 
-	h.fwRules = nil
+	m.fwRules = nil
 	return nberrors.FormatErrorOrNil(mErr)
 }