Merge branch 'feat/admin-cli' of github.com:netbirdio/netbird into feat/admin-cli

This commit is contained in:
jnfrati
2026-07-03 15:26:41 +02:00
378 changed files with 48704 additions and 4515 deletions
+21 -8
View File
@@ -358,7 +358,8 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
oldSettings.AutoUpdateVersion != newSettings.AutoUpdateVersion ||
oldSettings.AutoUpdateAlways != newSettings.AutoUpdateAlways ||
oldSettings.PeerLoginExpirationEnabled != newSettings.PeerLoginExpirationEnabled ||
oldSettings.PeerLoginExpiration != newSettings.PeerLoginExpiration {
oldSettings.PeerLoginExpiration != newSettings.PeerLoginExpiration ||
oldSettings.MetricsPushEnabled != newSettings.MetricsPushEnabled {
// Session deadline is derived from LastLogin + PeerLoginExpiration
// on every Login/Sync response. Without a fan-out push, connected
// peers keep the deadline they received at login time and only see
@@ -409,6 +410,7 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
am.handleAutoUpdateVersionSettings(ctx, oldSettings, newSettings, userID, accountID)
am.handleAutoUpdateAlwaysSettings(ctx, oldSettings, newSettings, userID, accountID)
am.handlePeerExposeSettings(ctx, oldSettings, newSettings, userID, accountID)
am.handleMetricsPushSettings(ctx, oldSettings, newSettings, userID, accountID)
if err = am.handleInactivityExpirationSettings(ctx, oldSettings, newSettings, userID, accountID); err != nil {
return nil, err
}
@@ -563,6 +565,16 @@ func (am *DefaultAccountManager) handleLazyConnectionSettings(ctx context.Contex
}
}
func (am *DefaultAccountManager) handleMetricsPushSettings(ctx context.Context, oldSettings, newSettings *types.Settings, userID, accountID string) {
if oldSettings.MetricsPushEnabled != newSettings.MetricsPushEnabled {
if newSettings.MetricsPushEnabled {
am.StoreEvent(ctx, userID, accountID, accountID, activity.AccountMetricsPushEnabled, nil)
} else {
am.StoreEvent(ctx, userID, accountID, accountID, activity.AccountMetricsPushDisabled, nil)
}
}
}
func (am *DefaultAccountManager) handlePeerLoginExpirationSettings(ctx context.Context, oldSettings, newSettings *types.Settings, userID, accountID string) {
if oldSettings.PeerLoginExpirationEnabled != newSettings.PeerLoginExpirationEnabled {
event := activity.AccountPeerLoginExpirationEnabled
@@ -689,7 +701,7 @@ func (am *DefaultAccountManager) peerLoginExpirationJob(ctx context.Context, acc
log.WithContext(ctx).Debugf("discovered %d peers to expire for account %s", len(peerIDs), accountID)
if err := am.expireAndUpdatePeers(ctx, accountID, expiredPeers); err != nil {
if err := am.expireAndUpdatePeers(ctx, accountID, expiredPeers, peerExpirationSessionExpired); err != nil {
log.WithContext(ctx).Errorf("failed updating account peers while expiring peers for account %s", accountID)
return peerSchedulerRetryInterval, true
}
@@ -724,7 +736,7 @@ func (am *DefaultAccountManager) peerInactivityExpirationJob(ctx context.Context
log.Debugf("discovered %d peers to expire for account %s", len(peerIDs), accountID)
if err := am.expireAndUpdatePeers(ctx, accountID, inactivePeers); err != nil {
if err := am.expireAndUpdatePeers(ctx, accountID, inactivePeers, peerExpirationInactivity); err != nil {
log.Errorf("failed updating account peers while expiring peers for account %s", accountID)
return peerSchedulerRetryInterval, true
}
@@ -1889,12 +1901,12 @@ func domainIsUpToDate(domain string, domainCategory string, userAuth auth.UserAu
// concurrent stream that started earlier loses the optimistic-lock race
// in MarkPeerConnected and bails without writing.
func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, int64, error) {
peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta}, accountID)
peer, netMap, postureChecks, dnsfwdPort, err := am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP}, accountID)
if err != nil {
return nil, nil, nil, 0, fmt.Errorf("error syncing peer: %w", err)
}
if err := am.MarkPeerConnected(ctx, peerPubKey, realIP, accountID, syncTime.UnixNano(), netMap); err != nil {
if err := am.MarkPeerConnected(ctx, peerPubKey, accountID, syncTime.UnixNano(), netMap); err != nil {
log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peerPubKey, err)
}
@@ -1914,13 +1926,13 @@ func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, account
return nil
}
func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error {
func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error {
accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peerPubKey)
if err != nil {
return err
}
_, _, _, _, err = am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, UpdateAccountPeers: true}, accountID)
_, _, _, _, err = am.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peerPubKey, Meta: meta, RealIP: realIP, UpdateAccountPeers: true}, accountID)
if err != nil {
return err
}
@@ -1949,7 +1961,7 @@ func (am *DefaultAccountManager) onPeersInvalidated(ctx context.Context, account
}
}
if len(peers) > 0 {
err := am.expireAndUpdatePeers(ctx, accountID, peers)
err := am.expireAndUpdatePeers(ctx, accountID, peers, peerExpirationValidationFailed)
if err != nil {
log.WithContext(ctx).Errorf("failed to expire and update invalidated peers for account %s: %v", accountID, err)
return
@@ -2045,6 +2057,7 @@ func newAccountWithId(ctx context.Context, accountID, userID, domain, email, nam
Extra: &types.ExtraSettings{
UserApprovalRequired: true,
},
LazyConnectionEnabled: true,
},
Onboarding: types.AccountOnboarding{
OnboardingFlowPending: true,
+2 -2
View File
@@ -62,7 +62,7 @@ type Manager interface {
GetUserFromUserAuth(ctx context.Context, userAuth auth.UserAuth) (*types.User, error)
ListUsers(ctx context.Context, accountID string) ([]*types.User, error)
GetPeers(ctx context.Context, accountID, userID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error)
MarkPeerConnected(ctx context.Context, peerKey string, realIP net.IP, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerConnected(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerDisconnected(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64) error
DeletePeer(ctx context.Context, accountID, peerID, userID string) error
UpdatePeer(ctx context.Context, accountID, userID string, p *nbpeer.Peer) (*nbpeer.Peer, error)
@@ -123,7 +123,7 @@ type Manager interface {
GetValidatedPeers(ctx context.Context, accountID string) (map[string]struct{}, map[string]string, error)
SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, int64, error)
OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string, streamStartTime time.Time) error
SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error
SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error
FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
GetAccountIDForPeerKey(ctx context.Context, peerKey string) (string, error)
GetAccountSettings(ctx context.Context, accountID string, userID string) (*types.Settings, error)
+8 -8
View File
@@ -1323,17 +1323,17 @@ func (mr *MockManagerMockRecorder) ExtendPeerSession(ctx, peerPubKey, userID int
}
// MarkPeerConnected mocks base method.
func (m *MockManager) MarkPeerConnected(ctx context.Context, peerKey string, realIP net.IP, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
func (m *MockManager) MarkPeerConnected(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "MarkPeerConnected", ctx, peerKey, realIP, accountID, sessionStartedAt, nmap)
ret := m.ctrl.Call(m, "MarkPeerConnected", ctx, peerKey, accountID, sessionStartedAt, nmap)
ret0, _ := ret[0].(error)
return ret0
}
// MarkPeerConnected indicates an expected call of MarkPeerConnected.
func (mr *MockManagerMockRecorder) MarkPeerConnected(ctx, peerKey, realIP, accountID, sessionStartedAt, nmap interface{}) *gomock.Call {
func (mr *MockManagerMockRecorder) MarkPeerConnected(ctx, peerKey, accountID, sessionStartedAt, nmap interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerConnected", reflect.TypeOf((*MockManager)(nil).MarkPeerConnected), ctx, peerKey, realIP, accountID, sessionStartedAt, nmap)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPeerConnected", reflect.TypeOf((*MockManager)(nil).MarkPeerConnected), ctx, peerKey, accountID, sessionStartedAt, nmap)
}
// MarkPeerDisconnected mocks base method.
@@ -1586,17 +1586,17 @@ func (mr *MockManagerMockRecorder) SyncPeer(ctx, sync, accountID interface{}) *g
}
// SyncPeerMeta mocks base method.
func (m *MockManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta peer.PeerSystemMeta) error {
func (m *MockManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta peer.PeerSystemMeta, realIP net.IP) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SyncPeerMeta", ctx, peerPubKey, meta)
ret := m.ctrl.Call(m, "SyncPeerMeta", ctx, peerPubKey, meta, realIP)
ret0, _ := ret[0].(error)
return ret0
}
// SyncPeerMeta indicates an expected call of SyncPeerMeta.
func (mr *MockManagerMockRecorder) SyncPeerMeta(ctx, peerPubKey, meta interface{}) *gomock.Call {
func (mr *MockManagerMockRecorder) SyncPeerMeta(ctx, peerPubKey, meta, realIP interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncPeerMeta", reflect.TypeOf((*MockManager)(nil).SyncPeerMeta), ctx, peerPubKey, meta)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SyncPeerMeta", reflect.TypeOf((*MockManager)(nil).SyncPeerMeta), ctx, peerPubKey, meta, realIP)
}
// SyncUserJWTGroups mocks base method.
+119 -8
View File
@@ -1836,7 +1836,7 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to get the account")
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), nil, accountID, time.Now().UTC().UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), accountID, time.Now().UTC().UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &types.Settings{
@@ -1907,7 +1907,7 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
require.NoError(t, err, "unable to get the account")
// when we mark peer as connected, the peer login expiration routine should trigger
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), nil, accountID, time.Now().UTC().UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), accountID, time.Now().UTC().UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
failed := waitTimeout(wg, time.Second)
@@ -1916,6 +1916,117 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.
}
}
func TestDefaultAccountManager_MarkPeerDisconnected_SchedulesInactivityExpiration(t *testing.T) {
manager, _, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to create an account")
key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
peerPubKey := key.PublicKey().String()
_, _, _, _, err = manager.AddPeer(context.Background(), "", "", userID, &nbpeer.Peer{
Key: peerPubKey,
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer"},
InactivityExpirationEnabled: true,
}, false)
require.NoError(t, err, "unable to add peer")
_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &types.Settings{
PeerLoginExpiration: time.Hour,
PeerLoginExpirationEnabled: true,
PeerInactivityExpiration: time.Hour,
PeerInactivityExpirationEnabled: true,
Extra: &types.ExtraSettings{},
})
require.NoError(t, err, "expecting to update account settings successfully but got error")
// Establish a session so the matching-token disconnect is actually applied.
streamStartTime := time.Now().UTC()
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
// Install the mock only now, so the assertion observes the disconnect, not
// the earlier connect.
scheduled := make(chan struct{}, 1)
manager.peerInactivityExpiry = &MockScheduler{
CancelFunc: func(ctx context.Context, IDs []string) {},
ScheduleFunc: func(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) {
select {
case scheduled <- struct{}{}:
default:
}
},
}
err = manager.MarkPeerDisconnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano())
require.NoError(t, err, "unable to mark peer disconnected")
select {
case <-scheduled:
// expected: disconnect re-armed the inactivity expiry timer
case <-time.After(time.Second):
t.Fatal("expected inactivity expiration to be rescheduled when an eligible peer disconnects")
}
}
func TestDefaultAccountManager_MarkPeerDisconnected_SkipsInactivityExpirationWhenDisabled(t *testing.T) {
manager, _, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
accountID, err := manager.GetAccountIDByUserID(context.Background(), auth.UserAuth{UserId: userID})
require.NoError(t, err, "unable to create an account")
key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
peerPubKey := key.PublicKey().String()
_, _, _, _, err = manager.AddPeer(context.Background(), "", "", userID, &nbpeer.Peer{
Key: peerPubKey,
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer"},
InactivityExpirationEnabled: true,
}, false)
require.NoError(t, err, "unable to add peer")
// Peer is eligible (SSO + inactivity enabled) but the account-level setting
// stays disabled, so disconnect must not schedule anything.
_, err = manager.UpdateAccountSettings(context.Background(), accountID, userID, &types.Settings{
PeerLoginExpiration: time.Hour,
PeerLoginExpirationEnabled: true,
PeerInactivityExpiration: time.Hour,
PeerInactivityExpirationEnabled: false,
Extra: &types.ExtraSettings{},
})
require.NoError(t, err, "expecting to update account settings successfully but got error")
streamStartTime := time.Now().UTC()
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
scheduled := make(chan struct{}, 1)
manager.peerInactivityExpiry = &MockScheduler{
CancelFunc: func(ctx context.Context, IDs []string) {},
ScheduleFunc: func(ctx context.Context, in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) {
select {
case scheduled <- struct{}{}:
default:
}
},
}
err = manager.MarkPeerDisconnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano())
require.NoError(t, err, "unable to mark peer disconnected")
select {
case <-scheduled:
t.Fatal("inactivity expiration must not be scheduled while the account-level setting is disabled")
case <-time.After(200 * time.Millisecond):
// expected: nothing scheduled
}
}
func TestDefaultAccountManager_OnPeerDisconnected_LastSeenCheck(t *testing.T) {
manager, _, err := createManager(t)
require.NoError(t, err, "unable to create account manager")
@@ -1935,7 +2046,7 @@ func TestDefaultAccountManager_OnPeerDisconnected_LastSeenCheck(t *testing.T) {
t.Run("disconnect peer when session token matches", func(t *testing.T) {
streamStartTime := time.Now().UTC()
err = manager.MarkPeerConnected(context.Background(), peerPubKey, nil, accountID, streamStartTime.UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
peer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, peerPubKey)
@@ -1956,7 +2067,7 @@ func TestDefaultAccountManager_OnPeerDisconnected_LastSeenCheck(t *testing.T) {
t.Run("skip disconnect when stored session is newer (zombie stream protection)", func(t *testing.T) {
// Newer stream wins on connect (sets SessionStartedAt = now ns).
streamStartTime := time.Now().UTC()
err = manager.MarkPeerConnected(context.Background(), peerPubKey, nil, accountID, streamStartTime.UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, streamStartTime.UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
peer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, peerPubKey)
@@ -1980,7 +2091,7 @@ func TestDefaultAccountManager_OnPeerDisconnected_LastSeenCheck(t *testing.T) {
t.Run("skip stale connect when stored session is newer (blocked goroutine protection)", func(t *testing.T) {
node2SyncTime := time.Now().UTC()
err = manager.MarkPeerConnected(context.Background(), peerPubKey, nil, accountID, node2SyncTime.UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, node2SyncTime.UnixNano(), nil)
require.NoError(t, err, "node 2 should connect peer")
peer, err := manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, peerPubKey)
@@ -1990,7 +2101,7 @@ func TestDefaultAccountManager_OnPeerDisconnected_LastSeenCheck(t *testing.T) {
"SessionStartedAt should equal node2SyncTime token")
node1StaleSyncTime := node2SyncTime.Add(-1 * time.Minute)
err = manager.MarkPeerConnected(context.Background(), peerPubKey, nil, accountID, node1StaleSyncTime.UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, node1StaleSyncTime.UnixNano(), nil)
require.NoError(t, err, "stale connect should not return error")
peer, err = manager.Store.GetPeerByPeerPubKey(context.Background(), store.LockingStrengthNone, peerPubKey)
@@ -2052,7 +2163,7 @@ func TestDefaultAccountManager_MarkPeerConnected_ConcurrentRace(t *testing.T) {
defer done.Done()
ready.Done()
start.Wait()
errs <- manager.MarkPeerConnected(context.Background(), peerPubKey, nil, accountID, token, nil)
errs <- manager.MarkPeerConnected(context.Background(), peerPubKey, accountID, token, nil)
}()
}
@@ -2093,7 +2204,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test
account, err := manager.Store.GetAccount(context.Background(), accountID)
require.NoError(t, err, "unable to get the account")
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), nil, accountID, time.Now().UTC().UnixNano(), nil)
err = manager.MarkPeerConnected(context.Background(), key.PublicKey().String(), accountID, time.Now().UTC().UnixNano(), nil)
require.NoError(t, err, "unable to mark peer connected")
wg := &sync.WaitGroup{}
+57
View File
@@ -245,6 +245,42 @@ const (
// tunnel. Distinct from UserLoggedInPeer (full interactive login).
UserExtendedPeerSession Activity = 125
// AgentNetworkProviderCreated indicates that a user created an Agent Network provider
AgentNetworkProviderCreated Activity = 126
// AgentNetworkProviderUpdated indicates that a user updated an Agent Network provider
AgentNetworkProviderUpdated Activity = 127
// AgentNetworkProviderDeleted indicates that a user deleted an Agent Network provider
AgentNetworkProviderDeleted Activity = 128
// AgentNetworkPolicyCreated indicates that a user created an Agent Network policy
AgentNetworkPolicyCreated Activity = 129
// AgentNetworkPolicyUpdated indicates that a user updated an Agent Network policy
AgentNetworkPolicyUpdated Activity = 130
// AgentNetworkPolicyDeleted indicates that a user deleted an Agent Network policy
AgentNetworkPolicyDeleted Activity = 131
// AgentNetworkGuardrailCreated indicates that a user created an Agent Network guardrail
AgentNetworkGuardrailCreated Activity = 132
// AgentNetworkGuardrailUpdated indicates that a user updated an Agent Network guardrail
AgentNetworkGuardrailUpdated Activity = 133
// AgentNetworkGuardrailDeleted indicates that a user deleted an Agent Network guardrail
AgentNetworkGuardrailDeleted Activity = 134
// AgentNetworkBudgetRuleCreated indicates that a user created an Agent Network budget rule
AgentNetworkBudgetRuleCreated Activity = 135
// AgentNetworkBudgetRuleUpdated indicates that a user updated an Agent Network budget rule
AgentNetworkBudgetRuleUpdated Activity = 136
// AgentNetworkBudgetRuleDeleted indicates that a user deleted an Agent Network budget rule
AgentNetworkBudgetRuleDeleted Activity = 137
// AgentNetworkSettingsUpdated indicates that a user updated Agent Network account settings
AgentNetworkSettingsUpdated Activity = 139
// AccountMetricsPushEnabled indicates that a user enabled metrics push for the account
AccountMetricsPushEnabled Activity = 140
// AccountMetricsPushDisabled indicates that a user disabled metrics push for the account
AccountMetricsPushDisabled Activity = 141
AccountDeleted Activity = 99999
)
@@ -400,6 +436,27 @@ var activityMap = map[Activity]Code{
UserExtendedPeerSession: {"User extended peer session", "user.peer.session.extend"},
AgentNetworkProviderCreated: {"Agent Network provider created", "agent_network.provider.create"},
AgentNetworkProviderUpdated: {"Agent Network provider updated", "agent_network.provider.update"},
AgentNetworkProviderDeleted: {"Agent Network provider deleted", "agent_network.provider.delete"},
AgentNetworkPolicyCreated: {"Agent Network policy created", "agent_network.policy.create"},
AgentNetworkPolicyUpdated: {"Agent Network policy updated", "agent_network.policy.update"},
AgentNetworkPolicyDeleted: {"Agent Network policy deleted", "agent_network.policy.delete"},
AgentNetworkGuardrailCreated: {"Agent Network guardrail created", "agent_network.guardrail.create"},
AgentNetworkGuardrailUpdated: {"Agent Network guardrail updated", "agent_network.guardrail.update"},
AgentNetworkGuardrailDeleted: {"Agent Network guardrail deleted", "agent_network.guardrail.delete"},
AgentNetworkBudgetRuleCreated: {"Agent Network budget rule created", "agent_network.budget_rule.create"},
AgentNetworkBudgetRuleUpdated: {"Agent Network budget rule updated", "agent_network.budget_rule.update"},
AgentNetworkBudgetRuleDeleted: {"Agent Network budget rule deleted", "agent_network.budget_rule.delete"},
AgentNetworkSettingsUpdated: {"Agent Network settings updated", "agent_network.settings.update"},
AccountMetricsPushEnabled: {"Account metrics push enabled", "account.setting.metrics.push.enable"},
AccountMetricsPushDisabled: {"Account metrics push disabled", "account.setting.metrics.push.disable"},
DomainAdded: {"Domain added", "domain.add"},
DomainDeleted: {"Domain deleted", "domain.delete"},
DomainValidated: {"Domain validated", "domain.validate"},
@@ -41,7 +41,7 @@ func TestAffectedPeers_DependencyCoverageMatrix(t *testing.T) {
_, err := s.manager.SavePolicy(ctx, s.accountID, userID, peerToResourcePolicyByGroup(s.sourceGroupID, s.resourceGroupID), true)
require.NoError(t, err)
return affectedpeers.Change{ChangedPeerIDs: []string{s.routerPeerID}},
[]string{s.sourcePeerID}, []string{s.unrelatedPeerID}
[]string{s.sourcePeerID, s.routerPeerID}, []string{s.unrelatedPeerID}
},
},
{
@@ -106,11 +106,9 @@ func TestAffectedPeers_DependencyCoverageMatrix(t *testing.T) {
change, mustContain, mustExclude := r.build(t, s, ctx)
affected := resolveAffected(t, s.manager.Store, s.accountID, change)
for _, id := range mustContain {
assert.Contains(t, affected, id, "expected peer to be affected")
}
for _, id := range mustExclude {
assert.NotContains(t, affected, id, "peer must not be affected")
assert.ElementsMatch(t, affected, mustContain, "expected peer to be affected")
for _, peerID := range mustExclude {
assert.NotContains(t, affected, peerID, "peer must not be affected")
}
})
}
@@ -251,7 +251,9 @@ func TestAffectedPeers_E2E_UpdateResource_DestinationResourcePolicy_RefreshesSou
}
}
func TestAffectedPeers_E2E_UpdateResource_DisabledSiblingRouter_StillBridged(t *testing.T) {
// A disabled sibling router routes to nobody, so updating a resource on its network
// must NOT refresh its peer (the enabled router carries the bridge instead).
func TestAffectedPeers_E2E_UpdateResource_DisabledSiblingRouterNotBridged(t *testing.T) {
s := setupRouterScenario(t, true)
ctx := context.Background()
@@ -274,13 +276,18 @@ func TestAffectedPeers_E2E_UpdateResource_DisabledSiblingRouter_StillBridged(t *
require.NoError(t, err)
disabledCh := s.updateManager.CreateChannel(ctx, disabledRouterPeer.ID)
t.Cleanup(func() { s.updateManager.CloseChannel(ctx, disabledRouterPeer.ID) })
enabledCh := s.updateManager.CreateChannel(ctx, s.routerPeerID)
t.Cleanup(func() {
s.updateManager.CloseChannel(ctx, disabledRouterPeer.ID)
s.updateManager.CloseChannel(ctx, s.routerPeerID)
})
settleAffectedUpdates(disabledCh)
settleAffectedUpdates(disabledCh, enabledCh)
done := make(chan struct{})
go func() {
peerShouldReceiveUpdate(t, disabledCh)
peerShouldReceiveUpdate(t, enabledCh)
peerShouldNotReceiveUpdate(t, disabledCh)
close(done)
}()
@@ -298,7 +305,7 @@ func TestAffectedPeers_E2E_UpdateResource_DisabledSiblingRouter_StillBridged(t *
select {
case <-done:
case <-time.After(peerUpdateTimeout):
t.Error("timeout: resource update did not refresh the disabled sibling router's peer")
t.Error("timeout")
}
}
@@ -682,6 +682,9 @@ func TestAffectedPeers_AllRoutingPeers_Network(t *testing.T) {
assert.Contains(t, affected, secondRouterPeer.ID, "second routing peer on the same network must also be affected")
}
// A disabled router in the snapshot routes to nobody, so it is skipped when the
// walk scans existing account data: a policy edit still folds the literal source
// group, but not the disabled router's peer.
func TestAffectedPeers_DisabledRouter(t *testing.T) {
s := setupRouterScenario(t, true)
ctx := context.Background()
@@ -694,11 +697,13 @@ func TestAffectedPeers_DisabledRouter(t *testing.T) {
affected := s.resolvePolicyAffected(ctx, peerToResourcePolicyByGroup(s.sourceGroupID, s.resourceGroupID))
assert.Contains(t, affected, s.sourcePeerID, "source peer must be affected")
assert.Contains(t, affected, s.routerPeerID,
"disabled router's peer must still be affected: Enabled must not gate affected-peers")
assert.Contains(t, affected, s.sourcePeerID, "source peer (literal policy source group) must be affected")
assert.NotContains(t, affected, s.routerPeerID,
"a disabled router routes to nobody, so its peer must not be folded from snapshot data")
}
// A disabled resource in the snapshot is skipped: the policy edit still folds the
// literal source group, but the resource no longer bridges to its network's router.
func TestAffectedPeers_DisabledResource(t *testing.T) {
s := setupRouterScenario(t, true)
ctx := context.Background()
@@ -710,9 +715,9 @@ func TestAffectedPeers_DisabledResource(t *testing.T) {
affected := s.resolvePolicyAffected(ctx, peerToResourcePolicyByGroup(s.sourceGroupID, s.resourceGroupID))
assert.Contains(t, affected, s.sourcePeerID, "source peer must be affected")
assert.Contains(t, affected, s.routerPeerID,
"disabled resource must still resolve the routing peer: Enabled must not gate affected-peers")
assert.Contains(t, affected, s.sourcePeerID, "source peer (literal policy source group) must be affected")
assert.NotContains(t, affected, s.routerPeerID,
"a disabled resource routes to nobody, so its network's router must not be folded from snapshot data")
}
func TestAffectedPeers_DisabledRule(t *testing.T) {
+96 -45
View File
@@ -96,33 +96,54 @@ func affectedGroupID(i int) string { return fmt.Sprintf("affected-grp-%d", i)
func affectedGroupName(i int) string { return fmt.Sprintf("AffectedGroup%d", i) }
func TestCollectGroupChange_PolicyLinked(t *testing.T) {
manager, s, accountID, _, groupIDs := setupAffectedPeersTest(t)
manager, s, accountID, peerIDs, groupIDs := setupAffectedPeersTest(t)
ctx := context.Background()
_, err := manager.SavePolicy(ctx, accountID, userID, &types.Policy{
Enabled: true,
Rules: []*types.PolicyRule{
{
Enabled: true,
Sources: []string{groupIDs[0]},
Destinations: []string{groupIDs[1]},
Bidirectional: true,
Action: types.PolicyTrafficActionAccept,
Enabled: true,
Sources: []string{groupIDs[0]},
Destinations: []string{groupIDs[1]},
SourceResource: types.Resource{ID: peerIDs[0], Type: types.ResourceTypePeer},
DestinationResource: types.Resource{ID: peerIDs[1], Type: types.ResourceTypePeer},
Bidirectional: true,
Action: types.PolicyTrafficActionAccept,
},
{
Enabled: true,
Sources: []string{groupIDs[0]},
Destinations: []string{groupIDs[1]},
SourceResource: types.Resource{ID: peerIDs[2], Type: types.ResourceTypeHost},
DestinationResource: types.Resource{ID: peerIDs[3], Type: types.ResourceTypeHost},
Bidirectional: true,
Action: types.PolicyTrafficActionAccept,
},
{
Enabled: true,
Sources: []string{groupIDs[0]},
Destinations: []string{groupIDs[1]},
SourceResource: types.Resource{ID: "", Type: types.ResourceTypePeer},
DestinationResource: types.Resource{ID: "", Type: types.ResourceTypePeer},
Bidirectional: true,
Action: types.PolicyTrafficActionAccept,
},
},
}, true)
require.NoError(t, err)
groups, _ := collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[0]})
assert.Contains(t, groups, groupIDs[0])
assert.Contains(t, groups, groupIDs[1])
groups, directPeers := collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[0]})
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.ElementsMatch(t, directPeers, []string{peerIDs[1]})
groups, _ = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[1]})
assert.Contains(t, groups, groupIDs[0])
assert.Contains(t, groups, groupIDs[1])
groups, directPeers = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[1]})
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.ElementsMatch(t, directPeers, []string{peerIDs[0]})
groups, _ = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[2]})
groups, directPeers = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[2]})
assert.Empty(t, groups)
assert.Empty(t, directPeers)
}
func TestCollectGroupChange_PolicyWithDirectPeerResource(t *testing.T) {
@@ -133,20 +154,44 @@ func TestCollectGroupChange_PolicyWithDirectPeerResource(t *testing.T) {
Enabled: true,
Rules: []*types.PolicyRule{
{
Enabled: true,
Sources: []string{groupIDs[0]},
SourceResource: types.Resource{ID: peerIDs[3], Type: types.ResourceTypePeer},
Destinations: []string{groupIDs[1]},
Action: types.PolicyTrafficActionAccept,
Enabled: true,
Sources: []string{groupIDs[0]},
SourceResource: types.Resource{ID: peerIDs[3], Type: types.ResourceTypePeer},
DestinationResource: types.Resource{ID: peerIDs[4], Type: types.ResourceTypePeer},
Destinations: []string{groupIDs[1]},
Action: types.PolicyTrafficActionAccept,
},
{
Enabled: true,
Sources: []string{groupIDs[0]},
SourceResource: types.Resource{ID: peerIDs[1], Type: types.ResourceTypeHost},
DestinationResource: types.Resource{ID: peerIDs[2], Type: types.ResourceTypeHost},
Destinations: []string{groupIDs[1]},
Action: types.PolicyTrafficActionAccept,
},
{
Enabled: true,
Sources: []string{groupIDs[0]},
SourceResource: types.Resource{ID: "", Type: types.ResourceTypePeer},
DestinationResource: types.Resource{ID: "", Type: types.ResourceTypePeer},
Destinations: []string{groupIDs[1]},
Action: types.PolicyTrafficActionAccept,
},
},
}, true)
require.NoError(t, err)
groups, directPeers := collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[0]})
assert.Contains(t, groups, groupIDs[0])
assert.Contains(t, groups, groupIDs[1])
assert.Contains(t, directPeers, peerIDs[3])
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.ElementsMatch(t, directPeers, []string{peerIDs[4]})
groups, directPeers = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[1]})
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.ElementsMatch(t, directPeers, []string{peerIDs[3]})
groups, directPeers = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[2]})
assert.Empty(t, groups)
assert.Empty(t, directPeers)
}
func TestCollectGroupChange_PolicyWithNonPeerResource_NoDirectPeers(t *testing.T) {
@@ -168,8 +213,7 @@ func TestCollectGroupChange_PolicyWithNonPeerResource_NoDirectPeers(t *testing.T
require.NoError(t, err)
groups, directPeers := collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[0]})
assert.Contains(t, groups, groupIDs[0])
assert.Contains(t, groups, groupIDs[1])
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.Empty(t, directPeers, "non-peer resources should not produce direct peer IDs")
}
@@ -294,6 +338,7 @@ func TestCollectGroupChange_NetworkRouterLinked(t *testing.T) {
AccountID: accountID,
PeerGroups: []string{groupIDs[0]},
Peer: peerIDs[3],
Enabled: true,
})
require.NoError(t, err)
@@ -324,6 +369,7 @@ func TestCollectGroupChange_NetworkRouterPeerOnlyNoGroups(t *testing.T) {
NetworkID: net1.ID,
AccountID: accountID,
Peer: peerIDs[4],
Enabled: true,
})
require.NoError(t, err)
@@ -373,17 +419,11 @@ func TestCollectGroupChange_MultipleEntities(t *testing.T) {
require.NoError(t, err)
groups, directPeers := collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[0]})
assert.Contains(t, groups, groupIDs[0])
assert.Contains(t, groups, groupIDs[1])
assert.NotContains(t, groups, groupIDs[2])
assert.NotContains(t, groups, groupIDs[3])
assert.ElementsMatch(t, groups, []string{groupIDs[0], groupIDs[1]})
assert.Empty(t, directPeers)
groups, directPeers = collectGroupChangeAffectedGroups(ctx, s, accountID, []string{groupIDs[3]})
assert.Contains(t, groups, groupIDs[2])
assert.Contains(t, groups, groupIDs[3])
assert.NotContains(t, groups, groupIDs[0])
assert.NotContains(t, groups, groupIDs[1])
assert.ElementsMatch(t, groups, []string{groupIDs[2], groupIDs[3]})
assert.Empty(t, directPeers)
}
@@ -452,8 +492,9 @@ func TestResolveAffectedPeers_PolicyBetweenTwoGroups(t *testing.T) {
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[1]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1]}, result)
// peerIDs[2] is unrelated to the route; only its own map can change.
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[2]})
assert.Empty(t, result)
assert.ElementsMatch(t, []string{peerIDs[2]}, result)
}
func TestResolveAffectedPeers_PolicyThreeGroups(t *testing.T) {
@@ -474,7 +515,7 @@ func TestResolveAffectedPeers_PolicyThreeGroups(t *testing.T) {
require.NoError(t, err)
result := manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[0]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1], peerIDs[2]}, result)
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[2]}, result)
}
func TestResolveAffectedPeers_RoutePeerGroups(t *testing.T) {
@@ -506,8 +547,9 @@ func TestResolveAffectedPeers_RoutePeerGroups(t *testing.T) {
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[1]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1]}, result)
// peerIDs[2] is in no policy; only its own map can change, so it refreshes itself.
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[2]})
assert.Empty(t, result)
assert.ElementsMatch(t, []string{peerIDs[2]}, result)
}
func TestResolveAffectedPeers_RouteWithDirectPeer(t *testing.T) {
@@ -564,9 +606,9 @@ func TestResolveAffectedPeers_RouteWithAccessControlGroups(t *testing.T) {
result := manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[2]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1], peerIDs[2]}, result)
// peer3 is unrelated
// peer3 is unrelated to the route; only its own map can change.
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[3]})
assert.Empty(t, result)
assert.ElementsMatch(t, []string{peerIDs[3]}, result)
}
func TestResolveAffectedPeers_NetworkRouter(t *testing.T) {
@@ -587,6 +629,7 @@ func TestResolveAffectedPeers_NetworkRouter(t *testing.T) {
AccountID: accountID,
PeerGroups: []string{groupIDs[0]},
Peer: peerIDs[3],
Enabled: true,
})
require.NoError(t, err)
@@ -659,9 +702,13 @@ func TestResolveAffectedPeers_PeerInMultipleGroups(t *testing.T) {
}, true)
require.NoError(t, err)
// peer0 is in group0 AND group1, so both policies apply
// peer0 is in group0 AND group1, so both policies apply. A peer change folds
// only the changed peer plus the opposite side of each rule: group2 (peer2) via
// the group0 policy and group3 (peer3) via the group1 policy. peer1, a co-member
// of group1, is a sibling of the changed peer and must NOT refresh.
result := manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[0]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1], peerIDs[2], peerIDs[3]}, result)
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[2], peerIDs[3]}, result)
assert.NotContains(t, result, peerIDs[1], "co-member of the changed peer's group must not refresh")
}
func TestResolveAffectedPeers_MultipleChangedPeers(t *testing.T) {
@@ -697,7 +744,7 @@ func TestResolveAffectedPeers_MultipleChangedPeers(t *testing.T) {
require.NoError(t, err)
result := manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[0], peerIDs[2]})
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[1], peerIDs[2], peerIDs[3]}, result)
assert.ElementsMatch(t, []string{peerIDs[0], peerIDs[2], peerIDs[1], peerIDs[3]}, result)
}
func TestResolveAffectedPeers_SharedGroupAcrossPolicyAndRoute(t *testing.T) {
@@ -854,8 +901,9 @@ func TestAffectedPeers_IsolatedPolicies(t *testing.T) {
assert.NotContains(t, result, peerIDs[0])
assert.NotContains(t, result, peerIDs[1])
// peerIDs[4] is in neither isolated policy; only its own map can change.
result = manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[4]})
assert.Empty(t, result)
assert.ElementsMatch(t, []string{peerIDs[4]}, result)
}
func TestAffectedPeers_IsolatedRouteAndPolicy(t *testing.T) {
@@ -977,12 +1025,13 @@ func TestAffectedPeers_GroupUpdateOnlyAffectsLinkedPeers(t *testing.T) {
})
}
func TestAffectedPeers_UnlinkedGroupChange_NoUpdates(t *testing.T) {
// A peer in no policy/route refreshes only itself — no other peer is affected.
func TestAffectedPeers_UnlinkedPeerChange_RefreshesSelfOnly(t *testing.T) {
manager, s, accountID, peerIDs, _ := setupAffectedPeersTest(t)
ctx := context.Background()
result := manager.resolveAffectedPeersForPeerChanges(ctx, s, accountID, []string{peerIDs[0]})
assert.Empty(t, result)
assert.ElementsMatch(t, []string{peerIDs[0]}, result)
}
// TestAffectedPeers_PolicyChange_UnrelatedPeerNoUpdate verifies that creating/deleting a
@@ -1332,6 +1381,7 @@ func TestAffectedPeers_NetworkRouterUnlinkedPeerNoUpdate(t *testing.T) {
NetworkID: net1.ID,
AccountID: accountID,
PeerGroups: []string{"nr-grpA"},
Enabled: true,
})
require.NoError(t, err)
@@ -1755,7 +1805,9 @@ func TestCollectAffectedFromProxyServices_GroupContainingTargetPeerChanged(t *te
assert.Contains(t, directPeers, peerIDs[1], "target peer must be refreshed")
}
func TestCollectAffectedFromProxyServices_DisabledServiceStillMatches(t *testing.T) {
// A disabled service in the snapshot proxies nothing, so it is skipped: a changed
// target peer does not pull in the service's proxy peer.
func TestCollectAffectedFromProxyServices_DisabledServiceSkipped(t *testing.T) {
manager, s, accountID, peerIDs, _ := setupAffectedPeersTest(t)
ctx := context.Background()
@@ -1781,8 +1833,7 @@ func TestCollectAffectedFromProxyServices_DisabledServiceStillMatches(t *testing
require.NoError(t, s.CreateService(ctx, svc))
_, directPeers := collectPeerChangeAffectedGroups(ctx, manager.Store, accountID, nil, []string{peerIDs[1]})
assert.Contains(t, directPeers, peerIDs[0], "disabled service should still trigger a refresh so peers are ready when re-enabled")
assert.Contains(t, directPeers, peerIDs[1], "disabled target should still trigger a refresh")
assert.NotContains(t, directPeers, peerIDs[0], "a disabled service proxies nothing, so its proxy peer must not be folded")
}
func TestCollectAffectedFromProxyServices_NonPeerTargetType(t *testing.T) {
@@ -0,0 +1,95 @@
package affectedpeers
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/server/store"
)
// fakeProxyStore implements only the two store methods loadProxyServices calls;
// the embedded nil store.Store panics if anything else is invoked, which keeps
// the test honest about the surface under test.
type fakeProxyStore struct {
store.Store
proxyByCluster map[string][]string
persisted []*rpservice.Service
}
func (f *fakeProxyStore) GetEmbeddedProxyPeerIDsByCluster(_ context.Context, _ string) (map[string][]string, error) {
return f.proxyByCluster, nil
}
func (f *fakeProxyStore) GetAccountServices(_ context.Context, _ store.LockingStrength, _ string) ([]*rpservice.Service, error) {
return f.persisted, nil
}
func serviceIDs(svcs []*rpservice.Service) []string {
ids := make([]string, 0, len(svcs))
for _, s := range svcs {
ids = append(ids, s.ID)
}
return ids
}
// loadProxyServices must merge the synthesised agent-network services (which are
// never persisted) with the persisted ones, so the proxy-affected expansion can
// see agent-network AccessGroups. Without this the embedded proxy peer is never
// flagged on a client group change and only a full resync (restart) recovers.
func TestLoadProxyServices_MergesSynthesizedAgentNetworkServices(t *testing.T) {
prev := agentNetworkSynthesizer
t.Cleanup(func() { agentNetworkSynthesizer = prev })
SetAgentNetworkSynthesizer(func(_ context.Context, _ store.Store, _ string) ([]*rpservice.Service, error) {
return []*rpservice.Service{
{ID: "agent-net-svc-acc", ProxyCluster: "proxy.netbird.local", Private: true, AccessGroups: []string{"gB"}},
}, nil
})
s := &fakeProxyStore{
proxyByCluster: map[string][]string{"proxy.netbird.local": {"proxy-peer-1"}},
persisted: []*rpservice.Service{{ID: "persisted-rp-svc", ProxyCluster: "proxy.netbird.local"}},
}
snap := &Snapshot{}
require.NoError(t, snap.loadProxyServices(context.Background(), s, "acc"))
ids := serviceIDs(snap.services)
assert.Contains(t, ids, "persisted-rp-svc", "persisted services must be kept")
assert.Contains(t, ids, "agent-net-svc-acc", "synthesised agent-network service must be merged in")
}
// With no synthesiser registered, loadProxyServices falls back to persisted
// services only (no panic, no behaviour change for non-agent-network builds).
func TestLoadProxyServices_NoSynthesizerRegistered(t *testing.T) {
prev := agentNetworkSynthesizer
t.Cleanup(func() { agentNetworkSynthesizer = prev })
agentNetworkSynthesizer = nil
s := &fakeProxyStore{
proxyByCluster: map[string][]string{"c": {"proxy-1"}},
persisted: []*rpservice.Service{{ID: "persisted"}},
}
snap := &Snapshot{}
require.NoError(t, snap.loadProxyServices(context.Background(), s, "acc"))
assert.Equal(t, []string{"persisted"}, serviceIDs(snap.services))
}
// No embedded proxy peers → skip entirely (don't even call the synthesiser).
func TestLoadProxyServices_NoEmbeddedProxyPeersSkips(t *testing.T) {
prev := agentNetworkSynthesizer
t.Cleanup(func() { agentNetworkSynthesizer = prev })
called := false
SetAgentNetworkSynthesizer(func(_ context.Context, _ store.Store, _ string) ([]*rpservice.Service, error) {
called = true
return nil, nil
})
s := &fakeProxyStore{proxyByCluster: map[string][]string{}}
snap := &Snapshot{}
require.NoError(t, snap.loadProxyServices(context.Background(), s, "acc"))
assert.False(t, called, "synthesiser must not run for accounts without embedded proxy peers")
assert.Empty(t, snap.services)
}
File diff suppressed because it is too large Load Diff
@@ -10,8 +10,8 @@ import (
"github.com/netbirdio/netbird/management/server/types"
)
// policyGroupsAndPeers mirrors the explicit-policy extraction (RuleGroups +
// direct peers) the resolver folds in, for asserting the pure logic.
// policyGroupsAndPeers mirrors the both-sides extraction (RuleGroups + direct peers)
// the resolver folds in for a changed policy, for asserting the pure logic.
func policyGroupsAndPeers(policies ...*types.Policy) (groups []string, peers []string) {
peerSet := map[string]struct{}{}
for _, p := range policies {
@@ -19,7 +19,14 @@ func policyGroupsAndPeers(policies ...*types.Policy) (groups []string, peers []s
continue
}
groups = append(groups, p.RuleGroups()...)
collectPolicyDirectPeers(p, peerSet)
for _, rule := range p.Rules {
if rule.SourceResource.Type == types.ResourceTypePeer && rule.SourceResource.ID != "" {
peerSet[rule.SourceResource.ID] = struct{}{}
}
if rule.DestinationResource.Type == types.ResourceTypePeer && rule.DestinationResource.ID != "" {
peerSet[rule.DestinationResource.ID] = struct{}{}
}
}
}
for id := range peerSet {
peers = append(peers, id)
@@ -80,26 +87,6 @@ func TestChangeIsEmpty(t *testing.T) {
assert.False(t, Change{PostureCheckIDs: []string{"pc"}}.isEmpty())
}
func TestPolicyReferencesGroups(t *testing.T) {
policy := &types.Policy{Rules: []*types.PolicyRule{{Sources: []string{"g1", "g2"}, Destinations: []string{"g3"}}}}
assert.True(t, policyReferencesGroups(policy, map[string]struct{}{"g1": {}}))
assert.True(t, policyReferencesGroups(policy, map[string]struct{}{"g3": {}}))
assert.False(t, policyReferencesGroups(policy, map[string]struct{}{"g4": {}}))
assert.False(t, policyReferencesGroups(policy, map[string]struct{}{}))
}
func TestPolicyReferencesDirectPeers(t *testing.T) {
policy := &types.Policy{Rules: []*types.PolicyRule{{
SourceResource: types.Resource{Type: types.ResourceTypePeer, ID: "p1"},
DestinationResource: types.Resource{Type: types.ResourceTypeHost, ID: "r1"},
}}}
assert.True(t, policyReferencesDirectPeers(policy, map[string]struct{}{"p1": {}}))
assert.False(t, policyReferencesDirectPeers(policy, map[string]struct{}{"r1": {}}))
assert.False(t, policyReferencesDirectPeers(policy, map[string]struct{}{"p2": {}}))
}
func TestPolicyReferencesPostureChecks(t *testing.T) {
policy := &types.Policy{SourcePostureChecks: []string{"pc1", "pc2"}}
@@ -107,24 +94,9 @@ func TestPolicyReferencesPostureChecks(t *testing.T) {
assert.False(t, policyReferencesPostureChecks(policy, map[string]struct{}{"pc3": {}}))
}
func TestCollectPolicyDirectPeers(t *testing.T) {
policy := &types.Policy{Rules: []*types.PolicyRule{{
SourceResource: types.Resource{Type: types.ResourceTypePeer, ID: "p1"},
DestinationResource: types.Resource{Type: types.ResourceTypePeer, ID: "p2"},
}, {
DestinationResource: types.Resource{Type: types.ResourceTypeHost, ID: "r1"},
}}}
peerSet := map[string]struct{}{}
collectPolicyDirectPeers(policy, peerSet)
assert.Contains(t, peerSet, "p1")
assert.Contains(t, peerSet, "p2")
assert.NotContains(t, peerSet, "r1")
}
func TestCollectPolicySources(t *testing.T) {
policy := &types.Policy{Rules: []*types.PolicyRule{{
Enabled: true,
Sources: []string{"g1"},
SourceResource: types.Resource{Type: types.ResourceTypePeer, ID: "p1"},
Destinations: []string{"g2"},
@@ -0,0 +1,126 @@
package server
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/store"
)
// TestAgentNetwork_BudgetRuleCRUD_RealManager is the GC-1 no-mock guard for the
// account budget-rule manager surface: real DefaultAccountManager, real store,
// real permissions. It exercises create/get/list/update/delete through the
// permission-gated manager (not the store directly) and asserts the reused
// PolicyLimits cap shape and targets survive each step.
func TestAgentNetwork_BudgetRuleCRUD_RealManager(t *testing.T) {
am, _, err := createManager(t)
require.NoError(t, err, "createManager must succeed")
ctx := context.Background()
const (
accountID = "agent-net-budget-acct"
adminUserID = "agent-net-budget-admin"
)
account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false)
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed")
mgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
created, err := mgr.CreateBudgetRule(ctx, adminUserID, &agenttypes.AccountBudgetRule{
AccountID: accountID,
Name: "eng-monthly",
Enabled: true,
TargetGroups: []string{"grp-eng"},
TargetUsers: []string{"user-alice"},
Limits: agenttypes.PolicyLimits{
TokenLimit: agenttypes.PolicyTokenLimit{Enabled: true, GroupCap: 100_000, UserCap: 10_000, WindowSeconds: 2_592_000},
BudgetLimit: agenttypes.PolicyBudgetLimit{Enabled: true, GroupCapUsd: 500, WindowSeconds: 2_592_000},
},
})
require.NoError(t, err, "CreateBudgetRule must succeed")
require.NotEmpty(t, created.ID, "create must mint an ID")
got, err := mgr.GetBudgetRule(ctx, accountID, adminUserID, created.ID)
require.NoError(t, err, "GetBudgetRule must succeed")
assert.Equal(t, "eng-monthly", got.Name, "name round-trips through the manager")
assert.Equal(t, []string{"grp-eng"}, got.TargetGroups, "target groups round-trip")
assert.Equal(t, int64(100_000), got.Limits.TokenLimit.GroupCap, "token group cap round-trips")
list, err := mgr.GetAllBudgetRules(ctx, accountID, adminUserID)
require.NoError(t, err, "GetAllBudgetRules must succeed")
require.Len(t, list, 1, "exactly the one created rule must be listed")
created.Limits.TokenLimit.GroupCap = 200_000
updated, err := mgr.UpdateBudgetRule(ctx, adminUserID, created)
require.NoError(t, err, "UpdateBudgetRule must succeed")
assert.Equal(t, int64(200_000), updated.Limits.TokenLimit.GroupCap, "updated cap must persist")
require.NoError(t, mgr.DeleteBudgetRule(ctx, accountID, adminUserID, created.ID), "DeleteBudgetRule must succeed")
_, err = mgr.GetBudgetRule(ctx, accountID, adminUserID, created.ID)
assert.Error(t, err, "get after delete must fail")
}
// TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection is the
// GC-1 guard for UpdateSettings: it must apply the collection toggles while
// preserving the immutable Cluster/Subdomain pinned at bootstrap.
func TestAgentNetwork_UpdateSettings_PreservesImmutableAndTogglesCollection(t *testing.T) {
am, _, err := createManager(t)
require.NoError(t, err, "createManager must succeed")
ctx := context.Background()
const (
accountID = "agent-net-settings-acct"
adminUserID = "agent-net-settings-admin"
clusterAddr = "eu.proxy.netbird.io"
)
account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false)
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed")
mgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
// Creating a provider bootstraps the settings row (cluster + subdomain).
_, err = mgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{
AccountID: accountID,
ProviderID: "openai_api",
Name: "openai",
UpstreamURL: "https://api.openai.com",
APIKey: "sk-test",
Enabled: true,
Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}},
}, clusterAddr)
require.NoError(t, err, "CreateProvider must bootstrap settings")
before, err := mgr.GetSettings(ctx, accountID, adminUserID)
require.NoError(t, err, "GetSettings must succeed after bootstrap")
require.Equal(t, clusterAddr, before.Cluster, "cluster pinned at bootstrap")
require.NotEmpty(t, before.Subdomain, "subdomain pinned at bootstrap")
assert.False(t, before.EnablePromptCollection, "prompt collection defaults off")
// Attempt to flip toggles AND smuggle a different cluster/subdomain — the
// immutable fields must be ignored.
updated, err := mgr.UpdateSettings(ctx, adminUserID, &agenttypes.Settings{
AccountID: accountID,
Cluster: "attacker.cluster",
Subdomain: "evil",
EnableLogCollection: true,
EnablePromptCollection: true,
RedactPii: true,
})
require.NoError(t, err, "UpdateSettings must succeed")
assert.Equal(t, before.Cluster, updated.Cluster, "cluster is immutable and must be preserved")
assert.Equal(t, before.Subdomain, updated.Subdomain, "subdomain is immutable and must be preserved")
assert.True(t, updated.EnableLogCollection, "log collection toggle must apply")
assert.True(t, updated.EnablePromptCollection, "prompt collection toggle must apply")
assert.True(t, updated.RedactPii, "redact toggle must apply")
reloaded, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
require.NoError(t, err)
assert.Equal(t, before.Cluster, reloaded.Cluster, "persisted cluster unchanged")
assert.True(t, reloaded.EnablePromptCollection, "persisted prompt collection toggled on")
}
@@ -0,0 +1,199 @@
package server
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/peers"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
)
// TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale is the no-mock
// regression guard for the bug the user reported: restarting the proxy creates
// a fresh embedded peer with a NEW WireGuard public key (the proxy generates
// the keypair on every startup at proxy/internal/roundtrip/netbird.go:312).
// The PRIOR embedded peer record is never deleted on management, so the
// account accumulates a stale peer holding a stale CGNAT IP. Other peers
// in the account either keep routing to the dead IP, or — if synth DNS
// picks the wrong record — never see the new IP at all.
//
// What this test exercises (no mocks):
// - real SQLite test store
// - real DefaultAccountManager, network-map controller, peer-update channels
// - real peers.Manager.CreateProxyPeer path (the very method the proxy
// invokes over gRPC on every startup)
// - real agentnetwork.Manager + synth chain so the client receives a
// concrete DNS record that must point at the LATEST proxy peer.
//
// Pre-fix expected behavior (red): two embedded peers exist after the
// "restart"; the synth DNS record points at the stale one; the client
// receives an update reflecting the new peer but the old one lingers.
// Post-fix expected behavior (green): exactly one embedded peer exists
// after restart (with the new key) AND the client's network map carries
// the synth DNS pointing at that new peer's CGNAT IP.
func TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale(t *testing.T) {
am, updateManager, err := createManager(t)
require.NoError(t, err, "createManager must succeed")
ctx := context.Background()
const (
accountID = "an-restart-acct"
adminUserID = "an-restart-admin"
groupAID = "an-restart-grp-A"
clusterAddr = "eu.proxy.netbird.io"
clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
// Two different proxy pubkeys — the "before" and "after" of a
// proxy-process restart with fresh-keypair generation.
proxyKey1 = "Aaaaa1aaaaYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
proxyKey2 = "Bbbbb2bbbbYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
)
// --- Account scaffold ---
account := newAccountWithId(ctx, accountID, adminUserID, "an-restart.test", "", "", false)
require.NoError(t, am.Store.SaveAccount(ctx, account))
clientPeer := &nbpeer.Peer{
Key: clientKey,
Name: "an-restart-client",
DNSLabel: "an-restart-client",
Meta: nbpeer.PeerSystemMeta{Hostname: "an-restart-client", GoOS: "linux", WtVersion: "development"},
}
addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false)
require.NoError(t, err, "AddPeer for client must succeed")
require.NoError(t, am.MarkPeerConnected(ctx, clientKey, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for the client peer must succeed (affected-peer fan-out skips disconnected peers)")
// Place the client in group A so the synth policy reaches it.
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}}
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist group A")
// --- Real peers + agent-network managers ---
permMgr := permissions.NewManager(am.Store)
peersMgr := peers.NewManager(am.Store, permMgr)
peersMgr.SetAccountManager(am)
peersMgr.SetNetworkMapController(am.networkMapController)
agentMgr := agentnetwork.NewManager(am.Store, permMgr, am, nil)
// Subscribe BEFORE any state-mutating call so we don't lose the update
// that contains the synth DNS record.
clientCh := updateManager.CreateChannel(ctx, addedClient.ID)
t.Cleanup(func() { updateManager.CloseChannel(ctx, addedClient.ID) })
drain(clientCh)
// --- First proxy startup: register peer key K1, then mark it
// connected. In production the proxy follows CreateProxyPeer with the
// regular sync stream which lands on MarkPeerConnected; the synth DNS
// path filters out peers that aren't Connected (types/account.go:323),
// so without this step no DNS record would be emitted.
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey1, clusterAddr),
"first CreateProxyPeer (proxy startup) must succeed")
peer1ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
require.NoError(t, err, "proxy peer for K1 must be persisted after CreateProxyPeer")
require.NotEmpty(t, peer1ID)
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey1, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for K1 must succeed")
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
proxyIP1 := account.Peers[peer1ID].IP.String()
require.NotEmpty(t, proxyIP1, "K1 must have an assigned overlay IP")
// --- Provider + policy. CreateProvider / CreatePolicy trigger the
// agentnetwork reconcile which runs UpdateAccountPeers; the resulting
// NetworkMap delivered to the client carries the synth DNS record
// pointing at K1's IP. ---
provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{
AccountID: accountID,
ProviderID: "openai_api",
Name: "openai-test",
UpstreamURL: "https://api.openai.com",
APIKey: "sk-test-key",
Enabled: true,
Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}},
}, clusterAddr)
require.NoError(t, err, "CreateProvider must succeed")
_, err = agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{
AccountID: accountID,
Name: "p1",
Enabled: true,
SourceGroups: []string{groupAID},
DestinationProviderIDs: []string{provider.ID},
})
require.NoError(t, err, "CreatePolicy must succeed")
settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
require.NoError(t, err)
fqdn := settings.Endpoint()
rdata1 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
require.Equal(t, proxyIP1, rdata1,
"client must receive a synth DNS record pointing at K1's overlay IP after the synth path runs")
drain(clientCh)
// --- Proxy restart: NEW keypair K2, same account, same cluster ---
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey2, clusterAddr),
"second CreateProxyPeer (proxy restart with fresh keypair) must succeed")
peer2ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey2)
require.NoError(t, err, "proxy peer for K2 must be persisted after restart")
require.NotEmpty(t, peer2ID)
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey2, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
"MarkPeerConnected for K2 must succeed")
// In production the agent's sync stream pulls a fresh NetworkMap as
// part of its normal reconcile cadence; in this isolated test
// MarkPeerConnected's affected-peer fan-out can race the channel-side
// buffer in a way that swallows the synth-DNS-bearing update before
// our await reads it. Trigger an explicit account-wide fan-out so the
// assertion below tests what production actually delivers, not the
// in-test buffer race.
am.UpdateAccountPeers(ctx, accountID, types.UpdateReason{Resource: types.UpdateResourcePeer, Operation: types.UpdateOperationUpdate})
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
proxyIP2 := account.Peers[peer2ID].IP.String()
require.NotEmpty(t, proxyIP2, "K2 must have an assigned overlay IP")
require.NotEqual(t, proxyIP1, proxyIP2, "K2 must get a different overlay IP than K1 (sanity)")
// CRITICAL ASSERTION 1: K1 must no longer be in the store. The SqlStore
// returns ("", nil) for a missing key rather than NotFound, so assert
// on the returned ID being empty.
staleID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
require.NoError(t, err, "GetPeerIDByKey for a missing peer must not error")
assert.Empty(t, staleID,
"stale embedded proxy peer K1 must be removed when a new embedded peer registers for the same (account, cluster); pre-fix this assertion fails because management never cleans up the prior peer record")
// CRITICAL ASSERTION 2: exactly one embedded proxy peer remains, and it
// is K2.
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
embeddedKeys := []string{}
for _, p := range account.Peers {
if p.ProxyMeta.Embedded {
embeddedKeys = append(embeddedKeys, p.Key)
}
}
assert.Equal(t, []string{proxyKey2}, embeddedKeys,
"after a proxy restart exactly one embedded proxy peer should remain — the one with the new key K2")
// CRITICAL ASSERTION 3: the synth DNS record the client receives now
// points at K2's IP, not K1's.
rdata2 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
assert.Equal(t, proxyIP2, rdata2,
"after proxy restart, the client's synth DNS record must point at the NEW embedded peer's IP, not the stale K1 IP")
}
@@ -0,0 +1,212 @@
package server
import (
"context"
"net/netip"
"testing"
"time"
"github.com/miekg/dns"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
networkmap "github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
nbproto "github.com/netbirdio/netbird/shared/management/proto"
)
// TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers is the no-mock
// integration test for the live propagation path: a provider/policy mutation
// through the real agentnetwork.Manager triggers the real
// DefaultAccountManager.UpdateAccountPeers, which runs the real network-map
// controller (including AN-2b's injectAllProxyPolicies), and a network map is
// computed and fanned out to BOTH the embedded proxy peer and the client peer.
//
// Unlike the synthesizer/reconcile unit tests, nothing here is mocked: real
// SQLite store, real account manager + network-map controller, real
// agentnetwork manager, real peer update channels. The client peer's delivered
// map is asserted to actually carry the synth DNS surface, and provider
// create/delete are exercised end to end.
func TestAgentNetwork_ProviderCRUD_FansOutToProxyAndClientPeers(t *testing.T) {
am, updateManager, err := createManager(t)
require.NoError(t, err, "createManager must succeed")
ctx := context.Background()
const (
accountID = "agent-net-acct-1"
adminUserID = "agent-net-admin-1"
groupAID = "agent-net-grp-A"
clusterAddr = "eu.proxy.netbird.io"
clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
proxyPeerID = "agent-net-proxy-peer-1"
proxyPeerKey = "/yF0+vCfv+mRR5k0dca0TrGdO/oiNeAI58gToZm5NyI="
proxyIP = "100.64.0.99"
)
account := newAccountWithId(ctx, accountID, adminUserID, "agent-net.test", "", "", false)
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must succeed")
// Real client peer through the production AddPeer path.
clientPeer := &nbpeer.Peer{
Key: clientKey,
Name: "agent-net-client",
DNSLabel: "agent-net-client",
Meta: nbpeer.PeerSystemMeta{Hostname: "agent-net-client", GoOS: "linux", WtVersion: "development"},
}
addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false)
require.NoError(t, err, "AddPeer must add the client peer")
// Inject a connected embedded proxy peer + put the client in the source group.
account, err = am.Store.GetAccount(ctx, accountID)
require.NoError(t, err)
account.Peers[proxyPeerID] = &nbpeer.Peer{
ID: proxyPeerID,
AccountID: accountID,
Key: proxyPeerKey,
IP: netip.MustParseAddr(proxyIP),
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
ProxyMeta: nbpeer.ProxyMeta{Embedded: true, Cluster: clusterAddr},
DNSLabel: "agent-net-proxy",
}
account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}}
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist proxy peer + group")
// Subscribe to BOTH peers' update channels — this is how we observe the
// real fan-out.
clientCh := updateManager.CreateChannel(ctx, addedClient.ID)
proxyCh := updateManager.CreateChannel(ctx, proxyPeerID)
t.Cleanup(func() {
updateManager.CloseChannel(ctx, addedClient.ID)
updateManager.CloseChannel(ctx, proxyPeerID)
})
drain(clientCh)
drain(proxyCh)
// Real agentnetwork manager wired to the real account manager. proxyController
// is nil (no gRPC cluster fan-out here) — the reconcile still fires
// UpdateAccountPeers, which is the path under test.
agentMgr := agentnetwork.NewManager(am.Store, permissions.NewManager(am.Store), am, nil)
provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{
AccountID: accountID,
ProviderID: "openai_api",
Name: "openai-test",
UpstreamURL: "https://api.openai.com",
APIKey: "sk-test-key",
Enabled: true,
Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}},
}, clusterAddr)
require.NoError(t, err, "CreateProvider must succeed")
policy, err := agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{
AccountID: accountID,
Name: "p1",
Enabled: true,
SourceGroups: []string{groupAID},
DestinationProviderIDs: []string{provider.ID},
})
require.NoError(t, err, "CreatePolicy must succeed")
settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
require.NoError(t, err)
fqdn := settings.Endpoint()
// Both peers must receive a fan-out. The provider-create reconcile fires
// before the policy exists (synth service then has no AccessGroups, so no
// zone), and the async update buffer can collapse/reorder updates — so we
// poll until the client's delivered map actually carries the synth record.
rdata := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
assert.Equal(t, proxyIP, rdata,
"client peer's delivered network map must contain the synth DNS record pointing at the embedded proxy peer")
require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after create")
// UPDATE the provider — a new model on the existing service must still
// reconcile and keep the private surface routable (the live MODIFIED path).
provider.Models = append(provider.Models, agenttypes.ProviderModel{ID: "gpt-5.4-mini"})
_, err = agentMgr.UpdateProvider(ctx, adminUserID, provider)
require.NoError(t, err, "UpdateProvider must succeed")
assert.Equal(t, proxyIP, awaitZoneRData(clientCh, clusterAddr, fqdn, true),
"client peer must still resolve the synth record after the provider is updated")
require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after update")
// DELETE: detach the policy first (provider is in use), then drop the
// provider. Both peers update again and the synth surface disappears.
require.NoError(t, agentMgr.DeletePolicy(ctx, accountID, adminUserID, policy.ID), "DeletePolicy must succeed")
require.NoError(t, agentMgr.DeleteProvider(ctx, accountID, adminUserID, provider.ID), "DeleteProvider must succeed")
require.True(t, awaitUpdate(proxyCh), "embedded proxy peer must also receive a netmap update after delete")
assert.Empty(t, awaitZoneRData(clientCh, clusterAddr, fqdn, false),
"synth DNS record must be gone from the client's map after the provider is deleted")
}
// awaitZoneRData drains the channel for up to 8s. When wantPresent is true it
// returns as soon as the synth record appears (its RData). When false it drains
// to quiescence and returns the RData of the last delivered map (expected empty
// once the provider is gone), tolerating stale buffered updates that still
// carry the zone.
func awaitZoneRData(ch <-chan *networkmap.UpdateMessage, clusterAddr, fqdn string, wantPresent bool) string {
deadline := time.After(8 * time.Second)
last := ""
for {
select {
case m := <-ch:
if m == nil {
continue
}
last = synthZoneRData(m.Update, clusterAddr, fqdn)
if wantPresent && last != "" {
return last
}
case <-time.After(750 * time.Millisecond):
return last
case <-deadline:
return last
}
}
}
// awaitUpdate reports whether at least one update arrives within the window.
func awaitUpdate(ch <-chan *networkmap.UpdateMessage) bool {
select {
case m := <-ch:
return m != nil
case <-time.After(5 * time.Second):
return false
}
}
// drain empties any buffered updates (e.g. from AddPeer/SaveAccount) so the
// next observation reflects the operation under test.
func drain(ch <-chan *networkmap.UpdateMessage) {
for {
select {
case <-ch:
case <-time.After(200 * time.Millisecond):
return
}
}
}
// synthZoneRData returns the RData of the synth A record (record name == fqdn)
// inside the cluster's custom zone, or "" when absent.
func synthZoneRData(sync *nbproto.SyncResponse, clusterAddr, fqdn string) string {
if sync == nil {
return ""
}
for _, zone := range sync.GetNetworkMap().GetDNSConfig().GetCustomZones() {
if zone.GetDomain() != dns.Fqdn(clusterAddr) {
continue
}
for _, rec := range zone.GetRecords() {
if rec.GetName() == dns.Fqdn(fqdn) {
return rec.GetRData()
}
}
}
return ""
}
+11 -7
View File
@@ -520,7 +520,12 @@ func collectDeletableGroups(ctx context.Context, transaction store.Store, accoun
// GroupAddPeer appends peer to the group
func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, groupID, peerID string) error {
var snap *affectedpeers.Snapshot
change := affectedpeers.Change{ChangedGroupIDs: []string{groupID}}
// A membership change affects only the peer itself and the opposite side of THIS
// group's policies — not the group's other members, and not the peer's other
// groups. LinkGroups walks only this group (matched, not expanded); OutputPeerIDs
// refreshes the peer without seeding its other group memberships. For an
// intra-group policy the opposite side is the group, so its members still refresh.
change := affectedpeers.Change{OutputPeerIDs: []string{peerID}, LinkGroups: []string{groupID}}
err := am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
if err := transaction.AddPeerToGroup(ctx, accountID, peerID, groupID); err != nil {
@@ -586,10 +591,11 @@ func (am *DefaultAccountManager) GroupAddResource(ctx context.Context, accountID
// GroupDeletePeer removes peer from the group
func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, groupID, peerID string) error {
var snap *affectedpeers.Snapshot
change := affectedpeers.Change{
ChangedGroupIDs: []string{groupID},
RemovedPeersByGroup: map[string][]string{groupID: {peerID}},
}
// Same as GroupAddPeer: the removed peer and the opposite side of THIS group's
// policies refresh, not the group's other members or the peer's other groups. The
// peer is no longer in the group's index, but LinkGroups still drives the
// opposite-side walk, and OutputPeerIDs refreshes the removed peer itself.
change := affectedpeers.Change{OutputPeerIDs: []string{peerID}, LinkGroups: []string{groupID}}
err := am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
if err := transaction.RemovePeerFromGroup(ctx, peerID, groupID); err != nil {
@@ -600,8 +606,6 @@ func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID,
return err
}
// The removed peer is carried in change.RemovedPeersByGroup and folded in
// only when the group is linked, so loading post-removal is correct.
var err error
if snap, err = affectedpeers.Load(ctx, transaction, accountID, change); err != nil {
return err
+6 -1
View File
@@ -23,6 +23,8 @@ import (
idpmanager "github.com/netbirdio/netbird/management/server/idp"
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentnetworkhandlers "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/handlers"
"github.com/netbirdio/netbird/management/internals/modules/zones"
zonesManager "github.com/netbirdio/netbird/management/internals/modules/zones/manager"
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
@@ -59,7 +61,7 @@ import (
)
// NewAPIHandler creates the Management service HTTP API handler registering all the available endpoints.
func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager account.Manager, networksManager nbnetworks.Manager, resourceManager resources.Manager, routerManager routers.Manager, groupsManager nbgroups.Manager, LocationManager geolocation.Geolocation, authManager auth.Manager, appMetrics telemetry.AppMetrics, permissionsManager permissions.Manager, settingsManager settings.Manager, zManager zones.Manager, rManager records.Manager, networkMapController network_map.Controller, idpManager idpmanager.Manager, serviceManager service.Manager, reverseProxyDomainManager *manager.Manager, reverseProxyAccessLogsManager accesslogs.Manager, proxyGRPCServer *nbgrpc.ProxyServiceServer, trustedHTTPProxies []netip.Prefix, rateLimiter *middleware.APIRateLimiter, isValidChildAccount middleware.IsValidChildAccountFunc) (http.Handler, error) {
func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager account.Manager, networksManager nbnetworks.Manager, resourceManager resources.Manager, routerManager routers.Manager, groupsManager nbgroups.Manager, LocationManager geolocation.Geolocation, authManager auth.Manager, appMetrics telemetry.AppMetrics, permissionsManager permissions.Manager, settingsManager settings.Manager, zManager zones.Manager, rManager records.Manager, networkMapController network_map.Controller, idpManager idpmanager.Manager, serviceManager service.Manager, reverseProxyDomainManager *manager.Manager, reverseProxyAccessLogsManager accesslogs.Manager, proxyGRPCServer *nbgrpc.ProxyServiceServer, trustedHTTPProxies []netip.Prefix, rateLimiter *middleware.APIRateLimiter, isValidChildAccount middleware.IsValidChildAccountFunc, agentNetworkManager agentnetwork.Manager) (http.Handler, error) {
// Register bypass paths for unauthenticated endpoints
if err := bypass.AddBypassPath("/api/instance"); err != nil {
@@ -124,6 +126,9 @@ func NewAPIHandler(ctx context.Context, router *mux.Router, accountManager accou
zonesManager.RegisterEndpoints(router, zManager)
recordsManager.RegisterEndpoints(router, rManager)
idp.AddEndpoints(accountManager, router)
if agentNetworkManager != nil {
agentnetworkhandlers.RegisterEndpoints(agentNetworkManager, router)
}
instance.AddEndpoints(instanceManager, accountManager, router)
instance.AddVersionEndpoint(instanceManager, router)
if serviceManager != nil && reverseProxyDomainManager != nil {
@@ -283,6 +283,9 @@ func (h *handler) updateAccountRequestSettings(req api.PutApiAccountsAccountIdJS
if req.Settings.Ipv6EnabledGroups != nil {
returnSettings.IPv6EnabledGroups = *req.Settings.Ipv6EnabledGroups
}
if req.Settings.MetricsPushEnabled != nil {
returnSettings.MetricsPushEnabled = *req.Settings.MetricsPushEnabled
}
return returnSettings, nil
}
@@ -413,6 +416,7 @@ func toAccountResponse(accountID string, settings *types.Settings, meta *types.A
AutoUpdateVersion: &settings.AutoUpdateVersion,
AutoUpdateAlways: &settings.AutoUpdateAlways,
Ipv6EnabledGroups: &settings.IPv6EnabledGroups,
MetricsPushEnabled: &settings.MetricsPushEnabled,
EmbeddedIdpEnabled: &settings.EmbeddedIdpEnabled,
LocalAuthDisabled: &settings.LocalAuthDisabled,
LocalMfaEnabled: &settings.LocalMfaEnabled,
@@ -129,6 +129,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr(""),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -156,6 +157,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr(""),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -183,6 +185,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr("latest"),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -210,6 +213,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr(""),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -237,6 +241,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr(""),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -264,6 +269,7 @@ func TestAccounts_AccountsHandler(t *testing.T) {
DnsDomain: sr(""),
AutoUpdateAlways: br(false),
AutoUpdateVersion: sr(""),
MetricsPushEnabled: br(false),
EmbeddedIdpEnabled: br(false),
LocalAuthDisabled: br(false),
LocalMfaEnabled: br(false),
@@ -220,7 +220,7 @@ func (h *handler) getAllUsers(w http.ResponseWriter, r *http.Request) {
}
includeServiceUser, err := strconv.ParseBool(serviceUser)
log.WithContext(r.Context()).Debugf("Should include service user: %v", includeServiceUser)
log.WithContext(r.Context()).Tracef("Should include service user: %v", includeServiceUser)
if err != nil {
util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid service_user query parameter"), w)
return
@@ -152,7 +152,11 @@ func (m *AuthMiddleware) checkJWTFromRequest(r *http.Request, authHeaderParts []
return err
}
err = m.syncUserJWTGroups(ctx, userAuth)
// Detach the group-sync write from the request's cancellation: the dashboard
// SPA aborts in-flight requests on re-render, which would otherwise cancel the
// DB transaction mid-write and silently drop the synced groups. Context values
// (request id, logger) are preserved; the store bounds the tx with its own timeout.
err = m.syncUserJWTGroups(context.WithoutCancel(ctx), userAuth)
if err != nil {
log.WithContext(ctx).Errorf("HTTP server failed to sync user JWT groups: %s", err)
}
@@ -241,6 +241,66 @@ func TestAuthMiddleware_Handler(t *testing.T) {
}
}
// TestAuthMiddleware_SyncUserJWTGroupsDetachedFromRequestCancellation ensures the
// JWT group sync write is not bound to the request context. The dashboard SPA
// routinely aborts in-flight requests on re-render/navigation; if the sync ran in
// the request context, the cancellation would roll back the DB transaction and the
// synced groups would silently never persist. The sync must receive a context that
// is not cancelled even when the originating request is.
func TestAuthMiddleware_SyncUserJWTGroupsDetachedFromRequestCancellation(t *testing.T) {
var (
syncCalled bool
syncCtxErr error
)
mockAuth := &auth.MockManager{
ValidateAndParseTokenFunc: mockValidateAndParseToken,
EnsureUserAccessByJWTGroupsFunc: mockEnsureUserAccessByJWTGroups,
MarkPATUsedFunc: mockMarkPATUsed,
GetPATInfoFunc: mockGetAccountInfoFromPAT,
}
disabledLimiter := NewAPIRateLimiter(nil)
disabledLimiter.SetEnabled(false)
authMiddleware := NewAuthMiddleware(
mockAuth,
func(ctx context.Context, userAuth nbauth.UserAuth) (string, string, error) {
return userAuth.AccountId, userAuth.UserId, nil
},
func(ctx context.Context, userAuth nbauth.UserAuth) error {
syncCalled = true
syncCtxErr = ctx.Err()
return nil
},
func(ctx context.Context, userAuth nbauth.UserAuth) (*types.User, error) {
return &types.User{}, nil
},
disabledLimiter,
nil,
func(_ context.Context, _, _, _ string) bool { return false },
)
handlerToTest := authMiddleware.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
// Simulate the dashboard aborting the request: it arrives already cancelled.
ctx, cancel := context.WithCancel(context.Background())
cancel()
req := httptest.NewRequest("GET", "http://testing/test", nil).WithContext(ctx)
req.Header.Set("Authorization", "Bearer "+JWT)
rec := httptest.NewRecorder()
handlerToTest.ServeHTTP(rec, req)
if !syncCalled {
t.Fatal("syncUserJWTGroups was not called")
}
if syncCtxErr != nil {
t.Fatalf("syncUserJWTGroups received a cancelled context (%v); the group-sync write must be detached from request cancellation", syncCtxErr)
}
}
func TestAuthMiddleware_RateLimiting(t *testing.T) {
mockAuth := &auth.MockManager{
ValidateAndParseTokenFunc: mockValidateAndParseToken,
@@ -137,7 +137,7 @@ func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPee
zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)
apiRouter := mux.NewRouter().PathPrefix("/api").Subrouter()
apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil)
apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil, nil)
if err != nil {
t.Fatalf("Failed to create API handler: %v", err)
}
@@ -267,7 +267,7 @@ func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile strin
zoneRecordsManager := recordsManager.NewManager(store, am, permissionsManager)
apiRouter := mux.NewRouter().PathPrefix("/api").Subrouter()
apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil)
apiHandler, err := http2.NewAPIHandler(context.Background(), apiRouter, am, networksManager, resourcesManager, routersManager, groupsManager, geoMock, authManagerMock, metrics, permissionsManager, settingsManager, customZonesManager, zoneRecordsManager, networkMapController, nil, serviceManager, nil, nil, nil, nil, nil, nil, nil)
if err != nil {
t.Fatalf("Failed to create API handler: %v", err)
}
+16
View File
@@ -55,6 +55,7 @@ type DataSource interface {
GetStoreEngine() types.Engine
GetCustomDomainsCounts(ctx context.Context) (total int64, validated int64, err error)
GetProxyMetrics(ctx context.Context) (store.ProxyMetrics, error)
GetAgentNetworkMetrics(ctx context.Context) (store.AgentNetworkMetrics, error)
}
// ConnManager peer connection manager that holds state for current active connections
@@ -413,6 +414,13 @@ func (w *Worker) generateProperties(ctx context.Context) properties {
log.WithContext(ctx).Debugf("collect proxy metrics: %v", err)
}
// Agent-network adoption + usage, aggregated across all accounts in a few
// cheap queries; nil on FileStore.
agentNetworkMetrics, err := w.dataSource.GetAgentNetworkMetrics(ctx)
if err != nil {
log.WithContext(ctx).Debugf("collect agent network metrics: %v", err)
}
minActivePeerVersion, maxActivePeerVersion := getMinMaxVersion(peerActiveVersions)
metricsProperties["uptime"] = uptime
metricsProperties["accounts"] = accounts
@@ -471,6 +479,14 @@ func (w *Worker) generateProperties(ctx context.Context) properties {
metricsProperties["proxies_connected"] = proxyMetrics.ProxiesConnected
metricsProperties["custom_domains"] = customDomains
metricsProperties["custom_domains_validated"] = customDomainsValidated
metricsProperties["agent_network_accounts"] = agentNetworkMetrics.Accounts
metricsProperties["agent_network_providers"] = agentNetworkMetrics.Providers
metricsProperties["agent_network_policies"] = agentNetworkMetrics.Policies
metricsProperties["agent_network_budget_rules"] = agentNetworkMetrics.BudgetRules
metricsProperties["agent_network_log_collection_enabled"] = agentNetworkMetrics.LogCollectionEnabled
metricsProperties["agent_network_input_tokens"] = agentNetworkMetrics.InputTokens
metricsProperties["agent_network_output_tokens"] = agentNetworkMetrics.OutputTokens
metricsProperties["agent_network_cost_usd"] = agentNetworkMetrics.CostUSD
for targetType, count := range servicesTargetType {
metricsProperties["services_target_type_"+string(targetType)] = count
@@ -277,6 +277,21 @@ func (mockDatasource) GetProxyMetrics(_ context.Context) (store.ProxyMetrics, er
}, nil
}
// GetAgentNetworkMetrics returns canned agent-network counts so the
// generateProperties test can assert the adoption/usage signals end-to-end.
func (mockDatasource) GetAgentNetworkMetrics(_ context.Context) (store.AgentNetworkMetrics, error) {
return store.AgentNetworkMetrics{
Accounts: 2,
Providers: 5,
Policies: 3,
BudgetRules: 1,
LogCollectionEnabled: 2,
InputTokens: 1000,
OutputTokens: 500,
CostUSD: 1.25,
}, nil
}
// TestGenerateProperties tests and validate the properties generation by using the mockDatasource for the Worker.generateProperties
func TestGenerateProperties(t *testing.T) {
ds := mockDatasource{}
@@ -39,7 +39,7 @@ type MockAccountManager struct {
GetUserFromUserAuthFunc func(ctx context.Context, userAuth auth.UserAuth) (*types.User, error)
ListUsersFunc func(ctx context.Context, accountID string) ([]*types.User, error)
GetPeersFunc func(ctx context.Context, accountID, userID, nameFilter, ipFilter string) ([]*nbpeer.Peer, error)
MarkPeerConnectedFunc func(ctx context.Context, peerKey string, realIP net.IP, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerConnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error
MarkPeerDisconnectedFunc func(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64) error
SyncAndMarkPeerFunc func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP, syncTime time.Time) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, int64, error)
DeletePeerFunc func(ctx context.Context, accountID, peerKey, userID string) error
@@ -114,7 +114,7 @@ type MockAccountManager struct {
GetIdpManagerFunc func() idp.Manager
UpdateIntegratedValidatorFunc func(ctx context.Context, accountID, userID, validator string, groups []string) error
GroupValidationFunc func(ctx context.Context, accountId string, groups []string) (bool, error)
SyncPeerMetaFunc func(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error
SyncPeerMetaFunc func(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error
FindExistingPostureCheckFunc func(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
GetAccountIDForPeerKeyFunc func(ctx context.Context, peerKey string) (string, error)
GetAccountByIDFunc func(ctx context.Context, accountID string, userID string) (*types.Account, error)
@@ -345,9 +345,9 @@ func (am *MockAccountManager) GetAccountIDByUserID(ctx context.Context, userAuth
}
// MarkPeerConnected mock implementation of MarkPeerConnected from server.AccountManager interface
func (am *MockAccountManager) MarkPeerConnected(ctx context.Context, peerKey string, realIP net.IP, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
func (am *MockAccountManager) MarkPeerConnected(ctx context.Context, peerKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
if am.MarkPeerConnectedFunc != nil {
return am.MarkPeerConnectedFunc(ctx, peerKey, realIP, accountID, sessionStartedAt, nmap)
return am.MarkPeerConnectedFunc(ctx, peerKey, accountID, sessionStartedAt, nmap)
}
return status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
}
@@ -975,9 +975,9 @@ func (am *MockAccountManager) GroupValidation(ctx context.Context, accountId str
}
// SyncPeerMeta mocks SyncPeerMeta of the AccountManager interface
func (am *MockAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error {
func (am *MockAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) error {
if am.SyncPeerMetaFunc != nil {
return am.SyncPeerMetaFunc(ctx, peerPubKey, meta)
return am.SyncPeerMetaFunc(ctx, peerPubKey, meta, realIP)
}
return status.Errorf(codes.Unimplemented, "method SyncPeerMeta is not implemented")
}
+72 -28
View File
@@ -34,7 +34,16 @@ import (
"github.com/netbirdio/netbird/version"
)
const remoteJobsMinVer = "0.64.0"
type peerExpirationReason string
const (
remoteJobsMinVer = "0.64.0"
peerExpirationSessionExpired peerExpirationReason = "session expiration"
peerExpirationInactivity peerExpirationReason = "inactivity timeout"
peerExpirationValidationFailed peerExpirationReason = "failed integration validation"
peerExpirationUserBlocked peerExpirationReason = "blocked owner account"
)
// GetPeers returns peers visible to the user within an account.
// Users with "peers:read" see all peers. Otherwise, users see only their own peers, or none if restricted by account settings.
@@ -74,7 +83,7 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
//
// Disconnects use MarkPeerDisconnected and require the session to match
// exactly; see PeerStatus.SessionStartedAt for the protocol.
func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, realIP net.IP, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubKey string, accountID string, sessionStartedAt int64, nmap *types.NetworkMap) error {
start := time.Now()
defer func() {
am.metrics.AccountManagerMetrics().RecordPeerStatusUpdateDuration(telemetry.PeerStatusConnect, time.Since(start))
@@ -102,10 +111,6 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
}
am.metrics.AccountManagerMetrics().CountPeerStatusUpdate(telemetry.PeerStatusConnect, telemetry.PeerStatusApplied)
if am.geo != nil && realIP != nil {
am.updatePeerLocationIfChanged(ctx, accountID, peer, realIP)
}
if err = am.schedulePeerExpirations(ctx, accountID, peer); err != nil {
return err
}
@@ -192,27 +197,40 @@ func (am *DefaultAccountManager) MarkPeerDisconnected(ctx context.Context, peerP
}
}
if peer.AddedWithSSOLogin() && peer.InactivityExpirationEnabled {
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
if err != nil {
log.WithContext(ctx).Warnf("failed getting account settings to schedule inactivity expiration for peer %s: %v", peer.ID, err)
} else if settings.PeerInactivityExpirationEnabled {
am.checkAndSchedulePeerInactivityExpiration(ctx, accountID)
}
}
return nil
}
// updatePeerLocationIfChanged refreshes the geolocation on a separate
// row update, only when the connection IP actually changed. Geo lookups
// are expensive so we skip same-IP reconnects.
func (am *DefaultAccountManager) updatePeerLocationIfChanged(ctx context.Context, accountID string, peer *nbpeer.Peer, realIP net.IP) {
if peer.Location.ConnectionIP != nil && peer.Location.ConnectionIP.Equal(realIP) {
return
// resolvePeerLocation looks up the geo location for realIP, returning nil when
// there is nothing to apply: geo disabled, no real IP, the IP is unchanged from
// what the peer already has, or the lookup failed. Geo lookups are skipped on
// same-IP reconnects since they are comparatively expensive. The returned value
// is applied by Peer.UpdateMetaIfNew so the change is persisted by its peer save.
func (am *DefaultAccountManager) resolvePeerLocation(ctx context.Context, peer *nbpeer.Peer, realIP net.IP) *nbpeer.Location {
if am.geo == nil || realIP == nil {
return nil
}
location, err := am.geo.Lookup(realIP)
if err != nil {
log.WithContext(ctx).Warnf("failed to get location for peer %s realip: [%s]: %v", peer.ID, realIP.String(), err)
return
return nil
}
peer.Location.ConnectionIP = realIP
peer.Location.CountryCode = location.Country.ISOCode
peer.Location.CityName = location.City.Names.En
peer.Location.GeoNameID = location.City.GeonameID
if err := am.Store.SavePeerLocation(ctx, accountID, peer); err != nil {
log.WithContext(ctx).Warnf("could not store location for peer %s: %s", peer.ID, err)
if peer.Location.ConnectionIP != nil && peer.Location.ConnectionIP.Equal(realIP) && peer.Location.GeoNameID == location.City.GeonameID {
return nil
}
return &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: location.Country.ISOCode,
CityName: location.City.Names.En,
GeoNameID: location.City.GeonameID,
}
}
@@ -721,7 +739,7 @@ func (am *DefaultAccountManager) handleSetupKeyAddedPeer(ctx context.Context, en
func (am *DefaultAccountManager) AddPeer(ctx context.Context, accountID, setupKey, userID string, peer *nbpeer.Peer, temporary bool) (*nbpeer.Peer, *types.Network, []*posture.Checks, bool, error) {
if setupKey == "" && userID == "" && !peer.ProxyMeta.Embedded {
// no auth method provided => reject access
return nil, nil, nil, false, status.Errorf(status.Unauthenticated, "no peer auth method provided, please use a setup key or interactive SSO login")
return nil, nil, nil, false, status.ErrNoAuthMethodProvided
}
upperKey := strings.ToUpper(setupKey)
@@ -980,7 +998,8 @@ func getPeerIPDNSLabel(ip netip.Addr, peerHostName string) (string, error) {
// SyncPeer checks whether peer is eligible for receiving NetworkMap (authenticated) and returns its NetworkMap if eligible
func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSync, accountID string) (*nbpeer.Peer, *types.NetworkMap, []*posture.Checks, int64, error) {
var peer *nbpeer.Peer
var updated, versionChanged, ipv6CapabilityChanged bool
var ipv6CapabilityChanged bool
var metaDiff nbpeer.MetaDiff
var err error
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
@@ -1010,9 +1029,10 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
}
oldHasIPv6Cap := peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
updated, versionChanged = peer.UpdateMetaIfNew(ctx, sync.Meta)
newLocation := am.resolvePeerLocation(ctx, peer, sync.RealIP)
metaDiff = peer.UpdateMetaIfNew(ctx, sync.Meta, newLocation)
ipv6CapabilityChanged = oldHasIPv6Cap != peer.HasCapability(nbpeer.PeerCapabilityIPv6Overlay)
if updated {
if metaDiff.Updated() {
am.metrics.AccountManagerMetrics().CountPeerMetUpdate()
log.WithContext(ctx).Tracef("peer %s metadata updated", peer.ID)
if err = transaction.SavePeer(ctx, accountID, peer); err != nil {
@@ -1040,9 +1060,10 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
return nil, nil, nil, 0, err
}
if isStatusChanged || sync.UpdateAccountPeers || ipv6CapabilityChanged || (updated && (len(resPostureChecks) > 0 || versionChanged)) {
metaDiffAffectsPosture := posture.AffectsPosture(ctx, &metaDiff, resPostureChecks)
if requiresPeerUpdate(ctx, isStatusChanged, sync.UpdateAccountPeers, ipv6CapabilityChanged, metaDiffAffectsPosture, metaDiff.VersionChanged(), metaDiff.HostnameChanged()) {
changedPeerIDs := []string{peer.ID}
affectedPeerIDs := am.syncPeerAffectedPeers(ctx, accountID, peer.ID, nmap, peerNotValid, updated, len(resPostureChecks) > 0)
affectedPeerIDs := am.syncPeerAffectedPeers(ctx, accountID, peer.ID, nmap, peerNotValid, metaDiffAffectsPosture)
if err = am.networkMapController.OnPeersUpdated(ctx, accountID, changedPeerIDs, affectedPeerIDs); err != nil {
return nil, nil, nil, 0, fmt.Errorf("notify network map controller of peer update: %w", err)
}
@@ -1051,6 +1072,29 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
return peer, nmap, resPostureChecks, dnsFwdPort, nil
}
func requiresPeerUpdate(ctx context.Context, isStatusChanged, updateAccountPeers, ipv6CapabilityChanged, metaDiffAffectsPosture, versionChanged, hostname bool) bool {
var reason string
switch {
case isStatusChanged:
reason = "status changed"
case updateAccountPeers:
reason = "update account peers"
case ipv6CapabilityChanged:
reason = "ipv6 capability changed"
case metaDiffAffectsPosture:
reason = "meta diff affects posture"
case versionChanged:
reason = "version changed"
case hostname:
reason = "hostname changed"
default:
return false
}
log.WithContext(ctx).Tracef("peer update required: %s", reason)
return true
}
// syncPeerAffectedPeers resolves the peers affected by a SyncPeer change. The
// peer's own validated network map is bidirectional for policy and routing
// reachability, so when the peer stays valid and no source-posture gate is in
@@ -1059,8 +1103,8 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync types.PeerSy
// metadata change that flips a posture result removes this peer from others'
// maps asymmetrically; that case (and an invalid peer, whose map is empty) falls
// back to the resolver.
func (am *DefaultAccountManager) syncPeerAffectedPeers(ctx context.Context, accountID, peerID string, nmap *types.NetworkMap, peerNotValid, metaUpdated, hasPostureChecks bool) []string {
if peerNotValid || (metaUpdated && hasPostureChecks) {
func (am *DefaultAccountManager) syncPeerAffectedPeers(ctx context.Context, accountID, peerID string, nmap *types.NetworkMap, peerNotValid, metaChangeAffectedPosture bool) []string {
if peerNotValid || metaChangeAffectedPosture {
return am.resolveAffectedPeersForPeerChanges(ctx, am.Store, accountID, []string{peerID})
}
return affectedPeerIDsFromNetworkMap(nmap, peerID)
@@ -1170,7 +1214,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login types.Peer
}
// This is needed to keep in memory for the peer config. Otherwise browser client will end in a retry loop
peer.UpdateMetaIfNew(ctx, login.Meta)
peer.Meta = login.Meta
peerGroupIDs, err = getPeerGroupIDs(ctx, am.Store, accountID, peer.ID)
if err != nil {
+78 -29
View File
@@ -107,6 +107,15 @@ type Location struct {
GeoNameID uint // city level geoname id
}
// equal reports whether two locations match. ConnectionIP is a net.IP slice, so it uses
// IP.Equal, not ==.
func (l Location) equal(other Location) bool {
return l.CountryCode == other.CountryCode &&
l.CityName == other.CityName &&
l.GeoNameID == other.GeoNameID &&
l.ConnectionIP.Equal(other.ConnectionIP)
}
// NetworkAddress is the IP address with network and MAC address of a network interface
type NetworkAddress struct {
NetIP netip.Prefix `gorm:"serializer:json"`
@@ -256,50 +265,88 @@ func (p *Peer) Copy() *Peer {
}
}
// UpdateMetaIfNew updates peer's system metadata if new information is provided
// returns true if meta was updated, false otherwise
func (p *Peer) UpdateMetaIfNew(ctx context.Context, meta PeerSystemMeta) (updated, versionChanged bool) {
// UpdateMetaIfNew updates peer's system metadata and connection geo location if
// new information is provided. newLocation is the geo location resolved from the
// peer's current connection IP, or nil when there is nothing to apply (geo
// disabled, no real IP, or the IP is unchanged); the caller owns the expensive
// lookup and the same-IP guard. It returns a MetaDiff describing what changed;
// diff.Updated() reports whether the peer needs to be persisted.
func (p *Peer) UpdateMetaIfNew(ctx context.Context, meta PeerSystemMeta, newLocation *Location) MetaDiff {
if meta.isEmpty() {
return updated, versionChanged
return MetaDiff{}
}
versionChanged = p.Meta.WtVersion != meta.WtVersion
// Avoid overwriting UIVersion if the update was triggered sole by the CLI client
if meta.UIVersion == "" {
meta.UIVersion = p.Meta.UIVersion
}
oldVersion := p.Meta.WtVersion
effectiveLocation := p.Location
if newLocation != nil {
effectiveLocation = *newLocation
}
diff := metaDiff(p.Meta, meta)
if len(diff) != 0 {
diff := diffMeta(p.Meta, meta, p.Location, effectiveLocation)
if diff.Updated() {
p.Meta = meta
updated = true
}
p.Location = effectiveLocation
if diff.Updated() {
log.WithContext(ctx).Debug(diff.LogSummary())
}
versionInfo := ""
if versionChanged {
versionInfo = fmt.Sprintf("version changed: %s -> %s, ", oldVersion, meta.WtVersion)
}
return diff
}
if len(diff) > 0 || versionChanged {
log.WithContext(ctx).
Debugf("peer meta updated, %s%d field(s) changed: %s", versionInfo, len(diff), strings.Join(diff, ", "))
}
// MetaDiff holds a peer's full before/after state across a sync: both metas and both
// connection locations (the location lives on Peer, not PeerSystemMeta, but posture
// checks read it). Changed lists what moved, for logging and the persistence decision;
// the snapshots let a posture check be replayed against old and new. Everything is derived
// from these fields, so there are no parallel per-field flags to keep in sync.
type MetaDiff struct {
OldMeta PeerSystemMeta
NewMeta PeerSystemMeta
OldLocation Location
NewLocation Location
return updated, versionChanged
Changed []string
}
// Updated reports whether anything changed and the peer must be persisted. diffMeta fills
// Changed in the pass that builds the diff, so this is a length check, not a re-comparison.
// Pointer receiver: MetaDiff embeds two metas, so copying it per call is wasteful.
func (d *MetaDiff) Updated() bool {
return len(d.Changed) != 0
}
// VersionChanged reports whether the WireGuard client version changed (a client upgrade).
func (d *MetaDiff) VersionChanged() bool {
return d.OldMeta.WtVersion != d.NewMeta.WtVersion
}
// HostnameChanged reports whether the peer's hostname changed.
func (d *MetaDiff) HostnameChanged() bool {
return d.OldMeta.Hostname != d.NewMeta.Hostname
}
// LogSummary renders the changed fields as a single human-readable line.
func (d *MetaDiff) LogSummary() string {
return fmt.Sprintf("peer meta updated, %d field(s) changed: %s",
len(d.Changed), strings.Join(d.Changed, ", "))
}
// metaDiff returns a human-readable list of the fields that differ between the
// old and new meta, each formatted as `field: <old> -> <new>`. It is the single
// source of truth for meta comparison: isEqual reports equality as an empty
// diff, so the log line can never disagree with the change decision. Slices are
// cloned before sorting, so callers' meta is not mutated.
func metaDiff(oldMeta, newMeta PeerSystemMeta) []string {
var diff []string
return diffMeta(oldMeta, newMeta, Location{}, Location{}).Changed
}
// diffMeta snapshots a peer's old and new state and records a Changed entry per field that
// moved. It is the single source of truth for the comparison: isEqual is an empty Changed
// list, so the log line and the persistence decision can never disagree.
func diffMeta(oldMeta, newMeta PeerSystemMeta, oldLocation, newLocation Location) MetaDiff {
d := MetaDiff{OldMeta: oldMeta, NewMeta: newMeta, OldLocation: oldLocation, NewLocation: newLocation}
add := func(field string, oldVal, newVal any) {
diff = append(diff, fmt.Sprintf("%s: %v -> %v", field, oldVal, newVal))
d.Changed = append(d.Changed, fmt.Sprintf("%s: %v -> %v", field, oldVal, newVal))
}
if oldMeta.Hostname != newMeta.Hostname {
@@ -353,16 +400,18 @@ func metaDiff(oldMeta, newMeta PeerSystemMeta) []string {
if !capabilitiesEqual(oldMeta.Capabilities, newMeta.Capabilities) {
add("capabilities", oldMeta.Capabilities, newMeta.Capabilities)
}
if !sameMultiset(oldMeta.NetworkAddresses, newMeta.NetworkAddresses) {
add("network_addresses", fmt.Sprintf("%v", oldMeta.NetworkAddresses), fmt.Sprintf("%v", newMeta.NetworkAddresses))
}
if !sameMultiset(oldMeta.Files, newMeta.Files) {
add("files", fmt.Sprintf("%v", oldMeta.Files), fmt.Sprintf("%v", newMeta.Files))
}
return diff
if !oldLocation.equal(newLocation) {
add("connection_ip", oldLocation.ConnectionIP, newLocation.ConnectionIP)
}
return d
}
// sameMultiset reports whether two slices contain the same elements with the
+140 -1
View File
@@ -49,6 +49,7 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/server/activity"
"github.com/netbirdio/netbird/management/server/geolocation"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/management/server/store"
@@ -1047,7 +1048,7 @@ func testUpdateAccountPeers(t *testing.T) {
for _, channel := range peerChannels {
update := <-channel
assert.Nil(t, update.Update.NetbirdConfig)
assert.Nil(t, update.Update.NetbirdConfig, "fan-out updates must not carry a NetbirdConfig; clients treat a config without relay as relay disabled and wipe their relay URLs")
assert.Equal(t, tc.peers, len(update.Update.NetworkMap.RemotePeers))
assert.Equal(t, tc.peers*2, len(update.Update.NetworkMap.FirewallRules))
}
@@ -2893,3 +2894,141 @@ func TestUpdatePeer_DnsLabelUniqueName(t *testing.T) {
require.NoError(t, err, "renaming to unique FQDN should succeed")
assert.Equal(t, "api-server", updated.DNSLabel, "DNS label should be first label of FQDN")
}
// fakeGeo is a configurable geolocation.Geolocation implementation for tests. It
// returns a record built from the configured city geoname id, or an error when set.
type fakeGeo struct {
geoNameID uint
isoCode string
cityName string
err error
}
func (g *fakeGeo) Lookup(net.IP) (*geolocation.Record, error) {
if g.err != nil {
return nil, g.err
}
record := &geolocation.Record{}
record.City.GeonameID = g.geoNameID
record.City.Names.En = g.cityName
record.Country.ISOCode = g.isoCode
return record, nil
}
func (g *fakeGeo) GetAllCountries() ([]geolocation.Country, error) { return nil, nil }
func (g *fakeGeo) GetCitiesByCountry(string) ([]geolocation.City, error) { return nil, nil }
func (g *fakeGeo) Stop() error { return nil }
func TestResolvePeerLocation(t *testing.T) {
realIP := net.ParseIP("203.0.113.10")
tests := []struct {
name string
geo geolocation.Geolocation
peer *nbpeer.Peer
realIP net.IP
want *nbpeer.Location
wantNil bool
}{
{
name: "no geo configured returns nil",
geo: nil,
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
wantNil: true,
},
{
name: "nil real IP returns nil",
geo: &fakeGeo{geoNameID: 100},
peer: &nbpeer.Peer{ID: "p1"},
realIP: nil,
wantNil: true,
},
{
name: "lookup error returns nil",
geo: &fakeGeo{err: fmt.Errorf("lookup boom")},
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
wantNil: true,
},
{
name: "same IP and same geoname returns nil",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: realIP,
GeoNameID: 100,
},
},
realIP: realIP,
wantNil: true,
},
{
name: "same IP but changed geoname returns location",
geo: &fakeGeo{geoNameID: 200, isoCode: "US", cityName: "City B"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: realIP,
GeoNameID: 100,
},
},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City B",
GeoNameID: 200,
},
},
{
name: "different IP returns location",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: net.ParseIP("198.51.100.7"),
GeoNameID: 100,
},
},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City A",
GeoNameID: 100,
},
},
{
name: "no prior location returns location",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City A",
GeoNameID: 100,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
am := &DefaultAccountManager{geo: tt.geo}
got := am.resolvePeerLocation(context.Background(), tt.peer, tt.realIP)
if tt.wantNil {
assert.Nil(t, got, "resolved location should be nil")
return
}
require.NotNil(t, got, "resolved location should not be nil")
assert.True(t, tt.want.ConnectionIP.Equal(got.ConnectionIP), "connection IP should match")
assert.Equal(t, tt.want.CountryCode, got.CountryCode, "country code should match")
assert.Equal(t, tt.want.CityName, got.CityName, "city name should match")
assert.Equal(t, tt.want.GeoNameID, got.GeoNameID, "geoname id should match")
})
}
}
@@ -19,6 +19,7 @@ const (
Pats Module = "pats"
IdentityProviders Module = "identity_providers"
Services Module = "services"
AgentNetwork Module = "agent_network"
)
var All = map[Module]struct{}{
@@ -38,4 +39,5 @@ var All = map[Module]struct{}{
Pats: {},
IdentityProviders: {},
Services: {},
AgentNetwork: {},
}
@@ -0,0 +1,202 @@
package posture
import (
"context"
"net"
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
)
// diffFrom builds a MetaDiff from the old/new snapshots AffectsPosture replays against.
func diffFrom(oldMeta, newMeta nbpeer.PeerSystemMeta, oldLoc, newLoc nbpeer.Location) *nbpeer.MetaDiff {
return &nbpeer.MetaDiff{
OldMeta: oldMeta,
NewMeta: newMeta,
OldLocation: oldLoc,
NewLocation: newLoc,
}
}
func checks(def ChecksDefinition) []*Checks {
return []*Checks{{Checks: def}}
}
func TestAffectsPosture_NilDiff(t *testing.T) {
assert.False(t, AffectsPosture(context.Background(), nil, checks(ChecksDefinition{
NBVersionCheck: &NBVersionCheck{MinVersion: "1.0.0"},
})))
}
func TestAffectsPosture_NBVersion(t *testing.T) {
c := checks(ChecksDefinition{NBVersionCheck: &NBVersionCheck{MinVersion: "1.2.0"}})
tests := []struct {
name string
oldVer, newVer string
want bool
}{
{"both above min, no flip", "1.3.0", "1.4.0", false},
{"both below min, no flip", "1.0.0", "1.1.0", false},
{"crosses up below->above", "1.1.0", "1.3.0", true},
{"crosses down above->below", "1.3.0", "1.1.0", true},
{"unparsable old only -> flip", "garbage", "1.3.0", true},
{"unparsable both -> no flip", "garbage", "junk", false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
diff := diffFrom(
nbpeer.PeerSystemMeta{WtVersion: tt.oldVer},
nbpeer.PeerSystemMeta{WtVersion: tt.newVer},
nbpeer.Location{}, nbpeer.Location{},
)
assert.Equal(t, tt.want, AffectsPosture(context.Background(), diff, c))
})
}
}
func TestAffectsPosture_OSVersion_KernelBumpWithinMin(t *testing.T) {
c := checks(ChecksDefinition{OSVersionCheck: &OSVersionCheck{
Linux: &MinKernelVersionCheck{MinKernelVersion: "5.0.0"},
}})
// Kernel moves but stays above the minimum: verdict stays pass -> not affected.
withinMin := diffFrom(
nbpeer.PeerSystemMeta{GoOS: "linux", KernelVersion: "5.10.0-arch1"},
nbpeer.PeerSystemMeta{GoOS: "linux", KernelVersion: "5.15.0-arch2"},
nbpeer.Location{}, nbpeer.Location{},
)
assert.False(t, AffectsPosture(context.Background(), withinMin, c))
// Kernel drops below the minimum: verdict flips pass -> fail -> affected.
crossesDown := diffFrom(
nbpeer.PeerSystemMeta{GoOS: "linux", KernelVersion: "5.10.0-arch1"},
nbpeer.PeerSystemMeta{GoOS: "linux", KernelVersion: "4.19.0-arch1"},
nbpeer.Location{}, nbpeer.Location{},
)
assert.True(t, AffectsPosture(context.Background(), crossesDown, c))
}
func TestAffectsPosture_OSVersion_GoOSSwitchFlipsVerdict(t *testing.T) {
// Only Linux is constrained. An OS outside the switch (freebsd) passes; switching to a
// failing linux kernel flips the verdict pass -> fail.
c := checks(ChecksDefinition{OSVersionCheck: &OSVersionCheck{
Linux: &MinKernelVersionCheck{MinKernelVersion: "6.0.0"},
}})
diff := diffFrom(
nbpeer.PeerSystemMeta{GoOS: "freebsd"},
nbpeer.PeerSystemMeta{GoOS: "linux", KernelVersion: "4.19.0"},
nbpeer.Location{}, nbpeer.Location{},
)
assert.True(t, AffectsPosture(context.Background(), diff, c))
}
func TestAffectsPosture_Process_GoOSSwitchFlipsVerdict(t *testing.T) {
// Process runs at a linux path. Switching GoOS to windows (no WindowsPath configured)
// flips the verdict.
c := checks(ChecksDefinition{ProcessCheck: &ProcessCheck{
Processes: []Process{{LinuxPath: "/usr/bin/foo"}},
}})
files := []nbpeer.File{{Path: "/usr/bin/foo", ProcessIsRunning: true}}
diff := diffFrom(
nbpeer.PeerSystemMeta{GoOS: "linux", Files: files},
nbpeer.PeerSystemMeta{GoOS: "windows", Files: files},
nbpeer.Location{}, nbpeer.Location{},
)
assert.True(t, AffectsPosture(context.Background(), diff, c))
}
func TestAffectsPosture_Process_UnrelatedFileChange(t *testing.T) {
// A tracked process stays running while an unrelated file is added: the verdict does
// not move, so posture is not affected.
c := checks(ChecksDefinition{ProcessCheck: &ProcessCheck{
Processes: []Process{{LinuxPath: "/usr/bin/foo"}},
}})
diff := diffFrom(
nbpeer.PeerSystemMeta{GoOS: "linux", Files: []nbpeer.File{
{Path: "/usr/bin/foo", ProcessIsRunning: true},
}},
nbpeer.PeerSystemMeta{GoOS: "linux", Files: []nbpeer.File{
{Path: "/usr/bin/foo", ProcessIsRunning: true},
{Path: "/usr/bin/bar", ProcessIsRunning: true},
}},
nbpeer.Location{}, nbpeer.Location{},
)
assert.False(t, AffectsPosture(context.Background(), diff, c))
}
func TestAffectsPosture_GeoLocation(t *testing.T) {
c := checks(ChecksDefinition{GeoLocationCheck: &GeoLocationCheck{
Action: CheckActionAllow,
Locations: []Location{{CountryCode: "DE"}},
}})
// Moving within allowed countries keeps the verdict; moving out flips it.
stayAllowed := diffFrom(
nbpeer.PeerSystemMeta{}, nbpeer.PeerSystemMeta{},
nbpeer.Location{CountryCode: "DE", CityName: "Berlin"},
nbpeer.Location{CountryCode: "DE", CityName: "Munich"},
)
assert.False(t, AffectsPosture(context.Background(), stayAllowed, c))
moveOut := diffFrom(
nbpeer.PeerSystemMeta{}, nbpeer.PeerSystemMeta{},
nbpeer.Location{CountryCode: "DE"},
nbpeer.Location{CountryCode: "FR"},
)
assert.True(t, AffectsPosture(context.Background(), moveOut, c))
}
func TestAffectsPosture_PeerNetworkRange_ConnectionIP(t *testing.T) {
// The check reads the connection IP. Moving out of the allowed range flips the verdict;
// moving within it does not.
_, allowed, _ := net.ParseCIDR("10.0.0.0/8")
c := checks(ChecksDefinition{PeerNetworkRangeCheck: &PeerNetworkRangeCheck{
Action: CheckActionAllow,
Ranges: []netip.Prefix{netip.MustParsePrefix(allowed.String())},
}})
movesOutOfRange := diffFrom(
nbpeer.PeerSystemMeta{}, nbpeer.PeerSystemMeta{},
nbpeer.Location{ConnectionIP: net.ParseIP("10.1.2.3")},
nbpeer.Location{ConnectionIP: net.ParseIP("8.8.8.8")},
)
assert.True(t, AffectsPosture(context.Background(), movesOutOfRange, c))
staysInRange := diffFrom(
nbpeer.PeerSystemMeta{}, nbpeer.PeerSystemMeta{},
nbpeer.Location{ConnectionIP: net.ParseIP("10.1.2.3")},
nbpeer.Location{ConnectionIP: net.ParseIP("10.9.9.9")},
)
assert.False(t, AffectsPosture(context.Background(), staysInRange, c))
}
func TestAffectsPosture_IrrelevantFieldChange(t *testing.T) {
// Hostname changes but no check reads it: not affected even with checks present.
c := checks(ChecksDefinition{
NBVersionCheck: &NBVersionCheck{MinVersion: "1.0.0"},
GeoLocationCheck: &GeoLocationCheck{Action: CheckActionAllow, Locations: []Location{{CountryCode: "DE"}}},
})
diff := diffFrom(
nbpeer.PeerSystemMeta{Hostname: "old", WtVersion: "1.5.0"},
nbpeer.PeerSystemMeta{Hostname: "new", WtVersion: "1.5.0"},
nbpeer.Location{CountryCode: "DE"}, nbpeer.Location{CountryCode: "DE"},
)
assert.False(t, AffectsPosture(context.Background(), diff, c))
}
func TestAffectsPosture_NoChecks(t *testing.T) {
diff := diffFrom(
nbpeer.PeerSystemMeta{WtVersion: "1.0.0"},
nbpeer.PeerSystemMeta{WtVersion: "2.0.0"},
nbpeer.Location{}, nbpeer.Location{},
)
assert.False(t, AffectsPosture(context.Background(), diff, nil))
}
+42
View File
@@ -7,6 +7,8 @@ import (
"regexp"
"github.com/hashicorp/go-version"
log "github.com/sirupsen/logrus"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/shared/management/http/api"
"github.com/netbirdio/netbird/shared/management/status"
@@ -51,6 +53,46 @@ type Checks struct {
Checks ChecksDefinition `gorm:"serializer:json"`
}
// AffectsPosture reports whether the change in diff flips the verdict of any check. It
// replays each check against the peer's old and new state and compares verdicts, so a
// change that moves a field but stays the right side of a threshold (e.g. a kernel bump
// still above the minimum) does not force a re-evaluation. See verdictChanged for how an
// evaluation error counts.
func AffectsPosture(ctx context.Context, diff *nbpeer.MetaDiff, checks []*Checks) bool {
if diff == nil {
return false
}
oldPeer := nbpeer.Peer{Meta: diff.OldMeta, Location: diff.OldLocation}
newPeer := nbpeer.Peer{Meta: diff.NewMeta, Location: diff.NewLocation}
for _, c := range checks {
for _, check := range c.GetChecks() {
if verdictChanged(ctx, check, oldPeer, newPeer) {
return true
}
}
}
return false
}
// verdictChanged replays check against old and new state and reports whether the verdict
// differs. Like callers, it treats an evaluation error as deny: two errors are the same
// verdict (no change), an error on one side only is a flip.
func verdictChanged(ctx context.Context, check Check, oldPeer, newPeer nbpeer.Peer) bool {
oldPass, oldErr := check.Check(ctx, oldPeer)
newPass, newErr := check.Check(ctx, newPeer)
oldVerdict := oldPass && (oldErr == nil)
newVerdict := newPass && (newErr == nil)
changed := oldVerdict != newVerdict
log.WithContext(ctx).Tracef("posture check %s replay: verdict %t -> %t (changed=%t), errs: %v -> %v",
check.Name(), oldVerdict, newVerdict, changed, oldErr, newErr)
return changed
}
// ChecksDefinition contains definition of actual check
type ChecksDefinition struct {
NBVersionCheck *NBVersionCheck `json:",omitempty"`
+1
View File
@@ -489,6 +489,7 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
policy := &types.Policy{
AccountID: account.Id,
Enabled: true,
Rules: []*types.PolicyRule{
{
Enabled: true,
+6
View File
@@ -280,3 +280,9 @@ func (s *FileStore) GetCustomDomainsCounts(_ context.Context) (int64, int64, err
func (s *FileStore) GetProxyMetrics(_ context.Context) (ProxyMetrics, error) {
return ProxyMetrics{}, nil
}
// GetAgentNetworkMetrics is a no-op for FileStore — agent-network state isn't
// persisted in the JSON file format.
func (s *FileStore) GetAgentNetworkMetrics(_ context.Context) (AgentNetworkMetrics, error) {
return AgentNetworkMetrics{}, nil
}
+345 -24
View File
@@ -33,6 +33,7 @@ import (
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/internals/modules/zones"
@@ -137,6 +138,10 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
&networkTypes.Network{}, &routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{}, &types.AccountOnboarding{},
&types.Job{}, &zones.Zone{}, &records.Record{}, &types.UserInviteRecord{}, &rpservice.Service{}, &rpservice.Target{}, &domain.Domain{},
&accesslogs.AccessLogEntry{}, &proxy.Proxy{},
&agentNetworkTypes.Provider{}, &agentNetworkTypes.Policy{}, &agentNetworkTypes.Guardrail{}, &agentNetworkTypes.Settings{},
&agentNetworkTypes.Consumption{}, &agentNetworkTypes.AccountBudgetRule{},
&agentNetworkTypes.AgentNetworkAccessLog{}, &agentNetworkTypes.AgentNetworkAccessLogGroup{},
&agentNetworkTypes.AgentNetworkUsage{}, &agentNetworkTypes.AgentNetworkUsageGroup{},
)
if err != nil {
return nil, fmt.Errorf("auto migratePreAuto: %w", err)
@@ -581,28 +586,6 @@ func (s *SqlStore) MarkPeerDisconnectedIfSameSession(ctx context.Context, accoun
return result.RowsAffected > 0, nil
}
func (s *SqlStore) SavePeerLocation(ctx context.Context, accountID string, peerWithLocation *nbpeer.Peer) error {
// To maintain data integrity, we create a copy of the peer's location to prevent unintended updates to other fields.
var peerCopy nbpeer.Peer
// Since the location field has been migrated to JSON serialization,
// updating the struct ensures the correct data format is inserted into the database.
peerCopy.Location = peerWithLocation.Location
result := s.db.Model(&nbpeer.Peer{}).
Where(accountAndIDQueryCondition, accountID, peerWithLocation.ID).
Updates(peerCopy)
if result.Error != nil {
return status.Errorf(status.Internal, "failed to save peer locations to store: %v", result.Error)
}
if result.RowsAffected == 0 {
return status.Errorf(status.NotFound, peerNotFoundFMT, peerWithLocation.ID)
}
return nil
}
// ApproveAccountPeers marks all peers that currently require approval in the given account as approved.
func (s *SqlStore) ApproveAccountPeers(ctx context.Context, accountID string) (int, error) {
result := s.db.Model(&nbpeer.Peer{}).
@@ -1622,7 +1605,7 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
settings_jwt_groups_enabled, settings_jwt_groups_claim_name, settings_jwt_allow_groups,
settings_routing_peer_dns_resolution_enabled, settings_dns_domain, settings_network_range,
settings_network_range_v6, settings_ipv6_enabled_groups, settings_lazy_connection_enabled,
settings_local_mfa_enabled,
settings_local_mfa_enabled, settings_metrics_push_enabled,
-- Embedded ExtraSettings
settings_extra_peer_approval_enabled, settings_extra_user_approval_required,
settings_extra_integrated_validator, settings_extra_integrated_validator_groups
@@ -1645,6 +1628,7 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
sIPv6EnabledGroups sql.NullString
sLazyConnectionEnabled sql.NullBool
sLocalMFAEnabled sql.NullBool
sMetricsPushEnabled sql.NullBool
sExtraPeerApprovalEnabled sql.NullBool
sExtraUserApprovalRequired sql.NullBool
sExtraIntegratedValidator sql.NullString
@@ -1667,7 +1651,7 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
&sJWTGroupsEnabled, &sJWTGroupsClaimName, &sJWTAllowGroups,
&sRoutingPeerDNSResolutionEnabled, &sDNSDomain, &sNetworkRange,
&sNetworkRangeV6, &sIPv6EnabledGroups, &sLazyConnectionEnabled,
&sLocalMFAEnabled,
&sLocalMFAEnabled, &sMetricsPushEnabled,
&sExtraPeerApprovalEnabled, &sExtraUserApprovalRequired,
&sExtraIntegratedValidator, &sExtraIntegratedValidatorGroups,
)
@@ -1733,6 +1717,9 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
if sLocalMFAEnabled.Valid {
account.Settings.LocalMfaEnabled = sLocalMFAEnabled.Bool
}
if sMetricsPushEnabled.Valid {
account.Settings.MetricsPushEnabled = sMetricsPushEnabled.Bool
}
if sJWTAllowGroups.Valid {
_ = json.Unmarshal([]byte(sJWTAllowGroups.String), &account.Settings.JWTAllowGroups)
}
@@ -5595,6 +5582,340 @@ func (s *SqlStore) CreateAccessLog(ctx context.Context, logEntry *accesslogs.Acc
return nil
}
// CreateAgentNetworkAccessLog persists a flattened agent-network access-log
// entry together with its authorising-group child rows in a single
// transaction.
func (s *SqlStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
err := s.db.Transaction(func(tx *gorm.DB) error {
// Idempotent on the log id / (log_id, group_id) so a proxy resend of the
// same entry can't fail the request.
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(entry).Error; err != nil {
return err
}
if len(groups) > 0 {
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
return err
}
}
return nil
})
if err != nil {
log.WithContext(ctx).WithFields(log.Fields{
"account_id": entry.AccountID,
"service_id": entry.ServiceID,
"model": entry.Model,
}).Errorf("failed to create agent-network access log entry in store: %v", err)
return status.Errorf(status.Internal, "failed to create agent-network access log entry in store")
}
return nil
}
// CreateAgentNetworkUsage persists a stripped agent-network usage record
// together with its authorising-group child rows in a single transaction.
func (s *SqlStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
err := s.db.Transaction(func(tx *gorm.DB) error {
// Idempotent on the usage id / (usage_id, group_id) so a proxy resend of
// the same entry can't fail the request.
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(usage).Error; err != nil {
return err
}
if len(groups) > 0 {
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&groups).Error; err != nil {
return err
}
}
return nil
})
if err != nil {
log.WithContext(ctx).WithFields(log.Fields{
"account_id": usage.AccountID,
"model": usage.Model,
}).Errorf("failed to create agent-network usage record in store: %v", err)
return status.Errorf(status.Internal, "failed to create agent-network usage record in store")
}
return nil
}
// DeleteOldAgentNetworkAccessLogs deletes an account's access-log rows (and
// their authorising-group child rows) older than the cutoff. Usage records are
// untouched — they are the long-term aggregate. Returns the number of log rows
// deleted.
func (s *SqlStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
var deleted int64
err := s.db.Transaction(func(tx *gorm.DB) error {
// Remove group child rows for the soon-to-be-deleted logs first.
if err := tx.Exec(
"DELETE FROM agent_network_access_log_group WHERE account_id = ? AND log_id IN (SELECT id FROM agent_network_access_log WHERE account_id = ? AND timestamp < ?)",
accountID, accountID, olderThan,
).Error; err != nil {
return err
}
res := tx.Where("account_id = ? AND timestamp < ?", accountID, olderThan).
Delete(&agentNetworkTypes.AgentNetworkAccessLog{})
if res.Error != nil {
return res.Error
}
deleted = res.RowsAffected
return nil
})
if err != nil {
log.WithContext(ctx).Errorf("failed to delete old agent-network access logs for account %s: %v", accountID, err)
return 0, status.Errorf(status.Internal, "failed to delete old agent-network access logs")
}
return deleted, nil
}
// GetAgentNetworkUsageRows returns the stripped usage rows for an account that
// match the filter (date / user / group / provider / model). Aggregation into
// time buckets happens in the manager so granularities stay engine-portable.
func (s *SqlStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
var rows []*agentNetworkTypes.AgentNetworkUsage
query := s.applyAgentNetworkUsageFilters(
s.db.Where(accountIDCondition, accountID),
filter,
).Order("timestamp ASC")
if lockStrength != LockingStrengthNone {
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
}
if err := query.Find(&rows).Error; err != nil {
log.WithContext(ctx).Errorf("failed to get agent-network usage rows from store: %v", err)
return nil, status.Errorf(status.Internal, "failed to get agent-network usage rows from store")
}
return rows, nil
}
// applyAgentNetworkUsageFilters applies the shared access-log filter's
// date/user/group/provider/model conditions to a usage-table query. Pagination,
// sort and free-text search are ignored — the overview is an aggregate.
func (s *SqlStore) applyAgentNetworkUsageFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
if filter.UserID != nil {
query = query.Where("user_id = ?", *filter.UserID)
}
if filter.SessionID != nil {
query = query.Where("session_id = ?", *filter.SessionID)
}
if len(filter.ProviderIDs) > 0 {
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
}
if len(filter.Models) > 0 {
query = query.Where("model IN ?", filter.Models)
}
if len(filter.GroupIDs) > 0 {
query = query.Where(
"id IN (SELECT usage_id FROM agent_network_request_usage_group WHERE group_id IN ?)",
filter.GroupIDs,
)
}
if filter.StartDate != nil {
query = query.Where("timestamp >= ?", *filter.StartDate)
}
if filter.EndDate != nil {
query = query.Where("timestamp <= ?", *filter.EndDate)
}
return query
}
// GetAgentNetworkAccessLogs retrieves flattened agent-network access logs for
// an account with server-side pagination, filtering and sorting. Authorising
// group ids are hydrated from the group child table for the returned page.
func (s *SqlStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
var logs []*agentNetworkTypes.AgentNetworkAccessLog
var totalCount int64
countQuery := s.applyAgentNetworkAccessLogFilters(
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
filter,
)
if err := countQuery.Count(&totalCount).Error; err != nil {
log.WithContext(ctx).Errorf("failed to count agent-network access logs: %v", err)
return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access logs")
}
query := s.applyAgentNetworkAccessLogFilters(
s.db.Where(accountIDCondition, accountID),
filter,
).
Order(filter.GetSortColumn() + " " + filter.GetSortOrder()).
Limit(filter.GetLimit()).
Offset(filter.GetOffset())
if lockStrength != LockingStrengthNone {
query = query.Clauses(clause.Locking{Strength: string(lockStrength)})
}
if err := query.Find(&logs).Error; err != nil {
log.WithContext(ctx).Errorf("failed to get agent-network access logs from store: %v", err)
return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access logs from store")
}
if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, logs); err != nil {
return nil, 0, err
}
return logs, totalCount, nil
}
// applyAgentNetworkAccessLogFilters applies the filter conditions to a query.
func (s *SqlStore) applyAgentNetworkAccessLogFilters(query *gorm.DB, filter agentNetworkTypes.AgentNetworkAccessLogFilter) *gorm.DB {
if filter.Search != nil {
p := "%" + *filter.Search + "%"
query = query.Where(
"id LIKE ? OR host LIKE ? OR path LIKE ? OR model LIKE ? OR user_id IN (SELECT id FROM users WHERE email LIKE ? OR name LIKE ?)",
p, p, p, p, p, p,
)
}
if filter.UserID != nil {
query = query.Where("user_id = ?", *filter.UserID)
}
if filter.SessionID != nil {
query = query.Where("session_id = ?", *filter.SessionID)
}
if filter.Decision != nil {
query = query.Where("decision = ?", *filter.Decision)
}
if filter.PathPrefix != nil {
query = query.Where("path LIKE ?", *filter.PathPrefix+"%")
}
if len(filter.ProviderIDs) > 0 {
query = query.Where("resolved_provider_id IN ?", filter.ProviderIDs)
}
if len(filter.Models) > 0 {
query = query.Where("model IN ?", filter.Models)
}
if len(filter.GroupIDs) > 0 {
query = query.Where(
"id IN (SELECT log_id FROM agent_network_access_log_group WHERE group_id IN ?)",
filter.GroupIDs,
)
}
if filter.StartDate != nil {
query = query.Where("timestamp >= ?", *filter.StartDate)
}
if filter.EndDate != nil {
query = query.Where("timestamp <= ?", *filter.EndDate)
}
return query
}
// hydrateAgentNetworkAccessLogGroups loads the authorising group ids for the
// given page of entries and assigns them onto each entry's GroupIDs field.
func (s *SqlStore) hydrateAgentNetworkAccessLogGroups(ctx context.Context, accountID string, logs []*agentNetworkTypes.AgentNetworkAccessLog) error {
if len(logs) == 0 {
return nil
}
ids := make([]string, 0, len(logs))
for _, l := range logs {
ids = append(ids, l.ID)
}
var rows []agentNetworkTypes.AgentNetworkAccessLogGroup
if err := s.db.
Where(accountIDCondition, accountID).
Where("log_id IN ?", ids).
Find(&rows).Error; err != nil {
log.WithContext(ctx).Errorf("failed to hydrate agent-network access log groups: %v", err)
return status.Errorf(status.Internal, "failed to hydrate agent-network access log groups")
}
byLog := make(map[string][]string, len(logs))
for _, r := range rows {
byLog[r.LogID] = append(byLog[r.LogID], r.GroupID)
}
for _, l := range logs {
l.GroupIDs = byLog[l.ID]
}
return nil
}
// agentNetworkSessionKeyExpr is the SQL group key for session-grouped access
// logs: the row's session id, or — when the client sent none — the row id, so
// session-less requests each form their own singleton group. COALESCE/NULLIF
// are standard SQL, so this stays portable across SQLite and Postgres.
const agentNetworkSessionKeyExpr = "COALESCE(NULLIF(session_id, ''), id)"
// GetAgentNetworkAccessLogSessions retrieves agent-network access logs grouped
// by session, with server-side pagination, filtering and sorting at the session
// level. It paginates over the distinct session keys (ordered by the requested
// session-level aggregate), fetches every entry for the page's sessions, and
// folds them into per-session summaries. The returned count is the number of
// matching sessions. Filters apply to the entries, so a session's summary
// reflects only its filter-matching requests.
func (s *SqlStore) GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error) {
// Count distinct sessions via a grouped subquery — portable and avoids
// relying on COUNT(DISTINCT <expr>) quoting quirks.
sessionsSubquery := s.applyAgentNetworkAccessLogFilters(
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
filter,
).
Select(agentNetworkSessionKeyExpr + " AS session_key").
Group(agentNetworkSessionKeyExpr)
var totalCount int64
if err := s.db.Table("(?) AS sessions", sessionsSubquery).Count(&totalCount).Error; err != nil {
log.WithContext(ctx).Errorf("failed to count agent-network access-log sessions: %v", err)
return nil, 0, status.Errorf(status.Internal, "failed to count agent-network access-log sessions")
}
// The page of session keys, ordered by the session-level aggregate. The
// session-key tiebreaker keeps pagination deterministic when the primary
// aggregate ties.
type sessionKeyRow struct {
SessionKey string
}
var keyRows []sessionKeyRow
keyQuery := s.applyAgentNetworkAccessLogFilters(
s.db.Model(&agentNetworkTypes.AgentNetworkAccessLog{}).Where(accountIDCondition, accountID),
filter,
).
Select(agentNetworkSessionKeyExpr + " AS session_key").
Group(agentNetworkSessionKeyExpr).
Order(filter.GetSessionSortExpr() + " " + filter.GetSortOrder()).
Order("session_key ASC").
Limit(filter.GetLimit()).
Offset(filter.GetOffset())
if err := keyQuery.Scan(&keyRows).Error; err != nil {
log.WithContext(ctx).Errorf("failed to list agent-network access-log session keys: %v", err)
return nil, 0, status.Errorf(status.Internal, "failed to list agent-network access-log session keys")
}
if len(keyRows) == 0 {
return nil, totalCount, nil
}
keys := make([]string, 0, len(keyRows))
for _, r := range keyRows {
keys = append(keys, r.SessionKey)
}
// All entries for the page's sessions, contiguous per session and oldest
// first within each — the fold relies on that ordering.
var entries []*agentNetworkTypes.AgentNetworkAccessLog
entriesQuery := s.applyAgentNetworkAccessLogFilters(
s.db.Where(accountIDCondition, accountID),
filter,
).
Where(agentNetworkSessionKeyExpr+" IN ?", keys).
Order(agentNetworkSessionKeyExpr + ", timestamp ASC")
if lockStrength != LockingStrengthNone {
entriesQuery = entriesQuery.Clauses(clause.Locking{Strength: string(lockStrength)})
}
if err := entriesQuery.Find(&entries).Error; err != nil {
log.WithContext(ctx).Errorf("failed to get agent-network access-log session entries: %v", err)
return nil, 0, status.Errorf(status.Internal, "failed to get agent-network access-log session entries")
}
if err := s.hydrateAgentNetworkAccessLogGroups(ctx, accountID, entries); err != nil {
return nil, 0, err
}
return agentNetworkTypes.FoldAccessLogSessions(keys, entries), totalCount, nil
}
// GetAccountAccessLogs retrieves access logs for a given account with pagination and filtering
func (s *SqlStore) GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error) {
var logs []*accesslogs.AccessLogEntry
@@ -0,0 +1,664 @@
package store
import (
"context"
"errors"
"fmt"
"math"
"time"
log "github.com/sirupsen/logrus"
"gorm.io/gorm"
"gorm.io/gorm/clause"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/shared/management/status"
)
// GetAllAgentNetworkProviders returns Agent Network providers across
// every account. Used by the synthesizer to build the global service map.
func (s *SqlStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var providers []*agentNetworkTypes.Provider
if result := tx.Find(&providers); result.Error != nil {
log.WithContext(ctx).Errorf("failed to get all agent network providers from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get all agent network providers from store")
}
for _, provider := range providers {
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
}
}
return providers, nil
}
// GetAgentNetworkMetrics returns aggregated agent-network adoption + usage
// counts for the self-hosted metrics worker. Each value is a single cheap
// aggregate; token/cost are summed over the always-collected per-request usage
// ledger (independent of the log-collection toggle) so they reflect real usage.
func (s *SqlStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetrics, error) {
var m AgentNetworkMetrics
db := s.db.WithContext(ctx)
// Providers + distinct adopting accounts in one round-trip.
provRow := db.Model(&agentNetworkTypes.Provider{}).
Select("COUNT(*) AS providers, COUNT(DISTINCT account_id) AS accounts").Row()
if err := provRow.Scan(&m.Providers, &m.Accounts); err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("scan agent network provider metrics: %w", err)
}
if err := db.Model(&agentNetworkTypes.Policy{}).Count(&m.Policies).Error; err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("count agent network policies: %w", err)
}
if err := db.Model(&agentNetworkTypes.AccountBudgetRule{}).Count(&m.BudgetRules).Error; err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("count agent network budget rules: %w", err)
}
if err := db.Model(&agentNetworkTypes.Settings{}).
Where("enable_log_collection = ?", true).Count(&m.LogCollectionEnabled).Error; err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("count agent network log-collection accounts: %w", err)
}
// COALESCE so an empty ledger scans as 0 instead of NULL.
usageRow := db.Model(&agentNetworkTypes.AgentNetworkUsage{}).
Select("COALESCE(SUM(input_tokens), 0) AS input_tokens, " +
"COALESCE(SUM(output_tokens), 0) AS output_tokens, " +
"COALESCE(SUM(cost_usd), 0) AS cost_usd").Row()
if err := usageRow.Scan(&m.InputTokens, &m.OutputTokens, &m.CostUSD); err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("scan agent network usage metrics: %w", err)
}
return m, nil
}
func (s *SqlStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var providers []*agentNetworkTypes.Provider
result := tx.Find(&providers, accountIDCondition, accountID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get agent network providers from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network providers from store")
}
for _, provider := range providers {
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
}
}
return providers, nil
}
func (s *SqlStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var provider *agentNetworkTypes.Provider
result := tx.Take(&provider, accountAndIDQueryCondition, accountID, providerID)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.NewAgentNetworkProviderNotFoundError(providerID)
}
log.WithContext(ctx).Errorf("failed to get agent network provider from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network provider from store")
}
if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
}
return provider, nil
}
func (s *SqlStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error {
providerCopy := provider.Copy()
if err := providerCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
log.WithContext(ctx).Errorf("failed to encrypt agent network provider %s: %v", provider.ID, err)
return status.Errorf(status.Internal, "failed to encrypt agent network provider")
}
result := s.db.Save(providerCopy)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to save agent network provider to store: %v", result.Error)
return status.Errorf(status.Internal, "failed to save agent network provider to store")
}
return nil
}
func (s *SqlStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error {
result := s.db.Delete(&agentNetworkTypes.Provider{}, accountAndIDQueryCondition, accountID, providerID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete agent network provider from store: %v", result.Error)
return status.Errorf(status.Internal, "failed to delete agent network provider from store")
}
if result.RowsAffected == 0 {
return status.NewAgentNetworkProviderNotFoundError(providerID)
}
return nil
}
func (s *SqlStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var policies []*agentNetworkTypes.Policy
result := tx.Find(&policies, accountIDCondition, accountID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get agent network policies from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network policies from store")
}
return policies, nil
}
func (s *SqlStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var policy *agentNetworkTypes.Policy
result := tx.Take(&policy, accountAndIDQueryCondition, accountID, policyID)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.NewAgentNetworkPolicyNotFoundError(policyID)
}
log.WithContext(ctx).Errorf("failed to get agent network policy from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network policy from store")
}
return policy, nil
}
func (s *SqlStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error {
result := s.db.Save(policy)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to save agent network policy to store: %v", result.Error)
return status.Errorf(status.Internal, "failed to save agent network policy to store")
}
return nil
}
func (s *SqlStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error {
result := s.db.Delete(&agentNetworkTypes.Policy{}, accountAndIDQueryCondition, accountID, policyID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete agent network policy from store: %v", result.Error)
return status.Errorf(status.Internal, "failed to delete agent network policy from store")
}
if result.RowsAffected == 0 {
return status.NewAgentNetworkPolicyNotFoundError(policyID)
}
return nil
}
func (s *SqlStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var guardrails []*agentNetworkTypes.Guardrail
result := tx.Find(&guardrails, accountIDCondition, accountID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get agent network guardrails from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network guardrails from store")
}
return guardrails, nil
}
func (s *SqlStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var guardrail *agentNetworkTypes.Guardrail
result := tx.Take(&guardrail, accountAndIDQueryCondition, accountID, guardrailID)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
}
log.WithContext(ctx).Errorf("failed to get agent network guardrail from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network guardrail from store")
}
return guardrail, nil
}
func (s *SqlStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error {
result := s.db.Save(guardrail)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to save agent network guardrail to store: %v", result.Error)
return status.Errorf(status.Internal, "failed to save agent network guardrail to store")
}
return nil
}
func (s *SqlStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error {
result := s.db.Delete(&agentNetworkTypes.Guardrail{}, accountAndIDQueryCondition, accountID, guardrailID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete agent network guardrail from store: %v", result.Error)
return status.Errorf(status.Internal, "failed to delete agent network guardrail from store")
}
if result.RowsAffected == 0 {
return status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
}
return nil
}
// GetAgentNetworkSettings returns the per-account Agent Network
// settings row. Returns status.NotFound when no row exists.
func (s *SqlStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var settings agentNetworkTypes.Settings
result := tx.Take(&settings, "account_id = ?", accountID)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.Errorf(status.NotFound, "agent network settings for account %s not found", accountID)
}
log.WithContext(ctx).Errorf("failed to get agent network settings from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network settings from store")
}
return &settings, nil
}
// GetAllAgentNetworkSettings returns every account's settings row. Used by the
// access-log retention sweep to learn each account's retention window.
func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var settings []*agentNetworkTypes.Settings
if err := tx.Find(&settings).Error; err != nil {
log.WithContext(ctx).Errorf("failed to list agent network settings: %v", err)
return nil, status.Errorf(status.Internal, "failed to list agent network settings")
}
return settings, nil
}
// GetAgentNetworkSettingsByCluster returns every Settings row pinned to
// the given proxy cluster. Used by the bootstrap label generator to
// build the set of subdomains already taken on a cluster.
func (s *SqlStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var settings []*agentNetworkTypes.Settings
result := tx.Find(&settings, "cluster = ?", cluster)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get agent network settings by cluster from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network settings by cluster from store")
}
return settings, nil
}
// SaveAgentNetworkSettings upserts the per-account Agent Network
// settings row.
func (s *SqlStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
result := s.db.Save(settings)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to save agent network settings to store: %v", result.Error)
return status.Errorf(status.Internal, "failed to save agent network settings to store")
}
return nil
}
// IncrementAgentNetworkConsumption atomically upserts the consumption
// row keyed on (account, dim_kind, dim_id, window_seconds, window_start)
// and adds the supplied deltas. Concurrent calls from multiple proxy
// nodes converge — the database performs the increment server-side via
// ON CONFLICT DO UPDATE so no read-modify-write race exists.
func (s *SqlStore) IncrementAgentNetworkConsumption(
ctx context.Context,
accountID string,
kind agentNetworkTypes.ConsumptionDimension,
dimID string,
windowSeconds int64,
windowStart time.Time,
tokensIn, tokensOut int64,
costUSD float64,
) error {
if accountID == "" || dimID == "" || windowSeconds <= 0 {
return status.Errorf(status.InvalidArgument, "account_id, dim_id and window_seconds must be set")
}
// Deltas are added server-side via ON CONFLICT; a negative or non-finite
// value would silently decrement / poison the persisted totals.
if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
}
row := agentNetworkTypes.Consumption{
AccountID: accountID,
DimensionKind: kind,
DimensionID: dimID,
WindowSeconds: windowSeconds,
WindowStartUTC: windowStart.UTC(),
TokensInput: tokensIn,
TokensOutput: tokensOut,
CostUSD: costUSD,
UpdatedAt: time.Now().UTC(),
}
const tbl = "agent_network_consumption"
err := s.db.Clauses(clause.OnConflict{
Columns: []clause.Column{
{Name: "account_id"},
{Name: "dim_kind"},
{Name: "dim_id"},
{Name: "window_seconds"},
{Name: "window_start_utc"},
},
DoUpdates: clause.Assignments(map[string]any{
"tokens_input": gorm.Expr(tbl+".tokens_input + ?", tokensIn),
"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
"cost_usd": gorm.Expr(tbl+".cost_usd + ?", costUSD),
"updated_at": time.Now().UTC(),
}),
}).Create(&row).Error
if err != nil {
log.WithContext(ctx).Errorf("failed to increment agent network consumption: %v", err)
return status.Errorf(status.Internal, "failed to increment agent network consumption")
}
return nil
}
// GetAgentNetworkConsumption returns the consumption row for the exact
// window key. Returns a zero-valued row (not found mapped to zero) so
// callers can use the result as the headroom basis without nil checks.
func (s *SqlStore) GetAgentNetworkConsumption(
ctx context.Context,
lockStrength LockingStrength,
accountID string,
kind agentNetworkTypes.ConsumptionDimension,
dimID string,
windowSeconds int64,
windowStart time.Time,
) (*agentNetworkTypes.Consumption, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var row agentNetworkTypes.Consumption
result := tx.Take(&row,
"account_id = ? AND dim_kind = ? AND dim_id = ? AND window_seconds = ? AND window_start_utc = ?",
accountID, kind, dimID, windowSeconds, windowStart.UTC())
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return &agentNetworkTypes.Consumption{
AccountID: accountID,
DimensionKind: kind,
DimensionID: dimID,
WindowSeconds: windowSeconds,
WindowStartUTC: windowStart.UTC(),
}, nil
}
log.WithContext(ctx).Errorf("failed to get agent network consumption: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
}
return &row, nil
}
// GetAgentNetworkConsumptionBatch reads many consumption counters for one
// account in a single query, returning a map keyed by the exact
// ConsumptionKey. Missing counters are simply absent from the map (callers
// treat absence as a zero counter). Replaces the per-cap point reads the
// policy selector previously issued one at a time.
func (s *SqlStore) GetAgentNetworkConsumptionBatch(
ctx context.Context,
lockStrength LockingStrength,
accountID string,
keys []agentNetworkTypes.ConsumptionKey,
) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) {
out := make(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, len(keys))
if len(keys) == 0 {
return out, nil
}
// Collect the distinct dim ids, windows and window starts so a single
// query scopes to exactly the current windows in play, then filter the
// returned rows down to the exact requested keys.
wanted := make(map[agentNetworkTypes.ConsumptionKey]struct{}, len(keys))
dimSet := make(map[string]struct{})
winSet := make(map[int64]struct{})
startSet := make(map[time.Time]struct{})
for _, k := range keys {
k.WindowStartUTC = k.WindowStartUTC.UTC()
wanted[k] = struct{}{}
dimSet[k.DimID] = struct{}{}
winSet[k.WindowSeconds] = struct{}{}
startSet[k.WindowStartUTC] = struct{}{}
}
dimIDs := make([]string, 0, len(dimSet))
for d := range dimSet {
dimIDs = append(dimIDs, d)
}
windows := make([]int64, 0, len(winSet))
for w := range winSet {
windows = append(windows, w)
}
starts := make([]time.Time, 0, len(startSet))
for t := range startSet {
starts = append(starts, t)
}
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var rows []*agentNetworkTypes.Consumption
result := tx.Find(&rows,
"account_id = ? AND dim_id IN ? AND window_seconds IN ? AND window_start_utc IN ?",
accountID, dimIDs, windows, starts)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to batch-get agent network consumption: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
}
for _, row := range rows {
k := agentNetworkTypes.ConsumptionKey{
Kind: row.DimensionKind,
DimID: row.DimensionID,
WindowSeconds: row.WindowSeconds,
WindowStartUTC: row.WindowStartUTC.UTC(),
}
if _, ok := wanted[k]; ok {
out[k] = row
}
}
return out, nil
}
// IncrementAgentNetworkConsumptionBatch applies the same usage delta to every
// supplied counter inside a single transaction, so all per-(dimension, window)
// counters a served request books are written atomically in one round-trip
// instead of one upsert per counter. Keys are deduplicated by the caller.
func (s *SqlStore) IncrementAgentNetworkConsumptionBatch(
ctx context.Context,
accountID string,
keys []agentNetworkTypes.ConsumptionKey,
tokensIn, tokensOut int64,
costUSD float64,
) error {
if accountID == "" {
return status.Errorf(status.InvalidArgument, "account_id must be set")
}
if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
}
if len(keys) == 0 {
return nil
}
const tbl = "agent_network_consumption"
err := s.db.Transaction(func(tx *gorm.DB) error {
for _, k := range keys {
if k.DimID == "" || k.WindowSeconds <= 0 {
return status.Errorf(status.InvalidArgument, "dim_id and window_seconds must be set")
}
now := time.Now().UTC()
row := agentNetworkTypes.Consumption{
AccountID: accountID,
DimensionKind: k.Kind,
DimensionID: k.DimID,
WindowSeconds: k.WindowSeconds,
WindowStartUTC: k.WindowStartUTC.UTC(),
TokensInput: tokensIn,
TokensOutput: tokensOut,
CostUSD: costUSD,
UpdatedAt: now,
}
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{
{Name: "account_id"},
{Name: "dim_kind"},
{Name: "dim_id"},
{Name: "window_seconds"},
{Name: "window_start_utc"},
},
DoUpdates: clause.Assignments(map[string]any{
"tokens_input": gorm.Expr(tbl+".tokens_input + ?", tokensIn),
"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
"cost_usd": gorm.Expr(tbl+".cost_usd + ?", costUSD),
"updated_at": now,
}),
}).Create(&row).Error; err != nil {
return err
}
}
return nil
})
if err != nil {
log.WithContext(ctx).Errorf("failed to batch-increment agent network consumption: %v", err)
return status.Errorf(status.Internal, "failed to increment agent network consumption")
}
return nil
}
// ListAgentNetworkConsumption returns every consumption row recorded
// for the account, ordered by window_start descending. Backs the
// dashboard's basic counter view.
func (s *SqlStore) ListAgentNetworkConsumption(
ctx context.Context,
lockStrength LockingStrength,
accountID string,
) ([]*agentNetworkTypes.Consumption, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var rows []*agentNetworkTypes.Consumption
result := tx.
Order("window_start_utc DESC").
Find(&rows, accountIDCondition, accountID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to list agent network consumption: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to list agent network consumption")
}
return rows, nil
}
// GetAccountAgentNetworkBudgetRules returns every account-level budget rule for
// the account.
func (s *SqlStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var rules []*agentNetworkTypes.AccountBudgetRule
result := tx.Find(&rules, accountIDCondition, accountID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get agent network budget rules from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network budget rules from store")
}
return rules, nil
}
// GetAgentNetworkBudgetRuleByID returns a single budget rule scoped to the
// account, or a NotFound error.
func (s *SqlStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) {
tx := s.db
if lockStrength != LockingStrengthNone {
tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
}
var rule *agentNetworkTypes.AccountBudgetRule
result := tx.Take(&rule, accountAndIDQueryCondition, accountID, ruleID)
if result.Error != nil {
if errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
}
log.WithContext(ctx).Errorf("failed to get agent network budget rule from store: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get agent network budget rule from store")
}
return rule, nil
}
// SaveAgentNetworkBudgetRule upserts a budget rule.
func (s *SqlStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error {
result := s.db.Save(rule)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to save agent network budget rule to store: %v", result.Error)
return status.Errorf(status.Internal, "failed to save agent network budget rule to store")
}
return nil
}
// DeleteAgentNetworkBudgetRule removes a budget rule scoped to the account.
func (s *SqlStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error {
result := s.db.Delete(&agentNetworkTypes.AccountBudgetRule{}, accountAndIDQueryCondition, accountID, ruleID)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to delete agent network budget rule from store: %v", result.Error)
return status.Errorf(status.Internal, "failed to delete agent network budget rule from store")
}
if result.RowsAffected == 0 {
return status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
}
return nil
}
@@ -0,0 +1,302 @@
package store
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// TestAgentNetworkUsage_RealStore_RoundTrip drives CreateAgentNetworkUsage and
// CreateAgentNetworkAccessLog through a real sqlite store to prove the schema
// migrates and the inserts succeed for both a populated (allowed) entry and a
// stripped (denied) entry.
func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err, "real sqlite test store must come up")
defer cleanup()
const accountID = "acc-anet-usage-1"
now := time.Now().UTC()
// Populated (allowed) usage row with two authorising groups.
usage := &agentNetworkTypes.AgentNetworkUsage{
ID: "log-allowed-1",
AccountID: accountID,
Timestamp: now,
UserID: "user-alice",
ResolvedProviderID: "prov-openai-1",
Provider: "openai",
Model: "gpt-4o",
SessionID: "sess-round-trip-1",
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
}
usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
{UsageID: usage.ID, GroupID: "grp-oncall", AccountID: accountID},
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "populated usage insert must succeed")
// Stripped (denied / 403) usage row: no provider/model/tokens, no groups.
denied := &agentNetworkTypes.AgentNetworkUsage{
ID: "log-denied-1",
AccountID: accountID,
Timestamp: now,
UserID: "user-bob",
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, denied, nil), "stripped usage insert must succeed")
// Idempotency: re-inserting the same id must not error.
require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "duplicate usage insert must be idempotent")
// Access-log row + group children.
entry := &agentNetworkTypes.AgentNetworkAccessLog{
ID: "log-allowed-1",
AccountID: accountID,
ServiceID: "agent-net-svc-1",
Timestamp: now,
UserID: "user-alice",
StatusCode: 200,
Provider: "openai",
Model: "gpt-4o",
SessionID: "sess-round-trip-1",
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
}
entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
{LogID: entry.ID, GroupID: "grp-oncall", AccountID: accountID},
}
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, entry, entryGroups), "access-log insert must succeed")
// Read back through the filtered list + verify group hydration.
logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
require.NoError(t, err, "list must succeed")
assert.Equal(t, int64(1), total, "one access-log row expected")
require.Len(t, logs, 1)
assert.ElementsMatch(t, []string{"grp-eng", "grp-oncall"}, logs[0].GroupIDs, "group ids must hydrate")
assert.Equal(t, "sess-round-trip-1", logs[0].SessionID, "session id must persist and read back on the access-log row")
// Session filter narrows the access-log listing to one conversation.
sessionID := "sess-round-trip-1"
sessLogs, sessTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &sessionID})
require.NoError(t, err)
assert.Equal(t, int64(1), sessTotal, "session filter must match the one row with that session id")
require.Len(t, sessLogs, 1)
assert.Equal(t, entry.ID, sessLogs[0].ID, "session filter must return the matching log row")
bogus := "no-such-session"
_, emptyTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &bogus})
require.NoError(t, err)
assert.Equal(t, int64(0), emptyTotal, "unknown session id must match nothing")
// Session filter also narrows the always-on usage rows.
sessUsage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{SessionID: &sessionID})
require.NoError(t, err)
require.Len(t, sessUsage, 1, "session filter must narrow usage rows to the matching session")
assert.Equal(t, "sess-round-trip-1", sessUsage[0].SessionID, "usage row must carry the session id")
}
// TestAgentNetworkUsageOverview_DailyAggregation drives GetAgentNetworkUsageRows
// + AggregateUsageByGranularity end-to-end against a real sqlite store, with
// two rows on the same day and one on another, plus a model filter.
func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err, "real sqlite test store must come up")
defer cleanup()
const accountID = "acc-anet-overview-1"
day1 := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
day1b := time.Date(2026, 5, 5, 22, 0, 0, 0, time.UTC)
day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
mk := func(id string, ts time.Time, model string, in, out int64, cost float64) *agentNetworkTypes.AgentNetworkUsage {
return &agentNetworkTypes.AgentNetworkUsage{
ID: id, AccountID: accountID, Timestamp: ts, Model: model,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, CostUSD: cost,
}
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u2", day1b, "gpt-4o", 200, 80, 0.20), nil))
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u3", day2, "claude-3", 10, 5, 0.01), nil))
rows, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
require.NoError(t, err)
require.Len(t, rows, 3, "all three usage rows expected")
buckets := agentNetworkTypes.AggregateUsageByGranularity(rows, agentNetworkTypes.UsageGranularityDay)
require.Len(t, buckets, 2, "two distinct days expected")
assert.Equal(t, "2026-05-05", buckets[0].PeriodStart, "oldest-first ordering")
assert.Equal(t, int64(300), buckets[0].InputTokens, "same-day input tokens summed")
assert.Equal(t, int64(130), buckets[0].OutputTokens)
assert.InDelta(t, 0.30, buckets[0].CostUSD, 1e-9, "same-day cost summed")
assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
assert.Equal(t, int64(15), buckets[1].TotalTokens)
// Model filter narrows to a single day.
model := "claude-3"
filtered, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Models: []string{model}})
require.NoError(t, err)
require.Len(t, filtered, 1, "model filter must narrow rows")
assert.Equal(t, "u3", filtered[0].ID)
}
// TestAgentNetworkAccessLogSessions_RealStore drives GetAgentNetworkAccessLogSessions
// against a real sqlite store: session grouping + aggregation, recency ordering,
// singleton groups for session-less requests, session pagination, the model
// filter narrowing sessions, and aggregate sorting.
func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err, "real sqlite test store must come up")
defer cleanup()
const accountID = "acc-anet-sessions-1"
base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
at := func(h int) time.Time { return base.Add(time.Duration(h) * time.Hour) }
mk := func(id, session, user, provider, model, decision string, ts time.Time, cost float64) *agentNetworkTypes.AgentNetworkAccessLog {
return &agentNetworkTypes.AgentNetworkAccessLog{
ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts,
UserID: user, StatusCode: 200, Provider: provider, Model: model,
SessionID: session, Decision: decision,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, CostUSD: cost,
}
}
// Two-request session s1 (alice), a one-request denied session s2 (bob), and
// two session-less requests (empty session id) that must each form their own
// singleton group.
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s1-a", "s1", "alice", "openai", "gpt-4o", "allow", at(1), 0.10),
[]agentNetworkTypes.AgentNetworkAccessLogGroup{{LogID: "s1-a", GroupID: "grp-eng", AccountID: accountID}}))
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s1-b", "s1", "alice", "openai", "gpt-4o", "allow", at(2), 0.20),
[]agentNetworkTypes.AgentNetworkAccessLogGroup{{LogID: "s1-b", GroupID: "grp-oncall", AccountID: accountID}}))
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s2-a", "s2", "bob", "anthropic", "claude-3", "deny", at(3), 0.05), nil))
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("se-old", "", "carol", "openai", "o1", "allow", at(0), 0.01), nil))
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("se-new", "", "dave", "mistral", "mistral-large", "allow", at(4), 0.02), nil))
// Default sort: last activity (MAX timestamp) descending.
sessions, total, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
require.NoError(t, err)
assert.Equal(t, int64(4), total, "four sessions: s1, s2, and two singletons")
require.Len(t, sessions, 4)
// se-new(t4) > s2(t3) > s1(t2) > se-old(t0)
assert.Equal(t, "", sessions[0].SessionID, "newest is a session-less singleton")
assert.Equal(t, "se-new", sessions[0].Entries[0].ID)
assert.Equal(t, "s2", sessions[1].SessionID)
assert.Equal(t, "s1", sessions[2].SessionID)
assert.Equal(t, "se-old", sessions[3].Entries[0].ID)
// s1 aggregation.
s1 := sessions[2]
assert.Equal(t, 2, s1.RequestCount, "s1 has two requests")
assert.Equal(t, int64(300), s1.TotalTokens, "tokens summed across the session")
assert.InDelta(t, 0.30, s1.CostUSD, 1e-9, "cost summed across the session")
assert.Equal(t, "alice", s1.UserID)
assert.Equal(t, "allow", s1.Decision)
// SQLite hands times back in time.Local; normalise to UTC so the instant is
// compared, not the (differing) *Location pointer.
assert.Equal(t, at(1), s1.StartedAt.UTC(), "started = earliest entry")
assert.Equal(t, at(2), s1.EndedAt.UTC(), "ended = latest entry")
assert.ElementsMatch(t, []string{"openai"}, s1.Providers)
assert.ElementsMatch(t, []string{"gpt-4o"}, s1.Models)
assert.ElementsMatch(t, []string{"grp-eng", "grp-oncall"}, s1.GroupIDs, "union of the entries' authorising groups")
// Denied session rolls up to deny.
assert.Equal(t, "deny", sessions[1].Decision, "any denied request makes the session deny")
// Pagination over sessions: 2 per page.
page1, total, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 2})
require.NoError(t, err)
assert.Equal(t, int64(4), total, "total still counts all sessions")
require.Len(t, page1, 2)
assert.Equal(t, "se-new", page1[0].Entries[0].ID)
assert.Equal(t, "s2", page1[1].SessionID)
page2, _, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 2, PageSize: 2})
require.NoError(t, err)
require.Len(t, page2, 2)
assert.Equal(t, "s1", page2[0].SessionID)
assert.Equal(t, "se-old", page2[1].Entries[0].ID)
// Model filter narrows to the session(s) with matching entries.
model := "claude-3"
filtered, fTotal, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, Models: []string{model}})
require.NoError(t, err)
assert.Equal(t, int64(1), fTotal, "only s2 has a claude-3 request")
require.Len(t, filtered, 1)
assert.Equal(t, "s2", filtered[0].SessionID)
// Sort by total session cost, descending: s1 (0.30) leads despite not being
// the most recent.
byCost, _, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SortBy: "cost_usd", SortOrder: "desc"})
require.NoError(t, err)
require.Len(t, byCost, 4)
assert.Equal(t, "s1", byCost[0].SessionID, "highest-cost session sorts first")
}
// TestDeleteOldAgentNetworkAccessLogs verifies the retention sweep removes only
// access-log rows (and their group children) older than the cutoff, leaving
// recent rows — and never touching usage records.
func TestDeleteOldAgentNetworkAccessLogs(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err, "real sqlite test store must come up")
defer cleanup()
const accountID = "acc-anet-retention-1"
old := time.Now().UTC().AddDate(0, 0, -40)
recent := time.Now().UTC().AddDate(0, 0, -1)
mkLog := func(id string, ts time.Time) (*agentNetworkTypes.AgentNetworkAccessLog, []agentNetworkTypes.AgentNetworkAccessLogGroup) {
return &agentNetworkTypes.AgentNetworkAccessLog{
ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts, StatusCode: 200, Model: "gpt-4o",
}, []agentNetworkTypes.AgentNetworkAccessLogGroup{
{LogID: id, GroupID: "grp-eng", AccountID: accountID},
}
}
oldEntry, oldGroups := mkLog("old-1", old)
recentEntry, recentGroups := mkLog("recent-1", recent)
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, oldEntry, oldGroups))
require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, recentEntry, recentGroups))
// A usage row for the old request must survive the access-log sweep.
require.NoError(t, s.CreateAgentNetworkUsage(ctx, &agentNetworkTypes.AgentNetworkUsage{
ID: "old-1", AccountID: accountID, Timestamp: old, Model: "gpt-4o", InputTokens: 10, TotalTokens: 10,
}, nil))
cutoff := time.Now().UTC().AddDate(0, 0, -30)
deleted, err := s.DeleteOldAgentNetworkAccessLogs(ctx, accountID, cutoff)
require.NoError(t, err)
assert.Equal(t, int64(1), deleted, "only the 40-day-old log is deleted")
logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
require.NoError(t, err)
assert.Equal(t, int64(1), total, "the recent log remains")
require.Len(t, logs, 1)
assert.Equal(t, "recent-1", logs[0].ID)
// Usage is untouched by the access-log retention sweep.
usage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
require.NoError(t, err)
require.Len(t, usage, 1, "usage record for the deleted log must survive")
}
@@ -0,0 +1,112 @@
package store
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// TestAgentNetworkBudgetRule_RealStore_RoundTrip is the GC-0 no-mock guard: it
// drives the budget-rule CRUD through a real sqlite store and asserts the full
// object — targets and the reused PolicyLimits cap shape — survives the
// save → gorm/JSON serialize → reload round-trip, then that delete removes it
// and a second delete reports NotFound.
func TestAgentNetworkBudgetRule_RealStore_RoundTrip(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err, "real sqlite test store must come up")
defer cleanup()
const accountID = "acc-budgetrule-1"
rule := agentNetworkTypes.NewAccountBudgetRule(accountID)
rule.Name = "eng-monthly"
rule.TargetGroups = []string{"grp-eng", "grp-oncall"}
rule.TargetUsers = []string{"user-alice"}
rule.Limits = agentNetworkTypes.PolicyLimits{
TokenLimit: agentNetworkTypes.PolicyTokenLimit{
Enabled: true, GroupCap: 100_000, UserCap: 10_000, WindowSeconds: 2_592_000,
},
BudgetLimit: agentNetworkTypes.PolicyBudgetLimit{
Enabled: true, GroupCapUsd: 500, UserCapUsd: 50, WindowSeconds: 2_592_000,
},
}
require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, rule), "save must succeed")
got, err := s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
require.NoError(t, err, "get by id must succeed after save")
assert.Equal(t, rule.Name, got.Name, "name must round-trip")
assert.Equal(t, []string{"grp-eng", "grp-oncall"}, got.TargetGroups, "target groups must round-trip")
assert.Equal(t, []string{"user-alice"}, got.TargetUsers, "target users must round-trip")
assert.Equal(t, rule.Limits, got.Limits, "the reused PolicyLimits cap shape must round-trip intact")
assert.True(t, got.Enabled, "enabled must round-trip")
list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, accountID)
require.NoError(t, err, "list must succeed")
require.Len(t, list, 1, "exactly the one saved rule must be listed")
assert.Equal(t, rule.ID, list[0].ID, "listed rule id must match")
require.NoError(t, s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID), "delete must succeed")
_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
assert.Error(t, err, "get after delete must report not found")
err = s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID)
assert.Error(t, err, "deleting an absent rule must report not found")
}
// TestAgentNetworkBudgetRule_RealStore_ScopedByAccount pins that rules are
// account-scoped: a rule under one account is invisible to another.
func TestAgentNetworkBudgetRule_RealStore_ScopedByAccount(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err)
defer cleanup()
ruleA := agentNetworkTypes.NewAccountBudgetRule("acc-A")
require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, ruleA))
list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, "acc-B")
require.NoError(t, err)
assert.Empty(t, list, "account B must not see account A's budget rule")
_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, "acc-B", ruleA.ID)
assert.Error(t, err, "cross-account get by id must not resolve")
}
// TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip pins the GC-0
// additive settings columns: the three collection toggles default off on a
// fresh row and survive a save/reload at their set values.
func TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip(t *testing.T) {
ctx := context.Background()
s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
require.NoError(t, err)
defer cleanup()
const accountID = "acc-settings-toggles"
require.NoError(t, s.SaveAgentNetworkSettings(ctx, &agentNetworkTypes.Settings{
AccountID: accountID,
Cluster: "eu.proxy.netbird.io",
Subdomain: "violet",
}))
got, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
require.NoError(t, err)
assert.False(t, got.EnableLogCollection, "log collection must default off")
assert.False(t, got.EnablePromptCollection, "prompt collection must default off")
assert.False(t, got.RedactPii, "redact pii must default off")
got.EnableLogCollection = true
got.EnablePromptCollection = true
got.RedactPii = true
require.NoError(t, s.SaveAgentNetworkSettings(ctx, got))
reloaded, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
require.NoError(t, err)
assert.True(t, reloaded.EnableLogCollection, "log collection must round-trip on")
assert.True(t, reloaded.EnablePromptCollection, "prompt collection must round-trip on")
assert.True(t, reloaded.RedactPii, "redact pii must round-trip on")
}
-50
View File
@@ -618,56 +618,6 @@ func TestSqlStore_SavePeerStatus(t *testing.T) {
assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
}
func TestSqlStore_SavePeerLocation(t *testing.T) {
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
t.Cleanup(cleanUp)
assert.NoError(t, err)
account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
require.NoError(t, err)
peer := &nbpeer.Peer{
AccountID: account.Id,
ID: "testpeer",
Location: nbpeer.Location{
ConnectionIP: net.ParseIP("0.0.0.0"),
CountryCode: "YY",
CityName: "City",
GeoNameID: 1,
},
CreatedAt: time.Now().UTC(),
Meta: nbpeer.PeerSystemMeta{},
}
// error is expected as peer is not in store yet
err = store.SavePeerLocation(context.Background(), account.Id, peer)
assert.Error(t, err)
account.Peers[peer.ID] = peer
err = store.SaveAccount(context.Background(), account)
require.NoError(t, err)
peer.Location.ConnectionIP = net.ParseIP("35.1.1.1")
peer.Location.CountryCode = "DE"
peer.Location.CityName = "Berlin"
peer.Location.GeoNameID = 2950159
err = store.SavePeerLocation(context.Background(), account.Id, account.Peers[peer.ID])
assert.NoError(t, err)
account, err = store.GetAccount(context.Background(), account.Id)
require.NoError(t, err)
actual := account.Peers[peer.ID].Location
assert.Equal(t, peer.Location, actual)
peer.ID = "non-existing-peer"
err = store.SavePeerLocation(context.Background(), account.Id, peer)
assert.Error(t, err)
parsedErr, ok := status.FromError(err)
require.True(t, ok)
require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
}
func Test_TestGetAccountByPrivateDomain(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("The SQLite store is not properly supported by Windows yet")
+66 -1
View File
@@ -37,6 +37,7 @@ import (
"github.com/netbirdio/netbird/util"
"github.com/netbirdio/netbird/util/crypt"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/migration"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
@@ -185,7 +186,6 @@ type Store interface {
// recorded by the database. Returns true when the update happened,
// false when a newer session has taken over.
MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
SavePeerLocation(ctx context.Context, accountID string, peer *nbpeer.Peer) error
ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
DeletePeer(ctx context.Context, accountID string, peerID string) error
@@ -301,6 +301,12 @@ type Store interface {
CreateAccessLog(ctx context.Context, log *accesslogs.AccessLogEntry) error
GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error)
DeleteOldAccessLogs(ctx context.Context, olderThan time.Time) (int64, error)
CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error
CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error
GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error)
GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error)
GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error)
DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error)
GetServiceTargetByTargetID(ctx context.Context, lockStrength LockingStrength, accountID string, targetID string) (*rpservice.Target, error)
GetTargetsByServiceID(ctx context.Context, lockStrength LockingStrength, accountID string, serviceID string) ([]*rpservice.Target, error)
DeleteTarget(ctx context.Context, accountID string, serviceID string, targetID uint) error
@@ -331,7 +337,40 @@ type Store interface {
// return a zero-valued struct.
GetProxyMetrics(ctx context.Context) (ProxyMetrics, error)
// GetAgentNetworkMetrics returns aggregated agent-network adoption + usage
// counts for the self-hosted metrics worker. Self-hosted only — file-based
// stores return a zero-valued struct.
GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetrics, error)
GetRoutingPeerNetworks(ctx context.Context, accountID, peerID string) ([]string, error)
// Agent Network persistence (providers, policies, guardrails, settings).
GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error)
GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error)
GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error)
SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error
DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error
GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error)
GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error)
SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error
DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error
GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error)
GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error)
SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error
DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error
GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error)
GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error)
GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error)
SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error
IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error
IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error
GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error)
GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error)
ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error)
GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error)
GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error)
SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error
DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error
}
// ProxyMetrics aggregates self-hosted proxy + cluster usage signals
@@ -358,6 +397,32 @@ type ProxyMetrics struct {
ProxiesConnected int64
}
// AgentNetworkMetrics aggregates self-hosted agent-network adoption + usage
// signals surfaced to the telemetry payload. Each field is best-effort: when a
// store cannot answer (e.g. FileStore) all fields are zero.
type AgentNetworkMetrics struct {
// Accounts is the number of distinct accounts with at least one provider
// configured (agent-network adoption).
Accounts int64
// Providers is the total number of configured providers across all accounts.
Providers int64
// Policies is the total number of agent-network policies across all accounts.
Policies int64
// BudgetRules is the total number of account-level budget rules ("budget
// limits") across all accounts.
BudgetRules int64
// LogCollectionEnabled is the number of accounts that have agent-network
// log collection turned on.
LogCollectionEnabled int64
// InputTokens / OutputTokens / CostUSD are summed over the always-collected
// per-request usage ledger (agent_network_request_usage), independent of the
// log-collection toggle. They reflect total metered LLM usage served through
// agent networks.
InputTokens int64
OutputTokens int64
CostUSD float64
}
const (
postgresDsnEnv = "NB_STORE_ENGINE_POSTGRES_DSN"
postgresDsnEnvLegacy = "NETBIRD_STORE_ENGINE_POSTGRES_DSN"
-14
View File
@@ -2998,20 +2998,6 @@ func (mr *MockStoreMockRecorder) SavePeer(ctx, accountID, peer interface{}) *gom
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SavePeer", reflect.TypeOf((*MockStore)(nil).SavePeer), ctx, accountID, peer)
}
// SavePeerLocation mocks base method.
func (m *MockStore) SavePeerLocation(ctx context.Context, accountID string, peer *peer.Peer) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SavePeerLocation", ctx, accountID, peer)
ret0, _ := ret[0].(error)
return ret0
}
// SavePeerLocation indicates an expected call of SavePeerLocation.
func (mr *MockStoreMockRecorder) SavePeerLocation(ctx, accountID, peer interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SavePeerLocation", reflect.TypeOf((*MockStore)(nil).SavePeerLocation), ctx, accountID, peer)
}
// SavePeerStatus mocks base method.
func (m *MockStore) SavePeerStatus(ctx context.Context, accountID, peerID string, status peer.PeerStatus) error {
m.ctrl.T.Helper()
@@ -0,0 +1,495 @@
package store
import (
context "context"
reflect "reflect"
time "time"
gomock "github.com/golang/mock/gomock"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// GetAllAgentNetworkProviders mocks base method.
func (m *MockStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAllAgentNetworkProviders", ctx, lockStrength)
ret0, _ := ret[0].([]*agentNetworkTypes.Provider)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAllAgentNetworkProviders indicates an expected call of GetAllAgentNetworkProviders.
func (mr *MockStoreMockRecorder) GetAllAgentNetworkProviders(ctx, lockStrength interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkProviders), ctx, lockStrength)
}
// GetAgentNetworkMetrics mocks base method.
func (m *MockStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetrics, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkMetrics", ctx)
ret0, _ := ret[0].(AgentNetworkMetrics)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkMetrics indicates an expected call of GetAgentNetworkMetrics.
func (mr *MockStoreMockRecorder) GetAgentNetworkMetrics(ctx interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkMetrics", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkMetrics), ctx)
}
// GetAccountAgentNetworkProviders mocks base method.
func (m *MockStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAccountAgentNetworkProviders", ctx, lockStrength, accountID)
ret0, _ := ret[0].([]*agentNetworkTypes.Provider)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAccountAgentNetworkProviders indicates an expected call of GetAccountAgentNetworkProviders.
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkProviders(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkProviders", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkProviders), ctx, lockStrength, accountID)
}
// GetAgentNetworkProviderByID mocks base method.
func (m *MockStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkProviderByID", ctx, lockStrength, accountID, providerID)
ret0, _ := ret[0].(*agentNetworkTypes.Provider)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkProviderByID indicates an expected call of GetAgentNetworkProviderByID.
func (mr *MockStoreMockRecorder) GetAgentNetworkProviderByID(ctx, lockStrength, accountID, providerID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkProviderByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkProviderByID), ctx, lockStrength, accountID, providerID)
}
// SaveAgentNetworkProvider mocks base method.
func (m *MockStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SaveAgentNetworkProvider", ctx, provider)
ret0, _ := ret[0].(error)
return ret0
}
// SaveAgentNetworkProvider indicates an expected call of SaveAgentNetworkProvider.
func (mr *MockStoreMockRecorder) SaveAgentNetworkProvider(ctx, provider interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkProvider), ctx, provider)
}
// DeleteAgentNetworkProvider mocks base method.
func (m *MockStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteAgentNetworkProvider", ctx, accountID, providerID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteAgentNetworkProvider indicates an expected call of DeleteAgentNetworkProvider.
func (mr *MockStoreMockRecorder) DeleteAgentNetworkProvider(ctx, accountID, providerID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkProvider", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkProvider), ctx, accountID, providerID)
}
// GetAccountAgentNetworkPolicies mocks base method.
func (m *MockStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAccountAgentNetworkPolicies", ctx, lockStrength, accountID)
ret0, _ := ret[0].([]*agentNetworkTypes.Policy)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAccountAgentNetworkPolicies indicates an expected call of GetAccountAgentNetworkPolicies.
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkPolicies(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkPolicies", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkPolicies), ctx, lockStrength, accountID)
}
// GetAgentNetworkPolicyByID mocks base method.
func (m *MockStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkPolicyByID", ctx, lockStrength, accountID, policyID)
ret0, _ := ret[0].(*agentNetworkTypes.Policy)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkPolicyByID indicates an expected call of GetAgentNetworkPolicyByID.
func (mr *MockStoreMockRecorder) GetAgentNetworkPolicyByID(ctx, lockStrength, accountID, policyID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkPolicyByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkPolicyByID), ctx, lockStrength, accountID, policyID)
}
// SaveAgentNetworkPolicy mocks base method.
func (m *MockStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SaveAgentNetworkPolicy", ctx, policy)
ret0, _ := ret[0].(error)
return ret0
}
// SaveAgentNetworkPolicy indicates an expected call of SaveAgentNetworkPolicy.
func (mr *MockStoreMockRecorder) SaveAgentNetworkPolicy(ctx, policy interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkPolicy), ctx, policy)
}
// DeleteAgentNetworkPolicy mocks base method.
func (m *MockStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteAgentNetworkPolicy", ctx, accountID, policyID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteAgentNetworkPolicy indicates an expected call of DeleteAgentNetworkPolicy.
func (mr *MockStoreMockRecorder) DeleteAgentNetworkPolicy(ctx, accountID, policyID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkPolicy", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkPolicy), ctx, accountID, policyID)
}
// GetAccountAgentNetworkGuardrails mocks base method.
func (m *MockStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAccountAgentNetworkGuardrails", ctx, lockStrength, accountID)
ret0, _ := ret[0].([]*agentNetworkTypes.Guardrail)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAccountAgentNetworkGuardrails indicates an expected call of GetAccountAgentNetworkGuardrails.
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkGuardrails(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkGuardrails", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkGuardrails), ctx, lockStrength, accountID)
}
// GetAgentNetworkGuardrailByID mocks base method.
func (m *MockStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkGuardrailByID", ctx, lockStrength, accountID, guardrailID)
ret0, _ := ret[0].(*agentNetworkTypes.Guardrail)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkGuardrailByID indicates an expected call of GetAgentNetworkGuardrailByID.
func (mr *MockStoreMockRecorder) GetAgentNetworkGuardrailByID(ctx, lockStrength, accountID, guardrailID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkGuardrailByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkGuardrailByID), ctx, lockStrength, accountID, guardrailID)
}
// SaveAgentNetworkGuardrail mocks base method.
func (m *MockStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SaveAgentNetworkGuardrail", ctx, guardrail)
ret0, _ := ret[0].(error)
return ret0
}
// SaveAgentNetworkGuardrail indicates an expected call of SaveAgentNetworkGuardrail.
func (mr *MockStoreMockRecorder) SaveAgentNetworkGuardrail(ctx, guardrail interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkGuardrail), ctx, guardrail)
}
// DeleteAgentNetworkGuardrail mocks base method.
func (m *MockStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteAgentNetworkGuardrail", ctx, accountID, guardrailID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteAgentNetworkGuardrail indicates an expected call of DeleteAgentNetworkGuardrail.
func (mr *MockStoreMockRecorder) DeleteAgentNetworkGuardrail(ctx, accountID, guardrailID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkGuardrail", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkGuardrail), ctx, accountID, guardrailID)
}
// GetAgentNetworkSettings mocks base method.
func (m *MockStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkSettings", ctx, lockStrength, accountID)
ret0, _ := ret[0].(*agentNetworkTypes.Settings)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkSettings indicates an expected call of GetAgentNetworkSettings.
func (mr *MockStoreMockRecorder) GetAgentNetworkSettings(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettings), ctx, lockStrength, accountID)
}
// GetAgentNetworkSettingsByCluster mocks base method.
func (m *MockStore) GetAgentNetworkSettingsByCluster(ctx context.Context, lockStrength LockingStrength, cluster string) ([]*agentNetworkTypes.Settings, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkSettingsByCluster", ctx, lockStrength, cluster)
ret0, _ := ret[0].([]*agentNetworkTypes.Settings)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkSettingsByCluster indicates an expected call of GetAgentNetworkSettingsByCluster.
func (mr *MockStoreMockRecorder) GetAgentNetworkSettingsByCluster(ctx, lockStrength, cluster interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkSettingsByCluster", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkSettingsByCluster), ctx, lockStrength, cluster)
}
// SaveAgentNetworkSettings mocks base method.
func (m *MockStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SaveAgentNetworkSettings", ctx, settings)
ret0, _ := ret[0].(error)
return ret0
}
// SaveAgentNetworkSettings indicates an expected call of SaveAgentNetworkSettings.
func (mr *MockStoreMockRecorder) SaveAgentNetworkSettings(ctx, settings interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkSettings), ctx, settings)
}
// IncrementAgentNetworkConsumption mocks base method.
func (m *MockStore) IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumption", ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD)
ret0, _ := ret[0].(error)
return ret0
}
// IncrementAgentNetworkConsumption indicates an expected call of IncrementAgentNetworkConsumption.
func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumption(ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumption), ctx, accountID, kind, dimID, windowSeconds, windowStart, tokensIn, tokensOut, costUSD)
}
// GetAgentNetworkConsumption mocks base method.
func (m *MockStore) GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkConsumption", ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart)
ret0, _ := ret[0].(*agentNetworkTypes.Consumption)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkConsumption indicates an expected call of GetAgentNetworkConsumption.
func (mr *MockStoreMockRecorder) GetAgentNetworkConsumption(ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumption), ctx, lockStrength, accountID, kind, dimID, windowSeconds, windowStart)
}
// GetAgentNetworkConsumptionBatch mocks base method.
func (m *MockStore) GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkConsumptionBatch", ctx, lockStrength, accountID, keys)
ret0, _ := ret[0].(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkConsumptionBatch indicates an expected call of GetAgentNetworkConsumptionBatch.
func (mr *MockStoreMockRecorder) GetAgentNetworkConsumptionBatch(ctx, lockStrength, accountID, keys interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkConsumptionBatch), ctx, lockStrength, accountID, keys)
}
// IncrementAgentNetworkConsumptionBatch mocks base method.
func (m *MockStore) IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "IncrementAgentNetworkConsumptionBatch", ctx, accountID, keys, tokensIn, tokensOut, costUSD)
ret0, _ := ret[0].(error)
return ret0
}
// IncrementAgentNetworkConsumptionBatch indicates an expected call of IncrementAgentNetworkConsumptionBatch.
func (mr *MockStoreMockRecorder) IncrementAgentNetworkConsumptionBatch(ctx, accountID, keys, tokensIn, tokensOut, costUSD interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementAgentNetworkConsumptionBatch", reflect.TypeOf((*MockStore)(nil).IncrementAgentNetworkConsumptionBatch), ctx, accountID, keys, tokensIn, tokensOut, costUSD)
}
// ListAgentNetworkConsumption mocks base method.
func (m *MockStore) ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListAgentNetworkConsumption", ctx, lockStrength, accountID)
ret0, _ := ret[0].([]*agentNetworkTypes.Consumption)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListAgentNetworkConsumption indicates an expected call of ListAgentNetworkConsumption.
func (mr *MockStoreMockRecorder) ListAgentNetworkConsumption(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAgentNetworkConsumption", reflect.TypeOf((*MockStore)(nil).ListAgentNetworkConsumption), ctx, lockStrength, accountID)
}
// GetAccountAgentNetworkBudgetRules mocks base method.
func (m *MockStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAccountAgentNetworkBudgetRules", ctx, lockStrength, accountID)
ret0, _ := ret[0].([]*agentNetworkTypes.AccountBudgetRule)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAccountAgentNetworkBudgetRules indicates an expected call of GetAccountAgentNetworkBudgetRules.
func (mr *MockStoreMockRecorder) GetAccountAgentNetworkBudgetRules(ctx, lockStrength, accountID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAccountAgentNetworkBudgetRules", reflect.TypeOf((*MockStore)(nil).GetAccountAgentNetworkBudgetRules), ctx, lockStrength, accountID)
}
// GetAgentNetworkBudgetRuleByID mocks base method.
func (m *MockStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkBudgetRuleByID", ctx, lockStrength, accountID, ruleID)
ret0, _ := ret[0].(*agentNetworkTypes.AccountBudgetRule)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkBudgetRuleByID indicates an expected call of GetAgentNetworkBudgetRuleByID.
func (mr *MockStoreMockRecorder) GetAgentNetworkBudgetRuleByID(ctx, lockStrength, accountID, ruleID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkBudgetRuleByID", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkBudgetRuleByID), ctx, lockStrength, accountID, ruleID)
}
// SaveAgentNetworkBudgetRule mocks base method.
func (m *MockStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SaveAgentNetworkBudgetRule", ctx, rule)
ret0, _ := ret[0].(error)
return ret0
}
// SaveAgentNetworkBudgetRule indicates an expected call of SaveAgentNetworkBudgetRule.
func (mr *MockStoreMockRecorder) SaveAgentNetworkBudgetRule(ctx, rule interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).SaveAgentNetworkBudgetRule), ctx, rule)
}
// DeleteAgentNetworkBudgetRule mocks base method.
func (m *MockStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteAgentNetworkBudgetRule", ctx, accountID, ruleID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteAgentNetworkBudgetRule indicates an expected call of DeleteAgentNetworkBudgetRule.
func (mr *MockStoreMockRecorder) DeleteAgentNetworkBudgetRule(ctx, accountID, ruleID interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAgentNetworkBudgetRule", reflect.TypeOf((*MockStore)(nil).DeleteAgentNetworkBudgetRule), ctx, accountID, ruleID)
}
// CreateAgentNetworkAccessLog mocks base method.
func (m *MockStore) CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "CreateAgentNetworkAccessLog", ctx, entry, groups)
ret0, _ := ret[0].(error)
return ret0
}
// CreateAgentNetworkAccessLog indicates an expected call of CreateAgentNetworkAccessLog.
func (mr *MockStoreMockRecorder) CreateAgentNetworkAccessLog(ctx, entry, groups interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkAccessLog", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkAccessLog), ctx, entry, groups)
}
// CreateAgentNetworkUsage mocks base method.
func (m *MockStore) CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "CreateAgentNetworkUsage", ctx, usage, groups)
ret0, _ := ret[0].(error)
return ret0
}
// CreateAgentNetworkUsage indicates an expected call of CreateAgentNetworkUsage.
func (mr *MockStoreMockRecorder) CreateAgentNetworkUsage(ctx, usage, groups interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateAgentNetworkUsage", reflect.TypeOf((*MockStore)(nil).CreateAgentNetworkUsage), ctx, usage, groups)
}
// GetAgentNetworkAccessLogs mocks base method.
func (m *MockStore) GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkAccessLogs", ctx, lockStrength, accountID, filter)
ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkAccessLog)
ret1, _ := ret[1].(int64)
ret2, _ := ret[2].(error)
return ret0, ret1, ret2
}
// GetAgentNetworkAccessLogs indicates an expected call of GetAgentNetworkAccessLogs.
func (mr *MockStoreMockRecorder) GetAgentNetworkAccessLogs(ctx, lockStrength, accountID, filter interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkAccessLogs), ctx, lockStrength, accountID, filter)
}
// GetAgentNetworkAccessLogSessions mocks base method.
func (m *MockStore) GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkAccessLogSessions", ctx, lockStrength, accountID, filter)
ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkAccessLogSession)
ret1, _ := ret[1].(int64)
ret2, _ := ret[2].(error)
return ret0, ret1, ret2
}
// GetAgentNetworkAccessLogSessions indicates an expected call of GetAgentNetworkAccessLogSessions.
func (mr *MockStoreMockRecorder) GetAgentNetworkAccessLogSessions(ctx, lockStrength, accountID, filter interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkAccessLogSessions", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkAccessLogSessions), ctx, lockStrength, accountID, filter)
}
// GetAgentNetworkUsageRows mocks base method.
func (m *MockStore) GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAgentNetworkUsageRows", ctx, lockStrength, accountID, filter)
ret0, _ := ret[0].([]*agentNetworkTypes.AgentNetworkUsage)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAgentNetworkUsageRows indicates an expected call of GetAgentNetworkUsageRows.
func (mr *MockStoreMockRecorder) GetAgentNetworkUsageRows(ctx, lockStrength, accountID, filter interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAgentNetworkUsageRows", reflect.TypeOf((*MockStore)(nil).GetAgentNetworkUsageRows), ctx, lockStrength, accountID, filter)
}
// DeleteOldAgentNetworkAccessLogs mocks base method.
func (m *MockStore) DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteOldAgentNetworkAccessLogs", ctx, accountID, olderThan)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// DeleteOldAgentNetworkAccessLogs indicates an expected call of DeleteOldAgentNetworkAccessLogs.
func (mr *MockStoreMockRecorder) DeleteOldAgentNetworkAccessLogs(ctx, accountID, olderThan interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldAgentNetworkAccessLogs", reflect.TypeOf((*MockStore)(nil).DeleteOldAgentNetworkAccessLogs), ctx, accountID, olderThan)
}
// GetAllAgentNetworkSettings mocks base method.
func (m *MockStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetAllAgentNetworkSettings", ctx, lockStrength)
ret0, _ := ret[0].([]*agentNetworkTypes.Settings)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetAllAgentNetworkSettings indicates an expected call of GetAllAgentNetworkSettings.
func (mr *MockStoreMockRecorder) GetAllAgentNetworkSettings(ctx, lockStrength interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAllAgentNetworkSettings", reflect.TypeOf((*MockStore)(nil).GetAllAgentNetworkSettings), ctx, lockStrength)
}
+3
View File
@@ -12,6 +12,9 @@ type PeerSync struct {
WireGuardPubKey string
// Meta is the system information passed by peer, must be always present
Meta nbpeer.PeerSystemMeta
// RealIP is the peer's connection IP, used to refresh its geo location.
// May be nil when the request has no associated connection IP.
RealIP net.IP
// UpdateAccountPeers indicate updating account peers,
// which occurs when the peer's metadata is updated
UpdateAccountPeers bool
+4
View File
@@ -73,6 +73,9 @@ type Settings struct {
// For new accounts this defaults to the All group.
IPv6EnabledGroups []string `gorm:"serializer:json"`
// MetricsPushEnabled globally enables or disables client metrics push for the account
MetricsPushEnabled bool `gorm:"default:false"`
// EmbeddedIdpEnabled indicates if the embedded identity provider is enabled.
// This is a runtime-only field, not stored in the database.
EmbeddedIdpEnabled bool `gorm:"-"`
@@ -110,6 +113,7 @@ func (s *Settings) Copy() *Settings {
AutoUpdateVersion: s.AutoUpdateVersion,
AutoUpdateAlways: s.AutoUpdateAlways,
IPv6EnabledGroups: slices.Clone(s.IPv6EnabledGroups),
MetricsPushEnabled: s.MetricsPushEnabled,
EmbeddedIdpEnabled: s.EmbeddedIdpEnabled,
LocalAuthDisabled: s.LocalAuthDisabled,
LocalMfaEnabled: s.LocalMfaEnabled,
+7 -5
View File
@@ -675,7 +675,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
}
if len(peersToExpire) > 0 {
if err := am.expireAndUpdatePeers(ctx, accountID, peersToExpire); err != nil {
if err := am.expireAndUpdatePeers(ctx, accountID, peersToExpire, peerExpirationUserBlocked); err != nil {
log.WithContext(ctx).Errorf("failed update expired peers: %s", err)
return nil, err
}
@@ -1059,8 +1059,8 @@ func (am *DefaultAccountManager) BuildUserInfosForAccount(ctx context.Context, a
if err != nil {
return nil, err
}
log.WithContext(ctx).Debugf("Got %d users from ExternalCache for account %s", len(usersFromIntegration), accountID)
log.WithContext(ctx).Debugf("Got %d users from InternalCache for account %s", len(queriedUsers), accountID)
log.WithContext(ctx).Tracef("Got %d users from ExternalCache for account %s", len(usersFromIntegration), accountID)
log.WithContext(ctx).Tracef("Got %d users from InternalCache for account %s", len(queriedUsers), accountID)
queriedUsers = append(queriedUsers, usersFromIntegration...)
}
@@ -1118,7 +1118,7 @@ func (am *DefaultAccountManager) BuildUserInfosForAccount(ctx context.Context, a
}
// expireAndUpdatePeers expires all peers of the given user and updates them in the account
func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accountID string, peers []*nbpeer.Peer) error {
func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accountID string, peers []*nbpeer.Peer, reason peerExpirationReason) error {
log.WithContext(ctx).Debugf("Expiring %d peers for account %s", len(peers), accountID)
settings, err := am.Store.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
if err != nil {
@@ -1145,10 +1145,12 @@ func (am *DefaultAccountManager) expireAndUpdatePeers(ctx context.Context, accou
if err := am.Store.SavePeerStatus(ctx, accountID, peer.ID, *peer.Status); err != nil {
return err
}
meta := peer.EventMeta(dnsDomain)
meta["reason"] = string(reason)
am.StoreEvent(
ctx,
peer.UserID, peer.ID, accountID,
activity.PeerLoginExpired, peer.EventMeta(dnsDomain),
activity.PeerLoginExpired, meta,
)
}