Refactor policy get and save account to use store methods

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
This commit is contained in:
bcmmbaga 2024-11-12 11:17:16 +03:00
parent ffce48ca5f
commit 0c0fd380bd
No known key found for this signature in database
GPG Key ID: 511EED5C928AD547
8 changed files with 308 additions and 176 deletions

View File

@ -113,7 +113,7 @@ type AccountManager interface {
GroupAddPeer(ctx context.Context, accountId, groupID, peerID string) error GroupAddPeer(ctx context.Context, accountId, groupID, peerID string) error
GroupDeletePeer(ctx context.Context, accountId, groupID, peerID string) error GroupDeletePeer(ctx context.Context, accountId, groupID, peerID string) error
GetPolicy(ctx context.Context, accountID, policyID, userID string) (*Policy, error) GetPolicy(ctx context.Context, accountID, policyID, userID string) (*Policy, error)
SavePolicy(ctx context.Context, accountID, userID string, policy *Policy, isUpdate bool) error SavePolicy(ctx context.Context, accountID, userID string, policy *Policy) (*Policy, error)
DeletePolicy(ctx context.Context, accountID, policyID, userID string) error DeletePolicy(ctx context.Context, accountID, policyID, userID string) error
ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error) ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error)
GetRoute(ctx context.Context, accountID string, routeID route.ID, userID string) (*route.Route, error) GetRoute(ctx context.Context, accountID string, routeID route.ID, userID string) (*route.Route, error)

View File

@ -1238,8 +1238,7 @@ func TestAccountManager_NetworkUpdates_SaveGroup(t *testing.T) {
return return
} }
policy := Policy{ _, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
ID: "policy",
Enabled: true, Enabled: true,
Rules: []*PolicyRule{ Rules: []*PolicyRule{
{ {
@ -1250,8 +1249,7 @@ func TestAccountManager_NetworkUpdates_SaveGroup(t *testing.T) {
Action: PolicyTrafficActionAccept, Action: PolicyTrafficActionAccept,
}, },
}, },
} })
err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false)
require.NoError(t, err) require.NoError(t, err)
updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID) updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
@ -1320,19 +1318,6 @@ func TestAccountManager_NetworkUpdates_SavePolicy(t *testing.T) {
updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID) updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
defer manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID) defer manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
policy := Policy{
Enabled: true,
Rules: []*PolicyRule{
{
Enabled: true,
Sources: []string{"groupA"},
Destinations: []string{"groupA"},
Bidirectional: true,
Action: PolicyTrafficActionAccept,
},
},
}
wg := sync.WaitGroup{} wg := sync.WaitGroup{}
wg.Add(1) wg.Add(1)
go func() { go func() {
@ -1345,7 +1330,19 @@ func TestAccountManager_NetworkUpdates_SavePolicy(t *testing.T) {
} }
}() }()
if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil { _, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
Enabled: true,
Rules: []*PolicyRule{
{
Enabled: true,
Sources: []string{"groupA"},
Destinations: []string{"groupA"},
Bidirectional: true,
Action: PolicyTrafficActionAccept,
},
},
})
if err != nil {
t.Errorf("delete default rule: %v", err) t.Errorf("delete default rule: %v", err)
return return
} }
@ -1366,7 +1363,7 @@ func TestAccountManager_NetworkUpdates_DeletePeer(t *testing.T) {
return return
} }
policy := Policy{ _, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
Enabled: true, Enabled: true,
Rules: []*PolicyRule{ Rules: []*PolicyRule{
{ {
@ -1377,9 +1374,8 @@ func TestAccountManager_NetworkUpdates_DeletePeer(t *testing.T) {
Action: PolicyTrafficActionAccept, Action: PolicyTrafficActionAccept,
}, },
}, },
} })
if err != nil {
if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil {
t.Errorf("save policy: %v", err) t.Errorf("save policy: %v", err)
return return
} }
@ -1421,7 +1417,12 @@ func TestAccountManager_NetworkUpdates_DeleteGroup(t *testing.T) {
require.NoError(t, err, "failed to save group") require.NoError(t, err, "failed to save group")
policy := Policy{ if err := manager.DeletePolicy(context.Background(), account.Id, account.Policies[0].ID, userID); err != nil {
t.Errorf("delete default rule: %v", err)
return
}
policy, err := manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
Enabled: true, Enabled: true,
Rules: []*PolicyRule{ Rules: []*PolicyRule{
{ {
@ -1432,14 +1433,8 @@ func TestAccountManager_NetworkUpdates_DeleteGroup(t *testing.T) {
Action: PolicyTrafficActionAccept, Action: PolicyTrafficActionAccept,
}, },
}, },
} })
if err != nil {
if err := manager.DeletePolicy(context.Background(), account.Id, account.Policies[0].ID, userID); err != nil {
t.Errorf("delete default rule: %v", err)
return
}
if err := manager.SavePolicy(context.Background(), account.Id, userID, &policy, false); err != nil {
t.Errorf("save policy: %v", err) t.Errorf("save policy: %v", err)
return return
} }

