mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-21 04:49:10 +02:00
Merge branch 'main' into fix-ssh-authorized-users-multi-rule
# Conflicts: # management/server/types/account.go # shared/management/types/networkmap_components.go
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"os"
|
||||
"os/signal"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
@@ -22,9 +23,11 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
agentnetworkpricing "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
|
||||
"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"
|
||||
)
|
||||
@@ -110,6 +113,29 @@ var (
|
||||
mgmtSingleAccModeDomain = ""
|
||||
}
|
||||
|
||||
// Load the management-side LLM pricing defaults file: an
|
||||
// explicitly configured path is required to load (a typo must
|
||||
// fail startup — the operator believes those rates are live);
|
||||
// otherwise <datadir>/defaults_llm_pricing.yaml is probed and
|
||||
// may be absent (compiled-in defaults serve). A relative path
|
||||
// is resolved against the datadir so a bare filename lands
|
||||
// alongside the store. Either way the path stays watched: the
|
||||
// reloader picks up edits — and the file appearing later —
|
||||
// without a restart.
|
||||
pricingPath := config.AgentNetwork.PricingDefaultsFile
|
||||
pricingRequired := pricingPath != ""
|
||||
if !pricingRequired {
|
||||
pricingPath = agentnetworkpricing.DefaultFileName
|
||||
}
|
||||
if !filepath.IsAbs(pricingPath) {
|
||||
pricingPath = filepath.Join(config.Datadir, pricingPath)
|
||||
}
|
||||
log.Infof("loading agent-network pricing defaults from %s (required: %v)", pricingPath, pricingRequired)
|
||||
if err := agentnetworkpricing.LoadFile(pricingPath, pricingRequired); err != nil {
|
||||
return fmt.Errorf("load agent-network pricing defaults: %v", err)
|
||||
}
|
||||
agentnetworkpricing.StartReloader(ctx, agentnetworkpricing.ReloadInterval)
|
||||
|
||||
srv := newServer(&server.Config{
|
||||
NbConfig: config,
|
||||
DNSDomain: dnsDomain,
|
||||
@@ -153,8 +179,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 +247,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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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.")
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -7,12 +7,28 @@ package catalog
|
||||
import "github.com/netbirdio/netbird/shared/management/http/api"
|
||||
|
||||
// Model is the in-memory representation of a catalog model.
|
||||
//
|
||||
// The three cache rates mirror the proxy cost meter's Entry semantics
|
||||
// (USD per 1k tokens; 0 = no rate configured, that bucket bills at
|
||||
// InputPer1k):
|
||||
// - CachedInputPer1k: OpenAI-shape rate for cached prompt tokens
|
||||
// (a SUBSET of input tokens). Typically 0.1-0.5x input.
|
||||
// - CacheReadPer1k / CacheCreationPer1k: Anthropic-shape rates for
|
||||
// the two ADDITIVE prompt-cache buckets. Typically 0.1x / 1.25x
|
||||
// input.
|
||||
//
|
||||
// The catalog is the single default-pricing source: the agentnetwork
|
||||
// pricing package folds these models into per-surface tables that the
|
||||
// synthesizer ships to the proxy's cost_meter.
|
||||
type Model struct {
|
||||
ID string
|
||||
Label string
|
||||
InputPer1k float64
|
||||
OutputPer1k float64
|
||||
ContextWindow int
|
||||
ID string
|
||||
Label string
|
||||
InputPer1k float64
|
||||
OutputPer1k float64
|
||||
CachedInputPer1k float64
|
||||
CacheReadPer1k float64
|
||||
CacheCreationPer1k float64
|
||||
ContextWindow int
|
||||
}
|
||||
|
||||
// ProviderKind groups catalog entries for UI presentation. The split
|
||||
@@ -65,6 +81,17 @@ type Provider struct {
|
||||
// surface — the proxy middleware then falls back to URL sniffing
|
||||
// or skips request-side enrichment.
|
||||
ParserID string
|
||||
// PricingSurfaces names the cost-meter pricing surfaces this
|
||||
// provider's Models are priced under ("openai", "anthropic",
|
||||
// "bedrock" — the llm.Parser surface the request parser stamps as
|
||||
// llm.provider at billing time). NOT derivable from ParserID:
|
||||
// bedrock_api and vertex_ai_api leave ParserID empty (URL-sniffed)
|
||||
// yet price under "bedrock" / "anthropic", and kimi_api serves two
|
||||
// body shapes so it prices under both. Nil for gateway/custom
|
||||
// entries, which declare no models. Same (surface, model) pair
|
||||
// contributed by two providers must carry identical rates — the
|
||||
// pricing package's tests enforce that.
|
||||
PricingSurfaces []string
|
||||
// IdentityInjection, when non-nil, instructs the proxy to stamp
|
||||
// the caller's NetBird identity onto upstream requests under the
|
||||
// configured header names. Used for gateways like LiteLLM that
|
||||
@@ -197,6 +224,12 @@ type JSONMetadataInjection struct {
|
||||
// enforces a 128-char limit per value; oversized values are
|
||||
// truncated rather than failing the request. 0 disables the cap.
|
||||
MaxValueLength int
|
||||
// Sanitize, when true, replaces characters outside the destination's
|
||||
// accepted set with '_' before emitting each value. AWS Bedrock's
|
||||
// X-Amzn-Bedrock-Request-Metadata restricts values to a limited character
|
||||
// class, so unsanitized group display names (e.g. containing spaces) would
|
||||
// make Bedrock reject the request with 400.
|
||||
Sanitize bool
|
||||
}
|
||||
|
||||
// providers is the canonical list of supported Agent Network providers.
|
||||
@@ -213,6 +246,7 @@ var providers = []Provider{
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#10A37F",
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
// Pricing + context windows cross-checked against LiteLLM's
|
||||
// model_prices_and_context_window.json. Notable corrections from
|
||||
// earlier values: o4-mini repriced from $4/$16 to $1.10/$4.40
|
||||
@@ -220,20 +254,20 @@ var providers = []Provider{
|
||||
// family context windows split between 1.05M for full-size
|
||||
// models and 272K for mini/nano/codex variants.
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.5", Label: "GPT-5.5", InputPer1k: 0.005, OutputPer1k: 0.030, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.5-pro", Label: "GPT-5.5 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4", Label: "GPT-5.4", InputPer1k: 0.0025, OutputPer1k: 0.015, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-pro", Label: "GPT-5.4 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini", InputPer1k: 0.00075, OutputPer1k: 0.0045, ContextWindow: 272000},
|
||||
{ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano", InputPer1k: 0.0002, OutputPer1k: 0.00125, ContextWindow: 272000},
|
||||
{ID: "gpt-5.3-codex", Label: "GPT-5.3 Codex", InputPer1k: 0.00175, OutputPer1k: 0.014, ContextWindow: 272000},
|
||||
{ID: "gpt-5.3-chat-latest", Label: "GPT-5.3 Chat", InputPer1k: 0.00175, OutputPer1k: 0.014, ContextWindow: 128000},
|
||||
{ID: "o4-mini", Label: "o4-mini", InputPer1k: 0.0011, OutputPer1k: 0.0044, ContextWindow: 200000},
|
||||
{ID: "gpt-4.1", Label: "GPT-4.1", InputPer1k: 0.002, OutputPer1k: 0.008, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-mini", Label: "GPT-4.1 mini", InputPer1k: 0.0004, OutputPer1k: 0.0016, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-nano", Label: "GPT-4.1 nano", InputPer1k: 0.0001, OutputPer1k: 0.0004, ContextWindow: 1047576},
|
||||
{ID: "gpt-4o", Label: "GPT-4o", InputPer1k: 0.0025, OutputPer1k: 0.010, ContextWindow: 128000},
|
||||
{ID: "gpt-4o-mini", Label: "GPT-4o mini", InputPer1k: 0.00015, OutputPer1k: 0.0006, ContextWindow: 128000},
|
||||
{ID: "gpt-5.5", Label: "GPT-5.5", InputPer1k: 0.005, OutputPer1k: 0.030, CachedInputPer1k: 0.0005, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.5-pro", Label: "GPT-5.5 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, CachedInputPer1k: 0.003, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4", Label: "GPT-5.4", InputPer1k: 0.0025, OutputPer1k: 0.015, CachedInputPer1k: 0.00025, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-pro", Label: "GPT-5.4 Pro", InputPer1k: 0.030, OutputPer1k: 0.180, CachedInputPer1k: 0.003, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini", InputPer1k: 0.00075, OutputPer1k: 0.0045, CachedInputPer1k: 0.000075, ContextWindow: 272000},
|
||||
{ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano", InputPer1k: 0.0002, OutputPer1k: 0.00125, CachedInputPer1k: 0.00002, ContextWindow: 272000},
|
||||
{ID: "gpt-5.3-codex", Label: "GPT-5.3 Codex", InputPer1k: 0.00175, OutputPer1k: 0.014, CachedInputPer1k: 0.000175, ContextWindow: 272000},
|
||||
{ID: "gpt-5.3-chat-latest", Label: "GPT-5.3 Chat", InputPer1k: 0.00175, OutputPer1k: 0.014, CachedInputPer1k: 0.000175, ContextWindow: 128000},
|
||||
{ID: "o4-mini", Label: "o4-mini", InputPer1k: 0.0011, OutputPer1k: 0.0044, CachedInputPer1k: 0.000275, ContextWindow: 200000},
|
||||
{ID: "gpt-4.1", Label: "GPT-4.1", InputPer1k: 0.002, OutputPer1k: 0.008, CachedInputPer1k: 0.0005, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-mini", Label: "GPT-4.1 mini", InputPer1k: 0.0004, OutputPer1k: 0.0016, CachedInputPer1k: 0.0001, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-nano", Label: "GPT-4.1 nano", InputPer1k: 0.0001, OutputPer1k: 0.0004, CachedInputPer1k: 0.000025, ContextWindow: 1047576},
|
||||
{ID: "gpt-4o", Label: "GPT-4o", InputPer1k: 0.0025, OutputPer1k: 0.010, CachedInputPer1k: 0.00125, ContextWindow: 128000},
|
||||
{ID: "gpt-4o-mini", Label: "GPT-4o mini", InputPer1k: 0.00015, OutputPer1k: 0.0006, CachedInputPer1k: 0.000075, ContextWindow: 128000},
|
||||
{ID: "gpt-4-turbo", Label: "GPT-4 Turbo", InputPer1k: 0.01, OutputPer1k: 0.03, ContextWindow: 128000},
|
||||
{ID: "gpt-3.5-turbo", Label: "GPT-3.5 Turbo", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 16385},
|
||||
{ID: "text-embedding-3-large", Label: "text-embedding-3-large", InputPer1k: 0.00013, OutputPer1k: 0, ContextWindow: 8191},
|
||||
@@ -251,6 +285,7 @@ var providers = []Provider{
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#D97757",
|
||||
ParserID: "anthropic",
|
||||
PricingSurfaces: []string{"anthropic"},
|
||||
// Per Anthropic's current model lineup. Pricing in USD per 1k
|
||||
// tokens. Context windows: 4.6+ family is 1M; Haiku 4.5 stays at
|
||||
// 200K. claude-3-7-sonnet and claude-3-5-haiku retired
|
||||
@@ -261,14 +296,14 @@ var providers = []Provider{
|
||||
// account to be on >= 30-day data retention or all requests
|
||||
// 400.
|
||||
Models: []Model{
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5", InputPer1k: 0.010, OutputPer1k: 0.050, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-6", Label: "Claude Opus 4.6", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (deprecated, retires 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000},
|
||||
{ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000},
|
||||
{ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000},
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-6", Label: "Claude Opus 4.6", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (deprecated, retires 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, CacheReadPer1k: 0.0015, CacheCreationPer1k: 0.01875, ContextWindow: 200000},
|
||||
{ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 200000},
|
||||
{ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5", InputPer1k: 0.001, OutputPer1k: 0.005, CacheReadPer1k: 0.0001, CacheCreationPer1k: 0.00125, ContextWindow: 200000},
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -282,18 +317,19 @@ var providers = []Provider{
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#0078D4",
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
// Mirrors openai_api pricing — Azure resells OpenAI models at the
|
||||
// same per-token rates, just under different deployment names.
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.5", Label: "GPT-5.5 (Azure)", InputPer1k: 0.005, OutputPer1k: 0.030, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4", Label: "GPT-5.4 (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.015, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini (Azure)", InputPer1k: 0.00075, OutputPer1k: 0.0045, ContextWindow: 272000},
|
||||
{ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano (Azure)", InputPer1k: 0.0002, OutputPer1k: 0.00125, ContextWindow: 272000},
|
||||
{ID: "o4-mini", Label: "o4-mini (Azure)", InputPer1k: 0.0011, OutputPer1k: 0.0044, ContextWindow: 200000},
|
||||
{ID: "gpt-4.1", Label: "GPT-4.1 (Azure)", InputPer1k: 0.002, OutputPer1k: 0.008, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-mini", Label: "GPT-4.1 mini (Azure)", InputPer1k: 0.0004, OutputPer1k: 0.0016, ContextWindow: 1047576},
|
||||
{ID: "gpt-4o", Label: "GPT-4o (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.010, ContextWindow: 128000},
|
||||
{ID: "gpt-4o-mini", Label: "GPT-4o mini (Azure)", InputPer1k: 0.00015, OutputPer1k: 0.0006, ContextWindow: 128000},
|
||||
{ID: "gpt-5.5", Label: "GPT-5.5 (Azure)", InputPer1k: 0.005, OutputPer1k: 0.030, CachedInputPer1k: 0.0005, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4", Label: "GPT-5.4 (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.015, CachedInputPer1k: 0.00025, ContextWindow: 1050000},
|
||||
{ID: "gpt-5.4-mini", Label: "GPT-5.4 Mini (Azure)", InputPer1k: 0.00075, OutputPer1k: 0.0045, CachedInputPer1k: 0.000075, ContextWindow: 272000},
|
||||
{ID: "gpt-5.4-nano", Label: "GPT-5.4 Nano (Azure)", InputPer1k: 0.0002, OutputPer1k: 0.00125, CachedInputPer1k: 0.00002, ContextWindow: 272000},
|
||||
{ID: "o4-mini", Label: "o4-mini (Azure)", InputPer1k: 0.0011, OutputPer1k: 0.0044, CachedInputPer1k: 0.000275, ContextWindow: 200000},
|
||||
{ID: "gpt-4.1", Label: "GPT-4.1 (Azure)", InputPer1k: 0.002, OutputPer1k: 0.008, CachedInputPer1k: 0.0005, ContextWindow: 1047576},
|
||||
{ID: "gpt-4.1-mini", Label: "GPT-4.1 mini (Azure)", InputPer1k: 0.0004, OutputPer1k: 0.0016, CachedInputPer1k: 0.0001, ContextWindow: 1047576},
|
||||
{ID: "gpt-4o", Label: "GPT-4o (Azure)", InputPer1k: 0.0025, OutputPer1k: 0.010, CachedInputPer1k: 0.00125, ContextWindow: 128000},
|
||||
{ID: "gpt-4o-mini", Label: "GPT-4o mini (Azure)", InputPer1k: 0.00015, OutputPer1k: 0.0006, CachedInputPer1k: 0.000075, ContextWindow: 128000},
|
||||
{ID: "gpt-35-turbo", Label: "GPT-3.5 Turbo (Azure)", InputPer1k: 0.0005, OutputPer1k: 0.0015, ContextWindow: 16385},
|
||||
},
|
||||
},
|
||||
@@ -307,6 +343,9 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#FF9900",
|
||||
// ParserID stays empty (path-style dispatch via IsBedrockPathStyle);
|
||||
// the request parser meters these under the "bedrock" surface.
|
||||
PricingSurfaces: []string{"bedrock"},
|
||||
// Anthropic models on Bedrock take the anthropic.* prefix and
|
||||
// follow the same lineup / pricing as the first-party Anthropic
|
||||
// catalog entry above. claude-3-7-sonnet and claude-3-5-haiku
|
||||
@@ -316,19 +355,31 @@ var providers = []Provider{
|
||||
// Llama 3.3 70B entry kept unchanged — LiteLLM tracks only
|
||||
// per-region Llama 3 entries; standalone 3.3 not yet listed.
|
||||
Models: []Model{
|
||||
{ID: "anthropic.claude-opus-4-8", Label: "Claude Opus 4.8 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-7", Label: "Claude Opus 4.7 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-6", Label: "Claude Opus 4.6 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-1", Label: "Claude Opus 4.1 (Bedrock, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000},
|
||||
{ID: "anthropic.claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000},
|
||||
{ID: "anthropic.claude-haiku-4-5", Label: "Claude Haiku 4.5 (Bedrock)", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000},
|
||||
{ID: "anthropic.claude-opus-4-8", Label: "Claude Opus 4.8 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-7", Label: "Claude Opus 4.7 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-6", Label: "Claude Opus 4.6 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-opus-4-1", Label: "Claude Opus 4.1 (Bedrock, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, CacheReadPer1k: 0.0015, CacheCreationPer1k: 0.01875, ContextWindow: 200000},
|
||||
{ID: "anthropic.claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 200000},
|
||||
{ID: "anthropic.claude-haiku-4-5", Label: "Claude Haiku 4.5 (Bedrock)", InputPer1k: 0.001, OutputPer1k: 0.005, CacheReadPer1k: 0.0001, CacheCreationPer1k: 0.00125, ContextWindow: 200000},
|
||||
{ID: "meta.llama3-3-70b-instruct", Label: "Llama 3.3 70B (Bedrock)", InputPer1k: 0.00072, OutputPer1k: 0.00072, ContextWindow: 128000},
|
||||
{ID: "amazon.nova-2-lite", Label: "Amazon Nova 2 Lite (Bedrock, preview)", InputPer1k: 0.0003, OutputPer1k: 0.0025, ContextWindow: 1000000},
|
||||
{ID: "amazon.nova-pro", Label: "Amazon Nova Pro (Bedrock)", InputPer1k: 0.0008, OutputPer1k: 0.0032, ContextWindow: 300000},
|
||||
{ID: "amazon.nova-lite", Label: "Amazon Nova Lite (Bedrock)", InputPer1k: 0.00006, OutputPer1k: 0.00024, ContextWindow: 300000},
|
||||
{ID: "amazon.nova-micro", Label: "Amazon Nova Micro (Bedrock)", InputPer1k: 0.000035, OutputPer1k: 0.00014, ContextWindow: 128000},
|
||||
},
|
||||
// Bedrock accepts a cost-allocation metadata header; stamp the caller's
|
||||
// user + authorizing group so spend can be attributed in AWS Cost
|
||||
// Management. Sanitized because Bedrock restricts the value character set.
|
||||
IdentityInjection: &IdentityInjection{
|
||||
JSONMetadata: &JSONMetadataInjection{
|
||||
Header: "X-Amzn-Bedrock-Request-Metadata",
|
||||
UserKey: "user",
|
||||
GroupsKey: "group",
|
||||
MaxValueLength: 256,
|
||||
Sanitize: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "vertex_ai_api",
|
||||
@@ -340,6 +391,10 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#4285F4",
|
||||
// ParserID stays empty (path-style dispatch via IsVertexPathStyle);
|
||||
// Anthropic-on-Vertex requests are metered under the "anthropic"
|
||||
// surface with the bare, unversioned model id.
|
||||
PricingSurfaces: []string{"anthropic"},
|
||||
// Vertex carries the model in the URL path and authenticates with a
|
||||
// service-account-minted OAuth token (api_key = "keyfile::<base64 SA>").
|
||||
// Only Anthropic-on-Vertex is metered today: the request parser maps the
|
||||
@@ -351,14 +406,14 @@ var providers = []Provider{
|
||||
// exists — the router denies unmeterable publishers rather than forward
|
||||
// them uncounted.
|
||||
Models: []Model{
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5 (Vertex)", InputPer1k: 0.010, OutputPer1k: 0.050, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-6", Label: "Claude Opus 4.6 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (Vertex, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, ContextWindow: 200000},
|
||||
{ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 200000},
|
||||
{ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5 (Vertex)", InputPer1k: 0.001, OutputPer1k: 0.005, ContextWindow: 200000},
|
||||
{ID: "claude-fable-5", Label: "Claude Fable 5 (Vertex)", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-8", Label: "Claude Opus 4.8 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-7", Label: "Claude Opus 4.7 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-6", Label: "Claude Opus 4.6 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-opus-4-1", Label: "Claude Opus 4.1 (Vertex, deprecated 2026-08-05)", InputPer1k: 0.015, OutputPer1k: 0.075, CacheReadPer1k: 0.0015, CacheCreationPer1k: 0.01875, ContextWindow: 200000},
|
||||
{ID: "claude-sonnet-4-6", Label: "Claude Sonnet 4.6 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-4-5", Label: "Claude Sonnet 4.5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 200000},
|
||||
{ID: "claude-haiku-4-5", Label: "Claude Haiku 4.5 (Vertex)", InputPer1k: 0.001, OutputPer1k: 0.005, CacheReadPer1k: 0.0001, CacheCreationPer1k: 0.00125, ContextWindow: 200000},
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -372,6 +427,7 @@ var providers = []Provider{
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#FF7000",
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
// Pricing + context windows cross-checked against LiteLLM. Key
|
||||
// gotchas the marketing page hides:
|
||||
// - `mistral-medium-latest` aliases to Medium 3.1 ($0.40/$2),
|
||||
@@ -402,6 +458,58 @@ 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: "",
|
||||
// Both body shapes are metered: /v1/chat/completions under
|
||||
// "openai", /anthropic/v1/messages under "anthropic" — so the
|
||||
// K3 entry is priced on both surfaces.
|
||||
PricingSurfaces: []string{"openai", "anthropic"},
|
||||
// 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{
|
||||
// Carries both cache shapes: Moonshot reports cache hits
|
||||
// OpenAI-style on /v1/chat/completions (CachedInputPer1k)
|
||||
// and Anthropic-style on /anthropic/v1/messages
|
||||
// (CacheReadPer1k) — $0.30/MTok either way. Each surface's
|
||||
// cost formula reads only its own field, so the superset
|
||||
// entry prices both endpoints correctly. No cache-creation
|
||||
// rate published; writes bill at the input rate.
|
||||
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, CachedInputPer1k: 0.0003, CacheReadPer1k: 0.0003, ContextWindow: 1000000},
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "litellm_proxy",
|
||||
Kind: KindGateway,
|
||||
@@ -699,13 +807,28 @@ func IsBedrockPathStyle(providerID string) bool {
|
||||
func (p Provider) ToAPIResponse() api.AgentNetworkCatalogProvider {
|
||||
models := make([]api.AgentNetworkCatalogModel, 0, len(p.Models))
|
||||
for _, m := range p.Models {
|
||||
models = append(models, api.AgentNetworkCatalogModel{
|
||||
am := api.AgentNetworkCatalogModel{
|
||||
Id: m.ID,
|
||||
Label: m.Label,
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
ContextWindow: m.ContextWindow,
|
||||
})
|
||||
}
|
||||
// Cache rates are emitted only when configured so the dashboard
|
||||
// can prefill them; 0 stays off the wire (absent = no rate).
|
||||
if m.CachedInputPer1k > 0 {
|
||||
v := m.CachedInputPer1k
|
||||
am.CachedInputPer1k = &v
|
||||
}
|
||||
if m.CacheReadPer1k > 0 {
|
||||
v := m.CacheReadPer1k
|
||||
am.CacheReadPer1k = &v
|
||||
}
|
||||
if m.CacheCreationPer1k > 0 {
|
||||
v := m.CacheCreationPer1k
|
||||
am.CacheCreationPer1k = &v
|
||||
}
|
||||
models = append(models, am)
|
||||
}
|
||||
kind := api.AgentNetworkCatalogProviderKindProvider
|
||||
switch p.Kind {
|
||||
@@ -725,6 +848,10 @@ func (p Provider) ToAPIResponse() api.AgentNetworkCatalogProvider {
|
||||
BrandColor: p.BrandColor,
|
||||
Models: models,
|
||||
}
|
||||
if len(p.PricingSurfaces) > 0 {
|
||||
surfaces := append([]string(nil), p.PricingSurfaces...)
|
||||
resp.PricingSurfaces = &surfaces
|
||||
}
|
||||
if len(p.ExtraHeaders) > 0 {
|
||||
extras := make([]api.AgentNetworkCatalogExtraHeader, 0, len(p.ExtraHeaders))
|
||||
for _, h := range p.ExtraHeaders {
|
||||
|
||||
@@ -7,6 +7,7 @@ package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
@@ -15,6 +16,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
nbcontext "github.com/netbirdio/netbird/management/server/context"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
@@ -52,11 +54,45 @@ func (h *handler) getCatalogProviders(w http.ResponseWriter, r *http.Request) {
|
||||
entries := catalog.All()
|
||||
out := make([]api.AgentNetworkCatalogProvider, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
out = append(out, e.ToAPIResponse())
|
||||
resp := e.ToAPIResponse()
|
||||
applyDefaultPricing(e, &resp)
|
||||
out = append(out, resp)
|
||||
}
|
||||
util.WriteJSONObject(r.Context(), w, out)
|
||||
}
|
||||
|
||||
// applyDefaultPricing overwrites the catalog response's model rates with
|
||||
// the LIVE default pricing table, which may differ from the compiled-in
|
||||
// catalog rates when the operator provides a defaults_llm_pricing.yaml.
|
||||
// This keeps the dashboard's model-row prefill identical to what the
|
||||
// proxy will actually bill — the same table the synthesizer ships.
|
||||
func applyDefaultPricing(cp catalog.Provider, resp *api.AgentNetworkCatalogProvider) {
|
||||
if len(cp.PricingSurfaces) == 0 {
|
||||
return
|
||||
}
|
||||
for i := range resp.Models {
|
||||
m := &resp.Models[i]
|
||||
e, ok := pricing.LookupDefault(cp.PricingSurfaces, m.Id)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
m.InputPer1k = e.InputPer1k
|
||||
m.OutputPer1k = e.OutputPer1k
|
||||
m.CachedInputPer1k = positiveRatePtr(e.CachedInputPer1k)
|
||||
m.CacheReadPer1k = positiveRatePtr(e.CacheReadPer1k)
|
||||
m.CacheCreationPer1k = positiveRatePtr(e.CacheCreationPer1k)
|
||||
}
|
||||
}
|
||||
|
||||
// positiveRatePtr renders a cache rate for the API: absent (nil) when
|
||||
// unset, matching the catalog response convention.
|
||||
func positiveRatePtr(v float64) *float64 {
|
||||
if v <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &v
|
||||
}
|
||||
|
||||
func (h *handler) getAllProviders(w http.ResponseWriter, r *http.Request) {
|
||||
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
|
||||
if err != nil {
|
||||
@@ -213,5 +249,38 @@ func validate(req *api.AgentNetworkProviderRequest, requireAPIKey bool) error {
|
||||
if requireAPIKey && (req.ApiKey == nil || strings.TrimSpace(*req.ApiKey) == "") {
|
||||
return status.Errorf(status.InvalidArgument, "api_key is required")
|
||||
}
|
||||
if req.Models != nil {
|
||||
for i, m := range *req.Models {
|
||||
if err := validateModel(i, m); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateModel is the single ingress guard for operator-entered pricing:
|
||||
// these rates are synthesized into the proxy's cost_meter config verbatim,
|
||||
// and a negative or non-finite rate there would poison every cost the
|
||||
// proxy records, so reject at the API boundary.
|
||||
func validateModel(i int, m api.AgentNetworkProviderModel) error {
|
||||
if strings.TrimSpace(m.Id) == "" {
|
||||
return status.Errorf(status.InvalidArgument, "models[%d]: id is required", i)
|
||||
}
|
||||
rates := map[string]*float64{
|
||||
"input_per_1k": &m.InputPer1k,
|
||||
"output_per_1k": &m.OutputPer1k,
|
||||
"cached_input_per_1k": m.CachedInputPer1k,
|
||||
"cache_read_per_1k": m.CacheReadPer1k,
|
||||
"cache_creation_per_1k": m.CacheCreationPer1k,
|
||||
}
|
||||
for field, v := range rates {
|
||||
if v == nil {
|
||||
continue
|
||||
}
|
||||
if *v < 0 || math.IsNaN(*v) || math.IsInf(*v, 0) {
|
||||
return status.Errorf(status.InvalidArgument, "models[%d] (%s): %s must be a finite, non-negative USD rate", i, m.Id, field)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
func f(v float64) *float64 { return &v }
|
||||
|
||||
// TestValidate_ModelRates guards the single ingress point for operator-entered
|
||||
// pricing. These rates flow verbatim into the proxy's cost_meter config at
|
||||
// synthesis time; the proxy treats a bad rate as a chain-build failure, so
|
||||
// rejecting here is what keeps an account's gateway from going down.
|
||||
func TestValidate_ModelRates(t *testing.T) {
|
||||
base := func(models ...api.AgentNetworkProviderModel) *api.AgentNetworkProviderRequest {
|
||||
key := "sk-test"
|
||||
return &api.AgentNetworkProviderRequest{
|
||||
ProviderId: "openai_api",
|
||||
Name: "OpenAI",
|
||||
UpstreamUrl: "https://api.openai.com",
|
||||
ApiKey: &key,
|
||||
Models: &models,
|
||||
}
|
||||
}
|
||||
|
||||
valid := api.AgentNetworkProviderModel{
|
||||
Id: "gpt-4o", InputPer1k: 0.0025, OutputPer1k: 0.01,
|
||||
CachedInputPer1k: f(0.00125),
|
||||
}
|
||||
require.NoError(t, validate(base(valid), true), "finite non-negative rates must pass")
|
||||
|
||||
zeroRates := api.AgentNetworkProviderModel{Id: "self-hosted-llama", InputPer1k: 0, OutputPer1k: 0}
|
||||
require.NoError(t, validate(base(zeroRates), true), "explicit zero prices are allowed (free / self-hosted models)")
|
||||
|
||||
cases := map[string]api.AgentNetworkProviderModel{
|
||||
"empty id": {Id: " ", InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
"negative input": {Id: "m", InputPer1k: -0.001, OutputPer1k: 0.002},
|
||||
"negative output": {Id: "m", InputPer1k: 0.001, OutputPer1k: -0.002},
|
||||
"NaN input": {Id: "m", InputPer1k: math.NaN(), OutputPer1k: 0.002},
|
||||
"Inf output": {Id: "m", InputPer1k: 0.001, OutputPer1k: math.Inf(1)},
|
||||
"negative cached": {Id: "m", InputPer1k: 0.001, OutputPer1k: 0.002, CachedInputPer1k: f(-1)},
|
||||
"NaN cache read": {Id: "m", InputPer1k: 0.001, OutputPer1k: 0.002, CacheReadPer1k: f(math.NaN())},
|
||||
"Inf cache creation": {Id: "m", InputPer1k: 0.001, OutputPer1k: 0.002, CacheCreationPer1k: f(math.Inf(-1))},
|
||||
}
|
||||
for name, m := range cases {
|
||||
assert.Error(t, validate(base(m), true), "case %q must be rejected", name)
|
||||
}
|
||||
}
|
||||
@@ -86,12 +86,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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
// Package pricing builds the default LLM pricing table the synthesizer
|
||||
// ships to the proxy's cost_meter middleware. The catalog is the single
|
||||
// source of default rates: every catalog provider's models are folded
|
||||
// into the pricing surfaces the provider declares (PricingSurfaces),
|
||||
// then a small supplemental list adds priced-but-not-operator-selectable
|
||||
// entries. Management is the sole pricing authority — the proxy carries
|
||||
// no embedded price list and bills exclusively from the table it is
|
||||
// sent.
|
||||
//go:generate go run gen.go
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
// Entry is a single model's pricing in USD per 1k tokens. This struct IS
|
||||
// the wire shape: the synthesizer marshals it verbatim into cost_meter's
|
||||
// ConfigJSON, and the proxy unmarshals the same field names.
|
||||
//
|
||||
// A zero rate means "no rate configured" — the proxy bills that cache
|
||||
// bucket at InputPer1k (identical semantics to the retired proxy-embedded
|
||||
// table). CachedInputPer1k is the OpenAI shape (cached prompt tokens are
|
||||
// a subset of input); CacheReadPer1k / CacheCreationPer1k are the
|
||||
// Anthropic shape (additive buckets).
|
||||
type Entry struct {
|
||||
InputPer1k float64 `json:"input_per_1k"`
|
||||
OutputPer1k float64 `json:"output_per_1k"`
|
||||
CachedInputPer1k float64 `json:"cached_input_per_1k,omitempty"`
|
||||
CacheReadPer1k float64 `json:"cache_read_per_1k,omitempty"`
|
||||
CacheCreationPer1k float64 `json:"cache_creation_per_1k,omitempty"`
|
||||
}
|
||||
|
||||
// supplementalDefaults are (surface, model) entries that are priced but
|
||||
// deliberately not operator-selectable in the catalog. Each carries a
|
||||
// reason; when one of these models joins a catalog lineup, delete the
|
||||
// row here — the collision test fails loudly if the rates ever disagree.
|
||||
var supplementalDefaults = map[string]map[string]Entry{
|
||||
"openai": {
|
||||
// GPT-5 (2025) family — kept for gateway requests using the
|
||||
// unsuffixed ids; the dashboard offers only the 5.x lineup.
|
||||
"gpt-5": {InputPer1k: 0.00125, OutputPer1k: 0.01, CachedInputPer1k: 0.000125},
|
||||
"gpt-5-mini": {InputPer1k: 0.00025, OutputPer1k: 0.002, CachedInputPer1k: 0.000025},
|
||||
"gpt-5-nano": {InputPer1k: 0.00005, OutputPer1k: 0.0004, CachedInputPer1k: 0.000005},
|
||||
},
|
||||
"anthropic": {
|
||||
// claude-opus-5 is not yet in the catalog lineup but gateway /
|
||||
// grandfathered traffic uses it; priced so it isn't skipped.
|
||||
"claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625},
|
||||
// "kimi-k3[1m]" is the 1M-context alias some Claude Code guides
|
||||
// configure against Moonshot's Anthropic-compatible endpoint;
|
||||
// priced identically to kimi-k3 so those requests aren't skipped.
|
||||
"kimi-k3[1m]": {InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003},
|
||||
},
|
||||
"bedrock": {
|
||||
"anthropic.claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625},
|
||||
},
|
||||
}
|
||||
|
||||
var (
|
||||
compiledOnce sync.Once
|
||||
compiledTable map[string]map[string]Entry
|
||||
// mergedTable holds the current live table when a pricing defaults
|
||||
// file is loaded: the file merged entry-whole over the compiled-in
|
||||
// base. Nil while no file is loaded (or after the file is removed),
|
||||
// in which case the compiled-in table serves. Swapped atomically by
|
||||
// the file loader/reloader; readers never block.
|
||||
mergedTable atomic.Pointer[map[string]map[string]Entry]
|
||||
)
|
||||
|
||||
// DefaultTable returns the current default pricing table keyed
|
||||
// surface -> model -> Entry: the management-side defaults file (see
|
||||
// LoadFile / StartReloader) when one is loaded, merged over the
|
||||
// compiled-in catalog table, which alone serves as the fallback when no
|
||||
// file exists. The snapshot may change between calls as the file is
|
||||
// re-read — consumers (the synthesizer on every reconcile, the catalog
|
||||
// endpoint on every request) pick up fresh rates automatically. Callers
|
||||
// must not mutate the returned maps.
|
||||
func DefaultTable() map[string]map[string]Entry {
|
||||
if t := mergedTable.Load(); t != nil {
|
||||
return *t
|
||||
}
|
||||
return compiledBase()
|
||||
}
|
||||
|
||||
// compiledBase returns the compiled-in table (catalog + supplementals),
|
||||
// built once.
|
||||
func compiledBase() map[string]map[string]Entry {
|
||||
compiledOnce.Do(func() {
|
||||
compiledTable = buildDefaultTable()
|
||||
})
|
||||
return compiledTable
|
||||
}
|
||||
|
||||
func buildDefaultTable() map[string]map[string]Entry {
|
||||
out := make(map[string]map[string]Entry)
|
||||
for _, p := range catalog.All() {
|
||||
for _, surface := range p.PricingSurfaces {
|
||||
inner, ok := out[surface]
|
||||
if !ok {
|
||||
inner = make(map[string]Entry)
|
||||
out[surface] = inner
|
||||
}
|
||||
for _, m := range p.Models {
|
||||
// First writer wins; providers contributing the same
|
||||
// (surface, model) must agree on rates — enforced by
|
||||
// TestDefaultTable_NoConflictingContributions.
|
||||
if _, dup := inner[m.ID]; dup {
|
||||
continue
|
||||
}
|
||||
inner[m.ID] = entryFromCatalogModel(m)
|
||||
}
|
||||
}
|
||||
}
|
||||
for surface, models := range supplementalDefaults {
|
||||
inner, ok := out[surface]
|
||||
if !ok {
|
||||
inner = make(map[string]Entry)
|
||||
out[surface] = inner
|
||||
}
|
||||
for id, e := range models {
|
||||
if _, dup := inner[id]; dup {
|
||||
continue
|
||||
}
|
||||
inner[id] = e
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func entryFromCatalogModel(m catalog.Model) Entry {
|
||||
return Entry{
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
CachedInputPer1k: m.CachedInputPer1k,
|
||||
CacheReadPer1k: m.CacheReadPer1k,
|
||||
CacheCreationPer1k: m.CacheCreationPer1k,
|
||||
}
|
||||
}
|
||||
|
||||
// LookupDefault returns the default entry for model on the first of the
|
||||
// given surfaces that prices it. Used by the synthesizer to seed a
|
||||
// per-provider entry with default cache rates before overlaying the
|
||||
// operator's stored prices.
|
||||
func LookupDefault(surfaces []string, model string) (Entry, bool) {
|
||||
table := DefaultTable()
|
||||
for _, s := range surfaces {
|
||||
if e, ok := table[s][model]; ok {
|
||||
return e, true
|
||||
}
|
||||
}
|
||||
return Entry{}, false
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
# Default LLM pricing used by NetBird's Agent Network cost metering.
|
||||
# GENERATED from the management catalog — do not edit this file in the
|
||||
# repository; regenerate with:
|
||||
#
|
||||
# go generate ./management/internals/modules/agentnetwork/pricing
|
||||
#
|
||||
# Operators: copy this file to <datadir>/defaults_llm_pricing.yaml (or
|
||||
# any path configured via management.json:
|
||||
#
|
||||
# { "AgentNetwork": { "PricingDefaultsFile": "/path/defaults_llm_pricing.yaml" } }
|
||||
#
|
||||
# ) and adjust the entries you want to change. Management re-reads the
|
||||
# file periodically (mtime poll, every minute): the live table feeds the
|
||||
# proxies' cost metering and the dashboard's model-price prefill, so
|
||||
# edits apply without a restart. Your file only needs the entries you
|
||||
# want to change — but each entry REPLACES the built-in entry for that
|
||||
# surface+model whole, so repeat the cache rates you want to keep.
|
||||
# Unknown fields and negative or non-finite rates are rejected: at
|
||||
# startup that fails boot (for an explicitly configured path); at
|
||||
# runtime the previous table is kept and a warning is logged. Deleting
|
||||
# the file reverts to the built-in defaults below.
|
||||
#
|
||||
# Top-level keys are pricing surfaces — the parser shape requests are
|
||||
# metered under: "openai" (also Azure, Mistral, and OpenAI-compatible
|
||||
# gateways), "anthropic" (also Anthropic-on-Vertex), "bedrock"
|
||||
# (normalized ids, e.g. anthropic.claude-sonnet-4-5). Model keys must be
|
||||
# the normalized id the proxy meters (version/region suffixes stripped).
|
||||
#
|
||||
# Values are USD per 1_000 tokens. Optional cache fields:
|
||||
# cached_input_per_1k OpenAI shape: rate for cached prompt tokens
|
||||
# (a SUBSET of input tokens). Absent -> cached
|
||||
# portion bills at input_per_1k.
|
||||
# cache_read_per_1k Anthropic shape: rate for cache_read tokens
|
||||
# (ADDITIVE to input). Absent -> input rate.
|
||||
# cache_creation_per_1k Anthropic shape: rate for cache_creation
|
||||
# tokens (ADDITIVE to input). Absent -> input
|
||||
# rate.
|
||||
|
||||
anthropic:
|
||||
claude-fable-5:
|
||||
input_per_1k: 0.01
|
||||
output_per_1k: 0.05
|
||||
cache_read_per_1k: 0.001
|
||||
cache_creation_per_1k: 0.0125
|
||||
claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
kimi-k3:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.0003
|
||||
cache_read_per_1k: 0.0003
|
||||
"kimi-k3[1m]":
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
|
||||
bedrock:
|
||||
amazon.nova-2-lite:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0025
|
||||
amazon.nova-lite:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00024
|
||||
amazon.nova-micro:
|
||||
input_per_1k: 0.000035
|
||||
output_per_1k: 0.00014
|
||||
amazon.nova-pro:
|
||||
input_per_1k: 0.0008
|
||||
output_per_1k: 0.0032
|
||||
anthropic.claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
anthropic.claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
anthropic.claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
meta.llama3-3-70b-instruct:
|
||||
input_per_1k: 0.00072
|
||||
output_per_1k: 0.00072
|
||||
|
||||
openai:
|
||||
codestral-2508:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0009
|
||||
codestral-latest:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.003
|
||||
devstral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
devstral-small-latest:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0003
|
||||
gpt-3.5-turbo:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
gpt-35-turbo:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
gpt-4-turbo:
|
||||
input_per_1k: 0.01
|
||||
output_per_1k: 0.03
|
||||
gpt-4.1:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.008
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-4.1-mini:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.0016
|
||||
cached_input_per_1k: 0.0001
|
||||
gpt-4.1-nano:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000025
|
||||
gpt-4o:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.00125
|
||||
gpt-4o-mini:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.0006
|
||||
cached_input_per_1k: 0.000075
|
||||
gpt-5:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.000125
|
||||
gpt-5-mini:
|
||||
input_per_1k: 0.00025
|
||||
output_per_1k: 0.002
|
||||
cached_input_per_1k: 0.000025
|
||||
gpt-5-nano:
|
||||
input_per_1k: 0.00005
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000005
|
||||
gpt-5.3-chat-latest:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
gpt-5.3-codex:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
gpt-5.4:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.00025
|
||||
gpt-5.4-mini:
|
||||
input_per_1k: 0.00075
|
||||
output_per_1k: 0.0045
|
||||
cached_input_per_1k: 0.000075
|
||||
gpt-5.4-nano:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.00125
|
||||
cached_input_per_1k: 0.00002
|
||||
gpt-5.4-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
gpt-5.5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.03
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-5.5-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
kimi-k3:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.0003
|
||||
cache_read_per_1k: 0.0003
|
||||
magistral-medium-latest:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.005
|
||||
magistral-small-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
ministral-3-14b-2512:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.0002
|
||||
ministral-3-3b-2512:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0001
|
||||
ministral-8b-latest:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.00015
|
||||
mistral-embed:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0
|
||||
mistral-large-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
mistral-medium-3-5:
|
||||
input_per_1k: 0.0015
|
||||
output_per_1k: 0.0075
|
||||
mistral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
mistral-small-latest:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00018
|
||||
o4-mini:
|
||||
input_per_1k: 0.0011
|
||||
output_per_1k: 0.0044
|
||||
cached_input_per_1k: 0.000275
|
||||
text-embedding-3-large:
|
||||
input_per_1k: 0.00013
|
||||
output_per_1k: 0
|
||||
text-embedding-3-small:
|
||||
input_per_1k: 0.00002
|
||||
output_per_1k: 0
|
||||
@@ -0,0 +1,148 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
// TestDefaultTable_CoversEveryCatalogModel replaces the proxy's old
|
||||
// hand-maintained coverage list: because the table is built FROM the
|
||||
// catalog, drift is impossible by construction — this test guards the
|
||||
// fold itself (every catalog model of every surfaced provider resolves,
|
||||
// with exactly the catalog's rates).
|
||||
func TestDefaultTable_CoversEveryCatalogModel(t *testing.T) {
|
||||
table := DefaultTable()
|
||||
for _, p := range catalog.All() {
|
||||
if len(p.PricingSurfaces) == 0 {
|
||||
assert.Empty(t, p.Models, "catalog entry %s declares models but no pricing surfaces — those models would never be priced", p.ID)
|
||||
continue
|
||||
}
|
||||
for _, surface := range p.PricingSurfaces {
|
||||
byModel, ok := table[surface]
|
||||
require.True(t, ok, "surface %q (provider %s) missing from default table", surface, p.ID)
|
||||
for _, m := range p.Models {
|
||||
e, ok := byModel[m.ID]
|
||||
require.True(t, ok, "%s/%s (provider %s) missing from default table", surface, m.ID, p.ID)
|
||||
assert.Equal(t, m.InputPer1k, e.InputPer1k, "%s/%s input rate", surface, m.ID)
|
||||
assert.Equal(t, m.OutputPer1k, e.OutputPer1k, "%s/%s output rate", surface, m.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultTable_NoConflictingContributions enforces the collision rule
|
||||
// documented on catalog.Provider.PricingSurfaces: when two catalog
|
||||
// providers contribute the same (surface, model) pair — azure/vertex
|
||||
// mirroring openai/anthropic, kimi on both surfaces — their rates must be
|
||||
// identical, because the surface-keyed table can only hold one entry.
|
||||
// If a provider ever diverges (e.g. Azure reprices a model), this fails
|
||||
// and the divergence must move to per-provider-record pricing.
|
||||
func TestDefaultTable_NoConflictingContributions(t *testing.T) {
|
||||
type contribution struct {
|
||||
providerID string
|
||||
entry Entry
|
||||
}
|
||||
seen := map[string]map[string]contribution{}
|
||||
for _, p := range catalog.All() {
|
||||
for _, surface := range p.PricingSurfaces {
|
||||
if seen[surface] == nil {
|
||||
seen[surface] = map[string]contribution{}
|
||||
}
|
||||
for _, m := range p.Models {
|
||||
e := entryFromCatalogModel(m)
|
||||
if prev, dup := seen[surface][m.ID]; dup {
|
||||
assert.Equal(t, prev.entry, e,
|
||||
"%s/%s: %s and %s contribute different rates", surface, m.ID, prev.providerID, p.ID)
|
||||
continue
|
||||
}
|
||||
seen[surface][m.ID] = contribution{providerID: p.ID, entry: e}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Supplemental entries must never shadow a catalog-contributed model —
|
||||
// they exist precisely because the catalog does NOT list them.
|
||||
for surface, models := range supplementalDefaults {
|
||||
for id := range models {
|
||||
_, fromCatalog := seen[surface][id]
|
||||
assert.False(t, fromCatalog, "supplemental %s/%s is now in the catalog — delete the supplemental row", surface, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultTable_AllRatesFiniteNonNegative mirrors the proxy-side
|
||||
// NewTable validation so a bad catalog edit is caught here, at unit-test
|
||||
// time, rather than as a chain-build failure on every proxy.
|
||||
func TestDefaultTable_AllRatesFiniteNonNegative(t *testing.T) {
|
||||
for surface, models := range DefaultTable() {
|
||||
for id, e := range models {
|
||||
for field, v := range map[string]float64{
|
||||
"input": e.InputPer1k,
|
||||
"output": e.OutputPer1k,
|
||||
"cached_input": e.CachedInputPer1k,
|
||||
"cache_read": e.CacheReadPer1k,
|
||||
"cache_creation": e.CacheCreationPer1k,
|
||||
} {
|
||||
assert.False(t, v < 0 || math.IsNaN(v) || math.IsInf(v, 0),
|
||||
"%s/%s: %s rate %v must be finite and non-negative", surface, id, field, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultTable_PinnedRates pins rates that previously drifted or are
|
||||
// easy to mis-enter (carried over from the proxy's retired
|
||||
// defaults_coverage_test), plus the supplemental entries.
|
||||
func TestDefaultTable_PinnedRates(t *testing.T) {
|
||||
table := DefaultTable()
|
||||
|
||||
gpt54 := table["openai"]["gpt-5.4"]
|
||||
assert.InDelta(t, 0.0025, gpt54.InputPer1k, 1e-9, "gpt-5.4 input")
|
||||
assert.InDelta(t, 0.015, gpt54.OutputPer1k, 1e-9, "gpt-5.4 output")
|
||||
assert.InDelta(t, 0.00025, gpt54.CachedInputPer1k, 1e-9, "gpt-5.4 cached input")
|
||||
|
||||
sonnet := table["bedrock"]["anthropic.claude-sonnet-4-5"]
|
||||
assert.InDelta(t, 0.003, sonnet.InputPer1k, 1e-9, "bedrock sonnet-4-5 input")
|
||||
assert.InDelta(t, 0.015, sonnet.OutputPer1k, 1e-9, "bedrock sonnet-4-5 output")
|
||||
assert.InDelta(t, 0.0003, sonnet.CacheReadPer1k, 1e-9, "bedrock sonnet-4-5 cache read")
|
||||
assert.InDelta(t, 0.00375, sonnet.CacheCreationPer1k, 1e-9, "bedrock sonnet-4-5 cache creation")
|
||||
|
||||
// Vertex Claude prices under "anthropic" with the bare id.
|
||||
fable := table["anthropic"]["claude-fable-5"]
|
||||
assert.InDelta(t, 0.010, fable.InputPer1k, 1e-9, "claude-fable-5 input")
|
||||
assert.InDelta(t, 0.0125, fable.CacheCreationPer1k, 1e-9, "claude-fable-5 cache creation")
|
||||
|
||||
// Supplementals present on their surfaces.
|
||||
for surface, ids := range map[string][]string{
|
||||
"openai": {"gpt-5", "gpt-5-mini", "gpt-5-nano"},
|
||||
"anthropic": {"claude-opus-5", "kimi-k3[1m]", "kimi-k3"},
|
||||
"bedrock": {"anthropic.claude-opus-5"},
|
||||
} {
|
||||
for _, id := range ids {
|
||||
_, ok := table[surface][id]
|
||||
assert.True(t, ok, "%s/%s must be priced", surface, id)
|
||||
}
|
||||
}
|
||||
|
||||
// Embeddings bill input-only — output stays zero.
|
||||
emb := table["openai"]["text-embedding-3-large"]
|
||||
assert.Zero(t, emb.OutputPer1k, "embedding output rate must be zero")
|
||||
assert.Positive(t, emb.InputPer1k, "embedding input rate must be set")
|
||||
}
|
||||
|
||||
func TestLookupDefault_SurfaceOrder(t *testing.T) {
|
||||
// kimi-k3 exists on both surfaces; first surface in the slice wins.
|
||||
e, ok := LookupDefault([]string{"openai", "anthropic"}, "kimi-k3")
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.003, e.InputPer1k, 1e-9)
|
||||
|
||||
_, ok = LookupDefault([]string{"bedrock"}, "gpt-4o")
|
||||
assert.False(t, ok, "gpt-4o is not a bedrock model")
|
||||
|
||||
_, ok = LookupDefault(nil, "gpt-4o")
|
||||
assert.False(t, ok, "no surfaces, no match")
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// MarshalDefaultsYAML renders the built-in default pricing table (catalog
|
||||
// + supplementals, WITHOUT any operator override) as the YAML schema
|
||||
// LoadOverrideFile consumes. It backs the generated
|
||||
// defaults_llm_pricing.example.yaml so operators start from a file that
|
||||
// matches the compiled-in rates exactly; a golden test keeps the two in
|
||||
// sync. Output is deterministic (sorted surfaces and models).
|
||||
func MarshalDefaultsYAML() []byte {
|
||||
var b bytes.Buffer
|
||||
b.WriteString(exampleHeader)
|
||||
|
||||
table := buildDefaultTable()
|
||||
surfaces := make([]string, 0, len(table))
|
||||
for s := range table {
|
||||
surfaces = append(surfaces, s)
|
||||
}
|
||||
sort.Strings(surfaces)
|
||||
|
||||
for _, surface := range surfaces {
|
||||
fmt.Fprintf(&b, "\n%s:\n", surface)
|
||||
models := table[surface]
|
||||
ids := make([]string, 0, len(models))
|
||||
for id := range models {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
sort.Strings(ids)
|
||||
for _, id := range ids {
|
||||
e := models[id]
|
||||
fmt.Fprintf(&b, " %s:\n", yamlKey(id))
|
||||
fmt.Fprintf(&b, " input_per_1k: %s\n", rate(e.InputPer1k))
|
||||
fmt.Fprintf(&b, " output_per_1k: %s\n", rate(e.OutputPer1k))
|
||||
if e.CachedInputPer1k > 0 {
|
||||
fmt.Fprintf(&b, " cached_input_per_1k: %s\n", rate(e.CachedInputPer1k))
|
||||
}
|
||||
if e.CacheReadPer1k > 0 {
|
||||
fmt.Fprintf(&b, " cache_read_per_1k: %s\n", rate(e.CacheReadPer1k))
|
||||
}
|
||||
if e.CacheCreationPer1k > 0 {
|
||||
fmt.Fprintf(&b, " cache_creation_per_1k: %s\n", rate(e.CacheCreationPer1k))
|
||||
}
|
||||
}
|
||||
}
|
||||
return b.Bytes()
|
||||
}
|
||||
|
||||
// rate renders a USD-per-1k rate without float noise ("0.00015", not
|
||||
// "0.000150000000...").
|
||||
func rate(v float64) string {
|
||||
return strconv.FormatFloat(v, 'f', -1, 64)
|
||||
}
|
||||
|
||||
// yamlKey quotes model ids that YAML would otherwise misparse (e.g.
|
||||
// "kimi-k3[1m]" starts a flow sequence unquoted).
|
||||
func yamlKey(id string) string {
|
||||
for _, r := range id {
|
||||
switch r {
|
||||
case '[', ']', '{', '}', ':', '#', ',', '&', '*', '!', '|', '>', '\'', '"', '%', '@', '`':
|
||||
return strconv.Quote(id)
|
||||
}
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
const exampleHeader = `# Default LLM pricing used by NetBird's Agent Network cost metering.
|
||||
# GENERATED from the management catalog — do not edit this file in the
|
||||
# repository; regenerate with:
|
||||
#
|
||||
# go generate ./management/internals/modules/agentnetwork/pricing
|
||||
#
|
||||
# Operators: copy this file to <datadir>/defaults_llm_pricing.yaml (or
|
||||
# any path configured via management.json:
|
||||
#
|
||||
# { "AgentNetwork": { "PricingDefaultsFile": "/path/defaults_llm_pricing.yaml" } }
|
||||
#
|
||||
# ) and adjust the entries you want to change. Management re-reads the
|
||||
# file periodically (mtime poll, every minute): the live table feeds the
|
||||
# proxies' cost metering and the dashboard's model-price prefill, so
|
||||
# edits apply without a restart. Your file only needs the entries you
|
||||
# want to change — but each entry REPLACES the built-in entry for that
|
||||
# surface+model whole, so repeat the cache rates you want to keep.
|
||||
# Unknown fields and negative or non-finite rates are rejected: at
|
||||
# startup that fails boot (for an explicitly configured path); at
|
||||
# runtime the previous table is kept and a warning is logged. Deleting
|
||||
# the file reverts to the built-in defaults below.
|
||||
#
|
||||
# Top-level keys are pricing surfaces — the parser shape requests are
|
||||
# metered under: "openai" (also Azure, Mistral, and OpenAI-compatible
|
||||
# gateways), "anthropic" (also Anthropic-on-Vertex), "bedrock"
|
||||
# (normalized ids, e.g. anthropic.claude-sonnet-4-5). Model keys must be
|
||||
# the normalized id the proxy meters (version/region suffixes stripped).
|
||||
#
|
||||
# Values are USD per 1_000 tokens. Optional cache fields:
|
||||
# cached_input_per_1k OpenAI shape: rate for cached prompt tokens
|
||||
# (a SUBSET of input tokens). Absent -> cached
|
||||
# portion bills at input_per_1k.
|
||||
# cache_read_per_1k Anthropic shape: rate for cache_read tokens
|
||||
# (ADDITIVE to input). Absent -> input rate.
|
||||
# cache_creation_per_1k Anthropic shape: rate for cache_creation
|
||||
# tokens (ADDITIVE to input). Absent -> input
|
||||
# rate.
|
||||
`
|
||||
@@ -0,0 +1,20 @@
|
||||
//go:build ignore
|
||||
|
||||
// Regenerates defaults_llm_pricing.example.yaml from the compiled-in
|
||||
// default pricing table. Run via:
|
||||
//
|
||||
// go generate ./management/internals/modules/agentnetwork/pricing
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := os.WriteFile("defaults_llm_pricing.example.yaml", pricing.MarshalDefaultsYAML(), 0o644); err != nil {
|
||||
log.Fatalf("write defaults_llm_pricing.example.yaml: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"math"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// DefaultFileName is the basename probed under management's datadir when
|
||||
// AgentNetwork.PricingDefaultsFile doesn't configure an explicit path.
|
||||
const DefaultFileName = "defaults_llm_pricing.yaml"
|
||||
|
||||
// ReloadInterval is the cadence at which the pricing file's mtime is
|
||||
// polled for changes.
|
||||
const ReloadInterval = time.Minute
|
||||
|
||||
// maxFileBytes bounds the pricing file read so a misconfigured path
|
||||
// (pointed at a huge file) cannot exhaust process memory.
|
||||
const maxFileBytes = 4 << 20
|
||||
|
||||
// pricingFile mirrors the on-disk YAML schema — the same schema the
|
||||
// proxy's retired embedded defaults_pricing.yaml used, so files written
|
||||
// for it keep working. Keys are pricing surfaces ("openai", "anthropic",
|
||||
// "bedrock"); nested keys are normalized model ids.
|
||||
type pricingFile map[string]map[string]struct {
|
||||
InputPer1k float64 `yaml:"input_per_1k"`
|
||||
OutputPer1k float64 `yaml:"output_per_1k"`
|
||||
CachedInputPer1k float64 `yaml:"cached_input_per_1k"`
|
||||
CacheReadPer1k float64 `yaml:"cache_read_per_1k"`
|
||||
CacheCreationPer1k float64 `yaml:"cache_creation_per_1k"`
|
||||
}
|
||||
|
||||
// fileState tracks the watched pricing file across reloads.
|
||||
var fileState struct {
|
||||
mu sync.Mutex
|
||||
path string
|
||||
mtime int64
|
||||
}
|
||||
|
||||
// LoadFile loads the management-side pricing defaults file and makes it
|
||||
// the live table (merged entry-whole over the compiled-in fallback; see
|
||||
// DefaultTable). The path stays registered for the periodic reloader, so
|
||||
// later edits — or the file (re)appearing after deletion — are picked up
|
||||
// without a restart.
|
||||
//
|
||||
// required governs the missing-file case: true for an explicitly
|
||||
// configured path (a typo must fail startup rather than silently bill
|
||||
// with built-ins the operator believes they replaced), false for the
|
||||
// conventional <datadir>/defaults_llm_pricing.yaml probe (absent file =
|
||||
// compiled-in defaults, still watched in case it appears). A file that
|
||||
// exists but is malformed is always an error at load time.
|
||||
func LoadFile(path string, required bool) error {
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
fileState.mu.Lock()
|
||||
fileState.path = path
|
||||
fileState.mu.Unlock()
|
||||
|
||||
table, mtime, err := readFile(path)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) && !required {
|
||||
log.Infof("agent-network pricing defaults file %s not present; serving built-in defaults", path)
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
storeFileTable(table, mtime)
|
||||
log.Infof("agent-network pricing defaults loaded from %s", path)
|
||||
return nil
|
||||
}
|
||||
|
||||
// StartReloader launches the periodic mtime-poll goroutine for the file
|
||||
// registered by LoadFile. Runtime failures are lenient — a parse error
|
||||
// keeps the previously loaded table and logs a warning; a deleted file
|
||||
// reverts to the compiled-in defaults — so a mid-edit save can never
|
||||
// take pricing down. Returns immediately when no path was registered.
|
||||
func StartReloader(ctx context.Context, interval time.Duration) {
|
||||
fileState.mu.Lock()
|
||||
path := fileState.path
|
||||
fileState.mu.Unlock()
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
if interval <= 0 {
|
||||
interval = ReloadInterval
|
||||
}
|
||||
go func() {
|
||||
t := time.NewTicker(interval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
reload()
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// reload performs one mtime check + reload cycle.
|
||||
func reload() {
|
||||
fileState.mu.Lock()
|
||||
path, lastMtime := fileState.path, fileState.mtime
|
||||
fileState.mu.Unlock()
|
||||
|
||||
log.Debugf("agent-network pricing defaults reload: checking %s for changes", path)
|
||||
|
||||
st, err := os.Stat(path)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
// File removed (or not yet created): serve compiled-in
|
||||
// defaults and reset mtime so a future (re)appearance loads.
|
||||
if mergedTable.Swap(nil) != nil {
|
||||
log.Warnf("agent-network pricing defaults file %s removed; reverting to built-in defaults", path)
|
||||
}
|
||||
setMtime(0)
|
||||
return
|
||||
}
|
||||
log.Warnf("agent-network pricing defaults reload: stat %s: %v", path, err)
|
||||
return
|
||||
}
|
||||
if st.ModTime().UnixNano() == lastMtime {
|
||||
log.Debugf("agent-network pricing defaults %s unchanged since last check", path)
|
||||
return
|
||||
}
|
||||
|
||||
table, mtime, err := readFile(path)
|
||||
if err != nil {
|
||||
// Keep the previously loaded table — never blank prices because
|
||||
// an operator saved mid-edit.
|
||||
log.Warnf("agent-network pricing defaults reload failed for %s (keeping previous table): %v", path, err)
|
||||
return
|
||||
}
|
||||
storeFileTable(table, mtime)
|
||||
log.Infof("agent-network pricing defaults reloaded from %s", path)
|
||||
}
|
||||
|
||||
func readFile(path string) (map[string]map[string]Entry, int64, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("open pricing defaults %s: %w", path, err)
|
||||
}
|
||||
defer func() { _ = f.Close() }()
|
||||
|
||||
st, err := f.Stat()
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("stat pricing defaults %s: %w", path, err)
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxFileBytes+1))
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("read pricing defaults %s: %w", path, err)
|
||||
}
|
||||
if len(data) > maxFileBytes {
|
||||
return nil, 0, fmt.Errorf("pricing defaults %s exceeds %d bytes", path, maxFileBytes)
|
||||
}
|
||||
table, err := parsePricingYAML(data)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("parse pricing defaults %s: %w", path, err)
|
||||
}
|
||||
return table, st.ModTime().UnixNano(), nil
|
||||
}
|
||||
|
||||
// storeFileTable merges the parsed file over the compiled-in base and
|
||||
// publishes the result as the live table. File entries replace the
|
||||
// built-in entry for the same (surface, model) whole — they are not
|
||||
// field-merged — and surfaces/models the file doesn't mention keep the
|
||||
// built-in rates, so a partial file only needs the entries it changes.
|
||||
func storeFileTable(table map[string]map[string]Entry, mtime int64) {
|
||||
base := compiledBase()
|
||||
merged := make(map[string]map[string]Entry, len(base)+len(table))
|
||||
for surface, models := range base {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for id, e := range models {
|
||||
inner[id] = e
|
||||
}
|
||||
merged[surface] = inner
|
||||
}
|
||||
for surface, models := range table {
|
||||
inner, ok := merged[surface]
|
||||
if !ok {
|
||||
inner = make(map[string]Entry, len(models))
|
||||
merged[surface] = inner
|
||||
}
|
||||
for id, e := range models {
|
||||
inner[id] = e
|
||||
}
|
||||
}
|
||||
mergedTable.Store(&merged)
|
||||
setMtime(mtime)
|
||||
}
|
||||
|
||||
func setMtime(v int64) {
|
||||
fileState.mu.Lock()
|
||||
fileState.mtime = v
|
||||
fileState.mu.Unlock()
|
||||
}
|
||||
|
||||
// parsePricingYAML decodes and validates the pricing YAML. Unknown
|
||||
// fields are rejected (typos surface instead of silently pricing at 0)
|
||||
// and every rate must be a finite, non-negative USD amount — the same
|
||||
// constraints the HTTP API enforces on operator per-provider prices.
|
||||
func parsePricingYAML(data []byte) (map[string]map[string]Entry, error) {
|
||||
dec := yaml.NewDecoder(bytes.NewReader(data))
|
||||
dec.KnownFields(true)
|
||||
|
||||
var raw pricingFile
|
||||
if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, fmt.Errorf("decode yaml: %w", err)
|
||||
}
|
||||
|
||||
out := make(map[string]map[string]Entry, len(raw))
|
||||
for surface, models := range raw {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for model, e := range models {
|
||||
for field, v := range map[string]float64{
|
||||
"input_per_1k": e.InputPer1k,
|
||||
"output_per_1k": e.OutputPer1k,
|
||||
"cached_input_per_1k": e.CachedInputPer1k,
|
||||
"cache_read_per_1k": e.CacheReadPer1k,
|
||||
"cache_creation_per_1k": e.CacheCreationPer1k,
|
||||
} {
|
||||
if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
|
||||
return nil, fmt.Errorf("%s/%s: %s must be a finite, non-negative rate, got %v", surface, model, field, v)
|
||||
}
|
||||
}
|
||||
inner[model] = Entry{
|
||||
InputPer1k: e.InputPer1k,
|
||||
OutputPer1k: e.OutputPer1k,
|
||||
CachedInputPer1k: e.CachedInputPer1k,
|
||||
CacheReadPer1k: e.CacheReadPer1k,
|
||||
CacheCreationPer1k: e.CacheCreationPer1k,
|
||||
}
|
||||
}
|
||||
out[surface] = inner
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// resetFileState snapshots and restores the package-level file state so
|
||||
// tests stay order-independent.
|
||||
func resetFileState(t *testing.T) {
|
||||
t.Helper()
|
||||
prevMerged := mergedTable.Load()
|
||||
fileState.mu.Lock()
|
||||
prevPath, prevMtime := fileState.path, fileState.mtime
|
||||
fileState.mu.Unlock()
|
||||
t.Cleanup(func() {
|
||||
mergedTable.Store(prevMerged)
|
||||
fileState.mu.Lock()
|
||||
fileState.path, fileState.mtime = prevPath, prevMtime
|
||||
fileState.mu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
func writePricing(t *testing.T, path, yml string) {
|
||||
t.Helper()
|
||||
require.NoError(t, os.WriteFile(path, []byte(yml), 0o600))
|
||||
}
|
||||
|
||||
func TestLoadFile_MergesOverCompiledDefaults(t *testing.T) {
|
||||
resetFileState(t)
|
||||
path := filepath.Join(t.TempDir(), DefaultFileName)
|
||||
writePricing(t, path, `
|
||||
openai:
|
||||
# Reprice a built-in model. The entry replaces the built-in WHOLE:
|
||||
# omitting the cache rate here drops the built-in 0.00125 discount.
|
||||
gpt-4o:
|
||||
input_per_1k: 0.9
|
||||
output_per_1k: 1.8
|
||||
# A model NetBird doesn't know at all.
|
||||
my-private-ft:
|
||||
input_per_1k: 0.01
|
||||
output_per_1k: 0.02
|
||||
cached_input_per_1k: 0.005
|
||||
gemini:
|
||||
gemini-pro:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.005
|
||||
`)
|
||||
require.NoError(t, LoadFile(path, true))
|
||||
table := DefaultTable()
|
||||
|
||||
gpt4o := table["openai"]["gpt-4o"]
|
||||
assert.InDelta(t, 0.9, gpt4o.InputPer1k, 1e-9, "file rate replaces the compiled-in rate")
|
||||
assert.Zero(t, gpt4o.CachedInputPer1k, "entries replace whole — omitted cache rate is dropped, not inherited")
|
||||
|
||||
ft := table["openai"]["my-private-ft"]
|
||||
assert.InDelta(t, 0.005, ft.CachedInputPer1k, 1e-9, "unknown models are added to the surface")
|
||||
_, ok := table["gemini"]["gemini-pro"]
|
||||
assert.True(t, ok, "a surface the catalog doesn't declare can be added")
|
||||
|
||||
// Untouched entries keep compiled-in rates (catalog, other surface,
|
||||
// supplemental).
|
||||
assert.InDelta(t, 0.00015, table["openai"]["gpt-4o-mini"].InputPer1k, 1e-9, "unlisted model keeps compiled rate")
|
||||
assert.InDelta(t, 0.003, table["anthropic"]["claude-sonnet-4-5"].InputPer1k, 1e-9, "unlisted surface untouched")
|
||||
assert.InDelta(t, 0.00125, table["openai"]["gpt-5"].InputPer1k, 1e-9, "supplemental entries untouched")
|
||||
|
||||
// The synthesizer-facing lookup reads the live table too.
|
||||
e, ok := LookupDefault([]string{"openai"}, "gpt-4o")
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.9, e.InputPer1k, 1e-9, "LookupDefault serves the file-backed rate")
|
||||
}
|
||||
|
||||
func TestLoadFile_MissingPath(t *testing.T) {
|
||||
resetFileState(t)
|
||||
missing := filepath.Join(t.TempDir(), DefaultFileName)
|
||||
|
||||
require.Error(t, LoadFile(missing, true),
|
||||
"explicitly configured path that doesn't exist must fail startup")
|
||||
|
||||
require.NoError(t, LoadFile(missing, false),
|
||||
"conventional datadir probe tolerates an absent file (compiled-in defaults serve)")
|
||||
assert.Nil(t, mergedTable.Load(), "no file, no merged table")
|
||||
fileState.mu.Lock()
|
||||
path := fileState.path
|
||||
fileState.mu.Unlock()
|
||||
assert.Equal(t, missing, path, "the path stays registered so the reloader picks the file up when it appears")
|
||||
}
|
||||
|
||||
func TestLoadFile_RejectsInvalid(t *testing.T) {
|
||||
resetFileState(t)
|
||||
dir := t.TempDir()
|
||||
cases := map[string]string{
|
||||
"unknown field (typo)": "openai:\n gpt-4o:\n input_per1k: 0.1\n",
|
||||
"negative rate": "openai:\n gpt-4o:\n input_per_1k: -0.1\n",
|
||||
"non-numeric rate": "openai:\n gpt-4o:\n input_per_1k: cheap\n",
|
||||
"not a mapping": "- just\n- a\n- list\n",
|
||||
}
|
||||
for name, yml := range cases {
|
||||
path := filepath.Join(dir, DefaultFileName)
|
||||
writePricing(t, path, yml)
|
||||
assert.Error(t, LoadFile(path, true), "case %q must be rejected", name)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReload_LifeCycle drives the reloader's single-shot reload through
|
||||
// its full lifecycle: file edit picked up on mtime change, a broken save
|
||||
// keeps the previous table, and file removal reverts to the compiled-in
|
||||
// defaults (then a re-created file loads again).
|
||||
func TestReload_LifeCycle(t *testing.T) {
|
||||
resetFileState(t)
|
||||
path := filepath.Join(t.TempDir(), DefaultFileName)
|
||||
writePricing(t, path, "openai:\n gpt-4o:\n input_per_1k: 0.5\n output_per_1k: 1\n")
|
||||
require.NoError(t, LoadFile(path, true))
|
||||
require.InDelta(t, 0.5, DefaultTable()["openai"]["gpt-4o"].InputPer1k, 1e-9)
|
||||
|
||||
// Edit: new mtime, new rates.
|
||||
writePricing(t, path, "openai:\n gpt-4o:\n input_per_1k: 0.7\n output_per_1k: 1.4\n")
|
||||
bumpMtime(t, path)
|
||||
reload()
|
||||
assert.InDelta(t, 0.7, DefaultTable()["openai"]["gpt-4o"].InputPer1k, 1e-9, "edit must be picked up")
|
||||
|
||||
// Broken save: previous table survives.
|
||||
writePricing(t, path, "openai:\n gpt-4o:\n input_per_1k: -1\n")
|
||||
bumpMtime(t, path)
|
||||
reload()
|
||||
assert.InDelta(t, 0.7, DefaultTable()["openai"]["gpt-4o"].InputPer1k, 1e-9,
|
||||
"a malformed save must keep the previously loaded table, never blank prices")
|
||||
|
||||
// Removal: compiled-in defaults serve again.
|
||||
require.NoError(t, os.Remove(path))
|
||||
reload()
|
||||
assert.InDelta(t, 0.0025, DefaultTable()["openai"]["gpt-4o"].InputPer1k, 1e-9,
|
||||
"file removal reverts to the compiled-in rate")
|
||||
|
||||
// Re-created file loads without a restart.
|
||||
writePricing(t, path, "openai:\n gpt-4o:\n input_per_1k: 0.9\n output_per_1k: 1.8\n")
|
||||
bumpMtime(t, path)
|
||||
reload()
|
||||
assert.InDelta(t, 0.9, DefaultTable()["openai"]["gpt-4o"].InputPer1k, 1e-9,
|
||||
"a file appearing after removal (or after a missing-probe boot) must load")
|
||||
}
|
||||
|
||||
// bumpMtime pushes the file's mtime forward past the previously recorded
|
||||
// value — timestamps can otherwise collide within the test's timescale.
|
||||
func bumpMtime(t *testing.T, path string) {
|
||||
t.Helper()
|
||||
st, err := os.Stat(path)
|
||||
require.NoError(t, err)
|
||||
next := st.ModTime().Add(2 * 1e9)
|
||||
require.NoError(t, os.Chtimes(path, next, next))
|
||||
}
|
||||
|
||||
// TestExampleYAML_InSyncWithBuiltins is the golden guard for
|
||||
// defaults_llm_pricing.example.yaml: the shipped example must stay
|
||||
// byte-identical to what the compiled-in table renders (catalog edits
|
||||
// require `go generate ./management/internals/modules/agentnetwork/pricing`)
|
||||
// and must round-trip through the same parser operators' files go
|
||||
// through, reproducing the compiled-in table exactly.
|
||||
func TestExampleYAML_InSyncWithBuiltins(t *testing.T) {
|
||||
onDisk, err := os.ReadFile("defaults_llm_pricing.example.yaml")
|
||||
require.NoError(t, err, "example file must exist next to the package")
|
||||
require.Equal(t, string(MarshalDefaultsYAML()), string(onDisk),
|
||||
"defaults_llm_pricing.example.yaml is stale — run: go generate ./management/internals/modules/agentnetwork/pricing")
|
||||
|
||||
parsed, err := parsePricingYAML(onDisk)
|
||||
require.NoError(t, err, "the example must be a valid pricing defaults file")
|
||||
assert.Equal(t, buildDefaultTable(), parsed,
|
||||
"parsing the example must reproduce the compiled-in table exactly")
|
||||
}
|
||||
@@ -233,9 +233,19 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
return nil, err
|
||||
}
|
||||
|
||||
costMeterJSON, err := buildCostMeterConfigJSON(enabledProviders, groupIndex)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -243,7 +253,7 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
// Use the merged decision (account settings OR policy-required redaction),
|
||||
// not the raw account flag, so a policy that mandates PII redaction is
|
||||
// honored by the capture parsers even when the account toggle is off.
|
||||
middlewares := buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON, mergedGuardrails.PromptCapture.RedactPii, mergedGuardrails.PromptCapture.Enabled)
|
||||
middlewares := buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON, costMeterJSON, mergedGuardrails.PromptCapture.RedactPii, mergedGuardrails.PromptCapture.Enabled)
|
||||
|
||||
priv, pub, err := pickServiceSessionKeys(enabledProviders)
|
||||
if err != nil {
|
||||
@@ -540,6 +550,7 @@ type identityInjectJSONMetadata struct {
|
||||
UserKey string `json:"user_key,omitempty"`
|
||||
GroupsKey string `json:"groups_key,omitempty"`
|
||||
MaxValueLength int `json:"max_value_length,omitempty"`
|
||||
Sanitize bool `json:"sanitize,omitempty"`
|
||||
}
|
||||
|
||||
// buildIdentityInjectConfigJSON walks the enabled providers and emits
|
||||
@@ -583,9 +594,11 @@ func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[s
|
||||
func buildIdentityInjectRule(p *types.Provider, entry catalog.Provider) (identityInjectProvider, bool) {
|
||||
rule := identityInjectProvider{ProviderID: p.ID}
|
||||
// Identity-stamping shape (one of HeaderPair / JSONMetadata). Skip the
|
||||
// shape silently when the catalog entry doesn't declare one — extras
|
||||
// can still apply, see below.
|
||||
if entry.IdentityInjection != nil {
|
||||
// shape silently when the catalog entry doesn't declare one, or when the
|
||||
// operator disabled metadata for this provider — extras can still apply,
|
||||
// see below. MetadataDisabled suppresses only the identity dimensions
|
||||
// (user + authorizing group), not the catalog's routing ExtraHeaders.
|
||||
if !p.MetadataDisabled && entry.IdentityInjection != nil {
|
||||
switch {
|
||||
case entry.IdentityInjection.HeaderPair != nil:
|
||||
rule.HeaderPair = buildIdentityHeaderPair(p, entry.IdentityInjection.HeaderPair)
|
||||
@@ -651,6 +664,7 @@ func buildIdentityJSONMetadata(p *types.Provider, jm *catalog.JSONMetadataInject
|
||||
UserKey: userKey,
|
||||
GroupsKey: groupsKey,
|
||||
MaxValueLength: jm.MaxValueLength,
|
||||
Sanitize: jm.Sanitize,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -691,7 +705,7 @@ func buildIdentityExtraHeaders(p *types.Provider, extras []catalog.ExtraHeader)
|
||||
// requests bound for gateways like LiteLLM that key budgets and
|
||||
// attribution off request headers. CanMutate is required so its
|
||||
// HeadersAdd / HeadersRemove pass the framework's mutation gate.
|
||||
func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byte, redactPii, capturePromptContent bool) []rpservice.MiddlewareConfig {
|
||||
func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON, costMeterJSON []byte, redactPii, capturePromptContent bool) []rpservice.MiddlewareConfig {
|
||||
// Both parsers receive an explicit capture flag derived from the account's
|
||||
// enable_prompt_collection toggle; nil/unset would default to the legacy
|
||||
// "always emit" behavior in the middleware, which is precisely what we
|
||||
@@ -760,10 +774,13 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
|
||||
ConfigJSON: []byte("{}"),
|
||||
},
|
||||
{
|
||||
// Carries the full pricing table (defaults + per-provider
|
||||
// operator prices) so the proxy bills without an embedded
|
||||
// price list; see buildCostMeterConfigJSON.
|
||||
ID: middlewareIDCostMeter,
|
||||
Enabled: true,
|
||||
Slot: rpservice.MiddlewareSlotOnResponse,
|
||||
ConfigJSON: []byte("{}"),
|
||||
ConfigJSON: costMeterJSON,
|
||||
},
|
||||
{
|
||||
ID: middlewareIDLLMResponseParser,
|
||||
@@ -776,10 +793,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 {
|
||||
@@ -824,13 +843,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 {
|
||||
@@ -839,6 +855,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
|
||||
@@ -982,38 +1066,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 {
|
||||
@@ -1021,64 +1078,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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// costMeterConfig is the JSON shape the proxy-side cost_meter middleware
|
||||
// expects (mirror-type pattern, same as routerConfig). The top-level
|
||||
// "pricing" wrapper is the feature-detection signal: an old proxy's config
|
||||
// struct ignores it as an unknown field, and a new proxy treats its
|
||||
// absence as "old management" (skips every cost computation and warns).
|
||||
type costMeterConfig struct {
|
||||
Pricing *costMeterPricing `json:"pricing,omitempty"`
|
||||
}
|
||||
|
||||
// costMeterPricing carries the full pricing table:
|
||||
// - Defaults: surface ("openai"/"anthropic"/"bedrock") -> normalized
|
||||
// model id -> rates. The full default table ships to every account —
|
||||
// it is small (~10 KB) and keeps gateway-style providers (which
|
||||
// enumerate no models) priced for every catalog model.
|
||||
// - Providers: provider record id (matched against the
|
||||
// llm.resolved_provider_id metadata llm_router stamps) -> normalized
|
||||
// model id -> rates. Entries are fully materialized here at synth
|
||||
// time — default cache rates already folded in — so the proxy does
|
||||
// two map lookups and no merging.
|
||||
type costMeterPricing struct {
|
||||
Defaults map[string]map[string]pricing.Entry `json:"defaults,omitempty"`
|
||||
Providers map[string]map[string]pricing.Entry `json:"providers,omitempty"`
|
||||
}
|
||||
|
||||
// buildCostMeterConfigJSON assembles the cost_meter middleware config
|
||||
// from the default pricing table plus the operator's stored per-provider
|
||||
// model prices. Same orphan rule as the router: a provider no enabled
|
||||
// policy authorises is unreachable, so its prices are not shipped.
|
||||
//
|
||||
// Overlay semantics per model row:
|
||||
// - The row's model id is normalized exactly the way the proxy's
|
||||
// request parser normalizes the ids it meters (bedrock ARN/region/
|
||||
// version stripping, vertex "@version" stripping), so the per-record
|
||||
// lookup key compares equal to llm.model at billing time.
|
||||
// - The entry starts from the default entry for that model (when one
|
||||
// exists) to inherit cache rates the operator didn't state.
|
||||
// - Operator input/output overlay verbatim — including an explicit 0,
|
||||
// which prices the model as free (self-hosted / internal endpoints)
|
||||
// rather than silently reverting to list price.
|
||||
// - Cache-rate pointers overlay only when non-nil: nil means "inherit
|
||||
// the default", an explicit 0 means "no discount, bill this bucket
|
||||
// at the input rate".
|
||||
func buildCostMeterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) {
|
||||
cfg := costMeterConfig{Pricing: &costMeterPricing{
|
||||
Defaults: pricing.DefaultTable(),
|
||||
}}
|
||||
|
||||
perRecord := make(map[string]map[string]pricing.Entry)
|
||||
for _, p := range providers {
|
||||
if _, hasPolicy := groupIndex[p.ID]; !hasPolicy {
|
||||
// Orphan: unreachable via the router, so unpriceable.
|
||||
continue
|
||||
}
|
||||
if len(p.Models) == 0 {
|
||||
// Gateway-style "claim every model" provider: the defaults
|
||||
// table is its price list.
|
||||
continue
|
||||
}
|
||||
entry, _ := catalog.Lookup(p.ProviderID)
|
||||
models := make(map[string]pricing.Entry, len(p.Models))
|
||||
for _, m := range p.Models {
|
||||
id := normalizePricingModelID(p.ProviderID, m.ID)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if _, dup := models[id]; dup {
|
||||
// First occurrence wins on post-normalization duplicates,
|
||||
// matching providerModelIDs' dedup order for routing.
|
||||
continue
|
||||
}
|
||||
models[id] = materializeEntry(entry.PricingSurfaces, id, m)
|
||||
}
|
||||
if len(models) > 0 {
|
||||
perRecord[p.ID] = models
|
||||
}
|
||||
}
|
||||
if len(perRecord) > 0 {
|
||||
cfg.Pricing.Providers = perRecord
|
||||
}
|
||||
|
||||
out, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal cost_meter middleware config: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// normalizePricingModelID maps an operator-entered model id onto the
|
||||
// normalized id the proxy's request parser emits as llm.model — the key
|
||||
// the cost meter looks up at billing time.
|
||||
func normalizePricingModelID(catalogProviderID, modelID string) string {
|
||||
switch {
|
||||
case catalog.IsBedrockPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
case catalog.IsVertexPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeVertexModel(modelID)
|
||||
default:
|
||||
return modelID
|
||||
}
|
||||
}
|
||||
|
||||
// materializeEntry folds the default entry for (surfaces, model) — when
|
||||
// one exists — under the operator's stored prices, producing the fully
|
||||
// materialized wire entry.
|
||||
func materializeEntry(surfaces []string, normalizedID string, m types.ProviderModel) pricing.Entry {
|
||||
e, _ := pricing.LookupDefault(surfaces, normalizedID) // zero Entry on miss
|
||||
e.InputPer1k = m.InputPer1k
|
||||
e.OutputPer1k = m.OutputPer1k
|
||||
if m.CachedInputPer1k != nil {
|
||||
e.CachedInputPer1k = *m.CachedInputPer1k
|
||||
}
|
||||
if m.CacheReadPer1k != nil {
|
||||
e.CacheReadPer1k = *m.CacheReadPer1k
|
||||
}
|
||||
if m.CacheCreationPer1k != nil {
|
||||
e.CacheCreationPer1k = *m.CacheCreationPer1k
|
||||
}
|
||||
return e
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
func fptr(v float64) *float64 { return &v }
|
||||
|
||||
func decodeCostMeterConfig(t *testing.T, raw []byte) costMeterConfig {
|
||||
t.Helper()
|
||||
var cfg costMeterConfig
|
||||
require.NoError(t, json.Unmarshal(raw, &cfg), "cost meter config must round-trip")
|
||||
require.NotNil(t, cfg.Pricing, "pricing wrapper must be present")
|
||||
return cfg
|
||||
}
|
||||
|
||||
// TestBuildCostMeterConfig_BedrockModelNormalization: the operator may
|
||||
// paste region-prefixed, versioned, or ARN-wrapped Bedrock ids; the
|
||||
// per-record map must be keyed by the normalized id the request parser
|
||||
// emits as llm.model, or the lookup never hits at billing time.
|
||||
func TestBuildCostMeterConfig_BedrockModelNormalization(t *testing.T) {
|
||||
bedrock := &types.Provider{
|
||||
ID: "prov-bedrock",
|
||||
ProviderID: "bedrock_api",
|
||||
Enabled: true,
|
||||
Models: []types.ProviderModel{
|
||||
{ID: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", InputPer1k: 0.0033, OutputPer1k: 0.0165},
|
||||
// Post-normalization duplicate of the row above under a
|
||||
// different regional spelling — first occurrence wins.
|
||||
{ID: "us.anthropic.claude-sonnet-4-5-20250929-v1:0", InputPer1k: 9.9, OutputPer1k: 9.9},
|
||||
},
|
||||
}
|
||||
raw, err := buildCostMeterConfigJSON([]*types.Provider{bedrock}, map[string][]string{"prov-bedrock": {"grp"}})
|
||||
require.NoError(t, err)
|
||||
cfg := decodeCostMeterConfig(t, raw)
|
||||
|
||||
models := cfg.Pricing.Providers["prov-bedrock"]
|
||||
require.Len(t, models, 1, "both spellings normalize to one model; first row wins")
|
||||
e, ok := models["anthropic.claude-sonnet-4-5"]
|
||||
require.True(t, ok, "key must be the normalized id the parser emits, not the operator's raw spelling")
|
||||
assert.InDelta(t, 0.0033, e.InputPer1k, 1e-9, "first row's rate wins the dedup")
|
||||
assert.InDelta(t, 0.0003, e.CacheReadPer1k, 1e-9, "cache read inherited from the bedrock default entry")
|
||||
assert.InDelta(t, 0.00375, e.CacheCreationPer1k, 1e-9, "cache creation inherited from the bedrock default entry")
|
||||
}
|
||||
|
||||
// TestBuildCostMeterConfig_CacheRateNilVsZero pins the pointer semantics:
|
||||
// nil inherits the default cache rate, explicit 0 clears it (that bucket
|
||||
// bills at the input rate on the proxy).
|
||||
func TestBuildCostMeterConfig_CacheRateNilVsZero(t *testing.T) {
|
||||
p := &types.Provider{
|
||||
ID: "prov-oai",
|
||||
ProviderID: "openai_api",
|
||||
Enabled: true,
|
||||
Models: []types.ProviderModel{
|
||||
{ID: "gpt-4o", InputPer1k: 0.002, OutputPer1k: 0.008}, // nil → inherit 0.00125
|
||||
{ID: "gpt-4o-mini", InputPer1k: 0.0001, OutputPer1k: 0.0005, CachedInputPer1k: fptr(0)}, // explicit 0 → no discount
|
||||
{ID: "my-custom-ft", InputPer1k: 0.01, OutputPer1k: 0.02, CachedInputPer1k: fptr(0.005)}, // unknown model, explicit rate
|
||||
},
|
||||
}
|
||||
raw, err := buildCostMeterConfigJSON([]*types.Provider{p}, map[string][]string{"prov-oai": {"grp"}})
|
||||
require.NoError(t, err)
|
||||
cfg := decodeCostMeterConfig(t, raw)
|
||||
models := cfg.Pricing.Providers["prov-oai"]
|
||||
|
||||
assert.InDelta(t, 0.00125, models["gpt-4o"].CachedInputPer1k, 1e-9, "nil cache pointer inherits the default rate")
|
||||
assert.Zero(t, models["gpt-4o-mini"].CachedInputPer1k, "explicit 0 overrides the default (0.000075) — bucket bills at input rate")
|
||||
custom := models["my-custom-ft"]
|
||||
assert.InDelta(t, 0.005, custom.CachedInputPer1k, 1e-9, "unknown model keeps the operator's explicit cache rate")
|
||||
assert.Zero(t, custom.CacheReadPer1k, "no default to inherit for a model outside the catalog")
|
||||
}
|
||||
|
||||
// TestBuildCostMeterConfig_OrphanAndGatewayProviders: an orphan (no
|
||||
// authorising policy) is unreachable so its prices must not ship; a
|
||||
// gateway with no model rows relies on the defaults table and gets no
|
||||
// per-record entry.
|
||||
func TestBuildCostMeterConfig_OrphanAndGatewayProviders(t *testing.T) {
|
||||
orphan := &types.Provider{
|
||||
ID: "prov-orphan",
|
||||
ProviderID: "openai_api",
|
||||
Enabled: true,
|
||||
Models: []types.ProviderModel{{ID: "gpt-4o", InputPer1k: 1, OutputPer1k: 1}},
|
||||
}
|
||||
gateway := &types.Provider{
|
||||
ID: "prov-litellm",
|
||||
ProviderID: "litellm_proxy",
|
||||
Enabled: true,
|
||||
Models: []types.ProviderModel{},
|
||||
}
|
||||
raw, err := buildCostMeterConfigJSON(
|
||||
[]*types.Provider{orphan, gateway},
|
||||
map[string][]string{"prov-litellm": {"grp"}}, // orphan has no policy
|
||||
)
|
||||
require.NoError(t, err)
|
||||
cfg := decodeCostMeterConfig(t, raw)
|
||||
|
||||
assert.NotContains(t, cfg.Pricing.Providers, "prov-orphan", "orphan provider prices must not ship")
|
||||
assert.NotContains(t, cfg.Pricing.Providers, "prov-litellm", "empty-models gateway needs no per-record entry")
|
||||
assert.NotEmpty(t, cfg.Pricing.Defaults["openai"], "defaults still ship so the gateway's catalog-model traffic is priced")
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
@@ -33,14 +33,17 @@ func newSynthTestSettings() *types.Settings {
|
||||
|
||||
func newSynthTestProvider() *types.Provider {
|
||||
return &types.Provider{
|
||||
ID: "prov-1",
|
||||
AccountID: testAccountID,
|
||||
ProviderID: "openai_api",
|
||||
Name: "OpenAI",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test-key",
|
||||
Enabled: true,
|
||||
Models: []types.ProviderModel{{ID: "gpt-5.4", InputPer1k: 0.0025, OutputPer1k: 0.015}},
|
||||
ID: "prov-1",
|
||||
AccountID: testAccountID,
|
||||
ProviderID: "openai_api",
|
||||
Name: "OpenAI",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test-key",
|
||||
Enabled: true,
|
||||
// Prices deliberately differ from the catalog's gpt-5.4 rates
|
||||
// (0.0025/0.015) so pricing tests can prove the operator's
|
||||
// stored price overlays the catalog default.
|
||||
Models: []types.ProviderModel{{ID: "gpt-5.4", InputPer1k: 0.004, OutputPer1k: 0.02}},
|
||||
SessionPrivateKey: "test-priv-key",
|
||||
SessionPublicKey: "test-pub-key",
|
||||
CreatedAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC),
|
||||
@@ -214,7 +217,27 @@ func TestSynthesizeServices_HappyPath(t *testing.T) {
|
||||
|
||||
assert.Equal(t, middlewareIDCostMeter, mws[6].ID, "seventh middleware is the cost meter")
|
||||
assert.Equal(t, rpservice.MiddlewareSlotOnResponse, mws[6].Slot, "cost meter runs on_response")
|
||||
assert.Equal(t, []byte("{}"), mws[6].ConfigJSON, "cost meter carries an explicit empty config")
|
||||
|
||||
var costCfg costMeterConfig
|
||||
require.NoError(t, json.Unmarshal(mws[6].ConfigJSON, &costCfg), "cost meter config must unmarshal")
|
||||
require.NotNil(t, costCfg.Pricing, "cost meter config must carry the pricing table — its absence tells the proxy management predates config-delivered pricing")
|
||||
|
||||
gpt4o, ok := costCfg.Pricing.Defaults["openai"]["gpt-4o"]
|
||||
require.True(t, ok, "the full default table ships regardless of the account's providers")
|
||||
assert.InDelta(t, 0.0025, gpt4o.InputPer1k, 1e-9, "default gpt-4o input rate comes from the catalog")
|
||||
|
||||
openaiPrices, ok := costCfg.Pricing.Providers[openai.ID]
|
||||
require.True(t, ok, "operator-priced provider must have a per-record entry")
|
||||
gpt54, ok := openaiPrices["gpt-5.4"]
|
||||
require.True(t, ok, "operator's model row keys the per-record map")
|
||||
assert.InDelta(t, 0.004, gpt54.InputPer1k, 1e-9, "operator input price overlays the catalog default (0.0025)")
|
||||
assert.InDelta(t, 0.02, gpt54.OutputPer1k, 1e-9, "operator output price overlays the catalog default (0.015)")
|
||||
assert.InDelta(t, 0.00025, gpt54.CachedInputPer1k, 1e-9, "cache rate the operator didn't state is inherited from the default entry")
|
||||
|
||||
opus, ok := costCfg.Pricing.Providers[anthropic.ID]["claude-opus-4-7"]
|
||||
require.True(t, ok, "anthropic's model row keys its per-record map")
|
||||
assert.Zero(t, opus.InputPer1k, "operator-stored zero prices ship verbatim — an explicit $0 model bills as free, it does not revert to list price")
|
||||
assert.InDelta(t, 0.0005, opus.CacheReadPer1k, 1e-9, "cache rates still inherit from the default entry")
|
||||
|
||||
assert.Equal(t, middlewareIDLLMResponseParser, mws[7].ID, "eighth middleware is the response parser")
|
||||
assert.Equal(t, rpservice.MiddlewareSlotOnResponse, mws[7].Slot, "response parser runs on_response")
|
||||
@@ -698,6 +721,94 @@ func TestSynthesizeServices_IdentityInject_Portkey_NotCustomizable(t *testing.T)
|
||||
"same fixed-schema guarantee for the groups dimension")
|
||||
}
|
||||
|
||||
// TestSynthesizeServices_IdentityInject_Bedrock pins Bedrock's cost-allocation
|
||||
// metadata: a JSONMetadata shape emitting X-Amzn-Bedrock-Request-Metadata with
|
||||
// the reserved user/group keys, sanitized to Bedrock's accepted charset.
|
||||
func TestSynthesizeServices_IdentityInject_Bedrock(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ctrl := gomock.NewController(t)
|
||||
defer ctrl.Finish()
|
||||
mockStore := store.NewMockStore(ctrl)
|
||||
|
||||
br := newSynthTestProvider()
|
||||
br.ID = "prov-bedrock"
|
||||
br.ProviderID = "bedrock_api"
|
||||
br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
br.APIKey = "bedrock-bearer"
|
||||
br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
policy := newSynthTestPolicy(br.ID, "grp-eng", "")
|
||||
policy.ID = "pol-bedrock"
|
||||
|
||||
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
|
||||
[]*types.Provider{br},
|
||||
[]*types.Policy{policy},
|
||||
[]*types.Guardrail{})
|
||||
|
||||
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, services, 1)
|
||||
|
||||
var injectCfg identityInjectConfig
|
||||
for _, m := range services[0].Targets[0].Options.Middlewares {
|
||||
if m.ID == middlewareIDLLMIdentityInject {
|
||||
require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
|
||||
break
|
||||
}
|
||||
}
|
||||
require.Len(t, injectCfg.Providers, 1)
|
||||
entry := injectCfg.Providers[0]
|
||||
require.NotNil(t, entry.JSONMetadata, "Bedrock uses the JSONMetadata shape for cost-allocation metadata")
|
||||
assert.Nil(t, entry.HeaderPair, "shapes are mutually exclusive")
|
||||
assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", entry.JSONMetadata.Header,
|
||||
"the caller identity lands in Bedrock's cost-allocation metadata header")
|
||||
assert.Equal(t, "user", entry.JSONMetadata.UserKey)
|
||||
assert.Equal(t, "group", entry.JSONMetadata.GroupsKey)
|
||||
assert.True(t, entry.JSONMetadata.Sanitize,
|
||||
"Bedrock restricts the metadata value charset, so values must be sanitized")
|
||||
}
|
||||
|
||||
// TestSynthesizeServices_MetadataDisabled_SuppressesInjection verifies the
|
||||
// per-provider opt-out: a provider with MetadataDisabled set emits no
|
||||
// identity-inject entry (Bedrock has no catalog ExtraHeaders, so the whole
|
||||
// entry is dropped).
|
||||
func TestSynthesizeServices_MetadataDisabled_SuppressesInjection(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ctrl := gomock.NewController(t)
|
||||
defer ctrl.Finish()
|
||||
mockStore := store.NewMockStore(ctrl)
|
||||
|
||||
br := newSynthTestProvider()
|
||||
br.ID = "prov-bedrock"
|
||||
br.ProviderID = "bedrock_api"
|
||||
br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
br.APIKey = "bedrock-bearer"
|
||||
br.MetadataDisabled = true
|
||||
br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
policy := newSynthTestPolicy(br.ID, "grp-eng", "")
|
||||
policy.ID = "pol-bedrock"
|
||||
|
||||
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
|
||||
[]*types.Provider{br},
|
||||
[]*types.Policy{policy},
|
||||
[]*types.Guardrail{})
|
||||
|
||||
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, services, 1)
|
||||
|
||||
var injectCfg identityInjectConfig
|
||||
for _, m := range services[0].Targets[0].Options.Middlewares {
|
||||
if m.ID == middlewareIDLLMIdentityInject {
|
||||
require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.Empty(t, injectCfg.Providers,
|
||||
"metadata_disabled must drop the provider's identity-inject entry")
|
||||
}
|
||||
|
||||
// TestSynthesizeServices_IdentityInject_Vercel pins Vercel AI
|
||||
// Gateway's wiring: HeaderPair shape with fixed wire names dictated
|
||||
// by Vercel's Custom Reporting API (ai-reporting-user /
|
||||
@@ -943,8 +1054,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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(*)",
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -14,10 +14,24 @@ import (
|
||||
// ProviderModel is one row in the provider's models list. The operator
|
||||
// pins the per-1k input/output price for cost tracking; ID is the
|
||||
// model identifier the upstream provider expects on the wire.
|
||||
//
|
||||
// The three cache rates are pointers because absence is meaningful: nil
|
||||
// means "inherit NetBird's default rate for this model" (folded in at
|
||||
// synthesis time), while an explicit 0 means "no discount — bill this
|
||||
// cache bucket at the input rate".
|
||||
type ProviderModel struct {
|
||||
ID string `json:"id"`
|
||||
InputPer1k float64 `json:"input_per_1k"`
|
||||
OutputPer1k float64 `json:"output_per_1k"`
|
||||
// CachedInputPer1k is the OpenAI-shape rate for cached prompt tokens
|
||||
// (a subset of input tokens).
|
||||
CachedInputPer1k *float64 `json:"cached_input_per_1k,omitempty"`
|
||||
// CacheReadPer1k is the Anthropic-shape rate for cache-read tokens
|
||||
// (additive to input tokens).
|
||||
CacheReadPer1k *float64 `json:"cache_read_per_1k,omitempty"`
|
||||
// CacheCreationPer1k is the Anthropic-shape rate for cache-creation
|
||||
// tokens (additive to input tokens).
|
||||
CacheCreationPer1k *float64 `json:"cache_creation_per_1k,omitempty"`
|
||||
}
|
||||
|
||||
// Provider is an Agent Network AI provider record persisted per account.
|
||||
@@ -51,6 +65,12 @@ type Provider struct {
|
||||
// private or self-signed certificate. The synthesiser propagates it into
|
||||
// the router route so the proxy dials that provider's upstream insecurely.
|
||||
SkipTLSVerification bool `gorm:"column:skip_tls_verification"`
|
||||
// MetadataDisabled suppresses identity metadata injection for this provider.
|
||||
// Metadata (the caller's user + authorizing group) is injected by default;
|
||||
// when true the synthesiser omits the provider's identity-inject shape, so no
|
||||
// user/group headers (e.g. Bedrock's X-Amzn-Bedrock-Request-Metadata) are
|
||||
// stamped. Catalog ExtraHeaders (routing config) are unaffected.
|
||||
MetadataDisabled bool `gorm:"column:metadata_disabled"`
|
||||
// SessionPrivateKey + SessionPublicKey are the ed25519 keypair the
|
||||
// synthesised reverse-proxy service uses to sign / verify session
|
||||
// JWTs after a successful OIDC handshake. Generated once on
|
||||
@@ -122,9 +142,12 @@ func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) {
|
||||
if req.Models != nil {
|
||||
for _, m := range *req.Models {
|
||||
p.Models = append(p.Models, ProviderModel{
|
||||
ID: m.Id,
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
ID: m.Id,
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
CachedInputPer1k: copyFloatPtr(m.CachedInputPer1k),
|
||||
CacheReadPer1k: copyFloatPtr(m.CacheReadPer1k),
|
||||
CacheCreationPer1k: copyFloatPtr(m.CacheCreationPer1k),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -137,6 +160,9 @@ func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) {
|
||||
if req.SkipTlsVerification != nil {
|
||||
p.SkipTLSVerification = *req.SkipTlsVerification
|
||||
}
|
||||
if req.MetadataDisabled != nil {
|
||||
p.MetadataDisabled = *req.MetadataDisabled
|
||||
}
|
||||
// Identity-header overrides for catalogs flagged Customizable.
|
||||
// nil pointer = "field omitted on the wire" → leave the stored
|
||||
// value untouched (per the openapi description). Empty string is
|
||||
@@ -155,9 +181,12 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
|
||||
models := make([]api.AgentNetworkProviderModel, 0, len(p.Models))
|
||||
for _, m := range p.Models {
|
||||
models = append(models, api.AgentNetworkProviderModel{
|
||||
Id: m.ID,
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
Id: m.ID,
|
||||
InputPer1k: m.InputPer1k,
|
||||
OutputPer1k: m.OutputPer1k,
|
||||
CachedInputPer1k: copyFloatPtr(m.CachedInputPer1k),
|
||||
CacheReadPer1k: copyFloatPtr(m.CacheReadPer1k),
|
||||
CacheCreationPer1k: copyFloatPtr(m.CacheCreationPer1k),
|
||||
})
|
||||
}
|
||||
created := p.CreatedAt
|
||||
@@ -170,6 +199,7 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
|
||||
Models: models,
|
||||
Enabled: p.Enabled,
|
||||
SkipTlsVerification: p.SkipTLSVerification,
|
||||
MetadataDisabled: p.MetadataDisabled,
|
||||
CreatedAt: &created,
|
||||
UpdatedAt: &updated,
|
||||
}
|
||||
@@ -191,11 +221,27 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
|
||||
return resp
|
||||
}
|
||||
|
||||
// copyFloatPtr returns a fresh pointer to the same value, or nil. Keeps
|
||||
// stored models and API payloads from aliasing each other's rate fields.
|
||||
func copyFloatPtr(v *float64) *float64 {
|
||||
if v == nil {
|
||||
return nil
|
||||
}
|
||||
out := *v
|
||||
return &out
|
||||
}
|
||||
|
||||
// Copy returns a deep copy of the provider.
|
||||
func (p *Provider) Copy() *Provider {
|
||||
clone := *p
|
||||
if p.Models != nil {
|
||||
clone.Models = append([]ProviderModel(nil), p.Models...)
|
||||
clone.Models = make([]ProviderModel, len(p.Models))
|
||||
for i, m := range p.Models {
|
||||
m.CachedInputPer1k = copyFloatPtr(m.CachedInputPer1k)
|
||||
m.CacheReadPer1k = copyFloatPtr(m.CacheReadPer1k)
|
||||
m.CacheCreationPer1k = copyFloatPtr(m.CacheCreationPer1k)
|
||||
clone.Models[i] = m
|
||||
}
|
||||
}
|
||||
if p.ExtraValues != nil {
|
||||
clone.ExtraValues = make(map[string]string, len(p.ExtraValues))
|
||||
|
||||
@@ -42,3 +42,38 @@ func TestProvider_SkipTLSVerification_RoundTrip(t *testing.T) {
|
||||
assert.False(t, p.SkipTLSVerification, "explicit false must clear skip_tls_verification")
|
||||
assert.False(t, p.ToAPIResponse().SkipTlsVerification, "response must reflect the cleared value")
|
||||
}
|
||||
|
||||
// TestProvider_MetadataDisabled_RoundTrip covers the request→provider→response
|
||||
// mapping of metadata_disabled, with the same update semantics: nil preserves,
|
||||
// explicit false clears.
|
||||
func TestProvider_MetadataDisabled_RoundTrip(t *testing.T) {
|
||||
enable := true
|
||||
disable := false
|
||||
|
||||
base := func() *api.AgentNetworkProviderRequest {
|
||||
return &api.AgentNetworkProviderRequest{
|
||||
ProviderId: "bedrock_api",
|
||||
Name: "bedrock",
|
||||
UpstreamUrl: "https://bedrock-runtime.us-east-1.amazonaws.com",
|
||||
}
|
||||
}
|
||||
|
||||
p := NewProvider("acc-1")
|
||||
|
||||
req := base()
|
||||
req.MetadataDisabled = &enable
|
||||
p.FromAPIRequest(req)
|
||||
assert.True(t, p.MetadataDisabled, "create with metadata_disabled=true must set the field")
|
||||
assert.True(t, p.ToAPIResponse().MetadataDisabled, "response must surface metadata_disabled")
|
||||
|
||||
// Omitting the field on update leaves the stored value untouched.
|
||||
p.FromAPIRequest(base())
|
||||
assert.True(t, p.MetadataDisabled, "omitting metadata_disabled on update must preserve it")
|
||||
|
||||
// Explicit false clears it (re-enables metadata).
|
||||
req = base()
|
||||
req.MetadataDisabled = &disable
|
||||
p.FromAPIRequest(req)
|
||||
assert.False(t, p.MetadataDisabled, "explicit false must clear metadata_disabled")
|
||||
assert.False(t, p.ToAPIResponse().MetadataDisabled, "response must reflect the cleared value")
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -103,6 +103,12 @@ func TestSynthesizedService_WireShape(t *testing.T) {
|
||||
|
||||
assert.Equal(t, middlewareIDCostMeter, mws[6].GetId(), "seventh middleware id")
|
||||
assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, mws[6].GetSlot(), "cost meter slot")
|
||||
var costCfg costMeterConfig
|
||||
require.NoError(t, json.Unmarshal(mws[6].GetConfigJson(), &costCfg), "cost meter config JSON must decode from the wire")
|
||||
require.NotNil(t, costCfg.Pricing, "the pricing table must travel on the wire — the proxy has no embedded price list to fall back to")
|
||||
assert.NotEmpty(t, costCfg.Pricing.Defaults["openai"], "default table rides in every mapping")
|
||||
assert.NotEmpty(t, costCfg.Pricing.Defaults["anthropic"], "default table covers all surfaces")
|
||||
assert.NotEmpty(t, costCfg.Pricing.Defaults["bedrock"], "default table covers all surfaces")
|
||||
|
||||
assert.Equal(t, middlewareIDLLMResponseParser, mws[7].GetId(), "eighth middleware id")
|
||||
assert.Equal(t, proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, mws[7].GetSlot(), "response parser slot")
|
||||
|
||||
@@ -36,7 +36,7 @@ func NewGRPCController(proxyGRPCServer *nbgrpc.ProxyServiceServer, meter metric.
|
||||
// SendServiceUpdateToCluster sends a service update to a specific proxy cluster.
|
||||
func (c *GRPCController) SendServiceUpdateToCluster(ctx context.Context, accountID string, update *proto.ProxyMapping, clusterAddr string) {
|
||||
c.proxyGRPCServer.SendServiceUpdateToCluster(ctx, update, clusterAddr)
|
||||
c.metrics.IncrementServiceUpdateSendCount(clusterAddr)
|
||||
c.metrics.IncrementServiceUpdateSendCount()
|
||||
}
|
||||
|
||||
// GetOIDCValidationConfig returns the OIDC validation configuration from the gRPC server.
|
||||
@@ -53,7 +53,7 @@ func (c *GRPCController) RegisterProxyToCluster(ctx context.Context, clusterAddr
|
||||
proxySet.(*sync.Map).Store(proxyID, struct{}{})
|
||||
log.WithContext(ctx).Debugf("Registered proxy %s to cluster %s", proxyID, clusterAddr)
|
||||
|
||||
c.metrics.IncrementProxyConnectionCount(clusterAddr)
|
||||
c.metrics.IncrementProxyConnectionCount()
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -67,7 +67,7 @@ func (c *GRPCController) UnregisterProxyFromCluster(ctx context.Context, cluster
|
||||
proxySet.(*sync.Map).Delete(proxyID)
|
||||
log.WithContext(ctx).Debugf("Unregistered proxy %s from cluster %s", proxyID, clusterAddr)
|
||||
|
||||
c.metrics.DecrementProxyConnectionCount(clusterAddr)
|
||||
c.metrics.DecrementProxyConnectionCount()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ package manager
|
||||
import (
|
||||
"context"
|
||||
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/metric"
|
||||
)
|
||||
|
||||
@@ -48,25 +47,16 @@ func newMetrics(meter metric.Meter) (*metrics, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *metrics) IncrementProxyConnectionCount(clusterAddr string) {
|
||||
m.proxyConnectionCount.Add(context.Background(), 1,
|
||||
metric.WithAttributes(
|
||||
attribute.String("cluster", clusterAddr),
|
||||
))
|
||||
func (m *metrics) IncrementProxyConnectionCount() {
|
||||
m.proxyConnectionCount.Add(context.Background(), 1)
|
||||
}
|
||||
|
||||
func (m *metrics) DecrementProxyConnectionCount(clusterAddr string) {
|
||||
m.proxyConnectionCount.Add(context.Background(), -1,
|
||||
metric.WithAttributes(
|
||||
attribute.String("cluster", clusterAddr),
|
||||
))
|
||||
func (m *metrics) DecrementProxyConnectionCount() {
|
||||
m.proxyConnectionCount.Add(context.Background(), -1)
|
||||
}
|
||||
|
||||
func (m *metrics) IncrementServiceUpdateSendCount(clusterAddr string) {
|
||||
m.serviceUpdateSendCount.Add(context.Background(), 1,
|
||||
metric.WithAttributes(
|
||||
attribute.String("cluster", clusterAddr),
|
||||
))
|
||||
func (m *metrics) IncrementServiceUpdateSendCount() {
|
||||
m.serviceUpdateSendCount.Add(context.Background(), 1)
|
||||
}
|
||||
|
||||
func (m *metrics) IncrementProxyHeartbeatCount() {
|
||||
|
||||
@@ -55,12 +55,18 @@ type Config struct {
|
||||
|
||||
ReverseProxy ReverseProxy
|
||||
|
||||
AgentNetwork AgentNetwork
|
||||
|
||||
// disable default all-to-all policy
|
||||
DisableDefaultPolicy bool
|
||||
|
||||
// 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
|
||||
@@ -181,6 +187,25 @@ type StoreConfig struct {
|
||||
Engine types.Engine
|
||||
}
|
||||
|
||||
// AgentNetwork contains agent-network (LLM gateway) configuration.
|
||||
type AgentNetwork struct {
|
||||
// PricingDefaultsFile is the path to the YAML file holding the default
|
||||
// LLM pricing table (defaults_llm_pricing.yaml). A relative path is
|
||||
// resolved against <Datadir>, so a bare filename lands alongside the
|
||||
// store. Empty falls back to probing <Datadir>/defaults_llm_pricing.yaml;
|
||||
// with no file present the compiled-in defaults serve. Schema: surface ("openai"/"anthropic"/
|
||||
// "bedrock") -> model -> rates in USD per 1k tokens (input_per_1k,
|
||||
// output_per_1k, and the optional cached_input_per_1k /
|
||||
// cache_read_per_1k / cache_creation_per_1k). File entries replace the
|
||||
// compiled-in entry for the same surface+model whole; everything else
|
||||
// keeps the compiled-in rates. The file is re-read periodically (mtime
|
||||
// poll), and the live table feeds both the synthesizer (what proxies
|
||||
// bill with) and the dashboard's catalog endpoint (what model rows
|
||||
// prefill with). An explicitly configured path that fails to load
|
||||
// fails startup; runtime reload errors keep the previous table.
|
||||
PricingDefaultsFile string
|
||||
}
|
||||
|
||||
// ReverseProxy contains reverse proxy configuration in front of management.
|
||||
type ReverseProxy struct {
|
||||
// TrustedHTTPProxies represents a list of trusted HTTP proxies by their IP prefixes.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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": {}}},
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -1648,6 +1648,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)
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
@@ -4269,7 +4279,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_NetworkRangePreserved(t *te
|
||||
}
|
||||
|
||||
// Sanity: an actually different range still triggers reallocation.
|
||||
newRange := netip.MustParsePrefix("100.99.0.0/16")
|
||||
newRange := netip.MustParsePrefix("100.60.0.0/16")
|
||||
_, err = manager.UpdateAccountSettings(ctx, account.Id, userID, &types.Settings{
|
||||
PeerLoginExpirationEnabled: true,
|
||||
PeerLoginExpiration: types.DefaultPeerLoginExpiration,
|
||||
|
||||
@@ -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
|
||||
})
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
@@ -76,6 +79,9 @@ type EmbeddedIdPConfig struct {
|
||||
DashboardPostLogoutRedirectURIs []string
|
||||
// StaticConnectors are additional connectors to seed during initialization
|
||||
StaticConnectors []dex.Connector
|
||||
// GrantTypes restricts allowed OAuth2 grants; empty means all (Dex default). Omit the
|
||||
// device_code grant to disable the device flow; keep authorization_code and refresh_token.
|
||||
GrantTypes []string
|
||||
}
|
||||
|
||||
// EmbeddedStorageConfig holds storage configuration for the embedded IdP.
|
||||
@@ -175,6 +181,7 @@ func (c *EmbeddedIdPConfig) ToYAMLConfig() (*dex.YAMLConfig, error) {
|
||||
},
|
||||
OAuth2: dex.OAuth2{
|
||||
SkipApprovalScreen: true,
|
||||
GrantTypes: c.GrantTypes,
|
||||
},
|
||||
Frontend: dex.Frontend{
|
||||
Issuer: "NetBird",
|
||||
@@ -185,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,
|
||||
@@ -254,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 == "" {
|
||||
@@ -736,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.
|
||||
@@ -775,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.
|
||||
@@ -792,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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -1606,7 +1625,8 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
|
||||
settings_routing_peer_dns_resolution_enabled, settings_dns_domain, settings_network_range,
|
||||
settings_network_range_v6, settings_ipv6_enabled_groups, settings_lazy_connection_enabled,
|
||||
settings_local_mfa_enabled, settings_metrics_push_enabled, settings_agent_network_only,
|
||||
settings_dashboard_features,
|
||||
settings_dashboard_features, settings_auto_update_version, settings_auto_update_always,
|
||||
settings_peer_expose_enabled, settings_peer_expose_groups,
|
||||
-- Embedded ExtraSettings
|
||||
settings_extra_peer_approval_enabled, settings_extra_user_approval_required,
|
||||
settings_extra_integrated_validator, settings_extra_integrated_validator_groups
|
||||
@@ -1632,6 +1652,10 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
|
||||
sMetricsPushEnabled sql.NullBool
|
||||
sAgentNetworkOnly sql.NullBool
|
||||
sDashboardFeatures sql.NullString
|
||||
autoUpdateVersion sql.NullString
|
||||
autoUpdateAlways sql.NullBool
|
||||
peerExposeEnabled sql.NullBool
|
||||
peerExposeGroups sql.NullString
|
||||
sExtraPeerApprovalEnabled sql.NullBool
|
||||
sExtraUserApprovalRequired sql.NullBool
|
||||
sExtraIntegratedValidator sql.NullString
|
||||
@@ -1655,7 +1679,8 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
|
||||
&sRoutingPeerDNSResolutionEnabled, &sDNSDomain, &sNetworkRange,
|
||||
&sNetworkRangeV6, &sIPv6EnabledGroups, &sLazyConnectionEnabled,
|
||||
&sLocalMFAEnabled, &sMetricsPushEnabled, &sAgentNetworkOnly,
|
||||
&sDashboardFeatures,
|
||||
&sDashboardFeatures, &autoUpdateVersion, &autoUpdateAlways,
|
||||
&peerExposeEnabled, &peerExposeGroups,
|
||||
&sExtraPeerApprovalEnabled, &sExtraUserApprovalRequired,
|
||||
&sExtraIntegratedValidator, &sExtraIntegratedValidatorGroups,
|
||||
)
|
||||
@@ -1747,6 +1772,18 @@ func (s *SqlStore) getAccount(ctx context.Context, accountID string) (*types.Acc
|
||||
if sIPv6EnabledGroups.Valid {
|
||||
_ = json.Unmarshal([]byte(sIPv6EnabledGroups.String), &account.Settings.IPv6EnabledGroups)
|
||||
}
|
||||
if autoUpdateAlways.Valid {
|
||||
account.Settings.AutoUpdateAlways = autoUpdateAlways.Bool
|
||||
}
|
||||
if autoUpdateVersion.Valid {
|
||||
account.Settings.AutoUpdateVersion = autoUpdateVersion.String
|
||||
}
|
||||
if peerExposeEnabled.Valid {
|
||||
account.Settings.PeerExposeEnabled = peerExposeEnabled.Bool
|
||||
}
|
||||
if peerExposeGroups.Valid {
|
||||
_ = json.Unmarshal([]byte(peerExposeGroups.String), &account.Settings.PeerExposeGroups)
|
||||
}
|
||||
|
||||
if sExtraPeerApprovalEnabled.Valid {
|
||||
account.Settings.Extra.PeerApprovalEnabled = sExtraPeerApprovalEnabled.Bool
|
||||
@@ -1833,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 {
|
||||
@@ -1855,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,
|
||||
@@ -1864,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 {
|
||||
@@ -1984,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
|
||||
})
|
||||
@@ -2039,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
|
||||
@@ -2049,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)
|
||||
@@ -2074,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
|
||||
@@ -2083,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
|
||||
@@ -2101,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
|
||||
@@ -2111,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
|
||||
@@ -2153,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
|
||||
@@ -2162,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
|
||||
@@ -2198,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
|
||||
@@ -2206,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)
|
||||
}
|
||||
@@ -2218,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
|
||||
}
|
||||
@@ -2339,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
|
||||
}
|
||||
@@ -2385,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
|
||||
@@ -2403,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
|
||||
@@ -2413,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
|
||||
@@ -2441,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
|
||||
@@ -2450,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
|
||||
@@ -3812,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")
|
||||
}
|
||||
@@ -3900,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")
|
||||
@@ -6100,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()
|
||||
@@ -6107,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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"sort"
|
||||
"sync"
|
||||
@@ -34,6 +35,7 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
"github.com/netbirdio/netbird/shared/testing_helpers"
|
||||
"github.com/netbirdio/netbird/util/crypt"
|
||||
)
|
||||
|
||||
@@ -45,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) {
|
||||
@@ -296,6 +299,53 @@ func Test_SaveAccount(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func Test_AccountSettings_SaveAndRetrieve(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
}
|
||||
|
||||
populateFields := testing_helpers.NewPopulateFields().WithCustomFieldSetter(
|
||||
reflect.PointerTo(reflect.TypeOf(types.ExtraSettings{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
|
||||
es := types.ExtraSettings{}
|
||||
reflectedEs := reflect.ValueOf(&es).Elem()
|
||||
n, err := this.PopulateAll(reflectedEs)
|
||||
if err != nil {
|
||||
return n, err
|
||||
}
|
||||
field.Set(reflectedEs.Addr())
|
||||
return n, nil
|
||||
}).WithCustomFieldSetter(
|
||||
reflect.PointerTo(reflect.TypeOf(types.DashboardFeatures{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
|
||||
t := true
|
||||
df := types.DashboardFeatures{AgentNetwork: &t}
|
||||
reflectedDf := reflect.ValueOf(&df).Elem()
|
||||
field.Set(reflectedDf.Addr())
|
||||
return 1, nil
|
||||
}).WithSkippedTag("gorm", "-")
|
||||
|
||||
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
|
||||
account := newAccountWithId(context.Background(), "account_id", "testuser", "")
|
||||
setupKey, _ := types.GenerateDefaultSetupKey()
|
||||
account.SetupKeys[setupKey.Key] = setupKey
|
||||
|
||||
settings := types.Settings{}
|
||||
numOfExportedFields, err := populateFields.PopulateAll(reflect.ValueOf(&settings).Elem())
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, 27, numOfExportedFields)
|
||||
account.Settings = &settings
|
||||
|
||||
err = store.SaveAccount(context.Background(), account)
|
||||
assert.NoError(t, err)
|
||||
|
||||
accountFromDb, err := store.GetAccount(context.Background(), account.Id)
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, accountFromDb)
|
||||
assert.NotNil(t, accountFromDb.Settings)
|
||||
|
||||
assert.True(t, reflect.DeepEqual(&settings, accountFromDb.Settings), "created settings and settings retrieved from the db should match")
|
||||
})
|
||||
}
|
||||
|
||||
func TestSqlite_DeleteAccount(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("The SQLite store is not properly supported by Windows yet")
|
||||
@@ -512,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) {
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -1076,7 +1056,7 @@ func (a *Account) GetPeerConnectionResources(ctx context.Context, peer *nbpeer.P
|
||||
default:
|
||||
authorizedUsers[auth.Wildcard] = a.getAllowedUserIDs()
|
||||
}
|
||||
} else if peerServesAuth && policyRuleImpliesLegacySSH(rule) && peer.SSHEnabled {
|
||||
} else if peerServesAuth && PolicyRuleImpliesLegacySSH(rule) && peer.SSHEnabled {
|
||||
sshEnabled = true
|
||||
authorizedUsers[auth.Wildcard] = a.getAllowedUserIDs()
|
||||
}
|
||||
@@ -1107,6 +1087,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 {
|
||||
@@ -1142,15 +1123,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) {
|
||||
@@ -1158,10 +1139,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 {
|
||||
@@ -1176,7 +1153,7 @@ func PeerSSHEnabledFromPolicies(policies []*Policy, peerID string, peerGroupIDs
|
||||
}
|
||||
|
||||
isSSHRule := rule.Protocol == PolicyRuleProtocolNetbirdSSH ||
|
||||
(policyRuleImpliesLegacySSH(rule) && peerSSHEnabled)
|
||||
(PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled)
|
||||
if !isSSHRule {
|
||||
continue
|
||||
}
|
||||
@@ -1203,24 +1180,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
|
||||
@@ -1320,14 +1279,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]
|
||||
@@ -1354,13 +1313,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
|
||||
}
|
||||
@@ -1563,6 +1522,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{})
|
||||
@@ -1813,96 +1820,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 {
|
||||
|
||||
@@ -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)
|
||||
@@ -358,7 +495,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
|
||||
default:
|
||||
sshReqs.needAllowedUserIDs = true
|
||||
}
|
||||
} else if policyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
|
||||
} else if PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
|
||||
sshReqs.needAllowedUserIDs = true
|
||||
}
|
||||
}
|
||||
@@ -491,7 +628,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 {
|
||||
@@ -500,17 +644,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
|
||||
}
|
||||
@@ -545,7 +687,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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user