Fix all management tests

management/server/account.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"math/rand"
@@ -600,24 +601,33 @@ func (a *Account) GetPeer(peerID string) *Peer {
 }
 
 // ruleToPolicy converts a Rule to a Policy query object
-func (a *Account) ruleToPolicy(rule *Rule) *Policy {
-	getGroups := func(groups []string) (result string) {
-		for i := range groups {
-			groups[i] = fmt.Sprintf(`"%s"`, strings.Trim(groups[i], `"`))
-		}
-		return strings.Join(groups, ",")
+func (a *Account) ruleToPolicy(rule *Rule) (*Policy, error) {
+	input := struct {
+		All         []string
+		Source      []string
+		Destination []string
+	}{
+		All:         append(append([]string{}, rule.Destination...), rule.Source...),
+		Source:      rule.Source,
+		Destination: rule.Destination,
+	}
+	buff := new(bytes.Buffer)
+	if err := defaultPolicyTemplate.Execute(buff, input); err != nil {
+		return nil, err
 	}
 	return &Policy{
 		ID:          rule.ID,
 		Name:        rule.Name,
 		Description: rule.Description,
-		Query: fmt.Sprintf(
-			defaultPolicy,
-			getGroups(rule.Destination),
-			getGroups(rule.Source),
-		),
-		Disabled: rule.Disabled,
-	}
+		Query:       buff.String(),
+		Disabled:    rule.Disabled,
+		Meta: &PolicyMeta{
+			Action:       PolicyTrafficActionAccept,
+			Destinations: rule.Destination,
+			Sources:      rule.Source,
+			Port:         "",
+		},
+	}, nil
 }
 
 // BuildManager creates a new DefaultAccountManager with a provided Store
@@ -656,7 +666,9 @@ func BuildManager(store Store, peersUpdateManager *PeersUpdateManager, idpManage
 
 		_, err := account.GetGroupAll()
 		if err != nil {
-			addAllGroup(account)
+			if err := addAllGroup(account); err != nil {
+				return nil, err
+			}
 			shouldSave = true
 		}
 
@@ -1267,7 +1279,7 @@ func (am *DefaultAccountManager) GetDNSDomain() string {
 }
 
 // addAllGroup to account object if it doesn't exists
-func addAllGroup(account *Account) {
+func addAllGroup(account *Account) error {
 	if len(account.Groups) == 0 {
 		allGroup := &Group{
 			ID:   xid.New().String(),
@@ -1288,8 +1300,13 @@ func addAllGroup(account *Account) {
 		}
 		account.Rules = map[string]*Rule{defaultRule.ID: defaultRule}
 
-		account.Policies = []*Policy{account.ruleToPolicy(defaultRule)}
+		defaultPolicy, err := account.ruleToPolicy(defaultRule)
+		if err != nil {
+			return fmt.Errorf("convert rule to policy: %w", err)
+		}
+		account.Policies = []*Policy{defaultPolicy}
 	}
+	return nil
 }
 
 // newAccountWithId creates a new Account with a default SetupKey (doesn't store in a Store) and provided id
@@ -1330,7 +1347,9 @@ func newAccountWithId(accountId, userId, domain string) *Account {
 		},
 	}
 
-	addAllGroup(acc)
+	if err := addAllGroup(acc); err != nil {
+		log.Errorf("error adding all group to account %s: %v", acc.Id, err)
+	}
 	return acc
 }
 

management/server/file_store.go

@@ -1,7 +1,6 @@
 package server
 
 import (
-	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
@@ -118,22 +117,10 @@ func restore(file string) (*FileStore, error) {
 		if len(account.Policies) == 0 {
 			account.Policies = make([]*Policy, 0)
 			for _, rule := range account.Rules {
-				policy := &Policy{
-					ID:          xid.New().String(),
-					Name:        rule.Name,
-					Disabled:    rule.Disabled,
-					Description: rule.Description,
-					Query: fmt.Sprintf(
-						defaultPolicy,
-						strings.Join(rule.Destination, ","),
-						strings.Join(rule.Source, ","),
-					),
-					Meta: &PolicyMeta{
-						Action:       PolicyTrafficActionAccept,
-						Destinations: rule.Destination,
-						Sources:      rule.Source,
-						Port:         "",
-					},
+				policy, err := account.ruleToPolicy(rule)
+				if err != nil {
+					log.Errorf("unable to migrate rule to policy: %v", err)
+					continue
 				}
 				account.Policies = append(account.Policies, policy)
 			}

management/server/file_store_test.go

@@ -1,7 +1,6 @@
 package server
 
 import (
-	"fmt"
 	"net"
 	"path/filepath"
 	"testing"
@@ -175,9 +174,9 @@ func TestRestorePolicies_Migration(t *testing.T) {
 	require.Equal(t, policy.Description,
 		"This is a default rule that allows connections between all the resources",
 		"failed to restore a FileStore file - missing Account Policies Description")
-	require.Equal(t, policy.Query,
-		fmt.Sprintf(defaultPolicy, "cfefqs706sqkneg59g3g", "cfefqs706sqkneg59g3g"),
-		"failed to restore a FileStore file - missing Account Policies Query")
+	expectedPolicy, err := account.ruleToPolicy(account.Rules["cfefqs706sqkneg59g40"])
+	require.NoError(t, err, "failed to restore a FileStore file - missing Account Policies Rule")
+	require.Equal(t, policy.Query, expectedPolicy.Query, "failed to restore a FileStore file - missing Account Policies Query")
 	require.NotNil(t, policy.Meta, "failed to restore a FileStore file - missing Account Policies Meta")
 	require.Equal(t, policy.Meta.Action, PolicyTrafficActionAccept, "failed to restore a FileStore file - missing Account Policies Action")
 	require.Equal(t, policy.Meta.Destinations,

management/server/policy.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	_ "embed"
 	"fmt"
+	"html/template"
 	"strings"
 
 	"github.com/netbirdio/netbird/management/server/activity"
@@ -47,7 +48,10 @@ type PolicyUpdateOperation struct {
 var defaultPolicyModule string
 
 //go:embed rego/default_policy.rego
-var defaultPolicy string
+var defaultPolicyText string
+
+// defaultPolicyTemplate is a template for the default policy
+var defaultPolicyTemplate = template.Must(template.New("policy").Parse(defaultPolicyText))
 
 // PolicyMeta is the metadata of the policy
 type PolicyMeta struct {
@@ -181,6 +185,9 @@ func (a *Account) getRegoQuery(policies ...*Policy) (rego.PreparedEvalQuery, err
 		rego.Module("netbird", defaultPolicyModule),
 	}
 	for i, p := range policies {
+		if p.Disabled {
+			continue
+		}
 		queries = append(queries, rego.Module(fmt.Sprintf("netbird-%d", i), p.Query))
 	}
 	return rego.New(queries...).PrepareForEval(context.TODO())

management/server/policy_test.go

@@ -1,7 +1,6 @@
 package server
 
 import (
-	"fmt"
 	"net"
 	"testing"
 
@@ -31,17 +30,23 @@ func TestAccount_getPeersByPolicy(t *testing.T) {
 				Peers: []string{"peer1", "peer2", "peer3"},
 			},
 		},
-		Policies: []*Policy{
-			{
+		Rules: map[string]*Rule{
+			"default": {
 				ID:          "default",
 				Name:        "default",
 				Description: "default",
 				Disabled:    false,
-				Query:       fmt.Sprintf(defaultPolicy, `"gid1"`, `"gid1"`),
+				Source:      []string{"gid1"},
+				Destination: []string{"gid1"},
 			},
 		},
 	}
 
+	rule, err := account.ruleToPolicy(account.Rules["default"])
+	assert.NoError(t, err)
+
+	account.Policies = append(account.Policies, rule)
+
 	peers, firewallRules := account.getPeersByPolicy("peer1")
 	expected := []*Peer{account.Peers["peer2"], account.Peers["peer3"]}
 	assert.Equal(t, peers, expected)

management/server/rego/default_policy.rego

@@ -2,7 +2,9 @@ package netbird
 
 # all rules for peer connectivity and firewall
 all[rule] {
-  rule := array.concat(
-    rules_from_groups([%s], "dst", "accept", ""),
-    rules_from_groups([%s], "src", "accept", ""))[_]
+	is_peer_in_any_group([{{range $i, $e := .All}}{{if $i}},{{end}}"{{$e}}"{{end}}])
+	rule := array.concat(
+		rules_from_groups([{{range $i, $e := .Destination}}{{if $i}},{{end}}"{{$e}}"{{end}}], "dst", "accept", ""),
+		rules_from_groups([{{range $i, $e := .Source}}{{if $i}},{{end}}"{{$e}}"{{end}}], "src", "accept", ""),
+	)[_]
 }

management/server/rego/default_policy_module.rego

@@ -6,32 +6,35 @@ import future.keywords.contains
 
 # get_rule builds a netbird rule object from given parameters
 get_rule(peer_id, direction, action, port) := rule if {
-	peer := input.peers[_]
+    peer := input.peers[_]
     peer.ID == peer_id
-	rule := {
-		"ID": peer.ID,
-		"IP": peer.IP,
-		"Direction": direction,
-		"Action": action,
-		"Port": port,
-	}
-}
-
-# is_peer_group returns group by id if peer_id present in group peers
-get_peer_group(peer_id, group_id) := group if {
-	group := input.groups[_]
-	group.ID == group_id
+    rule := {
+        "ID": peer.ID,
+        "IP": peer.IP,
+        "Direction": direction,
+        "Action": action,
+        "Port": port,
+    }
 }
 
 # peers_from_group returns a list of peer ids for a given group id
 peers_from_group(group_id) := peers if {
-  group := get_peer_group(input.peer_id, group_id)
-  group.ID == group_id
-  peers := [ peer | peer := group.Peers[_]; peer == input.peer_id ]
+	group := input.groups[_]
+	group.ID == group_id
+	peers := [peer | peer := group.Peers[_]; peer != input.peer_id]
 }
 
 # netbird_rules_from_groups returns a list of netbird rules for a given list of group names
 rules_from_groups(groups, direction, action, port) := rules if {
-  group_id := groups[_]
-  rules := [ get_rule(peer, direction, action, port) | peer := peers_from_group(group_id)[_] ]
+	group_id := groups[_]
+	rules := [get_rule(peer, direction, action, port) | peer := peers_from_group(group_id)[_]]
+}
+
+# is_peer_in_any_group checks that input peer present at least in one group
+is_peer_in_any_group(groups) := count([group_id]) > 0 if {
+	group_id := groups[_]
+	group := input.groups[_]
+	group.ID == group_id
+	peer := group.Peers[_]
+	peer == input.peer_id
 }

management/server/rule.go

@@ -139,7 +139,11 @@ func (am *DefaultAccountManager) SaveRule(accountID, userID string, rule *Rule)
 
 	account.Rules[rule.ID] = rule
 	// temporary solution until we drop rules support
-	_ = am.savePolicy(account, account.ruleToPolicy(rule))
+	policy, err := account.ruleToPolicy(rule)
+	if err != nil {
+		return err
+	}
+	_ = am.savePolicy(account, policy)
 
 	account.Network.IncSerial()
 	if err = am.Store.SaveAccount(account); err != nil {