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
-		if len(account.Policies) == 0 {
+		// TODO: policy query generated from the Go template and rule object.
+		//       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)
-				if err != nil {
-					log.Errorf("unable to migrate rule to policy: %v", err)
-					continue
-				}
+		}
+		for _, rule := range account.Rules {
+			policy, err := RuleToPolicy(rule)
+			if err != nil {
+				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
 
-	return nil
-}
+	// NOTE: update this id each time when new field added
+	f.id = peerID + peerIP + direction + action + port
 
-// getRegoQuery returns a initialized Rego object with default rule.
-func (a *Account) getRegoQuery() (rego.PreparedEvalQuery, error) {
-	queries := []func(*rego.Rego){
-		rego.Query("data.netbird.all"),
-		rego.Module("netbird", defaultPolicyModule),
-	}
-	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())
+	return nil
 }
 
-// getPeersByPolicy returns all peers that given peer has access to.
-func (a *Account) getPeersByPolicy(peerID string) ([]*Peer, []*FirewallRule) {
+// queryPeersAndFwRulesByRego returns a list associated Peers and firewall rules list for this peer.
+func (a *Account) queryPeersAndFwRulesByRego(
+	peerID string,
+	queryNumber int,
+	query 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)