From a366d9e20813298d212cb2cf9d283950ee73ff39 Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Wed, 28 Jun 2023 17:29:02 +0200
Subject: [PATCH] Prevent sending nameserver configuration when peer is set as
 NS (#962)

* Prevent sending nameserver configuration when peer is set as NS

* Add DNS filter tests
---
 management/server/dns.go      | 22 ++++++++++++++++++----
 management/server/dns_test.go | 26 ++++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/management/server/dns.go b/management/server/dns.go
index 045608650..427ba40d1 100644
--- a/management/server/dns.go
+++ b/management/server/dns.go
@@ -2,13 +2,15 @@ package server
 
 import (
 	"fmt"
+	"strconv"
+
 	"github.com/miekg/dns"
+	log "github.com/sirupsen/logrus"
+
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
-	log "github.com/sirupsen/logrus"
-	"strconv"
 )
 
 const defaultTTL = 300
@@ -199,8 +201,10 @@ func getPeerNSGroups(account *Account, peerID string) []*nbdns.NameServerGroup {
 		for _, gID := range nsGroup.Groups {
 			_, found := groupList[gID]
 			if found {
-				peerNSGroups = append(peerNSGroups, nsGroup.Copy())
-				break
+				if !peerIsNameserver(account.GetPeer(peerID), nsGroup) {
+					peerNSGroups = append(peerNSGroups, nsGroup.Copy())
+					break
+				}
 			}
 		}
 	}
@@ -208,6 +212,16 @@ func getPeerNSGroups(account *Account, peerID string) []*nbdns.NameServerGroup {
 	return peerNSGroups
 }
 
+// peerIsNameserver returns true if the peer is a nameserver for a nsGroup
+func peerIsNameserver(peer *Peer, nsGroup *nbdns.NameServerGroup) bool {
+	for _, ns := range nsGroup.NameServers {
+		if peer.IP.Equal(ns.IP.AsSlice()) {
+			return true
+		}
+	}
+	return false
+}
+
 func addPeerLabelsToAccount(account *Account, peerLabels lookupMap) {
 	for _, peer := range account.Peers {
 		label, err := getPeerHostLabel(peer.Name, peerLabels)
diff --git a/management/server/dns_test.go b/management/server/dns_test.go
index d3ef464e6..092c52afa 100644
--- a/management/server/dns_test.go
+++ b/management/server/dns_test.go
@@ -1,10 +1,12 @@
 package server
 
 import (
+	"net/netip"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
 )
@@ -17,6 +19,7 @@ const (
 	dnsAccountID     = "testingAcc"
 	dnsAdminUserID   = "testingAdminUser"
 	dnsRegularUserID = "testingRegularUser"
+	dnsNSGroup1      = "ns1"
 )
 
 func TestGetDNSSettings(t *testing.T) {
@@ -163,6 +166,7 @@ func TestGetNetworkMap_DNSConfigSync(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, newAccountDNSConfig.DNSConfig.CustomZones, 1, "default DNS config should have one custom zone for peers")
 	require.True(t, newAccountDNSConfig.DNSConfig.ServiceEnable, "default DNS config should have local DNS service enabled")
+	require.Len(t, newAccountDNSConfig.DNSConfig.NameServerGroups, 0, "updated DNS config should have no nameserver groups since peer 1 is NS for the only existing NS group")
 
 	dnsSettings := account.DNSSettings.Copy()
 	dnsSettings.DisabledManagementGroups = append(dnsSettings.DisabledManagementGroups, dnsGroup1ID)
@@ -174,11 +178,11 @@ func TestGetNetworkMap_DNSConfigSync(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, updatedAccountDNSConfig.DNSConfig.CustomZones, 0, "updated DNS config should have no custom zone when peer belongs to a disabled group")
 	require.False(t, updatedAccountDNSConfig.DNSConfig.ServiceEnable, "updated DNS config should have local DNS service disabled when peer belongs to a disabled group")
-
 	peer2AccountDNSConfig, err := am.GetNetworkMap(peer2.ID)
 	require.NoError(t, err)
 	require.Len(t, peer2AccountDNSConfig.DNSConfig.CustomZones, 1, "DNS config should have one custom zone for peers not in the disabled group")
 	require.True(t, peer2AccountDNSConfig.DNSConfig.ServiceEnable, "DNS config should have DNS service enabled for peers not in the disabled group")
+	require.Len(t, peer2AccountDNSConfig.DNSConfig.NameServerGroups, 1, "updated DNS config should have 1 nameserver groups since peer 2 is part of the group All")
 }
 
 func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
@@ -246,7 +250,7 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 		return nil, err
 	}
 
-	_, _, err = am.AddPeer("", dnsAdminUserID, peer1)
+	savedPeer1, _, err := am.AddPeer("", dnsAdminUserID, peer1)
 	if err != nil {
 		return nil, err
 	}
@@ -284,6 +288,24 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 	account.Groups[newGroup1.ID] = newGroup1
 	account.Groups[newGroup2.ID] = newGroup2
 
+	allGroup, err := account.GetGroupAll()
+	if err != nil {
+		return nil, err
+	}
+
+	account.NameServerGroups[dnsNSGroup1] = &dns.NameServerGroup{
+		ID:   dnsNSGroup1,
+		Name: "ns-group-1",
+		NameServers: []dns.NameServer{{
+			IP:     netip.MustParseAddr(savedPeer1.IP.String()),
+			NSType: dns.UDPNameServerType,
+			Port:   dns.DefaultDNSPort,
+		}},
+		Primary: true,
+		Enabled: true,
+		Groups:  []string{allGroup.ID},
+	}
+
 	err = am.Store.SaveAccount(account)
 	if err != nil {
 		return nil, err