Merge branch 'main' into fix/remove-math-rand

This commit is contained in:
pascal
2026-07-29 12:06:01 +02:00
292 changed files with 20363 additions and 4042 deletions

177
management/cmd/admin.go Normal file
View File

@@ -0,0 +1,177 @@
package cmd
import (
"context"
"fmt"
"path"
"path/filepath"
"github.com/dexidp/dex/storage"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
"github.com/netbirdio/netbird/formatter/hook"
admincmd "github.com/netbirdio/netbird/management/cmd/admin"
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
"github.com/netbirdio/netbird/management/server/activity"
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/util"
)
var adminDatadir string
// newAdminCommands creates the admin command tree with management-specific resource openers.
func newAdminCommands() *cobra.Command {
cmd := admincmd.NewCommands(admincmd.Openers{
Resources: withAdminResources,
Store: withAdminStoreOnly,
IDP: withAdminIDPOnly,
})
cmd.PersistentFlags().StringVar(&adminDatadir, "datadir", "", "Override the data directory from config (used for store.db and the default idp.db)")
return cmd
}
func newLegacyTokenCommand() *cobra.Command {
cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
cmd.Deprecated = "use 'admin token' instead"
cmd.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
return cmd
}
// withAdminResources initializes logging, loads config, opens the management store
// and embedded IdP storage, and calls fn.
func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
managementStore, err := openAdminStore(ctx, config, datadir)
if err != nil {
return err
}
defer admincmd.CloseStore(ctx, managementStore)
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
if err != nil {
return err
}
defer admincmd.CloseIDPStorage(idpStorage)
eventStore, esErr := openAdminEventStore(ctx, config, datadir)
if esErr != nil {
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
}
if eventStore != nil {
defer func() {
if err := eventStore.Close(ctx); err != nil {
log.Debugf("close activity event store: %v", err)
}
}()
}
return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
})
}
// withAdminStoreOnly opens only the management store for admin subcommands that do not
// need embedded IdP storage.
func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
return withAdminConfig(cmd, false, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
managementStore, err := openAdminStore(ctx, config, datadir)
if err != nil {
return err
}
defer admincmd.CloseStore(ctx, managementStore)
return fn(ctx, managementStore)
})
}
func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, _ string) error {
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
if err != nil {
return err
}
defer admincmd.CloseIDPStorage(idpStorage)
return fn(ctx, idpStorage, idpStorageFile)
})
}
func withAdminConfig(cmd *cobra.Command, applyIDPDefaults bool, fn func(ctx context.Context, config *nbconfig.Config, datadir string) error) error {
if err := util.InitLog("error", "console"); err != nil {
return fmt.Errorf("init log: %w", err)
}
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
config, datadir, err := loadAdminMgmtConfig(ctx, applyIDPDefaults)
if err != nil {
return fmt.Errorf("load config: %w", err)
}
return fn(ctx, config, datadir)
}
func loadAdminMgmtConfig(ctx context.Context, applyIDPDefaults bool) (*nbconfig.Config, string, error) {
config := &nbconfig.Config{}
if _, err := util.ReadJsonWithEnvSub(nbconfig.MgmtConfigPath, config); err != nil {
return nil, "", err
}
if applyIDPDefaults {
if err := ApplyEmbeddedIdPConfig(ctx, config); err != nil {
return nil, "", err
}
}
datadir := config.Datadir
applyAdminDatadirOverride(config, &datadir)
return config, datadir, nil
}
func applyAdminDatadirOverride(config *nbconfig.Config, datadir *string) {
if adminDatadir == "" {
return
}
oldDatadir := *datadir
*datadir = adminDatadir
if config.EmbeddedIdP != nil && config.EmbeddedIdP.Storage.Type == "sqlite3" && isDefaultIDPStorageFile(config.EmbeddedIdP.Storage.Config.File, oldDatadir) {
config.EmbeddedIdP.Storage.Config.File = filepath.Join(*datadir, "idp.db")
}
}
func isDefaultIDPStorageFile(file, datadir string) bool {
if file == "" {
return true
}
defaultFile := filepath.Join(datadir, "idp.db")
legacyDefaultFile := path.Join(datadir, "idp.db")
legacySlashDefaultFile := path.Join(filepath.ToSlash(datadir), "idp.db")
return filepath.Clean(file) == filepath.Clean(defaultFile) ||
file == legacyDefaultFile ||
filepath.ToSlash(file) == legacySlashDefaultFile
}
func openAdminStore(ctx context.Context, config *nbconfig.Config, datadir string) (store.Store, error) {
managementStore, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
if err != nil {
return nil, fmt.Errorf("create store: %w", err)
}
return managementStore, nil
}
func openAdminEventStore(ctx context.Context, config *nbconfig.Config, datadir string) (activity.Store, error) {
if config.DataStoreEncryptionKey == "" {
return nil, fmt.Errorf("data store encryption key is not configured")
}
eventStore, err := activitystore.NewSqlStore(ctx, datadir, config.DataStoreEncryptionKey)
if err != nil {
return nil, fmt.Errorf("open activity event store: %w", err)
}
if eventStore == nil {
return nil, fmt.Errorf("open activity event store: returned nil store")
}
return eventStore, nil
}

View File

@@ -0,0 +1,577 @@
// Package admincmd provides reusable cobra commands for self-hosted administrator helpers.
// Both the management and combined binaries use these commands, each providing
// their own opener to handle config loading and storage initialization.
package admincmd
import (
"context"
"errors"
"fmt"
"io"
"log/slog"
"os"
"strings"
"time"
"github.com/dexidp/dex/storage"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
"golang.org/x/crypto/bcrypt"
"github.com/netbirdio/netbird/formatter/hook"
nbdex "github.com/netbirdio/netbird/idp/dex"
"github.com/netbirdio/netbird/management/cmd/proxy"
"github.com/netbirdio/netbird/management/cmd/token"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
"github.com/netbirdio/netbird/management/server"
"github.com/netbirdio/netbird/management/server/activity"
"github.com/netbirdio/netbird/management/server/idp"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
)
// Resources contains the storages required by the admin commands.
type Resources struct {
Store store.Store
IDPStorage storage.Storage
IDPStorageFile string
EventStore activity.Store
}
// Opener initializes command resources from the command context and calls fn.
type Opener func(cmd *cobra.Command, fn func(ctx context.Context, resources Resources) error) error
// StoreOpener initializes only the management store from the command context and calls fn.
type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
// IDPOpener initializes only the embedded IdP storage from the command context and calls fn.
type IDPOpener func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error
// Openers contains the resource openers needed by the admin command tree.
type Openers struct {
Resources Opener
Store StoreOpener
IDP IDPOpener
}
type userSelector struct {
email string
userID string
}
func (s userSelector) normalized() userSelector {
return userSelector{
email: strings.TrimSpace(s.email),
userID: strings.TrimSpace(s.userID),
}
}
func (s userSelector) validate() error {
s = s.normalized()
if (s.email == "") == (s.userID == "") {
return fmt.Errorf("provide exactly one of --email or --user-id")
}
return nil
}
// NewCommands creates the admin command tree with the given resource openers.
func NewCommands(openers Openers) *cobra.Command {
adminCmd := &cobra.Command{
Use: "admin",
Short: "Self-hosted administrator helpers",
Long: "Administrative helpers for self-hosted deployments using the embedded identity provider.",
}
userCmd := &cobra.Command{
Use: "user",
Short: "Manage local embedded IdP users",
}
var passwordSelector userSelector
var password string
var passwordFile string
passwordCmd := &cobra.Command{
Use: "change-password (--email email | --user-id id) (--password password | --password-file path)",
Aliases: []string{"set-password"},
Short: "Change a local user's password",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
if err := passwordSelector.validate(); err != nil {
return err
}
newPassword, err := resolvePasswordInput(cmd, password, passwordFile)
if err != nil {
return err
}
return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
return runChangePassword(ctx, idpStorage, cmd.OutOrStdout(), passwordSelector, newPassword, storageFile)
})
},
}
addUserSelectorFlags(passwordCmd, &passwordSelector)
passwordCmd.Flags().StringVar(&password, "password", "", "New password for the user")
passwordCmd.Flags().StringVar(&passwordFile, "password-file", "", "Read new password from file ('-' for stdin)")
var resetSelector userSelector
resetMFACmd := &cobra.Command{
Use: "reset-mfa (--email email | --user-id id)",
Short: "Reset a local user's MFA enrollment",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
if err := resetSelector.validate(); err != nil {
return err
}
return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
return runResetMFA(ctx, idpStorage, cmd.OutOrStdout(), resetSelector, storageFile)
})
},
}
addUserSelectorFlags(resetMFACmd, &resetSelector)
userCmd.AddCommand(passwordCmd, resetMFACmd)
mfaCmd := &cobra.Command{
Use: "mfa",
Short: "Manage local MFA for embedded IdP users",
}
enableCmd := &cobra.Command{
Use: "enable",
Short: "Enable MFA for local embedded IdP users",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), true)
})
},
}
disableCmd := &cobra.Command{
Use: "disable",
Short: "Disable MFA for local embedded IdP users",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), false)
})
},
}
statusCmd := &cobra.Command{
Use: "status",
Short: "Show local MFA status",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
return runMFAStatus(ctx, resources, cmd.OutOrStdout())
})
},
}
mfaCmd.AddCommand(enableCmd, disableCmd, statusCmd)
adminCmd.AddCommand(userCmd, mfaCmd)
if openers.Store != nil {
adminCmd.AddCommand(tokencmd.NewCommands(tokencmd.StoreOpener(openers.Store)))
adminCmd.AddCommand(proxycmd.NewCommands(proxycmd.StoreOpener(openers.Store)))
}
return adminCmd
}
// OpenEmbeddedIDPStorage opens the Dex storage configured for the embedded IdP.
func OpenEmbeddedIDPStorage(cfg *idp.EmbeddedIdPConfig) (storage.Storage, error) {
if cfg == nil || !cfg.Enabled {
return nil, fmt.Errorf("admin commands require the embedded IdP to be enabled")
}
yamlConfig, err := cfg.ToYAMLConfig()
if err != nil {
return nil, fmt.Errorf("build embedded IdP config: %w", err)
}
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
st, err := yamlConfig.Storage.OpenStorage(logger)
if err != nil {
return nil, fmt.Errorf("open embedded IdP storage: %w", err)
}
return st, nil
}
// CloseStore closes the management store and logs cleanup errors at debug level.
func CloseStore(ctx context.Context, s store.Store) {
if s == nil {
return
}
if err := s.Close(ctx); err != nil {
log.Debugf("close store: %v", err)
}
}
// OpenIDPStorage opens embedded IdP storage and returns its sqlite file path when applicable.
func OpenIDPStorage(config *nbconfig.Config) (storage.Storage, string, error) {
if config == nil {
return nil, "", fmt.Errorf("management config is required")
}
idpStorage, err := OpenEmbeddedIDPStorage(config.EmbeddedIdP)
if err != nil {
return nil, "", err
}
return idpStorage, embeddedIDPStorageFile(config), nil
}
func embeddedIDPStorageFile(config *nbconfig.Config) string {
if config.EmbeddedIdP == nil || config.EmbeddedIdP.Storage.Type != "sqlite3" {
return ""
}
return config.EmbeddedIdP.Storage.Config.File
}
// CloseIDPStorage closes embedded IdP storage and logs cleanup errors at debug level.
func CloseIDPStorage(s storage.Storage) {
if s == nil {
return
}
if err := s.Close(); err != nil {
log.Debugf("close embedded IdP storage: %v", err)
}
}
func addUserSelectorFlags(cmd *cobra.Command, selector *userSelector) {
cmd.Flags().StringVar(&selector.email, "email", "", "User email")
cmd.Flags().StringVar(&selector.userID, "user-id", "", "User ID")
}
func resolvePasswordInput(cmd *cobra.Command, password, passwordFile string) (string, error) {
if password != "" && passwordFile != "" {
return "", fmt.Errorf("provide only one of --password or --password-file")
}
if passwordFile == "" {
return password, nil
}
var data []byte
var err error
if passwordFile == "-" {
data, err = io.ReadAll(cmd.InOrStdin())
} else {
data, err = os.ReadFile(passwordFile)
}
if err != nil {
return "", fmt.Errorf("read password: %w", err)
}
return strings.TrimRight(string(data), "\r\n"), nil
}
func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string, idpStorageFile string) error {
if idpStorage == nil {
return fmt.Errorf("embedded IdP storage is required")
}
selector = selector.normalized()
if err := selector.validate(); err != nil {
return err
}
if password == "" {
return fmt.Errorf("password is required")
}
if err := server.ValidatePassword(password); err != nil {
return fmt.Errorf("invalid password: %w", err)
}
user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
if err != nil {
return err
}
hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
if err != nil {
return fmt.Errorf("hash password: %w", err)
}
if err := idpStorage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
old.Hash = hash
return old, nil
}); err != nil {
return fmt.Errorf("update password for %s: %w", user.Email, err)
}
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
return err
}
_, _ = fmt.Fprintf(w, "Password updated for %s.\n", user.Email)
return nil
}
func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, idpStorageFile string) error {
if idpStorage == nil {
return fmt.Errorf("embedded IdP storage is required")
}
selector = selector.normalized()
if err := selector.validate(); err != nil {
return err
}
user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
if err != nil {
return err
}
reset := false
err = idpStorage.UpdateUserIdentity(ctx, user.UserID, idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
reset = reset || len(old.MFASecrets) > 0 || len(old.WebAuthnCredentials) > 0
old.MFASecrets = map[string]*storage.MFASecret{}
old.WebAuthnCredentials = map[string][]storage.WebAuthnCredential{}
return old, nil
})
if errors.Is(err, storage.ErrNotFound) {
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
return err
}
_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
return nil
}
if err != nil {
return fmt.Errorf("reset MFA for %s: %w", user.Email, err)
}
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
return err
}
if reset {
_, _ = fmt.Fprintf(w, "MFA reset for %s. The user will re-enroll at next login.\n", user.Email)
} else {
_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
}
return nil
}
func runSetMFAEnabled(ctx context.Context, resources Resources, w io.Writer, enabled bool) error {
if resources.Store == nil {
return fmt.Errorf("management store is required")
}
if resources.IDPStorage == nil {
return fmt.Errorf("embedded IdP storage is required")
}
accountID, settings, err := getSingleAccountSettings(ctx, resources.Store)
if err != nil {
return err
}
oldEnabled := settings.LocalMfaEnabled
newSettings := settings.Copy()
newSettings.LocalMfaEnabled = enabled
if err := setIDPClientsMFA(ctx, resources.IDPStorage, enabled); err != nil {
return err
}
if err := resources.Store.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
if rollbackErr := setIDPClientsMFA(ctx, resources.IDPStorage, oldEnabled); rollbackErr != nil {
return fmt.Errorf("save local MFA account setting: %w (also failed to roll back embedded IdP MFA state: %v)", err, rollbackErr)
}
return fmt.Errorf("save local MFA account setting: %w", err)
}
if err := storeMFAActivity(ctx, resources.EventStore, accountID, enabled); err != nil {
_, _ = fmt.Fprintf(w, "Warning: failed to record audit event: %v\n", err)
}
state := "disabled"
if enabled {
state = "enabled"
}
_, _ = fmt.Fprintf(w, "Local MFA %s.\n", state)
return nil
}
func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
if resources.Store == nil {
return fmt.Errorf("management store is required")
}
if resources.IDPStorage == nil {
return fmt.Errorf("embedded IdP storage is required")
}
_, settings, err := getSingleAccountSettings(ctx, resources.Store)
if err != nil {
return err
}
accountStatus := "disabled"
if settings.LocalMfaEnabled {
accountStatus = "enabled"
}
clientStatus, err := idpClientsMFAStatus(ctx, resources.IDPStorage)
if err != nil {
return err
}
_, _ = fmt.Fprintf(w, "Account setting: %s\n", accountStatus)
_, _ = fmt.Fprintf(w, "Embedded IdP clients: %s\n", clientStatus)
return nil
}
func getSingleAccountSettings(ctx context.Context, s store.Store) (string, *types.Settings, error) {
count, err := s.GetAccountsCounter(ctx)
if err != nil {
return "", nil, fmt.Errorf("count accounts: %w", err)
}
if count != 1 {
return "", nil, fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", count)
}
accountID, err := s.GetAnyAccountID(ctx)
if err != nil {
return "", nil, fmt.Errorf("get account ID: %w", err)
}
settings, err := s.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return "", nil, fmt.Errorf("get account settings: %w", err)
}
if settings == nil {
settings = &types.Settings{}
}
return accountID, settings, nil
}
func storeMFAActivity(ctx context.Context, eventStore activity.Store, accountID string, enabled bool) error {
if eventStore == nil {
return nil
}
event := activity.AccountLocalMfaDisabled
if enabled {
event = activity.AccountLocalMfaEnabled
}
_, err := eventStore.Save(ctx, &activity.Event{
Timestamp: time.Now().UTC(),
Activity: event,
InitiatorID: string(hook.SystemSource),
TargetID: accountID,
AccountID: accountID,
})
if err != nil {
return fmt.Errorf("save local MFA audit event: %w", err)
}
return nil
}
func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector, idpStorageFile string) (storage.Password, error) {
selector = selector.normalized()
if err := selector.validate(); err != nil {
return storage.Password{}, err
}
if selector.email != "" {
user, err := idpStorage.GetPassword(ctx, selector.email)
if errors.Is(err, storage.ErrNotFound) {
if empty, listErr := localUsersEmpty(ctx, idpStorage); listErr != nil {
return storage.Password{}, listErr
} else if empty {
return storage.Password{}, noLocalUsersError(idpStorageFile)
}
return storage.Password{}, fmt.Errorf("local user with email %q not found", selector.email)
}
if err != nil {
return storage.Password{}, fmt.Errorf("get local user by email %q: %w", selector.email, err)
}
return user, nil
}
rawUserID := selector.userID
if decodedUserID, _, err := nbdex.DecodeDexUserID(selector.userID); err == nil && decodedUserID != "" {
rawUserID = decodedUserID
}
users, err := idpStorage.ListPasswords(ctx)
if err != nil {
return storage.Password{}, fmt.Errorf("list local users: %w", err)
}
for _, user := range users {
if user.UserID == rawUserID || user.UserID == selector.userID {
return user, nil
}
}
if len(users) == 0 {
return storage.Password{}, noLocalUsersError(idpStorageFile)
}
return storage.Password{}, fmt.Errorf("local user with ID %q not found", selector.userID)
}
func localUsersEmpty(ctx context.Context, idpStorage storage.Storage) (bool, error) {
users, err := idpStorage.ListPasswords(ctx)
if err != nil {
return false, fmt.Errorf("list local users: %w", err)
}
return len(users) == 0, nil
}
func noLocalUsersError(idpStorageFile string) error {
location := ""
if idpStorageFile != "" {
location = fmt.Sprintf(" (%s)", idpStorageFile)
}
return fmt.Errorf("no local users exist in the embedded IdP storage%s; the management server may never have started with this config, or --datadir points at the wrong location", location)
}
func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, userID string) error {
err := idpStorage.DeleteAuthSession(ctx, userID, idp.LocalConnectorID)
if err == nil || errors.Is(err, storage.ErrNotFound) {
return nil
}
return fmt.Errorf("delete local auth session for user %s: %w", userID, err)
}
func setIDPClientsMFA(ctx context.Context, idpStorage storage.Storage, enabled bool) error {
var mfaChain []string
if enabled {
mfaChain = []string{idp.DefaultTOTPAuthenticatorID}
}
clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
if err := nbdex.SetClientsMFAChain(ctx, idpStorage, clientIDs, mfaChain); err != nil {
if errors.Is(err, storage.ErrNotFound) {
return fmt.Errorf("embedded IdP client not found; start the management server once before toggling MFA: %w", err)
}
return fmt.Errorf("update MFA chain on embedded IdP clients: %w", err)
}
return nil
}
func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (string, error) {
clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
enabledCount := 0
for _, clientID := range clientIDs {
client, err := idpStorage.GetClient(ctx, clientID)
if errors.Is(err, storage.ErrNotFound) {
return "unknown", fmt.Errorf("embedded IdP client %q not found", clientID)
}
if err != nil {
return "unknown", fmt.Errorf("get embedded IdP client %q: %w", clientID, err)
}
if hasAuthenticator(client.MFAChain, idp.DefaultTOTPAuthenticatorID) {
enabledCount++
}
}
switch enabledCount {
case 0:
return "disabled", nil
case len(clientIDs):
return "enabled", nil
default:
return "partially enabled", nil
}
}
func hasAuthenticator(chain []string, authenticatorID string) bool {
for _, id := range chain {
if id == authenticatorID {
return true
}
}
return false
}

View File

@@ -0,0 +1,250 @@
package admincmd
import (
"bytes"
"context"
"io"
"log/slog"
"strings"
"testing"
"time"
"github.com/dexidp/dex/storage"
"github.com/dexidp/dex/storage/memory"
"github.com/spf13/cobra"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/bcrypt"
nbdex "github.com/netbirdio/netbird/idp/dex"
"github.com/netbirdio/netbird/management/server/idp"
mgmtstore "github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"
)
func newTestIDPStorage(t *testing.T) storage.Storage {
t.Helper()
st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
hash, err := bcrypt.GenerateFromPassword([]byte("OldPass1!"), bcrypt.DefaultCost)
require.NoError(t, err)
require.NoError(t, st.CreatePassword(context.Background(), storage.Password{
Email: "user@example.com",
Username: "User",
UserID: "user-1",
Hash: hash,
}))
require.NoError(t, st.CreateUserIdentity(context.Background(), storage.UserIdentity{
UserID: "user-1",
ConnectorID: idp.LocalConnectorID,
MFASecrets: map[string]*storage.MFASecret{
idp.DefaultTOTPAuthenticatorID: {
AuthenticatorID: idp.DefaultTOTPAuthenticatorID,
Type: "TOTP",
Secret: "otpauth://totp/NetBird:user@example.com?secret=ABC",
Confirmed: true,
CreatedAt: time.Now(),
},
},
WebAuthnCredentials: map[string][]storage.WebAuthnCredential{
"webauthn": {{CredentialID: []byte("credential")}},
},
}))
require.NoError(t, st.CreateAuthSession(context.Background(), storage.AuthSession{
UserID: "user-1",
ConnectorID: idp.LocalConnectorID,
Nonce: "nonce",
}))
require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientCLI, Name: "CLI"}))
require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientDashboard, Name: "Dashboard"}))
return st
}
func TestRunChangePassword(t *testing.T) {
ctx := context.Background()
st := newTestIDPStorage(t)
var out bytes.Buffer
err := runChangePassword(ctx, st, &out, userSelector{email: "user@example.com"}, "NewPass1!", "")
require.NoError(t, err)
require.Contains(t, out.String(), "Password updated")
user, err := st.GetPassword(ctx, "user@example.com")
require.NoError(t, err)
require.NoError(t, bcrypt.CompareHashAndPassword(user.Hash, []byte("NewPass1!")))
_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
require.ErrorIs(t, err, storage.ErrNotFound)
}
func TestRunChangePasswordValidatesPassword(t *testing.T) {
st := newTestIDPStorage(t)
err := runChangePassword(context.Background(), st, io.Discard, userSelector{email: "user@example.com"}, "short", "")
require.Error(t, err)
require.Contains(t, err.Error(), "invalid password")
}
func TestRunResetMFA(t *testing.T) {
ctx := context.Background()
st := newTestIDPStorage(t)
var out bytes.Buffer
encodedUserID := nbdex.EncodeDexUserID("user-1", idp.LocalConnectorID)
err := runResetMFA(ctx, st, &out, userSelector{userID: encodedUserID}, "")
require.NoError(t, err)
require.Contains(t, out.String(), "MFA reset")
identity, err := st.GetUserIdentity(ctx, "user-1", idp.LocalConnectorID)
require.NoError(t, err)
require.Empty(t, identity.MFASecrets)
require.Empty(t, identity.WebAuthnCredentials)
_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
require.ErrorIs(t, err, storage.ErrNotFound)
}
func TestRunResetMFAWithoutEnrollment(t *testing.T) {
ctx := context.Background()
st := newTestIDPStorage(t)
require.NoError(t, st.UpdateUserIdentity(ctx, "user-1", idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
old.MFASecrets = nil
old.WebAuthnCredentials = nil
return old, nil
}))
var out bytes.Buffer
err := runResetMFA(ctx, st, &out, userSelector{email: "user@example.com"}, "")
require.NoError(t, err)
require.Contains(t, out.String(), "No MFA enrollment found")
}
func TestSetIDPClientsMFA(t *testing.T) {
ctx := context.Background()
st := newTestIDPStorage(t)
require.NoError(t, setIDPClientsMFA(ctx, st, true))
status, err := idpClientsMFAStatus(ctx, st)
require.NoError(t, err)
require.Equal(t, "enabled", status)
require.NoError(t, setIDPClientsMFA(ctx, st, false))
status, err = idpClientsMFAStatus(ctx, st)
require.NoError(t, err)
require.Equal(t, "disabled", status)
}
func newTestManagementStore(t *testing.T, localMFAEnabled bool) mgmtstore.Store {
t.Helper()
ctx := context.Background()
st, err := mgmtstore.NewStore(ctx, types.SqliteStoreEngine, t.TempDir(), nil, false)
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, st.Close(ctx)) })
require.NoError(t, st.SaveAccount(ctx, &types.Account{
Id: "account-1",
Settings: &types.Settings{LocalMfaEnabled: localMFAEnabled},
}))
return st
}
func TestRunSetMFAEnabledDoesNotSaveWhenIDPUpdateFails(t *testing.T) {
ctx := context.Background()
managementStore := newTestManagementStore(t, false)
idpStorage := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
require.Error(t, err)
require.Contains(t, err.Error(), "embedded IdP client")
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
require.NoError(t, err)
require.False(t, settings.LocalMfaEnabled)
}
func TestRunSetMFAEnabledUpdatesSettingsAfterIDP(t *testing.T) {
ctx := context.Background()
managementStore := newTestManagementStore(t, false)
idpStorage := newTestIDPStorage(t)
err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
require.NoError(t, err)
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
require.NoError(t, err)
require.True(t, settings.LocalMfaEnabled)
clientStatus, err := idpClientsMFAStatus(ctx, idpStorage)
require.NoError(t, err)
require.Equal(t, "enabled", clientStatus)
}
func TestRunSetMFAEnabledSucceedsWithNilEventStore(t *testing.T) {
ctx := context.Background()
managementStore := newTestManagementStore(t, false)
idpStorage := newTestIDPStorage(t)
var out bytes.Buffer
var err error
require.NotPanics(t, func() {
err = runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage, EventStore: nil}, &out, true)
})
require.NoError(t, err)
require.Contains(t, out.String(), "Local MFA enabled")
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
require.NoError(t, err)
require.True(t, settings.LocalMfaEnabled)
}
func TestUserSelectorValidate(t *testing.T) {
require.NoError(t, userSelector{email: " user@example.com "}.validate())
require.NoError(t, userSelector{userID: "user-1"}.validate())
require.Error(t, userSelector{}.validate())
require.Error(t, userSelector{email: "user@example.com", userID: "user-1"}.validate())
}
func TestFindLocalUserNotFound(t *testing.T) {
st := newTestIDPStorage(t)
_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "")
require.Error(t, err)
require.True(t, strings.Contains(err.Error(), "not found"))
}
func TestFindLocalUserZeroUsersIncludesStoragePath(t *testing.T) {
st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "/var/lib/netbird/idp.db")
require.Error(t, err)
require.Contains(t, err.Error(), "no local users exist")
require.Contains(t, err.Error(), "/var/lib/netbird/idp.db")
}
func TestUserCommandValidatesSelectorBeforeOpeningStorage(t *testing.T) {
opened := false
cmd := NewCommands(Openers{
IDP: func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
opened = true
return nil
},
})
cmd.SetArgs([]string{"user", "change-password", "--password", "NewPass1!"})
cmd.SetOut(io.Discard)
cmd.SetErr(io.Discard)
err := cmd.Execute()
require.Error(t, err)
require.Contains(t, err.Error(), "provide exactly one")
require.False(t, opened)
}
func TestResolvePasswordInputFromStdin(t *testing.T) {
cmd := &cobra.Command{}
cmd.SetIn(strings.NewReader("NewPass1!\n"))
password, err := resolvePasswordInput(cmd, "", "-")
require.NoError(t, err)
require.Equal(t, "NewPass1!", password)
}
func TestResolvePasswordInputRejectsMultipleSources(t *testing.T) {
_, err := resolvePasswordInput(&cobra.Command{}, "NewPass1!", "-")
require.Error(t, err)
}

View File

@@ -0,0 +1,80 @@
package cmd
import (
"context"
"path"
"path/filepath"
"testing"
"github.com/stretchr/testify/require"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
"github.com/netbirdio/netbird/management/server/idp"
)
func TestApplyAdminDatadirOverrideRelocatesDefaultIDPStorage(t *testing.T) {
oldDatadir := filepath.Join(t.TempDir(), "old")
newDatadir := filepath.Join(t.TempDir(), "new")
for _, defaultFile := range []string{
"",
filepath.Join(oldDatadir, "idp.db"),
path.Join(oldDatadir, "idp.db"),
} {
t.Run(defaultFile, func(t *testing.T) {
cfg := &nbconfig.Config{
EmbeddedIdP: &idp.EmbeddedIdPConfig{
Enabled: true,
Storage: idp.EmbeddedStorageConfig{
Type: "sqlite3",
Config: idp.EmbeddedStorageTypeConfig{
File: defaultFile,
},
},
},
}
datadir := oldDatadir
oldAdminDatadir := adminDatadir
adminDatadir = newDatadir
t.Cleanup(func() { adminDatadir = oldAdminDatadir })
applyAdminDatadirOverride(cfg, &datadir)
require.Equal(t, newDatadir, datadir)
require.Equal(t, filepath.Join(newDatadir, "idp.db"), cfg.EmbeddedIdP.Storage.Config.File)
})
}
}
func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
eventStore, err := openAdminEventStore(context.Background(), &nbconfig.Config{}, t.TempDir())
require.Error(t, err)
require.Contains(t, err.Error(), "encryption key")
require.Nil(t, eventStore)
}
func TestApplyAdminDatadirOverrideKeepsExplicitIDPStorage(t *testing.T) {
oldDatadir := filepath.Join(t.TempDir(), "old")
newDatadir := filepath.Join(t.TempDir(), "new")
explicitFile := filepath.Join(t.TempDir(), "custom-idp.db")
cfg := &nbconfig.Config{
EmbeddedIdP: &idp.EmbeddedIdPConfig{
Enabled: true,
Storage: idp.EmbeddedStorageConfig{
Type: "sqlite3",
Config: idp.EmbeddedStorageTypeConfig{
File: explicitFile,
},
},
},
}
datadir := oldDatadir
oldAdminDatadir := adminDatadir
adminDatadir = newDatadir
t.Cleanup(func() { adminDatadir = oldAdminDatadir })
applyAdminDatadirOverride(cfg, &datadir)
require.Equal(t, newDatadir, datadir)
require.Equal(t, explicitFile, cfg.EmbeddedIdP.Storage.Config.File)
}

View File

@@ -13,6 +13,7 @@ import (
"os"
"os/signal"
"path"
"path/filepath"
"strings"
"syscall"
@@ -25,6 +26,7 @@ import (
"github.com/netbirdio/netbird/management/internals/server"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbdomain "github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/grpc"
"github.com/netbirdio/netbird/util"
"github.com/netbirdio/netbird/util/crypt"
)
@@ -153,8 +155,20 @@ func LoadMgmtConfig(ctx context.Context, mgmtConfigPath string) (*nbconfig.Confi
ApplyCommandLineOverrides(loadedConfig)
err := grpc.ValidateSyncMessageVersion(loadedConfig.HighestSupportedSyncMessageVersion)
if err != nil {
return nil, err
}
for account, version := range loadedConfig.PerAccountHighestSupportedSyncMessageVersion {
err := grpc.ValidateSyncMessageVersion(&version)
if err != nil {
return nil, fmt.Errorf("unrecognized sync message version for account %s, %w", account, err)
}
}
// Apply EmbeddedIdP config to HttpConfig if embedded IdP is enabled
err := ApplyEmbeddedIdPConfig(ctx, loadedConfig)
err = ApplyEmbeddedIdPConfig(ctx, loadedConfig)
if err != nil {
return nil, err
}
@@ -209,7 +223,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config) error {
cfg.EmbeddedIdP.Storage.Type = "sqlite3"
}
if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
cfg.EmbeddedIdP.Storage.Config.File = filepath.Join(cfg.Datadir, "idp.db")
}
issuer := cfg.EmbeddedIdP.Issuer

View File

@@ -4,6 +4,9 @@ import (
"context"
"os"
"testing"
"github.com/netbirdio/netbird/shared/management/grpc"
"github.com/stretchr/testify/assert"
)
const (
@@ -20,34 +23,49 @@ const (
"AuthAudience": "https://stageapp/",
"AuthIssuer": "https://something.eu.auth0.com/",
"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
},
"HighestSupportedSyncMessageVersion": 1,
"PerAccountHighestSupportedSyncMessageVersion": {
"1": 0,
"2": 1
}
}`
)
func Test_loadMgmtConfig(t *testing.T) {
tmpFile, err := createConfig()
if err != nil {
t.Fatalf("failed to create config: %s", err)
}
func Test_LoadMgmtConfig(t *testing.T) {
tmpFile, err := createConfig(exampleConfig)
assert.NoError(t, err)
cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
if err != nil {
t.Fatalf("failed to load management config: %s", err)
}
if cfg.Relay == nil {
t.Fatalf("config is nil")
}
if len(cfg.Relay.Addresses) == 0 {
t.Fatalf("relay address is empty")
}
assert.NoError(t, err)
assert.NotEmpty(t, cfg.Relay)
assert.NotEmpty(t, cfg.Relay.Addresses)
assert.Equal(t, int(grpc.ComponentNetworkMap), *cfg.HighestSupportedSyncMessageVersion)
assert.Equal(t, map[string]int{"1": int(grpc.Base), "2": int(grpc.ComponentNetworkMap)}, cfg.PerAccountHighestSupportedSyncMessageVersion)
}
func createConfig() (string, error) {
func Test_LoadMgmtConfig_Empty(t *testing.T) {
tmpFile, err := createConfig(`{
"HttpConfig": {
"AuthAudience": "https://stageapp/",
"AuthIssuer": "https://something.eu.auth0.com/",
"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
}
}`)
assert.NoError(t, err)
cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
assert.NoError(t, err)
assert.Nil(t, cfg.HighestSupportedSyncMessageVersion)
assert.Nil(t, cfg.PerAccountHighestSupportedSyncMessageVersion)
}
func createConfig(config string) (string, error) {
tmpfile, err := os.CreateTemp("", "config.json")
if err != nil {
return "", err
}
_, err = tmpfile.Write([]byte(exampleConfig))
_, err = tmpfile.Write([]byte(config))
if err != nil {
return "", err
}

View File

@@ -0,0 +1,141 @@
// Package proxycmd provides reusable cobra commands for managing reverse proxy instances.
// Both the management and combined binaries use these commands, each providing
// their own StoreOpener to handle config loading and store initialization.
package proxycmd
import (
"bufio"
"context"
"fmt"
"io"
"strings"
"text/tabwriter"
"github.com/spf13/cobra"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
"github.com/netbirdio/netbird/management/server/store"
)
// StoreOpener initializes a store from the command context and calls fn.
type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
const disconnectAllConfirmation = "disconnect all proxies"
// NewCommands creates the proxy command tree with the given store opener.
// Returns the parent "proxy" command with the disconnect-all subcommand.
func NewCommands(opener StoreOpener) *cobra.Command {
var dryRun bool
var force bool
proxyCmd := &cobra.Command{
Use: "proxy",
Short: "Manage reverse proxy instances",
Long: "Commands for inspecting and repairing the reverse proxy instances registered with the management server.",
}
disconnectAllCmd := &cobra.Command{
Use: "disconnect-all",
Short: "Force-mark all reverse proxy instances as disconnected",
Long: "Lists all reverse proxy instances and force-marks them as disconnected, regardless of their session state. " +
"Use this to repair stale connection state, e.g. after an unclean management server shutdown. " +
"By default, it asks for manual confirmation before changing state. Use --dry-run to preview without changing state, or --force to skip confirmation. " +
"Run during a maintenance window; affected live proxies may stay hidden until their next heartbeat or reconnect/re-register.",
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, _ []string) error {
return opener(cmd, func(ctx context.Context, s store.Store) error {
return runDisconnectAll(ctx, s, cmd.OutOrStdout(), cmd.InOrStdin(), dryRun, force)
})
},
}
disconnectAllCmd.Flags().BoolVar(&dryRun, "dry-run", false, "List reverse proxy instances that would be disconnected without changing state")
disconnectAllCmd.Flags().BoolVar(&force, "force", false, "Skip the confirmation prompt and apply the repair")
proxyCmd.AddCommand(disconnectAllCmd)
return proxyCmd
}
func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.Reader, dryRun, force bool) error {
proxies, err := s.GetAllProxies(ctx)
if err != nil {
return fmt.Errorf("list proxies: %w", err)
}
if len(proxies) == 0 {
_, _ = fmt.Fprintln(out, "No reverse proxy instances found.")
return nil
}
toDisconnect := 0
w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tACCOUNT\tSTATUS\tLAST SEEN")
_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t------\t---------")
for _, p := range proxies {
if p.Status != rpproxy.StatusDisconnected {
toDisconnect++
}
account := "-"
if p.AccountID != nil {
account = *p.AccountID
}
_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
p.ID,
p.ClusterAddress,
p.IPAddress,
account,
p.Status,
p.LastSeen.Format("2006-01-02 15:04:05"),
)
}
if err := w.Flush(); err != nil {
return fmt.Errorf("write proxy list: %w", err)
}
if dryRun {
_, _ = fmt.Fprintf(out, "\nDry run: would force-mark %d of %d reverse proxy instance(s) as disconnected.\n", toDisconnect, len(proxies))
return nil
}
if !force {
confirmed, err := confirmDisconnectAll(out, in)
if err != nil {
return err
}
if !confirmed {
_, _ = fmt.Fprintln(out, "Aborted. No reverse proxy instances were changed.")
return nil
}
}
disconnected, err := s.DisconnectAllProxies(ctx)
if err != nil {
return fmt.Errorf("disconnect proxies: %w", err)
}
_, _ = fmt.Fprintf(out, "\nForce-marked %d of %d reverse proxy instance(s) as disconnected.\n", disconnected, len(proxies))
return nil
}
func confirmDisconnectAll(out io.Writer, in io.Reader) (bool, error) {
if in == nil {
in = strings.NewReader("")
}
_, _ = fmt.Fprintln(out, "\nWARNING: This command changes stored reverse proxy state for every non-disconnected instance.")
_, _ = fmt.Fprintln(out, "Run it during a maintenance window; affected live proxies may stay hidden until "+
"their next heartbeat or reconnect/re-register.")
_, _ = fmt.Fprintf(out, "Type %q to continue: ", disconnectAllConfirmation)
scanner := bufio.NewScanner(in)
if !scanner.Scan() {
if err := scanner.Err(); err != nil {
return false, fmt.Errorf("read confirmation: %w", err)
}
return false, nil
}
return strings.EqualFold(strings.TrimSpace(scanner.Text()), disconnectAllConfirmation), nil
}

View File

@@ -0,0 +1,180 @@
package proxycmd
import (
"bytes"
"context"
"strings"
"testing"
"time"
"github.com/spf13/cobra"
"github.com/stretchr/testify/require"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
"github.com/netbirdio/netbird/management/server/store"
)
func newTestStore(t *testing.T) store.Store {
t.Helper()
s, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
require.NoError(t, err)
t.Cleanup(cleanup)
return s
}
func seedProxies(t *testing.T, ctx context.Context, s store.Store) {
t.Helper()
accountID := "account-1"
alreadyDisconnectedAt := time.Now().Add(-time.Hour)
seed := []*rpproxy.Proxy{
{
ID: "proxy-1",
SessionID: "session-1",
ClusterAddress: "cluster-a.example.com",
IPAddress: "10.0.0.1",
LastSeen: time.Now(),
Status: rpproxy.StatusConnected,
},
{
ID: "proxy-2",
SessionID: "session-2",
ClusterAddress: "cluster-b.example.com",
IPAddress: "10.0.0.2",
AccountID: &accountID,
LastSeen: time.Now(),
Status: rpproxy.StatusConnected,
},
{
ID: "proxy-3",
SessionID: "session-3",
ClusterAddress: "cluster-a.example.com",
IPAddress: "10.0.0.3",
LastSeen: time.Now().Add(-time.Hour),
Status: rpproxy.StatusDisconnected,
DisconnectedAt: &alreadyDisconnectedAt,
},
}
for _, p := range seed {
require.NoError(t, s.SaveProxy(ctx, p))
}
}
func proxiesByID(t *testing.T, ctx context.Context, s store.Store) map[string]*rpproxy.Proxy {
t.Helper()
proxies, err := s.GetAllProxies(ctx)
require.NoError(t, err)
require.Len(t, proxies, 3)
byID := make(map[string]*rpproxy.Proxy, len(proxies))
for _, p := range proxies {
byID[p.ID] = p
}
return byID
}
func TestRunDisconnectAllWithConfirmation(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
seedProxies(t, ctx, s)
var out bytes.Buffer
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(disconnectAllConfirmation+"\n"), false, false))
output := out.String()
require.Contains(t, output, "proxy-1")
require.Contains(t, output, "proxy-2")
require.Contains(t, output, "proxy-3")
require.Contains(t, output, "cluster-a.example.com")
require.Contains(t, output, "account-1")
require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
for _, p := range proxiesByID(t, ctx, s) {
require.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", p.ID)
require.NotNil(t, p.DisconnectedAt, "proxy %s should have a disconnected timestamp", p.ID)
}
}
func TestRunDisconnectAllForceSkipsConfirmation(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
seedProxies(t, ctx, s)
var out bytes.Buffer
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, true))
output := out.String()
require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
}
func TestRunDisconnectAllAbortLeavesProxiesUnchanged(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
seedProxies(t, ctx, s)
var out bytes.Buffer
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader("no\n"), false, false))
output := out.String()
require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
require.Contains(t, output, "Aborted. No reverse proxy instances were changed.")
byID := proxiesByID(t, ctx, s)
require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
}
func TestRunDisconnectAllDryRunLeavesProxiesUnchanged(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
seedProxies(t, ctx, s)
var out bytes.Buffer
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), true, false))
output := out.String()
require.Contains(t, output, "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
byID := proxiesByID(t, ctx, s)
require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
}
func TestNewCommandsDisconnectAllDryRun(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
seedProxies(t, ctx, s)
opened := false
cmd := NewCommands(func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
opened = true
return fn(cmd.Context(), s)
})
var out bytes.Buffer
cmd.SetOut(&out)
cmd.SetErr(&out)
cmd.SetIn(strings.NewReader(""))
cmd.SetArgs([]string{"disconnect-all", "--dry-run"})
require.NoError(t, cmd.ExecuteContext(ctx))
require.True(t, opened)
require.Contains(t, out.String(), "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
}
func TestRunDisconnectAllEmpty(t *testing.T) {
ctx := context.Background()
s := newTestStore(t)
var out bytes.Buffer
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, false))
require.Contains(t, out.String(), "No reverse proxy instances found.")
}

View File

@@ -83,7 +83,8 @@ func init() {
rootCmd.AddCommand(migrationCmd)
tc := newTokenCommands()
tc.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
rootCmd.AddCommand(tc)
ac := newAdminCommands()
ac.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
rootCmd.AddCommand(ac)
rootCmd.AddCommand(newLegacyTokenCommand())
}

View File

@@ -1,55 +0,0 @@
package cmd
import (
"context"
"fmt"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
"github.com/netbirdio/netbird/formatter/hook"
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/util"
)
var tokenDatadir string
// newTokenCommands creates the token command tree with management-specific store opener.
func newTokenCommands() *cobra.Command {
cmd := tokencmd.NewCommands(withTokenStore)
cmd.PersistentFlags().StringVar(&tokenDatadir, "datadir", "", "Override the data directory from config (where store.db is located)")
return cmd
}
// withTokenStore initializes logging, loads config, opens the store, and calls fn.
func withTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
if err := util.InitLog("error", "console"); err != nil {
return fmt.Errorf("init log: %w", err)
}
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
config, err := LoadMgmtConfig(ctx, nbconfig.MgmtConfigPath)
if err != nil {
return fmt.Errorf("load config: %w", err)
}
datadir := config.Datadir
if tokenDatadir != "" {
datadir = tokenDatadir
}
s, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
if err != nil {
return fmt.Errorf("create store: %w", err)
}
defer func() {
if err := s.Close(ctx); err != nil {
log.Debugf("close store: %v", err)
}
}()
return fn(ctx, s)
}

View File

@@ -29,6 +29,7 @@ import (
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/management/server/types"
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/status"
"github.com/netbirdio/netbird/util"
@@ -56,6 +57,10 @@ type Controller struct {
proxyController port_forwarding.Controller
integratedPeerValidator integrated_validator.IntegratedValidator
serverSupportedSyncMessageVersion sharedgrpc.SyncMessageVersion
perAccountServerSupportedSyncMessageVersions map[string]sharedgrpc.SyncMessageVersion
}
type bufferUpdate struct {
@@ -90,8 +95,10 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
dnsDomain: dnsDomain,
config: config,
proxyController: proxyController,
EphemeralPeersManager: ephemeralPeersManager,
proxyController: proxyController,
EphemeralPeersManager: ephemeralPeersManager,
serverSupportedSyncMessageVersion: sharedgrpc.SyncMessageVersionFromConfig(config.HighestSupportedSyncMessageVersion),
perAccountServerSupportedSyncMessageVersions: sharedgrpc.SyncMessageVersionsFromMap(config.PerAccountHighestSupportedSyncMessageVersion),
}
}
@@ -222,18 +229,53 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
c.metrics.CountCalcPostureChecksDuration(time.Since(start))
start = time.Now()
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
peerGroups := account.GetPeerGroups(p.ID)
proxyNetworkMap := proxyNetworkMaps[p.ID]
var update *proto.SyncResponse
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
log.WithContext(ctx).
WithFields(log.Fields{
"sync_message_version": commonSyncMessageVersion,
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
}).Debug("common highest sync message version")
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
components := account.GetPeerNetworkMapComponents(
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
start = time.Now()
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
// the client merges it into Calculate()'s output the same
// way the legacy server did via NetworkMap.Merge.
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
Update: update,
MessageType: network_map.MessageTypeNetworkMap,
})
return
}
nmap := account.GetPeerNetworkMapFromComponents(
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
if ok {
remotePeerNetworkMap.Merge(proxyNetworkMap)
if proxyNetworkMap != nil {
nmap.Merge(proxyNetworkMap)
}
peerGroups := account.GetPeerGroups(p.ID)
start = time.Now()
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
c.metrics.CountToSyncResponseDuration(time.Since(start))
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -251,6 +293,13 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
return nil
}
func (c *Controller) perAccountOrGlobalSupportedSyncMessageVersions(accountId string) sharedgrpc.SyncMessageVersion {
if perAccount, ok := c.perAccountServerSupportedSyncMessageVersions[accountId]; ok {
return perAccount
}
return c.serverSupportedSyncMessageVersion
}
// UpdatePeers updates all peers that belong to an account.
// Should be called when changes have to be synced to peers.
func (c *Controller) UpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error {
@@ -352,18 +401,53 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
c.metrics.CountCalcPostureChecksDuration(time.Since(start))
start = time.Now()
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
peerGroups := account.GetPeerGroups(p.ID)
proxyNetworkMap := proxyNetworkMaps[p.ID]
var update *proto.SyncResponse
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
log.WithContext(ctx).
WithFields(log.Fields{
"sync_message_version": commonSyncMessageVersion,
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
}).Debug("common highest sync message version")
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
components := account.GetPeerNetworkMapComponents(
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
start = time.Now()
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
// the client merges it into Calculate()'s output the same
// way the legacy server did via NetworkMap.Merge.
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
Update: update,
MessageType: network_map.MessageTypeNetworkMap,
})
return
}
nmap := account.GetPeerNetworkMapFromComponents(
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
if ok {
remotePeerNetworkMap.Merge(proxyNetworkMap)
if proxyNetworkMap != nil {
nmap.Merge(proxyNetworkMap)
}
peerGroups := account.GetPeerGroups(p.ID)
start = time.Now()
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
c.metrics.CountToSyncResponseDuration(time.Since(start))
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -451,13 +535,7 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
return err
}
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, peerId, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
proxyNetworkMap, ok := proxyNetworkMaps[peer.ID]
if ok {
remotePeerNetworkMap.Merge(proxyNetworkMap)
}
proxyNetworkMap := proxyNetworkMaps[peer.ID]
extraSettings, err := c.settingsManager.GetExtraSettings(ctx, peer.AccountID)
if err != nil {
return fmt.Errorf("failed to get extra settings: %v", err)
@@ -466,7 +544,45 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
peerGroups := account.GetPeerGroups(peerId)
dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
var update *proto.SyncResponse
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
c.perAccountOrGlobalSupportedSyncMessageVersions(accountId),
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
log.WithContext(ctx).
WithFields(log.Fields{
"sync_message_version": commonSyncMessageVersion,
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
}).Debug("common highest sync message version")
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
components := account.GetPeerNetworkMapComponents(
ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
// the client merges it into Calculate()'s output the same
// way the legacy server did via NetworkMap.Merge.
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
Update: update,
MessageType: network_map.MessageTypeNetworkMap,
})
return nil
}
nmap := account.GetPeerNetworkMapFromComponents(
ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
if proxyNetworkMap != nil {
nmap.Merge(proxyNetworkMap)
}
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
Update: update,
MessageType: network_map.MessageTypeNetworkMap,
@@ -513,6 +629,65 @@ func (c *Controller) BufferUpdateAccountPeers(ctx context.Context, accountID str
return nil
}
// GetValidatedPeerWithComponents is the components-format counterpart of
// GetValidatedPeerWithMap. It returns raw NetworkMapComponents for capable
// peers along with the proxy NetworkMap fragment (BYOP / port-forwarding
// data the legacy server folds in via NetworkMap.Merge). The gRPC layer
// encodes both into the wire envelope. Callers must gate on capability
// themselves before dispatching here — this method does NOT branch on it.
func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
if isRequiresApproval {
network, err := c.repo.GetAccountNetwork(ctx, accountID)
if err != nil {
return nil, nil, nil, nil, 0, err
}
return peer, &types.NetworkMapComponents{Network: network.Copy()}, nil, nil, 0, nil
}
account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
if err != nil {
return nil, nil, nil, nil, 0, err
}
c.injectAllProxyPolicies(ctx, account)
approvedPeersMap, err := c.integratedPeerValidator.GetValidatedPeers(ctx, account.Id, maps.Values(account.Groups), maps.Values(account.Peers), account.Settings.Extra)
if err != nil {
return nil, nil, nil, nil, 0, err
}
postureChecks, err := c.getPeerPostureChecks(account, peer.ID)
if err != nil {
return nil, nil, nil, nil, 0, err
}
accountZones, err := c.repo.GetAccountZones(ctx, account.Id)
if err != nil {
return nil, nil, nil, nil, 0, err
}
// Fetch the proxy network map fragment for this peer alongside the
// components — same single-account-load path the streaming controller
// uses, so initial-sync delivers BYOP/forwarding patches synchronously
// instead of waiting for the next streaming push.
proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peer.ID, account.Peers)
if err != nil {
log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
return nil, nil, nil, nil, 0, err
}
dnsDomain := c.GetDNSDomain(account.Settings)
peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
components := account.GetPeerNetworkMapComponents(ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
return peer, components, proxyNetworkMaps[peer.ID], postureChecks, dnsFwdPort, nil
}
// BufferUpdateAffectedPeers accumulates peer IDs and flushes them after the buffer interval.
func (c *Controller) BufferUpdateAffectedPeers(ctx context.Context, accountID string, peerIDs []string, reason types.UpdateReason) error {
if len(peerIDs) == 0 {

View File

@@ -24,6 +24,7 @@ type Controller interface {
UpdateAccountPeer(ctx context.Context, accountId string, peerId string) error
BufferUpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error
GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error)
GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error)
GetDNSDomain(settings *types.Settings) string
StartWarmup(context.Context)
GetNetworkMap(ctx context.Context, peerID string) (*types.NetworkMap, error)

View File

@@ -1,9 +1,9 @@
// Code generated by MockGen. DO NOT EDIT.
// Source: management/internals/controllers/network_map/interface.go
// Source: ./interface.go
//
// Generated by this command:
//
// mockgen -package network_map -destination=management/internals/controllers/network_map/interface_mock.go -source=management/internals/controllers/network_map/interface.go -build_flags=-mod=mod
// mockgen -package network_map -destination=interface_mock.go -source=./interface.go -build_flags=-mod=mod
//
// Package network_map is a generated GoMock package.
@@ -126,8 +126,27 @@ func (mr *MockControllerMockRecorder) GetNetworkMap(ctx, peerID any) *gomock.Cal
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNetworkMap", reflect.TypeOf((*MockController)(nil).GetNetworkMap), ctx, peerID)
}
// GetValidatedPeerWithComponents mocks base method.
func (m *MockController) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *peer.Peer) (*peer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetValidatedPeerWithComponents", ctx, isRequiresApproval, accountID, p)
ret0, _ := ret[0].(*peer.Peer)
ret1, _ := ret[1].(*types.NetworkMapComponents)
ret2, _ := ret[2].(*types.NetworkMap)
ret3, _ := ret[3].([]*posture.Checks)
ret4, _ := ret[4].(int64)
ret5, _ := ret[5].(error)
return ret0, ret1, ret2, ret3, ret4, ret5
}
// GetValidatedPeerWithComponents indicates an expected call of GetValidatedPeerWithComponents.
func (mr *MockControllerMockRecorder) GetValidatedPeerWithComponents(ctx, isRequiresApproval, accountID, p any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetValidatedPeerWithComponents", reflect.TypeOf((*MockController)(nil).GetValidatedPeerWithComponents), ctx, isRequiresApproval, accountID, p)
}
// GetValidatedPeerWithMap mocks base method.
func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetValidatedPeerWithMap", ctx, isRequiresApproval, accountID, peerID)
ret0, _ := ret[0].(*types.NetworkMap)
@@ -171,7 +190,7 @@ func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID
}
// OnPeersAdded mocks base method.
func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "OnPeersAdded", ctx, accountID, peerIDs, affectedPeerIDs)
ret0, _ := ret[0].(error)
@@ -185,7 +204,7 @@ func (mr *MockControllerMockRecorder) OnPeersAdded(ctx, accountID, peerIDs, affe
}
// OnPeersDeleted mocks base method.
func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "OnPeersDeleted", ctx, accountID, peerIDs, affectedPeerIDs)
ret0, _ := ret[0].(error)
@@ -199,7 +218,7 @@ func (mr *MockControllerMockRecorder) OnPeersDeleted(ctx, accountID, peerIDs, af
}
// OnPeersUpdated mocks base method.
func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs []string, affectedPeerIDs []string) error {
func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs, affectedPeerIDs []string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "OnPeersUpdated", ctx, accountId, peerIDs, affectedPeerIDs)
ret0, _ := ret[0].(error)

View File

@@ -18,21 +18,26 @@ import (
// contract between the proxy and management; management flattens them into
// queryable columns. Keep in sync with the proxy side.
const (
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDTotal = "cost.usd_total"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCachedInputTokens = "llm.cached_input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDInput = "cost.usd_input"
metaKeyCostUSDCachedInput = "cost.usd_cached_input"
metaKeyCostUSDCacheCreate = "cost.usd_cache_creation"
metaKeyCostUSDOutput = "cost.usd_output"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
)
// IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -108,20 +113,25 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
BytesUpload: e.BytesUpload,
BytesDownload: e.BytesDownload,
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CachedInputTokens: parseMetaInt(meta, metaKeyCachedInputTokens),
CacheCreationTokens: parseMetaInt(meta, metaKeyCacheCreationTokens),
InputCostUSD: parseMetaFloat(meta, metaKeyCostUSDInput),
CachedInputCostUSD: parseMetaFloat(meta, metaKeyCostUSDCachedInput),
CacheCreationCostUSD: parseMetaFloat(meta, metaKeyCostUSDCacheCreate),
OutputCostUSD: parseMetaFloat(meta, metaKeyCostUSDOutput),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
}
var groups []types.AgentNetworkAccessLogGroup
@@ -140,18 +150,23 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
// log's ID so the two correlate.
func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
usage := &types.AgentNetworkUsage{
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CostUSD: e.CostUSD,
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CachedInputTokens: e.CachedInputTokens,
CacheCreationTokens: e.CacheCreationTokens,
InputCostUSD: e.InputCostUSD,
CachedInputCostUSD: e.CachedInputCostUSD,
CacheCreationCostUSD: e.CacheCreationCostUSD,
OutputCostUSD: e.OutputCostUSD,
}
usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))

View File

@@ -28,17 +28,22 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
UserId: "user-1",
AgentNetwork: true,
Metadata: map[string]string{
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "150",
metaKeyCostUSDTotal: "0.0123",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "1174",
metaKeyCachedInputTokens: "256",
metaKeyCacheCreationTokens: "768",
metaKeyCostUSDInput: "0.0071",
metaKeyCostUSDCachedInput: "0.0009",
metaKeyCostUSDCacheCreate: "0.0020",
metaKeyCostUSDOutput: "0.0023",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
// repeated id must be de-duplicated before the group rows insert.
metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
},
@@ -65,7 +70,19 @@ func TestIngestAccessLog_RealStore_LogCollectionOff(t *testing.T) {
require.Len(t, usage, 1, "usage row must be written even with log collection off")
assert.Equal(t, int64(100), usage[0].InputTokens, "input tokens must round-trip from metadata")
assert.Equal(t, int64(50), usage[0].OutputTokens, "output tokens must round-trip from metadata")
assert.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost must round-trip from metadata")
assert.Equal(t, int64(256), usage[0].CachedInputTokens, "cache-read tokens must round-trip from metadata")
assert.Equal(t, int64(768), usage[0].CacheCreationTokens, "cache-write tokens must round-trip from metadata")
// The per-bucket breakdown is the only cost state stored, and must survive
// the write/read cycle as real columns — usage rows are the only cost
// record for accounts with log collection off, so a dropped column here
// loses the split permanently.
assert.InDelta(t, 0.0071, usage[0].InputCostUSD, 1e-9, "input cost must round-trip from metadata")
assert.InDelta(t, 0.0009, usage[0].CachedInputCostUSD, 1e-9, "cache-read cost must round-trip from metadata")
assert.InDelta(t, 0.0020, usage[0].CacheCreationCostUSD, 1e-9, "cache-write cost must round-trip from metadata")
assert.InDelta(t, 0.0023, usage[0].OutputCostUSD, 1e-9, "output cost must round-trip from metadata")
// Aggregates are derived from the stored columns, never stored themselves.
assert.InDelta(t, 0.0123, usage[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0029, usage[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
require.NoError(t, err)
@@ -96,6 +113,14 @@ func TestIngestAccessLog_RealStore_LogCollectionOn(t *testing.T) {
require.Equal(t, int64(1), total, "exactly one access-log row expected")
require.Len(t, logs, 1, "full access-log row must be written when log collection is on")
assert.Equal(t, "gpt-5.4", logs[0].Model, "model must flatten from metadata")
assert.Equal(t, int64(256), logs[0].CachedInputTokens, "cache-read tokens must flatten from metadata")
assert.Equal(t, int64(768), logs[0].CacheCreationTokens, "cache-write tokens must flatten from metadata")
assert.InDelta(t, 0.0029, logs[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
assert.InDelta(t, 0.0123, logs[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0071, logs[0].InputCostUSD, 1e-9, "input cost must flatten from metadata")
assert.InDelta(t, 0.0009, logs[0].CachedInputCostUSD, 1e-9, "cache-read cost must flatten from metadata")
assert.InDelta(t, 0.0020, logs[0].CacheCreationCostUSD, 1e-9, "cache-write cost must flatten from metadata")
assert.InDelta(t, 0.0023, logs[0].OutputCostUSD, 1e-9, "output cost must flatten from metadata")
assert.Equal(t, "hello", logs[0].RequestPrompt, "prompt must be retained when log collection is on")
assert.Equal(t, "world", logs[0].ResponseCompletion, "completion must be retained when log collection is on")
assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")

View File

@@ -38,7 +38,7 @@ func accessLogRow(id, sessionID string, ts time.Time, opts ...func(*types.AgentN
InputTokens: 100,
OutputTokens: 50,
TotalTokens: 150,
CostUSD: 0.01,
InputCostUSD: 0.01,
}
for _, o := range opts {
o(e)
@@ -74,7 +74,7 @@ func withTokens(in, out, total int64, cost float64) func(*types.AgentNetworkAcce
e.InputTokens = in
e.OutputTokens = out
e.TotalTokens = total
e.CostUSD = cost
e.InputCostUSD = cost
}
}
@@ -155,7 +155,7 @@ func TestAccessLogSessions_FoldAndAggregate(t *testing.T) {
assert.Equal(t, int64(310), a.InputTokens, "input tokens summed")
assert.Equal(t, int64(135), a.OutputTokens, "output tokens summed")
assert.Equal(t, int64(445), a.TotalTokens, "total tokens summed")
assert.InDelta(t, 0.031, a.CostUSD, 1e-9, "cost summed")
assert.InDelta(t, 0.031, a.TotalCostUSD(), 1e-9, "cost summed")
assert.Equal(t, "deny", a.Decision, "any deny makes the session a deny")
assert.ElementsMatch(t, []string{"openai", "anthropic"}, a.Providers, "distinct providers")
assert.ElementsMatch(t, []string{"gpt-5.4", "claude-haiku-4-5"}, a.Models, "distinct models")

View File

@@ -420,6 +420,47 @@ var providers = []Provider{
{ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192},
},
},
{
ID: "kimi_api",
Kind: KindProvider,
Name: "Kimi (Moonshot AI) API",
Description: "Kimi K3 / K2 models via the Moonshot AI platform",
DefaultHost: "api.moonshot.ai",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#1A1A2E",
// ParserID empty on purpose: Moonshot serves two body shapes on
// the same host and key, and the proxy's URL sniffer dispatches
// both (same pattern as Bifrost). /v1/chat/completions matches
// OpenAIParser; the Anthropic-compatible endpoint the official
// Claude Code guide uses (/anthropic/v1/messages) contains
// "/v1/messages" and matches AnthropicParser. Pinning "openai"
// here would misparse the Claude Code path — the primary way
// teams consume Kimi for coding today. Both endpoints accept the
// same Moonshot key via Authorization: Bearer (Claude Code's
// ANTHROPIC_AUTH_TOKEN rides that header too).
//
// api.moonshot.ai is the international platform; mainland-China
// accounts live on api.moonshot.cn with separate billing —
// operators there override the host on the provider record. The
// kimi.com subscription coding endpoint (api.kimi.com/coding,
// model id "k3") is account-bound seat licensing rather than a
// meterable platform key, so it's deliberately not the default.
ParserID: "",
// Pricing per Moonshot's platform rates at K3 launch (July 2026):
// $3/$15 per MTok with $0.30 cached input, flat across the 1M-token
// window. kimi-k3 is the ONLY model the platform serves newer
// accounts — K2-era ids (kimi-k2-thinking) and even the kimi-latest
// alias return resource_not_found_error, verified live 2026-07-21 —
// so it's the only catalog entry. Grandfathered accounts with K2
// access can still type those ids on the provider's model rows.
// The consumer app's "K3 Swarm Max" mode is not an API SKU, so it
// doesn't appear here.
Models: []Model{
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
},
},
{
ID: "litellm_proxy",
Kind: KindGateway,

View File

@@ -85,12 +85,17 @@ type Manager interface {
// PolicySelectionInput is the per-request selection envelope. The
// proxy populates it from CapturedData (account, user, groups) plus
// the provider llm_router resolved.
// the provider llm_router resolved and the model it extracted.
type PolicySelectionInput struct {
AccountID string
UserID string
GroupIDs []string
ProviderID string
// Model is the already-normalised upstream model id the proxy extracted
// (parser strips Bedrock region/version, Vertex @version), so a
// case-insensitive compare suffices. Empty = undetermined → not permitted
// (fail closed).
Model string
}
// PolicySelectionResult names the policy that "pays" for this request

View File

@@ -5,6 +5,7 @@ import (
"fmt"
"math"
"sort"
"strings"
"time"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
@@ -35,6 +36,10 @@ const (
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
//nolint:gosec // account deny code label, not a credential
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
// denyCodeModelBlocked is returned when policies govern the request's
// (provider, caller-groups) but none permits the model. Matches the proxy
// guardrail's code so both layers surface the same label.
denyCodeModelBlocked = "llm_policy.model_blocked"
)
// consumptionCache holds the consumption counters prefetched for one
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
}
candidates := filterApplicablePolicies(policies, in)
// Model-allowlist gate scoped to the matched policies: keep candidates whose
// guardrails permit the model (none enabled = unrestricted), deny when
// policies apply but none permits it. Skip the load when none has a guardrail.
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
if gErr != nil {
return nil, gErr
}
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
if len(permitted) == 0 {
return &PolicySelectionResult{
Allow: false,
DenyCode: denyCodeModelBlocked,
DenyReason: modelBlockedReason(in.Model),
}, nil
}
candidates = permitted
}
// Prefetch every consumption counter the ceiling + candidate policies will
// read, in a single store round-trip, then score against the cache.
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
return out
}
// anyPolicyHasGuardrails reports whether any policy references at least one
// guardrail, so the selector can skip loading guardrails when none do.
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
for _, p := range policies {
if p != nil && len(p.GuardrailIDs) > 0 {
return true
}
}
return false
}
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
// model-allowlist gate to resolve each candidate policy's attached guardrails.
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return nil, fmt.Errorf("list account guardrails: %w", err)
}
byID := make(map[string]*types.Guardrail, len(guardrails))
for _, g := range guardrails {
if g != nil {
byID[g.ID] = g
}
}
return byID, nil
}
// filterModelPermittedPolicies returns the subset of policies whose guardrails
// permit the model. Order is preserved so downstream scoring is unaffected.
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
out := make([]*types.Policy, 0, len(policies))
for _, p := range policies {
if policyPermitsModel(p, byID, model) {
out = append(out, p)
}
}
return out
}
// policyPermitsModel reports whether a policy permits the model. No
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
// otherwise the model must be in the union of its allowlists, so an
// empty/undetermined model fails closed.
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
if p == nil {
return false
}
wanted := normaliseModelID(model)
restricted := false
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
if wanted == "" {
continue
}
for _, allowed := range g.Checks.ModelAllowlist.Models {
if normaliseModelID(allowed) == wanted {
return true
}
}
}
return !restricted
}
// normaliseModelID lowercases and trims a model identifier so the allowlist
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
// normaliseModel so both layers agree on what "same model" means.
func normaliseModelID(model string) string {
return strings.ToLower(strings.TrimSpace(model))
}
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
// rejection. The model is quoted when known; an undetermined model is reported
// as such so the access log distinguishes "wrong model" from "no model".
func modelBlockedReason(model string) string {
if normaliseModelID(model) == "" {
return "request model could not be determined for the policy allowlist"
}
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
}
// candidate is the per-policy intermediate the selector ranks. A
// policy that's been exhausted on any enabled cap never makes it
// into this slice; the selector's deny envelope carries the latest

View File

@@ -0,0 +1,329 @@
package agentnetwork
import (
"context"
"errors"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/store"
)
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
// to reach providerID under the given guardrails. Uncapped keeps the selector's
// headroom scoring trivial so these tests isolate the model-allowlist gate.
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
AccountID: account,
Enabled: true,
SourceGroups: sourceGroups,
DestinationProviderIDs: []string{providerID},
GuardrailIDs: guardrailIDs,
CreatedAt: time.Now().UTC(),
}
}
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
// carries the given models.
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
return &types.Guardrail{
ID: id,
AccountID: account,
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
},
}
}
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
mockStore.EXPECT().
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
Return(policies, nil)
}
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
Return(guardrails, nil)
}
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
// decision: a policy authorises the (provider, group) but restricts the model,
// and the requested model isn't on the list, so the request is denied.
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
}
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
// is on the applicable policy's allowlist, so selection proceeds normally.
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
// and surrounding whitespace, matching the proxy guardrail's normalisation.
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.NoError(t, err)
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
}
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
// policies authorise the same (provider, group) and one has no guardrail, that
// policy makes the request unrestricted — not caught by the other's allowlist.
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
expectPolicies(mockStore, "acc-1", restricted, open)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
}
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
// model allowlisted only for grp-b must not be usable by a grp-a caller. The
// selector considers only policies applicable to the caller's groups.
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
expectPolicies(mockStore, "acc-1", polA, polB)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-a"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only allowed for grp-b
})
require.NoError(t, err)
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
// mirrors the proxy: with a restricted applicable policy and an empty model
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "", // undetermined
})
require.NoError(t, err)
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
// model allowlist is disabled imposes no model restriction, even though the
// policy references it.
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
disabled := &types.Guardrail{
ID: "g-1",
AccountID: "acc-1",
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
},
}
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", disabled)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
// allowlist guardrails permits the union of their models (not just the first).
func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only in the second guardrail's list
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
// resolving the candidate policies' guardrails surfaces as an error, not a
// silent allow/deny.
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
Return(nil, errors.New("store unavailable"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
assert.Nil(t, res)
}
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
// policy referencing a guardrail ID absent from the account's set (a stale
// reference) imposes no model restriction — same as no guardrail.
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1")
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
// gate narrows candidates before cap scoring: the permitting policy is selected
// even though the blocked one has a larger, more attractive cap.
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
polBig.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
}
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
polSmall.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
}
expectPolicies(mockStore, "acc-1", polBig, polSmall)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only pol-small's guardrail permits this
})
require.NoError(t, err)
assert.True(t, res.Allow)
assert.Equal(t, "pol-small", res.SelectedPolicyID,
"the model filter must exclude pol-big before cap scoring")
}

View File

@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
applyAccountCollectionControls(&mergedGuardrails, settings)
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
// The proxy guardrail is a per-provider fail-closed backstop; the
// authoritative per-policy/group decision is management's
// SelectPolicyForRequest. A provider lands in this map only when every
// authorising policy restricts models.
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
if err != nil {
return nil, err
}
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
// middleware expects. Mirrors the proxy registration documented in
// the management→proxy contract.
// the management→proxy contract. provider_allowlists is keyed by the
// resolved provider id llm_router stamps; a provider absent from the map is
// unrestricted at the proxy layer.
type guardrailConfig struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
}
type guardrailPromptCapture struct {
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
}
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
cfg := guardrailConfig{
ModelAllowlist: merged.ModelAllowlist,
PromptCapture: guardrailPromptCapture{
Enabled: merged.PromptCapture.Enabled,
RedactPii: merged.PromptCapture.RedactPii,
},
ProviderAllowlists: providerAllowlists,
PromptCapture: guardrailPromptCapture(capture),
}
out, err := json.Marshal(cfg)
if err != nil {
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
return out, nil
}
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
// is included only when every authorising policy restricts models (their union);
// if any leaves it unrestricted it is omitted, so management decides per group.
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
type providerAcc struct {
models map[string]struct{}
anyUnrestricted bool
}
accs := make(map[string]*providerAcc)
for _, p := range policies {
if p == nil {
continue
}
restricted, models := policyModelAllowlist(p, byID)
for _, providerID := range p.DestinationProviderIDs {
if providerID == "" {
continue
}
acc, ok := accs[providerID]
if !ok {
acc = &providerAcc{models: make(map[string]struct{})}
accs[providerID] = acc
}
if !restricted {
acc.anyUnrestricted = true
continue
}
for _, m := range models {
acc.models[m] = struct{}{}
}
}
}
out := make(map[string][]string, len(accs))
for providerID, acc := range accs {
if acc.anyUnrestricted {
continue
}
models := make([]string, 0, len(acc.models))
for m := range acc.models {
models = append(models, m)
}
sort.Strings(models)
out[providerID] = models
}
return out
}
// policyModelAllowlist reports whether a policy restricts models (has an
// allowlist-enabled guardrail) and the union of allowed models. Models are
// verbatim; the proxy factory lowercases/trims them at decode time.
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
restricted := false
var models []string
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
models = append(models, m)
}
}
}
return restricted, models
}
// buildAccountService composes the per-account gateway Service. The
// target carries the noop placeholder URL — the router middleware
// rewrites every request to the matched provider's upstream before the
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
return out
}
// MergedGuardrails is the JSON shape passed to the proxy via the
// guardrail middleware's config_json. Mirrors the proxy-side
// expectations and is intentionally distinct from
// types.GuardrailChecks so we can evolve either side independently.
// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
// merged here — the model allowlist is emitted per-provider, and
// token/budget/retention moved onto Policy.Limits and account Settings.
type MergedGuardrails struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
TokenLimits MergedTokenLimits `json:"token_limits"`
Budget MergedBudget `json:"budget"`
PromptCapture MergedPromptCapture `json:"prompt_capture"`
Retention MergedRetention `json:"retention"`
}
type MergedTokenLimits struct {
Hourly *MergedTokenWindow `json:"hourly,omitempty"`
Daily *MergedTokenWindow `json:"daily,omitempty"`
Monthly *MergedTokenWindow `json:"monthly,omitempty"`
}
type MergedTokenWindow struct {
MaxInputTokens int `json:"max_input_tokens,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
}
type MergedBudget struct {
Hourly *MergedBudgetWindow `json:"hourly,omitempty"`
Daily *MergedBudgetWindow `json:"daily,omitempty"`
Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
}
type MergedBudgetWindow struct {
SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
PromptCapture MergedPromptCapture
}
type MergedPromptCapture struct {
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
RedactPii bool `json:"redact_pii"`
}
type MergedRetention struct {
Enabled bool `json:"enabled"`
Days int `json:"days"`
}
// mergeGuardrails computes the effective guardrail spec applied at the
// proxy, given the referencing policies and the account's guardrail
// catalogue. Policy enabled-ness is the caller's responsibility — only
// enabled policies should be passed in.
// mergeGuardrails folds the referencing policies' guardrails into the
// prompt-capture decision only. The model allowlist is enforced per-policy/group
// in management and shipped per-provider; token/budget/retention live off
// guardrails now.
//
// Merge rules:
// - Model allowlist: union of allowlists across policies that enable it.
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
// if any enabling policy turns it on.
// - Retention: enabled if any enables it; smallest non-zero days wins.
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
// sticks if any enabling policy turns it on.
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
merged := MergedGuardrails{}
allowlist := make(map[string]struct{})
allowlistEnabled := false
for _, policy := range policies {
for _, gID := range policy.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil {
continue
}
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
mergeGuardrail(g, &merged)
}
}
if allowlistEnabled {
merged.ModelAllowlist = make([]string, 0, len(allowlist))
for m := range allowlist {
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
}
sort.Strings(merged.ModelAllowlist)
}
return merged
}
// mergeGuardrail folds a single guardrail's enabled checks into the
// running merge: model-allowlist models join the shared set (and flip
// allowlistEnabled), and prompt-capture / redact-pii stick once any
// enabling guardrail turns them on.
//
// TokenLimits, Budget, and Retention have moved off guardrails — token
// and budget caps now live on the Policy itself (Policy.Limits) and
// retention moves to account-level Settings — so they are not merged here.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
if g.Checks.ModelAllowlist.Enabled {
*allowlistEnabled = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
allowlist[m] = struct{}{}
}
}
}
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
// running merge: enabled / redact-pii stick once any enabling guardrail turns
// them on.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
if g.Checks.PromptCapture.Enabled {
merged.PromptCapture.Enabled = true
if g.Checks.PromptCapture.RedactPii {

View File

@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
require.Len(t, services, 1)
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
"model allowlist is a pure policy guardrail and must always reach the config")
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
"model allowlist is a pure policy guardrail and must reach the per-provider config")
assert.False(t, cfg.PromptCapture.Enabled,
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
}
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
require.Len(t, services, 1, "exactly one synth service expected")
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
}

View File

@@ -0,0 +1,95 @@
package agentnetwork
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// policyForProviders builds an enabled policy authorising the given providers
// under the given guardrails (both optional). Groups are irrelevant to
// buildProviderAllowlists, which keys purely on destination provider.
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
Enabled: true,
DestinationProviderIDs: providerIDs,
GuardrailIDs: guardrailIDs,
}
}
func TestBuildProviderAllowlists(t *testing.T) {
byID := map[string]*types.Guardrail{
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
}
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
"a provider every policy restricts carries the sorted union of their models")
})
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", nil, "prov-x"), // no guardrail
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
})
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a policy whose only guardrail has a disabled allowlist is unrestricted")
})
t.Run("providers are isolated from one another", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
})
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
})
t.Run("union across a single policy's guardrails", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
"a policy's own multiple allowlist guardrails union together")
})
t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
got := buildProviderAllowlists([]*types.Policy{
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
}, empty)
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
})
}

View File

@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
var cfg guardrailConfig
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
"model allowlist union must keep both models")
// Both policies restrict the same provider, so the per-provider backstop
// carries the union of their models — a coarse gate that management's
// per-policy/group check narrows; it only blocks models outside the union
// when management is down.
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
"per-provider allowlist union must keep both models")
}
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {

View File

@@ -41,8 +41,24 @@ type AgentNetworkAccessLog struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
Stream bool
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown — one column per token bucket the provider
// bills separately. These four are the only cost state stored: the total
// and the cache portion are derived on read (TotalCostUSD / CacheCostUSD)
// rather than stored alongside, so a stored aggregate can never drift out
// of step with the components it summarises.
//
// default:0 matters on upgrade: these columns are ALTER TABLE ADD COLUMN
// on an existing table, and without it every historical row holds NULL —
// which a raw SUM()/scan into float64 can't read. The default backfills
// them as 0, so pre-upgrade rows report an unknown split, not an error.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
Stream bool
// Prompt capture. Only populated when prompt collection is enabled
// (account master switch AND policy guardrail). Heavy free text.
@@ -60,19 +76,44 @@ type AgentNetworkAccessLog struct {
// the reverse-proxy AccessLogEntry table.
func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_log" }
// CostUSDSQLExpr is the SQL sum of the per-bucket cost columns — the total cost
// of a row. Used wherever a query has to sort or aggregate on total cost now
// that no cost_usd column is stored. Plain arithmetic over NOT NULL columns, so
// it stays portable across SQLite and Postgres.
const CostUSDSQLExpr = "(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd)"
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (a *AgentNetworkAccessLog) TotalCostUSD() float64 {
return a.InputCostUSD + a.CachedInputCostUSD + a.CacheCreationCostUSD + a.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets:
// cache reads plus cache writes.
func (a *AgentNetworkAccessLog) CacheCostUSD() float64 {
return a.CachedInputCostUSD + a.CacheCreationCostUSD
}
// ToAPIResponse renders the flattened entry as the API representation.
func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
out := api.AgentNetworkAccessLog{
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CostUsd: a.CostUSD,
Stream: &a.Stream,
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CachedInputTokens: a.CachedInputTokens,
CacheCreationTokens: a.CacheCreationTokens,
InputCostUsd: a.InputCostUSD,
CachedInputCostUsd: a.CachedInputCostUSD,
CacheCreationCostUsd: a.CacheCreationCostUSD,
OutputCostUsd: a.OutputCostUSD,
CostUsd: a.TotalCostUSD(),
CacheCostUsd: a.CacheCostUSD(),
Stream: &a.Stream,
}
out.UserId = strPtr(a.UserID)
@@ -112,20 +153,36 @@ func strPtr(s string) *string {
// summary plus its ordered entries. Assembled in Go from a page of entries — it
// is not a stored table.
type AgentNetworkAccessLogSession struct {
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
}
// TotalCostUSD is the session's total cost: the sum of the four per-bucket
// costs accumulated across its entries.
func (sess *AgentNetworkAccessLogSession) TotalCostUSD() float64 {
return sess.InputCostUSD + sess.CachedInputCostUSD + sess.CacheCreationCostUSD + sess.OutputCostUSD
}
// CacheCostUSD is the session's prompt-cache spend: cache reads plus writes.
func (sess *AgentNetworkAccessLogSession) CacheCostUSD() float64 {
return sess.CachedInputCostUSD + sess.CacheCreationCostUSD
}
// sessionKey is the grouping key for an entry: its session id, or — when the
@@ -205,7 +262,12 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
sess.InputTokens += e.InputTokens
sess.OutputTokens += e.OutputTokens
sess.TotalTokens += e.TotalTokens
sess.CostUSD += e.CostUSD
sess.CachedInputTokens += e.CachedInputTokens
sess.CacheCreationTokens += e.CacheCreationTokens
sess.InputCostUSD += e.InputCostUSD
sess.CachedInputCostUSD += e.CachedInputCostUSD
sess.CacheCreationCostUSD += e.CacheCreationCostUSD
sess.OutputCostUSD += e.OutputCostUSD
if e.Timestamp.Before(sess.StartedAt) {
sess.StartedAt = e.Timestamp
}
@@ -248,15 +310,22 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
}
out := api.AgentNetworkAccessLogSession{
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CostUsd: sess.CostUSD,
Decision: sess.Decision,
Entries: entries,
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CachedInputTokens: sess.CachedInputTokens,
CacheCreationTokens: sess.CacheCreationTokens,
InputCostUsd: sess.InputCostUSD,
CachedInputCostUsd: sess.CachedInputCostUSD,
CacheCreationCostUsd: sess.CacheCreationCostUSD,
OutputCostUsd: sess.OutputCostUSD,
CostUsd: sess.TotalCostUSD(),
CacheCostUsd: sess.CacheCostUSD(),
Decision: sess.Decision,
Entries: entries,
}
out.SessionId = strPtr(sess.SessionID)
out.UserId = strPtr(sess.UserID)

View File

@@ -54,7 +54,7 @@ var accessLogSortFields = map[string]string{
"provider": "provider",
"status_code": "status_code",
"duration": "duration",
"cost_usd": "cost_usd",
"cost_usd": CostUSDSQLExpr,
"total_tokens": "total_tokens",
"user_id": "user_id",
"decision": "decision",
@@ -70,7 +70,7 @@ var accessLogSortFields = map[string]string{
var sessionSortExprs = map[string]string{ //nolint:gosec // G101 false positive: "total_tokens" sort key, not a credential
"timestamp": "MAX(timestamp)",
"started_at": "MIN(timestamp)",
"cost_usd": "SUM(cost_usd)",
"cost_usd": "SUM" + CostUSDSQLExpr,
"total_tokens": "SUM(total_tokens)",
"duration": "SUM(duration)",
"request_count": "COUNT(*)",

View File

@@ -0,0 +1,124 @@
package types
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// costRow builds an access-log entry carrying only a cost breakdown — the rest
// of the row is irrelevant to the summation identities under test.
func costRow(id, session string, ts time.Time, in, cachedIn, cacheCreate, out float64) *AgentNetworkAccessLog {
return &AgentNetworkAccessLog{
ID: id,
SessionID: session,
Timestamp: ts,
InputCostUSD: in,
CachedInputCostUSD: cachedIn,
CacheCreationCostUSD: cacheCreate,
OutputCostUSD: out,
}
}
// TestAPIResponse_CostComponentsSumToAggregates is the contract a client adding
// up an API response depends on: within a single rendered object, the four
// per-bucket fields sum to cost_usd, and the two cache fields sum to
// cache_cost_usd. Uses rates that are not exactly representable in binary
// floating point, so the identity is checked against real arithmetic rather
// than round numbers.
func TestAPIResponse_CostComponentsSumToAggregates(t *testing.T) {
row := costRow("r1", "s1", time.Now(), 0.000768, 0.0002304, 0.00192, 0.003)
api := row.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "rendered buckets must sum to the rendered cost_usd")
assert.InDelta(t, api.CachedInputCostUsd+api.CacheCreationCostUsd, api.CacheCostUsd, 1e-12,
"rendered cache buckets must sum to the rendered cache_cost_usd")
assert.InDelta(t, 0.0059184, api.CostUsd, 1e-12, "total is the exact sum, not a separately rounded figure")
assert.InDelta(t, 0.0021504, api.CacheCostUsd, 1e-12, "cache cost is the exact sum of the two cache buckets")
}
// TestSessionSummary_SumsMatchSummedEntries proves a session summary equals the
// sum of the entries it renders: a client that adds up the entries itself must
// land on the same number the summary reports, per bucket and in total.
func TestSessionSummary_SumsMatchSummedEntries(t *testing.T) {
base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
entries := []*AgentNetworkAccessLog{
costRow("r1", "s1", base, 0.000768, 0.0002304, 0.00192, 0.003),
costRow("r2", "s1", base.Add(time.Minute), 0.000625, 0.0009375, 0, 0.005),
costRow("r3", "s1", base.Add(2*time.Minute), 0.0000016, 0, 0, 0.0000032),
}
sessions := FoldAccessLogSessions([]string{"s1"}, entries)
require.Len(t, sessions, 1)
sess := sessions[0].ToAPIResponse()
var wantInput, wantCachedInput, wantCacheCreation, wantOutput float64
for _, e := range entries {
wantInput += e.InputCostUSD
wantCachedInput += e.CachedInputCostUSD
wantCacheCreation += e.CacheCreationCostUSD
wantOutput += e.OutputCostUSD
}
assert.InDelta(t, wantInput, sess.InputCostUsd, 1e-12, "session input cost is the sum of its entries")
assert.InDelta(t, wantCachedInput, sess.CachedInputCostUsd, 1e-12, "session cache-read cost is the sum of its entries")
assert.InDelta(t, wantCacheCreation, sess.CacheCreationCostUsd, 1e-12, "session cache-write cost is the sum of its entries")
assert.InDelta(t, wantOutput, sess.OutputCostUsd, 1e-12, "session output cost is the sum of its entries")
assert.InDelta(t, wantInput+wantCachedInput+wantCacheCreation+wantOutput, sess.CostUsd, 1e-12,
"session total equals the summed entry buckets")
// Summing the rendered entries must give the same answer as reading the
// summary — the property a UI relies on when it totals a table itself.
var fromEntries float64
for _, e := range sess.Entries {
fromEntries += e.CostUsd
}
assert.InDelta(t, sess.CostUsd, fromEntries, 1e-12, "summary total must match the summed rendered entries")
// The sub-microdollar row must still contribute; it would vanish under
// 6-decimal quantisation.
assert.Greater(t, sess.InputCostUsd, 0.001393, "small-cost rows must not be quantised away")
}
// TestUsageBuckets_SumsMatchSummedRows proves the same identity one level up:
// a usage bucket equals the sum of the ledger rows folded into it, and the
// buckets together equal the whole range.
func TestUsageBuckets_SumsMatchSummedRows(t *testing.T) {
day1 := time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC)
day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
rows := []*AgentNetworkUsage{
{ID: "u1", Timestamp: day1, InputCostUSD: 0.000768, CachedInputCostUSD: 0.0002304, CacheCreationCostUSD: 0.00192, OutputCostUSD: 0.003},
{ID: "u2", Timestamp: day1.Add(time.Hour), InputCostUSD: 0.000625, CachedInputCostUSD: 0.0009375, OutputCostUSD: 0.005},
{ID: "u3", Timestamp: day2, InputCostUSD: 0.0000016, OutputCostUSD: 0.0000032},
}
buckets := AggregateUsageByGranularity(rows, UsageGranularityDay)
require.Len(t, buckets, 2, "two distinct days expected")
var total, cache float64
for _, b := range buckets {
api := b.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "each bucket's components must sum to its cost_usd")
total += api.CostUsd
cache += api.CacheCostUsd
}
var wantTotal, wantCache float64
for _, r := range rows {
wantTotal += r.TotalCostUSD()
wantCache += r.CacheCostUSD()
}
assert.InDelta(t, wantTotal, total, 1e-12, "buckets must sum to the total across all ledger rows")
assert.InDelta(t, wantCache, cache, 1e-12, "buckets must sum to the cache spend across all ledger rows")
// A month bucket over the same rows must total identically — regrouping
// changes the partition, never the sum.
monthly := AggregateUsageByGranularity(rows, UsageGranularityMonth)
require.Len(t, monthly, 1)
assert.InDelta(t, wantTotal, monthly[0].ToAPIResponse().CostUsd, 1e-12,
"re-bucketing at a different granularity must preserve the total")
}

View File

@@ -25,8 +25,19 @@ type AgentNetworkUsage struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
CreatedAt time.Time
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown, mirroring AgentNetworkAccessLog — the only
// cost state stored; total and cache portion are derived on read. Kept on
// the usage ledger too so spend can be attributed per bucket even for
// accounts with log collection turned off. See AgentNetworkAccessLog for
// why the columns carry a zero default.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
CreatedAt time.Time
}
// TableName keeps usage records in their own stripped table. Named
@@ -34,6 +45,17 @@ type AgentNetworkUsage struct {
// agent_network_usage table in a shared database.
func (AgentNetworkUsage) TableName() string { return "agent_network_request_usage" }
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (u *AgentNetworkUsage) TotalCostUSD() float64 {
return u.InputCostUSD + u.CachedInputCostUSD + u.CacheCreationCostUSD + u.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets.
func (u *AgentNetworkUsage) CacheCostUSD() float64 {
return u.CachedInputCostUSD + u.CacheCreationCostUSD
}
// AgentNetworkUsageGroup is the normalised many-to-many row linking a usage
// record to one authorising group, mirroring AgentNetworkAccessLogGroup so the
// usage overview can filter by group with a `group_id IN (...)` join.

View File

@@ -33,21 +33,45 @@ func ParseUsageGranularity(s string) UsageGranularity {
// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
// the UTC start of the bucket as YYYY-MM-DD.
type AgentNetworkUsageBucket struct {
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
}
// TotalCostUSD is the bucket's total spend: the sum of the four per-bucket
// costs. Derived rather than accumulated separately so it cannot disagree with
// the components.
func (b *AgentNetworkUsageBucket) TotalCostUSD() float64 {
return b.InputCostUSD + b.CachedInputCostUSD + b.CacheCreationCostUSD + b.OutputCostUSD
}
// CacheCostUSD is the bucket's prompt-cache spend: cache reads plus writes.
func (b *AgentNetworkUsageBucket) CacheCostUSD() float64 {
return b.CachedInputCostUSD + b.CacheCreationCostUSD
}
// ToAPIResponse renders the bucket as the API representation.
func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
return api.AgentNetworkUsageBucket{
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CostUsd: b.CostUSD,
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CachedInputTokens: b.CachedInputTokens,
CacheCreationTokens: b.CacheCreationTokens,
InputCostUsd: b.InputCostUSD,
CachedInputCostUsd: b.CachedInputCostUSD,
CacheCreationCostUsd: b.CacheCreationCostUSD,
OutputCostUsd: b.OutputCostUSD,
CostUsd: b.TotalCostUSD(),
CacheCostUsd: b.CacheCostUSD(),
}
}
@@ -84,7 +108,12 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
b.InputTokens += r.InputTokens
b.OutputTokens += r.OutputTokens
b.TotalTokens += r.TotalTokens
b.CostUSD += r.CostUSD
b.CachedInputTokens += r.CachedInputTokens
b.CacheCreationTokens += r.CacheCreationTokens
b.InputCostUSD += r.InputCostUSD
b.CachedInputCostUSD += r.CachedInputCostUSD
b.CacheCreationCostUSD += r.CacheCreationCostUSD
b.OutputCostUSD += r.OutputCostUSD
}
out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))

View File

@@ -61,6 +61,10 @@ type Config struct {
// EmbeddedIdP contains configuration for the embedded Dex OIDC provider.
// When set, Dex will be embedded in the management server and serve requests at /oauth2/
EmbeddedIdP *idp.EmbeddedIdPConfig
HighestSupportedSyncMessageVersion *int
PerAccountHighestSupportedSyncMessageVersion map[string]int
}
// GetAuthAudiences returns the audience from the http config and device authorization flow config

View File

@@ -0,0 +1,766 @@
package grpc
import (
"encoding/base64"
"strconv"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/proto"
)
// wgKeyRawLen is the raw byte length of a WireGuard public key.
const wgKeyRawLen = 32
// ComponentsEnvelopeInput bundles the data the component-format encoder needs.
// The envelope is fully self-contained — every field needed by the client's
// local Calculate() comes from the components struct itself. The only
// externally-supplied data is the receiving peer's PeerConfig (which is
// computed alongside the components in the network_map controller and reused
// from the legacy proto path) and the dns_domain string.
type ComponentsEnvelopeInput struct {
Components *types.NetworkMapComponents
PeerConfig *proto.PeerConfig
DNSDomain string
DNSForwarderPort int64
// UserIDClaim is the OIDC claim name the client should embed in
// SshAuth.UserIDClaim when reconstructing the NetworkMap. Empty value
// is OK — client treats empty as "no SshAuth to build".
UserIDClaim string
// ProxyPatch carries pre-expanded NetworkMap fragments injected by
// external controllers (BYOP/port-forwarding). Nil when no proxy data
// is present; encoder skips the field in that case.
ProxyPatch *proto.ProxyPatch
}
// EncodeNetworkMapEnvelope converts NetworkMapComponents into the component
// wire envelope. The encoder is intentionally non-deterministic: it iterates
// Go maps in their native (random) order. Indexes inside the envelope
// (peer_indexes, source_group_ids, agent_version_idx, router_peer_indexes)
// are self-consistent within a single encode, so the decoder reconstructs
// the same typed objects regardless of emit order. Tests that need to
// compare envelopes do so semantically via proto round-trip + canonicalize,
// not byte-equal.
//
// Callers must NOT concatenate or merge envelopes from different encodes —
// index spaces are local to a single envelope.
func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvelope {
c := in.Components
// Graceful degrade when components is nil — matches the legacy path's
// behaviour for missing/unvalidated peers (return a NetworkMap with only
// Network populated). The receiver gets an envelope it can decode
// without crashing; AccountSettings stays non-nil so client-side
// dereferences are safe.
if c.IsEmpty() {
// Match legacy missing-peer minimum: a NetworkMap with only Network
// populated. The receiver gets enough to bootstrap (Network
// identifier, dns_domain, account_settings) and the peer itself.
return &proto.NetworkMapEnvelope{
Payload: &proto.NetworkMapEnvelope_Full{
Full: &proto.NetworkMapComponentsFull{
PeerConfig: in.PeerConfig,
// components.Peers always contains the target peer
Peers: []*proto.PeerCompact{toPeerCompact(c.Peers[c.PeerID])},
DnsDomain: in.DNSDomain,
DnsForwarderPort: in.DNSForwarderPort,
UserIdClaim: in.UserIDClaim,
AccountSettings: &proto.AccountSettingsCompact{},
ProxyPatch: in.ProxyPatch,
},
},
}
}
// Phase 1: build dedup tables. Every routing peer (in c.RouterPeers) and
// every regular peer (in c.Peers) must be indexed before any encoder
// looks up indexes via e.peerOrder — otherwise routes / routers_map for
// peers that exist only in c.RouterPeers would silently lose their
// peer_index reference.
enc := newComponentEncoder(c)
enc.indexAllPeers()
routerIdxs := enc.indexRouterPeers(c.RouterPeers)
// Phase 2: gather every policy that any consumer references (peer-pair
// policies + resource-only policies) so encodeResourcePoliciesMap can
// translate every *Policy pointer to a wire index.
allPolicies := unionPolicies(c.Policies, c.ResourcePoliciesMap)
policies := enc.encodePolicies(allPolicies)
// Phase 3: emit. Order of struct field expressions no longer matters:
// every encoder either reads from the dedup tables or works on
// independent input.
full := &proto.NetworkMapComponentsFull{
Serial: networkSerial(c.Network),
PeerConfig: in.PeerConfig,
Network: toAccountNetwork(c.Network),
AccountSettings: toAccountSettingsCompact(c.AccountSettings),
DnsForwarderPort: in.DNSForwarderPort,
UserIdClaim: in.UserIDClaim,
ProxyPatch: in.ProxyPatch,
DnsSettings: enc.encodeDNSSettings(c.DNSSettings),
DnsDomain: in.DNSDomain,
CustomZoneDomain: c.CustomZoneDomain,
AgentVersions: enc.agentVersions,
Peers: enc.peers,
RouterPeerIndexes: routerIdxs,
Policies: policies,
Groups: enc.encodeGroups(),
Routes: enc.encodeRoutes(c.Routes),
NameserverGroups: enc.encodeNameServerGroups(c.NameServerGroups),
AllDnsRecords: encodeSimpleRecords(c.AllDNSRecords),
AccountZones: encodeCustomZones(c.AccountZones),
NetworkResources: enc.encodeNetworkResources(c.NetworkResources),
RoutersMap: enc.encodeRoutersMap(c.RoutersMap),
ResourcePoliciesMap: enc.encodeResourcePoliciesMap(c.ResourcePoliciesMap),
GroupIdToUserIds: enc.encodeGroupIDToUserIDs(c.GroupIDToUserIDs),
AllowedUserIds: stringSetToSlice(c.AllowedUserIDs),
PostureFailedPeers: enc.encodePostureFailedPeers(c.PostureFailedPeers),
}
return &proto.NetworkMapEnvelope{
Payload: &proto.NetworkMapEnvelope_Full{Full: full},
}
}
// networkSerial returns c.Network.CurrentSerial() with a nil guard. The
// production path always populates c.Network, but the encoder is exported
// and a hand-built components struct may omit it.
func networkSerial(n *types.Network) uint64 {
if n == nil {
return 0
}
return n.CurrentSerial()
}
type componentEncoder struct {
components *types.NetworkMapComponents
peerOrder map[string]uint32
peers []*proto.PeerCompact
agentVersionOrder map[string]uint32
agentVersions []string
}
func newComponentEncoder(c *types.NetworkMapComponents) *componentEncoder {
return &componentEncoder{
components: c,
peerOrder: make(map[string]uint32, len(c.Peers)),
peers: make([]*proto.PeerCompact, 0, len(c.Peers)),
agentVersionOrder: make(map[string]uint32),
}
}
func (e *componentEncoder) indexAllPeers() {
for _, p := range e.components.Peers {
if p == nil {
continue
}
e.appendPeer(p)
}
}
func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
if idx, ok := e.peerOrder[p.ID]; ok {
return idx
}
idx := uint32(len(e.peers))
e.peerOrder[p.ID] = idx
e.peers = append(e.peers, toPeerCompact(p))
return idx
}
// indexRouterPeers ensures every router peer is in the peer dedup table
// (c.RouterPeers may contain peers not in c.Peers when validation rules drop
// them) and returns their wire indexes for the RouterPeerIndexes field. Must
// run before any encoder that resolves peer ids via e.peerOrder.
func (e *componentEncoder) indexRouterPeers(routers map[string]*types.ComponentPeer) []uint32 {
if len(routers) == 0 {
return nil
}
out := make([]uint32, 0, len(routers))
for _, p := range routers {
if p == nil {
continue
}
out = append(out, e.appendPeer(p))
}
return out
}
func (e *componentEncoder) encodeGroups() []*proto.GroupCompact {
if len(e.components.Groups) == 0 {
return nil
}
out := make([]*proto.GroupCompact, 0, len(e.components.Groups))
for _, g := range e.components.Groups {
peerIdxs := make([]uint32, 0, len(g.Peers))
for _, peerID := range g.Peers {
if idx, ok := e.peerOrder[peerID]; ok {
peerIdxs = append(peerIdxs, idx)
}
}
out = append(out, &proto.GroupCompact{
Id: g.PublicID,
PeerIndexes: peerIdxs,
IsAll: g.IsGroupAll(),
})
}
return out
}
// encodePolicies flattens Policy{Rules} → []PolicyCompact. Returns the wire
// list and a map from policy pointer to the indexes of its emitted rules in
// that list — used by encodeResourcePoliciesMap to translate
// ResourcePoliciesMap[resourceID][]*Policy into wire-side indexes.
func (e *componentEncoder) encodePolicies(policies []*types.Policy) []*proto.PolicyCompact {
if len(policies) == 0 {
return nil
}
out := make([]*proto.PolicyCompact, 0, len(policies))
for _, pol := range policies {
if !pol.Enabled {
continue
}
for _, r := range pol.Rules {
if r == nil || !r.Enabled {
continue
}
out = append(out, e.encodePolicyRule(pol, r))
}
}
return out
}
// encodePolicyRule maps a single PolicyRule under pol to a PolicyCompact entry.
func (e *componentEncoder) encodePolicyRule(pol *types.Policy, r *types.PolicyRule) *proto.PolicyCompact {
return &proto.PolicyCompact{
Id: pol.PublicID,
Action: networkmap.GetProtoAction(string(r.Action)),
Protocol: networkmap.GetProtoProtocol(string(r.Protocol)),
Bidirectional: r.Bidirectional,
Ports: portsToUint32(r.Ports),
PortRanges: portRangesToProto(r.PortRanges),
SourceGroupIds: e.groupPublicXids(r.Sources),
DestinationGroupIds: e.groupPublicXids(r.Destinations),
AuthorizedUser: r.AuthorizedUser,
AuthorizedGroups: e.encodeAuthorizedGroups(r.AuthorizedGroups),
SourceResource: e.resourceToProto(r.SourceResource),
DestinationResource: e.resourceToProto(r.DestinationResource),
SourcePostureCheckIds: e.postureCheckSeqs(pol.SourcePostureChecks),
}
}
// groupPublicXids maps the xid group IDs in src to their public xids,
// dropping any group with invalid public xid.
func (e *componentEncoder) groupPublicXids(src []string) []string {
if len(src) == 0 {
return nil
}
out := make([]string, 0, len(src))
for _, gid := range src {
if id, ok := e.groupPublicXid(gid); ok {
out = append(out, id)
}
}
return out
}
// unionPolicies merges c.Policies with every policy referenced by
// c.ResourcePoliciesMap, deduplicating by pointer identity. Resource-only
// policies (relevant to a NetworkResource but not to peer-pair traffic)
// only live in ResourcePoliciesMap; without this union step they'd be lost
// from the wire and the client's resource-policy lookup would come back
// empty.
func unionPolicies(policies []*types.Policy, resourcePolicies map[string][]*types.Policy) []*types.Policy {
// Fast path: non-router peers have no resource-only policies, so the
// "union" is identical to `policies`. Skip the dedup map allocation.
if len(resourcePolicies) == 0 {
return policies
}
seen := make(map[string]struct{}, len(policies))
out := make([]*types.Policy, 0, len(policies))
for _, p := range policies {
if p == nil {
continue
}
if _, ok := seen[p.ID]; ok {
continue
}
seen[p.ID] = struct{}{}
out = append(out, p)
}
for _, list := range resourcePolicies {
for _, p := range list {
if p == nil {
continue
}
if _, ok := seen[p.ID]; ok {
continue
}
seen[p.ID] = struct{}{}
out = append(out, p)
}
}
return out
}
// encodeAuthorizedGroups translates rule.AuthorizedGroups (map keyed by
// group xid → local-user names) to the wire form (map keyed by group
// account_seq_id → UserNameList). Groups without a seq id are dropped —
// matches how source/destination group references handle the same case.
func (e *componentEncoder) encodeAuthorizedGroups(m map[string][]string) map[string]*proto.UserNameList {
if len(m) == 0 {
return nil
}
out := make(map[string]*proto.UserNameList, len(m))
for groupID, names := range m {
id, ok := e.groupPublicXid(groupID)
if !ok {
continue
}
out[id] = &proto.UserNameList{Names: names}
}
return out
}
func (e *componentEncoder) groupPublicXid(groupID string) (string, bool) {
g, ok := e.components.Groups[groupID]
if !ok {
return "", false
}
return g.PublicID, true
}
// resourceToProto translates types.Resource for the wire. For peer-typed
// resources the peer id is converted to a peer index into the envelope's
// peers array. For other resource types only the type string is shipped
// today (Calculate's resource-typed rule path consults SourceResource only
// for "peer" — other types fall through to group-based lookup).
func (e *componentEncoder) resourceToProto(r types.Resource) *proto.ResourceCompact {
if r.ID == "" && r.Type == "" {
return nil
}
out := &proto.ResourceCompact{Type: string(r.Type)}
if r.Type == types.ResourceTypePeer && r.ID != "" {
if idx, ok := e.peerOrder[r.ID]; ok {
out.PeerIndexSet = true
out.PeerIndex = idx
}
}
return out
}
// postureCheckSeqs translates a slice of posture-check xids to their
// public xids. Unresolvable xids are silently dropped — matches how group/peer
// references handle the same case.
func (e *componentEncoder) postureCheckSeqs(xids []string) []string {
if len(xids) == 0 || len(e.components.PostureCheckXIDToPublicID) == 0 {
return nil
}
out := make([]string, 0, len(xids))
for _, xid := range xids {
if seq, ok := e.components.PostureCheckXIDToPublicID[xid]; ok {
out = append(out, seq)
}
}
return out
}
// networkSeq translates a Network xid to its public id using
// the NetworkMapComponents.NetworkXIDToPublicID lookup. Returns (0,false) when
// the xid isn't known — callers decide whether to skip the parent record.
func (e *componentEncoder) networkPublicId(xid string) (string, bool) {
if xid == "" {
return "", false
}
id, ok := e.components.NetworkXIDToPublicID[xid]
if !ok {
return "", false
}
return id, true
}
func (e *componentEncoder) encodeDNSSettings(s *types.DNSSettings) *proto.DNSSettingsCompact {
if s == nil || len(s.DisabledManagementGroups) == 0 {
return nil
}
out := &proto.DNSSettingsCompact{
DisabledManagementGroupIds: make([]string, 0, len(s.DisabledManagementGroups)),
}
for _, gid := range s.DisabledManagementGroups {
if id, ok := e.groupPublicXid(gid); ok {
out.DisabledManagementGroupIds = append(out.DisabledManagementGroupIds, id)
}
}
return out
}
func (e *componentEncoder) encodeRoutes(routes []*nbroute.Route) []*proto.RouteRaw {
if len(routes) == 0 {
return nil
}
out := make([]*proto.RouteRaw, 0, len(routes))
for _, r := range routes {
if r == nil {
continue
}
rr := &proto.RouteRaw{
Id: r.PublicID,
NetId: string(r.NetID),
Description: r.Description,
KeepRoute: r.KeepRoute,
NetworkType: int32(r.NetworkType),
Masquerade: r.Masquerade,
Metric: int32(r.Metric),
Enabled: r.Enabled,
SkipAutoApply: r.SkipAutoApply,
Domains: r.Domains.ToPunycodeList(),
GroupIds: e.groupPublicXids(r.Groups),
AccessControlGroupIds: e.groupPublicXids(r.AccessControlGroups),
PeerGroupIds: e.groupPublicXids(r.PeerGroups),
}
if r.Network.IsValid() {
rr.NetworkCidr = r.Network.String()
}
if r.Peer != "" {
if idx, ok := e.peerOrder[r.Peer]; ok {
rr.PeerIndexSet = true
rr.PeerIndex = idx
}
}
out = append(out, rr)
}
return out
}
func (e *componentEncoder) encodeNameServerGroups(nsgs []*nbdns.NameServerGroup) []*proto.NameServerGroupRaw {
if len(nsgs) == 0 {
return nil
}
out := make([]*proto.NameServerGroupRaw, 0, len(nsgs))
for _, nsg := range nsgs {
if nsg == nil {
continue
}
entry := &proto.NameServerGroupRaw{
Id: nsg.PublicID,
Nameservers: encodeNameServers(nsg.NameServers),
GroupIds: e.groupPublicXids(nsg.Groups),
Primary: nsg.Primary,
Domains: nsg.Domains,
Enabled: nsg.Enabled,
SearchDomainsEnabled: nsg.SearchDomainsEnabled,
}
out = append(out, entry)
}
return out
}
func encodeNameServers(servers []nbdns.NameServer) []*proto.NameServer {
if len(servers) == 0 {
return nil
}
out := make([]*proto.NameServer, 0, len(servers))
for _, s := range servers {
out = append(out, &proto.NameServer{
IP: s.IP.String(),
NSType: int64(s.NSType),
Port: int64(s.Port),
})
}
return out
}
func encodeSimpleRecords(records []nbdns.SimpleRecord) []*proto.SimpleRecord {
if len(records) == 0 {
return nil
}
out := make([]*proto.SimpleRecord, 0, len(records))
for _, r := range records {
out = append(out, &proto.SimpleRecord{
Name: r.Name,
Type: int64(r.Type),
Class: r.Class,
TTL: int64(r.TTL),
RData: r.RData,
})
}
return out
}
func encodeCustomZones(zones []nbdns.CustomZone) []*proto.CustomZone {
if len(zones) == 0 {
return nil
}
out := make([]*proto.CustomZone, 0, len(zones))
for _, z := range zones {
out = append(out, &proto.CustomZone{
Domain: z.Domain,
Records: encodeSimpleRecords(z.Records),
SearchDomainDisabled: z.SearchDomainDisabled,
NonAuthoritative: z.NonAuthoritative,
})
}
return out
}
func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentResource) []*proto.NetworkResourceRaw {
if len(resources) == 0 {
return nil
}
out := make([]*proto.NetworkResourceRaw, 0, len(resources))
for _, r := range resources {
if r == nil {
continue
}
entry := &proto.NetworkResourceRaw{
Id: r.PublicID,
Name: r.Name,
Description: r.Description,
Type: string(r.Type),
Address: r.Address,
DomainValue: r.Domain,
Enabled: r.Enabled,
}
if id, ok := e.networkPublicId(r.NetworkID); ok {
entry.NetworkSeq = id
}
if r.Prefix.IsValid() {
entry.PrefixCidr = r.Prefix.String()
}
out = append(out, entry)
}
return out
}
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*types.ComponentRouter) map[string]*proto.NetworkRouterList {
if len(routersMap) == 0 {
return nil
}
out := make(map[string]*proto.NetworkRouterList, len(routersMap))
for networkXID, routers := range routersMap {
if len(routers) == 0 {
continue
}
id, ok := e.networkPublicId(networkXID)
if !ok {
continue
}
entries := make([]*proto.NetworkRouterEntry, 0, len(routers))
for peerID, r := range routers {
if r == nil {
continue
}
entry := &proto.NetworkRouterEntry{
Id: r.PublicID,
PeerGroupIds: e.groupPublicXids(r.PeerGroups),
Masquerade: r.Masquerade,
Metric: int32(r.Metric),
Enabled: r.Enabled,
}
if idx, ok := e.peerOrder[peerID]; ok {
entry.PeerIndexSet = true
entry.PeerIndex = idx
}
entries = append(entries, entry)
}
out[id] = &proto.NetworkRouterList{Entries: entries}
}
return out
}
func (e *componentEncoder) encodeResourcePoliciesMap(rpm map[string][]*types.Policy) map[string]*proto.PolicyIds {
if len(rpm) == 0 {
return nil
}
// resourceXIDToPublicID is local to one encode — built from components.NetworkResources
// (small slice). Network resources without seq id are dropped, matching how
// other components-without-seq are silently filtered.
resourceXIDToPublicID := make(map[string]string, len(e.components.NetworkResources))
for _, r := range e.components.NetworkResources {
if r != nil {
resourceXIDToPublicID[r.ID] = r.PublicID
}
}
out := make(map[string]*proto.PolicyIds, len(rpm))
for resourceXID, policies := range rpm {
resId, ok := resourceXIDToPublicID[resourceXID]
if !ok {
continue
}
ids := make([]string, 0, len(policies))
for _, pol := range policies {
ids = append(ids, pol.PublicID)
}
if len(ids) == 0 {
continue
}
out[resId] = &proto.PolicyIds{Ids: ids}
}
return out
}
func (e *componentEncoder) encodeGroupIDToUserIDs(m map[string][]string) map[string]*proto.UserIDList {
if len(m) == 0 {
return nil
}
out := make(map[string]*proto.UserIDList, len(m))
for groupID, userIDs := range m {
id, ok := e.groupPublicXid(groupID)
if !ok || len(userIDs) == 0 {
continue
}
out[id] = &proto.UserIDList{UserIds: userIDs}
}
return out
}
func stringSetToSlice(s map[string]struct{}) []string {
if len(s) == 0 {
return nil
}
out := make([]string, 0, len(s))
for k := range s {
out = append(out, k)
}
return out
}
func (e *componentEncoder) encodePostureFailedPeers(m map[string]map[string]struct{}) map[string]*proto.PeerIndexSet {
if len(m) == 0 {
return nil
}
out := make(map[string]*proto.PeerIndexSet, len(m))
for checkXID, failedPeerIDs := range m {
id, ok := e.components.PostureCheckXIDToPublicID[checkXID]
if !ok {
continue
}
idxs := make([]uint32, 0, len(failedPeerIDs))
for peerID := range failedPeerIDs {
if idx, ok := e.peerOrder[peerID]; ok {
idxs = append(idxs, idx)
}
}
if len(idxs) == 0 {
continue
}
out[id] = &proto.PeerIndexSet{PeerIndexes: idxs}
}
return out
}
// toAccountSettingsCompact always returns a non-nil message — the client
// dereferences it unconditionally during Calculate(), so a nil here would
// crash the receiver. A missing types.AccountSettingsInfo on the server
// (which shouldn't happen in production but the encoder is exported)
// degrades to login_expiration_enabled = false, which makes
// LoginExpired() return false for every peer.
func toAccountSettingsCompact(s *types.AccountSettingsInfo) *proto.AccountSettingsCompact {
if s == nil {
return &proto.AccountSettingsCompact{}
}
return &proto.AccountSettingsCompact{
PeerLoginExpirationEnabled: s.PeerLoginExpirationEnabled,
PeerLoginExpirationNs: int64(s.PeerLoginExpiration),
}
}
func toAccountNetwork(n *types.Network) *proto.AccountNetwork {
if n == nil {
return nil
}
out := &proto.AccountNetwork{
Identifier: n.Identifier,
NetCidr: n.Net.String(),
Dns: n.Dns,
Serial: n.CurrentSerial(),
}
if len(n.NetV6.IP) > 0 {
out.NetV6Cidr = n.NetV6.String()
}
return out
}
func toPeerCompact(p *types.ComponentPeer) *proto.PeerCompact {
pc := &proto.PeerCompact{
WgPubKey: decodeWgKey(p.Key),
SshPubKey: []byte(p.SSHKey),
DnsLabel: p.DNSLabel,
AgentVersion: p.AgentVersion,
AddedWithSsoLogin: p.AddedWithSSOLogin,
LoginExpirationEnabled: p.LoginExpirationEnabled,
SshEnabled: p.SSHEnabled,
SupportsIpv6: p.SupportsIPv6,
SupportsSourcePrefixes: p.SupportsSourcePrefixes,
ServerSshAllowed: p.ServerSSHAllowed,
}
if !p.LastLogin.IsZero() {
pc.LastLoginUnixNano = p.LastLogin.UnixNano()
}
switch {
case !p.IP.IsValid():
// leave Ip nil
case p.IP.Is4() || p.IP.Is4In6():
ip := p.IP.Unmap().As4()
pc.Ip = ip[:]
default:
ip := p.IP.As16()
pc.Ip = ip[:]
}
if p.IPv6.IsValid() {
ip := p.IPv6.As16()
pc.Ipv6 = ip[:]
}
return pc
}
// decodeWgKey returns the raw 32 bytes of a base64-encoded WireGuard public
// key, or nil for an empty / malformed key.
func decodeWgKey(s string) []byte {
if s == "" {
return nil
}
out := make([]byte, wgKeyRawLen)
n, err := base64.StdEncoding.Decode(out, []byte(s))
if err != nil || n != wgKeyRawLen {
return nil
}
return out
}
func portsToUint32(ports []string) []uint32 {
if len(ports) == 0 {
return nil
}
out := make([]uint32, 0, len(ports))
for _, p := range ports {
v, err := strconv.ParseUint(p, 10, 16)
if err != nil {
continue
}
out = append(out, uint32(v))
}
return out
}
func portRangesToProto(ranges []types.RulePortRange) []*proto.PortInfo_Range {
if len(ranges) == 0 {
return nil
}
out := make([]*proto.PortInfo_Range, 0, len(ranges))
for _, r := range ranges {
out = append(out, &proto.PortInfo_Range{
Start: uint32(r.Start),
End: uint32(r.End),
})
}
return out
}

View File

@@ -0,0 +1,781 @@
package grpc
import (
"bytes"
"cmp"
"net"
"net/netip"
"slices"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
goproto "google.golang.org/protobuf/proto"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
)
const testWgKeyA = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
const testWgKeyB = "BBCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
const testWgKeyC = "CBCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
// canonicalize rewrites a NetworkMapComponentsFull in place into a canonical
// form: peers reordered by wg_pub_key, with the rest of the message rewritten
// to reference the new peer indexes. Groups, policies, and router indexes are
// also sorted. After canonicalize, two envelopes built from the same logical
// input compare byte-equal via proto.Equal.
//
// This lives on the test side — the encoder itself emits in map-iteration
// order. Test-side normalization is the contract for "two encodes are
// equivalent".
func canonicalize(full *proto.NetworkMapComponentsFull) {
if full == nil {
return
}
type peerEntry struct {
peer *proto.PeerCompact
oldIdx uint32
}
entries := make([]peerEntry, len(full.Peers))
for i, p := range full.Peers {
entries[i] = peerEntry{peer: p, oldIdx: uint32(i)}
}
// DnsLabel is unique per peer; it tiebreaks on equal WgPubKey (e.g. both
// nil from malformed keys, or both empty for placeholders).
slices.SortFunc(entries, func(a, b peerEntry) int {
if c := bytes.Compare(a.peer.WgPubKey, b.peer.WgPubKey); c != 0 {
return c
}
return cmp.Compare(a.peer.DnsLabel, b.peer.DnsLabel)
})
remap := make(map[uint32]uint32, len(entries))
newPeers := make([]*proto.PeerCompact, len(entries))
for newIdx, e := range entries {
remap[e.oldIdx] = uint32(newIdx)
newPeers[newIdx] = e.peer
}
full.Peers = newPeers
full.RouterPeerIndexes = remapAndSort(full.RouterPeerIndexes, remap)
for _, g := range full.Groups {
g.PeerIndexes = remapAndSort(g.PeerIndexes, remap)
}
slices.SortFunc(full.Groups, func(a, b *proto.GroupCompact) int { return cmp.Compare(a.Id, b.Id) })
for _, r := range full.Routes {
if r.PeerIndexSet {
if newIdx, ok := remap[r.PeerIndex]; ok {
r.PeerIndex = newIdx
}
}
slices.Sort(r.GroupIds)
slices.Sort(r.AccessControlGroupIds)
slices.Sort(r.PeerGroupIds)
}
slices.SortFunc(full.Routes, func(a, b *proto.RouteRaw) int { return cmp.Compare(a.Id, b.Id) })
for _, list := range full.RoutersMap {
for _, entry := range list.Entries {
if entry.PeerIndexSet {
if newIdx, ok := remap[entry.PeerIndex]; ok {
entry.PeerIndex = newIdx
}
}
slices.Sort(entry.PeerGroupIds)
}
slices.SortFunc(list.Entries, func(a, b *proto.NetworkRouterEntry) int { return cmp.Compare(a.Id, b.Id) })
}
for _, set := range full.PostureFailedPeers {
set.PeerIndexes = remapAndSort(set.PeerIndexes, remap)
}
for _, p := range full.Policies {
slices.Sort(p.SourceGroupIds)
slices.Sort(p.DestinationGroupIds)
}
// Sort policies by (Id, source_group_ids, destination_group_ids) so that
// multiple PolicyCompact entries sharing the same Id (one per rule, when
// a Policy has multiple rules) still get a deterministic order. After
// sorting we remap indexes in ResourcePoliciesMap.
policyOldOrder := make(map[*proto.PolicyCompact]uint32, len(full.Policies))
for i, p := range full.Policies {
policyOldOrder[p] = uint32(i)
}
slices.SortFunc(full.Policies, func(a, b *proto.PolicyCompact) int {
if c := cmp.Compare(a.Id, b.Id); c != 0 {
return c
}
if c := slices.Compare(a.SourceGroupIds, b.SourceGroupIds); c != 0 {
return c
}
return slices.Compare(a.DestinationGroupIds, b.DestinationGroupIds)
})
policyRemap := make(map[uint32]uint32, len(full.Policies))
for newIdx, p := range full.Policies {
policyRemap[policyOldOrder[p]] = uint32(newIdx)
}
for _, idxs := range full.ResourcePoliciesMap {
slices.Sort(idxs.Ids)
}
for _, list := range full.GroupIdToUserIds {
slices.Sort(list.UserIds)
}
slices.Sort(full.AllowedUserIds)
}
func remapAndSort(idxs []uint32, remap map[uint32]uint32) []uint32 {
out := make([]uint32, 0, len(idxs))
for _, i := range idxs {
if newIdx, ok := remap[i]; ok {
out = append(out, newIdx)
}
}
slices.Sort(out)
return out
}
// envelopesEquivalent decodes both envelopes, canonicalizes them, and reports
// whether they're proto.Equal. Use instead of byte-comparing marshaled output:
// the encoder is intentionally non-deterministic.
func envelopesEquivalent(a, b *proto.NetworkMapEnvelope) bool {
canonicalize(a.GetFull())
canonicalize(b.GetFull())
return goproto.Equal(a, b)
}
func newTestComponents() *types.NetworkMapComponents {
peerA := &types.ComponentPeer{
ID: "peer-a",
Key: testWgKeyA,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peera",
SSHKey: "ssh-a",
AgentVersion: "0.40.0",
}
peerB := &types.ComponentPeer{
ID: "peer-b",
Key: testWgKeyB,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
DNSLabel: "peerb",
AgentVersion: "0.25.0",
}
peerC := &types.ComponentPeer{
ID: "peer-c",
Key: testWgKeyC,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc",
AgentVersion: "0.40.0",
}
return &types.NetworkMapComponents{
PeerID: "peer-a",
Network: &types.Network{
Identifier: "net-test",
Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)},
Serial: 7,
},
AccountSettings: &types.AccountSettingsInfo{
PeerLoginExpirationEnabled: true,
PeerLoginExpiration: 2 * time.Hour,
},
Peers: map[string]*types.ComponentPeer{
"peer-a": peerA,
"peer-b": peerB,
"peer-c": peerC,
},
Groups: map[string]*types.ComponentGroup{
"group-src": {ID: "group-src", PublicID: "1", Name: "Src", Peers: []string{"peer-a"}},
"group-dst": {ID: "group-dst", PublicID: "2", Name: "Dst", Peers: []string{"peer-b", "peer-c"}},
},
Policies: []*types.Policy{
{
ID: "pol-1",
PublicID: "10",
Enabled: true,
Rules: []*types.PolicyRule{{
ID: "rule-1", Enabled: true, Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolTCP, Bidirectional: true,
Ports: []string{"22", "80"},
PortRanges: []types.RulePortRange{{Start: 8000, End: 8100}},
Sources: []string{"group-src"},
Destinations: []string{"group-dst"},
}},
},
},
RouterPeers: map[string]*types.ComponentPeer{"peer-c": peerC},
}
}
func TestEncodeNetworkMapEnvelope_Basic(t *testing.T) {
c := newTestComponents()
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
Components: c,
DNSDomain: "netbird.cloud",
})
require.NotNil(t, env)
full := env.GetFull()
require.NotNil(t, full, "envelope must contain Full payload")
assert.EqualValues(t, 7, full.Serial)
assert.Equal(t, "netbird.cloud", full.DnsDomain)
require.NotNil(t, full.Network)
assert.Equal(t, "net-test", full.Network.Identifier)
assert.Equal(t, "100.64.0.0/10", full.Network.NetCidr)
require.NotNil(t, full.AccountSettings)
assert.True(t, full.AccountSettings.PeerLoginExpirationEnabled)
assert.EqualValues(t, (2 * time.Hour).Nanoseconds(), full.AccountSettings.PeerLoginExpirationNs)
require.Len(t, full.Peers, 3)
byLabel := map[string]*proto.PeerCompact{}
for _, p := range full.Peers {
assert.Len(t, p.WgPubKey, 32, "wg key must be raw 32 bytes")
assert.Len(t, p.Ip, 4, "ipv4 must be raw 4 bytes")
byLabel[p.DnsLabel] = p
}
assert.Len(t, byLabel["peerb"].Ipv6, 16, "peer-b has ipv6 → 16 bytes")
}
func TestEncodeNetworkMapEnvelope_RepeatEncodesEquivalent(t *testing.T) {
c := newTestComponents()
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
// Hammer it 100 times — Go map iteration is randomized per call, so each
// run produces different wire bytes, but the canonicalized form must
// match.
for i := 0; i < 100; i++ {
got := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
require.True(t, envelopesEquivalent(expected, got),
"encode #%d must be semantically equivalent to first encode", i)
}
}
func TestEncodeNetworkMapEnvelope_ConcurrentEncodesEquivalent(t *testing.T) {
c := newTestComponents()
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
const goroutines = 50
var wg sync.WaitGroup
wg.Add(goroutines)
results := make([]*proto.NetworkMapEnvelope, goroutines)
for i := 0; i < goroutines; i++ {
i := i
go func() {
defer wg.Done()
results[i] = EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
}()
}
wg.Wait()
for i, got := range results {
require.NotNil(t, got, "goroutine %d returned nil", i)
require.True(t, envelopesEquivalent(expected, got),
"goroutine %d produced inequivalent envelope", i)
}
}
func TestEncodeNetworkMapEnvelope_GroupsByAccountPublicId(t *testing.T) {
c := newTestComponents()
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Groups, 2)
groupByID := map[string]*proto.GroupCompact{}
for _, g := range full.Groups {
groupByID[g.Id] = g
}
require.Contains(t, groupByID, "1")
require.Contains(t, groupByID, "2")
assert.Len(t, groupByID["1"].PeerIndexes, 1)
assert.Len(t, groupByID["2"].PeerIndexes, 2)
}
func TestEncodeNetworkMapEnvelope_PolicyExpansion(t *testing.T) {
c := newTestComponents()
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Policies, 1)
pc := full.Policies[0]
assert.EqualValues(t, "10", pc.Id)
assert.Equal(t, proto.RuleAction_ACCEPT, pc.Action)
assert.Equal(t, proto.RuleProtocol_TCP, pc.Protocol)
assert.True(t, pc.Bidirectional)
assert.Equal(t, []uint32{22, 80}, pc.Ports)
require.Len(t, pc.PortRanges, 1)
assert.EqualValues(t, 8000, pc.PortRanges[0].Start)
assert.EqualValues(t, 8100, pc.PortRanges[0].End)
assert.Equal(t, []string{"1"}, pc.SourceGroupIds)
assert.Equal(t, []string{"2"}, pc.DestinationGroupIds)
}
func TestEncodeNetworkMapEnvelope_RouterIndexes(t *testing.T) {
c := newTestComponents()
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.RouterPeerIndexes, 1)
idx := full.RouterPeerIndexes[0]
require.Less(t, int(idx), len(full.Peers))
assert.Equal(t, "peerc", full.Peers[idx].DnsLabel)
}
func TestEncodeNetworkMapEnvelope_DisabledPolicySkipped(t *testing.T) {
c := newTestComponents()
c.Policies[0].Enabled = false
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
assert.Empty(t, full.Policies)
}
func TestEncodeNetworkMapEnvelope_TwoPeersSameMalformedKey(t *testing.T) {
// Both peers have nil WgPubKey after decode; canonicalize must still
// produce a stable order using DnsLabel as a tiebreaker, so 100 encodes
// canonicalize identically.
c := newTestComponents()
c.Peers["peer-a"].Key = "garbage-a-!!!"
c.Peers["peer-b"].Key = "garbage-b-!!!"
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
for i := 0; i < 100; i++ {
got := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
require.True(t, envelopesEquivalent(expected, got),
"encode #%d with two same-key peers must canonicalize equivalently", i)
}
}
func TestEncodeNetworkMapEnvelope_MalformedWgKey(t *testing.T) {
c := newTestComponents()
c.Peers["peer-a"].Key = "not-base64-!!!"
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Peers, 3)
var byLabel = map[string]*proto.PeerCompact{}
for _, p := range full.Peers {
byLabel[p.DnsLabel] = p
}
assert.Nil(t, byLabel["peera"].WgPubKey, "peer with malformed key encodes nil WgPubKey")
assert.Len(t, byLabel["peerb"].WgPubKey, 32, "other peers retain their key")
}
func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
c := newTestComponents()
v6Only := &types.ComponentPeer{
ID: "peer-v6",
Key: testWgKeyA,
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
DNSLabel: "peerv6",
AgentVersion: "0.40.0",
}
c.Peers["peer-v6"] = v6Only
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
var found *proto.PeerCompact
for _, p := range full.Peers {
if p.DnsLabel == "peerv6" {
found = p
}
}
require.NotNil(t, found, "ipv6-only peer must be present")
assert.Empty(t, found.Ip, "no IPv4 address → empty Ip")
assert.Len(t, found.Ipv6, 16)
}
func TestEncodeNetworkMapEnvelope_PeerWithoutIP(t *testing.T) {
c := newTestComponents()
c.Peers["peer-noip"] = &types.ComponentPeer{
ID: "peer-noip",
Key: testWgKeyA,
DNSLabel: "peernoip",
AgentVersion: "0.40.0",
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
var found *proto.PeerCompact
for _, p := range full.Peers {
if p.DnsLabel == "peernoip" {
found = p
}
}
require.NotNil(t, found)
assert.Empty(t, found.Ip)
assert.Empty(t, found.Ipv6)
}
func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
c := &types.NetworkMapComponents{
Network: &types.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
}
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
full := env.GetFull()
require.NotNil(t, full)
assert.Empty(t, full.Peers)
assert.Empty(t, full.Groups)
assert.Empty(t, full.Policies)
assert.Empty(t, full.RouterPeerIndexes)
require.NotNil(t, full.AccountSettings, "AccountSettingsCompact must always be emitted (client dereferences it unconditionally)")
}
func TestEncodeNetworkMapEnvelope_PeerLoginExpirationFields(t *testing.T) {
c := newTestComponents()
now := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
c.Peers["peer-a"].AddedWithSSOLogin = true
c.Peers["peer-a"].LoginExpirationEnabled = true
c.Peers["peer-a"].LastLogin = now
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
var pa *proto.PeerCompact
for _, p := range full.Peers {
if p.DnsLabel == "peera" {
pa = p
}
}
require.NotNil(t, pa)
assert.True(t, pa.AddedWithSsoLogin)
assert.True(t, pa.LoginExpirationEnabled)
assert.Equal(t, now.UnixNano(), pa.LastLoginUnixNano)
// peer-b has no UserID and no LastLogin → all fields zero-value.
var pb *proto.PeerCompact
for _, p := range full.Peers {
if p.DnsLabel == "peerb" {
pb = p
}
}
require.NotNil(t, pb)
assert.False(t, pb.AddedWithSsoLogin)
assert.False(t, pb.LoginExpirationEnabled)
assert.Zero(t, pb.LastLoginUnixNano)
}
func TestEncodeNetworkMapEnvelope_RoutesRoundTrip(t *testing.T) {
c := newTestComponents()
c.Routes = []*nbroute.Route{
{
ID: "route-peer",
PublicID: "100",
NetID: "net-A",
Description: "via peer-c",
Network: netip.MustParsePrefix("10.0.0.0/16"),
Peer: "peer-c", // peer ID, not WG key
Groups: []string{"group-src"},
AccessControlGroups: []string{"group-dst"},
Enabled: true,
},
{
ID: "route-peergroup",
PublicID: "101",
NetID: "net-B",
Network: netip.MustParsePrefix("10.1.0.0/16"),
PeerGroups: []string{"group-src", "group-dst"},
Enabled: true,
},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Routes, 2)
byNetID := map[string]*proto.RouteRaw{}
for _, r := range full.Routes {
byNetID[r.NetId] = r
}
r1 := byNetID["net-A"]
require.NotNil(t, r1)
assert.True(t, r1.PeerIndexSet, "route with peer must set peer_index_set")
require.Less(t, int(r1.PeerIndex), len(full.Peers))
assert.Equal(t, "peerc", full.Peers[r1.PeerIndex].DnsLabel)
assert.Equal(t, []string{"1"}, r1.GroupIds, "group-src has AccountSeqID 1")
assert.Equal(t, []string{"2"}, r1.AccessControlGroupIds, "group-dst has AccountSeqID 2")
assert.Empty(t, r1.PeerGroupIds)
r2 := byNetID["net-B"]
require.NotNil(t, r2)
assert.False(t, r2.PeerIndexSet, "route with peer_groups must NOT set peer_index_set")
assert.ElementsMatch(t, []string{"1", "2"}, r2.PeerGroupIds)
}
func TestEncodeNetworkMapEnvelope_RouteWithMissingPeerLeavesIndexUnset(t *testing.T) {
c := newTestComponents()
c.Routes = []*nbroute.Route{{
ID: "route-x",
PublicID: "100",
Peer: "peer-not-in-components",
Network: netip.MustParsePrefix("10.0.0.0/16"),
Enabled: true,
}}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Routes, 1)
assert.False(t, full.Routes[0].PeerIndexSet,
"missing peer reference must not pretend to point at peer index 0")
}
func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing.T) {
c := newTestComponents()
// Policy that exists ONLY in ResourcePoliciesMap, not in c.Policies. This
// is the I1 case — without unionPolicies the encoder would silently
// drop it from the wire.
resourceOnlyPolicy := &types.Policy{
ID: "pol-resource", PublicID: "99", Enabled: true,
Rules: []*types.PolicyRule{{
ID: "rule-r", Enabled: true, Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolTCP,
Sources: []string{"group-src"},
Destinations: []string{"group-dst"},
}},
}
c.ResourcePoliciesMap = map[string][]*types.Policy{
"resource-x": {c.Policies[0], resourceOnlyPolicy}, // shared + resource-only
}
// Resource must appear in components.NetworkResources with a seq id —
// encoder uses that to translate the xid map key to uint32.
c.NetworkResources = []*types.ComponentResource{
{ID: "resource-x", PublicID: "77", Name: "res-x", Enabled: true},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.Policies, 2, "encoded policies must include both peer-traffic and resource-only")
policyByID := map[string]*proto.PolicyCompact{}
policyIds := make([]string, 0)
for _, p := range full.Policies {
policyByID[p.Id] = p
policyIds = append(policyIds, p.Id)
}
require.Contains(t, policyByID, "10", "original peer-traffic policy id 10")
require.Contains(t, policyByID, "99", "resource-only policy id 99")
require.Contains(t, full.ResourcePoliciesMap, "77")
ids := full.ResourcePoliciesMap["77"].Ids
require.Len(t, ids, 2)
assert.ElementsMatch(t, policyIds, ids,
"resource policies map must reference both wire policy indexes")
}
func TestEncodeNetworkMapEnvelope_NameServerGroups(t *testing.T) {
c := newTestComponents()
c.NameServerGroups = []*nbdns.NameServerGroup{{
ID: "nsg-1", PublicID: "50", Name: "Main", Description: "primary",
NameServers: []nbdns.NameServer{{
IP: netip.MustParseAddr("8.8.8.8"), NSType: nbdns.UDPNameServerType, Port: 53,
}},
Groups: []string{"group-src", "group-not-persisted"},
Primary: true, Enabled: true,
Domains: []string{"corp.example"},
}}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.NameserverGroups, 1)
nsg := full.NameserverGroups[0]
assert.EqualValues(t, "50", nsg.Id)
assert.True(t, nsg.Primary)
require.Len(t, nsg.Nameservers, 1)
assert.Equal(t, "8.8.8.8", nsg.Nameservers[0].IP)
assert.Equal(t, []string{"1"}, nsg.GroupIds)
}
func TestEncodeNetworkMapEnvelope_PostureFailedPeers(t *testing.T) {
c := newTestComponents()
c.PostureCheckXIDToPublicID = map[string]string{"check-1": "33"}
c.PostureFailedPeers = map[string]map[string]struct{}{
"check-1": {
"peer-a": {},
"peer-b": {},
"peer-not-in-account": {},
},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Contains(t, full.PostureFailedPeers, "33")
idxs := full.PostureFailedPeers["33"].PeerIndexes
assert.Len(t, idxs, 2, "missing peer is silently dropped (filterPostureFailedPeers guarantees presence in real data)")
}
func TestEncodeNetworkMapEnvelope_RoutersMap(t *testing.T) {
c := newTestComponents()
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
"net-1": {
"peer-c": {
PublicID: "200",
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
},
},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Contains(t, full.RoutersMap, "5")
entries := full.RoutersMap["5"].Entries
require.Len(t, entries, 1)
e := entries[0]
assert.EqualValues(t, "200", e.Id)
assert.True(t, e.PeerIndexSet)
require.Less(t, int(e.PeerIndex), len(full.Peers))
assert.Equal(t, "peerc", full.Peers[e.PeerIndex].DnsLabel)
assert.True(t, e.Masquerade)
assert.EqualValues(t, 10, e.Metric)
assert.True(t, e.Enabled)
}
func TestEncodeNetworkMapEnvelope_RouterPeerNotInComponentsPeers(t *testing.T) {
// Router peer in c.RouterPeers but NOT in c.Peers (validation may have
// filtered it). indexRouterPeers runs before encodeRoutersMap, so the
// peer_index reference must still resolve.
c := newTestComponents()
delete(c.Peers, "peer-c")
routerPeer := &types.ComponentPeer{
ID: "peer-c", Key: testWgKeyC, IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc", AgentVersion: "0.40.0",
}
c.RouterPeers = map[string]*types.ComponentPeer{"peer-c": routerPeer}
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
"net-1": {"peer-c": {PublicID: "1", Peer: "peer-c", Enabled: true}},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Contains(t, full.RoutersMap, "5")
require.Len(t, full.RoutersMap["5"].Entries, 1)
e := full.RoutersMap["5"].Entries[0]
assert.True(t, e.PeerIndexSet, "router peer must be indexed even when not in c.Peers")
}
func TestEncodeNetworkMapEnvelope_GroupIDToUserIDs(t *testing.T) {
c := newTestComponents()
c.GroupIDToUserIDs = map[string][]string{
"group-src": {"user-1", "user-2"},
"group-missing": {"user-4"}, // group not in components → drop
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.Len(t, full.GroupIdToUserIds, 1, "only present groups survive")
require.Contains(t, full.GroupIdToUserIds, "1")
assert.ElementsMatch(t, []string{"user-1", "user-2"}, full.GroupIdToUserIds["1"].UserIds)
}
func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
assert.Nil(t, toProxyPatch(nil, "netbird.cloud", false, false))
assert.Nil(t, toProxyPatch(&types.NetworkMap{}, "netbird.cloud", false, false),
"empty NetworkMap (no peers, rules, routes etc) → nil patch so proto3 omits the field")
}
func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
nm := &types.NetworkMap{
Peers: []*types.ComponentPeer{{
ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
DNSLabel: "extpeer", AgentVersion: "0.40.0",
}},
FirewallRules: []*types.FirewallRule{{
PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
}},
}
patch := toProxyPatch(nm, "netbird.cloud", false, false)
require.NotNil(t, patch)
assert.Len(t, patch.Peers, 1)
assert.Len(t, patch.FirewallRules, 1)
}
// TestEncodeNetworkMapEnvelope_ProxyPatchPropagated covers the ProxyPatch
// pass-through in both encoder branches (normal path + nil-Components
// graceful-degrade). Guards against a regression that drops `ProxyPatch:`
// from one of the envelope struct literals.
func TestEncodeNetworkMapEnvelope_ProxyPatchPropagated(t *testing.T) {
patch := &proto.ProxyPatch{
ForwardingRules: []*proto.ForwardingRule{{
Protocol: proto.RuleProtocol_TCP,
DestinationPort: &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 80}},
TranslatedAddress: net.IPv4(10, 0, 0, 1).To4(),
TranslatedPort: &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 8080}},
}},
}
t.Run("normal_path", func(t *testing.T) {
c := newTestComponents()
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
Components: c,
ProxyPatch: patch,
}).GetFull()
require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the normal encode path")
assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
})
t.Run("empty_components_graceful_degrade", func(t *testing.T) {
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
Components: emptyNetworkMapComponents(),
ProxyPatch: patch,
}).GetFull()
require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the nil-Components branch too")
assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
})
}
func TestEncodeNetworkMapEnvelope_NilComponentsGracefulDegrade(t *testing.T) {
// nil Components → minimal envelope, no crash. Matches the legacy
// behaviour for missing/unvalidated peers.
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
Components: emptyNetworkMapComponents(),
DNSDomain: "netbird.cloud",
})
require.NotNil(t, env)
full := env.GetFull()
require.NotNil(t, full)
require.NotNil(t, full.AccountSettings, "AccountSettings must always be non-nil")
assert.Equal(t, "netbird.cloud", full.DnsDomain)
assert.Len(t, full.Peers, 1)
assert.Empty(t, full.Policies)
}
func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
c := &types.NetworkMapComponents{
Network: &types.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
// AccountSettings deliberately nil
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
require.NotNil(t, full.AccountSettings, "client dereferences AccountSettings unconditionally during Calculate(); a nil here would crash the receiver")
assert.False(t, full.AccountSettings.PeerLoginExpirationEnabled)
assert.Zero(t, full.AccountSettings.PeerLoginExpirationNs)
}
func emptyNetworkMapComponents() *types.NetworkMapComponents {
return types.EmptyNetworkMapComponents(
&types.NetworkMapComponents{
PeerID: "peer-id", Peers: map[string]*types.ComponentPeer{"peer-id": {}}},
)
}

View File

@@ -0,0 +1,200 @@
package grpc
import (
"context"
integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
"github.com/netbirdio/netbird/client/ssh/auth"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/management/server/types"
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
"github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/proto"
)
// ToComponentSyncResponse builds a SyncResponse carrying the compact
// NetworkMapEnvelope for capability-aware peers. The legacy proto.NetworkMap
// field is intentionally left empty — capable peers ignore it and the
// envelope alone is the authoritative wire shape.
//
// PeerConfig is computed once server-side using the receiving peer's own
// account-level network metadata. EnableSSH inside PeerConfig is left at
// peer.SSHEnabled (the peer's local setting); account-policy-driven SSH is
// computed by the client from the envelope's GroupIDToUserIDs / AllowedUserIDs
// inside Calculate(), so the SshConfig.SshEnabled bit may flip true on the
// client even though the server-side PeerConfig reports false.
func ToComponentSyncResponse(
ctx context.Context,
config *nbconfig.Config,
httpConfig *nbconfig.HttpServerConfig,
deviceFlowConfig *nbconfig.DeviceAuthorizationFlow,
peer *nbpeer.Peer,
turnCredentials *Token,
relayCredentials *Token,
components *types.NetworkMapComponents,
proxyPatch *types.NetworkMap,
dnsName string,
checks []*posture.Checks,
settings *types.Settings,
extraSettings *types.ExtraSettings,
peerGroups []string,
dnsFwdPort int64,
) *proto.SyncResponse {
//
// 'component' parameter is expected to never be nil
// 'peer' parameter is expected to never be nil
//
// TODO (dmitri) consider using invariants?
//
enableSSH := computeSSHEnabledForPeer(components, peer)
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH, components.ForceRoutingPeerDNSResolution)
includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
useSourcePrefixes := peer.SupportsSourcePrefixes()
userIDClaim := auth.DefaultUserIDClaim
if httpConfig != nil && httpConfig.AuthUserIDClaim != "" {
userIDClaim = httpConfig.AuthUserIDClaim
}
envelope := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
Components: components,
PeerConfig: peerConfig,
DNSDomain: dnsName,
DNSForwarderPort: dnsFwdPort,
UserIDClaim: userIDClaim,
ProxyPatch: toProxyPatch(proxyPatch, dnsName, includeIPv6, useSourcePrefixes),
})
resp := &proto.SyncResponse{
PeerConfig: peerConfig,
NetworkMapEnvelope: envelope,
Checks: toProtocolChecks(ctx, checks),
Version: int32(sharedgrpc.ComponentNetworkMap),
}
nbConfig := toNetbirdConfig(config, turnCredentials, relayCredentials, extraSettings, settings)
resp.NetbirdConfig = integrationsConfig.ExtendNetBirdConfig(peer.ID, peerGroups, nbConfig, extraSettings)
// settings == nil → field stays nil → "no info in this snapshot", client
// preserves the deadline it already had. settings non-nil → emit either a
// valid deadline or the explicit-zero "disabled" sentinel via
// encodeSessionExpiresAt.
if settings != nil {
resp.SessionExpiresAt = encodeSessionExpiresAt(
peer.SessionExpiresAt(settings.PeerLoginExpirationEnabled, settings.PeerLoginExpiration),
)
}
return resp
}
// toProxyPatch converts a proxy-injected *types.NetworkMap into the wire
// patch the components envelope ships alongside. Returns nil when there are
// no fragments to merge — proto3 omits a nil message field, so the receiver
// sees no patch and skips the merge step entirely.
//
// We reuse the legacy proto-conversion helpers (toProtocolRoutes,
// toProtocolFirewallRules, toProtocolRoutesFirewallRules,
// appendRemotePeerConfig, ForwardingRule.ToProto) because the proxy
// delivers fragments pre-expanded — there's no raw component shape to
// derive them from. Components purity isn't violated: proxy data isn't
// policy-graph-derived, it's externally injected post-Calculate, so the
// client merges it on top of its locally-computed NetworkMap.
func toProxyPatch(nm *types.NetworkMap, dnsName string, includeIPv6, useSourcePrefixes bool) *proto.ProxyPatch {
if nm == nil {
return nil
}
if len(nm.Peers) == 0 && len(nm.OfflinePeers) == 0 && len(nm.FirewallRules) == 0 &&
len(nm.Routes) == 0 && len(nm.RoutesFirewallRules) == 0 && len(nm.ForwardingRules) == 0 {
return nil
}
patch := &proto.ProxyPatch{
Peers: networkmap.AppendRemotePeerConfig(nil, nm.Peers, dnsName, includeIPv6),
OfflinePeers: networkmap.AppendRemotePeerConfig(nil, nm.OfflinePeers, dnsName, includeIPv6),
FirewallRules: networkmap.ToProtocolFirewallRules(nm.FirewallRules, includeIPv6, useSourcePrefixes),
Routes: networkmap.ToProtocolRoutes(nm.Routes),
RouteFirewallRules: networkmap.ToProtocolRoutesFirewallRules(nm.RoutesFirewallRules),
}
if len(nm.ForwardingRules) > 0 {
patch.ForwardingRules = make([]*proto.ForwardingRule, 0, len(nm.ForwardingRules))
for _, r := range nm.ForwardingRules {
patch.ForwardingRules = append(patch.ForwardingRules, r.ToProto())
}
}
return patch
}
// computeSSHEnabledForPeer mirrors the SSH-server-activation bit that
// Calculate() folds into NetworkMap.EnableSSH. Components-format peers
// receive a freshly-computed PeerConfig.SshConfig.SshEnabled at sync time;
// without this helper the field would be incorrectly false for any peer
// that's the destination of an SSH-enabling policy without having
// peer.SSHEnabled set locally.
//
// Mirrors the two activation paths Calculate() uses:
// 1. Explicit: rule.Protocol == NetbirdSSH and peer is in the rule's
// destinations.
// 2. Legacy implicit: rule covers TCP/22 or TCP/22022 (or ALL), peer is in
// destinations, AND the peer has SSHEnabled set locally — this is the
// "allow-all/TCP-22 implies SSH activation for SSH-capable peers" path.
//
// The full SSH AuthorizedUsers map is still produced by the client when it
// runs Calculate() over the envelope.
func computeSSHEnabledForPeer(c *types.NetworkMapComponents, peer *nbpeer.Peer) bool {
if c == nil || peer == nil {
return false
}
// Mirror Calculate's `getAllPeersFromGroups` invariant: target peer must
// exist in c.Peers, otherwise no rule applies to it.
if _, ok := c.Peers[peer.ID]; !ok {
return false
}
for _, policy := range c.Policies {
if policy == nil || !policy.Enabled {
continue
}
for _, rule := range policy.Rules {
if ruleEnablesSSHForPeer(c, rule, peer) {
return true
}
}
}
return false
}
// ruleEnablesSSHForPeer returns true when rule is active, targets peer, and
// either explicitly authorises SSH or covers the legacy TCP/22 path while the
// peer itself has SSH enabled locally.
func ruleEnablesSSHForPeer(c *types.NetworkMapComponents, rule *types.PolicyRule, peer *nbpeer.Peer) bool {
if rule == nil || !rule.Enabled {
return false
}
if !peerInDestinations(c, rule, peer.ID) {
return false
}
if rule.Protocol == types.PolicyRuleProtocolNetbirdSSH {
return true
}
return peer.SSHEnabled && types.PolicyRuleImpliesLegacySSH(rule)
}
// peerInDestinations reports whether peerID is in any of rule.Destinations'
// groups (or matches DestinationResource if it's a peer-typed resource —
// for non-peer types Calculate falls through to group lookup, so we mirror
// that exactly to avoid silent divergence).
func peerInDestinations(c *types.NetworkMapComponents, rule *types.PolicyRule, peerID string) bool {
if rule.DestinationResource.Type == types.ResourceTypePeer && rule.DestinationResource.ID != "" {
return rule.DestinationResource.ID == peerID
}
for _, groupID := range rule.Destinations {
if c.IsPeerInGroup(peerID, groupID) {
return true
}
}
return false
}

View File

@@ -0,0 +1,184 @@
package grpc
import (
"testing"
"github.com/stretchr/testify/assert"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
// TestComputeSSHEnabledForPeer covers both Calculate-mirroring branches:
// explicit NetbirdSSH protocol, and the legacy implicit case where a
// TCP/22 (or 22022 / ALL / port-range-covering-22) rule activates SSH when
// the destination peer has SSHEnabled=true locally.
func TestComputeSSHEnabledForPeer(t *testing.T) {
const targetPeerID = "target"
const targetGroupID = "g_dst"
mkComponents := func(rule *types.PolicyRule, sshEnabled bool) (*types.NetworkMapComponents, *nbpeer.Peer) {
peer := &nbpeer.Peer{ID: targetPeerID, SSHEnabled: sshEnabled}
group := &types.ComponentGroup{ID: targetGroupID, Name: "dst", Peers: []string{targetPeerID}}
return &types.NetworkMapComponents{
Peers: map[string]*types.ComponentPeer{targetPeerID: peer.ToComponent()},
Groups: map[string]*types.ComponentGroup{targetGroupID: group},
Policies: []*types.Policy{{
ID: "p",
Enabled: true,
Rules: []*types.PolicyRule{rule},
}},
}, peer
}
cases := []struct {
name string
peerSSH bool
rule types.PolicyRule
wantEnabled bool
}{
{
name: "explicit-netbird-ssh-activates-regardless-of-peer-ssh",
peerSSH: false,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
Destinations: []string{targetGroupID},
},
wantEnabled: true,
},
{
name: "implicit-tcp-22-with-peer-ssh",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
Destinations: []string{targetGroupID},
},
wantEnabled: true,
},
{
name: "implicit-tcp-22-without-peer-ssh-disabled",
peerSSH: false,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
Destinations: []string{targetGroupID},
},
wantEnabled: false,
},
{
name: "implicit-tcp-22022-with-peer-ssh",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22022"},
Destinations: []string{targetGroupID},
},
wantEnabled: true,
},
{
name: "implicit-all-protocol-with-peer-ssh",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolALL,
Destinations: []string{targetGroupID},
},
wantEnabled: true,
},
{
name: "implicit-port-range-covers-22",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true,
Protocol: types.PolicyRuleProtocolTCP,
PortRanges: []types.RulePortRange{{Start: 20, End: 30}},
Destinations: []string{targetGroupID},
},
wantEnabled: true,
},
{
name: "tcp-80-no-ssh",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"80"},
Destinations: []string{targetGroupID},
},
wantEnabled: false,
},
{
name: "disabled-rule-skipped",
peerSSH: true,
rule: types.PolicyRule{
Enabled: false, Protocol: types.PolicyRuleProtocolNetbirdSSH,
Destinations: []string{targetGroupID},
},
wantEnabled: false,
},
{
name: "peer-not-in-destinations",
peerSSH: true,
rule: types.PolicyRule{
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
Destinations: []string{"g_other"}, // target not in this group
},
wantEnabled: false,
},
{
name: "peer-typed-destination-resource-matches",
peerSSH: false,
rule: types.PolicyRule{
Enabled: true,
Protocol: types.PolicyRuleProtocolNetbirdSSH,
DestinationResource: types.Resource{ID: targetPeerID, Type: types.ResourceTypePeer},
},
wantEnabled: true,
},
{
name: "non-peer-destination-resource-falls-through-to-groups",
peerSSH: false,
rule: types.PolicyRule{
Enabled: true,
Protocol: types.PolicyRuleProtocolNetbirdSSH,
DestinationResource: types.Resource{ID: targetPeerID, Type: "host"}, // wrong type
Destinations: []string{targetGroupID}, // saved by group fallback
},
wantEnabled: true,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
c, peer := mkComponents(&tc.rule, tc.peerSSH)
got := computeSSHEnabledForPeer(c, peer)
assert.Equal(t, tc.wantEnabled, got)
})
}
}
// TestComputeSSHEnabledForPeer_TargetMissingFromComponents covers the
// belt-and-suspenders presence guard mirroring Calculate's
// getAllPeersFromGroups invariant.
func TestComputeSSHEnabledForPeer_TargetMissingFromComponents(t *testing.T) {
peer := &nbpeer.Peer{ID: "missing", SSHEnabled: true}
c := &types.NetworkMapComponents{
Peers: map[string]*types.ComponentPeer{}, // target peer NOT present
Groups: map[string]*types.ComponentGroup{
"g": {ID: "g", Peers: []string{"missing"}},
},
Policies: []*types.Policy{{
ID: "p", Enabled: true,
Rules: []*types.PolicyRule{{
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
Destinations: []string{"g"},
}},
}},
}
assert.False(t, computeSSHEnabledForPeer(c, peer),
"missing target peer must short-circuit to false, not consult policies")
}
// TestComputeSSHEnabledForPeer_NilInputs guards the cheap nil-checks at
// function entry — Calculate doesn't accept nil either, but the helper is
// exported indirectly via ToComponentSyncResponse and may receive nil
// components on graceful-degrade paths.
func TestComputeSSHEnabledForPeer_NilInputs(t *testing.T) {
assert.False(t, computeSSHEnabledForPeer(nil, &nbpeer.Peer{ID: "x"}))
assert.False(t, computeSSHEnabledForPeer(&types.NetworkMapComponents{}, nil))
}

View File

@@ -10,24 +10,20 @@ import (
"github.com/hashicorp/go-version"
nbversion "github.com/netbirdio/netbird/version"
log "github.com/sirupsen/logrus"
goproto "google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
"github.com/netbirdio/netbird/client/ssh/auth"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/netiputil"
"github.com/netbirdio/netbird/shared/sshauth"
)
const (
@@ -123,7 +119,7 @@ func toNetbirdConfig(config *nbconfig.Config, turnCredentials *Token, relayToken
return nbConfig
}
func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool) *proto.PeerConfig {
func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool, forceRoutingPeerDNS bool) *proto.PeerConfig {
netmask, _ := network.Net.Mask.Size()
fqdn := peer.FQDN(dnsName)
@@ -139,7 +135,7 @@ func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, set
Address: fmt.Sprintf("%s/%d", peer.IP.String(), netmask),
SshConfig: sshConfig,
Fqdn: fqdn,
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled,
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled || peer.ProxyMeta.Embedded || forceRoutingPeerDNS,
LazyConnectionEnabled: settings.LazyConnectionEnabled,
AutoUpdate: &proto.AutoUpdateSettings{
Version: settings.AutoUpdateVersion,
@@ -166,12 +162,12 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
useSourcePrefixes := peer.SupportsSourcePrefixes()
response := &proto.SyncResponse{
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
NetworkMap: &proto.NetworkMap{
Serial: networkMap.Network.CurrentSerial(),
Routes: toProtocolRoutes(networkMap.Routes),
DNSConfig: toProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
Routes: networkmap.ToProtocolRoutes(networkMap.Routes),
DNSConfig: networkmap.ToProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
},
Checks: toProtocolChecks(ctx, checks),
}
@@ -183,7 +179,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
response.NetworkMap.PeerConfig = response.PeerConfig
remotePeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers))
remotePeers = appendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6)
remotePeers = networkmap.AppendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6)
if !shouldSkipSendingDeprecatedRemotePeers(peer.Meta.WtVersion) {
response.RemotePeers = remotePeers
@@ -193,13 +189,13 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
response.RemotePeersIsEmpty = len(remotePeers) == 0
response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty
response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6)
response.NetworkMap.OfflinePeers = networkmap.AppendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6)
firewallRules := toProtocolFirewallRules(networkMap.FirewallRules, includeIPv6, useSourcePrefixes)
firewallRules := networkmap.ToProtocolFirewallRules(networkMap.FirewallRules, includeIPv6, useSourcePrefixes)
response.NetworkMap.FirewallRules = firewallRules
response.NetworkMap.FirewallRulesIsEmpty = len(firewallRules) == 0
routesFirewallRules := toProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules)
routesFirewallRules := networkmap.ToProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules)
response.NetworkMap.RoutesFirewallRules = routesFirewallRules
response.NetworkMap.RoutesFirewallRulesIsEmpty = len(routesFirewallRules) == 0
@@ -212,7 +208,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
}
if networkMap.AuthorizedUsers != nil {
hashedUsers, machineUsers := buildAuthorizedUsersProto(ctx, networkMap.AuthorizedUsers)
hashedUsers, machineUsers := networkmap.BuildAuthorizedUsersProto(ctx, networkMap.AuthorizedUsers)
userIDClaim := auth.DefaultUserIDClaim
if httpConfig != nil && httpConfig.AuthUserIDClaim != "" {
userIDClaim = httpConfig.AuthUserIDClaim
@@ -252,33 +248,6 @@ func encodeSessionExpiresAt(deadline time.Time) *timestamppb.Timestamp {
return timestamppb.New(deadline)
}
func buildAuthorizedUsersProto(ctx context.Context, authorizedUsers map[string]map[string]struct{}) ([][]byte, map[string]*proto.MachineUserIndexes) {
userIDToIndex := make(map[string]uint32)
var hashedUsers [][]byte
machineUsers := make(map[string]*proto.MachineUserIndexes, len(authorizedUsers))
for machineUser, users := range authorizedUsers {
indexes := make([]uint32, 0, len(users))
for userID := range users {
idx, exists := userIDToIndex[userID]
if !exists {
hash, err := sshauth.HashUserID(userID)
if err != nil {
log.WithContext(ctx).Errorf("failed to hash user id %s: %v", userID, err)
continue
}
idx = uint32(len(hashedUsers))
userIDToIndex[userID] = idx
hashedUsers = append(hashedUsers, hash[:])
}
indexes = append(indexes, idx)
}
machineUsers[machineUser] = &proto.MachineUserIndexes{Indexes: indexes}
}
return hashedUsers, machineUsers
}
func shouldSkipSendingDeprecatedRemotePeers(peerVersion string) bool {
if nbversion.IsDevelopmentVersion(peerVersion) {
return true
@@ -292,51 +261,6 @@ func shouldSkipSendingDeprecatedRemotePeers(peerVersion string) bool {
return precomputedDeprecatedRemotePeersConstraint.Check(peerNBVersion)
}
func appendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
for _, rPeer := range peers {
allowedIPs := []string{rPeer.IP.String() + "/32"}
if includeIPv6 && rPeer.IPv6.IsValid() {
allowedIPs = append(allowedIPs, rPeer.IPv6.String()+"/128")
}
dst = append(dst, &proto.RemotePeerConfig{
WgPubKey: rPeer.Key,
AllowedIps: allowedIPs,
SshConfig: &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
Fqdn: rPeer.FQDN(dnsName),
AgentVersion: rPeer.Meta.WtVersion,
})
}
return dst
}
// toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache
func toProtocolDNSConfig(update nbdns.Config, cache *cache.DNSConfigCache, forwardPort int64) *proto.DNSConfig {
protoUpdate := &proto.DNSConfig{
ServiceEnable: update.ServiceEnable,
CustomZones: make([]*proto.CustomZone, 0, len(update.CustomZones)),
NameServerGroups: make([]*proto.NameServerGroup, 0, len(update.NameServerGroups)),
ForwarderPort: forwardPort,
}
for _, zone := range update.CustomZones {
protoZone := convertToProtoCustomZone(zone)
protoUpdate.CustomZones = append(protoUpdate.CustomZones, protoZone)
}
for _, nsGroup := range update.NameServerGroups {
cacheKey := nsGroup.ID
if cachedGroup, exists := cache.GetNameServerGroup(cacheKey); exists {
protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, cachedGroup)
} else {
protoGroup := convertToProtoNameServerGroup(nsGroup)
cache.SetNameServerGroup(cacheKey, protoGroup)
protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, protoGroup)
}
}
return protoUpdate
}
func ToResponseProto(configProto nbconfig.Protocol) proto.HostConfig_Protocol {
switch configProto {
case nbconfig.UDP:
@@ -354,203 +278,6 @@ func ToResponseProto(configProto nbconfig.Protocol) proto.HostConfig_Protocol {
}
}
func toProtocolRoutes(routes []*nbroute.Route) []*proto.Route {
protoRoutes := make([]*proto.Route, 0, len(routes))
for _, r := range routes {
protoRoutes = append(protoRoutes, toProtocolRoute(r))
}
return protoRoutes
}
func toProtocolRoute(route *nbroute.Route) *proto.Route {
return &proto.Route{
ID: string(route.ID),
NetID: string(route.NetID),
Network: route.Network.String(),
Domains: route.Domains.ToPunycodeList(),
NetworkType: int64(route.NetworkType),
Peer: route.Peer,
Metric: int64(route.Metric),
Masquerade: route.Masquerade,
KeepRoute: route.KeepRoute,
SkipAutoApply: route.SkipAutoApply,
}
}
// toProtocolFirewallRules converts the firewall rules to the protocol firewall rules.
// When useSourcePrefixes is true, the compact SourcePrefixes field is populated
// alongside the deprecated PeerIP for forward compatibility.
// Wildcard rules ("0.0.0.0") are expanded into separate v4 and v6 SourcePrefixes
// when includeIPv6 is true.
func toProtocolFirewallRules(rules []*types.FirewallRule, includeIPv6, useSourcePrefixes bool) []*proto.FirewallRule {
result := make([]*proto.FirewallRule, 0, len(rules))
for i := range rules {
rule := rules[i]
fwRule := &proto.FirewallRule{
PolicyID: []byte(rule.PolicyID),
PeerIP: rule.PeerIP, //nolint:staticcheck // populated for backward compatibility
Direction: getProtoDirection(rule.Direction),
Action: getProtoAction(rule.Action),
Protocol: getProtoProtocol(rule.Protocol),
Port: rule.Port,
}
if useSourcePrefixes && rule.PeerIP != "" {
result = append(result, populateSourcePrefixes(fwRule, rule, includeIPv6)...)
}
if shouldUsePortRange(fwRule) {
fwRule.PortInfo = rule.PortRange.ToProto()
}
result = append(result, fwRule)
}
return result
}
// populateSourcePrefixes sets SourcePrefixes on fwRule and returns any
// additional rules needed (e.g. a v6 wildcard clone when the peer IP is unspecified).
func populateSourcePrefixes(fwRule *proto.FirewallRule, rule *types.FirewallRule, includeIPv6 bool) []*proto.FirewallRule {
addr, err := netip.ParseAddr(rule.PeerIP)
if err != nil {
return nil
}
if !addr.IsUnspecified() {
fwRule.SourcePrefixes = [][]byte{netiputil.EncodeAddr(addr.Unmap())}
return nil
}
// IPv4Unspecified/0 is always valid, error is impossible.
v4Wildcard, _ := netiputil.EncodePrefix(netip.PrefixFrom(netip.IPv4Unspecified(), 0))
fwRule.SourcePrefixes = [][]byte{v4Wildcard}
if !includeIPv6 {
return nil
}
v6Rule := goproto.Clone(fwRule).(*proto.FirewallRule)
v6Rule.PeerIP = "::" //nolint:staticcheck // populated for backward compatibility
// IPv6Unspecified/0 is always valid, error is impossible.
v6Wildcard, _ := netiputil.EncodePrefix(netip.PrefixFrom(netip.IPv6Unspecified(), 0))
v6Rule.SourcePrefixes = [][]byte{v6Wildcard}
if shouldUsePortRange(v6Rule) {
v6Rule.PortInfo = rule.PortRange.ToProto()
}
return []*proto.FirewallRule{v6Rule}
}
// getProtoDirection converts the direction to proto.RuleDirection.
func getProtoDirection(direction int) proto.RuleDirection {
if direction == types.FirewallRuleDirectionOUT {
return proto.RuleDirection_OUT
}
return proto.RuleDirection_IN
}
func toProtocolRoutesFirewallRules(rules []*types.RouteFirewallRule) []*proto.RouteFirewallRule {
result := make([]*proto.RouteFirewallRule, len(rules))
for i := range rules {
rule := rules[i]
result[i] = &proto.RouteFirewallRule{
SourceRanges: rule.SourceRanges,
Action: getProtoAction(rule.Action),
Destination: rule.Destination,
Protocol: getProtoProtocol(rule.Protocol),
PortInfo: getProtoPortInfo(rule),
IsDynamic: rule.IsDynamic,
Domains: rule.Domains.ToPunycodeList(),
PolicyID: []byte(rule.PolicyID),
RouteID: string(rule.RouteID),
}
}
return result
}
// getProtoAction converts the action to proto.RuleAction.
func getProtoAction(action string) proto.RuleAction {
if action == string(types.PolicyTrafficActionDrop) {
return proto.RuleAction_DROP
}
return proto.RuleAction_ACCEPT
}
// getProtoProtocol converts the protocol to proto.RuleProtocol.
func getProtoProtocol(protocol string) proto.RuleProtocol {
switch types.PolicyRuleProtocolType(protocol) {
case types.PolicyRuleProtocolALL:
return proto.RuleProtocol_ALL
case types.PolicyRuleProtocolTCP:
return proto.RuleProtocol_TCP
case types.PolicyRuleProtocolUDP:
return proto.RuleProtocol_UDP
case types.PolicyRuleProtocolICMP:
return proto.RuleProtocol_ICMP
default:
return proto.RuleProtocol_UNKNOWN
}
}
// getProtoPortInfo converts the port info to proto.PortInfo.
func getProtoPortInfo(rule *types.RouteFirewallRule) *proto.PortInfo {
var portInfo proto.PortInfo
if rule.Port != 0 {
portInfo.PortSelection = &proto.PortInfo_Port{Port: uint32(rule.Port)}
} else if portRange := rule.PortRange; portRange.Start != 0 && portRange.End != 0 {
portInfo.PortSelection = &proto.PortInfo_Range_{
Range: &proto.PortInfo_Range{
Start: uint32(portRange.Start),
End: uint32(portRange.End),
},
}
}
return &portInfo
}
func shouldUsePortRange(rule *proto.FirewallRule) bool {
return rule.Port == "" && (rule.Protocol == proto.RuleProtocol_UDP || rule.Protocol == proto.RuleProtocol_TCP)
}
// Helper function to convert nbdns.CustomZone to proto.CustomZone
func convertToProtoCustomZone(zone nbdns.CustomZone) *proto.CustomZone {
protoZone := &proto.CustomZone{
Domain: zone.Domain,
Records: make([]*proto.SimpleRecord, 0, len(zone.Records)),
SearchDomainDisabled: zone.SearchDomainDisabled,
NonAuthoritative: zone.NonAuthoritative,
}
for _, record := range zone.Records {
protoZone.Records = append(protoZone.Records, &proto.SimpleRecord{
Name: record.Name,
Type: int64(record.Type),
Class: record.Class,
TTL: int64(record.TTL),
RData: record.RData,
})
}
return protoZone
}
// Helper function to convert nbdns.NameServerGroup to proto.NameServerGroup
func convertToProtoNameServerGroup(nsGroup *nbdns.NameServerGroup) *proto.NameServerGroup {
protoGroup := &proto.NameServerGroup{
Primary: nsGroup.Primary,
Domains: nsGroup.Domains,
SearchDomainsEnabled: nsGroup.SearchDomainsEnabled,
NameServers: make([]*proto.NameServer, 0, len(nsGroup.NameServers)),
}
for _, ns := range nsGroup.NameServers {
protoGroup.NameServers = append(protoGroup.NameServers, &proto.NameServer{
IP: ns.IP.String(),
Port: int64(ns.Port),
NSType: int64(ns.NSType),
})
}
return protoGroup
}
// buildJWTConfig constructs JWT configuration for SSH servers from management server config
func buildJWTConfig(config *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow) *proto.JWTConfig {
if config == nil || config.AuthAudience == "" {

View File

@@ -2,6 +2,7 @@ package grpc
import (
"fmt"
"net"
"net/netip"
"reflect"
"testing"
@@ -14,7 +15,9 @@ import (
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/shared/management/networkmap"
)
func TestToProtocolDNSConfigWithCache(t *testing.T) {
@@ -64,13 +67,13 @@ func TestToProtocolDNSConfigWithCache(t *testing.T) {
}
// First run with config1
result1 := toProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
result1 := networkmap.ToProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
// Second run with config2
result2 := toProtocolDNSConfig(config2, &cache, int64(network_map.DnsForwarderPort))
result2 := networkmap.ToProtocolDNSConfig(config2, &cache, int64(network_map.DnsForwarderPort))
// Third run with config1 again
result3 := toProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
result3 := networkmap.ToProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
// Verify that result1 and result3 are identical
if !reflect.DeepEqual(result1, result3) {
@@ -102,15 +105,14 @@ func BenchmarkToProtocolDNSConfig(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
toProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
networkmap.ToProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
}
})
b.Run(fmt.Sprintf("WithoutCache-Size%d", size), func(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
cache := &cache.DNSConfigCache{}
toProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
networkmap.ToProtocolDNSConfig(testData, nil, int64(network_map.DnsForwarderPort))
}
})
}
@@ -301,3 +303,35 @@ func TestToNetbirdConfig_RelayInvariant(t *testing.T) {
assert.True(t, nbCfg.Metrics.Enabled, "metrics flag should carry the settings value")
})
}
func TestToPeerConfig_RoutingPeerDNSResolution(t *testing.T) {
network := &types.Network{Net: net.IPNet{IP: net.IPv4(100, 0, 0, 0), Mask: net.CIDRMask(8, 32)}}
newPeer := func(embedded bool) *nbpeer.Peer {
p := &nbpeer.Peer{IP: netip.MustParseAddr("100.0.0.1")}
p.ProxyMeta.Embedded = embedded
return p
}
tests := []struct {
name string
globalFlag bool
embedded bool
forceParam bool
wantEnabled bool
}{
{name: "global off, regular peer, no force", wantEnabled: false},
{name: "global on wins", globalFlag: true, wantEnabled: true},
{name: "embedded proxy peer forced", embedded: true, wantEnabled: true},
{name: "routing peer forced via param", forceParam: true, wantEnabled: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
settings := &types.Settings{RoutingPeerDNSResolutionEnabled: tt.globalFlag}
cfg := toPeerConfig(newPeer(tt.embedded), network, "netbird.selfhosted", settings, nil, nil, false, tt.forceParam)
assert.Equal(t, tt.wantEnabled, cfg.RoutingPeerDnsResolutionEnabled,
"RoutingPeerDnsResolutionEnabled should reflect global || embedded || forced")
})
}
}

View File

@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
UserID: req.GetUserId(),
GroupIDs: req.GetGroupIds(),
ProviderID: req.GetProviderId(),
Model: req.GetModel(),
})
if err != nil {
log.WithContext(ctx).Errorf("select policy for request: %v", err)
@@ -608,11 +609,11 @@ func (s *ProxyServiceServer) disconnectProxy(conn *proxyConnection) {
if err := s.proxyController.UnregisterProxyFromCluster(context.Background(), conn.address, conn.proxyID); err != nil {
log.Warnf("Failed to unregister proxy %s from cluster: %v", conn.proxyID, err)
}
conn.cancel()
if err := s.proxyManager.Disconnect(context.Background(), conn.proxyID, conn.sessionID); err != nil {
log.Warnf("Failed to mark proxy %s as disconnected: %v", conn.proxyID, err)
}
conn.cancel()
log.Infof("Proxy %s session %s disconnected", conn.proxyID, conn.sessionID)
}

View File

@@ -0,0 +1,138 @@
package grpc
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/codes"
grpcstatus "google.golang.org/grpc/status"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/shared/management/proto"
)
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
// and returns a pre-programmed result, so tests can assert what the handler
// forwards to the selector.
type fakeAgentNetworkLimits struct {
gotInput agentnetwork.PolicySelectionInput
result *agentnetwork.PolicySelectionResult
err error
}
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
f.gotInput = in
if f.err != nil {
return nil, f.err
}
return f.result, nil
}
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
return nil
}
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
// the model the proxy extracted must reach the selector's Model unchanged,
// alongside the account/user/group/provider fields.
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
UserId: "user-1",
GroupIds: []string{"grp-a", "grp-b"},
ProviderId: "prov-1",
Model: "claude-opus-4",
}
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
assert.Equal(t, "user-1", fake.gotInput.UserID)
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
"the request's model must be forwarded to the selector")
}
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
// model (empty string) is forwarded as-is; the selector decides how to treat it.
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
}
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
}
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
// envelope surfaces the model-allowlist deny code + reason through the response.
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
Allow: false,
DenyCode: "llm_policy.model_blocked",
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "deny", resp.Decision)
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
assert.NotEmpty(t, resp.DenyReason)
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
}
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
// failure surfaces as an Internal gRPC error rather than a silent allow.
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
}
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
// fallback: with no limits service wired the RPC returns Unimplemented.
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
s := &ProxyServiceServer{}
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Unimplemented, st.Code())
}

View File

@@ -25,6 +25,7 @@ import (
"google.golang.org/grpc/status"
"github.com/netbirdio/netbird/shared/management/client/common"
"github.com/netbirdio/netbird/shared/management/grpc"
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
@@ -245,6 +246,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
realIP := getRealIP(ctx)
sRealIP := realIP.String()
peerMeta := extractPeerMeta(ctx, syncReq.GetMeta())
userID, err := s.accountManager.GetUserIDByPeerKey(ctx, peerKey.String())
if err != nil {
s.syncSem.Add(-1)
@@ -683,8 +685,9 @@ func extractPeerMeta(ctx context.Context, meta *proto.PeerSystemMeta) nbpeer.Pee
LazyConnectionEnabled: meta.GetFlags().GetLazyConnectionEnabled(),
DisableIPv6: meta.GetFlags().GetDisableIPv6(),
},
Files: files,
Capabilities: capabilitiesToInt32(meta.GetCapabilities()),
Files: files,
Capabilities: capabilitiesToInt32(meta.GetCapabilities()),
SyncMessageVersion: int(meta.GetSyncMessageVersion()),
}
}
@@ -918,7 +921,7 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
// if peer has reached this point then it has logged in
loginResp := &proto.LoginResponse{
NetbirdConfig: toNetbirdConfig(s.config, nil, relayToken, nil, settings),
PeerConfig: toPeerConfig(peer, network, s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH),
PeerConfig: toPeerConfig(peer, network, s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH, false),
Checks: toProtocolChecks(ctx, postureChecks),
}
@@ -1016,7 +1019,43 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
return status.Errorf(codes.Internal, "failed to get peer groups %s", err)
}
plainResp := ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, peer, turnToken, relayToken, networkMap, s.networkMapController.GetDNSDomain(settings), postureChecks, nil, settings, settings.Extra, peerGroups, dnsFwdPort)
dnsName := s.networkMapController.GetDNSDomain(settings)
var plainResp *proto.SyncResponse
commonSyncMessageVersion := grpc.HighestCommonSyncMessageVersion(
s.perAccountOrGlobalSyncMessageVersions(peer.AccountID),
grpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
log.WithContext(ctx).
WithFields(log.Fields{
"sync_message_version": commonSyncMessageVersion,
"server_sync_message_version": s.perAccountOrGlobalSyncMessageVersions(peer.AccountID),
"peer_sync_message_version": grpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
}).Debug("common highest sync message version")
if commonSyncMessageVersion == grpc.ComponentNetworkMap {
// Capable peer: discard the legacy NetworkMap that SyncAndMarkPeer
// computed and recompute the raw components instead. This wastes one
// Calculate() call per initial-sync — the component-based wire
// format is what the peer actually consumes. The streaming path
// (network_map.Controller.UpdateAccountPeers) skips this duplication
// because it dispatches by capability before computing.
//
// TODO: refactor SyncPeer / SyncAndMarkPeer / their mocks + manager
// interfaces to return PeerNetworkMapResult so the initial-sync path
// stops doing duplicate work. Deferred until the client-side
// decoder lands and there's a real deployment of capability=3 peers
// worth optimizing for.
freshPeer, components, proxyPatch, freshPostureChecks, freshDnsFwdPort, err := s.networkMapController.GetValidatedPeerWithComponents(ctx, false, peer.AccountID, peer)
if err != nil {
log.WithContext(ctx).Errorf("failed to build components for peer %s on initial sync: %v", peer.ID, err)
return status.Errorf(codes.Internal, "failed to build initial sync envelope")
}
plainResp = ToComponentSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, freshPeer, turnToken, relayToken, components, proxyPatch, dnsName, freshPostureChecks, settings, settings.Extra, peerGroups, freshDnsFwdPort)
} else {
plainResp = ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, peer, turnToken, relayToken, networkMap, dnsName, postureChecks, nil, settings, settings.Extra, peerGroups, dnsFwdPort)
}
key, err := s.secretsManager.GetWGKey()
if err != nil {
@@ -1041,6 +1080,13 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
return nil
}
func (s *Server) perAccountOrGlobalSyncMessageVersions(accountId string) grpc.SyncMessageVersion {
if version, ok := s.config.PerAccountHighestSupportedSyncMessageVersion[accountId]; ok {
return grpc.SyncMessageVersionFromConfig(&version)
}
return grpc.SyncMessageVersionFromConfig(s.config.HighestSupportedSyncMessageVersion)
}
// GetDeviceAuthorizationFlow returns a device authorization flow information
// This is used for initiating an Oauth 2 device authorization grant flow
// which will be used by our clients to Login

View File

@@ -1647,6 +1647,10 @@ func (am *DefaultAccountManager) SyncUserJWTGroups(ctx context.Context, userAuth
return nil
}
for _, g := range newGroupsToCreate {
g.PublicID = xid.New().String()
}
if err = transaction.CreateGroups(ctx, userAuth.AccountId, newGroupsToCreate); err != nil {
return fmt.Errorf("error saving groups: %w", err)
}

View File

@@ -3170,6 +3170,16 @@ func TestAccount_SetJWTGroups(t *testing.T) {
user, err := manager.Store.GetUserByUserID(context.Background(), store.LockingStrengthNone, "user2")
assert.NoError(t, err, "unable to get user")
assert.Len(t, user.AutoGroups, 1, "new group should be added")
var newJWTGroup *types.Group
for _, g := range groups {
if g.Name == "group3" {
newJWTGroup = g
break
}
}
require.NotNil(t, newJWTGroup, "JIT-created JWT group not found")
assert.NotEqual(t, "", newJWTGroup.PublicID, "JIT-created JWT group must have a non-empty PublicID")
})
t.Run("remove all JWT groups when list is empty", func(t *testing.T) {

View File

@@ -93,6 +93,8 @@ func (am *DefaultAccountManager) CreateGroup(ctx context.Context, accountID, use
events := am.prepareGroupEvents(ctx, transaction, accountID, userID, newGroup)
eventsToStore = append(eventsToStore, events...)
newGroup.PublicID = xid.New().String()
if err := transaction.CreateGroup(ctx, newGroup); err != nil {
return status.Errorf(status.Internal, "failed to create group: %v", err)
}
@@ -158,6 +160,8 @@ func (am *DefaultAccountManager) UpdateGroup(ctx context.Context, accountID, use
return err
}
newGroup.PublicID = oldGroup.PublicID
if err = transaction.UpdateGroup(ctx, newGroup); err != nil {
return err
}
@@ -235,6 +239,7 @@ func (am *DefaultAccountManager) CreateGroups(ctx context.Context, accountID, us
}
newGroup.AccountID = accountID
newGroup.PublicID = xid.New().String()
if err = transaction.CreateGroup(ctx, newGroup); err != nil {
return err
@@ -327,6 +332,12 @@ func (am *DefaultAccountManager) updateSingleGroup(ctx context.Context, accountI
newGroup.AccountID = accountID
oldGroup, err := transaction.GetGroupByID(ctx, store.LockingStrengthNone, accountID, newGroup.ID)
if err != nil {
return err
}
newGroup.PublicID = oldGroup.PublicID
if err := transaction.UpdateGroup(ctx, newGroup); err != nil {
return err
}
@@ -341,7 +352,6 @@ func (am *DefaultAccountManager) updateSingleGroup(ctx context.Context, accountI
events = am.prepareGroupEvents(ctx, transaction, accountID, userID, newGroup)
var err error
snap, err = affectedpeers.Load(ctx, transaction, accountID, change)
return err
})

View File

@@ -6,6 +6,7 @@ import (
"github.com/netbirdio/netbird/management/server/account"
"github.com/netbirdio/netbird/management/server/activity"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/permissions/modules"
"github.com/netbirdio/netbird/management/server/permissions/operations"
@@ -30,6 +31,10 @@ type managerImpl struct {
accountManager account.Manager
}
func eventMetaResource(group *types.Group, resource *resourceTypes.NetworkResource) map[string]any {
return map[string]any{"name": group.Name, "id": group.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
type mockManager struct {
}
@@ -109,7 +114,7 @@ func (m *managerImpl) AddResourceToGroupInTransaction(ctx context.Context, trans
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, group.EventMetaResource(networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, eventMetaResource(group, networkResource))
}
return event, nil
@@ -133,7 +138,7 @@ func (m *managerImpl) RemoveResourceFromGroupInTransaction(ctx context.Context,
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, group.EventMetaResource(networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, eventMetaResource(group, networkResource))
}
return event, nil

View File

@@ -446,7 +446,7 @@ func (h *Handler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request) {
netMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, dns.CustomZone{}, nil, validPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil, account.GetActiveGroupUsers())
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, dnsDomain))
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, account.Peers, dnsDomain))
}
func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request) {
@@ -534,15 +534,20 @@ func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request)
util.WriteJSONObject(r.Context(), w, resp)
}
func toAccessiblePeers(netMap *types.NetworkMap, dnsDomain string) []api.AccessiblePeer {
// toAccessiblePeers rehydrates the calculated map's component peers into the
// account's full peer objects, which carry the location/status/meta fields
// the API response needs.
func toAccessiblePeers(netMap *types.NetworkMap, accountPeers map[string]*nbpeer.Peer, dnsDomain string) []api.AccessiblePeer {
accessiblePeers := make([]api.AccessiblePeer, 0, len(netMap.Peers)+len(netMap.OfflinePeers))
for _, p := range netMap.Peers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
}
for _, p := range netMap.OfflinePeers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
add := func(peers []*types.ComponentPeer) {
for _, p := range peers {
if peer := accountPeers[p.ID]; peer != nil {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(peer, dnsDomain))
}
}
}
add(netMap.Peers)
add(netMap.OfflinePeers)
return accessiblePeers
}

View File

@@ -21,8 +21,11 @@ import (
)
const (
staticClientDashboard = "netbird-dashboard"
staticClientCLI = "netbird-cli"
StaticClientDashboard = "netbird-dashboard"
StaticClientCLI = "netbird-cli"
DefaultTOTPAuthenticatorID = "default-totp"
LocalConnectorID = dex.LocalConnectorID
defaultCLIRedirectURL1 = "http://localhost:53000/"
defaultCLIRedirectURL2 = "http://localhost:54000/"
defaultScopes = "openid profile email groups"
@@ -189,14 +192,14 @@ func (c *EmbeddedIdPConfig) ToYAMLConfig() (*dex.YAMLConfig, error) {
EnablePasswordDB: true,
StaticClients: []storage.Client{
{
ID: staticClientDashboard,
ID: StaticClientDashboard,
Name: "NetBird Dashboard",
Public: true,
RedirectURIs: redirectURIs,
PostLogoutRedirectURIs: sanitizePostLogoutRedirectURIs(dashboardPostLogoutRedirectURIs),
},
{
ID: staticClientCLI,
ID: StaticClientCLI,
Name: "NetBird CLI",
Public: true,
RedirectURIs: redirectURIs,
@@ -258,13 +261,13 @@ func sanitizePostLogoutRedirectURIs(uris []string) []string {
func configureMFA(cfg *dex.YAMLConfig, sessionMaxLifetime, sessionIdleTimeout string, rememberMe bool, sessionCookieEncryptionKey string) error {
cfg.MFA.Authenticators = []dex.MFAAuthenticator{{
ID: "default-totp",
ID: DefaultTOTPAuthenticatorID,
// Has to be caps otherwise it will fail
Type: "TOTP",
Config: map[string]interface{}{
"issuer": "NetBird",
},
ConnectorTypes: []string{"local"},
ConnectorTypes: []string{LocalConnectorID},
}}
if sessionMaxLifetime == "" {
@@ -740,7 +743,7 @@ func (m *EmbeddedIdPManager) GetDefaultScopes() string {
// GetCLIClientID returns the client ID for CLI authentication.
func (m *EmbeddedIdPManager) GetCLIClientID() string {
return staticClientCLI
return StaticClientCLI
}
// GetCLIRedirectURLs returns the redirect URLs configured for the CLI client.
@@ -779,7 +782,7 @@ func (m *EmbeddedIdPManager) GetLocalKeysLocation() string {
// GetClientIDs returns the OAuth2 client IDs configured for this provider.
func (m *EmbeddedIdPManager) GetClientIDs() []string {
return []string{staticClientDashboard, staticClientCLI}
return []string{StaticClientDashboard, StaticClientCLI}
}
// GetUserIDClaim returns the JWT claim name used for user identification.
@@ -796,11 +799,11 @@ func (m *EmbeddedIdPManager) IsLocalAuthDisabled() bool {
func (m *EmbeddedIdPManager) SetMFAEnabled(ctx context.Context, enabled bool) error {
var mfaChain []string
if enabled {
mfaChain = []string{"default-totp"}
mfaChain = []string{DefaultTOTPAuthenticatorID}
}
if err := m.provider.SetClientsMFAChain(ctx, []string{
staticClientCLI,
staticClientDashboard,
StaticClientCLI,
StaticClientDashboard,
}, mfaChain); err != nil {
return fmt.Errorf("failed to set MFA enabled=%v: %w", enabled, err)
}

View File

@@ -331,7 +331,7 @@ func TestEmbeddedIdPConfig_ToYAMLConfig_IncludesDeviceCallbackRedirectURI(t *tes
var cliRedirectURIs []string
for _, client := range yamlConfig.StaticClients {
if client.ID == staticClientCLI {
if client.ID == StaticClientCLI {
cliRedirectURIs = client.RedirectURIs
break
}

View File

@@ -13,6 +13,7 @@ import (
"strings"
"unicode/utf8"
"github.com/rs/xid"
log "github.com/sirupsen/logrus"
"gorm.io/gorm"
"gorm.io/gorm/clause"
@@ -635,3 +636,128 @@ func RemoveDuplicatePeerKeys(ctx context.Context, db *gorm.DB) error {
return nil
}
func BackfillPublicIDs[T any](ctx context.Context, db *gorm.DB) error {
var model T
if !db.Migrator().HasTable(&model) {
log.WithContext(ctx).Debugf("Table for %T does not exist, no backfill needed", model)
return nil
}
stmt := &gorm.Statement{DB: db}
err := stmt.Parse(&model)
if err != nil {
return fmt.Errorf("parse model: %w", err)
}
tableName := stmt.Schema.Table
if err := db.Transaction(func(tx *gorm.DB) error {
if !tx.Migrator().HasColumn(&model, "public_id") {
log.WithContext(ctx).Infof("Column public_id does not exist in table %s, adding it", tableName)
if err := tx.Migrator().AddColumn(&model, "public_id"); err != nil {
return fmt.Errorf("add column public_id: %w", err)
}
}
var rows []map[string]any
if err := tx.Table(tableName).Select("id", "public_id").Where("public_id IS NULL").Or("public_id = ''").Find(&rows).Error; err != nil {
return fmt.Errorf("failed to find rows with empty public_id: %w", err)
}
if len(rows) == 0 {
log.WithContext(ctx).Infof("No rows with empty public_id found in table %s, no migration needed", tableName)
return nil
}
for _, row := range rows {
if err := tx.Table(tableName).Where("id = ?", row["id"]).Update("public_id", xid.New().String()).Error; err != nil {
return fmt.Errorf("failed to update row with id %v: %w", row["id"], err)
}
}
return nil
}); err != nil {
return err
}
log.WithContext(ctx).Infof("Backfill of empty public_id in table %s completed", tableName)
return nil
}
// FoldCostAggregatesIntoBuckets migrates a per-request cost table from the old
// "stored aggregate" shape (cost_usd + cache_cost_usd columns) to the per-bucket
// breakdown, where the total and cache portion are derived on read instead.
//
// The fold preserves both aggregates exactly for historical rows: the cache
// total moves into cached_input_cost_usd and the remainder into
// input_cost_usd, so a row's derived total and cache cost still match what it
// reported before the upgrade. The finer split is genuinely unknown for those
// rows — the old schema never recorded a read/write or input/output division —
// so it is lumped rather than guessed; only rows written after the upgrade
// carry a true four-way split.
//
// Dropping the columns before folding would zero every historical row's cost,
// so the update runs first and the drop only happens once it succeeds. A table
// with no cost_usd column has already been migrated (or was created fresh) and
// is skipped.
func FoldCostAggregatesIntoBuckets[T any](ctx context.Context, db *gorm.DB) error {
var model T
if !db.Migrator().HasTable(&model) {
log.WithContext(ctx).Debugf("table for %T does not exist, no cost-bucket migration needed", model)
return nil
}
if !db.Migrator().HasColumn(&model, "cost_usd") {
log.WithContext(ctx).Debugf("table for %T has no cost_usd column, cost buckets already migrated", model)
return nil
}
stmt := &gorm.Statement{DB: db}
if err := stmt.Parse(&model); err != nil {
return fmt.Errorf("parse model schema: %w", err)
}
tableName := stmt.Schema.Table
// COALESCE guards rows whose new columns were added as NULL by an earlier
// AutoMigrate run that predates the NOT NULL default.
hasCacheColumn := db.Migrator().HasColumn(&model, "cache_cost_usd")
cacheExpr := "0"
if hasCacheColumn {
cacheExpr = "COALESCE(cache_cost_usd, 0)"
}
if err := db.Transaction(func(tx *gorm.DB) error {
// Only touch rows that carry a legacy total and no breakdown yet, so
// the migration is idempotent and never overwrites a true split.
update := fmt.Sprintf(`UPDATE %s
SET input_cost_usd = COALESCE(cost_usd, 0) - %s,
cached_input_cost_usd = %s,
cache_creation_cost_usd = 0,
output_cost_usd = 0
WHERE COALESCE(cost_usd, 0) <> 0
AND COALESCE(input_cost_usd, 0) = 0
AND COALESCE(cached_input_cost_usd, 0) = 0
AND COALESCE(cache_creation_cost_usd, 0) = 0
AND COALESCE(output_cost_usd, 0) = 0`, tableName, cacheExpr, cacheExpr)
res := tx.Exec(update)
if res.Error != nil {
return fmt.Errorf("fold legacy cost aggregates in %s: %w", tableName, res.Error)
}
log.WithContext(ctx).Infof("folded legacy cost aggregates into per-bucket columns for %d rows in table %s", res.RowsAffected, tableName)
if err := tx.Migrator().DropColumn(&model, "cost_usd"); err != nil {
return fmt.Errorf("drop cost_usd from %s: %w", tableName, err)
}
if hasCacheColumn {
if err := tx.Migrator().DropColumn(&model, "cache_cost_usd"); err != nil {
return fmt.Errorf("drop cache_cost_usd from %s: %w", tableName, err)
}
}
return nil
}); err != nil {
return err
}
log.WithContext(ctx).Infof("migration of stored cost aggregates to per-bucket columns in table %s completed", tableName)
return nil
}

View File

@@ -16,6 +16,7 @@ import (
"gorm.io/driver/sqlite"
"gorm.io/gorm"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/migration"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/testutil"
@@ -639,3 +640,99 @@ func TestCleanupOrphanedResources_SkipsWhenForeignKeyExists(t *testing.T) {
db.Model(&testChildWithFK{}).Count(&count)
assert.Equal(t, int64(2), count, "Both rows should survive — migration must skip when FK constraint exists")
}
// legacyCostRow is the pre-breakdown shape of the usage table: cost was stored
// as a total plus a cache portion, with no per-bucket columns. Used to build a
// realistic pre-upgrade table for the fold migration to run against.
type legacyCostRow struct {
ID string `gorm:"primaryKey"`
AccountID string
Model string
CostUSD float64
CacheCostUSD float64
}
func (legacyCostRow) TableName() string { return "agent_network_request_usage" }
// TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost covers the upgrade
// path: a table written under the old schema must come out with its per-row
// total and cache cost unchanged, because dropping cost_usd without folding it
// forward would silently zero every historical row's spend.
func TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
// setupDatabase hands back a process-shared database, so start from a clean
// table rather than inheriting rows from another test.
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&legacyCostRow{}), "legacy table must be created")
require.NoError(t, db.Create(&legacyCostRow{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6", CostUSD: 0.0123, CacheCostUSD: 0.0029,
}).Error)
require.NoError(t, db.Create(&legacyCostRow{
ID: "u2", AccountID: "acct-1", Model: "gpt-4o", CostUSD: 0.5, CacheCostUSD: 0,
}).Error)
// A zero-cost row (denied / unpriced request) must stay zero, not be touched.
require.NoError(t, db.Create(&legacyCostRow{ID: "u3", AccountID: "acct-1", Model: "gw/unpriced"}).Error)
// AutoMigrate adds the per-bucket columns alongside the legacy ones, exactly
// as a real upgrade does before the post-auto migrations run.
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}), "new columns must be added")
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db))
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cost_usd"),
"legacy cost_usd column must be dropped once folded")
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cache_cost_usd"),
"legacy cache_cost_usd column must be dropped once folded")
var rows []*agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.Order("id").Find(&rows).Error)
require.Len(t, rows, 3)
// u1: total and cache portion both preserved; the read/write and
// input/output splits are unknowable for a legacy row, so the cache total
// lands on cached_input and the remainder on input.
assert.InDelta(t, 0.0123, rows[0].TotalCostUSD(), 1e-9, "historical total must survive the fold")
assert.InDelta(t, 0.0029, rows[0].CacheCostUSD(), 1e-9, "historical cache cost must survive the fold")
assert.InDelta(t, 0.0094, rows[0].InputCostUSD, 1e-9, "non-cache remainder lands on input")
assert.InDelta(t, 0.0029, rows[0].CachedInputCostUSD, 1e-9, "legacy cache total lands on cached input")
assert.Zero(t, rows[0].CacheCreationCostUSD, "legacy rows carry no read/write split to recover")
assert.Zero(t, rows[0].OutputCostUSD, "legacy rows carry no input/output split to recover")
// u2: no cache spend — the whole total is the non-cache remainder.
assert.InDelta(t, 0.5, rows[1].TotalCostUSD(), 1e-9, "cache-free historical total must survive")
assert.Zero(t, rows[1].CacheCostUSD(), "a cache-free row must stay cache-free")
// u3: zero stays zero rather than being rewritten.
assert.Zero(t, rows[2].TotalCostUSD(), "an unpriced row must remain unpriced")
}
// TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated proves the migration is
// safe to re-run: with no legacy column present it is a no-op that leaves a
// true four-way split untouched.
func TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}))
// Timestamp must be set explicitly: a zero time.Time serialises as
// '0000-00-00 00:00:00', which MySQL rejects under strict mode.
require.NoError(t, db.Create(&agentNetworkTypes.AgentNetworkUsage{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6",
Timestamp: time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC),
InputCostUSD: 0.001, CachedInputCostUSD: 0.002, CacheCreationCostUSD: 0.003, OutputCostUSD: 0.004,
}).Error)
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db),
"running against an already-migrated table must be a no-op, not an error")
var row agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.First(&row, "id = ?", "u1").Error)
assert.InDelta(t, 0.001, row.InputCostUSD, 1e-9, "a true split must not be rewritten")
assert.InDelta(t, 0.002, row.CachedInputCostUSD, 1e-9)
assert.InDelta(t, 0.003, row.CacheCreationCostUSD, 1e-9)
assert.InDelta(t, 0.004, row.OutputCostUSD, 1e-9)
assert.InDelta(t, 0.01, row.TotalCostUSD(), 1e-9, "derived total sums the four buckets")
}

View File

@@ -67,6 +67,8 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
return err
}
newNSGroup.PublicID = xid.New().String()
if err = transaction.SaveNameServerGroup(ctx, newNSGroup); err != nil {
return err
}
@@ -116,6 +118,8 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
return err
}
nsGroupToSave.PublicID = oldNSGroup.PublicID
if err = transaction.SaveNameServerGroup(ctx, nsGroupToSave); err != nil {
return err
}

View File

@@ -71,9 +71,16 @@ func (m *managerImpl) CreateNetwork(ctx context.Context, userID string, network
network.ID = xid.New().String()
err = m.store.SaveNetwork(ctx, network)
err = m.store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
network.PublicID = xid.New().String()
if err := transaction.SaveNetwork(ctx, network); err != nil {
return fmt.Errorf("failed to save network: %w", err)
}
return nil
})
if err != nil {
return nil, fmt.Errorf("failed to save network: %w", err)
return nil, err
}
m.accountManager.StoreEvent(ctx, userID, network.ID, network.AccountID, activity.NetworkCreated, network.EventMeta())
@@ -102,14 +109,25 @@ func (m *managerImpl) UpdateNetwork(ctx context.Context, userID string, network
return nil, status.NewPermissionDeniedError()
}
_, err = m.store.GetNetworkByID(ctx, store.LockingStrengthUpdate, network.AccountID, network.ID)
err = m.store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
existing, err := transaction.GetNetworkByID(ctx, store.LockingStrengthUpdate, network.AccountID, network.ID)
if err != nil {
return fmt.Errorf("failed to get network: %w", err)
}
network.PublicID = existing.PublicID
if err := transaction.SaveNetwork(ctx, network); err != nil {
return fmt.Errorf("failed to save network: %w", err)
}
return nil
})
if err != nil {
return nil, fmt.Errorf("failed to get network: %w", err)
return nil, err
}
m.accountManager.StoreEvent(ctx, userID, network.ID, network.AccountID, activity.NetworkUpdated, network.EventMeta())
return network, m.store.SaveNetwork(ctx, network)
return network, nil
}
func (m *managerImpl) DeleteNetwork(ctx context.Context, accountID, userID, networkID string) error {

View File

@@ -255,3 +255,73 @@ func Test_UpdateNetworkFailsWithPermissionDenied(t *testing.T) {
require.Error(t, err)
require.Nil(t, updatedNetwork)
}
// Test_CreateNetworkAllocatesSeqID verifies that CreateNetwork sets a
// non-zero AccountSeqID on the persisted network (allocated through the
// account_seq_counters table).
func Test_CreateNetworkSetsPublicId(t *testing.T) {
ctx := context.Background()
const accountID = "testAccountId"
const userID = "testAdminId"
s, cleanUp, err := store.NewTestStoreFromSQL(ctx, "../testdata/networks.sql", t.TempDir())
require.NoError(t, err)
t.Cleanup(cleanUp)
am := mock_server.MockAccountManager{}
permissionsManager := permissions.NewManager(s)
groupsManager := groups.NewManagerMock()
routerManager := routers.NewManagerMock()
resourcesManager := resources.NewManager(s, permissionsManager, groupsManager, &am, nil)
manager := NewManager(s, permissionsManager, resourcesManager, routerManager, &am)
created, err := manager.CreateNetwork(ctx, userID, &types.Network{
AccountID: accountID,
Name: "seq-allocation-test",
})
require.NoError(t, err)
require.NotEqual(t, "", created.PublicID, "CreateNetwork must allocate a non-zero AccountSeqID")
}
// Test_UpdateNetworkPreservesSeqID verifies UpdateNetwork does not reset
// AccountSeqID even when the caller passes a zero value (the shape REST
// handlers produce because the field is `json:"-"`).
func Test_UpdateNetworkPreservesPublicId(t *testing.T) {
ctx := context.Background()
const accountID = "testAccountId"
const userID = "testAdminId"
s, cleanUp, err := store.NewTestStoreFromSQL(ctx, "../testdata/networks.sql", t.TempDir())
require.NoError(t, err)
t.Cleanup(cleanUp)
am := mock_server.MockAccountManager{}
permissionsManager := permissions.NewManager(s)
groupsManager := groups.NewManagerMock()
routerManager := routers.NewManagerMock()
resourcesManager := resources.NewManager(s, permissionsManager, groupsManager, &am, nil)
manager := NewManager(s, permissionsManager, resourcesManager, routerManager, &am)
created, err := manager.CreateNetwork(ctx, userID, &types.Network{
AccountID: accountID,
Name: "seq-preserve-original",
})
require.NoError(t, err)
originalPublicId := created.PublicID
require.NotZero(t, originalPublicId)
update := &types.Network{
AccountID: accountID,
ID: created.ID,
Name: "seq-preserve-renamed",
}
require.Equal(t, "", update.PublicID, "incoming struct must mirror an HTTP handler shape")
_, err = manager.UpdateNetwork(ctx, userID, update)
require.NoError(t, err)
got, err := manager.GetNetwork(ctx, accountID, userID, created.ID)
require.NoError(t, err)
require.Equal(t, originalPublicId, got.PublicID, "PublicID must survive UpdateNetwork")
require.Equal(t, "seq-preserve-renamed", got.Name)
}

View File

@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"github.com/rs/xid"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
@@ -146,6 +147,8 @@ func (m *managerImpl) createResourceInTransaction(ctx context.Context, transacti
return nil, nil, fmt.Errorf("failed to get network: %w", err)
}
resource.PublicID = xid.New().String()
if err = transaction.SaveNetworkResource(ctx, resource); err != nil {
return nil, nil, fmt.Errorf("failed to save network resource: %w", err)
}
@@ -245,6 +248,7 @@ func (m *managerImpl) UpdateResource(ctx context.Context, userID string, resourc
if err != nil {
return fmt.Errorf("failed to get network resource: %w", err)
}
resource.PublicID = oldResource.PublicID
oldGroups, err := m.groupsManager.GetResourceGroupsInTransaction(ctx, transaction, store.LockingStrengthNone, resource.AccountID, resource.ID)
if err != nil {

View File

@@ -14,6 +14,7 @@ import (
nbDomain "github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkResourceType string
@@ -32,6 +33,7 @@ type NetworkResource struct {
ID string `gorm:"primaryKey"`
NetworkID string `gorm:"index"`
AccountID string `gorm:"index"`
PublicID string `json:"-"`
Name string
Description string
Type NetworkResourceType
@@ -63,6 +65,27 @@ func NewNetworkResource(accountID, networkID, name, description, address string,
}, nil
}
// ToComponent converts the resource to its self-contained components
// representation. Returns nil for a nil resource.
func (n *NetworkResource) ToComponent() *sharedTypes.ComponentResource {
if n == nil {
return nil
}
return &sharedTypes.ComponentResource{
ID: n.ID,
PublicID: n.PublicID,
NetworkID: n.NetworkID,
AccountID: n.AccountID,
Name: n.Name,
Description: n.Description,
Type: sharedTypes.ComponentResourceType(n.Type),
Address: n.Address,
Domain: n.Domain,
Prefix: n.Prefix,
Enabled: n.Enabled,
}
}
func (n *NetworkResource) ToAPIResponse(groups []api.GroupMinimum) *api.NetworkResource {
addr := n.Prefix.String()
if n.Type == Domain {
@@ -96,6 +119,7 @@ func (n *NetworkResource) Copy() *NetworkResource {
ID: n.ID,
AccountID: n.AccountID,
NetworkID: n.NetworkID,
PublicID: n.PublicID,
Name: n.Name,
Description: n.Description,
Type: n.Type,

View File

@@ -104,6 +104,8 @@ func (m *managerImpl) CreateRouter(ctx context.Context, userID string, router *t
router.ID = xid.New().String()
router.PublicID = xid.New().String()
err = transaction.CreateNetworkRouter(ctx, router)
if err != nil {
return fmt.Errorf("failed to create network router: %w", err)
@@ -199,6 +201,11 @@ func (m *managerImpl) updateRouterInTransaction(ctx context.Context, transaction
return nil, nil, affectedpeers.Change{}, status.NewRouterNotPartOfNetworkError(router.ID, router.NetworkID)
}
// Preserve PublicID from the existing router so the upstream
// UpdateNetworkRouter (which does Updates(router) with Select("*"))
// doesn't clobber it with the request's zero value.
router.PublicID = existing.PublicID
if err = transaction.UpdateNetworkRouter(ctx, router); err != nil {
return nil, nil, affectedpeers.Change{}, fmt.Errorf("failed to update network router: %w", err)
}

View File

@@ -7,12 +7,14 @@ import (
"github.com/netbirdio/netbird/management/server/networks/types"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkRouter struct {
ID string `gorm:"primaryKey"`
NetworkID string `gorm:"index"`
AccountID string `gorm:"index"`
PublicID string `json:"-"`
Peer string
PeerGroups []string `gorm:"serializer:json"`
Masquerade bool
@@ -20,6 +22,36 @@ type NetworkRouter struct {
Enabled bool
}
// ToComponent converts the router to its self-contained components
// representation. Returns nil for a nil router.
func (n *NetworkRouter) ToComponent() *sharedTypes.ComponentRouter {
if n == nil {
return nil
}
return &sharedTypes.ComponentRouter{
NetworkID: n.NetworkID,
PublicID: n.PublicID,
Peer: n.Peer,
PeerGroups: n.PeerGroups,
Masquerade: n.Masquerade,
Metric: n.Metric,
Enabled: n.Enabled,
}
}
// ToComponentMap converts a peer-keyed router map to its components
// representation.
func ToComponentMap(routers map[string]*NetworkRouter) map[string]*sharedTypes.ComponentRouter {
if routers == nil {
return nil
}
out := make(map[string]*sharedTypes.ComponentRouter, len(routers))
for id, r := range routers {
out[id] = r.ToComponent()
}
return out
}
func NewNetworkRouter(accountID string, networkID string, peer string, peerGroups []string, masquerade bool, metric int, enabled bool) (*NetworkRouter, error) {
r := &NetworkRouter{
ID: xid.New().String(),
@@ -81,6 +113,7 @@ func (n *NetworkRouter) Copy() *NetworkRouter {
ID: n.ID,
NetworkID: n.NetworkID,
AccountID: n.AccountID,
PublicID: n.PublicID,
Peer: n.Peer,
PeerGroups: n.PeerGroups,
Masquerade: n.Masquerade,

View File

@@ -7,8 +7,11 @@ import (
)
type Network struct {
ID string `gorm:"primaryKey"`
AccountID string `gorm:"index"`
ID string `gorm:"primaryKey"`
AccountID string `gorm:"index"`
PublicID string `json:"-"`
Name string
Description string
}
@@ -41,11 +44,12 @@ func (n *Network) FromAPIRequest(req *api.NetworkRequest) {
}
}
// Copy returns a copy of a posture checks.
// Copy returns a copy of a network.
func (n *Network) Copy() *Network {
return &Network{
ID: n.ID,
AccountID: n.AccountID,
PublicID: n.PublicID,
Name: n.Name,
Description: n.Description,
}

View File

@@ -405,7 +405,7 @@ func (am *DefaultAccountManager) CreatePeerJob(ctx context.Context, accountID, p
return status.NewPeerNotPartOfAccountError()
}
meetMinVer, err := posture.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
meetMinVer, err := version.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
if !version.IsDevelopmentVersion(p.Meta.WtVersion) && (!meetMinVer || err != nil) {
return status.Errorf(status.PreconditionFailed, "peer version %s does not meet the minimum required version %s for remote jobs", p.Meta.WtVersion, remoteJobsMinVer)
}
@@ -1588,7 +1588,7 @@ func affectedPeerIDsFromNetworkMap(nmap *types.NetworkMap, selfPeerID string) []
}
seen := make(map[string]struct{}, len(nmap.Peers)+len(nmap.OfflinePeers))
ids := make([]string, 0, len(nmap.Peers)+len(nmap.OfflinePeers))
add := func(peers []*nbpeer.Peer) {
add := func(peers []*types.ComponentPeer) {
for _, p := range peers {
if p == nil || p.ID == "" || p.ID == selfPeerID {
continue

View File

@@ -13,12 +13,14 @@ import (
"github.com/netbirdio/netbird/management/server/util"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
// Peer capability constants mirror the proto enum values.
const (
PeerCapabilitySourcePrefixes int32 = 1
PeerCapabilityIPv6Overlay int32 = 2
PeerCapabilitySourcePrefixes int32 = 1
PeerCapabilityIPv6Overlay int32 = 2
PeerCapabilityComponentNetworkMap int32 = 3
)
// Peer represents a machine connected to the network.
@@ -172,6 +174,7 @@ type PeerSystemMeta struct { //nolint:revive
Flags Flags `gorm:"serializer:json"`
Files []File `gorm:"serializer:json"`
Capabilities []int32 `gorm:"serializer:json"`
SyncMessageVersion int
}
func (p PeerSystemMeta) isEqual(other PeerSystemMeta) bool {
@@ -203,6 +206,35 @@ func (p *Peer) AddedWithSSOLogin() bool {
return p.UserID != ""
}
// ToComponent converts the peer to its self-contained components
// representation, carrying exactly the subset of peer data that crosses the
// components wire format. Returns nil for a nil peer so callers can convert
// possibly-missing peers without guarding.
func (p *Peer) ToComponent() *sharedTypes.ComponentPeer {
if p == nil {
return nil
}
cp := &sharedTypes.ComponentPeer{
ID: p.ID,
Key: p.Key,
IP: p.IP,
IPv6: p.IPv6,
DNSLabel: p.DNSLabel,
SSHKey: p.SSHKey,
SSHEnabled: p.SSHEnabled,
ServerSSHAllowed: p.Meta.Flags.ServerSSHAllowed,
AgentVersion: p.Meta.WtVersion,
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
SupportsIPv6: p.SupportsIPv6(),
LoginExpirationEnabled: p.LoginExpirationEnabled,
AddedWithSSOLogin: p.AddedWithSSOLogin(),
}
if p.LastLogin != nil {
cp.LastLogin = *p.LastLogin
}
return cp
}
// HasCapability reports whether the peer has the given capability.
func (p *Peer) HasCapability(capability int32) bool {
return slices.Contains(p.Meta.Capabilities, capability)
@@ -218,6 +250,14 @@ func (p *Peer) SupportsSourcePrefixes() bool {
return p.HasCapability(PeerCapabilitySourcePrefixes)
}
// SupportsComponentNetworkMap reports whether the peer assembles its
// NetworkMap from server-shipped components instead of consuming a fully
// expanded NetworkMap. Determines whether the network_map controller skips
// Calculate() server-side and emits the components envelope.
func (p *Peer) SupportsComponentNetworkMap() bool {
return p.HasCapability(PeerCapabilityComponentNetworkMap)
}
func capabilitiesEqual(a, b []int32) bool {
if len(a) != len(b) {
return false
@@ -406,6 +446,9 @@ func diffMeta(oldMeta, newMeta PeerSystemMeta, oldLocation, newLocation Location
if !sameMultiset(oldMeta.Files, newMeta.Files) {
add("files", fmt.Sprintf("%v", oldMeta.Files), fmt.Sprintf("%v", newMeta.Files))
}
if oldMeta.SyncMessageVersion != newMeta.SyncMessageVersion {
add("sync_meta_version", fmt.Sprintf("%d", oldMeta.SyncMessageVersion), fmt.Sprintf("%d", newMeta.SyncMessageVersion))
}
if !oldLocation.equal(newLocation) {
add("connection_ip", oldLocation.ConnectionIP, newLocation.ConnectionIP)

View File

@@ -1092,14 +1092,14 @@ func TestToSyncResponse(t *testing.T) {
}
networkMap := &types.NetworkMap{
Network: &types.Network{Net: *ipnet, Serial: 1000},
Peers: []*nbpeer.Peer{{
Peers: []*types.ComponentPeer{{
IP: netip.MustParseAddr("192.168.1.2"),
IPv6: netip.MustParseAddr("fd00::2"),
Key: "peer2-key",
DNSLabel: "peer2",
SSHEnabled: true,
SSHKey: "peer2-ssh-key"}},
OfflinePeers: []*nbpeer.Peer{{
OfflinePeers: []*types.ComponentPeer{{
IP: netip.MustParseAddr("192.168.1.3"),
IPv6: netip.MustParseAddr("fd00::3"),
Key: "peer3-key",

View File

@@ -67,10 +67,13 @@ func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, user
action = activity.PolicyUpdated
policy.PublicID = existingPolicy.PublicID
if err = transaction.SavePolicy(ctx, policy); err != nil {
return err
}
} else {
policy.PublicID = xid.New().String()
if err = transaction.CreatePolicy(ctx, policy); err != nil {
return err
}

View File

@@ -49,6 +49,8 @@ type Checks struct {
// AccountID is a reference to the Account that this object belongs
AccountID string `json:"-" gorm:"index"`
PublicID string `json:"-"`
// Checks is a set of objects that perform the actual checks
Checks ChecksDefinition `gorm:"serializer:json"`
}
@@ -167,6 +169,7 @@ func (pc *Checks) Copy() *Checks {
Name: pc.Name,
Description: pc.Description,
AccountID: pc.AccountID,
PublicID: pc.PublicID,
Checks: pc.Checks.Copy(),
}
return checks

View File

@@ -3,11 +3,9 @@ package posture
import (
"context"
"fmt"
"strings"
"github.com/hashicorp/go-version"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbversion "github.com/netbirdio/netbird/version"
)
type NBVersionCheck struct {
@@ -16,14 +14,8 @@ type NBVersionCheck struct {
var _ Check = (*NBVersionCheck)(nil)
// sanitizeVersion removes anything after the pre-release tag (e.g., "-dev", "-alpha", etc.)
func sanitizeVersion(version string) string {
parts := strings.Split(version, "-")
return parts[0]
}
func (n *NBVersionCheck) Check(ctx context.Context, peer nbpeer.Peer) (bool, error) {
meetsMin, err := MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
meetsMin, err := nbversion.MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
if err != nil {
return false, err
}
@@ -48,21 +40,3 @@ func (n *NBVersionCheck) Validate() error {
}
return nil
}
// MeetsMinVersion checks if the peer's version meets or exceeds the minimum required version
func MeetsMinVersion(minVer, peerVer string) (bool, error) {
peerVer = sanitizeVersion(peerVer)
minVer = sanitizeVersion(minVer)
peerNBVer, err := version.NewVersion(peerVer)
if err != nil {
return false, err
}
constraints, err := version.NewConstraint(">= " + minVer)
if err != nil {
return false, err
}
return constraints.Check(peerNBVer), nil
}

View File

@@ -139,68 +139,3 @@ func TestNBVersionCheck_Validate(t *testing.T) {
})
}
}
func TestMeetsMinVersion(t *testing.T) {
tests := []struct {
name string
minVer string
peerVer string
want bool
wantErr bool
}{
{
name: "Peer version greater than min version",
minVer: "0.26.0",
peerVer: "0.60.1",
want: true,
wantErr: false,
},
{
name: "Peer version equals min version",
minVer: "1.0.0",
peerVer: "1.0.0",
want: true,
wantErr: false,
},
{
name: "Peer version less than min version",
minVer: "1.0.0",
peerVer: "0.9.9",
want: false,
wantErr: false,
},
{
name: "Peer version with pre-release tag greater than min version",
minVer: "1.0.0",
peerVer: "1.0.1-alpha",
want: true,
wantErr: false,
},
{
name: "Invalid peer version format",
minVer: "1.0.0",
peerVer: "dev",
want: false,
wantErr: true,
},
{
name: "Invalid min version format",
minVer: "invalid.version",
peerVer: "1.0.0",
want: false,
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := MeetsMinVersion(tt.minVer, tt.peerVer)
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
assert.Equal(t, tt.want, got)
})
}
}

View File

@@ -52,7 +52,15 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
}
if isUpdate {
existing, err := transaction.GetPostureChecksByID(ctx, store.LockingStrengthNone, accountID, postureChecks.ID)
if err != nil {
return err
}
postureChecks.PublicID = existing.PublicID
action = activity.PostureCheckUpdated
} else {
postureChecks.PublicID = xid.New().String()
}
postureChecks.AccountID = accountID

View File

@@ -563,3 +563,61 @@ func TestArePostureCheckChangesAffectPeers(t *testing.T) {
assert.Empty(t, directPeerIDs)
})
}
// TestSavePostureChecks_AllocatesSeqIDOnCreate verifies that the create path
// (no incoming ID) allocates a non-zero AccountSeqID via the
// account_seq_counters table.
func TestSavePostureChecks_AllocatesSeqIDOnCreate(t *testing.T) {
am, _, err := createManager(t)
require.NoError(t, err)
account, err := initTestPostureChecksAccount(am)
require.NoError(t, err)
created, err := am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
Name: "seq-allocation-test",
Checks: posture.ChecksDefinition{
NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.26.0"},
},
}, true)
require.NoError(t, err)
require.NotEqual(t, "", created.PublicID, "SavePostureChecks on create must create PublicID")
}
// TestSavePostureChecks_PreservesSeqIDOnUpdate verifies the update path does
// not reset AccountSeqID even when the caller passes a zero value (REST
// handler shape, because the field is `json:"-"`).
func TestSavePostureChecks_PreservesSeqIDOnUpdate(t *testing.T) {
am, _, err := createManager(t)
require.NoError(t, err)
account, err := initTestPostureChecksAccount(am)
require.NoError(t, err)
created, err := am.SavePostureChecks(context.Background(), account.Id, adminUserID, &posture.Checks{
Name: "seq-preserve-original",
Checks: posture.ChecksDefinition{
NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.26.0"},
},
}, true)
require.NoError(t, err)
originalPublicID := created.PublicID
require.NotEqual(t, "", originalPublicID)
update := &posture.Checks{
ID: created.ID,
Name: "seq-preserve-renamed",
Checks: posture.ChecksDefinition{
NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.27.0"},
},
}
require.Equal(t, "", update.PublicID, "incoming struct must mirror an HTTP handler shape")
_, err = am.SavePostureChecks(context.Background(), account.Id, adminUserID, update, false)
require.NoError(t, err)
got, err := am.GetPostureChecks(context.Background(), account.Id, created.ID, adminUserID)
require.NoError(t, err)
require.Equal(t, originalPublicID, got.PublicID, "PublicID must survive SavePostureChecks update")
require.Equal(t, "seq-preserve-renamed", got.Name)
}

View File

@@ -175,6 +175,8 @@ func (am *DefaultAccountManager) CreateRoute(ctx context.Context, accountID stri
return err
}
newRoute.PublicID = xid.New().String()
if err = transaction.SaveRoute(ctx, newRoute); err != nil {
return err
}
@@ -222,6 +224,7 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
}
routeToSave.AccountID = accountID
routeToSave.PublicID = oldRoute.PublicID
if err = transaction.SaveRoute(ctx, routeToSave); err != nil {
return err

View File

@@ -6,6 +6,7 @@ import (
"encoding/json"
"errors"
"fmt"
"math"
"net"
"net/netip"
"net/url"
@@ -642,6 +643,22 @@ func (s *SqlStore) SaveUser(ctx context.Context, user *types.User) error {
}
// CreateGroups creates the given list of groups to the database.
// groupUpsertColumns is the explicit allowlist of columns that get updated when
// CreateGroups / UpdateGroups hit a PK conflict. public_id is intentionally
// omitted so a caller passing an entity with the zero value (e.g. an HTTP
// handler-built struct) cannot reset the persisted public_id during an upsert.
// Keep this in sync with the Group schema in management/server/types/group.go.
func groupUpsertColumns() clause.Set {
return clause.AssignmentColumns([]string{
"account_id",
"name",
"issued",
"integration_ref_id",
"integration_ref_integration_type",
"resources",
})
}
func (s *SqlStore) CreateGroups(ctx context.Context, accountID string, groups []*types.Group) error {
if len(groups) == 0 {
return nil
@@ -651,8 +668,9 @@ func (s *SqlStore) CreateGroups(ctx context.Context, accountID string, groups []
result := tx.
Clauses(
clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
UpdateAll: true,
DoUpdates: groupUpsertColumns(),
},
).
Omit(clause.Associations).
@@ -676,8 +694,9 @@ func (s *SqlStore) UpdateGroups(ctx context.Context, accountID string, groups []
result := tx.
Clauses(
clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}},
UpdateAll: true,
DoUpdates: groupUpsertColumns(),
},
).
Omit(clause.Associations).
@@ -1851,7 +1870,7 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
meta_kernel_version, meta_network_addresses, meta_system_serial_number, meta_system_product_name, meta_system_manufacturer,
meta_environment, meta_flags, meta_files, meta_capabilities, peer_status_last_seen, peer_status_session_started_at,
peer_status_connected, peer_status_login_expired, peer_status_requires_approval, location_connection_ip,
location_country_code, location_city_name, location_geo_name_id, proxy_meta_embedded, proxy_meta_cluster, ipv6
location_country_code, location_city_name, location_geo_name_id, proxy_meta_embedded, proxy_meta_cluster, ipv6, meta_sync_message_version
FROM peers WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
@@ -1873,6 +1892,7 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
metaSystemSerialNumber, metaSystemProductName, metaSystemManufacturer sql.NullString
locationCountryCode, locationCityName, proxyCluster sql.NullString
locationGeoNameID sql.NullInt64
metaSyncMessageVersion sql.NullInt32
)
err := row.Scan(&p.ID, &p.AccountID, &p.Key, &ip, &p.Name, &p.DNSLabel, &p.UserID, &p.SSHKey, &sshEnabled,
@@ -1882,7 +1902,7 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
&metaSystemSerialNumber, &metaSystemProductName, &metaSystemManufacturer, &env, &flags, &files, &capabilities,
&peerStatusLastSeen, &peerStatusSessionStartedAt, &peerStatusConnected, &peerStatusLoginExpired,
&peerStatusRequiresApproval, &connIP, &locationCountryCode, &locationCityName, &locationGeoNameID,
&proxyEmbedded, &proxyCluster, &ipv6)
&proxyEmbedded, &proxyCluster, &ipv6, &metaSyncMessageVersion)
if err == nil {
if lastLogin.Valid {
@@ -2002,6 +2022,9 @@ func (s *SqlStore) getPeers(ctx context.Context, accountID string) ([]nbpeer.Pee
if connIP != nil {
_ = json.Unmarshal(connIP, &p.Location.ConnectionIP)
}
if metaSyncMessageVersion.Valid {
p.Meta.SyncMessageVersion = int(metaSyncMessageVersion.Int32)
}
}
return p, err
})
@@ -2057,7 +2080,7 @@ func (s *SqlStore) getUsers(ctx context.Context, accountID string) ([]types.User
}
func (s *SqlStore) getGroups(ctx context.Context, accountID string) ([]*types.Group, error) {
const query = `SELECT id, account_id, name, issued, resources, integration_ref_id, integration_ref_integration_type FROM groups WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, name, issued, resources, integration_ref_id, integration_ref_integration_type FROM groups WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2067,7 +2090,7 @@ func (s *SqlStore) getGroups(ctx context.Context, accountID string) ([]*types.Gr
var resources []byte
var refID sql.NullInt64
var refType sql.NullString
err := row.Scan(&g.ID, &g.AccountID, &g.Name, &g.Issued, &resources, &refID, &refType)
err := row.Scan(&g.ID, &g.AccountID, &g.PublicID, &g.Name, &g.Issued, &resources, &refID, &refType)
if err == nil {
if refID.Valid {
g.IntegrationReference.ID = int(refID.Int64)
@@ -2092,7 +2115,7 @@ func (s *SqlStore) getGroups(ctx context.Context, accountID string) ([]*types.Gr
}
func (s *SqlStore) getPolicies(ctx context.Context, accountID string) ([]*types.Policy, error) {
const query = `SELECT id, account_id, name, description, enabled, source_posture_checks FROM policies WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, name, description, enabled, source_posture_checks FROM policies WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2101,7 +2124,7 @@ func (s *SqlStore) getPolicies(ctx context.Context, accountID string) ([]*types.
var p types.Policy
var checks []byte
var enabled sql.NullBool
err := row.Scan(&p.ID, &p.AccountID, &p.Name, &p.Description, &enabled, &checks)
err := row.Scan(&p.ID, &p.AccountID, &p.PublicID, &p.Name, &p.Description, &enabled, &checks)
if err == nil {
if enabled.Valid {
p.Enabled = enabled.Bool
@@ -2119,7 +2142,7 @@ func (s *SqlStore) getPolicies(ctx context.Context, accountID string) ([]*types.
}
func (s *SqlStore) getRoutes(ctx context.Context, accountID string) ([]route.Route, error) {
const query = `SELECT id, account_id, network, domains, keep_route, net_id, description, peer, peer_groups, network_type, masquerade, metric, enabled, groups, access_control_groups, skip_auto_apply FROM routes WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, network, domains, keep_route, net_id, description, peer, peer_groups, network_type, masquerade, metric, enabled, groups, access_control_groups, skip_auto_apply FROM routes WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2129,7 +2152,7 @@ func (s *SqlStore) getRoutes(ctx context.Context, accountID string) ([]route.Rou
var network, domains, peerGroups, groups, accessGroups []byte
var keepRoute, masquerade, enabled, skipAutoApply sql.NullBool
var metric sql.NullInt64
err := row.Scan(&r.ID, &r.AccountID, &network, &domains, &keepRoute, &r.NetID, &r.Description, &r.Peer, &peerGroups, &r.NetworkType, &masquerade, &metric, &enabled, &groups, &accessGroups, &skipAutoApply)
err := row.Scan(&r.ID, &r.AccountID, &r.PublicID, &network, &domains, &keepRoute, &r.NetID, &r.Description, &r.Peer, &peerGroups, &r.NetworkType, &masquerade, &metric, &enabled, &groups, &accessGroups, &skipAutoApply)
if err == nil {
if keepRoute.Valid {
r.KeepRoute = keepRoute.Bool
@@ -2171,7 +2194,7 @@ func (s *SqlStore) getRoutes(ctx context.Context, accountID string) ([]route.Rou
}
func (s *SqlStore) getNameServerGroups(ctx context.Context, accountID string) ([]nbdns.NameServerGroup, error) {
const query = `SELECT id, account_id, name, description, name_servers, groups, "primary", domains, enabled, search_domains_enabled FROM name_server_groups WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, name, description, name_servers, groups, "primary", domains, enabled, search_domains_enabled FROM name_server_groups WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2180,7 +2203,7 @@ func (s *SqlStore) getNameServerGroups(ctx context.Context, accountID string) ([
var n nbdns.NameServerGroup
var ns, groups, domains []byte
var primary, enabled, searchDomainsEnabled sql.NullBool
err := row.Scan(&n.ID, &n.AccountID, &n.Name, &n.Description, &ns, &groups, &primary, &domains, &enabled, &searchDomainsEnabled)
err := row.Scan(&n.ID, &n.AccountID, &n.PublicID, &n.Name, &n.Description, &ns, &groups, &primary, &domains, &enabled, &searchDomainsEnabled)
if err == nil {
if primary.Valid {
n.Primary = primary.Bool
@@ -2216,7 +2239,7 @@ func (s *SqlStore) getNameServerGroups(ctx context.Context, accountID string) ([
}
func (s *SqlStore) getPostureChecks(ctx context.Context, accountID string) ([]*posture.Checks, error) {
const query = `SELECT id, account_id, name, description, checks FROM posture_checks WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, name, description, checks FROM posture_checks WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2224,7 +2247,7 @@ func (s *SqlStore) getPostureChecks(ctx context.Context, accountID string) ([]*p
checks, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*posture.Checks, error) {
var c posture.Checks
var checksDef []byte
err := row.Scan(&c.ID, &c.AccountID, &c.Name, &c.Description, &checksDef)
err := row.Scan(&c.ID, &c.AccountID, &c.PublicID, &c.Name, &c.Description, &checksDef)
if err == nil && checksDef != nil {
_ = json.Unmarshal(checksDef, &c.Checks)
}
@@ -2236,117 +2259,30 @@ func (s *SqlStore) getPostureChecks(ctx context.Context, accountID string) ([]*p
return checks, nil
}
func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpservice.Service, error) {
const serviceQuery = `SELECT id, account_id, name, domain, enabled, auth,
meta_created_at, meta_certificate_issued_at, meta_status, proxy_cluster,
pass_host_header, rewrite_redirects, session_private_key, session_public_key,
mode, listen_port, port_auto_assigned, source, source_peer, terminated,
private, access_groups
FROM services WHERE account_id = $1`
// serviceSelectColumns and targetSelectColumns are the column lists the Postgres
// pgx read path scans. They must stay in sync with the rpservice.Service and
// rpservice.Target gorm models; TestPgxServiceColumnsMatchGorm enforces this.
const serviceSelectColumns = `id, account_id, name, domain, enabled, auth, restrictions,
meta_created_at, meta_certificate_issued_at, meta_last_renewed_at, meta_status, proxy_cluster,
pass_host_header, rewrite_redirects, session_private_key, session_public_key,
mode, listen_port, port_auto_assigned, source, source_peer, terminated,
private, access_groups`
const targetsQuery = `SELECT id, account_id, service_id, path, host, port, protocol,
target_id, target_type, enabled
FROM targets WHERE service_id = ANY($1)`
const targetSelectColumns = `id, account_id, service_id, path, host, port, protocol,
target_id, target_type, enabled, proxy_protocol,
skip_tls_verify, request_timeout, session_idle_timeout, path_rewrite, custom_headers,
direct_upstream, middlewares, capture_max_request_bytes, capture_max_response_bytes,
capture_content_types, agent_network, disable_access_log`
func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpservice.Service, error) {
const serviceQuery = `SELECT ` + serviceSelectColumns + ` FROM services WHERE account_id = $1`
serviceRows, err := s.pool.Query(ctx, serviceQuery, accountID)
if err != nil {
return nil, err
}
services, err := pgx.CollectRows(serviceRows, func(row pgx.CollectableRow) (*rpservice.Service, error) {
var s rpservice.Service
var auth []byte
var accessGroups []byte
var createdAt, certIssuedAt sql.NullTime
var status, proxyCluster, sessionPrivateKey, sessionPublicKey sql.NullString
var mode, source, sourcePeer sql.NullString
var terminated, portAutoAssigned, private sql.NullBool
var listenPort sql.NullInt64
err := row.Scan(
&s.ID,
&s.AccountID,
&s.Name,
&s.Domain,
&s.Enabled,
&auth,
&createdAt,
&certIssuedAt,
&status,
&proxyCluster,
&s.PassHostHeader,
&s.RewriteRedirects,
&sessionPrivateKey,
&sessionPublicKey,
&mode,
&listenPort,
&portAutoAssigned,
&source,
&sourcePeer,
&terminated,
&private,
&accessGroups,
)
if err != nil {
return nil, err
}
if auth != nil {
if err := json.Unmarshal(auth, &s.Auth); err != nil {
return nil, err
}
}
if len(accessGroups) > 0 {
if err := json.Unmarshal(accessGroups, &s.AccessGroups); err != nil {
return nil, fmt.Errorf("unmarshal access_groups: %w", err)
}
}
if private.Valid {
s.Private = private.Bool
}
s.Meta = rpservice.Meta{}
if createdAt.Valid {
s.Meta.CreatedAt = createdAt.Time
}
if certIssuedAt.Valid {
t := certIssuedAt.Time
s.Meta.CertificateIssuedAt = &t
}
if status.Valid {
s.Meta.Status = status.String
}
if proxyCluster.Valid {
s.ProxyCluster = proxyCluster.String
}
if sessionPrivateKey.Valid {
s.SessionPrivateKey = sessionPrivateKey.String
}
if sessionPublicKey.Valid {
s.SessionPublicKey = sessionPublicKey.String
}
if mode.Valid {
s.Mode = mode.String
}
if source.Valid {
s.Source = source.String
}
if sourcePeer.Valid {
s.SourcePeer = sourcePeer.String
}
if terminated.Valid {
s.Terminated = terminated.Bool
}
if portAutoAssigned.Valid {
s.PortAutoAssigned = portAutoAssigned.Bool
}
if listenPort.Valid {
s.ListenPort = uint16(listenPort.Int64)
}
s.Targets = []*rpservice.Target{}
return &s, nil
})
services, err := pgx.CollectRows(serviceRows, scanService)
if err != nil {
return nil, err
}
@@ -2357,39 +2293,12 @@ func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpserv
serviceIDs := make([]string, len(services))
serviceMap := make(map[string]*rpservice.Service)
for i, s := range services {
serviceIDs[i] = s.ID
serviceMap[s.ID] = s
for i, svc := range services {
serviceIDs[i] = svc.ID
serviceMap[svc.ID] = svc
}
targetRows, err := s.pool.Query(ctx, targetsQuery, serviceIDs)
if err != nil {
return nil, err
}
targets, err := pgx.CollectRows(targetRows, func(row pgx.CollectableRow) (*rpservice.Target, error) {
var t rpservice.Target
var path sql.NullString
err := row.Scan(
&t.ID,
&t.AccountID,
&t.ServiceID,
&path,
&t.Host,
&t.Port,
&t.Protocol,
&t.TargetId,
&t.TargetType,
&t.Enabled,
)
if err != nil {
return nil, err
}
if path.Valid {
t.Path = &path.String
}
return &t, nil
})
targets, err := s.getServiceTargets(ctx, serviceIDs)
if err != nil {
return nil, err
}
@@ -2403,8 +2312,203 @@ func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpserv
return services, nil
}
func scanService(row pgx.CollectableRow) (*rpservice.Service, error) {
var s rpservice.Service
var auth []byte
var restrictions []byte
var accessGroups []byte
var createdAt, certIssuedAt, lastRenewedAt sql.NullTime
var status, proxyCluster, sessionPrivateKey, sessionPublicKey sql.NullString
var mode, source, sourcePeer sql.NullString
var terminated, portAutoAssigned, private sql.NullBool
var listenPort sql.NullInt64
err := row.Scan(
&s.ID,
&s.AccountID,
&s.Name,
&s.Domain,
&s.Enabled,
&auth,
&restrictions,
&createdAt,
&certIssuedAt,
&lastRenewedAt,
&status,
&proxyCluster,
&s.PassHostHeader,
&s.RewriteRedirects,
&sessionPrivateKey,
&sessionPublicKey,
&mode,
&listenPort,
&portAutoAssigned,
&source,
&sourcePeer,
&terminated,
&private,
&accessGroups,
)
if err != nil {
return nil, err
}
if auth != nil {
if err := json.Unmarshal(auth, &s.Auth); err != nil {
return nil, err
}
}
if len(restrictions) > 0 {
if err := json.Unmarshal(restrictions, &s.Restrictions); err != nil {
return nil, fmt.Errorf("unmarshal restrictions: %w", err)
}
}
if len(accessGroups) > 0 {
if err := json.Unmarshal(accessGroups, &s.AccessGroups); err != nil {
return nil, fmt.Errorf("unmarshal access_groups: %w", err)
}
}
if private.Valid {
s.Private = private.Bool
}
s.Meta = serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt, status)
if proxyCluster.Valid {
s.ProxyCluster = proxyCluster.String
}
if sessionPrivateKey.Valid {
s.SessionPrivateKey = sessionPrivateKey.String
}
if sessionPublicKey.Valid {
s.SessionPublicKey = sessionPublicKey.String
}
if mode.Valid {
s.Mode = mode.String
}
if source.Valid {
s.Source = source.String
}
if sourcePeer.Valid {
s.SourcePeer = sourcePeer.String
}
if terminated.Valid {
s.Terminated = terminated.Bool
}
if portAutoAssigned.Valid {
s.PortAutoAssigned = portAutoAssigned.Bool
}
if listenPort.Valid {
if listenPort.Int64 < 0 || listenPort.Int64 > math.MaxUint16 {
return nil, fmt.Errorf("listen_port %d out of range", listenPort.Int64)
}
s.ListenPort = uint16(listenPort.Int64)
}
s.Targets = []*rpservice.Target{}
return &s, nil
}
func serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt sql.NullTime, status sql.NullString) rpservice.Meta {
meta := rpservice.Meta{}
if createdAt.Valid {
meta.CreatedAt = createdAt.Time
}
if certIssuedAt.Valid {
t := certIssuedAt.Time
meta.CertificateIssuedAt = &t
}
if lastRenewedAt.Valid {
t := lastRenewedAt.Time
meta.LastRenewedAt = &t
}
if status.Valid {
meta.Status = status.String
}
return meta
}
func (s *SqlStore) getServiceTargets(ctx context.Context, serviceIDs []string) ([]*rpservice.Target, error) {
const targetsQuery = `SELECT ` + targetSelectColumns + ` FROM targets WHERE service_id = ANY($1)`
rows, err := s.pool.Query(ctx, targetsQuery, serviceIDs)
if err != nil {
return nil, err
}
return pgx.CollectRows(rows, scanTarget)
}
func scanTarget(row pgx.CollectableRow) (*rpservice.Target, error) {
var t rpservice.Target
var path sql.NullString
var pathRewrite sql.NullString
var proxyProtocol, skipTLSVerify, directUpstream, agentNetwork, disableAccessLog sql.NullBool
var requestTimeout, sessionIdleTimeout, captureMaxRequestBytes, captureMaxResponseBytes sql.NullInt64
var customHeaders, middlewares, captureContentTypes []byte
err := row.Scan(
&t.ID,
&t.AccountID,
&t.ServiceID,
&path,
&t.Host,
&t.Port,
&t.Protocol,
&t.TargetId,
&t.TargetType,
&t.Enabled,
&proxyProtocol,
&skipTLSVerify,
&requestTimeout,
&sessionIdleTimeout,
&pathRewrite,
&customHeaders,
&directUpstream,
&middlewares,
&captureMaxRequestBytes,
&captureMaxResponseBytes,
&captureContentTypes,
&agentNetwork,
&disableAccessLog,
)
if err != nil {
return nil, err
}
if path.Valid {
t.Path = &path.String
}
t.ProxyProtocol = proxyProtocol.Bool
t.Options.SkipTLSVerify = skipTLSVerify.Bool
t.Options.RequestTimeout = time.Duration(requestTimeout.Int64)
t.Options.SessionIdleTimeout = time.Duration(sessionIdleTimeout.Int64)
t.Options.PathRewrite = rpservice.PathRewriteMode(pathRewrite.String)
t.Options.DirectUpstream = directUpstream.Bool
t.Options.CaptureMaxRequestBytes = captureMaxRequestBytes.Int64
t.Options.CaptureMaxResponseBytes = captureMaxResponseBytes.Int64
t.Options.AgentNetwork = agentNetwork.Bool
t.Options.DisableAccessLog = disableAccessLog.Bool
if len(customHeaders) > 0 {
if err := json.Unmarshal(customHeaders, &t.Options.CustomHeaders); err != nil {
return nil, fmt.Errorf("unmarshal custom_headers: %w", err)
}
}
if len(middlewares) > 0 {
if err := json.Unmarshal(middlewares, &t.Options.Middlewares); err != nil {
return nil, fmt.Errorf("unmarshal middlewares: %w", err)
}
}
if len(captureContentTypes) > 0 {
if err := json.Unmarshal(captureContentTypes, &t.Options.CaptureContentTypes); err != nil {
return nil, fmt.Errorf("unmarshal capture_content_types: %w", err)
}
}
return &t, nil
}
func (s *SqlStore) getNetworks(ctx context.Context, accountID string) ([]*networkTypes.Network, error) {
const query = `SELECT id, account_id, name, description FROM networks WHERE account_id = $1`
const query = `SELECT id, account_id, public_id, name, description FROM networks WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2421,7 +2525,7 @@ func (s *SqlStore) getNetworks(ctx context.Context, accountID string) ([]*networ
}
func (s *SqlStore) getNetworkRouters(ctx context.Context, accountID string) ([]*routerTypes.NetworkRouter, error) {
const query = `SELECT id, network_id, account_id, peer, peer_groups, masquerade, metric, enabled FROM network_routers WHERE account_id = $1`
const query = `SELECT id, network_id, account_id, public_id, peer, peer_groups, masquerade, metric, enabled FROM network_routers WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2431,7 +2535,7 @@ func (s *SqlStore) getNetworkRouters(ctx context.Context, accountID string) ([]*
var peerGroups []byte
var masquerade, enabled sql.NullBool
var metric sql.NullInt64
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.Peer, &peerGroups, &masquerade, &metric, &enabled)
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Peer, &peerGroups, &masquerade, &metric, &enabled)
if err == nil {
if masquerade.Valid {
r.Masquerade = masquerade.Bool
@@ -2459,7 +2563,7 @@ func (s *SqlStore) getNetworkRouters(ctx context.Context, accountID string) ([]*
}
func (s *SqlStore) getNetworkResources(ctx context.Context, accountID string) ([]*resourceTypes.NetworkResource, error) {
const query = `SELECT id, network_id, account_id, name, description, type, domain, prefix, enabled FROM network_resources WHERE account_id = $1`
const query = `SELECT id, network_id, account_id, public_id, name, description, type, domain, prefix, enabled FROM network_resources WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)
if err != nil {
return nil, err
@@ -2468,7 +2572,7 @@ func (s *SqlStore) getNetworkResources(ctx context.Context, accountID string) ([
var r resourceTypes.NetworkResource
var prefix []byte
var enabled sql.NullBool
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.Name, &r.Description, &r.Type, &r.Domain, &prefix, &enabled)
err := row.Scan(&r.ID, &r.NetworkID, &r.AccountID, &r.PublicID, &r.Name, &r.Description, &r.Type, &r.Domain, &prefix, &enabled)
if err == nil {
if enabled.Valid {
r.Enabled = enabled.Bool
@@ -3830,7 +3934,7 @@ func (s *SqlStore) UpdateGroup(ctx context.Context, group *types.Group) error {
return status.Errorf(status.InvalidArgument, "group is nil")
}
if err := s.db.Omit(clause.Associations).Save(group).Error; err != nil {
if err := s.db.Omit(clause.Associations, "public_id").Save(group).Error; err != nil {
log.WithContext(ctx).Errorf("failed to save group to store: %v", err)
return status.Errorf(status.Internal, "failed to save group to store")
}
@@ -3918,7 +4022,7 @@ func (s *SqlStore) CreatePolicy(ctx context.Context, policy *types.Policy) error
// SavePolicy saves a policy to the database.
func (s *SqlStore) SavePolicy(ctx context.Context, policy *types.Policy) error {
result := s.db.Session(&gorm.Session{FullSaveAssociations: true}).Save(policy)
result := s.db.Session(&gorm.Session{FullSaveAssociations: true}).Omit("public_id").Save(policy)
if err := result.Error; err != nil {
log.WithContext(ctx).Errorf("failed to save policy to the store: %s", err)
return status.Errorf(status.Internal, "failed to save policy to store")
@@ -6118,6 +6222,37 @@ func (s *SqlStore) DisconnectProxy(ctx context.Context, proxyID, sessionID strin
return nil
}
// GetAllProxies returns all reverse proxy instance rows.
func (s *SqlStore) GetAllProxies(ctx context.Context) ([]*proxy.Proxy, error) {
var proxies []*proxy.Proxy
result := s.db.Order("cluster_address, id").Find(&proxies)
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to get proxies: %v", result.Error)
return nil, status.Errorf(status.Internal, "failed to get proxies")
}
return proxies, nil
}
// DisconnectAllProxies force-marks every proxy that is not already disconnected
// as disconnected, regardless of session ID. Unlike DisconnectProxy it is not
// session-guarded: it is an administrative repair helper, not part of the
// connection lifecycle. last_seen is left untouched so the stale-proxy reaper
// keeps working off the real last heartbeat. Returns the number of proxies updated.
func (s *SqlStore) DisconnectAllProxies(ctx context.Context) (int64, error) {
result := s.db.
Model(&proxy.Proxy{}).
Where("status != ?", proxy.StatusDisconnected).
Updates(map[string]any{
"status": proxy.StatusDisconnected,
"disconnected_at": time.Now(),
})
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to disconnect all proxies: %v", result.Error)
return 0, status.Errorf(status.Internal, "failed to disconnect all proxies")
}
return result.RowsAffected, nil
}
// UpdateProxyHeartbeat updates the last_seen timestamp for the proxy's current session.
func (s *SqlStore) UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error {
now := time.Now()
@@ -6125,7 +6260,11 @@ func (s *SqlStore) UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) err
result := s.db.
Model(&proxy.Proxy{}).
Where("id = ? AND session_id = ?", p.ID, p.SessionID).
Update("last_seen", now)
Updates(map[string]any{
"last_seen": now,
"status": proxy.StatusConnected,
"disconnected_at": nil,
})
if result.Error != nil {
log.WithContext(ctx).Errorf("failed to update proxy heartbeat: %v", result.Error)

View File

@@ -71,7 +71,7 @@ func (s *SqlStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetr
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()
"COALESCE(SUM" + agentNetworkTypes.CostUSDSQLExpr + ", 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)
}

View File

@@ -37,7 +37,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
InputCostUSD: 0.0231,
}
usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -71,7 +71,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
InputCostUSD: 0.0231,
}
entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -127,7 +127,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
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,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, InputCostUSD: cost,
}
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
@@ -143,7 +143,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
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.InDelta(t, 0.30, buckets[0].TotalCostUSD(), 1e-9, "same-day cost summed")
assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
assert.Equal(t, int64(15), buckets[1].TotalTokens)
@@ -174,7 +174,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
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,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, InputCostUSD: cost,
}
}
@@ -207,7 +207,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
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.InDelta(t, 0.30, s1.TotalCostUSD(), 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

View File

@@ -0,0 +1,74 @@
package store
import (
"strings"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm/schema"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
)
// TestPgxServiceColumnsMatchGorm guards the Postgres pgx read path against
// drifting from the gorm model. The SQLite/MySQL gorm path loads rows by struct,
// so a new column on a model is picked up automatically, but the hand-written
// pgx SELECT in sql_store.go must be updated by hand. This test fails when a
// gorm column is missing from the pgx column list, which otherwise silently
// returns zero-valued on Postgres with no compile error.
func TestPgxServiceColumnsMatchGorm(t *testing.T) {
tests := []struct {
name string
model any
selectColumns string
// excluded lists gorm columns intentionally not loaded by the pgx path.
excluded map[string]struct{}
}{
{
name: "service",
model: &rpservice.Service{},
selectColumns: serviceSelectColumns,
},
{
name: "target",
model: &rpservice.Target{},
selectColumns: targetSelectColumns,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
selected := parseColumnList(tc.selectColumns)
for _, col := range gormColumnNames(t, tc.model) {
if _, ok := tc.excluded[col]; ok {
continue
}
_, ok := selected[col]
assert.Truef(t, ok,
"gorm column %q is not read by the Postgres pgx SELECT; add it to %sSelectColumns in sql_store.go (or to the test's excluded set if it is intentionally not loaded)",
col, tc.name)
}
})
}
}
func parseColumnList(cols string) map[string]struct{} {
set := make(map[string]struct{})
for _, c := range strings.Split(cols, ",") {
if c = strings.TrimSpace(c); c != "" {
set[c] = struct{}{}
}
}
return set
}
// gormColumnNames returns the DB column names gorm would migrate for the model,
// using the same default naming strategy the store configures.
func gormColumnNames(t *testing.T, model any) []string {
t.Helper()
sch, err := schema.Parse(model, &sync.Map{}, schema.NamingStrategy{})
require.NoError(t, err)
return sch.DBNames
}

View File

@@ -0,0 +1,156 @@
package store
import (
"context"
"os"
"runtime"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
)
// TestSqlStore_DisconnectAllProxies guards the administrative
// force-disconnect helper:
//
// 1. Every proxy that is not already disconnected is marked
// disconnected regardless of its session ID (unlike
// DisconnectProxy, which is session-guarded).
// 2. Rows that are already disconnected are left untouched, so their
// original disconnected_at is preserved and the returned count
// reflects only the rows that actually changed.
// 3. last_seen is not modified — the stale-proxy reaper keeps working
// off the real last heartbeat.
func TestSqlStore_DisconnectAllProxies(t *testing.T) {
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
t.Skip("skip CI tests on darwin and windows")
}
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
lastSeenFresh := time.Now().Add(-30 * time.Second)
lastSeenStale := time.Now().Add(-30 * time.Minute)
oldDisconnectedAt := time.Now().Add(-time.Hour)
accountID := "acct-disconnect"
proxies := []*rpproxy.Proxy{
{
ID: "p-connected-fresh",
SessionID: "sess-1",
ClusterAddress: "cluster-a.example.com",
IPAddress: "10.0.0.1",
LastSeen: lastSeenFresh,
Status: rpproxy.StatusConnected,
},
{
ID: "p-connected-stale",
SessionID: "sess-2",
ClusterAddress: "cluster-b.example.com",
IPAddress: "10.0.0.2",
AccountID: &accountID,
LastSeen: lastSeenStale,
Status: rpproxy.StatusConnected,
},
{
ID: "p-already-disconnected",
SessionID: "sess-3",
ClusterAddress: "cluster-a.example.com",
IPAddress: "10.0.0.3",
LastSeen: lastSeenStale,
Status: rpproxy.StatusDisconnected,
DisconnectedAt: &oldDisconnectedAt,
},
}
for _, p := range proxies {
require.NoError(t, store.SaveProxy(ctx, p))
}
all, err := store.GetAllProxies(ctx)
require.NoError(t, err)
require.Len(t, all, 3)
disconnected, err := store.DisconnectAllProxies(ctx)
require.NoError(t, err)
assert.Equal(t, int64(2), disconnected)
all, err = store.GetAllProxies(ctx)
require.NoError(t, err)
require.Len(t, all, 3)
byID := make(map[string]*rpproxy.Proxy, len(all))
for _, p := range all {
byID[p.ID] = p
}
for id, p := range byID {
assert.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", id)
require.NotNil(t, p.DisconnectedAt, "proxy %s should have disconnected_at set", id)
}
// force-marked rows carry a fresh disconnected_at; the untouched row keeps its original one
assert.WithinDuration(t, time.Now(), *byID["p-connected-fresh"].DisconnectedAt, 10*time.Second)
assert.WithinDuration(t, time.Now(), *byID["p-connected-stale"].DisconnectedAt, 10*time.Second)
assert.WithinDuration(t, oldDisconnectedAt, *byID["p-already-disconnected"].DisconnectedAt, time.Second)
// last_seen is preserved so the stale reaper schedule is unaffected
assert.WithinDuration(t, lastSeenFresh, byID["p-connected-fresh"].LastSeen, time.Second)
assert.WithinDuration(t, lastSeenStale, byID["p-connected-stale"].LastSeen, time.Second)
// idempotent: a second run has nothing left to update
disconnected, err = store.DisconnectAllProxies(ctx)
require.NoError(t, err)
assert.Equal(t, int64(0), disconnected)
})
}
func TestSqlStore_UpdateProxyHeartbeatRestoresDisconnectedCurrentSession(t *testing.T) {
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
t.Skip("skip CI tests on darwin and windows")
}
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
proxy := &rpproxy.Proxy{
ID: "p-heartbeat",
SessionID: "sess-heartbeat",
ClusterAddress: "cluster-heartbeat.example.com",
IPAddress: "10.0.0.10",
LastSeen: time.Now().Add(-30 * time.Second),
Status: rpproxy.StatusConnected,
}
require.NoError(t, store.SaveProxy(ctx, proxy))
disconnected, err := store.DisconnectAllProxies(ctx)
require.NoError(t, err)
require.Equal(t, int64(1), disconnected)
require.NoError(t, store.UpdateProxyHeartbeat(ctx, &rpproxy.Proxy{ID: proxy.ID, SessionID: proxy.SessionID}))
all, err := store.GetAllProxies(ctx)
require.NoError(t, err)
require.Len(t, all, 1)
assert.Equal(t, rpproxy.StatusConnected, all[0].Status)
assert.Nil(t, all[0].DisconnectedAt)
assert.WithinDuration(t, time.Now(), all[0].LastSeen, 10*time.Second)
addresses, err := store.GetActiveProxyClusterAddresses(ctx)
require.NoError(t, err)
assert.Contains(t, addresses, proxy.ClusterAddress)
})
}
func TestSqlStore_GetAllProxies_Empty(t *testing.T) {
if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
t.Skip("skip CI tests on darwin and windows")
}
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
all, err := store.GetAllProxies(context.Background())
require.NoError(t, err)
assert.Empty(t, all)
})
}

View File

@@ -5,6 +5,7 @@ import (
"os"
"runtime"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -44,3 +45,91 @@ func TestSqlStore_GetAccount_PrivateServiceRoundtrip(t *testing.T) {
assert.Equal(t, []string{"grp-admins", "grp-ops"}, got.AccessGroups)
})
}
// TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip guards the Postgres pgx
// read path (getServices) against silently dropping columns present on the gorm
// model. Before the fix these fields loaded correctly on SQLite but came back
// zero-valued on Postgres because the hand-written SELECT and scan omitted them.
func TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip(t *testing.T) {
if os.Getenv("CI") == "true" && (runtime.GOOS == "darwin" || runtime.GOOS == "windows") {
t.Skip("skip CI tests on darwin and windows")
}
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
account := newAccountWithId(ctx, "account_svc_opts", "testuser", "")
require.NoError(t, store.SaveAccount(ctx, account))
renewedAt := time.Now().UTC().Truncate(time.Second)
targetPath := "/api"
svc := &rpservice.Service{
ID: "svc-opts",
AccountID: account.Id,
Name: "opts-svc",
Domain: "opts.example",
Enabled: true,
Mode: rpservice.ModeHTTP,
Restrictions: rpservice.AccessRestrictions{
AllowedCIDRs: []string{"10.0.0.0/8"},
BlockedCountries: []string{"XX"},
CrowdSecMode: "block",
},
Meta: rpservice.Meta{
LastRenewedAt: &renewedAt,
},
Targets: []*rpservice.Target{
{
AccountID: account.Id,
ServiceID: "svc-opts",
Path: &targetPath,
Host: "backend.internal",
Port: 8080,
Protocol: "http",
TargetId: "tgt-1",
Enabled: true,
ProxyProtocol: true,
Options: rpservice.TargetOptions{
SkipTLSVerify: true,
RequestTimeout: 30 * time.Second,
SessionIdleTimeout: 5 * time.Minute,
PathRewrite: rpservice.PathRewritePreserve,
CustomHeaders: map[string]string{"X-Foo": "bar"},
DirectUpstream: true,
CaptureMaxRequestBytes: 1024,
CaptureMaxResponseBytes: 2048,
CaptureContentTypes: []string{"application/json"},
AgentNetwork: true,
DisableAccessLog: true,
},
},
},
}
require.NoError(t, store.CreateService(ctx, svc))
loaded, err := store.GetAccount(ctx, account.Id)
require.NoError(t, err)
require.Len(t, loaded.Services, 1)
got := loaded.Services[0]
assert.Equal(t, []string{"10.0.0.0/8"}, got.Restrictions.AllowedCIDRs, "restrictions allowed CIDRs")
assert.Equal(t, []string{"XX"}, got.Restrictions.BlockedCountries, "restrictions blocked countries")
assert.Equal(t, "block", got.Restrictions.CrowdSecMode, "restrictions crowdsec mode")
require.NotNil(t, got.Meta.LastRenewedAt, "meta last renewed at")
assert.WithinDuration(t, renewedAt, *got.Meta.LastRenewedAt, time.Second, "meta last renewed at")
require.Len(t, got.Targets, 1)
tg := got.Targets[0]
assert.True(t, tg.ProxyProtocol, "target proxy protocol")
assert.True(t, tg.Options.SkipTLSVerify, "options skip TLS verify")
assert.Equal(t, 30*time.Second, tg.Options.RequestTimeout, "options request timeout")
assert.Equal(t, 5*time.Minute, tg.Options.SessionIdleTimeout, "options session idle timeout")
assert.Equal(t, rpservice.PathRewritePreserve, tg.Options.PathRewrite, "options path rewrite")
assert.Equal(t, map[string]string{"X-Foo": "bar"}, tg.Options.CustomHeaders, "options custom headers")
assert.True(t, tg.Options.DirectUpstream, "options direct upstream")
assert.Equal(t, int64(1024), tg.Options.CaptureMaxRequestBytes, "options capture max request bytes")
assert.Equal(t, int64(2048), tg.Options.CaptureMaxResponseBytes, "options capture max response bytes")
assert.Equal(t, []string{"application/json"}, tg.Options.CaptureContentTypes, "options capture content types")
assert.True(t, tg.Options.AgentNetwork, "options agent network")
assert.True(t, tg.Options.DisableAccessLog, "options disable access log")
})
}

View File

@@ -47,6 +47,7 @@ func runTestForAllEngines(t *testing.T, testDataFile string, f func(t *testing.T
}
t.Setenv("NETBIRD_STORE_ENGINE", string(engine))
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), testDataFile, t.TempDir())
assert.NoError(t, err, "engine: ", string(engine))
t.Cleanup(cleanUp)
assert.NoError(t, err)
t.Run(string(engine), func(t *testing.T) {
@@ -561,53 +562,60 @@ func TestSqlStore_GetPeerByIP_NotFound(t *testing.T) {
}
func TestSqlStore_SavePeer(t *testing.T) {
store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
t.Cleanup(cleanUp)
assert.NoError(t, err)
populateFields := testing_helpers.NewPopulateFields()
account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
require.NoError(t, err)
runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
require.NoError(t, err)
// save status of non-existing peer
peer := &nbpeer.Peer{
Key: "peerkey",
ID: "testpeer",
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
IPv6: netip.MustParseAddr("fd00::1"),
Meta: nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
Name: "peer name",
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
CreatedAt: time.Now().UTC(),
}
ctx := context.Background()
err = store.SavePeer(ctx, 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")
metadata := nbpeer.PeerSystemMeta{}
reflectedMetadata := reflect.ValueOf(&metadata).Elem()
// save new status of existing peer
account.Peers[peer.ID] = peer
numOfFields, err := populateFields.PopulateAll(reflectedMetadata)
assert.NoError(t, err)
assert.Equal(t, 32, numOfFields)
err = store.SaveAccount(context.Background(), account)
require.NoError(t, err)
// save status of non-existing peer
peer := &nbpeer.Peer{
Key: "peerkey",
ID: "testpeer",
IP: netip.AddrFrom4([4]byte{127, 0, 0, 1}),
IPv6: netip.MustParseAddr("fd00::1"),
Meta: metadata, //nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
Name: "peer name",
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
CreatedAt: time.Now().UTC(),
}
ctx := context.Background()
err = store.SavePeer(ctx, 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")
updatedPeer := peer.Copy()
updatedPeer.Status.Connected = false
updatedPeer.Meta.Hostname = "updatedpeer"
// save new status of existing peer
account.Peers[peer.ID] = peer
err = store.SavePeer(ctx, account.Id, updatedPeer)
require.NoError(t, err)
err = store.SaveAccount(context.Background(), account)
require.NoError(t, err)
account, err = store.GetAccount(context.Background(), account.Id)
require.NoError(t, err)
updatedPeer := peer.Copy()
updatedPeer.Status.Connected = false
updatedPeer.Meta.Hostname = "updatedpeer"
actual := account.Peers[peer.ID]
assert.Equal(t, updatedPeer.Meta, actual.Meta)
assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
err = store.SavePeer(ctx, account.Id, updatedPeer)
require.NoError(t, err)
account, err = store.GetAccount(context.Background(), account.Id)
require.NoError(t, err)
actual := account.Peers[peer.ID]
assert.Equal(t, updatedPeer.Meta, actual.Meta)
assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
})
}
func TestSqlStore_SavePeerStatus(t *testing.T) {

View File

@@ -323,6 +323,8 @@ type Store interface {
GetClusterSupportsCrowdSec(ctx context.Context, clusterAddr string) *bool
GetClusterSupportsPrivate(ctx context.Context, clusterAddr string) *bool
CleanupStaleProxies(ctx context.Context, inactivityDuration time.Duration) error
GetAllProxies(ctx context.Context) ([]*proxy.Proxy, error)
DisconnectAllProxies(ctx context.Context) (int64, error)
GetProxyByAccountID(ctx context.Context, accountID string) (*proxy.Proxy, error)
CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error)
IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error)
@@ -582,6 +584,30 @@ func getMigrationsPreAuto(ctx context.Context) []migrationFunc {
func(db *gorm.DB) error {
return migration.CleanupOrphanedResources[domain.Domain, types.Account](ctx, db, "account_id")
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[types.Policy](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[types.Group](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[route.Route](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[resourceTypes.NetworkResource](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[routerTypes.NetworkRouter](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[dns.NameServerGroup](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[networkTypes.Network](ctx, db)
},
func(db *gorm.DB) error {
return migration.BackfillPublicIDs[posture.Checks](ctx, db)
},
}
}
@@ -624,6 +650,14 @@ func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
func(db *gorm.DB) error {
return migration.DropIndex[proxy.Proxy](ctx, db, "idx_proxy_account_id_unique")
},
// Post-auto so the per-bucket cost columns already exist when the legacy
// aggregates are folded into them and dropped.
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkAccessLog](ctx, db)
},
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db)
},
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -1,495 +0,0 @@
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)
}

View File

@@ -10,19 +10,20 @@ import (
// UpdateChannelMetrics represents all metrics related to the UpdateChannel
type UpdateChannelMetrics struct {
createChannelDurationMicro metric.Int64Histogram
closeChannelDurationMicro metric.Int64Histogram
closeChannelsDurationMicro metric.Int64Histogram
closeChannels metric.Int64Histogram
sendUpdateDurationMicro metric.Int64Histogram
getAllConnectedPeersDurationMicro metric.Int64Histogram
getAllConnectedPeers metric.Int64Histogram
hasChannelDurationMicro metric.Int64Histogram
calcPostureChecksDurationMicro metric.Int64Histogram
calcPeerNetworkMapDurationMs metric.Int64Histogram
mergeNetworkMapDurationMicro metric.Int64Histogram
toSyncResponseDurationMicro metric.Int64Histogram
ctx context.Context
createChannelDurationMicro metric.Int64Histogram
closeChannelDurationMicro metric.Int64Histogram
closeChannelsDurationMicro metric.Int64Histogram
closeChannels metric.Int64Histogram
sendUpdateDurationMicro metric.Int64Histogram
getAllConnectedPeersDurationMicro metric.Int64Histogram
getAllConnectedPeers metric.Int64Histogram
hasChannelDurationMicro metric.Int64Histogram
calcPostureChecksDurationMicro metric.Int64Histogram
calcPeerNetworkMapDurationMs metric.Int64Histogram
mergeNetworkMapDurationMicro metric.Int64Histogram
toSyncResponseDurationMicro metric.Int64Histogram
toComponentSyncResponseDurationMicro metric.Int64Histogram
ctx context.Context
}
// NewUpdateChannelMetrics creates an instance of UpdateChannel
@@ -125,20 +126,29 @@ func NewUpdateChannelMetrics(ctx context.Context, meter metric.Meter) (*UpdateCh
return nil, err
}
toComponentSyncResponseDurationMicro, err := meter.Int64Histogram("management.updatechannel.tocomponentsyncresponse.duration.micro",
metric.WithUnit("microseconds"),
metric.WithDescription("Duration of how long it takes to convert components to component sync response"),
)
if err != nil {
return nil, err
}
return &UpdateChannelMetrics{
createChannelDurationMicro: createChannelDurationMicro,
closeChannelDurationMicro: closeChannelDurationMicro,
closeChannelsDurationMicro: closeChannelsDurationMicro,
closeChannels: closeChannels,
sendUpdateDurationMicro: sendUpdateDurationMicro,
getAllConnectedPeersDurationMicro: getAllConnectedPeersDurationMicro,
getAllConnectedPeers: getAllConnectedPeers,
hasChannelDurationMicro: hasChannelDurationMicro,
calcPostureChecksDurationMicro: calcPostureChecksDurationMicro,
calcPeerNetworkMapDurationMs: calcPeerNetworkMapDurationMs,
mergeNetworkMapDurationMicro: mergeNetworkMapDurationMicro,
toSyncResponseDurationMicro: toSyncResponseDurationMicro,
ctx: ctx,
createChannelDurationMicro: createChannelDurationMicro,
closeChannelDurationMicro: closeChannelDurationMicro,
closeChannelsDurationMicro: closeChannelsDurationMicro,
closeChannels: closeChannels,
sendUpdateDurationMicro: sendUpdateDurationMicro,
getAllConnectedPeersDurationMicro: getAllConnectedPeersDurationMicro,
getAllConnectedPeers: getAllConnectedPeers,
hasChannelDurationMicro: hasChannelDurationMicro,
calcPostureChecksDurationMicro: calcPostureChecksDurationMicro,
calcPeerNetworkMapDurationMs: calcPeerNetworkMapDurationMs,
mergeNetworkMapDurationMicro: mergeNetworkMapDurationMicro,
toSyncResponseDurationMicro: toSyncResponseDurationMicro,
toComponentSyncResponseDurationMicro: toComponentSyncResponseDurationMicro,
ctx: ctx,
}, nil
}
@@ -193,3 +203,7 @@ func (metrics *UpdateChannelMetrics) CountMergeNetworkMapDuration(duration time.
func (metrics *UpdateChannelMetrics) CountToSyncResponseDuration(duration time.Duration) {
metrics.toSyncResponseDurationMicro.Record(metrics.ctx, duration.Microseconds())
}
func (metrics *UpdateChannelMetrics) CountToComponentSyncResponseDuration(duration time.Duration) {
metrics.toComponentSyncResponseDurationMicro.Record(metrics.ctx, duration.Microseconds())
}

View File

@@ -29,7 +29,6 @@ import (
"github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/status"
"github.com/netbirdio/netbird/version"
)
const (
@@ -42,27 +41,8 @@ const (
PublicCategory = "public"
PrivateCategory = "private"
UnknownCategory = "unknown"
// firewallRuleMinPortRangesVer defines the minimum peer version that supports port range rules.
firewallRuleMinPortRangesVer = "0.48.0"
// firewallRuleMinNativeSSHVer defines the minimum peer version that supports native SSH features in the firewall rules.
firewallRuleMinNativeSSHVer = "0.60.0"
// nativeSSHPortString defines the default port number as a string used for native SSH connections; this port is used by clients when hijacking ssh connections.
nativeSSHPortString = "22022"
nativeSSHPortNumber = 22022
// defaultSSHPortString defines the standard SSH port number as a string, commonly used for default SSH connections.
defaultSSHPortString = "22"
defaultSSHPortNumber = 22
)
type supportedFeatures struct {
nativeSSH bool
portRanges bool
}
type LookupMap map[string]struct{}
// AccountMeta is a struct that contains a stripped down version of the Account object.
// It doesn't carry any peers, groups, policies, or routes, etc. Just some metadata (e.g. ID, created by, created at, etc).
type AccountMeta struct {
@@ -303,8 +283,8 @@ func (a *Account) SynthesizePrivateServiceZones(peerID string) []nbdns.CustomZon
// it, adding a single private service would black-hole every
// other name under the zone apex.
zone = &nbdns.CustomZone{
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
NonAuthoritative: true,
SearchDomainDisabled: true,
}
@@ -1071,7 +1051,7 @@ func (a *Account) GetPeerConnectionResources(ctx context.Context, peer *nbpeer.P
default:
authorizedUsers[auth.Wildcard] = a.getAllowedUserIDs()
}
} else if peerInDestinations && policyRuleImpliesLegacySSH(rule) && peer.SSHEnabled {
} else if peerInDestinations && PolicyRuleImpliesLegacySSH(rule) && peer.SSHEnabled {
sshEnabled = true
authorizedUsers[auth.Wildcard] = a.getAllowedUserIDs()
}
@@ -1102,6 +1082,7 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
peersExists := make(map[string]struct{})
rules := make([]*FirewallRule, 0)
peers := make([]*nbpeer.Peer, 0)
targetComponent := targetPeer.ToComponent()
return func(rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int) {
for _, peer := range groupPeers {
@@ -1137,15 +1118,15 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
if len(rule.Ports) == 0 && len(rule.PortRanges) == 0 {
rules = append(rules, &fr)
} else {
rules = append(rules, expandPortsAndRanges(fr, rule, targetPeer)...)
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetComponent)...)
}
rules = appendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, firewallRuleContext{
direction: direction,
dirStr: strconv.Itoa(direction),
protocolStr: string(protocol),
actionStr: string(rule.Action),
portsJoined: strings.Join(rule.Ports, ","),
rules = AppendIPv6FirewallRule(rules, rulesExists, peer.ToComponent(), targetComponent, rule, FirewallRuleContext{
Direction: direction,
DirStr: strconv.Itoa(direction),
ProtocolStr: string(protocol),
ActionStr: string(rule.Action),
PortsJoined: strings.Join(rule.Ports, ","),
})
}
}, func() ([]*nbpeer.Peer, []*FirewallRule) {
@@ -1153,10 +1134,6 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
}
}
func policyRuleImpliesLegacySSH(rule *PolicyRule) bool {
return rule.Protocol == PolicyRuleProtocolALL || (rule.Protocol == PolicyRuleProtocolTCP && (portsIncludesSSH(rule.Ports) || portRangeIncludesSSH(rule.PortRanges)))
}
// PeerSSHEnabledFromPolicies is the network-map-free equivalent of the sshEnabled
// determination in GetPeerConnectionResources / CalculateNetworkMapFromComponents.
func PeerSSHEnabledFromPolicies(policies []*Policy, peerID string, peerGroupIDs map[string]struct{}, peerSSHEnabled bool) bool {
@@ -1171,7 +1148,7 @@ func PeerSSHEnabledFromPolicies(policies []*Policy, peerID string, peerGroupIDs
}
isSSHRule := rule.Protocol == PolicyRuleProtocolNetbirdSSH ||
(policyRuleImpliesLegacySSH(rule) && peerSSHEnabled)
(PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled)
if !isSSHRule {
continue
}
@@ -1198,24 +1175,6 @@ func ruleHasDestination(rule *PolicyRule, peerID string, peerGroupIDs map[string
return false
}
func portRangeIncludesSSH(portRanges []RulePortRange) bool {
for _, pr := range portRanges {
if (pr.Start <= defaultSSHPortNumber && pr.End >= defaultSSHPortNumber) || (pr.Start <= nativeSSHPortNumber && pr.End >= nativeSSHPortNumber) {
return true
}
}
return false
}
func portsIncludesSSH(ports []string) bool {
for _, port := range ports {
if port == defaultSSHPortString || port == nativeSSHPortString {
return true
}
}
return false
}
// getAllPeersFromGroups for given peer ID and list of groups
//
// Returns a list of peers from specified groups that pass specified posture checks
@@ -1315,14 +1274,14 @@ func (a *Account) getRouteFirewallRules(ctx context.Context, peerID string, poli
}
rulePeers := a.getRulePeers(rule, policy.SourcePostureChecks, peerID, distributionPeers, validatedPeersMap)
rules := generateRouteFirewallRules(ctx, route, rule, rulePeers, FirewallRuleDirectionIN, includeIPv6)
rules := GenerateRouteFirewallRules(ctx, route, rule, rulePeers, FirewallRuleDirectionIN, includeIPv6)
fwRules = append(fwRules, rules...)
}
}
return fwRules
}
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*nbpeer.Peer {
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*ComponentPeer {
distPeersWithPolicy := make(map[string]struct{})
for _, id := range rule.Sources {
group := a.Groups[id]
@@ -1349,13 +1308,13 @@ func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID
}
}
distributionGroupPeers := make([]*nbpeer.Peer, 0, len(distPeersWithPolicy))
distributionGroupPeers := make([]*ComponentPeer, 0, len(distPeersWithPolicy))
for pID := range distPeersWithPolicy {
peer := a.Peers[pID]
if peer == nil {
continue
}
distributionGroupPeers = append(distributionGroupPeers, peer)
distributionGroupPeers = append(distributionGroupPeers, peer.ToComponent())
}
return distributionGroupPeers
}
@@ -1558,6 +1517,54 @@ func (a *Account) GetResourceRoutersMap() map[string]map[string]*routerTypes.Net
return routers
}
// forcesRoutingPeerDNSResolution reports whether the given peer must run
// routing-peer DNS resolution regardless of the account-global
// RoutingPeerDNSResolutionEnabled setting. It returns true when the peer is a
// router for a domain network resource that is targeted by an enabled
// reverse-proxy service, so the peer's DNS forwarder starts and can resolve
// the target for the embedded proxy peers. Embedded proxy peers themselves are
// handled at PeerConfig build time.
func (a *Account) forcesRoutingPeerDNSResolution(peerID string, routers map[string]map[string]*routerTypes.NetworkRouter) bool {
targeted := a.proxyTargetedDomainResourceIDs()
if len(targeted) == 0 {
return false
}
for _, resource := range a.NetworkResources {
if resource == nil || !resource.Enabled || resource.Type != resourceTypes.Domain {
continue
}
if _, ok := targeted[resource.ID]; !ok {
continue
}
if _, isRouter := routers[resource.NetworkID][peerID]; isRouter {
return true
}
}
return false
}
// proxyTargetedDomainResourceIDs returns the set of domain network resource IDs
// targeted by an enabled, non-terminated reverse-proxy service.
func (a *Account) proxyTargetedDomainResourceIDs() map[string]struct{} {
ids := make(map[string]struct{})
for _, svc := range a.Services {
if svc == nil || !svc.Enabled || svc.Terminated {
continue
}
for _, target := range svc.Targets {
if target == nil || !target.Enabled {
continue
}
if target.TargetType == service.TargetTypeDomain {
ids[target.TargetId] = struct{}{}
}
}
}
return ids
}
// getPoliciesSourcePeers collects all unique peers from the source groups defined in the given policies.
func getPoliciesSourcePeers(policies []*Policy, groups map[string]*Group) map[string]struct{} {
sourcePeers := make(map[string]struct{})
@@ -1808,96 +1815,6 @@ func (a *Account) createProxyPolicy(svc *service.Service, target *service.Target
}
}
// expandPortsAndRanges expands Ports and PortRanges of a rule into individual firewall rules
func expandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *nbpeer.Peer) []*FirewallRule {
features := peerSupportedFirewallFeatures(peer.Meta.WtVersion)
var expanded []*FirewallRule
for _, port := range rule.Ports {
fr := base
fr.Port = port
expanded = append(expanded, &fr)
}
for _, portRange := range rule.PortRanges {
// prefer PolicyRule.Ports
if len(rule.Ports) > 0 {
break
}
fr := base
if features.portRanges {
fr.PortRange = portRange
} else {
// Peer doesn't support port ranges, only allow single-port ranges
if portRange.Start != portRange.End {
continue
}
fr.Port = strconv.FormatUint(uint64(portRange.Start), 10)
}
expanded = append(expanded, &fr)
}
if shouldCheckRulesForNativeSSH(features.nativeSSH, rule, peer) || rule.Protocol == PolicyRuleProtocolNetbirdSSH {
expanded = addNativeSSHRule(base, expanded)
}
return expanded
}
// addNativeSSHRule adds a native SSH rule (port 22022) to the expanded rules if the base rule has port 22 configured.
func addNativeSSHRule(base FirewallRule, expanded []*FirewallRule) []*FirewallRule {
shouldAdd := false
for _, fr := range expanded {
if isPortInRule(nativeSSHPortString, 22022, fr) {
return expanded
}
if isPortInRule(defaultSSHPortString, 22, fr) {
shouldAdd = true
}
}
if !shouldAdd {
return expanded
}
fr := base
fr.Port = nativeSSHPortString
return append(expanded, &fr)
}
func isPortInRule(portString string, portInt uint16, rule *FirewallRule) bool {
return rule.Port == portString || (rule.PortRange.Start <= portInt && portInt <= rule.PortRange.End)
}
// shouldCheckRulesForNativeSSH determines whether specific policy rules should be checked for native SSH support.
// While users can add the nativeSSHPortString, we look for cases when they used port 22 and based on SSH enabled
// in both management and client, we indicate to add the native port.
func shouldCheckRulesForNativeSSH(supportsNative bool, rule *PolicyRule, peer *nbpeer.Peer) bool {
return supportsNative && peer.SSHEnabled && peer.Meta.Flags.ServerSSHAllowed && rule.Protocol == PolicyRuleProtocolTCP
}
// peerSupportedFirewallFeatures checks if the peer version supports port ranges.
func peerSupportedFirewallFeatures(peerVer string) supportedFeatures {
if version.IsDevelopmentVersion(peerVer) {
return supportedFeatures{true, true}
}
var features supportedFeatures
meetMinVer, err := posture.MeetsMinVersion(firewallRuleMinNativeSSHVer, peerVer)
features.nativeSSH = err == nil && meetMinVer
if features.nativeSSH {
features.portRanges = true
} else {
meetMinVer, err = posture.MeetsMinVersion(firewallRuleMinPortRangesVer, peerVer)
features.portRanges = err == nil && meetMinVer
}
return features
}
// filterZoneRecordsForPeers filters DNS records to only include peers to connect.
// AAAA records are excluded when the requesting peer lacks IPv6 capability.
func filterZoneRecordsForPeers(peer *nbpeer.Peer, customZone nbdns.CustomZone, peersToConnect, expiredPeers []*nbpeer.Peer) []nbdns.SimpleRecord {

View File

@@ -9,13 +9,44 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/modules/zones"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/route"
)
// GetPeerNetworkMapResult dispatches to either the legacy-NetworkMap path or
// the components path based on the peer's capability and the kill switch.
// Capable peers (PeerCapabilityComponentNetworkMap) get the raw components
// shape — the server skips Calculate() entirely for them, saving CPU
// proportional to the number of capable peers in the account. Legacy peers
// (or any peer when componentsDisabled is true) get the fully-expanded
// NetworkMap as before.
func (a *Account) GetPeerNetworkMapResult(
ctx context.Context,
peerID string,
componentsDisabled bool,
peersCustomZone nbdns.CustomZone,
accountZones []*zones.Zone,
validatedPeersMap map[string]struct{},
resourcePolicies map[string][]*Policy,
routers map[string]map[string]*routerTypes.NetworkRouter,
metrics *telemetry.AccountManagerMetrics,
groupIDToUserIDs map[string][]string,
) PeerNetworkMapResult {
peer := a.Peers[peerID]
if !componentsDisabled && peer != nil && peer.SupportsComponentNetworkMap() {
components := a.GetPeerNetworkMapComponents(
ctx, peerID, peersCustomZone, accountZones, validatedPeersMap, resourcePolicies, routers, groupIDToUserIDs,
)
return PeerNetworkMapResult{Components: components}
}
return PeerNetworkMapResult{
NetworkMap: a.GetPeerNetworkMapFromComponents(
ctx, peerID, peersCustomZone, accountZones, validatedPeersMap, resourcePolicies, routers, metrics, groupIDToUserIDs,
),
}
}
func (a *Account) GetPeerNetworkMapFromComponents(
ctx context.Context,
peerID string,
@@ -40,8 +71,8 @@ func (a *Account) GetPeerNetworkMapFromComponents(
groupIDToUserIDs,
)
if components == nil {
return &NetworkMap{Network: a.Network.Copy()}
if components.IsEmpty() {
return &NetworkMap{Network: components.Network}
}
nm := CalculateNetworkMapFromComponents(ctx, components)
@@ -71,26 +102,56 @@ func (a *Account) GetPeerNetworkMapComponents(
routers map[string]map[string]*routerTypes.NetworkRouter,
groupIDToUserIDs map[string][]string,
) *NetworkMapComponents {
peer := a.Peers[peerID]
// this can never happen, things are very wrong if it did
// TODO (dmitri) maybe consider using invariants?
if peer == nil {
return nil
log.WithField("peer id", peerID).Error("NetworkMapComponents are computed for a peer missing from the account")
return EmptyNetworkMapComponents(&NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
if _, ok := validatedPeersMap[peerID]; !ok {
return nil
// Mirror legacy graceful-degrade: GetPeerNetworkMapFromComponents
// returns &NetworkMap{Network: a.Network.Copy()} when components is
// nil. Match that floor so the receiving client always sees the
// account Network identifier, not a fully-empty envelope.
return EmptyNetworkMapComponents(&NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
components := &NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
NameServerGroups: make([]*nbdns.NameServerGroup, 0),
CustomZoneDomain: peersCustomZone.Domain,
ResourcePoliciesMap: make(map[string][]*Policy),
RoutersMap: make(map[string]map[string]*routerTypes.NetworkRouter),
NetworkResources: make([]*resourceTypes.NetworkResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(a.PostureChecks)),
RouterPeers: make(map[string]*nbpeer.Peer),
PeerID: peerID,
Network: a.Network.Copy(),
NameServerGroups: make([]*nbdns.NameServerGroup, 0),
CustomZoneDomain: peersCustomZone.Domain,
ResourcePoliciesMap: make(map[string][]*Policy),
RoutersMap: make(map[string]map[string]*ComponentRouter),
NetworkResources: make([]*ComponentResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(a.PostureChecks)),
RouterPeers: make(map[string]*ComponentPeer),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
ForceRoutingPeerDNSResolution: a.forcesRoutingPeerDNSResolution(peerID, routers),
}
for _, n := range a.Networks {
if n != nil {
components.NetworkXIDToPublicID[n.ID] = n.PublicID
}
}
for _, pc := range a.PostureChecks {
if pc != nil {
components.PostureCheckXIDToPublicID[pc.ID] = pc.PublicID
}
}
components.AccountSettings = &AccountSettingsInfo{
@@ -102,6 +163,7 @@ func (a *Account) GetPeerNetworkMapComponents(
components.DNSSettings = &a.DNSSettings
// relevantPeers always contains the target peer (peerID)
relevantPeers, relevantGroups, relevantPolicies, relevantRoutes, sshReqs := a.getPeersGroupsPoliciesRoutes(ctx, peerID, peer.SSHEnabled, validatedPeersMap, &components.PostureFailedPeers)
if len(sshReqs.neededGroupIDs) > 0 {
@@ -112,7 +174,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
components.Peers = relevantPeers
components.Groups = relevantGroups
components.Groups = GroupsToComponent(relevantGroups)
components.Policies = relevantPolicies
components.Routes = relevantRoutes
components.AllDNSRecords = filterDNSRecordsByPeers(peersCustomZone.Records, relevantPeers, peer.SupportsIPv6() && peer.IPv6.IsValid())
@@ -161,7 +223,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
for _, pID := range a.getPostureValidPeersSaveFailed(peers, policy.SourcePostureChecks, validatedPeersMap, &components.PostureFailedPeers) {
if _, exists := components.Peers[pID]; !exists {
components.Peers[pID] = a.GetPeer(pID)
components.Peers[pID] = a.GetPeer(pID).ToComponent()
}
}
} else {
@@ -194,14 +256,14 @@ func (a *Account) GetPeerNetworkMapComponents(
for _, srcGroupID := range rule.Sources {
if g := a.Groups[srcGroupID]; g != nil {
if _, exists := components.Groups[srcGroupID]; !exists {
components.Groups[srcGroupID] = g
components.Groups[srcGroupID] = g.ToComponent()
}
}
}
for _, dstGroupID := range rule.Destinations {
if g := a.Groups[dstGroupID]; g != nil {
if _, exists := components.Groups[dstGroupID]; !exists {
components.Groups[dstGroupID] = g
components.Groups[dstGroupID] = g.ToComponent()
}
}
}
@@ -209,22 +271,29 @@ func (a *Account) GetPeerNetworkMapComponents(
components.ResourcePoliciesMap[resource.ID] = policies
}
components.RoutersMap[resource.NetworkID] = networkRoutingPeers
for peerIDKey := range networkRoutingPeers {
if p := a.Peers[peerIDKey]; p != nil {
if _, exists := components.RouterPeers[peerIDKey]; !exists {
components.RouterPeers[peerIDKey] = p
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := validatedPeersMap[peerIDKey]; validated {
components.Peers[peerIDKey] = p
// Only expose router peers and the per-network routers_map when this
// target peer actually has access to the resource (either as a router
// itself or via a policy that includes it as a source). Without this
// gate, every peer's envelope was leaking router peers of every
// network in the account — accounts with many tenants/networks
// shipped tens of unrelated peers in `peers[]` and `routers_map`.
if addSourcePeers {
components.RoutersMap[resource.NetworkID] = routerTypes.ToComponentMap(networkRoutingPeers)
for peerIDKey := range networkRoutingPeers {
if p := a.Peers[peerIDKey]; p != nil {
cp := components.RouterPeers[peerIDKey]
if cp == nil {
cp = p.ToComponent()
components.RouterPeers[peerIDKey] = cp
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := validatedPeersMap[peerIDKey]; validated {
components.Peers[peerIDKey] = cp
}
}
}
}
}
if addSourcePeers {
components.NetworkResources = append(components.NetworkResources, resource)
components.NetworkResources = append(components.NetworkResources, resource.ToComponent())
}
}
@@ -245,27 +314,53 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
peerSSHEnabled bool,
validatedPeersMap map[string]struct{},
postureFailedPeers *map[string]map[string]struct{},
) (map[string]*nbpeer.Peer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*nbpeer.Peer, len(a.Peers)/4)
) (map[string]*ComponentPeer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*ComponentPeer, len(a.Peers)/4)
relevantGroupIDs := make(map[string]*Group, len(a.Groups)/4)
relevantPolicies := make([]*Policy, 0, len(a.Policies))
relevantRoutes := make([]*route.Route, 0, len(a.Routes))
sshReqs := sshRequirements{neededGroupIDs: make(map[string]struct{})}
relevantPeerIDs[peerID] = a.GetPeer(peerID)
relevantPeerIDs[peerID] = a.GetPeer(peerID).ToComponent()
peerGroupSet := make(map[string]struct{}, 8)
for groupID, group := range a.Groups {
if slices.Contains(group.Peers, peerID) {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
peerGroupSet[groupID] = struct{}{}
}
}
routeAccessControlGroups := make(map[string]struct{})
for _, r := range a.Routes {
for _, groupID := range r.Groups {
if r == nil {
continue
}
relevant := r.Peer == peerID
if !relevant {
for _, groupID := range r.PeerGroups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant && r.Enabled {
for _, groupID := range r.Groups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant {
continue
}
for _, groupID := range r.PeerGroups {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
}
for _, groupID := range r.PeerGroups {
for _, groupID := range r.Groups {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
}
if r.Enabled {
@@ -274,6 +369,44 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
routeAccessControlGroups[groupID] = struct{}{}
}
}
// Include route advertisers in relevantPeerIDs. The envelope
// encoder writes route.peer_index by looking up r.Peer in the
// shipped peers list; if the advertiser is policy-isolated from
// the target peer (no rule edge between them), it would otherwise
// be omitted and the decoder would fail to resolve r.Peer, leaving
// the client without a WG tunnel target for this route. Legacy
// NetworkMap.Routes shipped the WG public key inline, so the
// equivalence path doesn't surface this — but the dependency is
// real once a client actually tries to use the route.
// Gate by validatedPeersMap so non-validated advertisers stay out
// (matches the network-resource router behaviour at the bottom of
// this loop, and the legacy invariant that only validated peers
// reach a client's view).
if r.Peer != "" {
if _, ok := validatedPeersMap[r.Peer]; ok {
if p := a.GetPeer(r.Peer); p != nil {
relevantPeerIDs[r.Peer] = p.ToComponent()
}
}
}
for _, groupID := range r.PeerGroups {
g := a.GetGroup(groupID)
if g == nil {
continue
}
for _, pid := range g.Peers {
if _, exists := relevantPeerIDs[pid]; exists {
continue
}
if _, ok := validatedPeersMap[pid]; !ok {
continue
}
if p := a.GetPeer(pid); p != nil {
relevantPeerIDs[pid] = p.ToComponent()
}
}
}
relevantRoutes = append(relevantRoutes, r)
}
@@ -327,7 +460,9 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInSources {
policyRelevant = true
for _, pid := range destinationPeers {
relevantPeerIDs[pid] = a.GetPeer(pid)
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = a.GetGroup(dstGroupID)
@@ -337,7 +472,9 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInDestinations {
policyRelevant = true
for _, pid := range sourcePeers {
relevantPeerIDs[pid] = a.GetPeer(pid)
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = a.GetGroup(srcGroupID)
@@ -353,7 +490,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
default:
sshReqs.needAllowedUserIDs = true
}
} else if policyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
} else if PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
sshReqs.needAllowedUserIDs = true
}
}
@@ -486,7 +623,14 @@ func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChe
return dest
}
func filterGroupPeers(groups *map[string]*Group, peers map[string]*nbpeer.Peer) {
// filterGroupPeers trims each group's Peers slice to only those peers that
// also appear in `peers`. Groups whose filtered list is empty are NOT
// deleted from the map — they're kept so the components wire encoder can
// still resolve seq references from routes/policies/access-control groups
// that name them. Calculate() tolerates groups with empty Peers (the inner
// loops simply iterate zero times), so retaining them is behaviourally a
// no-op for the legacy path that consumes the same NetworkMapComponents.
func filterGroupPeers(groups *map[string]*ComponentGroup, peers map[string]*ComponentPeer) {
for groupID, groupInfo := range *groups {
filteredPeers := make([]string, 0, len(groupInfo.Peers))
for _, pid := range groupInfo.Peers {
@@ -495,17 +639,15 @@ func filterGroupPeers(groups *map[string]*Group, peers map[string]*nbpeer.Peer)
}
}
if len(filteredPeers) == 0 {
delete(*groups, groupID)
} else if len(filteredPeers) != len(groupInfo.Peers) {
ng := groupInfo.Copy()
if len(filteredPeers) != len(groupInfo.Peers) {
ng := *groupInfo
ng.Peers = filteredPeers
(*groups)[groupID] = ng
(*groups)[groupID] = &ng
}
}
}
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*nbpeer.Peer) {
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*ComponentPeer) {
if len(*postureFailedPeers) == 0 {
return
}
@@ -540,7 +682,7 @@ func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}
}
}
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*nbpeer.Peer, includeIPv6 bool) []nbdns.SimpleRecord {
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*ComponentPeer, includeIPv6 bool) []nbdns.SimpleRecord {
if len(records) == 0 || len(peers) == 0 {
return nil
}

View File

@@ -9,7 +9,6 @@ import (
"github.com/stretchr/testify/require"
nbdns "github.com/netbirdio/netbird/dns"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
)
func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
@@ -49,7 +48,7 @@ func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
})
}
func netmapPeerIDs(peers []*nbpeer.Peer) []string {
func netmapPeerIDs(peers []*ComponentPeer) []string {
ids := make([]string, 0, len(peers))
for _, p := range peers {
if p == nil {

View File

@@ -666,7 +666,7 @@ func Test_ExpandPortsAndRanges_SSHRuleExpansion(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := expandPortsAndRanges(tt.base, tt.rule, tt.peer)
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer.ToComponent())
var ports []string
for _, fr := range result {
@@ -1751,3 +1751,71 @@ func hasPrivateAccessPolicy(account *Account, serviceID string) bool {
}
return false
}
func TestForcesRoutingPeerDNSResolution(t *testing.T) {
buildAccountRes := func(serviceEnabled, targetEnabled, resourceEnabled bool, targetType service.TargetType, resType resourceTypes.NetworkResourceType) *Account {
return &Account{
Id: "accountID",
Groups: map[string]*Group{
"router-group": {ID: "router-group", Peers: []string{"router-peer-grp"}},
},
NetworkRouters: []*routerTypes.NetworkRouter{
{ID: "r1", NetworkID: "net-1", AccountID: "accountID", Peer: "router-peer", Enabled: true},
{ID: "r2", NetworkID: "net-1", AccountID: "accountID", PeerGroups: []string{"router-group"}, Enabled: true},
},
NetworkResources: []*resourceTypes.NetworkResource{
{ID: "res-domain", AccountID: "accountID", NetworkID: "net-1", Type: resType, Domain: "example.org", Enabled: resourceEnabled},
},
Services: []*service.Service{
{
ID: "svc-1", AccountID: "accountID", Enabled: serviceEnabled,
Targets: []*service.Target{
{TargetId: "res-domain", TargetType: targetType, Enabled: targetEnabled},
},
},
},
}
}
buildAccount := func(serviceEnabled, targetEnabled, resourceEnabled bool, targetType service.TargetType) *Account {
return buildAccountRes(serviceEnabled, targetEnabled, resourceEnabled, targetType, resourceTypes.Domain)
}
t.Run("router peer for RP-targeted domain resource is forced", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypeDomain)
routers := account.GetResourceRoutersMap()
assert.True(t, account.forcesRoutingPeerDNSResolution("router-peer", routers), "direct router peer should be forced")
assert.True(t, account.forcesRoutingPeerDNSResolution("router-peer-grp", routers), "group-member router peer should be forced")
})
t.Run("non-router peer is not forced", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("other-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when service disabled", func(t *testing.T) {
account := buildAccount(false, true, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when target disabled", func(t *testing.T) {
account := buildAccount(true, false, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when resource disabled", func(t *testing.T) {
account := buildAccount(true, true, false, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced for non-domain target type", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypePeer)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when targeted resource is not a domain", func(t *testing.T) {
account := buildAccountRes(true, true, true, service.TargetTypeDomain, resourceTypes.Host)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()),
"a domain target pointing at a non-domain resource must not force resolution")
})
}

View File

@@ -0,0 +1,148 @@
package types
import (
"context"
"math/rand"
"net"
"net/netip"
nbroute "github.com/netbirdio/netbird/route"
sharedtypes "github.com/netbirdio/netbird/shared/management/types"
)
// Type aliases for types relocated to shared/management/types so that the
// client-side compute path can depend on them
type DNSSettings = sharedtypes.DNSSettings
type FirewallRule = sharedtypes.FirewallRule
type Network = sharedtypes.Network
type NetworkMap = sharedtypes.NetworkMap
type ForwardingRule = sharedtypes.ForwardingRule
type Policy = sharedtypes.Policy
type PolicyUpdateOperation = sharedtypes.PolicyUpdateOperation
type PolicyRule = sharedtypes.PolicyRule
type PolicyUpdateOperationType = sharedtypes.PolicyUpdateOperationType
type PolicyTrafficActionType = sharedtypes.PolicyTrafficActionType
type PolicyRuleProtocolType = sharedtypes.PolicyRuleProtocolType
type PolicyRuleDirection = sharedtypes.PolicyRuleDirection
type RulePortRange = sharedtypes.RulePortRange
type Resource = sharedtypes.Resource
type ResourceType = sharedtypes.ResourceType
type RouteFirewallRule = sharedtypes.RouteFirewallRule
type NetworkMapComponents = sharedtypes.NetworkMapComponents
type ComponentPeer = sharedtypes.ComponentPeer
type ComponentGroup = sharedtypes.ComponentGroup
type ComponentRouter = sharedtypes.ComponentRouter
type ComponentResource = sharedtypes.ComponentResource
type ComponentResourceType = sharedtypes.ComponentResourceType
const (
ComponentResourceHost = sharedtypes.ComponentResourceHost
ComponentResourceSubnet = sharedtypes.ComponentResourceSubnet
ComponentResourceDomain = sharedtypes.ComponentResourceDomain
)
var EmptyNetworkMapComponents = sharedtypes.EmptyNetworkMapComponents
type AccountSettingsInfo = sharedtypes.AccountSettingsInfo
type GroupCompact = sharedtypes.GroupCompact
type NetworkMapComponentsCompact = sharedtypes.NetworkMapComponentsCompact
type LookupMap = sharedtypes.LookupMap
type FirewallRuleContext = sharedtypes.FirewallRuleContext
const GroupAllName = sharedtypes.GroupAllName
// Function forwarders preserve types.X(...) call sites that previously
// resolved to package-local funcs. Plain forwarders (not var aliases) keep
// the symbol immutable and allow the inliner to flatten the call.
func PolicyRuleImpliesLegacySSH(rule *PolicyRule) bool {
return sharedtypes.PolicyRuleImpliesLegacySSH(rule)
}
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *ComponentPeer) []*FirewallRule {
return sharedtypes.ExpandPortsAndRanges(base, rule, peer)
}
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *ComponentPeer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
return sharedtypes.AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, rc)
}
func CalculateNetworkMapFromComponents(ctx context.Context, components *NetworkMapComponents) *NetworkMap {
return sharedtypes.CalculateNetworkMapFromComponents(ctx, components)
}
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*ComponentPeer, direction int, includeIPv6 bool) []*RouteFirewallRule {
return sharedtypes.GenerateRouteFirewallRules(ctx, route, rule, groupPeers, direction, includeIPv6)
}
func AllocateIPv6Subnet(r *rand.Rand) net.IPNet {
return sharedtypes.AllocateIPv6Subnet(r)
}
func NewNetwork() *Network {
return sharedtypes.NewNetwork()
}
func AllocatePeerIP(prefix netip.Prefix, takenIps []netip.Addr) (netip.Addr, error) {
return sharedtypes.AllocatePeerIP(prefix, takenIps)
}
func AllocateRandomPeerIP(prefix netip.Prefix) (netip.Addr, error) {
return sharedtypes.AllocateRandomPeerIP(prefix)
}
func AllocateRandomPeerIPv6(prefix netip.Prefix) (netip.Addr, error) {
return sharedtypes.AllocateRandomPeerIPv6(prefix)
}
func ParseRuleString(rule string) (PolicyRuleProtocolType, RulePortRange, error) {
return sharedtypes.ParseRuleString(rule)
}
const (
FirewallRuleDirectionIN = sharedtypes.FirewallRuleDirectionIN
FirewallRuleDirectionOUT = sharedtypes.FirewallRuleDirectionOUT
)
const (
ResourceTypePeer = sharedtypes.ResourceTypePeer
ResourceTypeDomain = sharedtypes.ResourceTypeDomain
ResourceTypeHost = sharedtypes.ResourceTypeHost
ResourceTypeSubnet = sharedtypes.ResourceTypeSubnet
)
const (
PolicyTrafficActionAccept = sharedtypes.PolicyTrafficActionAccept
PolicyTrafficActionDrop = sharedtypes.PolicyTrafficActionDrop
)
const (
PolicyRuleProtocolALL = sharedtypes.PolicyRuleProtocolALL
PolicyRuleProtocolTCP = sharedtypes.PolicyRuleProtocolTCP
PolicyRuleProtocolUDP = sharedtypes.PolicyRuleProtocolUDP
PolicyRuleProtocolICMP = sharedtypes.PolicyRuleProtocolICMP
PolicyRuleProtocolNetbirdSSH = sharedtypes.PolicyRuleProtocolNetbirdSSH
)
const (
PolicyRuleFlowDirect = sharedtypes.PolicyRuleFlowDirect
PolicyRuleFlowBidirect = sharedtypes.PolicyRuleFlowBidirect
)
const (
DefaultRuleName = sharedtypes.DefaultRuleName
DefaultRuleDescription = sharedtypes.DefaultRuleDescription
DefaultPolicyName = sharedtypes.DefaultPolicyName
DefaultPolicyDescription = sharedtypes.DefaultPolicyDescription
)

View File

@@ -1,16 +0,0 @@
package types
// DNSSettings defines dns settings at the account level
type DNSSettings struct {
// DisabledManagementGroups groups whose DNS management is disabled
DisabledManagementGroups []string `gorm:"serializer:json"`
}
// Copy returns a copy of the DNS settings
func (d DNSSettings) Copy() DNSSettings {
settings := DNSSettings{
DisabledManagementGroups: make([]string, len(d.DisabledManagementGroups)),
}
copy(settings.DisabledManagementGroups, d.DisabledManagementGroups)
return settings
}

View File

@@ -1,182 +0,0 @@
package types
import (
"context"
"fmt"
"reflect"
"strconv"
"strings"
log "github.com/sirupsen/logrus"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
)
const (
FirewallRuleDirectionIN = 0
FirewallRuleDirectionOUT = 1
)
// FirewallRule is a rule of the firewall.
type FirewallRule struct {
// PolicyID is the ID of the policy this rule is derived from
PolicyID string
// PeerIP of the peer
PeerIP string
// Direction of the traffic
Direction int
// Action of the traffic
Action string
// Protocol of the traffic
Protocol string
// Port of the traffic
Port string
// PortRange represents the range of ports for a firewall rule
PortRange RulePortRange
}
// Equal checks if two firewall rules are equal.
func (r *FirewallRule) Equal(other *FirewallRule) bool {
return reflect.DeepEqual(r, other)
}
// generateRouteFirewallRules generates a list of firewall rules for a given route.
// For static routes, source ranges match the destination family (v4 or v6).
// For dynamic routes (domain-based), separate v4 and v6 rules are generated
// so the routing peer's forwarding chain allows both address families.
func generateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int, includeIPv6 bool) []*RouteFirewallRule {
rulesExists := make(map[string]struct{})
rules := make([]*RouteFirewallRule, 0)
v4Sources, v6Sources := splitPeerSourcesByFamily(groupPeers)
isV6Route := route.Network.Addr().Is6()
// Skip v6 destination routes entirely for peers without IPv6 support
if isV6Route && !includeIPv6 {
return rules
}
// Pick sources matching the destination family
sourceRanges := v4Sources
if isV6Route {
sourceRanges = v6Sources
}
baseRule := RouteFirewallRule{
PolicyID: rule.PolicyID,
RouteID: route.ID,
SourceRanges: sourceRanges,
Action: string(rule.Action),
Destination: route.Network.String(),
Protocol: string(rule.Protocol),
Domains: route.Domains,
IsDynamic: route.IsDynamic(),
}
if len(rule.Ports) == 0 {
rules = append(rules, generateRulesWithPortRanges(baseRule, rule, rulesExists)...)
} else {
rules = append(rules, generateRulesWithPorts(ctx, baseRule, rule, rulesExists)...)
}
// Generate v6 counterpart for dynamic routes and 0.0.0.0/0 exit node routes.
isDefaultV4 := !isV6Route && route.Network.Bits() == 0
if includeIPv6 && (route.IsDynamic() || isDefaultV4) && len(v6Sources) > 0 {
v6Rule := baseRule
v6Rule.SourceRanges = v6Sources
if isDefaultV4 {
v6Rule.Destination = "::/0"
v6Rule.RouteID = route.ID + "-v6-default"
}
if len(rule.Ports) == 0 {
rules = append(rules, generateRulesWithPortRanges(v6Rule, rule, rulesExists)...)
} else {
rules = append(rules, generateRulesWithPorts(ctx, v6Rule, rule, rulesExists)...)
}
}
return rules
}
// splitPeerSourcesByFamily separates peer IPs into v4 (/32) and v6 (/128) source ranges.
func splitPeerSourcesByFamily(groupPeers []*nbpeer.Peer) (v4, v6 []string) {
v4 = make([]string, 0, len(groupPeers))
v6 = make([]string, 0, len(groupPeers))
for _, peer := range groupPeers {
if peer == nil {
continue
}
v4 = append(v4, fmt.Sprintf(AllowedIPsFormat, peer.IP))
if peer.IPv6.IsValid() {
v6 = append(v6, fmt.Sprintf(AllowedIPsV6Format, peer.IPv6))
}
}
return
}
// generateRulesForPeer generates rules for a given peer based on ports and port ranges.
func generateRulesWithPortRanges(baseRule RouteFirewallRule, rule *PolicyRule, rulesExists map[string]struct{}) []*RouteFirewallRule {
rules := make([]*RouteFirewallRule, 0)
ruleIDBase := generateRuleIDBase(rule, baseRule)
if len(rule.Ports) == 0 {
if len(rule.PortRanges) == 0 {
if _, ok := rulesExists[ruleIDBase]; !ok {
rulesExists[ruleIDBase] = struct{}{}
rules = append(rules, &baseRule)
}
} else {
for _, portRange := range rule.PortRanges {
ruleID := fmt.Sprintf("%s%d-%d", ruleIDBase, portRange.Start, portRange.End)
if _, ok := rulesExists[ruleID]; !ok {
rulesExists[ruleID] = struct{}{}
pr := baseRule
pr.PortRange = portRange
rules = append(rules, &pr)
}
}
}
return rules
}
return rules
}
// generateRulesWithPorts generates rules when specific ports are provided.
func generateRulesWithPorts(ctx context.Context, baseRule RouteFirewallRule, rule *PolicyRule, rulesExists map[string]struct{}) []*RouteFirewallRule {
rules := make([]*RouteFirewallRule, 0)
ruleIDBase := generateRuleIDBase(rule, baseRule)
for _, port := range rule.Ports {
ruleID := ruleIDBase + port
if _, ok := rulesExists[ruleID]; ok {
continue
}
rulesExists[ruleID] = struct{}{}
pr := baseRule
p, err := strconv.ParseUint(port, 10, 16)
if err != nil {
log.WithContext(ctx).Errorf("failed to parse port %s for rule: %s", port, rule.ID)
continue
}
pr.Port = uint16(p)
rules = append(rules, &pr)
}
return rules
}
// generateRuleIDBase generates the base rule ID for checking duplicates.
func generateRuleIDBase(rule *PolicyRule, baseRule RouteFirewallRule) string {
return rule.ID + strings.Join(baseRule.SourceRanges, ",") + strconv.Itoa(FirewallRuleDirectionIN) + baseRule.Protocol + baseRule.Action
}

View File

@@ -1,197 +0,0 @@
package types
import (
"context"
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/domain"
)
func TestSplitPeerSourcesByFamily(t *testing.T) {
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
},
{
IP: netip.MustParseAddr("100.64.0.2"),
},
{
IP: netip.MustParseAddr("100.64.0.3"),
IPv6: netip.MustParseAddr("fd00::3"),
},
nil,
}
v4, v6 := splitPeerSourcesByFamily(peers)
assert.Equal(t, []string{"100.64.0.1/32", "100.64.0.2/32", "100.64.0.3/32"}, v4)
assert.Equal(t, []string{"fd00::1/128", "fd00::3/128"}, v6)
}
func TestGenerateRouteFirewallRules_V4Route(t *testing.T) {
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
},
{
IP: netip.MustParseAddr("100.64.0.2"),
},
}
r := &route.Route{
ID: "route1",
Network: netip.MustParsePrefix("10.0.0.0/24"),
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, true)
require.Len(t, rules, 1)
assert.Equal(t, []string{"100.64.0.1/32", "100.64.0.2/32"}, rules[0].SourceRanges, "v4 route should only have v4 sources")
assert.Equal(t, "10.0.0.0/24", rules[0].Destination)
}
func TestGenerateRouteFirewallRules_V6Route(t *testing.T) {
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
},
{
IP: netip.MustParseAddr("100.64.0.2"),
},
}
r := &route.Route{
ID: "route1",
Network: netip.MustParsePrefix("2001:db8::/32"),
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, true)
require.Len(t, rules, 1)
assert.Equal(t, []string{"fd00::1/128"}, rules[0].SourceRanges, "v6 route should only have v6 sources")
}
func TestGenerateRouteFirewallRules_DynamicRoute_DualStack(t *testing.T) {
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
},
{
IP: netip.MustParseAddr("100.64.0.2"),
},
}
r := &route.Route{
ID: "route1",
NetworkType: route.DomainNetwork,
Domains: domain.List{"example.com"},
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, true)
require.Len(t, rules, 2, "dynamic route should produce both v4 and v6 rules")
assert.Equal(t, []string{"100.64.0.1/32", "100.64.0.2/32"}, rules[0].SourceRanges)
assert.Equal(t, []string{"fd00::1/128"}, rules[1].SourceRanges)
assert.Equal(t, rules[0].Domains, rules[1].Domains)
assert.True(t, rules[0].IsDynamic)
assert.True(t, rules[1].IsDynamic)
}
func TestGenerateRouteFirewallRules_DynamicRoute_NoV6Peers(t *testing.T) {
peers := []*nbpeer.Peer{
{IP: netip.MustParseAddr("100.64.0.1")},
{IP: netip.MustParseAddr("100.64.0.2")},
}
r := &route.Route{
ID: "route1",
NetworkType: route.DomainNetwork,
Domains: domain.List{"example.com"},
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, true)
require.Len(t, rules, 1, "no v6 peers means only v4 rule")
assert.Equal(t, []string{"100.64.0.1/32", "100.64.0.2/32"}, rules[0].SourceRanges)
}
func TestGenerateRouteFirewallRules_IncludeIPv6False(t *testing.T) {
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
},
{
IP: netip.MustParseAddr("100.64.0.2"),
IPv6: netip.MustParseAddr("fd00::2"),
},
}
t.Run("v6 route excluded", func(t *testing.T) {
r := &route.Route{
ID: "route1",
Network: netip.MustParsePrefix("2001:db8::/32"),
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, false)
assert.Empty(t, rules, "v6 route should produce no rules when includeIPv6 is false")
})
t.Run("dynamic route only v4", func(t *testing.T) {
r := &route.Route{
ID: "route1",
NetworkType: route.DomainNetwork,
Domains: domain.List{"example.com"},
}
rule := &PolicyRule{
PolicyID: "policy1",
ID: "rule1",
Action: PolicyTrafficActionAccept,
Protocol: PolicyRuleProtocolALL,
}
rules := generateRouteFirewallRules(context.Background(), r, rule, peers, FirewallRuleDirectionIN, false)
require.Len(t, rules, 1, "dynamic route with includeIPv6=false should produce only v4 rule")
assert.Equal(t, []string{"100.64.0.1/32", "100.64.0.2/32"}, rules[0].SourceRanges)
})
}

View File

@@ -2,7 +2,6 @@ package types
import (
"github.com/netbirdio/netbird/management/server/integration_reference"
"github.com/netbirdio/netbird/management/server/networks/resources/types"
)
const (
@@ -19,6 +18,8 @@ type Group struct {
// AccountID is a reference to Account that this object belongs
AccountID string `json:"-" gorm:"index"`
PublicID string `json:"-"`
// Name visible in the UI
Name string
@@ -66,14 +67,11 @@ func (g *Group) EventMeta() map[string]any {
return map[string]any{"name": g.Name}
}
func (g *Group) EventMetaResource(resource *types.NetworkResource) map[string]any {
return map[string]any{"name": g.Name, "id": g.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
func (g *Group) Copy() *Group {
group := &Group{
ID: g.ID,
AccountID: g.AccountID,
PublicID: g.PublicID,
Name: g.Name,
Issued: g.Issued,
Peers: make([]string, len(g.Peers)),
@@ -92,14 +90,39 @@ func (g *Group) HasPeers() bool {
return len(g.Peers) > 0
}
// GroupAllName is the reserved name of the default group that contains every peer in an account.
const GroupAllName = "All"
// IsGroupAll checks if the group is a default "All" group.
func (g *Group) IsGroupAll() bool {
return g.Name == GroupAllName
}
// ToComponent converts the group to its self-contained components
// representation. The Peers slice is shared, not copied — components are
// treated as immutable snapshots. Returns nil for a nil group.
func (g *Group) ToComponent() *ComponentGroup {
if g == nil {
return nil
}
return &ComponentGroup{
ID: g.ID,
PublicID: g.PublicID,
Name: g.Name,
Peers: g.Peers,
}
}
// GroupsToComponent converts an id-keyed group map to its components
// representation, preserving nil entries.
func GroupsToComponent(groups map[string]*Group) map[string]*ComponentGroup {
if groups == nil {
return nil
}
out := make(map[string]*ComponentGroup, len(groups))
for id, g := range groups {
out[id] = g.ToComponent()
}
return out
}
// AddPeer adds peerID to Peers if not present, returning true if added.
func (g *Group) AddPeer(peerID string) bool {
if peerID == "" {

View File

@@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/require"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
func TestNetworkMapComponents_IPv6EndToEnd(t *testing.T) {
@@ -104,7 +105,7 @@ func TestNetworkMapComponents_RemotePeerWithoutCapability(t *testing.T) {
require.NotNil(t, nm)
t.Run("AllowedIPs include remote v6", func(t *testing.T) {
var dst *nbpeer.Peer
var dst *types.ComponentPeer
for _, p := range nm.Peers {
if p.ID == "peer-dst-1" {
dst = p

View File

@@ -1,373 +0,0 @@
package types
import (
"crypto/rand"
"encoding/binary"
"fmt"
"net"
"net/netip"
"slices"
"sync"
"github.com/c-robinson/iplib"
"github.com/rs/xid"
"golang.org/x/exp/maps"
nbdns "github.com/netbirdio/netbird/dns"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/util"
"github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/status"
)
const (
// SubnetSize is a size of the subnet of the global network, e.g. 100.77.0.0/16
SubnetSize = 16
// NetSize is a global network size 100.64.0.0/10
NetSize = 10
// AllowedIPsFormat generates Wireguard AllowedIPs format (e.g. 100.64.30.1/32)
AllowedIPsFormat = "%s/32"
// AllowedIPsV6Format generates AllowedIPs format for v6 (e.g. fd12:3456:7890::1/128)
AllowedIPsV6Format = "%s/128"
// IPv6SubnetSize is the prefix length of per-account IPv6 subnets.
// Each account gets a /64 from its unique /48 ULA prefix.
IPv6SubnetSize = 64
)
type NetworkMap struct {
Peers []*nbpeer.Peer
Network *Network
Routes []*route.Route
DNSConfig nbdns.Config
OfflinePeers []*nbpeer.Peer
FirewallRules []*FirewallRule
RoutesFirewallRules []*RouteFirewallRule
ForwardingRules []*ForwardingRule
AuthorizedUsers map[string]map[string]struct{}
EnableSSH bool
}
func (nm *NetworkMap) Merge(other *NetworkMap) {
nm.Peers = mergeUniquePeersByID(nm.Peers, other.Peers)
nm.Routes = util.MergeUnique(nm.Routes, other.Routes)
nm.OfflinePeers = mergeUniquePeersByID(nm.OfflinePeers, other.OfflinePeers)
nm.FirewallRules = util.MergeUnique(nm.FirewallRules, other.FirewallRules)
nm.RoutesFirewallRules = util.MergeUnique(nm.RoutesFirewallRules, other.RoutesFirewallRules)
nm.ForwardingRules = util.MergeUnique(nm.ForwardingRules, other.ForwardingRules)
}
func mergeUniquePeersByID(peers1, peers2 []*nbpeer.Peer) []*nbpeer.Peer {
result := make(map[string]*nbpeer.Peer)
for _, peer := range peers1 {
result[peer.ID] = peer
}
for _, peer := range peers2 {
if _, ok := result[peer.ID]; !ok {
result[peer.ID] = peer
}
}
return maps.Values(result)
}
type ForwardingRule struct {
RuleProtocol string
DestinationPorts RulePortRange
TranslatedAddress net.IP
TranslatedPorts RulePortRange
}
func (f *ForwardingRule) ToProto() *proto.ForwardingRule {
var protocol proto.RuleProtocol
switch f.RuleProtocol {
case "icmp":
protocol = proto.RuleProtocol_ICMP
case "tcp":
protocol = proto.RuleProtocol_TCP
case "udp":
protocol = proto.RuleProtocol_UDP
case "all":
protocol = proto.RuleProtocol_ALL
default:
protocol = proto.RuleProtocol_UNKNOWN
}
return &proto.ForwardingRule{
Protocol: protocol,
DestinationPort: f.DestinationPorts.ToProto(),
TranslatedAddress: ipToBytes(f.TranslatedAddress),
TranslatedPort: f.TranslatedPorts.ToProto(),
}
}
func (f *ForwardingRule) Equal(other *ForwardingRule) bool {
return f.RuleProtocol == other.RuleProtocol &&
f.DestinationPorts.Equal(&other.DestinationPorts) &&
f.TranslatedAddress.Equal(other.TranslatedAddress) &&
f.TranslatedPorts.Equal(&other.TranslatedPorts)
}
func ipToBytes(ip net.IP) []byte {
if ip4 := ip.To4(); ip4 != nil {
return ip4
}
return ip.To16()
}
type Network struct {
Identifier string `json:"id"`
Net net.IPNet `gorm:"serializer:json"`
// NetV6 is the IPv6 ULA subnet for this account's overlay. Empty if not yet allocated.
NetV6 net.IPNet `gorm:"serializer:json"`
Dns string
// Serial is an ID that increments by 1 when any change to the network happened (e.g. new peer has been added).
// Used to synchronize state to the client apps.
Serial uint64
Mu sync.Mutex `json:"-" gorm:"-"`
}
// NewNetwork creates a new Network initializing it with a Serial=0
// It takes a random /16 subnet from 100.64.0.0/10 (64 different subnets)
// and a random /64 subnet from fd00:4e42::/32 for IPv6.
func NewNetwork() *Network {
n := iplib.NewNet4(net.ParseIP("100.64.0.0"), NetSize)
sub, _ := n.Subnet(SubnetSize)
intn := util.RandIntn(len(sub))
return &Network{
Identifier: xid.New().String(),
Net: sub[intn].IPNet,
NetV6: AllocateIPv6Subnet(),
Dns: "",
Serial: 0,
}
}
// AllocateIPv6Subnet generates a random RFC 4193 ULA /64 prefix.
// The format follows RFC 4193 section 3.1: fd + 40-bit Global ID + 16-bit Subnet ID.
// The Global ID and Subnet ID are randomized (simplified from the SHA-1 algorithm
// in section 3.2.2), giving 2^56 possible /64 subnets across all accounts.
func AllocateIPv6Subnet() net.IPNet {
ip := make(net.IP, 16)
ip[0] = 0xfd
// Bytes 1-5: 40-bit random Global ID, bytes 6-7: 16-bit random Subnet ID
if _, err := rand.Read(ip[1:8]); err != nil {
panic(err)
}
return net.IPNet{
IP: ip,
Mask: net.CIDRMask(IPv6SubnetSize, 128),
}
}
// IncSerial increments Serial by 1 reflecting that the network state has been changed
func (n *Network) IncSerial() {
n.Mu.Lock()
defer n.Mu.Unlock()
n.Serial++
}
// CurrentSerial returns the Network.Serial of the network (latest state id)
func (n *Network) CurrentSerial() uint64 {
n.Mu.Lock()
defer n.Mu.Unlock()
return n.Serial
}
func (n *Network) Copy() *Network {
n.Mu.Lock()
defer n.Mu.Unlock()
return &Network{
Identifier: n.Identifier,
Net: n.Net,
NetV6: n.NetV6,
Dns: n.Dns,
Serial: n.Serial,
}
}
// validateIPv4Prefix ensures the prefix is an IPv4 network with assignable host addresses.
func validateIPv4Prefix(prefix netip.Prefix) error {
if !prefix.IsValid() || !prefix.Addr().Is4() || prefix.Bits() < 1 || prefix.Bits() >= 31 {
return fmt.Errorf("invalid IPv4 subnet: %s", prefix.String())
}
return nil
}
// AllocatePeerIP picks an available IP from a netip.Prefix.
// This method considers already taken IPs and reuses IPs if there are gaps in takenIps.
// E.g. if prefix=100.30.0.0/16 and takenIps=[100.30.0.1, 100.30.0.4] then the result would be 100.30.0.2 or 100.30.0.3.
func AllocatePeerIP(prefix netip.Prefix, takenIps []netip.Addr) (netip.Addr, error) {
if err := validateIPv4Prefix(prefix); err != nil {
return netip.Addr{}, err
}
b := prefix.Masked().Addr().As4()
baseIP := binary.BigEndian.Uint32(b[:])
hostBits := 32 - prefix.Bits()
totalIPs := uint32(1 << hostBits)
taken := make(map[uint32]struct{}, len(takenIps)+1)
taken[baseIP] = struct{}{} // reserve network IP
taken[baseIP+totalIPs-1] = struct{}{} // reserve broadcast IP
for _, ip := range takenIps {
if !ip.Is4() {
continue
}
ab := ip.As4()
taken[binary.BigEndian.Uint32(ab[:])] = struct{}{}
}
maxAttempts := (int(totalIPs) - len(taken)) / 100
for i := 0; i < maxAttempts; i++ {
offset := uint32(util.RandIntn(int(totalIPs-2))) + 1
candidate := baseIP + offset
if _, exists := taken[candidate]; !exists {
return uint32ToIP(candidate), nil
}
}
for offset := uint32(1); offset < totalIPs-1; offset++ {
candidate := baseIP + offset
if _, exists := taken[candidate]; !exists {
return uint32ToIP(candidate), nil
}
}
return netip.Addr{}, status.Errorf(status.PreconditionFailed, "network %s is out of IPs", prefix.String())
}
// AllocateRandomPeerIP picks a random available IP from a netip.Prefix.
func AllocateRandomPeerIP(prefix netip.Prefix) (netip.Addr, error) {
if err := validateIPv4Prefix(prefix); err != nil {
return netip.Addr{}, err
}
b := prefix.Masked().Addr().As4()
baseIP := binary.BigEndian.Uint32(b[:])
hostBits := 32 - prefix.Bits()
totalIPs := uint32(1 << hostBits)
offset := uint32(util.RandIntn(int(totalIPs-2))) + 1
candidate := baseIP + offset
return uint32ToIP(candidate), nil
}
// AllocateRandomPeerIPv6 picks a random host address within the given IPv6 prefix.
// Only the host bits (after the prefix length) are randomized.
func AllocateRandomPeerIPv6(prefix netip.Prefix) (netip.Addr, error) {
ones := prefix.Bits()
if ones == 0 || ones > 126 || !prefix.Addr().Is6() {
return netip.Addr{}, fmt.Errorf("invalid IPv6 subnet: %s", prefix.String())
}
ip := prefix.Addr().As16()
// Determine which byte the host bits start in
firstHostByte := ones / 8
// If the prefix doesn't end on a byte boundary, handle the partial byte
partialBits := ones % 8
var rnd [16]byte
if _, err := rand.Read(rnd[firstHostByte:]); err != nil {
return netip.Addr{}, err
}
if partialBits > 0 {
// Keep the network bits in the partial byte, randomize the rest
hostMask := byte(0xff >> partialBits)
ip[firstHostByte] = (ip[firstHostByte] & ^hostMask) | (rnd[firstHostByte] & hostMask)
firstHostByte++
}
// Randomize remaining full host bytes
for i := firstHostByte; i < 16; i++ {
ip[i] = rnd[i]
}
// Avoid all-zeros and all-ones host parts by checking only host bits.
if isHostAllZeroOrOnes(ip[:], ones) {
ip = prefix.Masked().Addr().As16()
ip[15] |= 0x01
}
return netip.AddrFrom16(ip).Unmap(), nil
}
// isHostAllZeroOrOnes checks whether all host bits (after prefixLen) are zero or all ones.
func isHostAllZeroOrOnes(ip []byte, prefixLen int) bool {
hostStart := prefixLen / 8
partialBits := prefixLen % 8
hostSlice := slices.Clone(ip[hostStart:])
if partialBits > 0 {
hostSlice[0] &= 0xff >> partialBits
}
allZero := !slices.ContainsFunc(hostSlice, func(v byte) bool { return v != 0 })
if allZero {
return true
}
// Build the all-ones mask for host bits
onesMask := make([]byte, len(hostSlice))
for i := range onesMask {
onesMask[i] = 0xff
}
if partialBits > 0 {
onesMask[0] = 0xff >> partialBits
}
return slices.Equal(hostSlice, onesMask)
}
func uint32ToIP(n uint32) netip.Addr {
var b [4]byte
binary.BigEndian.PutUint32(b[:], n)
return netip.AddrFrom4(b)
}
// generateIPs generates a list of all possible IPs of the given network excluding IPs specified in the exclusion list
func generateIPs(ipNet *net.IPNet, exclusions map[string]struct{}) ([]net.IP, int) {
var ips []net.IP
for ip := ipNet.IP.Mask(ipNet.Mask); ipNet.Contains(ip); incIP(ip) {
if _, ok := exclusions[ip.String()]; !ok && ip[3] != 0 {
ips = append(ips, copyIP(ip))
}
}
// remove network address, broadcast and Fake DNS resolver address
lenIPs := len(ips)
switch {
case lenIPs < 2:
return ips, lenIPs
case lenIPs < 3:
return ips[1 : len(ips)-1], lenIPs - 2
default:
return ips[1 : len(ips)-2], lenIPs - 3
}
}
func copyIP(ip net.IP) net.IP {
dup := make(net.IP, len(ip))
copy(dup, ip)
return dup
}
func incIP(ip net.IP) {
for j := len(ip) - 1; j >= 0; j-- {
ip[j]++
if ip[j] > 0 {
break
}
}
}

View File

@@ -1,292 +0,0 @@
package types
import (
"encoding/binary"
"net"
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewNetwork(t *testing.T) {
network := NewNetwork()
// generated net should be a subnet of a larger 100.64.0.0/10 net
ipNet := net.IPNet{IP: net.ParseIP("100.64.0.0"), Mask: net.IPMask{255, 192, 0, 0}}
assert.Equal(t, ipNet.Contains(network.Net.IP), true)
}
func TestAllocatePeerIP(t *testing.T) {
prefix := netip.MustParsePrefix("100.64.0.0/24")
var ips []netip.Addr
for i := 0; i < 252; i++ {
ip, err := AllocatePeerIP(prefix, ips)
if err != nil {
t.Fatal(err)
}
ips = append(ips, ip)
}
assert.Len(t, ips, 252)
uniq := make(map[string]struct{})
for _, ip := range ips {
if _, ok := uniq[ip.String()]; !ok {
uniq[ip.String()] = struct{}{}
} else {
t.Errorf("found duplicate IP %s", ip.String())
}
}
}
func TestAllocatePeerIPSmallSubnet(t *testing.T) {
// Test /27 network (10.0.0.0/27) - should only have 30 usable IPs (10.0.0.1 to 10.0.0.30)
prefix := netip.MustParsePrefix("10.0.0.0/27")
var ips []netip.Addr
// Allocate all available IPs in the /27 network
for i := 0; i < 30; i++ {
ip, err := AllocatePeerIP(prefix, ips)
if err != nil {
t.Fatal(err)
}
// Verify IP is within the correct range
if !prefix.Contains(ip) {
t.Errorf("allocated IP %s is not within network %s", ip.String(), prefix.String())
}
ips = append(ips, ip)
}
assert.Len(t, ips, 30)
// Verify all IPs are unique
uniq := make(map[string]struct{})
for _, ip := range ips {
if _, ok := uniq[ip.String()]; !ok {
uniq[ip.String()] = struct{}{}
} else {
t.Errorf("found duplicate IP %s", ip.String())
}
}
// Try to allocate one more IP - should fail as network is full
_, err := AllocatePeerIP(prefix, ips)
if err == nil {
t.Error("expected error when network is full, but got none")
}
}
func TestAllocatePeerIPVariousCIDRs(t *testing.T) {
testCases := []struct {
name string
cidr string
expectedUsable int
}{
{"/30 network", "192.168.1.0/30", 2}, // 4 total - 2 reserved = 2 usable
{"/29 network", "192.168.1.0/29", 6}, // 8 total - 2 reserved = 6 usable
{"/28 network", "192.168.1.0/28", 14}, // 16 total - 2 reserved = 14 usable
{"/27 network", "192.168.1.0/27", 30}, // 32 total - 2 reserved = 30 usable
{"/26 network", "192.168.1.0/26", 62}, // 64 total - 2 reserved = 62 usable
{"/25 network", "192.168.1.0/25", 126}, // 128 total - 2 reserved = 126 usable
{"/16 network", "10.0.0.0/16", 65534}, // 65536 total - 2 reserved = 65534 usable
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
prefix, err := netip.ParsePrefix(tc.cidr)
require.NoError(t, err)
prefix = prefix.Masked()
var ips []netip.Addr
// For larger networks, test only a subset to avoid long test runs
testCount := tc.expectedUsable
if testCount > 1000 {
testCount = 1000
}
// Allocate IPs and verify they're within the correct range
for i := 0; i < testCount; i++ {
ip, err := AllocatePeerIP(prefix, ips)
require.NoError(t, err, "failed to allocate IP %d", i)
// Verify IP is within the correct range
assert.True(t, prefix.Contains(ip), "allocated IP %s is not within network %s", ip.String(), prefix.String())
// Verify IP is not network or broadcast address
networkAddr := prefix.Masked().Addr()
hostBits := 32 - prefix.Bits()
b := networkAddr.As4()
baseIP := binary.BigEndian.Uint32(b[:])
broadcastIP := uint32ToIP(baseIP + (1 << hostBits) - 1)
assert.NotEqual(t, networkAddr, ip, "allocated network address %s", ip.String())
assert.NotEqual(t, broadcastIP, ip, "allocated broadcast address %s", ip.String())
ips = append(ips, ip)
}
assert.Len(t, ips, testCount)
// Verify all IPs are unique
uniq := make(map[string]struct{})
for _, ip := range ips {
ipStr := ip.String()
assert.NotContains(t, uniq, ipStr, "found duplicate IP %s", ipStr)
uniq[ipStr] = struct{}{}
}
})
}
}
func TestAllocateIPv4InvalidPrefixes(t *testing.T) {
prefixes := []netip.Prefix{
{},
netip.MustParsePrefix("0.0.0.0/0"),
netip.MustParsePrefix("192.168.1.0/31"),
netip.MustParsePrefix("192.168.1.1/32"),
netip.MustParsePrefix("fd12:3456:7890:abcd::/64"),
}
for _, prefix := range prefixes {
t.Run(prefix.String(), func(t *testing.T) {
_, err := AllocatePeerIP(prefix, nil)
assert.Error(t, err)
_, err = AllocateRandomPeerIP(prefix)
assert.Error(t, err)
})
}
}
func TestAllocatePeerIPIgnoresNonIPv4TakenIPs(t *testing.T) {
prefix := netip.MustParsePrefix("192.168.1.0/29")
ip, err := AllocatePeerIP(prefix, []netip.Addr{netip.MustParseAddr("fd12:3456:7890:abcd::1")})
require.NoError(t, err)
assert.True(t, prefix.Contains(ip))
}
func TestGenerateIPs(t *testing.T) {
ipNet := net.IPNet{IP: net.ParseIP("100.64.0.0"), Mask: net.IPMask{255, 255, 255, 0}}
ips, ipsLen := generateIPs(&ipNet, map[string]struct{}{"100.64.0.0": {}})
if ipsLen != 252 {
t.Errorf("expected 252 ips, got %d", len(ips))
return
}
if ips[len(ips)-1].String() != "100.64.0.253" {
t.Errorf("expected last ip to be: 100.64.0.253, got %s", ips[len(ips)-1].String())
}
}
func TestNewNetworkHasIPv6(t *testing.T) {
network := NewNetwork()
assert.NotNil(t, network.NetV6.IP, "v6 subnet should be allocated")
assert.True(t, network.NetV6.IP.To4() == nil, "v6 subnet should be IPv6")
assert.Equal(t, byte(0xfd), network.NetV6.IP[0], "v6 subnet should be ULA (fd prefix)")
ones, bits := network.NetV6.Mask.Size()
assert.Equal(t, 64, ones, "v6 subnet should be /64")
assert.Equal(t, 128, bits)
}
func TestAllocateIPv6SubnetUniqueness(t *testing.T) {
seen := make(map[string]struct{})
for i := 0; i < 100; i++ {
network := NewNetwork()
key := network.NetV6.IP.String()
_, duplicate := seen[key]
assert.False(t, duplicate, "duplicate v6 subnet: %s", key)
seen[key] = struct{}{}
}
}
func TestAllocateRandomPeerIPv6(t *testing.T) {
prefix := netip.MustParsePrefix("fd12:3456:7890:abcd::/64")
ip, err := AllocateRandomPeerIPv6(prefix)
require.NoError(t, err)
assert.True(t, ip.Is6(), "should be IPv6")
assert.True(t, prefix.Contains(ip), "should be within subnet")
// First 8 bytes (network prefix) should match
b := ip.As16()
prefixBytes := prefix.Addr().As16()
assert.Equal(t, prefixBytes[:8], b[:8], "prefix should match")
// Interface ID should not be all zeros
allZero := true
for _, v := range b[8:] {
if v != 0 {
allZero = false
break
}
}
assert.False(t, allZero, "interface ID should not be all zeros")
}
func TestAllocateRandomPeerIPv6_VariousPrefixes(t *testing.T) {
tests := []struct {
name string
cidr string
prefix int
}{
{"standard /64", "fd00:1234:5678:abcd::/64", 64},
{"small /112", "fd00:1234:5678:abcd::/112", 112},
{"large /48", "fd00:1234::/48", 48},
{"non-boundary /60", "fd00:1234:5670::/60", 60},
{"non-boundary /52", "fd00:1230::/52", 52},
{"minimum /120", "fd00:1234:5678:abcd::100/120", 120},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
prefix, err := netip.ParsePrefix(tt.cidr)
require.NoError(t, err)
prefix = prefix.Masked()
assert.Equal(t, tt.prefix, prefix.Bits())
for i := 0; i < 50; i++ {
ip, err := AllocateRandomPeerIPv6(prefix)
require.NoError(t, err)
assert.True(t, prefix.Contains(ip), "IP %s should be within %s", ip, prefix)
}
})
}
}
func TestAllocateRandomPeerIPv6_PreservesNetworkBits(t *testing.T) {
// For a /112, bytes 0-13 should be preserved, only bytes 14-15 should vary
prefix := netip.MustParsePrefix("fd00:1234:5678:abcd:ef01:2345:6789:0/112")
prefixBytes := prefix.Addr().As16()
for i := 0; i < 20; i++ {
ip, err := AllocateRandomPeerIPv6(prefix)
require.NoError(t, err)
// First 14 bytes (112 bits = 14 bytes) must match the network
b := ip.As16()
assert.Equal(t, prefixBytes[:14], b[:14], "network bytes should be preserved for /112")
}
}
func TestAllocateRandomPeerIPv6_NonByteBoundary(t *testing.T) {
// For a /60, the first 7.5 bytes are network, so byte 7 is partial
prefix := netip.MustParsePrefix("fd00:1234:5678:abc0::/60")
prefixBytes := prefix.Addr().As16()
for i := 0; i < 50; i++ {
ip, err := AllocateRandomPeerIPv6(prefix)
require.NoError(t, err)
b := ip.As16()
assert.True(t, prefix.Contains(ip), "IP %s should be within %s", ip, prefix)
// First 7 bytes must match exactly
assert.Equal(t, prefixBytes[:7], b[:7], "full network bytes should match for /60")
// Byte 7: top 4 bits (0xc = 1100) must be preserved
assert.Equal(t, prefixBytes[7]&0xf0, b[7]&0xf0, "partial byte network bits should be preserved for /60")
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -1,230 +0,0 @@
package types
import (
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/route"
)
type GroupCompact struct {
Name string
PeerIndexes []int
}
type NetworkMapComponentsCompact struct {
PeerID string
Network *Network
AccountSettings *AccountSettingsInfo
DNSSettings *DNSSettings
CustomZoneDomain string
AllPeers []*nbpeer.Peer
PeerIndexes []int
RouterPeerIndexes []int
Groups map[string]*GroupCompact
AllPolicies []*Policy
PolicyIndexes []int
ResourcePoliciesMap map[string][]int
Routes []*route.Route
NameServerGroups []*nbdns.NameServerGroup
AllDNSRecords []nbdns.SimpleRecord
AccountZones []nbdns.CustomZone
RoutersMap map[string]map[string]*routerTypes.NetworkRouter
NetworkResources []*resourceTypes.NetworkResource
GroupIDToUserIDs map[string][]string
AllowedUserIDs map[string]struct{}
PostureFailedPeers map[string]map[string]struct{}
}
func (c *NetworkMapComponents) ToCompact() *NetworkMapComponentsCompact {
peerToIndex := make(map[string]int)
var allPeers []*nbpeer.Peer
for id, peer := range c.Peers {
if _, exists := peerToIndex[id]; !exists {
peerToIndex[id] = len(allPeers)
allPeers = append(allPeers, peer)
}
}
for id, peer := range c.RouterPeers {
if _, exists := peerToIndex[id]; !exists {
peerToIndex[id] = len(allPeers)
allPeers = append(allPeers, peer)
}
}
peerIndexes := make([]int, 0, len(c.Peers))
for id := range c.Peers {
peerIndexes = append(peerIndexes, peerToIndex[id])
}
routerPeerIndexes := make([]int, 0, len(c.RouterPeers))
for id := range c.RouterPeers {
routerPeerIndexes = append(routerPeerIndexes, peerToIndex[id])
}
groups := make(map[string]*GroupCompact, len(c.Groups))
for id, group := range c.Groups {
peerIdxs := make([]int, 0, len(group.Peers))
for _, peerID := range group.Peers {
if idx, ok := peerToIndex[peerID]; ok {
peerIdxs = append(peerIdxs, idx)
}
}
groups[id] = &GroupCompact{
Name: group.Name,
PeerIndexes: peerIdxs,
}
}
policyToIndex := make(map[*Policy]int)
var allPolicies []*Policy
for _, policy := range c.Policies {
if _, exists := policyToIndex[policy]; !exists {
policyToIndex[policy] = len(allPolicies)
allPolicies = append(allPolicies, policy)
}
}
for _, policies := range c.ResourcePoliciesMap {
for _, policy := range policies {
if _, exists := policyToIndex[policy]; !exists {
policyToIndex[policy] = len(allPolicies)
allPolicies = append(allPolicies, policy)
}
}
}
policyIndexes := make([]int, len(c.Policies))
for i, policy := range c.Policies {
policyIndexes[i] = policyToIndex[policy]
}
var resourcePoliciesMap map[string][]int
if len(c.ResourcePoliciesMap) > 0 {
resourcePoliciesMap = make(map[string][]int, len(c.ResourcePoliciesMap))
for resID, policies := range c.ResourcePoliciesMap {
indexes := make([]int, len(policies))
for i, policy := range policies {
indexes[i] = policyToIndex[policy]
}
resourcePoliciesMap[resID] = indexes
}
}
return &NetworkMapComponentsCompact{
PeerID: c.PeerID,
Network: c.Network,
AccountSettings: c.AccountSettings,
DNSSettings: c.DNSSettings,
CustomZoneDomain: c.CustomZoneDomain,
AllPeers: allPeers,
PeerIndexes: peerIndexes,
RouterPeerIndexes: routerPeerIndexes,
Groups: groups,
AllPolicies: allPolicies,
PolicyIndexes: policyIndexes,
ResourcePoliciesMap: resourcePoliciesMap,
Routes: c.Routes,
NameServerGroups: c.NameServerGroups,
AllDNSRecords: c.AllDNSRecords,
AccountZones: c.AccountZones,
RoutersMap: c.RoutersMap,
NetworkResources: c.NetworkResources,
GroupIDToUserIDs: c.GroupIDToUserIDs,
AllowedUserIDs: c.AllowedUserIDs,
PostureFailedPeers: c.PostureFailedPeers,
}
}
func (c *NetworkMapComponentsCompact) ToFull() *NetworkMapComponents {
peers := make(map[string]*nbpeer.Peer, len(c.PeerIndexes))
for _, idx := range c.PeerIndexes {
if idx >= 0 && idx < len(c.AllPeers) {
peer := c.AllPeers[idx]
peers[peer.ID] = peer
}
}
routerPeers := make(map[string]*nbpeer.Peer, len(c.RouterPeerIndexes))
for _, idx := range c.RouterPeerIndexes {
if idx >= 0 && idx < len(c.AllPeers) {
peer := c.AllPeers[idx]
routerPeers[peer.ID] = peer
}
}
groups := make(map[string]*Group, len(c.Groups))
for id, gc := range c.Groups {
peerIDs := make([]string, 0, len(gc.PeerIndexes))
for _, idx := range gc.PeerIndexes {
if idx >= 0 && idx < len(c.AllPeers) {
peerIDs = append(peerIDs, c.AllPeers[idx].ID)
}
}
groups[id] = &Group{
ID: id,
Name: gc.Name,
Peers: peerIDs,
}
}
policies := make([]*Policy, len(c.PolicyIndexes))
for i, idx := range c.PolicyIndexes {
if idx >= 0 && idx < len(c.AllPolicies) {
policies[i] = c.AllPolicies[idx]
}
}
var resourcePoliciesMap map[string][]*Policy
if len(c.ResourcePoliciesMap) > 0 {
resourcePoliciesMap = make(map[string][]*Policy, len(c.ResourcePoliciesMap))
for resID, indexes := range c.ResourcePoliciesMap {
pols := make([]*Policy, 0, len(indexes))
for _, idx := range indexes {
if idx >= 0 && idx < len(c.AllPolicies) {
pols = append(pols, c.AllPolicies[idx])
}
}
resourcePoliciesMap[resID] = pols
}
}
return &NetworkMapComponents{
PeerID: c.PeerID,
Network: c.Network,
AccountSettings: c.AccountSettings,
DNSSettings: c.DNSSettings,
CustomZoneDomain: c.CustomZoneDomain,
Peers: peers,
RouterPeers: routerPeers,
Groups: groups,
Policies: policies,
Routes: c.Routes,
NameServerGroups: c.NameServerGroups,
AllDNSRecords: c.AllDNSRecords,
AccountZones: c.AccountZones,
ResourcePoliciesMap: resourcePoliciesMap,
RoutersMap: c.RoutersMap,
NetworkResources: c.NetworkResources,
GroupIDToUserIDs: c.GroupIDToUserIDs,
AllowedUserIDs: c.AllowedUserIDs,
PostureFailedPeers: c.PostureFailedPeers,
}
}

View File

@@ -9,6 +9,7 @@ import (
"testing"
"time"
"github.com/rs/xid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -88,13 +89,13 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
for i := start; i < end; i++ {
groupPeers = append(groupPeers, fmt.Sprintf("peer-%d", i))
}
groups[groupID] = &types.Group{ID: groupID, Name: fmt.Sprintf("Group %d", g), Peers: groupPeers}
groups[groupID] = &types.Group{ID: groupID, PublicID: xid.New().String(), Name: fmt.Sprintf("Group %d", g), Peers: groupPeers}
}
policies := make([]*types.Policy, 0, numGroups+2)
if withDefaultPolicy {
policies = append(policies, &types.Policy{
ID: "policy-all", Name: "Default-Allow", Enabled: true,
ID: "policy-all", PublicID: xid.New().String(), Name: "Default-Allow", Enabled: true,
Rules: []*types.PolicyRule{{
ID: "rule-all", Name: "Allow All", Enabled: true, Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolALL, Bidirectional: true,
@@ -107,7 +108,7 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
groupID := fmt.Sprintf("group-%d", g)
dstGroup := fmt.Sprintf("group-%d", (g+1)%numGroups)
policies = append(policies, &types.Policy{
ID: fmt.Sprintf("policy-%d", g), Name: fmt.Sprintf("Policy %d", g), Enabled: true,
ID: fmt.Sprintf("policy-%d", g), PublicID: xid.New().String(), Name: fmt.Sprintf("Policy %d", g), Enabled: true,
Rules: []*types.PolicyRule{{
ID: fmt.Sprintf("rule-%d", g), Name: fmt.Sprintf("Rule %d", g), Enabled: true,
Action: types.PolicyTrafficActionAccept, Protocol: types.PolicyRuleProtocolTCP,
@@ -120,7 +121,7 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
if numGroups >= 2 {
policies = append(policies, &types.Policy{
ID: "policy-drop", Name: "Drop DB traffic", Enabled: true,
ID: "policy-drop", PublicID: xid.New().String(), Name: "Drop DB traffic", Enabled: true,
Rules: []*types.PolicyRule{{
ID: "rule-drop", Name: "Drop DB", Enabled: true, Action: types.PolicyTrafficActionDrop,
Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"5432"}, Bidirectional: true,
@@ -144,6 +145,7 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
groupID := fmt.Sprintf("group-%d", r%numGroups)
routes[routeID] = &route.Route{
ID: routeID,
PublicID: xid.New().String(),
Network: netip.MustParsePrefix(fmt.Sprintf("10.%d.0.0/16", r)),
Peer: peers[routePeerID].Key,
PeerID: routePeerID,
@@ -178,18 +180,18 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
}
routerPeerID := fmt.Sprintf("peer-%d", routerPeerIdx)
networksList = append(networksList, &networkTypes.Network{ID: netID, Name: fmt.Sprintf("Network %d", nr), AccountID: "test-account"})
networksList = append(networksList, &networkTypes.Network{ID: netID, PublicID: xid.New().String(), Name: fmt.Sprintf("Network %d", nr), AccountID: "test-account"})
networkResources = append(networkResources, &resourceTypes.NetworkResource{
ID: resID, NetworkID: netID, AccountID: "test-account", Enabled: true,
ID: resID, PublicID: xid.New().String(), NetworkID: netID, AccountID: "test-account", Enabled: true,
Address: fmt.Sprintf("svc-%d.netbird.cloud", nr),
})
networkRouters = append(networkRouters, &routerTypes.NetworkRouter{
ID: fmt.Sprintf("router-%d", nr), NetworkID: netID, Peer: routerPeerID,
ID: fmt.Sprintf("router-%d", nr), PublicID: xid.New().String(), NetworkID: netID, Peer: routerPeerID,
Enabled: true, AccountID: "test-account",
})
policies = append(policies, &types.Policy{
ID: fmt.Sprintf("policy-res-%d", nr), Name: fmt.Sprintf("Resource Policy %d", nr), Enabled: true,
ID: fmt.Sprintf("policy-res-%d", nr), PublicID: xid.New().String(), Name: fmt.Sprintf("Resource Policy %d", nr), Enabled: true,
SourcePostureChecks: []string{"posture-check-ver"},
Rules: []*types.PolicyRule{{
ID: fmt.Sprintf("rule-res-%d", nr), Name: fmt.Sprintf("Allow Resource %d", nr), Enabled: true,
@@ -215,12 +217,12 @@ func buildScalableTestAccount(numPeers, numGroups int, withDefaultPolicy bool) (
DNSSettings: types.DNSSettings{DisabledManagementGroups: []string{}},
NameServerGroups: map[string]*nbdns.NameServerGroup{
"ns-group-main": {
ID: "ns-group-main", Name: "Main NS", Enabled: true, Groups: []string{"group-all"},
ID: "ns-group-main", PublicID: xid.New().String(), Name: "Main NS", Enabled: true, Groups: []string{"group-all"},
NameServers: []nbdns.NameServer{{IP: netip.MustParseAddr("8.8.8.8"), NSType: nbdns.UDPNameServerType, Port: 53}},
},
},
PostureChecks: []*posture.Checks{
{ID: "posture-check-ver", Name: "Check version", Checks: posture.ChecksDefinition{
{ID: "posture-check-ver", PublicID: xid.New().String(), Name: "Check version", Checks: posture.ChecksDefinition{
NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.26.0"},
}},
},

View File

@@ -49,7 +49,7 @@ func allPeersValidated(account *types.Account, excludePeerIDs ...string) map[str
return validated
}
func peerIDs(peers []*nbpeer.Peer) []string {
func peerIDs(peers []*types.ComponentPeer) []string {
ids := make([]string, len(peers))
for i, p := range peers {
ids[i] = p.ID

View File

@@ -0,0 +1,163 @@
package types_test
import (
"context"
"crypto/rand"
"encoding/base64"
"fmt"
"testing"
goproto "google.golang.org/protobuf/proto"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/types"
)
// wireBenchScales — trimmed scale set for wire-size measurements. Encoding
// and marshalling are linear, so the largest extremes don't add signal.
var wireBenchScales = []benchmarkScale{
{"100peers_5groups", 100, 5},
{"500peers_20groups", 500, 20},
{"1000peers_50groups", 1000, 50},
{"5000peers_100groups", 5000, 100},
}
// assignValidWgKeys overwrites every peer's Key with a valid base64-encoded
// 32-byte string. The default scalableTestAccount uses unparsable strings
// like "key-peer-0", which makes the components encoder emit a nil WgPubKey
// and the legacy encoder ship 10-char placeholders — both shrink the wire
// size in unrealistic ways. Production peers always have valid 44-char base64
// keys, so any benchmark/breakdown that wants honest numbers must call this.
func assignValidWgKeys(account *types.Account) {
for _, p := range account.Peers {
var raw [32]byte
_, _ = rand.Read(raw[:])
p.Key = base64.StdEncoding.EncodeToString(raw[:])
}
}
// BenchmarkNetworkMapWireEncode reports per-call ns and the marshaled wire
// size for both encoding paths. Run with:
//
// go test -run=^$ -bench=BenchmarkNetworkMapWireEncode -benchmem ./management/server/types/
func BenchmarkNetworkMapWireEncode(b *testing.B) {
skipCIBenchmark(b)
for _, scale := range wireBenchScales {
account, validatedPeers := scalableTestAccount(scale.peers, scale.groups)
// populateAccountSeqIDs(account)
assignValidWgKeys(account)
ctx := context.Background()
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
peerID := "peer-0"
peer := account.Peers[peerID]
networkMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
components := account.GetPeerNetworkMapComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, groupIDToUserIDs)
dnsCache := &cache.DNSConfigCache{}
settings := &types.Settings{}
// Pre-encode once so the size metric is identical for every run inside
// the same scale; the b.Loop call only re-runs encode + Marshal.
legacyResp := mgmtgrpc.ToSyncResponse(ctx, nil, nil, nil, peer, nil, nil, networkMap, "netbird.cloud", nil, dnsCache, settings, nil, nil, 0)
legacyBytes, err := goproto.Marshal(legacyResp.NetworkMap)
if err != nil {
b.Fatalf("marshal legacy networkmap: %v", err)
}
envelopeInput := mgmtgrpc.ComponentsEnvelopeInput{
Components: components,
PeerConfig: legacyResp.NetworkMap.PeerConfig,
DNSDomain: "netbird.cloud",
}
envelope := mgmtgrpc.EncodeNetworkMapEnvelope(envelopeInput)
envelopeBytes, err := goproto.Marshal(envelope)
if err != nil {
b.Fatalf("marshal envelope: %v", err)
}
b.Run(fmt.Sprintf("legacy/%s", scale.name), func(b *testing.B) {
b.ReportAllocs()
b.ReportMetric(float64(len(legacyBytes)), "bytes/msg")
b.ResetTimer()
for range b.N {
resp := mgmtgrpc.ToSyncResponse(ctx, nil, nil, nil, peer, nil, nil, networkMap, "netbird.cloud", nil, dnsCache, settings, nil, nil, 0)
if _, err := goproto.Marshal(resp.NetworkMap); err != nil {
b.Fatal(err)
}
}
})
b.Run(fmt.Sprintf("components/%s", scale.name), func(b *testing.B) {
b.ReportAllocs()
b.ReportMetric(float64(len(envelopeBytes)), "bytes/msg")
b.ResetTimer()
for range b.N {
env := mgmtgrpc.EncodeNetworkMapEnvelope(envelopeInput)
if _, err := goproto.Marshal(env); err != nil {
b.Fatal(err)
}
}
})
}
}
// BenchmarkNetworkMapWireSize is a fast snapshot of the wire size by scale
// without a tight encode loop. Run with -bench to see one ns/op + bytes per
// scale (treat the timing as informational; the sample is one Marshal per
// scale, not the full b.N loop).
func BenchmarkNetworkMapWireSize(b *testing.B) {
skipCIBenchmark(b)
for _, scale := range wireBenchScales {
account, validatedPeers := scalableTestAccount(scale.peers, scale.groups)
// populateAccountSeqIDs(account)
assignValidWgKeys(account)
ctx := context.Background()
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
peerID := "peer-0"
peer := account.Peers[peerID]
networkMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
components := account.GetPeerNetworkMapComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, groupIDToUserIDs)
dnsCache := &cache.DNSConfigCache{}
settings := &types.Settings{}
legacyResp := mgmtgrpc.ToSyncResponse(ctx, nil, nil, nil, peer, nil, nil, networkMap, "netbird.cloud", nil, dnsCache, settings, nil, nil, 0)
legacyBytes, err := goproto.Marshal(legacyResp.NetworkMap)
if err != nil {
b.Fatalf("marshal legacy networkmap: %v", err)
}
env := mgmtgrpc.EncodeNetworkMapEnvelope(mgmtgrpc.ComponentsEnvelopeInput{
Components: components,
PeerConfig: legacyResp.NetworkMap.PeerConfig,
DNSDomain: "netbird.cloud",
})
envBytes, err := goproto.Marshal(env)
if err != nil {
b.Fatalf("marshal envelope: %v", err)
}
b.Run(fmt.Sprintf("size/%s", scale.name), func(b *testing.B) {
b.ReportMetric(float64(len(legacyBytes)), "legacy_bytes")
b.ReportMetric(float64(len(envBytes)), "components_bytes")
ratio := float64(len(envBytes)) / float64(len(legacyBytes))
b.ReportMetric(ratio, "components/legacy")
for range b.N {
}
})
}
}

View File

@@ -0,0 +1,149 @@
package types_test
import (
"context"
"fmt"
"os"
"testing"
goproto "google.golang.org/protobuf/proto"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// TestNetworkMapWireBreakdown is a one-shot diagnostic: it computes the wire
// size attributable to each top-level field of both the legacy NetworkMap and
// the components NetworkMapEnvelope at the 5000-peer scale, so the migration
// docs can attribute the size reduction to each optimization. Runs only on
// demand via -run TestNetworkMapWireBreakdown.
func TestNetworkMapWireBreakdown(t *testing.T) {
if testing.Short() {
t.Skip("size diagnostic, skipped with -short")
}
if os.Getenv("NB_RUN_WIRE_BREAKDOWN") != "1" {
t.Skip("set NB_RUN_WIRE_BREAKDOWN=1 to run wire breakdown diagnostic")
}
const peerCount, groupCount = 5000, 100
account, validatedPeers := scalableTestAccount(peerCount, groupCount)
assignValidWgKeys(account)
ctx := context.Background()
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
peerID := "peer-0"
peer := account.Peers[peerID]
networkMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
components := account.GetPeerNetworkMapComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, groupIDToUserIDs)
dnsCache := &cache.DNSConfigCache{}
settings := &types.Settings{}
legacyResp := mgmtgrpc.ToSyncResponse(ctx, nil, nil, nil, peer, nil, nil, networkMap, "netbird.cloud", nil, dnsCache, settings, nil, nil, 0)
legacyTotal := mustMarshalSize(t, legacyResp.NetworkMap)
envelope := mgmtgrpc.EncodeNetworkMapEnvelope(mgmtgrpc.ComponentsEnvelopeInput{
Components: components,
PeerConfig: legacyResp.NetworkMap.PeerConfig,
DNSDomain: "netbird.cloud",
})
componentsTotal := mustMarshalSize(t, envelope)
t.Logf("\n=== LEGACY NetworkMap (%d peers, %d groups) ===", peerCount, groupCount)
t.Logf(" Total: %d bytes\n", legacyTotal)
legacyBreakdown := []struct {
name string
nm *proto.NetworkMap
}{
{"RemotePeers", &proto.NetworkMap{RemotePeers: legacyResp.NetworkMap.RemotePeers}},
{"OfflinePeers", &proto.NetworkMap{OfflinePeers: legacyResp.NetworkMap.OfflinePeers}},
{"FirewallRules", &proto.NetworkMap{FirewallRules: legacyResp.NetworkMap.FirewallRules}},
{"Routes", &proto.NetworkMap{Routes: legacyResp.NetworkMap.Routes}},
{"RoutesFirewallRules", &proto.NetworkMap{RoutesFirewallRules: legacyResp.NetworkMap.RoutesFirewallRules}},
{"DNSConfig", &proto.NetworkMap{DNSConfig: legacyResp.NetworkMap.DNSConfig}},
{"PeerConfig", &proto.NetworkMap{PeerConfig: legacyResp.NetworkMap.PeerConfig}},
{"SshAuth", &proto.NetworkMap{SshAuth: legacyResp.NetworkMap.SshAuth}},
}
for _, e := range legacyBreakdown {
size := mustMarshalSize(t, e.nm)
t.Logf(" %-22s %8d bytes %5.1f%%", e.name, size, pct(size, legacyTotal))
}
full := envelope.GetFull()
if full == nil {
t.Fatalf("expected full network map envelope payload, got nil")
}
t.Logf("\n=== COMPONENTS NetworkMapEnvelope (%d peers, %d groups) ===", peerCount, groupCount)
t.Logf(" Total: %d bytes (%.1f%% of legacy)\n", componentsTotal, pct(componentsTotal, legacyTotal))
componentsBreakdown := []struct {
name string
nm *proto.NetworkMapComponentsFull
}{
{"Peers", &proto.NetworkMapComponentsFull{Peers: full.Peers}},
{"Policies", &proto.NetworkMapComponentsFull{Policies: full.Policies}},
{"Groups", &proto.NetworkMapComponentsFull{Groups: full.Groups}},
{"Routes (raw)", &proto.NetworkMapComponentsFull{Routes: full.Routes}},
{"NameServerGroups", &proto.NetworkMapComponentsFull{NameserverGroups: full.NameserverGroups}},
{"AllDNSRecords", &proto.NetworkMapComponentsFull{AllDnsRecords: full.AllDnsRecords}},
{"AccountZones", &proto.NetworkMapComponentsFull{AccountZones: full.AccountZones}},
{"NetworkResources", &proto.NetworkMapComponentsFull{NetworkResources: full.NetworkResources}},
{"RoutersMap", &proto.NetworkMapComponentsFull{RoutersMap: full.RoutersMap}},
{"ResourcePoliciesMap", &proto.NetworkMapComponentsFull{ResourcePoliciesMap: full.ResourcePoliciesMap}},
{"GroupIDToUserIDs", &proto.NetworkMapComponentsFull{GroupIdToUserIds: full.GroupIdToUserIds}},
{"AllowedUserIDs", &proto.NetworkMapComponentsFull{AllowedUserIds: full.AllowedUserIds}},
{"PostureFailedPeers", &proto.NetworkMapComponentsFull{PostureFailedPeers: full.PostureFailedPeers}},
{"DNSSettings", &proto.NetworkMapComponentsFull{DnsSettings: full.DnsSettings}},
{"PeerConfig", &proto.NetworkMapComponentsFull{PeerConfig: full.PeerConfig}},
{"AgentVersions", &proto.NetworkMapComponentsFull{AgentVersions: full.AgentVersions}},
}
for _, e := range componentsBreakdown {
size := mustMarshalSize(t, e.nm)
t.Logf(" %-22s %8d bytes %5.1f%%", e.name, size, pct(size, componentsTotal))
}
t.Logf("\n=== Per-PeerCompact average ===")
if len(full.Peers) > 0 {
t.Logf(" PeerCompact avg: %d bytes/peer", mustMarshalSize(t, &proto.NetworkMapComponentsFull{Peers: full.Peers})/len(full.Peers))
}
if len(legacyResp.NetworkMap.RemotePeers) > 0 {
t.Logf(" RemotePeer avg: %d bytes/peer",
mustMarshalSize(t, &proto.NetworkMap{RemotePeers: legacyResp.NetworkMap.RemotePeers})/len(legacyResp.NetworkMap.RemotePeers))
}
t.Logf("\n=== FirewallRule expansion footprint ===")
t.Logf(" legacy FirewallRules count: %d", len(legacyResp.NetworkMap.FirewallRules))
t.Logf(" components Policies count: %d", len(full.Policies))
t.Logf(" components Groups count: %d", len(full.Groups))
totalGroupPeerIdxs := 0
for _, g := range full.Groups {
totalGroupPeerIdxs += len(g.PeerIndexes)
}
t.Logf(" components peer-index refs across all groups: %d", totalGroupPeerIdxs)
}
func mustMarshalSize(t *testing.T, m goproto.Message) int {
b, err := goproto.Marshal(m)
if err != nil {
t.Fatalf("marshal: %v", err)
}
return len(b)
}
func pct(part, total int) float64 {
if total == 0 {
return 0
}
return 100 * float64(part) / float64(total)
}
// Stops fmt being unused if the breakdown loop above is later commented out.
var _ = fmt.Sprintf

View File

@@ -0,0 +1,25 @@
package types
// PeerNetworkMapResult is what the network_map controller produces for a
// single peer. Exactly one of NetworkMap or Components is populated depending
// on the peer's capability:
//
// - Components-capable peers (PeerCapabilityComponentNetworkMap) get
// Components: the raw types.NetworkMapComponents the client decodes and
// runs Calculate() on locally. NetworkMap stays nil — the server skips
// the expansion entirely.
// - Legacy peers (or any peer when the kill switch is set) get NetworkMap:
// the fully-expanded view the legacy gRPC path consumes.
//
// The gRPC layer (ToSyncResponseForPeer) dispatches by which field is
// non-nil; callers must not rely on both being set.
type PeerNetworkMapResult struct {
NetworkMap *NetworkMap
Components *NetworkMapComponents
}
// IsComponents reports whether the result carries the components shape.
// Use this in preference to direct nil checks on the fields.
func (r PeerNetworkMapResult) IsComponents() bool {
return r.Components != nil
}

View File

@@ -0,0 +1,104 @@
package types_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
nbdns "github.com/netbirdio/netbird/dns"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
// helper: marks the given peer as components-capable.
func markCapable(p *nbpeer.Peer) {
p.Meta.Capabilities = append(p.Meta.Capabilities, nbpeer.PeerCapabilityComponentNetworkMap)
}
func TestGetPeerNetworkMapResult_CapablePeerGetsComponents(t *testing.T) {
account, validatedPeers := scalableTestAccount(10, 2)
markCapable(account.Peers["peer-0"])
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
result := account.GetPeerNetworkMapResult(
context.Background(),
"peer-0",
false, // componentsDisabled
nbdns.CustomZone{},
nil,
validatedPeers,
resourcePolicies,
routers,
nil,
groupIDToUserIDs,
)
require.True(t, result.IsComponents(), "capable peer must get the components shape")
assert.Nil(t, result.NetworkMap)
require.NotNil(t, result.Components)
assert.Equal(t, "peer-0", result.Components.PeerID)
}
func TestGetPeerNetworkMapResult_LegacyPeerGetsNetworkMap(t *testing.T) {
account, validatedPeers := scalableTestAccount(10, 2)
// peer-0 left without the component capability
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
result := account.GetPeerNetworkMapResult(
context.Background(),
"peer-0",
false,
nbdns.CustomZone{},
nil,
validatedPeers,
resourcePolicies,
routers,
nil,
groupIDToUserIDs,
)
assert.False(t, result.IsComponents())
assert.Nil(t, result.Components)
require.NotNil(t, result.NetworkMap, "legacy peer must get a NetworkMap")
}
func TestGetPeerNetworkMapResult_KillSwitchOverridesCapability(t *testing.T) {
// Capable peer + componentsDisabled=true → falls back to legacy.
account, validatedPeers := scalableTestAccount(10, 2)
markCapable(account.Peers["peer-0"])
resourcePolicies := account.GetResourcePoliciesMap()
routers := account.GetResourceRoutersMap()
groupIDToUserIDs := account.GetActiveGroupUsers()
result := account.GetPeerNetworkMapResult(
context.Background(),
"peer-0",
true, // componentsDisabled = true (kill switch)
nbdns.CustomZone{},
nil,
validatedPeers,
resourcePolicies,
routers,
nil,
groupIDToUserIDs,
)
assert.False(t, result.IsComponents(), "kill switch must force legacy NetworkMap path")
assert.Nil(t, result.Components)
require.NotNil(t, result.NetworkMap)
}
func TestPeerNetworkMapResult_IsComponents(t *testing.T) {
assert.True(t, types.PeerNetworkMapResult{Components: &types.NetworkMapComponents{}}.IsComponents())
assert.False(t, types.PeerNetworkMapResult{NetworkMap: &types.NetworkMap{}}.IsComponents())
assert.False(t, types.PeerNetworkMapResult{}.IsComponents())
}

View File

@@ -1,265 +0,0 @@
package types
import (
"errors"
"fmt"
"strconv"
"strings"
)
const (
// PolicyTrafficActionAccept indicates that the traffic is accepted
PolicyTrafficActionAccept = PolicyTrafficActionType("accept")
// PolicyTrafficActionDrop indicates that the traffic is dropped
PolicyTrafficActionDrop = PolicyTrafficActionType("drop")
)
const (
// PolicyRuleProtocolALL type of traffic
PolicyRuleProtocolALL = PolicyRuleProtocolType("all")
// PolicyRuleProtocolTCP type of traffic
PolicyRuleProtocolTCP = PolicyRuleProtocolType("tcp")
// PolicyRuleProtocolUDP type of traffic
PolicyRuleProtocolUDP = PolicyRuleProtocolType("udp")
// PolicyRuleProtocolICMP type of traffic
PolicyRuleProtocolICMP = PolicyRuleProtocolType("icmp")
// PolicyRuleProtocolNetbirdSSH type of traffic
PolicyRuleProtocolNetbirdSSH = PolicyRuleProtocolType("netbird-ssh")
)
const (
// PolicyRuleFlowDirect allows traffic from source to destination
PolicyRuleFlowDirect = PolicyRuleDirection("direct")
// PolicyRuleFlowBidirect allows traffic to both directions
PolicyRuleFlowBidirect = PolicyRuleDirection("bidirect")
)
const (
// DefaultRuleName is a name for the Default rule that is created for every account
DefaultRuleName = "Default"
// DefaultRuleDescription is a description for the Default rule that is created for every account
DefaultRuleDescription = "This is a default rule that allows connections between all the resources"
// DefaultPolicyName is a name for the Default policy that is created for every account
DefaultPolicyName = "Default"
// DefaultPolicyDescription is a description for the Default policy that is created for every account
DefaultPolicyDescription = "This is a default policy that allows connections between all the resources"
)
// PolicyUpdateOperation operation object with type and values to be applied
type PolicyUpdateOperation struct {
Type PolicyUpdateOperationType
Values []string
}
// Policy of the Rego query
type Policy struct {
// ID of the policy'
ID string `gorm:"primaryKey"`
// AccountID is a reference to Account that this object belongs
AccountID string `json:"-" gorm:"index"`
// Name of the Policy
Name string
// Description of the policy visible in the UI
Description string
// Enabled status of the policy
Enabled bool
// Rules of the policy
Rules []*PolicyRule `gorm:"foreignKey:PolicyID;references:id;constraint:OnDelete:CASCADE;"`
// SourcePostureChecks are ID references to Posture checks for policy source groups
SourcePostureChecks []string `gorm:"serializer:json"`
}
// Copy returns a copy of the policy.
func (p *Policy) Copy() *Policy {
c := &Policy{
ID: p.ID,
AccountID: p.AccountID,
Name: p.Name,
Description: p.Description,
Enabled: p.Enabled,
Rules: make([]*PolicyRule, len(p.Rules)),
SourcePostureChecks: make([]string, len(p.SourcePostureChecks)),
}
for i, r := range p.Rules {
c.Rules[i] = r.Copy()
}
copy(c.SourcePostureChecks, p.SourcePostureChecks)
return c
}
func (p *Policy) Equal(other *Policy) bool {
if p == nil || other == nil {
return p == other
}
if p.ID != other.ID ||
p.AccountID != other.AccountID ||
p.Name != other.Name ||
p.Description != other.Description ||
p.Enabled != other.Enabled {
return false
}
if !stringSlicesEqualUnordered(p.SourcePostureChecks, other.SourcePostureChecks) {
return false
}
if len(p.Rules) != len(other.Rules) {
return false
}
otherRules := make(map[string]*PolicyRule, len(other.Rules))
for _, r := range other.Rules {
otherRules[r.ID] = r
}
for _, r := range p.Rules {
otherRule, ok := otherRules[r.ID]
if !ok {
return false
}
if !r.Equal(otherRule) {
return false
}
}
return true
}
// EventMeta returns activity event meta related to this policy
func (p *Policy) EventMeta() map[string]any {
return map[string]any{"name": p.Name}
}
// UpgradeAndFix different version of policies to latest version
func (p *Policy) UpgradeAndFix() {
for _, r := range p.Rules {
// start migrate from version v0.20.3
if r.Protocol == "" {
r.Protocol = PolicyRuleProtocolALL
}
if r.Protocol == PolicyRuleProtocolALL && !r.Bidirectional {
r.Bidirectional = true
}
// -- v0.20.4
}
}
// RuleGroups returns a list of all groups referenced in the policy's rules,
// including sources and destinations.
func (p *Policy) RuleGroups() []string {
groups := make([]string, 0)
for _, rule := range p.Rules {
groups = append(groups, rule.Sources...)
groups = append(groups, rule.Destinations...)
}
return groups
}
// SourceGroups returns a slice of all unique source groups referenced in the policy's rules.
func (p *Policy) SourceGroups() []string {
if len(p.Rules) == 1 {
return p.Rules[0].Sources
}
groups := make(map[string]struct{}, len(p.Rules))
for _, rule := range p.Rules {
for _, source := range rule.Sources {
groups[source] = struct{}{}
}
}
groupIDs := make([]string, 0, len(groups))
for groupID := range groups {
groupIDs = append(groupIDs, groupID)
}
return groupIDs
}
func ParseRuleString(rule string) (PolicyRuleProtocolType, RulePortRange, error) {
rule = strings.TrimSpace(strings.ToLower(rule))
if rule == "all" {
return PolicyRuleProtocolALL, RulePortRange{}, nil
}
if rule == "icmp" {
return PolicyRuleProtocolICMP, RulePortRange{}, nil
}
split := strings.Split(rule, "/")
if len(split) != 2 {
return "", RulePortRange{}, errors.New("invalid rule format: expected protocol/port or protocol/port-range")
}
protoStr := strings.TrimSpace(split[0])
portStr := strings.TrimSpace(split[1])
var protocol PolicyRuleProtocolType
switch protoStr {
case "tcp":
protocol = PolicyRuleProtocolTCP
case "udp":
protocol = PolicyRuleProtocolUDP
case "icmp":
return "", RulePortRange{}, errors.New("icmp does not accept ports; use 'icmp' without '/…'")
case "netbird-ssh":
return PolicyRuleProtocolNetbirdSSH, RulePortRange{Start: nativeSSHPortNumber, End: nativeSSHPortNumber}, nil
default:
return "", RulePortRange{}, fmt.Errorf("invalid protocol: %q", protoStr)
}
portRange, err := parsePortRange(portStr)
if err != nil {
return "", RulePortRange{}, err
}
return protocol, portRange, nil
}
func parsePortRange(portStr string) (RulePortRange, error) {
if strings.Contains(portStr, "-") {
rangeParts := strings.Split(portStr, "-")
if len(rangeParts) != 2 {
return RulePortRange{}, fmt.Errorf("invalid port range %q", portStr)
}
start, err := parsePort(strings.TrimSpace(rangeParts[0]))
if err != nil {
return RulePortRange{}, err
}
end, err := parsePort(strings.TrimSpace(rangeParts[1]))
if err != nil {
return RulePortRange{}, err
}
if start > end {
return RulePortRange{}, fmt.Errorf("invalid port range: start %d > end %d", start, end)
}
return RulePortRange{Start: uint16(start), End: uint16(end)}, nil
}
p, err := parsePort(portStr)
if err != nil {
return RulePortRange{}, err
}
return RulePortRange{Start: uint16(p), End: uint16(p)}, nil
}
func parsePort(portStr string) (int, error) {
if portStr == "" {
return 0, errors.New("empty port")
}
p, err := strconv.Atoi(portStr)
if err != nil {
return 0, fmt.Errorf("invalid port %q: %w", portStr, err)
}
if p < 1 || p > 65535 {
return 0, fmt.Errorf("port out of range (165535): %d", p)
}
return p, nil
}

View File

@@ -1,225 +0,0 @@
package types
import (
"slices"
"github.com/netbirdio/netbird/shared/management/proto"
)
// PolicyUpdateOperationType operation type
type PolicyUpdateOperationType int
// PolicyTrafficActionType action type for the firewall
type PolicyTrafficActionType string
// PolicyRuleProtocolType type of traffic
type PolicyRuleProtocolType string
// PolicyRuleDirection direction of traffic
type PolicyRuleDirection string
// RulePortRange represents a range of ports for a firewall rule.
type RulePortRange struct {
Start uint16
End uint16
}
func (r *RulePortRange) ToProto() *proto.PortInfo {
return &proto.PortInfo{
PortSelection: &proto.PortInfo_Range_{
Range: &proto.PortInfo_Range{
Start: uint32(r.Start),
End: uint32(r.End),
},
},
}
}
func (r *RulePortRange) Equal(other *RulePortRange) bool {
return r.Start == other.Start && r.End == other.End
}
// PolicyRule is the metadata of the policy
type PolicyRule struct {
// ID of the policy rule
ID string `gorm:"primaryKey"`
// PolicyID is a reference to Policy that this object belongs
PolicyID string `json:"-" gorm:"index"`
// Name of the rule visible in the UI
Name string
// Description of the rule visible in the UI
Description string
// Enabled status of rule in the system
Enabled bool
// Action policy accept or drops packets
Action PolicyTrafficActionType
// Destinations policy destination groups
Destinations []string `gorm:"serializer:json"`
// DestinationResource policy destination resource that the rule is applied to
DestinationResource Resource `gorm:"serializer:json"`
// Sources policy source groups
Sources []string `gorm:"serializer:json"`
// SourceResource policy source resource that the rule is applied to
SourceResource Resource `gorm:"serializer:json"`
// Bidirectional define if the rule is applicable in both directions, sources, and destinations
Bidirectional bool
// Protocol type of the traffic
Protocol PolicyRuleProtocolType
// Ports or it ranges list
Ports []string `gorm:"serializer:json"`
// PortRanges a list of port ranges.
PortRanges []RulePortRange `gorm:"serializer:json"`
// AuthorizedGroups is a map of groupIDs and their respective access to local users via ssh
AuthorizedGroups map[string][]string `gorm:"serializer:json"`
// AuthorizedUser is a list of userIDs that are authorized to access local resources via ssh
AuthorizedUser string
}
// Copy returns a copy of a policy rule
func (pm *PolicyRule) Copy() *PolicyRule {
rule := &PolicyRule{
ID: pm.ID,
PolicyID: pm.PolicyID,
Name: pm.Name,
Description: pm.Description,
Enabled: pm.Enabled,
Action: pm.Action,
Destinations: make([]string, len(pm.Destinations)),
DestinationResource: pm.DestinationResource,
Sources: make([]string, len(pm.Sources)),
SourceResource: pm.SourceResource,
Bidirectional: pm.Bidirectional,
Protocol: pm.Protocol,
Ports: make([]string, len(pm.Ports)),
PortRanges: make([]RulePortRange, len(pm.PortRanges)),
AuthorizedGroups: make(map[string][]string, len(pm.AuthorizedGroups)),
AuthorizedUser: pm.AuthorizedUser,
}
copy(rule.Destinations, pm.Destinations)
copy(rule.Sources, pm.Sources)
copy(rule.Ports, pm.Ports)
copy(rule.PortRanges, pm.PortRanges)
for k, v := range pm.AuthorizedGroups {
rule.AuthorizedGroups[k] = make([]string, len(v))
copy(rule.AuthorizedGroups[k], v)
}
return rule
}
func (pm *PolicyRule) Equal(other *PolicyRule) bool {
if pm == nil || other == nil {
return pm == other
}
if pm.ID != other.ID ||
pm.PolicyID != other.PolicyID ||
pm.Name != other.Name ||
pm.Description != other.Description ||
pm.Enabled != other.Enabled ||
pm.Action != other.Action ||
pm.Bidirectional != other.Bidirectional ||
pm.Protocol != other.Protocol ||
pm.SourceResource != other.SourceResource ||
pm.DestinationResource != other.DestinationResource ||
pm.AuthorizedUser != other.AuthorizedUser {
return false
}
if !stringSlicesEqualUnordered(pm.Sources, other.Sources) {
return false
}
if !stringSlicesEqualUnordered(pm.Destinations, other.Destinations) {
return false
}
if !stringSlicesEqualUnordered(pm.Ports, other.Ports) {
return false
}
if !portRangeSlicesEqualUnordered(pm.PortRanges, other.PortRanges) {
return false
}
if !authorizedGroupsEqual(pm.AuthorizedGroups, other.AuthorizedGroups) {
return false
}
return true
}
func stringSlicesEqualUnordered(a, b []string) bool {
if len(a) != len(b) {
return false
}
if len(a) == 0 {
return true
}
sorted1 := make([]string, len(a))
sorted2 := make([]string, len(b))
copy(sorted1, a)
copy(sorted2, b)
slices.Sort(sorted1)
slices.Sort(sorted2)
return slices.Equal(sorted1, sorted2)
}
func portRangeSlicesEqualUnordered(a, b []RulePortRange) bool {
if len(a) != len(b) {
return false
}
if len(a) == 0 {
return true
}
cmp := func(x, y RulePortRange) int {
if x.Start != y.Start {
if x.Start < y.Start {
return -1
}
return 1
}
if x.End != y.End {
if x.End < y.End {
return -1
}
return 1
}
return 0
}
sorted1 := make([]RulePortRange, len(a))
sorted2 := make([]RulePortRange, len(b))
copy(sorted1, a)
copy(sorted2, b)
slices.SortFunc(sorted1, cmp)
slices.SortFunc(sorted2, cmp)
return slices.EqualFunc(sorted1, sorted2, func(x, y RulePortRange) bool {
return x.Start == y.Start && x.End == y.End
})
}
func authorizedGroupsEqual(a, b map[string][]string) bool {
if len(a) != len(b) {
return false
}
for k, va := range a {
vb, ok := b[k]
if !ok {
return false
}
if !stringSlicesEqualUnordered(va, vb) {
return false
}
}
return true
}

View File

@@ -1,39 +0,0 @@
package types
import (
"github.com/netbirdio/netbird/shared/management/http/api"
)
type ResourceType string
const (
ResourceTypePeer ResourceType = "peer"
ResourceTypeDomain ResourceType = "domain"
ResourceTypeHost ResourceType = "host"
ResourceTypeSubnet ResourceType = "subnet"
)
type Resource struct {
ID string
Type ResourceType
}
func (r *Resource) ToAPIResponse() *api.Resource {
if r.ID == "" && r.Type == "" {
return nil
}
return &api.Resource{
Id: r.ID,
Type: api.ResourceType(r.Type),
}
}
func (r *Resource) FromAPIRequest(req *api.Resource) {
if req == nil {
return
}
r.ID = req.Id
r.Type = ResourceType(req.Type)
}

Some files were not shown because too many files have changed in this diff Show More