Fix broken auto-generated Rego rule (#769)

Default Rego policy generated from the rules in some cases is broken.
This change fixes the Rego template for rules to generate policies.

Also, file store load constantly regenerates policy objects from rules.
It allows updating/fixing of the default Rego template during releases.

go.mod

@@ -56,6 +56,7 @@ require (
 	go.opentelemetry.io/otel/exporters/prometheus v0.33.0
 	go.opentelemetry.io/otel/metric v0.33.0
 	go.opentelemetry.io/otel/sdk/metric v0.33.0
+	golang.org/x/exp v0.0.0-20220518171630-0b5c67f07fdf
 	golang.org/x/net v0.8.0
 	golang.org/x/sync v0.1.0
 	golang.org/x/term v0.6.0
@@ -126,7 +127,6 @@ require (
 	go.opentelemetry.io/otel v1.11.1 // indirect
 	go.opentelemetry.io/otel/sdk v1.11.1 // indirect
 	go.opentelemetry.io/otel/trace v1.11.1 // indirect
-	golang.org/x/exp v0.0.0-20220518171630-0b5c67f07fdf // indirect
 	golang.org/x/image v0.0.0-20200430140353-33d19683fad8 // indirect
 	golang.org/x/mod v0.8.0 // indirect
 	golang.org/x/text v0.8.0 // indirect

management/server/file_store.go

@@ -121,15 +121,25 @@ func restore(file string) (*FileStore, error) {
 			store.PrivateDomain2AccountID[account.Domain] = accountID
 		}
 
-		// if no policies are defined, that means we need to migrate Rules to policies
+		// TODO: policy query generated from the Go template and rule object.
-		if len(account.Policies) == 0 {
+		//       We need to refactor this part to avoid using templating for policies queries building
+		//       and drop this migration part.
+		policies := make(map[string]int, len(account.Policies))
+		for i, policy := range account.Policies {
+			policies[policy.ID] = i
+		}
+		if account.Policies == nil {
 			account.Policies = make([]*Policy, 0)
-			for _, rule := range account.Rules {
+		}
-				policy, err := RuleToPolicy(rule)
+		for _, rule := range account.Rules {
-				if err != nil {
+			policy, err := RuleToPolicy(rule)
-					log.Errorf("unable to migrate rule to policy: %v", err)
+			if err != nil {
-					continue
+				log.Errorf("unable to migrate rule to policy: %v", err)
-				}
+				continue
+			}
+			if i, ok := policies[policy.ID]; ok {
+				account.Policies[i] = policy
+			} else {
 				account.Policies = append(account.Policies, policy)
 			}
 		}

management/server/policy.go

@@ -178,6 +178,9 @@ type FirewallRule struct {
 
 	// Port of the traffic
 	Port string
+
+	// id for internal purposes
+	id string
 }
 
 // parseFromRegoResult parses the Rego result to a FirewallRule.
@@ -218,39 +221,35 @@ func (f *FirewallRule) parseFromRegoResult(value interface{}) error {
 	f.Action = action
 	f.Port = port
 
+	// NOTE: update this id each time when new field added
+	f.id = peerID + peerIP + direction + action + port
+
 	return nil
 }
 
-// getRegoQuery returns a initialized Rego object with default rule.
+// queryPeersAndFwRulesByRego returns a list associated Peers and firewall rules list for this peer.
-func (a *Account) getRegoQuery() (rego.PreparedEvalQuery, error) {
+func (a *Account) queryPeersAndFwRulesByRego(
-	queries := []func(*rego.Rego){
+	peerID string,
-		rego.Query("data.netbird.all"),
+	queryNumber int,
-		rego.Module("netbird", defaultPolicyModule),
+	query string,
-	}
+) ([]*Peer, []*FirewallRule) {
-	for i, p := range a.Policies {
-		if !p.Enabled {
-			continue
-		}
-		queries = append(queries, rego.Module(fmt.Sprintf("netbird-%d", i), p.Query))
-	}
-	return rego.New(queries...).PrepareForEval(context.TODO())
-}
-
-// getPeersByPolicy returns all peers that given peer has access to.
-func (a *Account) getPeersByPolicy(peerID string) ([]*Peer, []*FirewallRule) {
 	input := map[string]interface{}{
 		"peer_id": peerID,
 		"peers":   a.Peers,
 		"groups":  a.Groups,
 	}
 
-	query, err := a.getRegoQuery()
+	stmt, err := rego.New(
+		rego.Query("data.netbird.all"),
+		rego.Module("netbird", defaultPolicyModule),
+		rego.Module(fmt.Sprintf("netbird-%d", queryNumber), query),
+	).PrepareForEval(context.TODO())
 	if err != nil {
 		log.WithError(err).Error("get Rego query")
 		return nil, nil
 	}
 
-	evalResult, err := query.Eval(
+	evalResult, err := stmt.Eval(
 		context.TODO(),
 		rego.EvalInput(input),
 	)
@@ -318,6 +317,33 @@ func (a *Account) getPeersByPolicy(peerID string) ([]*Peer, []*FirewallRule) {
 	return peers, rules
 }
 
+// getPeersByPolicy returns all peers that given peer has access to.
+func (a *Account) getPeersByPolicy(peerID string) (peers []*Peer, rules []*FirewallRule) {
+	peersSeen := make(map[string]struct{})
+	ruleSeen := make(map[string]struct{})
+	for i, policy := range a.Policies {
+		if !policy.Enabled {
+			continue
+		}
+		p, r := a.queryPeersAndFwRulesByRego(peerID, i, policy.Query)
+		for _, peer := range p {
+			if _, ok := peersSeen[peer.ID]; ok {
+				continue
+			}
+			peers = append(peers, peer)
+			peersSeen[peer.ID] = struct{}{}
+		}
+		for _, rule := range r {
+			if _, ok := ruleSeen[rule.id]; ok {
+				continue
+			}
+			rules = append(rules, rule)
+			ruleSeen[rule.id] = struct{}{}
+		}
+	}
+	return
+}
+
 // GetPolicy from the store
 func (am *DefaultAccountManager) GetPolicy(accountID, policyID, userID string) (*Policy, error) {
 	unlock := am.Store.AcquireAccountLock(accountID)

management/server/policy_test.go

@@ -5,63 +5,268 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/exp/slices"
 )
 
 func TestAccount_getPeersByPolicy(t *testing.T) {
 	account := &Account{
 		Peers: map[string]*Peer{
-			"peer1": {
+			"cfif97at2r9s73au3q00": {
-				ID: "peer1",
+				ID: "cfif97at2r9s73au3q00",
-				IP: net.IPv4(10, 20, 0, 1),
+				IP: net.ParseIP("100.65.14.88"),
 			},
-			"peer2": {
+			"cfif97at2r9s73au3q0g": {
-				ID: "peer2",
+				ID: "cfif97at2r9s73au3q0g",
-				IP: net.IPv4(10, 20, 0, 2),
+				IP: net.ParseIP("100.65.80.39"),
 			},
-			"peer3": {
+			"cfif97at2r9s73au3q10": {
-				ID: "peer3",
+				ID: "cfif97at2r9s73au3q10",
-				IP: net.IPv4(10, 20, 0, 3),
+				IP: net.ParseIP("100.65.254.139"),
+			},
+			"cfif97at2r9s73au3q20": {
+				ID: "cfif97at2r9s73au3q20",
+				IP: net.ParseIP("100.65.62.5"),
+			},
+			"cfj4tiqt2r9s73dmeun0": {
+				ID: "cfj4tiqt2r9s73dmeun0",
+				IP: net.ParseIP("100.65.32.206"),
+			},
+			"cg7h032t2r9s73cg5fk0": {
+				ID: "cg7h032t2r9s73cg5fk0",
+				IP: net.ParseIP("100.65.250.202"),
+			},
+			"cgcnkj2t2r9s73cg5vv0": {
+				ID: "cgcnkj2t2r9s73cg5vv0",
+				IP: net.ParseIP("100.65.13.186"),
+			},
+			"cgcol4qt2r9s73cg601g": {
+				ID: "cgcol4qt2r9s73cg601g",
+				IP: net.ParseIP("100.65.29.55"),
 			},
 		},
 		Groups: map[string]*Group{
-			"gid1": {
+			"cet9e92t2r9s7383ns20": {
-				ID:    "gid1",
+				ID:   "cet9e92t2r9s7383ns20",
-				Name:  "all",
+				Name: "All",
-				Peers: []string{"peer1", "peer2", "peer3"},
+				Peers: []string{
+					"cfif97at2r9s73au3q0g",
+					"cfif97at2r9s73au3q00",
+					"cfif97at2r9s73au3q20",
+					"cfif97at2r9s73au3q10",
+					"cfj4tiqt2r9s73dmeun0",
+					"cg7h032t2r9s73cg5fk0",
+					"cgcnkj2t2r9s73cg5vv0",
+					"cgcol4qt2r9s73cg601g",
+				},
+			},
+			"cev90bat2r9s7383o150": {
+				ID:   "cev90bat2r9s7383o150",
+				Name: "swarm",
+				Peers: []string{
+					"cfif97at2r9s73au3q0g",
+					"cfif97at2r9s73au3q00",
+					"cfif97at2r9s73au3q20",
+					"cfj4tiqt2r9s73dmeun0",
+					"cgcnkj2t2r9s73cg5vv0",
+					"cgcol4qt2r9s73cg601g",
+				},
 			},
 		},
 		Rules: map[string]*Rule{
-			"default": {
+			"cet9e92t2r9s7383ns2g": {
-				ID:          "default",
+				ID:          "cet9e92t2r9s7383ns2g",
-				Name:        "default",
+				Name:        "Default",
-				Description: "default",
+				Description: "This is a default rule that allows connections between all the resources",
-				Disabled:    false,
+				Source: []string{
-				Source:      []string{"gid1"},
+					"cet9e92t2r9s7383ns20",
-				Destination: []string{"gid1"},
+				},
+				Destination: []string{
+					"cet9e92t2r9s7383ns20",
+				},
+			},
+			"cev90bat2r9s7383o15g": {
+				ID:          "cev90bat2r9s7383o15g",
+				Name:        "Swarm",
+				Description: "",
+				Source: []string{
+					"cev90bat2r9s7383o150",
+					"cet9e92t2r9s7383ns20",
+				},
+				Destination: []string{
+					"cev90bat2r9s7383o150",
+				},
 			},
 		},
 	}
 
-	rule, err := RuleToPolicy(account.Rules["default"])
+	rule1, err := RuleToPolicy(account.Rules["cet9e92t2r9s7383ns2g"])
 	assert.NoError(t, err)
 
-	account.Policies = append(account.Policies, rule)
+	rule2, err := RuleToPolicy(account.Rules["cev90bat2r9s7383o15g"])
+	assert.NoError(t, err)
 
-	peers, firewallRules := account.getPeersByPolicy("peer1")
+	account.Policies = append(account.Policies, rule1, rule2)
-	assert.Len(t, peers, 2)
-	assert.Contains(t, peers, account.Peers["peer2"])
-	assert.Contains(t, peers, account.Peers["peer3"])
 
-	epectedFirewallRules := []*FirewallRule{
+	t.Run("check that all peers get map", func(t *testing.T) {
-		{PeerID: "peer1", PeerIP: "10.20.0.1", Direction: "dst", Action: "accept", Port: ""},
+		for _, p := range account.Peers {
-		{PeerID: "peer2", PeerIP: "10.20.0.2", Direction: "dst", Action: "accept", Port: ""},
+			peers, firewallRules := account.getPeersByPolicy(p.ID)
-		{PeerID: "peer3", PeerIP: "10.20.0.3", Direction: "dst", Action: "accept", Port: ""},
+			assert.GreaterOrEqual(t, len(peers), 2, "mininum number peers should present")
-		{PeerID: "peer1", PeerIP: "10.20.0.1", Direction: "src", Action: "accept", Port: ""},
+			assert.GreaterOrEqual(t, len(firewallRules), 2, "mininum number of firewall rules should present")
-		{PeerID: "peer2", PeerIP: "10.20.0.2", Direction: "src", Action: "accept", Port: ""},
+		}
-		{PeerID: "peer3", PeerIP: "10.20.0.3", Direction: "src", Action: "accept", Port: ""},
+	})
-	}
+
-	assert.Len(t, firewallRules, len(epectedFirewallRules))
+	t.Run("check first peer map details", func(t *testing.T) {
-	for i := range firewallRules {
+		peers, firewallRules := account.getPeersByPolicy("cfif97at2r9s73au3q0g")
-		assert.Equal(t, firewallRules[i], epectedFirewallRules[i])
+		assert.Len(t, peers, 7)
-	}
+		assert.Contains(t, peers, account.Peers["cfif97at2r9s73au3q00"])
+		assert.Contains(t, peers, account.Peers["cfif97at2r9s73au3q10"])
+		assert.Contains(t, peers, account.Peers["cfif97at2r9s73au3q20"])
+		assert.Contains(t, peers, account.Peers["cfj4tiqt2r9s73dmeun0"])
+		assert.Contains(t, peers, account.Peers["cg7h032t2r9s73cg5fk0"])
+
+		epectedFirewallRules := []*FirewallRule{
+			{
+				PeerID:    "cfif97at2r9s73au3q00",
+				PeerIP:    "100.65.14.88",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q00100.65.14.88srcaccept",
+			},
+			{
+				PeerID:    "cfif97at2r9s73au3q00",
+				PeerIP:    "100.65.14.88",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q00100.65.14.88dstaccept",
+			},
+
+			{
+				PeerID:    "cfif97at2r9s73au3q0g",
+				PeerIP:    "100.65.80.39",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q0g100.65.80.39dstaccept",
+			},
+			{
+				PeerID:    "cfif97at2r9s73au3q0g",
+				PeerIP:    "100.65.80.39",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q0g100.65.80.39srcaccept",
+			},
+
+			{
+				PeerID:    "cfif97at2r9s73au3q10",
+				PeerIP:    "100.65.254.139",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q10100.65.254.139dstaccept",
+			},
+			{
+				PeerID:    "cfif97at2r9s73au3q10",
+				PeerIP:    "100.65.254.139",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q10100.65.254.139srcaccept",
+			},
+
+			{
+				PeerID:    "cfif97at2r9s73au3q20",
+				PeerIP:    "100.65.62.5",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q20100.65.62.5dstaccept",
+			},
+			{
+				PeerID:    "cfif97at2r9s73au3q20",
+				PeerIP:    "100.65.62.5",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfif97at2r9s73au3q20100.65.62.5srcaccept",
+			},
+
+			{
+				PeerID:    "cfj4tiqt2r9s73dmeun0",
+				PeerIP:    "100.65.32.206",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfj4tiqt2r9s73dmeun0100.65.32.206dstaccept",
+			},
+			{
+				PeerID:    "cfj4tiqt2r9s73dmeun0",
+				PeerIP:    "100.65.32.206",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cfj4tiqt2r9s73dmeun0100.65.32.206srcaccept",
+			},
+
+			{
+				PeerID:    "cg7h032t2r9s73cg5fk0",
+				PeerIP:    "100.65.250.202",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cg7h032t2r9s73cg5fk0100.65.250.202dstaccept",
+			},
+			{
+				PeerID:    "cg7h032t2r9s73cg5fk0",
+				PeerIP:    "100.65.250.202",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cg7h032t2r9s73cg5fk0100.65.250.202srcaccept",
+			},
+
+			{
+				PeerID:    "cgcnkj2t2r9s73cg5vv0",
+				PeerIP:    "100.65.13.186",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cgcnkj2t2r9s73cg5vv0100.65.13.186dstaccept",
+			},
+			{
+				PeerID:    "cgcnkj2t2r9s73cg5vv0",
+				PeerIP:    "100.65.13.186",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cgcnkj2t2r9s73cg5vv0100.65.13.186srcaccept",
+			},
+
+			{
+				PeerID:    "cgcol4qt2r9s73cg601g",
+				PeerIP:    "100.65.29.55",
+				Direction: "dst",
+				Action:    "accept",
+				Port:      "",
+				id:        "cgcol4qt2r9s73cg601g100.65.29.55dstaccept",
+			},
+			{
+				PeerID:    "cgcol4qt2r9s73cg601g",
+				PeerIP:    "100.65.29.55",
+				Direction: "src",
+				Action:    "accept",
+				Port:      "",
+				id:        "cgcol4qt2r9s73cg601g100.65.29.55srcaccept",
+			},
+		}
+		assert.Len(t, firewallRules, len(epectedFirewallRules))
+		slices.SortFunc(firewallRules, func(a, b *FirewallRule) bool {
+			return a.PeerID < b.PeerID
+		})
+		for i := range firewallRules {
+			assert.Equal(t, epectedFirewallRules[i], firewallRules[i])
+		}
+	})
 }

management/server/rego/default_policy.rego

@@ -1,9 +1,9 @@
 package netbird
 
 all[rule] {
-	is_peer_in_any_group([{{range $i, $e := .All}}{{if $i}},{{end}}"{{$e}}"{{end}}])
+    is_peer_in_any_group([{{range $i, $e := .All}}{{if $i}},{{end}}"{{$e}}"{{end}}])
-	rule := array.concat(
+    rule := {
-		rules_from_groups([{{range $i, $e := .Destination}}{{if $i}},{{end}}"{{$e}}"{{end}}], "dst", "accept", ""),
+        {{range $i, $e := .Destination}}rules_from_group("{{$e}}", "dst", "accept", ""),{{end}}
-		rules_from_groups([{{range $i, $e := .Source}}{{if $i}},{{end}}"{{$e}}"{{end}}], "src", "accept", ""),
+        {{range $i, $e := .Source}}rules_from_group("{{$e}}", "src", "accept", ""),{{end}}
-	)[_]
+    }[_][_]
 }

management/server/rego/default_policy_module.rego

@@ -17,17 +17,11 @@ get_rule(peer_id, direction, action, port) := rule if {
     }
 }
 
-# peers_from_group returns a list of peer ids for a given group id
+# netbird_rules_from_group returns a list of netbird rules for a given group_id
-peers_from_group(group_id) := peers if {
+rules_from_group(group_id, direction, action, port) := rules if {
 	group := input.groups[_]
 	group.ID == group_id
-	peers := [peer | peer := group.Peers[_]]
+	rules := [get_rule(peer, direction, action, port) | peer := group.Peers[_]]
-}
-
-# 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)[_]]
 }
 
 # is_peer_in_any_group checks that input peer present at least in one group