mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-22 08:21:30 +02:00
Compare commits
12 Commits
lazy-conn-
...
windows-sl
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
25fe7d13d9 | ||
|
|
c533dead28 | ||
|
|
21fc5b81f6 | ||
|
|
67af1b2ef2 | ||
|
|
9906b9b1a1 | ||
|
|
877e889250 | ||
|
|
d15830a2d0 | ||
|
|
141f3d0390 | ||
|
|
e1a24376ab | ||
|
|
8f901f8899 | ||
|
|
c6bf5fbbfb | ||
|
|
e70a69bbcf |
@@ -17,7 +17,9 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal"
|
||||
"github.com/netbirdio/netbird/client/internal/auth"
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
nbnet "github.com/netbirdio/netbird/client/net"
|
||||
"github.com/netbirdio/netbird/client/proto"
|
||||
"github.com/netbirdio/netbird/client/server"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
@@ -331,6 +333,14 @@ func doForegroundLogin(ctx context.Context, cmd *cobra.Command, setupKey string,
|
||||
return fmt.Errorf("read config file %s: %v", configFilePath, err)
|
||||
}
|
||||
|
||||
// Mirror runInForegroundMode: recover residual state (DNS, firewall,
|
||||
// ssh config, legacy routing) from a previous unclean shutdown and
|
||||
// enable advanced routing before dialing management.
|
||||
if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configFilePath).GetStatePath()); err != nil {
|
||||
log.Warnf("failed to restore residual state: %v", err)
|
||||
}
|
||||
nbnet.Init()
|
||||
|
||||
err = foregroundLogin(ctx, cmd, config, setupKey, activeProf.ID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("foreground login failed: %v", err)
|
||||
|
||||
@@ -22,6 +22,8 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
"github.com/netbirdio/netbird/client/proto"
|
||||
nbnet "github.com/netbirdio/netbird/client/net"
|
||||
"github.com/netbirdio/netbird/client/server"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
@@ -229,6 +231,24 @@ func runInForegroundMode(ctx context.Context, cmd *cobra.Command, activeProf *pr
|
||||
|
||||
_, _ = profilemanager.UpdateOldManagementURL(ctx, config, configFilePath)
|
||||
|
||||
// Restore residual state left by a previous run that did not shut down
|
||||
// cleanly, mirroring what the daemon does before connecting: it recovers
|
||||
// DNS config (a stale resolv.conf takeover can make the management
|
||||
// hostname unresolvable), firewall rules, ssh config and legacy routing.
|
||||
// Route cleanup itself happens at engine start; nbnet.Init() below lets
|
||||
// the management dial bypass a leftover fwmark rule until then.
|
||||
// Foreground mode is particularly exposed in containers: a crashed
|
||||
// container restarts inside the same (pod) network namespace, so stale
|
||||
// state survives while the process does not.
|
||||
if err := server.RestoreResidualState(ctx, profilemanager.NewServiceManager(configPath).GetStatePath()); err != nil {
|
||||
log.Warnf("failed to restore residual state: %v", err)
|
||||
}
|
||||
|
||||
// Enable advanced routing (as the daemon does on startup) so the
|
||||
// management dial bypasses a leftover fwmark rule instead of being
|
||||
// shunted into a stale routing table.
|
||||
nbnet.Init()
|
||||
|
||||
err = foregroundLogin(ctx, cmd, config, providedSetupKey, activeProf.ID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("foreground login failed: %v", err)
|
||||
|
||||
@@ -292,18 +292,16 @@ func (s *serviceViaListener) generateFreePort() (uint16, error) {
|
||||
return customPort, nil
|
||||
}
|
||||
|
||||
udpAddr := net.UDPAddrFromAddrPort(netip.MustParseAddrPort("0.0.0.0:0"))
|
||||
probeListener, err := net.ListenUDP("udp", udpAddr)
|
||||
probeListener, err := net.ListenUDP("udp4", &net.UDPAddr{})
|
||||
if err != nil {
|
||||
log.Debugf("failed to bind random port for DNS: %s", err)
|
||||
return 0, err
|
||||
}
|
||||
|
||||
addrPort := netip.MustParseAddrPort(probeListener.LocalAddr().String()) // might panic if address is incorrect
|
||||
err = probeListener.Close()
|
||||
if err != nil {
|
||||
port := uint16(probeListener.LocalAddr().(*net.UDPAddr).Port)
|
||||
if err = probeListener.Close(); err != nil {
|
||||
log.Debugf("failed to free up DNS port: %s", err)
|
||||
return 0, err
|
||||
}
|
||||
return addrPort.Port(), nil
|
||||
return port, nil
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build !darwin || ios
|
||||
//go:build (!darwin && !windows) || ios
|
||||
|
||||
package sleep
|
||||
|
||||
|
||||
255
client/internal/sleep/detector_windows.go
Normal file
255
client/internal/sleep/detector_windows.go
Normal file
@@ -0,0 +1,255 @@
|
||||
//go:build windows
|
||||
|
||||
package sleep
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// Power broadcast event types delivered to the notification callback.
|
||||
// https://learn.microsoft.com/en-us/windows/win32/power/power-management-events
|
||||
const (
|
||||
pbtAPMSuspend uintptr = 0x0004
|
||||
pbtAPMResumeSuspend uintptr = 0x0007
|
||||
pbtAPMResumeAutomatic uintptr = 0x0012
|
||||
deviceNotifyCallback uintptr = 0x00000002
|
||||
)
|
||||
|
||||
var (
|
||||
// The callback flavor of these functions is exported by user32.dll, not
|
||||
// powrprof.dll (which exports the differently-shaped Power* variants).
|
||||
user32 = windows.NewLazySystemDLL("user32.dll")
|
||||
|
||||
// https://learn.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-registersuspendresumenotification
|
||||
registerSuspendResumeNotification = user32.NewProc("RegisterSuspendResumeNotification")
|
||||
// https://learn.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-unregistersuspendresumenotification
|
||||
unregisterSuspendResumeNotification = user32.NewProc("UnregisterSuspendResumeNotification")
|
||||
|
||||
libInitOnce sync.Once
|
||||
libInitErr error
|
||||
|
||||
// callbackThunk is the single C-callable trampoline registered with the OS.
|
||||
// windows.NewCallback slots are a finite, non-reclaimable resource, so a
|
||||
// single thunk dispatches to the Detector identified by the callback Context.
|
||||
callbackThunk uintptr
|
||||
|
||||
// registry maps the Context value handed to the OS callback back to the
|
||||
// Detector that registered it, mirroring darwin's serviceRegistry.
|
||||
registry = make(map[int]*Detector)
|
||||
registryMu sync.Mutex
|
||||
nextHandle int
|
||||
|
||||
// lifecycleMu serializes Register/Deregister so concurrent lifecycle calls
|
||||
// can't race on the shared registry or the OS registration handle.
|
||||
lifecycleMu sync.Mutex
|
||||
)
|
||||
|
||||
// deviceNotifySubscribeParameters is DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS.
|
||||
// https://learn.microsoft.com/en-us/windows/win32/api/winuser/ns-winuser-device_notify_subscribe_parameters
|
||||
type deviceNotifySubscribeParameters struct {
|
||||
Callback uintptr
|
||||
Context uintptr
|
||||
}
|
||||
|
||||
// Detector delivers sleep and wake events to a registered callback.
|
||||
type Detector struct {
|
||||
callback func(event EventType)
|
||||
done chan struct{}
|
||||
|
||||
// handle keys this detector in the package registry and is passed to the OS
|
||||
// as the callback Context. Zero means the detector is not registered.
|
||||
handle int
|
||||
// hPowerNotify is the HPOWERNOTIFY returned by RegisterSuspendResumeNotification.
|
||||
hPowerNotify uintptr
|
||||
// params is kept alive for the lifetime of the registration so the OS never
|
||||
// dereferences freed memory.
|
||||
params deviceNotifySubscribeParameters
|
||||
}
|
||||
|
||||
// NewDetector resolves powrprof.dll symbols and returns a Detector.
|
||||
func NewDetector() (*Detector, error) {
|
||||
if err := initLibs(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Detector{}, nil
|
||||
}
|
||||
|
||||
func initLibs() error {
|
||||
libInitOnce.Do(func() {
|
||||
if err := registerSuspendResumeNotification.Find(); err != nil {
|
||||
libInitErr = fmt.Errorf("resolve RegisterSuspendResumeNotification: %w", err)
|
||||
return
|
||||
}
|
||||
if err := unregisterSuspendResumeNotification.Find(); err != nil {
|
||||
libInitErr = fmt.Errorf("resolve UnregisterSuspendResumeNotification: %w", err)
|
||||
return
|
||||
}
|
||||
callbackThunk = windows.NewCallback(powerCallback)
|
||||
})
|
||||
return libInitErr
|
||||
}
|
||||
|
||||
// Register installs callback for power events and subscribes to suspend/resume
|
||||
// notifications via powrprof.dll.
|
||||
func (d *Detector) Register(callback func(event EventType)) error {
|
||||
lifecycleMu.Lock()
|
||||
defer lifecycleMu.Unlock()
|
||||
|
||||
registryMu.Lock()
|
||||
if d.handle != 0 {
|
||||
registryMu.Unlock()
|
||||
return fmt.Errorf("detector service already registered")
|
||||
}
|
||||
d.callback = callback
|
||||
d.done = make(chan struct{})
|
||||
nextHandle++
|
||||
handle := nextHandle
|
||||
d.handle = handle
|
||||
registry[handle] = d
|
||||
registryMu.Unlock()
|
||||
|
||||
d.params = deviceNotifySubscribeParameters{
|
||||
Callback: callbackThunk,
|
||||
Context: uintptr(handle),
|
||||
}
|
||||
|
||||
ret, _, callErr := registerSuspendResumeNotification.Call(
|
||||
uintptr(unsafe.Pointer(&d.params)),
|
||||
deviceNotifyCallback,
|
||||
)
|
||||
if ret == 0 {
|
||||
registryMu.Lock()
|
||||
delete(registry, handle)
|
||||
close(d.done)
|
||||
d.done = nil
|
||||
d.handle = 0
|
||||
registryMu.Unlock()
|
||||
return fmt.Errorf("RegisterSuspendResumeNotification failed: %w", callErr)
|
||||
}
|
||||
d.hPowerNotify = ret
|
||||
|
||||
log.Info("sleep detection service started on Windows")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Deregister unsubscribes from power notifications and removes the detector.
|
||||
func (d *Detector) Deregister() error {
|
||||
lifecycleMu.Lock()
|
||||
defer lifecycleMu.Unlock()
|
||||
|
||||
registryMu.Lock()
|
||||
if d.handle == 0 {
|
||||
registryMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
handle := d.handle
|
||||
hPowerNotify := d.hPowerNotify
|
||||
done := d.done
|
||||
registryMu.Unlock()
|
||||
|
||||
log.Info("sleep detection service stopping (deregister)")
|
||||
|
||||
// Unregister the OS subscription first. If it fails, leave handle and
|
||||
// hPowerNotify intact so a later call can retry the cleanup.
|
||||
if hPowerNotify != 0 {
|
||||
ret, _, callErr := unregisterSuspendResumeNotification.Call(hPowerNotify)
|
||||
if ret == 0 {
|
||||
return fmt.Errorf("UnregisterSuspendResumeNotification failed: %w", callErr)
|
||||
}
|
||||
}
|
||||
|
||||
registryMu.Lock()
|
||||
close(done)
|
||||
delete(registry, handle)
|
||||
d.handle = 0
|
||||
d.hPowerNotify = 0
|
||||
registryMu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *Detector) triggerCallback(event EventType, cb func(event EventType), done <-chan struct{}) {
|
||||
if cb == nil || done == nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
doneChan := make(chan struct{})
|
||||
// The OS invokes this callback synchronously on the suspend path, so run the
|
||||
// teardown inline with a bounded budget (mirroring the macOS detector) so
|
||||
// Down completes before the machine suspends without blocking indefinitely.
|
||||
timeout := time.NewTimer(20 * time.Second)
|
||||
defer timeout.Stop()
|
||||
|
||||
go func() {
|
||||
defer close(doneChan)
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Errorf("panic in sleep callback: %v", r)
|
||||
}
|
||||
}()
|
||||
log.Info("sleep detection event fired")
|
||||
cb(event)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-doneChan:
|
||||
case <-done:
|
||||
case <-timeout.C:
|
||||
log.Warn("sleep callback timed out")
|
||||
}
|
||||
}
|
||||
|
||||
// powerCallback is the DEVICE_NOTIFY_CALLBACK_ROUTINE trampoline, invoked by the
|
||||
// OS on a system thread. A Go panic crossing the syscall boundary has undefined
|
||||
// behavior, so contain it here. It must return ERROR_SUCCESS (0).
|
||||
func powerCallback(context uintptr, msgType uintptr, setting uintptr) uintptr {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Errorf("panic in sleep powerCallback: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
var event EventType
|
||||
switch msgType {
|
||||
case pbtAPMSuspend:
|
||||
event = EventTypeSleep
|
||||
case pbtAPMResumeAutomatic, pbtAPMResumeSuspend:
|
||||
event = EventTypeWakeUp
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
|
||||
dispatchEvent(int(context), event)
|
||||
return 0
|
||||
}
|
||||
|
||||
func dispatchEvent(handle int, event EventType) {
|
||||
registryMu.Lock()
|
||||
d := registry[handle]
|
||||
var (
|
||||
cb func(event EventType)
|
||||
done <-chan struct{}
|
||||
)
|
||||
if d != nil {
|
||||
cb = d.callback
|
||||
done = d.done
|
||||
}
|
||||
registryMu.Unlock()
|
||||
|
||||
if d == nil {
|
||||
return
|
||||
}
|
||||
d.triggerCallback(event, cb, done)
|
||||
}
|
||||
83
client/internal/sleep/detector_windows_test.go
Normal file
83
client/internal/sleep/detector_windows_test.go
Normal file
@@ -0,0 +1,83 @@
|
||||
//go:build windows
|
||||
|
||||
package sleep
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// registerFake inserts a detector into the registry without touching the OS,
|
||||
// so dispatch routing can be exercised in isolation. It returns the handle and
|
||||
// a cleanup that removes the entry.
|
||||
func registerFake(t *testing.T, cb func(EventType)) (int, func()) {
|
||||
t.Helper()
|
||||
|
||||
registryMu.Lock()
|
||||
nextHandle++
|
||||
handle := nextHandle
|
||||
d := &Detector{callback: cb, done: make(chan struct{}), handle: handle}
|
||||
registry[handle] = d
|
||||
registryMu.Unlock()
|
||||
|
||||
return handle, func() {
|
||||
registryMu.Lock()
|
||||
delete(registry, handle)
|
||||
registryMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func TestPowerCallback_MapsMessageTypes(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
msgType uintptr
|
||||
want EventType
|
||||
fires bool
|
||||
}{
|
||||
{"suspend", pbtAPMSuspend, EventTypeSleep, true},
|
||||
{"resume automatic", pbtAPMResumeAutomatic, EventTypeWakeUp, true},
|
||||
{"resume suspend", pbtAPMResumeSuspend, EventTypeWakeUp, true},
|
||||
{"unknown", 0x9999, EventTypeUnknown, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := make(chan EventType, 1)
|
||||
handle, cleanup := registerFake(t, func(e EventType) { got <- e })
|
||||
defer cleanup()
|
||||
|
||||
ret := powerCallback(uintptr(handle), tt.msgType, 0)
|
||||
require.Equal(t, uintptr(0), ret, "callback must return ERROR_SUCCESS")
|
||||
|
||||
if !tt.fires {
|
||||
assert.Empty(t, got, "no event should fire for unhandled message type")
|
||||
return
|
||||
}
|
||||
select {
|
||||
case e := <-got:
|
||||
assert.Equal(t, tt.want, e, "mapped event type should match")
|
||||
default:
|
||||
t.Fatal("expected callback to fire")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchEvent_UnknownHandleNoPanic(t *testing.T) {
|
||||
require.NotPanics(t, func() {
|
||||
dispatchEvent(-1, EventTypeSleep)
|
||||
}, "dispatch for an unregistered handle must be a no-op")
|
||||
}
|
||||
|
||||
func TestTriggerCallback_SkipsAfterDone(t *testing.T) {
|
||||
done := make(chan struct{})
|
||||
close(done)
|
||||
|
||||
fired := false
|
||||
d := &Detector{}
|
||||
d.triggerCallback(EventTypeSleep, func(EventType) { fired = true }, done)
|
||||
|
||||
assert.False(t, fired, "callback must not run once the detector is done")
|
||||
}
|
||||
@@ -44,10 +44,25 @@ type Auth struct {
|
||||
// NewAuth instantiate Auth struct and validate the management URL
|
||||
func NewAuth(cfgPath string, mgmURL string) (*Auth, error) {
|
||||
inputCfg := profilemanager.ConfigInput{
|
||||
ConfigPath: cfgPath,
|
||||
ManagementURL: mgmURL,
|
||||
}
|
||||
|
||||
cfg, err := profilemanager.CreateInMemoryConfig(inputCfg)
|
||||
// Load the existing config when a config file is already present so an
|
||||
// interactive re-login reuses the peer's persisted WireGuard private key
|
||||
// (and thus its identity) instead of generating a fresh one. Generating a
|
||||
// new key registers a brand-new peer on the management server on every
|
||||
// re-auth (named after the fallback hostname). Only fall back to a fresh
|
||||
// in-memory config for the first-time login when no config file exists yet.
|
||||
// DirectUpdateOrCreateConfig uses non-atomic writes so it also works inside
|
||||
// the tvOS App Group sandbox where atomic temp-file+rename is blocked.
|
||||
var cfg *profilemanager.Config
|
||||
var err error
|
||||
if cfgPath != "" {
|
||||
cfg, err = profilemanager.DirectUpdateOrCreateConfig(inputCfg)
|
||||
} else {
|
||||
cfg, err = profilemanager.CreateInMemoryConfig(inputCfg)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -181,7 +181,7 @@ func (s *Server) Start() error {
|
||||
log.Warnf("failed to redirect stderr: %v", err)
|
||||
}
|
||||
|
||||
if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
log.Warnf(errRestoreResidualState, err)
|
||||
}
|
||||
|
||||
@@ -551,7 +551,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
|
||||
s.actCancel = cancel
|
||||
s.mutex.Unlock()
|
||||
|
||||
if err := restoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
if err := RestoreResidualState(s.rootCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
log.Warnf(errRestoreResidualState, err)
|
||||
}
|
||||
|
||||
@@ -858,7 +858,7 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
|
||||
|
||||
return s.waitForUp(callerCtx)
|
||||
}
|
||||
if err := restoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
if err := RestoreResidualState(callerCtx, s.profileManager.GetStatePath()); err != nil {
|
||||
log.Warnf(errRestoreResidualState, err)
|
||||
}
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ func (s *Server) CleanState(ctx context.Context, req *proto.CleanStateRequest) (
|
||||
|
||||
if req.All {
|
||||
// Reuse existing cleanup logic for all states
|
||||
if err := restoreResidualState(ctx, statePath); err != nil {
|
||||
if err := RestoreResidualState(ctx, statePath); err != nil {
|
||||
return nil, status.Errorf(codes.Internal, "failed to clean all states: %v", err)
|
||||
}
|
||||
|
||||
@@ -113,9 +113,9 @@ func (s *Server) DeleteState(ctx context.Context, req *proto.DeleteStateRequest)
|
||||
}, nil
|
||||
}
|
||||
|
||||
// restoreResidualState checks if the client was not shut down in a clean way and restores residual if required.
|
||||
// RestoreResidualState checks if the client was not shut down in a clean way and restores residual if required.
|
||||
// Otherwise, we might not be able to connect to the management server to retrieve new config.
|
||||
func restoreResidualState(ctx context.Context, statePath string) error {
|
||||
func RestoreResidualState(ctx context.Context, statePath string) error {
|
||||
if statePath == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ func availableProviders() []providerCase {
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
}
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireMessages})
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock})
|
||||
}
|
||||
return ps
|
||||
}
|
||||
@@ -224,9 +224,12 @@ func TestProvidersMatrix(t *testing.T) {
|
||||
var c int
|
||||
var b string
|
||||
var cerr error
|
||||
if pc.kind == harness.WireVertex {
|
||||
switch pc.kind {
|
||||
case harness.WireVertex:
|
||||
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
|
||||
} else {
|
||||
case harness.WireBedrock:
|
||||
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
|
||||
default:
|
||||
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
|
||||
}
|
||||
if cerr == nil {
|
||||
|
||||
168
e2e/agentnetwork/guardrail_test.go
Normal file
168
e2e/agentnetwork/guardrail_test.go
Normal file
@@ -0,0 +1,168 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// catalogModel returns the normalized catalog id the proxy stamps for a
|
||||
// path-routed provider's configured model — the form the guardrail allowlist is
|
||||
// compared against (region prefix / @version stripped).
|
||||
func catalogModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
return strings.TrimPrefix(pc.model, "us.")
|
||||
case harness.WireVertex:
|
||||
return strings.SplitN(pc.model, "@", 2)[0]
|
||||
default:
|
||||
return pc.model
|
||||
}
|
||||
}
|
||||
|
||||
// disallowedModel returns a valid-shaped model id for the provider that is NOT
|
||||
// the configured/allowed one, so the guardrail must reject it before the
|
||||
// request ever reaches the upstream.
|
||||
func disallowedModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
return "us.anthropic.claude-opus-4-8"
|
||||
case harness.WireVertex:
|
||||
return "claude-opus-4-8@20250101"
|
||||
default:
|
||||
return "unlisted-model"
|
||||
}
|
||||
}
|
||||
|
||||
// sendModel drives one request for the given model through the provider's native
|
||||
// wire shape and returns the HTTP status.
|
||||
func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint, proxyIP string, pc providerCase, model string) int {
|
||||
t.Helper()
|
||||
var code int
|
||||
var err error
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
code, _, err = cl.Bedrock(ctx, endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
case harness.WireVertex:
|
||||
code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
|
||||
return code
|
||||
}
|
||||
|
||||
// TestModelAllowlistEnforced provisions a Model Allowlist guardrail limiting each
|
||||
// path-routed provider (Bedrock, Vertex) to its configured model, then drives
|
||||
// requests over the tunnel: the allowed model returns 200 while a model outside
|
||||
// the allowlist is denied 403 by the guardrail before it reaches the upstream.
|
||||
// This is the coverage missing for #6751 — the model for these providers travels
|
||||
// in the URL path, and the allowlist must be enforced there.
|
||||
func TestModelAllowlistEnforced(t *testing.T) {
|
||||
var providers []providerCase
|
||||
for _, pc := range availableProviders() {
|
||||
if pc.kind == harness.WireBedrock || pc.kind == harness.WireVertex {
|
||||
providers = append(providers, pc)
|
||||
}
|
||||
}
|
||||
if len(providers) == 0 {
|
||||
t.Skip("no path-routed provider keys set (AWS_BEARER_TOKEN_BEDROCK / GOOGLE_VERTEX_*); source ~/.llm-keys")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-allowlist"})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-allowlist-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
|
||||
// Providers with their configured (allowed) models; the first bootstraps the cluster.
|
||||
ids := make([]string, 0, len(providers))
|
||||
allowed := make([]string, 0, len(providers))
|
||||
for i, pc := range providers {
|
||||
req := providerRequest(pc)
|
||||
if i == 0 {
|
||||
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
|
||||
}
|
||||
prov, perr := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, perr, "create provider %s", pc.name)
|
||||
id := prov.Id
|
||||
ids = append(ids, id)
|
||||
allowed = append(allowed, catalogModel(pc))
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) })
|
||||
}
|
||||
|
||||
// Guardrail allowlisting exactly the configured models.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-allowlist"
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = allowed
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-allowlist",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: ids,
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings for endpoint")
|
||||
require.NotEmpty(t, settings.Endpoint, "agent-network endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-proxy-allowlist")
|
||||
require.NoError(t, err, "mint proxy token via CLI")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
|
||||
for _, pc := range providers {
|
||||
pc := pc
|
||||
t.Run(pc.name, func(t *testing.T) {
|
||||
// The admin's allowlisted model is served end to end.
|
||||
assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, pc.model),
|
||||
"allowlisted model must be permitted for %s", pc.name)
|
||||
// A model outside the allowlist is rejected by the guardrail (before
|
||||
// the upstream), regardless of whether it is a real catalog model.
|
||||
assert.Equal(t, 403, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, disallowedModel(pc)),
|
||||
"model outside the allowlist must be denied for %s", pc.name)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -107,6 +107,17 @@ func (c *Combined) DeletePolicy(ctx context.Context, id string) error {
|
||||
return anDelete(ctx, c, "/api/agent-network/policies/"+id)
|
||||
}
|
||||
|
||||
// CreateGuardrail creates an agent-network guardrail (e.g. a model allowlist)
|
||||
// that can then be attached to a policy via its GuardrailIds.
|
||||
func (c *Combined) CreateGuardrail(ctx context.Context, req api.AgentNetworkGuardrailRequest) (api.AgentNetworkGuardrail, error) {
|
||||
return anRequest[api.AgentNetworkGuardrail](ctx, c, http.MethodPost, "/api/agent-network/guardrails", req)
|
||||
}
|
||||
|
||||
// DeleteGuardrail removes a guardrail by id.
|
||||
func (c *Combined) DeleteGuardrail(ctx context.Context, id string) error {
|
||||
return anDelete(ctx, c, "/api/agent-network/guardrails/"+id)
|
||||
}
|
||||
|
||||
// GetSettings returns the account's agent-network settings row. It exists only
|
||||
// after the first provider create bootstraps it.
|
||||
func (c *Combined) GetSettings(ctx context.Context) (api.AgentNetworkSettings, error) {
|
||||
|
||||
@@ -194,6 +194,11 @@ const (
|
||||
// WireVertex is the Anthropic-on-Vertex rawPredict shape: the client posts
|
||||
// the full Vertex model path and the proxy mints the SA OAuth token.
|
||||
WireVertex = "vertex"
|
||||
// WireBedrock is the native AWS Bedrock InvokeModel shape: the model id
|
||||
// travels in the URL path (/model/{id}/invoke), not the body, so the proxy
|
||||
// routes by path. This is what a Bedrock SDK client sends and the shape the
|
||||
// model-allowlist guardrail must enforce.
|
||||
WireBedrock = "bedrock"
|
||||
)
|
||||
|
||||
// Chat issues a chat-completion POST to the agent-network endpoint over the
|
||||
@@ -226,6 +231,17 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
|
||||
}
|
||||
|
||||
// Bedrock issues a native AWS Bedrock InvokeModel POST over the tunnel. The
|
||||
// model id is carried in the request path (/model/{id}/invoke), so the proxy
|
||||
// routes by path; the body uses the bedrock anthropic_version rather than a
|
||||
// model field. A non-empty sessionID is sent as the universal x-session-id
|
||||
// header the proxy records.
|
||||
func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) {
|
||||
path := "/model/" + model + "/invoke"
|
||||
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
|
||||
}
|
||||
|
||||
// withSessionID appends the x-session-id header when sessionID is non-empty.
|
||||
func withSessionID(headers []string, sessionID string) []string {
|
||||
if sessionID == "" {
|
||||
|
||||
2
go.mod
2
go.mod
@@ -335,7 +335,7 @@ replace github.com/cloudflare/circl => codeberg.org/cunicu/circl v0.0.0-20230801
|
||||
|
||||
replace github.com/pion/ice/v4 => github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51
|
||||
|
||||
replace github.com/dexidp/dex => github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1
|
||||
replace github.com/dexidp/dex => github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5
|
||||
|
||||
replace github.com/dexidp/dex/api/v2 => github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1
|
||||
|
||||
|
||||
4
go.sum
4
go.sum
@@ -476,8 +476,8 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
|
||||
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1 h1:4TaYr9O4xX0D2kszeOLclTiCbA3eHq3xWV+9ILJbIYs=
|
||||
github.com/netbirdio/dex v0.244.1-0.20260512110716-8d70ad8647c1/go.mod h1:IHH+H8vK2GfqtIt5u/5OdPh18yk0oDHuj2vz5+Goetg=
|
||||
github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5 h1:3PwQv8aR46qN2u16+Dv6udnH3sbVKX5KrGwF35CKSI0=
|
||||
github.com/netbirdio/dex v0.244.1-0.20260716205454-a163de3129e5/go.mod h1:IHH+H8vK2GfqtIt5u/5OdPh18yk0oDHuj2vz5+Goetg=
|
||||
github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1 h1:neE7z+FPUkldl3faK/Jt+hJK2L+1XfQ1W33TQhU9m88=
|
||||
github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1/go.mod h1:awuTyT29CYALpEyET0S307EgNlPWrc7fFKRAyhsO45M=
|
||||
github.com/netbirdio/easyjson v0.9.0 h1:6Nw2lghSVuy8RSkAYDhDv1thBVEmfVbKZnV7T7Z6Aus=
|
||||
|
||||
@@ -613,6 +613,10 @@ func (c *YAMLConfig) ToServerConfig(stor storage.Storage, logger *slog.Logger) s
|
||||
cfg.SupportedResponseTypes = c.OAuth2.ResponseTypes
|
||||
}
|
||||
|
||||
if len(c.OAuth2.GrantTypes) > 0 {
|
||||
cfg.AllowedGrantTypes = c.OAuth2.GrantTypes
|
||||
}
|
||||
|
||||
// Apply expiry settings
|
||||
if c.Expiry.IDTokens != "" {
|
||||
if d, err := parseDuration(c.Expiry.IDTokens); err == nil {
|
||||
|
||||
@@ -21,7 +21,7 @@ import (
|
||||
"github.com/dexidp/dex/server/signer"
|
||||
"github.com/dexidp/dex/storage"
|
||||
"github.com/dexidp/dex/storage/sql"
|
||||
jose "github.com/go-jose/go-jose/v4"
|
||||
"github.com/go-jose/go-jose/v4"
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
|
||||
@@ -595,3 +595,90 @@ enablePasswordDB: true
|
||||
assert.True(t, cfg.ContinueOnConnectorFailure,
|
||||
"buildDexConfig must set ContinueOnConnectorFailure to true so management starts even if an external IdP is down")
|
||||
}
|
||||
|
||||
func TestToServerConfig_WiresGrantTypes(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "dex-grants-*")
|
||||
require.NoError(t, err)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
stor := openTestStorage(t, tmpDir)
|
||||
defer stor.Close()
|
||||
|
||||
logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
|
||||
|
||||
grants := []string{"authorization_code", "refresh_token"}
|
||||
cfg := &YAMLConfig{Issuer: "http://localhost:5599/oauth2", OAuth2: OAuth2{GrantTypes: grants}}
|
||||
assert.Equal(t, grants, cfg.ToServerConfig(stor, logger).AllowedGrantTypes)
|
||||
|
||||
empty := &YAMLConfig{Issuer: "http://localhost:5599/oauth2"}
|
||||
assert.Empty(t, empty.ToServerConfig(stor, logger).AllowedGrantTypes)
|
||||
}
|
||||
|
||||
func newDeviceGuardProvider(t *testing.T, grantTypesYAML string) *Provider {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "dex-devguard-*")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = os.RemoveAll(tmpDir) })
|
||||
|
||||
yamlContent := `
|
||||
issuer: http://localhost:5599/oauth2
|
||||
storage:
|
||||
type: sqlite3
|
||||
config:
|
||||
file: ` + filepath.Join(tmpDir, "dex.db") + `
|
||||
web:
|
||||
http: 127.0.0.1:5599
|
||||
enablePasswordDB: true
|
||||
` + grantTypesYAML
|
||||
|
||||
configPath := filepath.Join(tmpDir, "config.yaml")
|
||||
require.NoError(t, os.WriteFile(configPath, []byte(yamlContent), 0644))
|
||||
|
||||
yamlConfig, err := LoadConfig(configPath)
|
||||
require.NoError(t, err)
|
||||
|
||||
provider, err := NewProviderFromYAML(context.Background(), yamlConfig)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = provider.Stop(context.Background()) })
|
||||
return provider
|
||||
}
|
||||
|
||||
func TestHandler_BlocksDeviceEndpointsWhenDeviceGrantDisabled(t *testing.T) {
|
||||
provider := newDeviceGuardProvider(t, `
|
||||
oauth2:
|
||||
grantTypes:
|
||||
- authorization_code
|
||||
- refresh_token
|
||||
`)
|
||||
|
||||
devicePaths := []string{
|
||||
"/oauth2/device",
|
||||
"/oauth2/device/code",
|
||||
"/oauth2/device/token",
|
||||
"/oauth2/device/auth/verify_code",
|
||||
"/oauth2/device/callback",
|
||||
}
|
||||
for _, path := range devicePaths {
|
||||
for _, method := range []string{http.MethodGet, http.MethodPost} {
|
||||
req := httptest.NewRequest(method, path, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
provider.Handler().ServeHTTP(rec, req)
|
||||
assert.Equal(t, http.StatusNotFound, rec.Code, "%s %s must be blocked", method, path)
|
||||
}
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/oauth2/.well-known/openid-configuration", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
provider.Handler().ServeHTTP(rec, req)
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestHandler_AllowsDeviceEndpointsWhenGrantsDefault(t *testing.T) {
|
||||
provider := newDeviceGuardProvider(t, "")
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/oauth2/device/code", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
provider.Handler().ServeHTTP(rec, req)
|
||||
assert.NotEqual(t, http.StatusNotFound, rec.Code)
|
||||
}
|
||||
|
||||
@@ -226,30 +226,6 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee
|
||||
return nil
|
||||
}
|
||||
|
||||
// Dedupe stale embedded peer records for the same (account, cluster).
|
||||
// The proxy generates a fresh WireGuard keypair on every startup
|
||||
// (proxy/internal/roundtrip/netbird.go), so without this sweep the
|
||||
// prior embedded peer would linger forever — holding its CGNAT IP
|
||||
// allocation, polluting other peers' rosters, and (most visibly)
|
||||
// leaving the synth DNS pointing at the dead address. The
|
||||
// (account, cluster) tuple identifies "the embedded peer for this
|
||||
// proxy instance at this cluster"; any record matching that tuple
|
||||
// with a different pubkey is by definition stale and must go.
|
||||
staleIDs, err := m.findStaleEmbeddedProxyPeers(ctx, accountID, cluster, peerKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("scan for stale embedded proxy peers: %w", err)
|
||||
}
|
||||
if len(staleIDs) > 0 {
|
||||
// userID="" + checkConnected=false: the deletion is initiated
|
||||
// by management itself on behalf of the freshly-registering
|
||||
// proxy, not by an end user; the stale peer may still be
|
||||
// marked Connected from its prior session, but its session is
|
||||
// dead by definition (its key no longer exists).
|
||||
if err := m.DeletePeers(ctx, accountID, staleIDs, "", false); err != nil {
|
||||
return fmt.Errorf("delete stale embedded proxy peers %v: %w", staleIDs, err)
|
||||
}
|
||||
}
|
||||
|
||||
name := fmt.Sprintf("proxy-%s", xid.New().String())
|
||||
newPeer := &peer.Peer{
|
||||
Ephemeral: true,
|
||||
@@ -275,29 +251,3 @@ func (m *managerImpl) CreateProxyPeer(ctx context.Context, accountID string, pee
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// findStaleEmbeddedProxyPeers returns the peer IDs of embedded proxy peer
|
||||
// records in accountID that target the same cluster but carry a different
|
||||
// WireGuard pubkey than the freshly-registering one. Used by CreateProxyPeer
|
||||
// to garbage-collect stale records left behind when the proxy restarts with a
|
||||
// regenerated keypair.
|
||||
func (m *managerImpl) findStaleEmbeddedProxyPeers(ctx context.Context, accountID, cluster, newKey string) ([]string, error) {
|
||||
account, err := m.store.GetAccount(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var stale []string
|
||||
for _, p := range account.Peers {
|
||||
if p == nil || !p.ProxyMeta.Embedded {
|
||||
continue
|
||||
}
|
||||
if p.ProxyMeta.Cluster != cluster {
|
||||
continue
|
||||
}
|
||||
if p.Key == newKey {
|
||||
continue
|
||||
}
|
||||
stale = append(stale, p.ID)
|
||||
}
|
||||
return stale, nil
|
||||
}
|
||||
|
||||
@@ -1,19 +1,24 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
// nolint:gosec
|
||||
_ "net/http/pprof"
|
||||
"os"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/cmd"
|
||||
)
|
||||
|
||||
func main() {
|
||||
go func() {
|
||||
log.Println(http.ListenAndServe("localhost:6060", nil))
|
||||
}()
|
||||
if pprofAddr := os.Getenv("NB_PPROF_ADDR"); pprofAddr != "" {
|
||||
log.Infof("pprof enabled, listening on: %s", pprofAddr)
|
||||
go func() {
|
||||
log.Println(http.ListenAndServe(pprofAddr, nil))
|
||||
}()
|
||||
}
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
@@ -4269,7 +4269,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,
|
||||
|
||||
@@ -1,199 +0,0 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/peers"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
agenttypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/permissions"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
// TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale is the no-mock
|
||||
// regression guard for the bug the user reported: restarting the proxy creates
|
||||
// a fresh embedded peer with a NEW WireGuard public key (the proxy generates
|
||||
// the keypair on every startup at proxy/internal/roundtrip/netbird.go:312).
|
||||
// The PRIOR embedded peer record is never deleted on management, so the
|
||||
// account accumulates a stale peer holding a stale CGNAT IP. Other peers
|
||||
// in the account either keep routing to the dead IP, or — if synth DNS
|
||||
// picks the wrong record — never see the new IP at all.
|
||||
//
|
||||
// What this test exercises (no mocks):
|
||||
// - real SQLite test store
|
||||
// - real DefaultAccountManager, network-map controller, peer-update channels
|
||||
// - real peers.Manager.CreateProxyPeer path (the very method the proxy
|
||||
// invokes over gRPC on every startup)
|
||||
// - real agentnetwork.Manager + synth chain so the client receives a
|
||||
// concrete DNS record that must point at the LATEST proxy peer.
|
||||
//
|
||||
// Pre-fix expected behavior (red): two embedded peers exist after the
|
||||
// "restart"; the synth DNS record points at the stale one; the client
|
||||
// receives an update reflecting the new peer but the old one lingers.
|
||||
// Post-fix expected behavior (green): exactly one embedded peer exists
|
||||
// after restart (with the new key) AND the client's network map carries
|
||||
// the synth DNS pointing at that new peer's CGNAT IP.
|
||||
func TestAgentNetwork_ProxyRestart_PropagatesNewPeerAndDropsStale(t *testing.T) {
|
||||
am, updateManager, err := createManager(t)
|
||||
require.NoError(t, err, "createManager must succeed")
|
||||
ctx := context.Background()
|
||||
|
||||
const (
|
||||
accountID = "an-restart-acct"
|
||||
adminUserID = "an-restart-admin"
|
||||
groupAID = "an-restart-grp-A"
|
||||
clusterAddr = "eu.proxy.netbird.io"
|
||||
clientKey = "BhRPtynAAYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
|
||||
// Two different proxy pubkeys — the "before" and "after" of a
|
||||
// proxy-process restart with fresh-keypair generation.
|
||||
proxyKey1 = "Aaaaa1aaaaYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
|
||||
proxyKey2 = "Bbbbb2bbbbYRDy08+q4HTMsos8fs4plTP4NOSh7C1ry8="
|
||||
)
|
||||
|
||||
// --- Account scaffold ---
|
||||
account := newAccountWithId(ctx, accountID, adminUserID, "an-restart.test", "", "", false)
|
||||
require.NoError(t, am.Store.SaveAccount(ctx, account))
|
||||
|
||||
clientPeer := &nbpeer.Peer{
|
||||
Key: clientKey,
|
||||
Name: "an-restart-client",
|
||||
DNSLabel: "an-restart-client",
|
||||
Meta: nbpeer.PeerSystemMeta{Hostname: "an-restart-client", GoOS: "linux", WtVersion: "development"},
|
||||
}
|
||||
addedClient, _, _, _, err := am.AddPeer(ctx, "", "", adminUserID, clientPeer, false)
|
||||
require.NoError(t, err, "AddPeer for client must succeed")
|
||||
require.NoError(t, am.MarkPeerConnected(ctx, clientKey, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
|
||||
"MarkPeerConnected for the client peer must succeed (affected-peer fan-out skips disconnected peers)")
|
||||
|
||||
// Place the client in group A so the synth policy reaches it.
|
||||
account, err = am.Store.GetAccount(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
account.Groups[groupAID] = &types.Group{ID: groupAID, Name: "groupA", Peers: []string{addedClient.ID}}
|
||||
require.NoError(t, am.Store.SaveAccount(ctx, account), "SaveAccount must persist group A")
|
||||
|
||||
// --- Real peers + agent-network managers ---
|
||||
permMgr := permissions.NewManager(am.Store)
|
||||
peersMgr := peers.NewManager(am.Store, permMgr)
|
||||
peersMgr.SetAccountManager(am)
|
||||
peersMgr.SetNetworkMapController(am.networkMapController)
|
||||
agentMgr := agentnetwork.NewManager(am.Store, permMgr, am, nil)
|
||||
|
||||
// Subscribe BEFORE any state-mutating call so we don't lose the update
|
||||
// that contains the synth DNS record.
|
||||
clientCh := updateManager.CreateChannel(ctx, addedClient.ID)
|
||||
t.Cleanup(func() { updateManager.CloseChannel(ctx, addedClient.ID) })
|
||||
drain(clientCh)
|
||||
|
||||
// --- First proxy startup: register peer key K1, then mark it
|
||||
// connected. In production the proxy follows CreateProxyPeer with the
|
||||
// regular sync stream which lands on MarkPeerConnected; the synth DNS
|
||||
// path filters out peers that aren't Connected (types/account.go:323),
|
||||
// so without this step no DNS record would be emitted.
|
||||
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey1, clusterAddr),
|
||||
"first CreateProxyPeer (proxy startup) must succeed")
|
||||
|
||||
peer1ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
|
||||
require.NoError(t, err, "proxy peer for K1 must be persisted after CreateProxyPeer")
|
||||
require.NotEmpty(t, peer1ID)
|
||||
|
||||
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey1, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
|
||||
"MarkPeerConnected for K1 must succeed")
|
||||
|
||||
account, err = am.Store.GetAccount(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
proxyIP1 := account.Peers[peer1ID].IP.String()
|
||||
require.NotEmpty(t, proxyIP1, "K1 must have an assigned overlay IP")
|
||||
|
||||
// --- Provider + policy. CreateProvider / CreatePolicy trigger the
|
||||
// agentnetwork reconcile which runs UpdateAccountPeers; the resulting
|
||||
// NetworkMap delivered to the client carries the synth DNS record
|
||||
// pointing at K1's IP. ---
|
||||
provider, err := agentMgr.CreateProvider(ctx, adminUserID, &agenttypes.Provider{
|
||||
AccountID: accountID,
|
||||
ProviderID: "openai_api",
|
||||
Name: "openai-test",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test-key",
|
||||
Enabled: true,
|
||||
Models: []agenttypes.ProviderModel{{ID: "gpt-5.4"}},
|
||||
}, clusterAddr)
|
||||
require.NoError(t, err, "CreateProvider must succeed")
|
||||
|
||||
_, err = agentMgr.CreatePolicy(ctx, adminUserID, &agenttypes.Policy{
|
||||
AccountID: accountID,
|
||||
Name: "p1",
|
||||
Enabled: true,
|
||||
SourceGroups: []string{groupAID},
|
||||
DestinationProviderIDs: []string{provider.ID},
|
||||
})
|
||||
require.NoError(t, err, "CreatePolicy must succeed")
|
||||
|
||||
settings, err := am.Store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
|
||||
require.NoError(t, err)
|
||||
fqdn := settings.Endpoint()
|
||||
|
||||
rdata1 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
|
||||
require.Equal(t, proxyIP1, rdata1,
|
||||
"client must receive a synth DNS record pointing at K1's overlay IP after the synth path runs")
|
||||
drain(clientCh)
|
||||
|
||||
// --- Proxy restart: NEW keypair K2, same account, same cluster ---
|
||||
require.NoError(t, peersMgr.CreateProxyPeer(ctx, accountID, proxyKey2, clusterAddr),
|
||||
"second CreateProxyPeer (proxy restart with fresh keypair) must succeed")
|
||||
|
||||
peer2ID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey2)
|
||||
require.NoError(t, err, "proxy peer for K2 must be persisted after restart")
|
||||
require.NotEmpty(t, peer2ID)
|
||||
|
||||
require.NoError(t, am.MarkPeerConnected(ctx, proxyKey2, accountID, time.Now().UnixNano(), &types.NetworkMap{}),
|
||||
"MarkPeerConnected for K2 must succeed")
|
||||
|
||||
// In production the agent's sync stream pulls a fresh NetworkMap as
|
||||
// part of its normal reconcile cadence; in this isolated test
|
||||
// MarkPeerConnected's affected-peer fan-out can race the channel-side
|
||||
// buffer in a way that swallows the synth-DNS-bearing update before
|
||||
// our await reads it. Trigger an explicit account-wide fan-out so the
|
||||
// assertion below tests what production actually delivers, not the
|
||||
// in-test buffer race.
|
||||
am.UpdateAccountPeers(ctx, accountID, types.UpdateReason{Resource: types.UpdateResourcePeer, Operation: types.UpdateOperationUpdate})
|
||||
|
||||
account, err = am.Store.GetAccount(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
proxyIP2 := account.Peers[peer2ID].IP.String()
|
||||
require.NotEmpty(t, proxyIP2, "K2 must have an assigned overlay IP")
|
||||
require.NotEqual(t, proxyIP1, proxyIP2, "K2 must get a different overlay IP than K1 (sanity)")
|
||||
|
||||
// CRITICAL ASSERTION 1: K1 must no longer be in the store. The SqlStore
|
||||
// returns ("", nil) for a missing key rather than NotFound, so assert
|
||||
// on the returned ID being empty.
|
||||
staleID, err := am.Store.GetPeerIDByKey(ctx, store.LockingStrengthNone, proxyKey1)
|
||||
require.NoError(t, err, "GetPeerIDByKey for a missing peer must not error")
|
||||
assert.Empty(t, staleID,
|
||||
"stale embedded proxy peer K1 must be removed when a new embedded peer registers for the same (account, cluster); pre-fix this assertion fails because management never cleans up the prior peer record")
|
||||
|
||||
// CRITICAL ASSERTION 2: exactly one embedded proxy peer remains, and it
|
||||
// is K2.
|
||||
account, err = am.Store.GetAccount(ctx, accountID)
|
||||
require.NoError(t, err)
|
||||
embeddedKeys := []string{}
|
||||
for _, p := range account.Peers {
|
||||
if p.ProxyMeta.Embedded {
|
||||
embeddedKeys = append(embeddedKeys, p.Key)
|
||||
}
|
||||
}
|
||||
assert.Equal(t, []string{proxyKey2}, embeddedKeys,
|
||||
"after a proxy restart exactly one embedded proxy peer should remain — the one with the new key K2")
|
||||
|
||||
// CRITICAL ASSERTION 3: the synth DNS record the client receives now
|
||||
// points at K2's IP, not K1's.
|
||||
rdata2 := awaitZoneRData(clientCh, clusterAddr, fqdn, true)
|
||||
assert.Equal(t, proxyIP2, rdata2,
|
||||
"after proxy restart, the client's synth DNS record must point at the NEW embedded peer's IP, not the stale K1 IP")
|
||||
}
|
||||
@@ -76,6 +76,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 +178,7 @@ func (c *EmbeddedIdPConfig) ToYAMLConfig() (*dex.YAMLConfig, error) {
|
||||
},
|
||||
OAuth2: dex.OAuth2{
|
||||
SkipApprovalScreen: true,
|
||||
GrantTypes: c.GrantTypes,
|
||||
},
|
||||
Frontend: dex.Frontend{
|
||||
Issuer: "NetBird",
|
||||
|
||||
@@ -1606,7 +1606,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 +1633,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 +1660,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 +1753,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
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -296,6 +298,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")
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ssh/auth"
|
||||
@@ -42,6 +43,14 @@ type NetworkMapComponents struct {
|
||||
PostureFailedPeers map[string]map[string]struct{}
|
||||
|
||||
RouterPeers map[string]*nbpeer.Peer
|
||||
|
||||
routesByPeerOnce sync.Once
|
||||
routesByPeerIdx map[string][]routeIndexEntry
|
||||
}
|
||||
|
||||
type routeIndexEntry struct {
|
||||
route *route.Route
|
||||
viaGroup bool
|
||||
}
|
||||
|
||||
type AccountSettingsInfo struct {
|
||||
@@ -530,33 +539,43 @@ func (c *NetworkMapComponents) getRoutingPeerRoutes(peerID string) (enabledRoute
|
||||
disabledRoutes = append(disabledRoutes, r)
|
||||
}
|
||||
|
||||
for _, r := range c.Routes {
|
||||
for _, groupID := range r.PeerGroups {
|
||||
group := c.GetGroupInfo(groupID)
|
||||
if group == nil {
|
||||
continue
|
||||
}
|
||||
for _, id := range group.Peers {
|
||||
if id != peerID {
|
||||
continue
|
||||
}
|
||||
|
||||
newPeerRoute := r.Copy()
|
||||
newPeerRoute.Peer = id
|
||||
newPeerRoute.PeerGroups = nil
|
||||
newPeerRoute.ID = route.ID(string(r.ID) + ":" + id)
|
||||
takeRoute(newPeerRoute)
|
||||
break
|
||||
}
|
||||
}
|
||||
if r.Peer == peerID {
|
||||
takeRoute(r.Copy())
|
||||
for _, entry := range c.routesByPeer()[peerID] {
|
||||
if entry.viaGroup {
|
||||
newPeerRoute := entry.route.Copy()
|
||||
newPeerRoute.PeerGroups = nil
|
||||
newPeerRoute.ID = route.ID(string(entry.route.ID) + ":" + peerID)
|
||||
takeRoute(newPeerRoute)
|
||||
continue
|
||||
}
|
||||
takeRoute(entry.route.Copy())
|
||||
}
|
||||
|
||||
return enabledRoutes, disabledRoutes
|
||||
}
|
||||
|
||||
func (c *NetworkMapComponents) routesByPeer() map[string][]routeIndexEntry {
|
||||
c.routesByPeerOnce.Do(func() {
|
||||
idx := make(map[string][]routeIndexEntry)
|
||||
for _, r := range c.Routes {
|
||||
for _, groupID := range r.PeerGroups {
|
||||
group := c.GetGroupInfo(groupID)
|
||||
if group == nil {
|
||||
continue
|
||||
}
|
||||
for _, id := range group.Peers {
|
||||
idx[id] = append(idx[id], routeIndexEntry{route: r, viaGroup: true})
|
||||
}
|
||||
}
|
||||
if r.Peer != "" {
|
||||
idx[r.Peer] = append(idx[r.Peer], routeIndexEntry{route: r})
|
||||
}
|
||||
}
|
||||
c.routesByPeerIdx = idx
|
||||
})
|
||||
|
||||
return c.routesByPeerIdx
|
||||
}
|
||||
|
||||
func (c *NetworkMapComponents) filterRoutesByGroups(routes []*route.Route, groupListMap LookupMap) []*route.Route {
|
||||
var filteredRoutes []*route.Route
|
||||
for _, r := range routes {
|
||||
|
||||
@@ -25,6 +25,14 @@ const (
|
||||
denyCodeModel = "llm_policy.model_blocked"
|
||||
denyReasonModel = "model_blocked"
|
||||
denyMessageModel = "model is not in the policy allowlist"
|
||||
// Deny reason used when an allowlist is configured but the request model
|
||||
// could not be determined. URL/path-routed providers (AWS Bedrock, Google
|
||||
// Vertex, ...) carry the model outside the JSON body, so a request shape the
|
||||
// parser does not recognise reaches the guardrail with no model. Such a
|
||||
// request must be denied (fail closed), never waved through.
|
||||
denyCodeModelUnknown = "llm_policy.model_unknown"
|
||||
denyReasonModelUnknown = "model_unknown"
|
||||
denyMessageModelUnknown = "request model could not be determined for the policy allowlist"
|
||||
)
|
||||
|
||||
// Middleware enforces the model allowlist and optionally captures the
|
||||
@@ -108,23 +116,37 @@ func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middlew
|
||||
if len(m.cfg.ModelAllowlist) == 0 {
|
||||
return nil
|
||||
}
|
||||
if !modelPresent {
|
||||
return nil
|
||||
// Fail closed: with an allowlist configured, a request whose model the
|
||||
// upstream parser could not extract (absent or empty) must be denied rather
|
||||
// than allowed. This is what enforces the allowlist for URL/path-routed
|
||||
// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
|
||||
if !modelPresent || normaliseModel(model) == "" {
|
||||
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
}
|
||||
if m.modelInAllowlist(model) {
|
||||
return nil
|
||||
}
|
||||
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
|
||||
}
|
||||
|
||||
// denyModel builds a 403 deny Output for a model-allowlist rejection. model is
|
||||
// included in the details only when non-empty.
|
||||
func denyModel(model, code, message, reason string) *middleware.Output {
|
||||
details := map[string]string{}
|
||||
if model != "" {
|
||||
details["model"] = model
|
||||
}
|
||||
return &middleware.Output{
|
||||
Decision: middleware.DecisionDeny,
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Code: denyCodeModel,
|
||||
Message: denyMessageModel,
|
||||
Details: map[string]string{"model": model},
|
||||
Code: code,
|
||||
Message: message,
|
||||
Details: details,
|
||||
},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
|
||||
{Key: middleware.KeyLLMPolicyReason, Value: denyReasonModel},
|
||||
{Key: middleware.KeyLLMPolicyReason, Value: reason},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -102,13 +102,44 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllowlistMissingModelKeyAllows(t *testing.T) {
|
||||
func TestAllowlistMissingModelKeyDenies(t *testing.T) {
|
||||
// Fail closed: with an allowlist configured, a request whose model the
|
||||
// parser could not extract (URL/path-routed providers such as Bedrock or
|
||||
// Vertex whose shape wasn't recognised) must be denied, not allowed.
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "missing model key must allow even with non-empty allowlist")
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
|
||||
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
|
||||
require.NotNil(t, out.DenyReason, "deny reason must be populated")
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
|
||||
assert.Equal(t, "allow", dec, "decision must be allow when model key is absent")
|
||||
assert.Equal(t, "deny", dec, "decision must be deny when model key is absent")
|
||||
reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason)
|
||||
assert.Equal(t, "model_unknown", reason, "reason metadata must be model_unknown")
|
||||
}
|
||||
|
||||
func TestAllowlistEmptyModelValueDenies(t *testing.T) {
|
||||
// A present-but-empty model is as undeterminable as an absent one.
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: " "},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
|
||||
require.NotNil(t, out.DenyReason, "deny reason must be populated")
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
}
|
||||
|
||||
func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
|
||||
// Without an allowlist there is nothing to enforce, so a missing model is
|
||||
// still allowed — the fail-closed rule only applies when a list is set.
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInput())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
|
||||
}
|
||||
|
||||
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
package llm_request_parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware"
|
||||
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail"
|
||||
)
|
||||
|
||||
// runParserGuardrail runs the request parser then the model-allowlist guardrail
|
||||
// in SlotOnRequest order, threading the parser's metadata into the guardrail the
|
||||
// same way the real chain does. It returns the guardrail decision so tests can
|
||||
// assert allowlist enforcement for URL/path-routed providers end to end.
|
||||
func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []string) *middleware.Output {
|
||||
t.Helper()
|
||||
parser := newMiddleware(t)
|
||||
parsed, err := parser.Invoke(context.Background(), &middleware.Input{
|
||||
Slot: middleware.SlotOnRequest,
|
||||
URL: url,
|
||||
Body: body,
|
||||
})
|
||||
require.NoError(t, err, "parser must not error")
|
||||
|
||||
guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
|
||||
out, err := guard.Invoke(context.Background(), &middleware.Input{
|
||||
Slot: middleware.SlotOnRequest,
|
||||
Metadata: parsed.Metadata,
|
||||
})
|
||||
require.NoError(t, err, "guardrail must not error")
|
||||
require.NotNil(t, out, "guardrail must return an output")
|
||||
return out
|
||||
}
|
||||
|
||||
// TestModelAllowlist_URLRoutedProviders validates that the model allowlist is
|
||||
// enforced for providers whose model travels in the URL path (AWS Bedrock,
|
||||
// Google Vertex) rather than the JSON body. The "unknown action" case is the
|
||||
// regression guard for #6751: a Bedrock request shape the parser cannot map to a
|
||||
// model must fail closed under an allowlist instead of bypassing it.
|
||||
func TestModelAllowlist_URLRoutedProviders(t *testing.T) {
|
||||
const bedrockBody = `{"anthropic_version":"bedrock-2023-05-31","messages":[{"role":"user","content":"hi"}]}`
|
||||
const vertexBody = `{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
body string
|
||||
allowlist []string
|
||||
decision middleware.Decision
|
||||
denyCode string
|
||||
}{
|
||||
{
|
||||
name: "bedrock allowed model passes",
|
||||
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-haiku-4-5-v1:0/invoke",
|
||||
body: bedrockBody,
|
||||
allowlist: []string{"anthropic.claude-haiku-4-5"},
|
||||
decision: middleware.DecisionAllow,
|
||||
},
|
||||
{
|
||||
name: "bedrock disallowed model denied",
|
||||
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/invoke",
|
||||
body: bedrockBody,
|
||||
allowlist: []string{"anthropic.claude-haiku-4-5"},
|
||||
decision: middleware.DecisionDeny,
|
||||
denyCode: "llm_policy.model_blocked",
|
||||
},
|
||||
{
|
||||
name: "bedrock unknown action fails closed",
|
||||
url: "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-8-v1:0/some-future-action",
|
||||
body: bedrockBody,
|
||||
allowlist: []string{"anthropic.claude-haiku-4-5"},
|
||||
decision: middleware.DecisionDeny,
|
||||
denyCode: "llm_policy.model_unknown",
|
||||
},
|
||||
{
|
||||
name: "vertex disallowed model denied",
|
||||
url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-opus-4-8@20250101:rawPredict",
|
||||
body: vertexBody,
|
||||
allowlist: []string{"claude-haiku-4-5"},
|
||||
decision: middleware.DecisionDeny,
|
||||
denyCode: "llm_policy.model_blocked",
|
||||
},
|
||||
{
|
||||
name: "vertex allowed model passes",
|
||||
url: "/v1/projects/p/locations/global/publishers/anthropic/models/claude-haiku-4-5@20250101:rawPredict",
|
||||
body: vertexBody,
|
||||
allowlist: []string{"claude-haiku-4-5"},
|
||||
decision: middleware.DecisionAllow,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
out := runParserGuardrail(t, tt.url, []byte(tt.body), tt.allowlist)
|
||||
assert.Equal(t, tt.decision, out.Decision, "unexpected decision for %s", tt.name)
|
||||
if tt.decision == middleware.DecisionDeny {
|
||||
require.NotNil(t, out.DenyReason, "deny reason must be set for %s", tt.name)
|
||||
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403 for %s", tt.name)
|
||||
assert.Equal(t, tt.denyCode, out.DenyReason.Code, "deny code for %s", tt.name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
101
shared/testing_helpers/populate_fields.go
Normal file
101
shared/testing_helpers/populate_fields.go
Normal file
@@ -0,0 +1,101 @@
|
||||
package testing_helpers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"reflect"
|
||||
)
|
||||
|
||||
type PopulateFields struct {
|
||||
CustomFieldSetters map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error)
|
||||
TagsToSkip map[string]string
|
||||
}
|
||||
|
||||
func NewPopulateFields() *PopulateFields {
|
||||
return &PopulateFields{CustomFieldSetters: defaultCustomFieldSetters(), TagsToSkip: make(map[string]string)}
|
||||
}
|
||||
|
||||
func (p *PopulateFields) WithCustomFieldSetter(t reflect.Type, f func(this *PopulateFields, field reflect.Value) (int, error)) *PopulateFields {
|
||||
p.CustomFieldSetters[t] = f
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *PopulateFields) WithSkippedTag(tag, value string) *PopulateFields {
|
||||
p.TagsToSkip[tag] = value
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *PopulateFields) PopulateAll(v reflect.Value) (int, error) {
|
||||
typ := v.Type()
|
||||
totalExportedFields := 0
|
||||
for i := 0; i < typ.NumField(); i++ {
|
||||
f := typ.Field(i)
|
||||
if f.PkgPath != "" { // unexported
|
||||
continue
|
||||
}
|
||||
|
||||
if p.skippedTagPresent(f.Tag) {
|
||||
continue
|
||||
}
|
||||
|
||||
numOfExportedFields, err := p.setNonZero(v.Field(i))
|
||||
totalExportedFields += numOfExportedFields
|
||||
if err != nil {
|
||||
return totalExportedFields, err
|
||||
}
|
||||
}
|
||||
return totalExportedFields, nil
|
||||
}
|
||||
|
||||
// setNonZero assigns a deterministic non-zero value to a field based on its kind,
|
||||
// recursing into nested structs and populating one element of slice fields.
|
||||
func (p *PopulateFields) setNonZero(field reflect.Value) (int, error) {
|
||||
if f, ok := p.CustomFieldSetters[field.Type()]; ok {
|
||||
return f(p, field)
|
||||
}
|
||||
|
||||
switch field.Kind() {
|
||||
case reflect.String:
|
||||
field.SetString("non-zero")
|
||||
case reflect.Bool:
|
||||
field.SetBool(true)
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
field.SetInt(7)
|
||||
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
|
||||
field.SetUint(7)
|
||||
case reflect.Float32, reflect.Float64:
|
||||
field.SetFloat(7)
|
||||
case reflect.Struct:
|
||||
n, err := p.PopulateAll(field)
|
||||
return n + 1, err
|
||||
case reflect.Slice:
|
||||
s := reflect.MakeSlice(field.Type(), 1, 1)
|
||||
_, err := p.setNonZero(s.Index(0))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
field.Set(s)
|
||||
default:
|
||||
return 0, fmt.Errorf("unhandled field kind %s; extend setNonZero", field.Kind())
|
||||
}
|
||||
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
func defaultCustomFieldSetters() map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error) {
|
||||
return map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error){
|
||||
reflect.TypeOf(netip.Prefix{}): func(_ *PopulateFields, field reflect.Value) (int, error) {
|
||||
field.Set(reflect.ValueOf(netip.MustParsePrefix("10.0.0.0/24")))
|
||||
return 1, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PopulateFields) skippedTagPresent(t reflect.StructTag) bool {
|
||||
for tag, value := range p.TagsToSkip {
|
||||
if v := t.Get(tag); v == value {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
Reference in New Issue
Block a user