View File

@ -6,10 +6,8 @@ import (
"strconv" "strconv"
"github.com/gorilla/mux" "github.com/gorilla/mux"
nbgroup "github.com/netbirdio/netbird/management/server/group"
"github.com/rs/xid"
"github.com/netbirdio/netbird/management/server" "github.com/netbirdio/netbird/management/server"
nbgroup "github.com/netbirdio/netbird/management/server/group"
"github.com/netbirdio/netbird/management/server/http/api" "github.com/netbirdio/netbird/management/server/http/api"
"github.com/netbirdio/netbird/management/server/http/util" "github.com/netbirdio/netbird/management/server/http/util"
"github.com/netbirdio/netbird/management/server/jwtclaims" "github.com/netbirdio/netbird/management/server/jwtclaims"
@ -122,14 +120,9 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
return return
} }
isUpdate := policyID != "" policy := &server.Policy{
if policyID == "" {
policyID = xid.New().String()
}
policy := server.Policy{
ID: policyID, ID: policyID,
AccountID: accountID,
Name: req.Name, Name: req.Name,
Enabled: req.Enabled, Enabled: req.Enabled,
Description: req.Description, Description: req.Description,
@ -137,6 +130,7 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
for _, rule := range req.Rules { for _, rule := range req.Rules {
pr := server.PolicyRule{ pr := server.PolicyRule{
ID: policyID, // TODO: when policy can contain multiple rules, need refactor ID: policyID, // TODO: when policy can contain multiple rules, need refactor
PolicyID: policyID,
Name: rule.Name, Name: rule.Name,
Destinations: rule.Destinations, Destinations: rule.Destinations,
Sources: rule.Sources, Sources: rule.Sources,
@ -225,7 +219,8 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
policy.SourcePostureChecks = *req.SourcePostureChecks policy.SourcePostureChecks = *req.SourcePostureChecks
} }
if err := h.accountManager.SavePolicy(r.Context(), accountID, userID, &policy, isUpdate); err != nil { policy, err := h.accountManager.SavePolicy(r.Context(), accountID, userID, policy)
if err != nil {
util.WriteError(r.Context(), err, w) util.WriteError(r.Context(), err, w)
return return
} }
@ -236,7 +231,7 @@ func (h *Policies) savePolicy(w http.ResponseWriter, r *http.Request, accountID
return return
} }
resp := toPolicyResponse(allGroups, &policy) resp := toPolicyResponse(allGroups, policy)
if len(resp.Rules) == 0 { if len(resp.Rules) == 0 {
util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w) util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w)
return return

View File

@ -49,7 +49,7 @@ type MockAccountManager struct {
GroupDeletePeerFunc func(ctx context.Context, accountID, groupID, peerID string) error GroupDeletePeerFunc func(ctx context.Context, accountID, groupID, peerID string) error
DeleteRuleFunc func(ctx context.Context, accountID, ruleID, userID string) error DeleteRuleFunc func(ctx context.Context, accountID, ruleID, userID string) error
GetPolicyFunc func(ctx context.Context, accountID, policyID, userID string) (*server.Policy, error) GetPolicyFunc func(ctx context.Context, accountID, policyID, userID string) (*server.Policy, error)
SavePolicyFunc func(ctx context.Context, accountID, userID string, policy *server.Policy, isUpdate bool) error SavePolicyFunc func(ctx context.Context, accountID, userID string, policy *server.Policy) (*server.Policy, error)
DeletePolicyFunc func(ctx context.Context, accountID, policyID, userID string) error DeletePolicyFunc func(ctx context.Context, accountID, policyID, userID string) error
ListPoliciesFunc func(ctx context.Context, accountID, userID string) ([]*server.Policy, error) ListPoliciesFunc func(ctx context.Context, accountID, userID string) ([]*server.Policy, error)
GetUsersFromAccountFunc func(ctx context.Context, accountID, userID string) ([]*server.UserInfo, error) GetUsersFromAccountFunc func(ctx context.Context, accountID, userID string) ([]*server.UserInfo, error)
@ -386,11 +386,11 @@ func (am *MockAccountManager) GetPolicy(ctx context.Context, accountID, policyID
} }
// SavePolicy mock implementation of SavePolicy from server.AccountManager interface // SavePolicy mock implementation of SavePolicy from server.AccountManager interface
func (am *MockAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *server.Policy, isUpdate bool) error { func (am *MockAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *server.Policy) (*server.Policy, error) {
if am.SavePolicyFunc != nil { if am.SavePolicyFunc != nil {
return am.SavePolicyFunc(ctx, accountID, userID, policy, isUpdate) return am.SavePolicyFunc(ctx, accountID, userID, policy)
} }
return status.Errorf(codes.Unimplemented, "method SavePolicy is not implemented") return nil, status.Errorf(codes.Unimplemented, "method SavePolicy is not implemented")
} }
// DeletePolicy mock implementation of DeletePolicy from server.AccountManager interface // DeletePolicy mock implementation of DeletePolicy from server.AccountManager interface

View File

@ -3,13 +3,13 @@ package server
import ( import (
"context" "context"
_ "embed" _ "embed"
"slices"
"strconv" "strconv"
"strings" "strings"
"github.com/netbirdio/netbird/management/proto"
"github.com/rs/xid"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/proto"
"github.com/netbirdio/netbird/management/server/activity" "github.com/netbirdio/netbird/management/server/activity"
nbgroup "github.com/netbirdio/netbird/management/server/group" nbgroup "github.com/netbirdio/netbird/management/server/group"
nbpeer "github.com/netbirdio/netbird/management/server/peer" nbpeer "github.com/netbirdio/netbird/management/server/peer"
@ -125,6 +125,7 @@ type PolicyRule struct {
func (pm *PolicyRule) Copy() *PolicyRule { func (pm *PolicyRule) Copy() *PolicyRule {
rule := &PolicyRule{ rule := &PolicyRule{
ID: pm.ID, ID: pm.ID,
PolicyID: pm.PolicyID,
Name: pm.Name, Name: pm.Name,
Description: pm.Description, Description: pm.Description,
Enabled: pm.Enabled, Enabled: pm.Enabled,
@ -171,6 +172,7 @@ type Policy struct {
func (p *Policy) Copy() *Policy { func (p *Policy) Copy() *Policy {
c := &Policy{ c := &Policy{
ID: p.ID, ID: p.ID,
AccountID: p.AccountID,
Name: p.Name, Name: p.Name,
Description: p.Description, Description: p.Description,
Enabled: p.Enabled, Enabled: p.Enabled,
@ -343,157 +345,209 @@ func (am *DefaultAccountManager) GetPolicy(ctx context.Context, accountID, polic
return nil, err return nil, err
} }
if !user.IsAdminOrServiceUser() || user.AccountID != accountID { if user.AccountID != accountID {
return nil, status.Errorf(status.PermissionDenied, "only users with admin power are allowed to view policies") return nil, status.NewUserNotPartOfAccountError()
} }
return am.Store.GetPolicyByID(ctx, LockingStrengthShare, policyID, accountID) if user.IsRegularUser() {
return nil, status.NewAdminPermissionError()
}
return am.Store.GetPolicyByID(ctx, LockingStrengthShare, accountID, policyID)
} }
// SavePolicy in the store // SavePolicy in the store
func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *Policy, isUpdate bool) error { func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *Policy) (*Policy, error) {
unlock := am.Store.AcquireWriteLockByUID(ctx, accountID) unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
defer unlock() defer unlock()
account, err := am.Store.GetAccount(ctx, accountID) user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
if err != nil {
return nil, err
}
if user.AccountID != accountID {
return nil, status.NewUserNotPartOfAccountError()
}
if user.IsRegularUser() {
return nil, status.NewAdminPermissionError()
}
var isUpdate = policy.ID != ""
var updateAccountPeers bool
var action = activity.PolicyAdded
err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
if err = validatePolicy(ctx, transaction, accountID, policy); err != nil {
return err
}
updateAccountPeers, err = arePolicyChangesAffectPeers(ctx, transaction, accountID, policy, isUpdate)
if err != nil { if err != nil {
return err return err
} }
updateAccountPeers, err := am.savePolicy(account, policy, isUpdate) if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
if err != nil {
return err return err
} }
account.Network.IncSerial() saveFunc := transaction.CreatePolicy
if err = am.Store.SaveAccount(ctx, account); err != nil {
return err
}
action := activity.PolicyAdded
if isUpdate { if isUpdate {
action = activity.PolicyUpdated action = activity.PolicyUpdated
saveFunc = transaction.SavePolicy
} }
return saveFunc(ctx, LockingStrengthUpdate, policy)
})
if err != nil {
return nil, err
}
am.StoreEvent(ctx, userID, policy.ID, accountID, action, policy.EventMeta()) am.StoreEvent(ctx, userID, policy.ID, accountID, action, policy.EventMeta())
if updateAccountPeers { if updateAccountPeers {
am.updateAccountPeers(ctx, accountID) am.updateAccountPeers(ctx, accountID)
} }
return policy, nil
}
// DeletePolicy from the store
func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, policyID, userID string) error {
unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
defer unlock()
user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
if err != nil {
return err
}
if user.AccountID != accountID {
return status.NewUserNotPartOfAccountError()
}
if user.IsRegularUser() {
return status.NewAdminPermissionError()
}
var policy *Policy
var updateAccountPeers bool
err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
policy, err = transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policyID)
if err != nil {
return err
}
updateAccountPeers, err = arePolicyChangesAffectPeers(ctx, transaction, accountID, policy, false)
if err != nil {
return err
}
if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
return err
}
return transaction.DeletePolicy(ctx, LockingStrengthUpdate, accountID, policyID)
})
if err != nil {
return err
}
am.StoreEvent(ctx, userID, policyID, accountID, activity.PolicyRemoved, policy.EventMeta())
if updateAccountPeers {
am.updateAccountPeers(ctx, accountID)
}
return nil return nil
} }
// DeletePolicy from the store // ListPolicies from the store.
func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, policyID, userID string) error {
unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
defer unlock()
account, err := am.Store.GetAccount(ctx, accountID)
if err != nil {
return err
}
policy, err := am.deletePolicy(account, policyID)
if err != nil {
return err
}
account.Network.IncSerial()
if err = am.Store.SaveAccount(ctx, account); err != nil {
return err
}
am.StoreEvent(ctx, userID, policy.ID, accountID, activity.PolicyRemoved, policy.EventMeta())
if am.anyGroupHasPeers(account, policy.ruleGroups()) {
am.updateAccountPeers(ctx, accountID)
}
return nil
}
// ListPolicies from the store
func (am *DefaultAccountManager) ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error) { func (am *DefaultAccountManager) ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error) {
user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID) user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
if !user.IsAdminOrServiceUser() || user.AccountID != accountID { if user.AccountID != accountID {
return nil, status.Errorf(status.PermissionDenied, "only users with admin power are allowed to view policies") return nil, status.NewUserNotPartOfAccountError()
}
if user.IsRegularUser() {
return nil, status.NewAdminPermissionError()
} }
return am.Store.GetAccountPolicies(ctx, LockingStrengthShare, accountID) return am.Store.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
} }
func (am *DefaultAccountManager) deletePolicy(account *Account, policyID string) (*Policy, error) { // arePolicyChangesAffectPeers checks if changes to a policy will affect any associated peers.
policyIdx := -1 func arePolicyChangesAffectPeers(ctx context.Context, transaction Store, accountID string, policy *Policy, isUpdate bool) (bool, error) {
for i, policy := range account.Policies {
if policy.ID == policyID {
policyIdx = i
break
}
}
if policyIdx < 0 {
return nil, status.Errorf(status.NotFound, "rule with ID %s doesn't exist", policyID)
}
policy := account.Policies[policyIdx]
account.Policies = append(account.Policies[:policyIdx], account.Policies[policyIdx+1:]...)
return policy, nil
}
// savePolicy saves or updates a policy in the given account.
// If isUpdate is true, the function updates the existing policy; otherwise, it adds a new policy.
func (am *DefaultAccountManager) savePolicy(account *Account, policyToSave *Policy, isUpdate bool) (bool, error) {
for index, rule := range policyToSave.Rules {
rule.Sources = filterValidGroupIDs(account, rule.Sources)
rule.Destinations = filterValidGroupIDs(account, rule.Destinations)
policyToSave.Rules[index] = rule
}
if policyToSave.SourcePostureChecks != nil {
policyToSave.SourcePostureChecks = filterValidPostureChecks(account, policyToSave.SourcePostureChecks)
}
if isUpdate { if isUpdate {
policyIdx := slices.IndexFunc(account.Policies, func(policy *Policy) bool { return policy.ID == policyToSave.ID }) existingPolicy, err := transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policy.ID)
if policyIdx < 0 { if err != nil {
return false, status.Errorf(status.NotFound, "couldn't find policy id %s", policyToSave.ID) return false, err
} }
oldPolicy := account.Policies[policyIdx] if !policy.Enabled && !existingPolicy.Enabled {
// Update the existing policy
account.Policies[policyIdx] = policyToSave
if !policyToSave.Enabled && !oldPolicy.Enabled {
return false, nil return false, nil
} }
updateAccountPeers := am.anyGroupHasPeers(account, oldPolicy.ruleGroups()) || am.anyGroupHasPeers(account, policyToSave.ruleGroups())
return updateAccountPeers, nil hasPeers, err := anyGroupHasPeers(ctx, transaction, policy.AccountID, existingPolicy.ruleGroups())
if err != nil {
return false, err
} }
// Add the new policy to the account if hasPeers {
account.Policies = append(account.Policies, policyToSave) return true, nil
}
return am.anyGroupHasPeers(account, policyToSave.ruleGroups()), nil return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
}
return anyGroupHasPeers(ctx, transaction, policy.AccountID, policy.ruleGroups())
} }
func toProtocolFirewallRules(rules []*FirewallRule) []*proto.FirewallRule { // validatePolicy validates the policy and its rules.
result := make([]*proto.FirewallRule, len(rules)) func validatePolicy(ctx context.Context, transaction Store, accountID string, policy *Policy) error {
for i := range rules { if policy.ID != "" {
rule := rules[i] _, err := transaction.GetPolicyByID(ctx, LockingStrengthShare, accountID, policy.ID)
if err != nil {
return err
}
} else {
policy.ID = xid.New().String()
policy.AccountID = accountID
}
result[i] = &proto.FirewallRule{ groups, err := transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
PeerIP: rule.PeerIP, if err != nil {
Direction: getProtoDirection(rule.Direction), return err
Action: getProtoAction(rule.Action),
Protocol: getProtoProtocol(rule.Protocol),
Port: rule.Port,
} }
postureChecks, err := transaction.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
if err != nil {
return err
} }
return result
for i, rule := range policy.Rules {
ruleCopy := rule.Copy()
if ruleCopy.ID == "" {
ruleCopy.ID = xid.New().String()
ruleCopy.PolicyID = policy.ID
}
ruleCopy.Sources = getValidGroupIDs(groups, ruleCopy.Sources)
ruleCopy.Destinations = getValidGroupIDs(groups, ruleCopy.Destinations)
policy.Rules[i] = ruleCopy
}
if policy.SourcePostureChecks != nil {
policy.SourcePostureChecks = getValidPostureCheckIDs(postureChecks, policy.SourcePostureChecks)
}
return nil
} }
// getAllPeersFromGroups for given peer ID and list of groups // getAllPeersFromGroups for given peer ID and list of groups
@ -574,27 +628,52 @@ func (a *Account) getPostureChecks(postureChecksID string) *posture.Checks {
return nil return nil
} }
// filterValidPostureChecks filters and returns the posture check IDs from the given list // getValidPostureCheckIDs filters and returns only the valid posture check IDs from the provided list.
// that are valid within the provided account. func getValidPostureCheckIDs(postureChecks []*posture.Checks, postureChecksIds []string) []string {
func filterValidPostureChecks(account *Account, postureChecksIds []string) []string { validPostureCheckIDs := make(map[string]struct{})
result := make([]string, 0, len(postureChecksIds)) for _, check := range postureChecks {
validPostureCheckIDs[check.ID] = struct{}{}
}
validIDs := make([]string, 0, len(postureChecksIds))
for _, id := range postureChecksIds { for _, id := range postureChecksIds {
for _, postureCheck := range account.PostureChecks { if _, exists := validPostureCheckIDs[id]; exists {
if id == postureCheck.ID { validIDs = append(validIDs, id)
result = append(result, id)
continue
} }
} }
}
return result return validIDs
} }
// filterValidGroupIDs filters a list of group IDs and returns only the ones present in the account's group map. // getValidGroupIDs filters and returns only the valid group IDs from the provided list.
func filterValidGroupIDs(account *Account, groupIDs []string) []string { func getValidGroupIDs(groups []*nbgroup.Group, groupIDs []string) []string {
result := make([]string, 0, len(groupIDs)) validGroupIDs := make(map[string]struct{})
for _, groupID := range groupIDs { for _, group := range groups {
if _, exists := account.Groups[groupID]; exists { validGroupIDs[group.ID] = struct{}{}
result = append(result, groupID) }
validIDs := make([]string, 0, len(groupIDs))
for _, id := range groupIDs {
if _, exists := validGroupIDs[id]; exists {
validIDs = append(validIDs, id)
}
}
return validIDs
}
// toProtocolFirewallRules converts the firewall rules to the protocol firewall rules.
func toProtocolFirewallRules(rules []*FirewallRule) []*proto.FirewallRule {
result := make([]*proto.FirewallRule, len(rules))
for i := range rules {
rule := rules[i]
result[i] = &proto.FirewallRule{
PeerIP: rule.PeerIP,
Direction: getProtoDirection(rule.Direction),
Action: getProtoAction(rule.Action),
Protocol: getProtoProtocol(rule.Protocol),
Port: rule.Port,
} }
} }
return result return result

View File

@ -1286,12 +1286,67 @@ func (s *SqlStore) DeleteGroups(ctx context.Context, strength LockingStrength, a
// GetAccountPolicies retrieves policies for an account. // GetAccountPolicies retrieves policies for an account.
func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) { func (s *SqlStore) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) {
return getRecords[*Policy](s.db.Preload(clause.Associations), lockStrength, accountID) var policies []*Policy
result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
Preload(clause.Associations).Find(&policies, accountIDCondition, accountID)
if err := result.Error; err != nil {
log.WithContext(ctx).Errorf("failed to get policies from the store: %s", result.Error)
return nil, status.Errorf(status.Internal, "failed to get policies from store")
}
return policies, nil
} }
// GetPolicyByID retrieves a policy by its ID and account ID. // GetPolicyByID retrieves a policy by its ID and account ID.
func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, policyID string, accountID string) (*Policy, error) { func (s *SqlStore) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*Policy, error) {
return getRecordByID[Policy](s.db.Preload(clause.Associations), lockStrength, policyID, accountID) var policy *Policy
result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Preload(clause.Associations).
First(&policy, accountAndIDQueryCondition, accountID, policyID)
if err := result.Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, status.NewPolicyNotFoundError(policyID)
}
log.WithContext(ctx).Errorf("failed to get policy from store: %s", err)
return nil, status.Errorf(status.Internal, "failed to get policy from store")
}
return policy, nil
}
func (s *SqlStore) CreatePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error {
result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Create(policy)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to create policy in store: %s", result.Error)
return status.Errorf(status.Internal, "failed to create policy in store")
}
return nil
}
// SavePolicy saves a policy to the database.
func (s *SqlStore) SavePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error {
result := s.db.Session(&gorm.Session{FullSaveAssociations: true}).
Clauses(clause.Locking{Strength: string(lockStrength)}).Save(policy)
if err := result.Error; err != nil {
log.WithContext(ctx).Errorf("failed to save policy to the store: %s", err)
return status.Errorf(status.Internal, "failed to save policy to store")
}
return nil
}
func (s *SqlStore) DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error {
result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
Delete(&Policy{}, accountAndIDQueryCondition, accountID, policyID)
if err := result.Error; err != nil {
log.WithContext(ctx).Errorf("failed to delete policy from store: %s", err)
return status.Errorf(status.Internal, "failed to delete policy from store")
}
if result.RowsAffected == 0 {
return status.NewPolicyNotFoundError(policyID)
}
return nil
} }
// GetAccountPostureChecks retrieves posture checks for an account. // GetAccountPostureChecks retrieves posture checks for an account.
@ -1324,7 +1379,7 @@ func (s *SqlStore) GetPostureChecksByID(ctx context.Context, lockStrength Lockin
// SavePostureChecks saves a posture checks to the database. // SavePostureChecks saves a posture checks to the database.
func (s *SqlStore) SavePostureChecks(ctx context.Context, lockStrength LockingStrength, postureCheck *posture.Checks) error { func (s *SqlStore) SavePostureChecks(ctx context.Context, lockStrength LockingStrength, postureCheck *posture.Checks) error {
result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Save(postureCheck) result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(postureCheck)
if result.Error != nil { if result.Error != nil {
if errors.Is(result.Error, gorm.ErrDuplicatedKey) { if errors.Is(result.Error, gorm.ErrDuplicatedKey) {
return status.Errorf(status.InvalidArgument, "name should be unique") return status.Errorf(status.InvalidArgument, "name should be unique")

View File

@ -139,3 +139,8 @@ func NewGroupNotFoundError(groupID string) error {
func NewPostureChecksNotFoundError(postureChecksID string) error { func NewPostureChecksNotFoundError(postureChecksID string) error {
return Errorf(NotFound, "posture checks: %s not found", postureChecksID) return Errorf(NotFound, "posture checks: %s not found", postureChecksID)
} }
// NewPolicyNotFoundError creates a new Error with NotFound type for a missing policy
func NewPolicyNotFoundError(policyID string) error {
return Errorf(NotFound, "policy: %s not found", policyID)
}

View File

@ -80,7 +80,10 @@ type Store interface {
DeleteGroups(ctx context.Context, strength LockingStrength, accountID string, groupIDs []string) error DeleteGroups(ctx context.Context, strength LockingStrength, accountID string, groupIDs []string) error
GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error) GetAccountPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*Policy, error)
GetPolicyByID(ctx context.Context, lockStrength LockingStrength, policyID string, accountID string) (*Policy, error) GetPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*Policy, error)
CreatePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error
SavePolicy(ctx context.Context, lockStrength LockingStrength, policy *Policy) error
DeletePolicy(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) error
GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error) GetPostureCheckByChecksDefinition(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error) GetAccountPostureChecks(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*posture.Checks, error)