mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-19 14:11:28 +02:00
Compare commits
33 Commits
main
...
agent-netw
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
494b22848b | ||
|
|
a33fab103d | ||
|
|
3439500cc6 | ||
|
|
aa72df1654 | ||
|
|
da155bd60f | ||
|
|
52bed3ec26 | ||
|
|
378fe13257 | ||
|
|
520c912c07 | ||
|
|
9cd7e27027 | ||
|
|
0783605690 | ||
|
|
dc91325ac1 | ||
|
|
1796b2a1d8 | ||
|
|
dd87760f3d | ||
|
|
63d16a7cd4 | ||
|
|
fa691721d0 | ||
|
|
d46e2574a7 | ||
|
|
649a867cd3 | ||
|
|
a39b3c4af4 | ||
|
|
b1337f09d0 | ||
|
|
08e187ccb7 | ||
|
|
93cdf64a19 | ||
|
|
c04fb1c388 | ||
|
|
5353cab54f | ||
|
|
a64a417ea5 | ||
|
|
652c5c8b68 | ||
|
|
4d2b8b407b | ||
|
|
03e02c86ce | ||
|
|
789d416215 | ||
|
|
6415215126 | ||
|
|
1ae352a08d | ||
|
|
d928bcb630 | ||
|
|
47b2667653 | ||
|
|
875dda1708 |
@@ -1,6 +1,6 @@
|
||||
# NetBird Agent Guidelines
|
||||
|
||||
**NetBird** is an open source connectivity platform: a WireGuard®-based overlay
|
||||
**NetBird** is an open-source connectivity platform: a WireGuard®-based overlay
|
||||
network with a control plane. The **agent** (`client/`) runs on user machines as
|
||||
a privileged daemon and manages the WireGuard interface, routing, firewall, and
|
||||
DNS. **Management** (`management/`) is the control plane and REST/gRPC API,
|
||||
|
||||
@@ -479,7 +479,7 @@ go test -race ./client/internal/dns/...
|
||||
|
||||
## Checklist before submitting a PR
|
||||
|
||||
As a critical network service and open source project, we must enforce a few
|
||||
As a critical network service and open-source project, we must enforce a few
|
||||
things before submitting a pull request. The
|
||||
[pull request template](/.github/pull_request_template.md) mirrors this list —
|
||||
fill it in rather than deleting it.
|
||||
|
||||
@@ -130,7 +130,7 @@ In November 2022, NetBird joined the [StartUpSecure program](https://www.forschu
|
||||

|
||||
|
||||
### Acknowledgements
|
||||
We build on open source technologies like [WireGuard®](https://www.wireguard.com/), [Pion ICE](https://github.com/pion/ice), and [Rosenpass](https://rosenpass.eu). We greatly appreciate the work these projects are doing, and we'd love it if you could support them too (e.g., by starring or contributing).
|
||||
We build on open-source technologies like [WireGuard®](https://www.wireguard.com/), [Pion ICE](https://github.com/pion/ice), and [Rosenpass](https://rosenpass.eu). We greatly appreciate the work these projects are doing, and we'd love it if you could support them too (e.g., by starring or contributing).
|
||||
|
||||
### Legal
|
||||
This repository is licensed under the BSD-3-Clause license, which applies to all parts of the repository except for the directories management/, signal/ and relay/.
|
||||
|
||||
@@ -14,7 +14,7 @@ Report security issues one of these two ways:
|
||||
on this repository. This is the preferred route: it keeps the discussion, the draft advisory, and the credit in one place.
|
||||
- **Email** — `security@netbird.io`.
|
||||
|
||||
If the finding affects NetBird Cloud or our hosted infrastructure rather than the open source code, email us rather than
|
||||
If the finding affects NetBird Cloud or our hosted infrastructure rather than the open-source code, email us rather than
|
||||
filing a repository report.
|
||||
|
||||
### What to include
|
||||
|
||||
@@ -40,6 +40,35 @@ You can then use this private endpoint to configure your AI agents, whether that
|
||||
Full step-by-step setup:
|
||||
**https://docs.netbird.io/agent-network/quickstart**
|
||||
|
||||
## Client settings that don't follow the endpoint
|
||||
|
||||
Most of an agent's traffic follows the base URL you hand it, but a few
|
||||
client-side checks call their vendor directly and never reach the proxy. On a
|
||||
network that blocks direct egress they fail even though inference works, so
|
||||
they are worth setting once when you roll the endpoint out.
|
||||
|
||||
For Claude Code:
|
||||
|
||||
- **Fast mode** checks availability against `api.anthropic.com` rather than the
|
||||
configured base URL. Set `CLAUDE_CODE_SKIP_FAST_MODE_ORG_CHECK=1` when the
|
||||
agent authenticates with `ANTHROPIC_AUTH_TOKEN` alone (the usual shape when
|
||||
the proxy injects the real provider key) or when a TLS-inspecting proxy
|
||||
answers the check itself. Set
|
||||
`CLAUDE_CODE_SKIP_FAST_MODE_NETWORK_ERRORS=1` when the network refuses the
|
||||
connection outright. Fast mode is an Anthropic-API feature, so it is
|
||||
unavailable on a Bedrock- or Vertex-backed endpoint whatever you set.
|
||||
- **Model discovery** is off by default. Set
|
||||
`CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` for the picker to list the
|
||||
models your policies authorise; the proxy filters the response to that set.
|
||||
The client gives discovery a three-second budget and treats any redirect as
|
||||
a failure, so the endpoint must serve `/v1/models` directly.
|
||||
- **The WebFetch domain safety check** also calls `api.anthropic.com` directly
|
||||
and is unaffected by the variables above.
|
||||
|
||||
Allowing direct egress to `api.anthropic.com` covers the network cases but not
|
||||
the credential one, where the check reaches Anthropic and is rejected because
|
||||
the agent presents a proxy-issued key.
|
||||
|
||||
## Architecture
|
||||
|
||||
Agent Network is built on two existing NetBird capabilities:
|
||||
|
||||
@@ -26,8 +26,6 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager"
|
||||
"github.com/netbirdio/netbird/client/internal/stdnet"
|
||||
"github.com/netbirdio/netbird/client/net"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
"github.com/netbirdio/netbird/client/netsweep"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
"github.com/netbirdio/netbird/formatter"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
@@ -42,6 +40,11 @@ const (
|
||||
AnonymizeLevelStrict = nbAnonymize.LevelStrictString
|
||||
)
|
||||
|
||||
// ConnectionListener export internal Listener for mobile
|
||||
type ConnectionListener interface {
|
||||
peer.Listener
|
||||
}
|
||||
|
||||
// TunAdapter export internal TunAdapter for mobile
|
||||
type TunAdapter interface {
|
||||
device.TunAdapter
|
||||
@@ -82,13 +85,6 @@ type Client struct {
|
||||
deviceName string
|
||||
uiVersion string
|
||||
networkChangeListener listener.NetworkChangeListener
|
||||
// netState outlives engine restarts: it mirrors the OS connectivity, not
|
||||
// the engine lifecycle. Run and RunWithoutLogin inject it into each new
|
||||
// ConnectClient, which distributes it to every reconnection loop.
|
||||
netState *netstate.State
|
||||
|
||||
// sweeper also outlives engine restarts; NotifyNetworkChange sweeps it.
|
||||
sweeper *netsweep.Sweeper
|
||||
|
||||
stateMu sync.RWMutex
|
||||
connectClient *internal.ConnectClient
|
||||
@@ -152,7 +148,6 @@ func NewClient(androidSDKVersion int, deviceName string, uiVersion string, tunAd
|
||||
execWorkaround(androidSDKVersion)
|
||||
|
||||
net.SetAndroidProtectSocketFn(tunAdapter.ProtectSocket)
|
||||
system.SetIFaceDiscover(iFaceDiscover)
|
||||
return &Client{
|
||||
deviceName: deviceName,
|
||||
uiVersion: uiVersion,
|
||||
@@ -161,8 +156,6 @@ func NewClient(androidSDKVersion int, deviceName string, uiVersion string, tunAd
|
||||
recorder: peer.NewRecorder(""),
|
||||
ctxCancelLock: &sync.Mutex{},
|
||||
networkChangeListener: networkChangeListener,
|
||||
netState: netstate.New(),
|
||||
sweeper: netsweep.New(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,8 +196,7 @@ func (c *Client) Run(platformFiles PlatformFiles, urlOpener URLOpener, isAndroid
|
||||
}
|
||||
// todo do not throw error in case of cancelled context
|
||||
ctx = internal.CtxInitState(ctx)
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder,
|
||||
internal.WithNetworkState(c.netState), internal.WithSweeper(c.sweeper))
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder)
|
||||
c.setState(cfg, cacheDir, cfgFile, connectClient)
|
||||
// This path runs the interactive SSO flow, so reaching here means the peer
|
||||
// is authenticated again — release the latch Status() reports from. Clear
|
||||
@@ -245,8 +237,7 @@ func (c *Client) RunWithoutLogin(platformFiles PlatformFiles, dns *DNSList, dnsR
|
||||
|
||||
// todo do not throw error in case of cancelled context
|
||||
ctx = internal.CtxInitState(ctx)
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder,
|
||||
internal.WithNetworkState(c.netState), internal.WithSweeper(c.sweeper))
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder)
|
||||
c.setState(cfg, cacheDir, cfgFile, connectClient)
|
||||
return connectClient.RunOnAndroid(c.tunAdapter, c.iFaceDiscover, c.networkChangeListener, slices.Clone(dns.items), dnsReadyListener, stateFile, cacheDir)
|
||||
}
|
||||
@@ -294,24 +285,6 @@ func (c *Client) GetTunSettings() (*TunSettings, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
// SetNetworkAvailable feeds OS-reported network availability into the client.
|
||||
// While unavailable, the internal reconnect loops suspend their attempts and
|
||||
// the connection listener reports NoNetwork instead of Connecting; when
|
||||
// availability returns, the loops resume immediately with a fresh backoff.
|
||||
func (c *Client) SetNetworkAvailable(available bool) {
|
||||
c.netState.Set(available)
|
||||
c.recorder.SetNetworkAvailable(available)
|
||||
}
|
||||
|
||||
// NotifyNetworkChange marks the management, signal and relay connections
|
||||
// stale after the OS switched networks and schedules a sweep that cuts
|
||||
// whatever has not redialed on the new network by then. The engine and the
|
||||
// TUN device stay untouched.
|
||||
func (c *Client) NotifyNetworkChange() {
|
||||
c.sweeper.MarkNetworkChange()
|
||||
log.Infof("network change: connections marked stale")
|
||||
}
|
||||
|
||||
// DebugBundle generates a debug bundle, uploads it, and returns the upload key.
|
||||
// It works both with and without a running engine. anonymizeLevel is "default"
|
||||
// or "strict"; strict also anonymizes internal IP ranges, peer names, and
|
||||
@@ -552,11 +525,7 @@ func (c *Client) OnUpdatedHostDNS(list *DNSList) error {
|
||||
|
||||
// SetConnectionListener set the network connection listener
|
||||
func (c *Client) SetConnectionListener(listener ConnectionListener) {
|
||||
if listener == nil {
|
||||
c.recorder.RemoveConnectionListener()
|
||||
return
|
||||
}
|
||||
c.recorder.SetConnectionListener(connectionListenerAdapter{listener})
|
||||
c.recorder.SetConnectionListener(listener)
|
||||
}
|
||||
|
||||
// RemoveConnectionListener remove connection listener
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
import (
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
)
|
||||
|
||||
// Client state values delivered via ConnectionListener.OnStateChanged,
|
||||
// re-exported as basic constants so gomobile emits them into the generated
|
||||
// Java bindings. They mirror peer.ClientState*: append-only, never reorder.
|
||||
const (
|
||||
ClientStateDisconnected = int(peer.ClientStateDisconnected)
|
||||
ClientStateConnected = int(peer.ClientStateConnected)
|
||||
ClientStateConnecting = int(peer.ClientStateConnecting)
|
||||
ClientStateDisconnecting = int(peer.ClientStateDisconnecting)
|
||||
ClientStateNoNetwork = int(peer.ClientStateNoNetwork)
|
||||
)
|
||||
|
||||
// ConnectionListener export internal Listener for mobile. It mirrors
|
||||
// peer.Listener with OnStateChanged taking a plain int (one of the
|
||||
// ClientState* constants), because gomobile cannot bind named types.
|
||||
type ConnectionListener interface {
|
||||
OnStateChanged(state int)
|
||||
OnConnected()
|
||||
OnDisconnected()
|
||||
OnConnecting()
|
||||
OnDisconnecting()
|
||||
OnAddressChanged(string, string)
|
||||
OnPeersListChanged(int)
|
||||
}
|
||||
|
||||
// connectionListenerAdapter adapts the gomobile-facing ConnectionListener to
|
||||
// peer.Listener, converting the typed state to the int the binding carries.
|
||||
type connectionListenerAdapter struct {
|
||||
ConnectionListener
|
||||
}
|
||||
|
||||
func (a connectionListenerAdapter) OnStateChanged(state peer.ClientState) {
|
||||
a.ConnectionListener.OnStateChanged(int(state))
|
||||
}
|
||||
@@ -191,49 +191,40 @@ func (a *Auth) login(urlOpener URLOpener, isAndroidTV bool) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// loginHintSetter is implemented by both concrete flows (PKCE and device code)
|
||||
// but absent from the OAuthFlow interface, hence the assertion below — the same
|
||||
// way internal/auth wires it in authenticateWithPKCEFlow.
|
||||
type loginHintSetter interface {
|
||||
SetLoginHint(hint string)
|
||||
}
|
||||
|
||||
func (a *Auth) foregroundGetTokenInfo(authClient *auth.Auth, urlOpener URLOpener, isAndroidTV bool) (*auth.TokenInfo, error) {
|
||||
oAuthFlow, err := authClient.GetOAuthFlow(a.ctx, isAndroidTV, profileLoginHint(a.cfgPath))
|
||||
oAuthFlow, err := authClient.GetOAuthFlow(a.ctx, isAndroidTV)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get OAuth flow: %v", err)
|
||||
}
|
||||
|
||||
return runOAuthFlow(a.ctx, oAuthFlow, urlOpener, nil)
|
||||
}
|
||||
|
||||
// profileLoginHint returns the stored account email for the profile at cfgPath.
|
||||
// An empty hint is deliberate, not a fallback: a fresh profile leaves the
|
||||
// choice to the IdP. Switching accounts is done by switching or removing
|
||||
// profiles, not by logging out — logout keeps the email.
|
||||
func profileLoginHint(cfgPath string) string {
|
||||
if cfgPath == "" {
|
||||
return ""
|
||||
// An empty hint is deliberate, not a fallback: a fresh profile leaves the
|
||||
// choice to the IdP. Switching accounts is done by switching or removing
|
||||
// profiles, not by logging out — logout keeps the email.
|
||||
if a.cfgPath != "" {
|
||||
if hint := readProfileEmail(a.cfgPath); hint != "" {
|
||||
if setter, ok := oAuthFlow.(loginHintSetter); ok {
|
||||
setter.SetLoginHint(hint)
|
||||
}
|
||||
}
|
||||
}
|
||||
return readProfileEmail(cfgPath)
|
||||
}
|
||||
|
||||
// runOAuthFlow drives an already acquired OAuth flow to a token: requests the
|
||||
// flow info, presents the verification URL through the opener and waits for
|
||||
// the browser round-trip. Open is called synchronously — it is what marks the
|
||||
// surface as opened on the client side, and a fast token's OnLoginSuccess is
|
||||
// a no-op until it has, so the dismissal would be dropped rather than
|
||||
// delayed. Openers must therefore not block: they post their UI work and
|
||||
// return. onWaiting, when set, runs after the URL is shown, right before the
|
||||
// blocking wait.
|
||||
func runOAuthFlow(ctx context.Context, flow auth.OAuthFlow, urlOpener URLOpener, onWaiting func()) (*auth.TokenInfo, error) {
|
||||
flowInfo, err := flow.RequestAuthInfo(ctx)
|
||||
flowInfo, err := oAuthFlow.RequestAuthInfo(context.TODO())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request auth info: %w", err)
|
||||
return nil, fmt.Errorf("getting a request OAuth flow info failed: %v", err)
|
||||
}
|
||||
|
||||
urlOpener.Open(flowInfo.VerificationURIComplete, flowInfo.UserCode)
|
||||
go urlOpener.Open(flowInfo.VerificationURIComplete, flowInfo.UserCode)
|
||||
|
||||
if onWaiting != nil {
|
||||
onWaiting()
|
||||
}
|
||||
|
||||
tokenInfo, err := flow.WaitToken(ctx, flowInfo)
|
||||
tokenInfo, err := oAuthFlow.WaitToken(a.ctx, flowInfo)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("wait for token: %w", err)
|
||||
return nil, fmt.Errorf("waiting for browser login failed: %v", err)
|
||||
}
|
||||
|
||||
return &tokenInfo, nil
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
)
|
||||
|
||||
type prefsStore interface {
|
||||
Get(namespace string, v any) (bool, error)
|
||||
Put(namespace string, v any) error
|
||||
}
|
||||
|
||||
type profilePrefs struct {
|
||||
prefs *profilemanager.Prefs
|
||||
}
|
||||
|
||||
func newProfilePrefs(configDir, profileID string) (*profilePrefs, error) {
|
||||
if configDir == "" || profileID == "" {
|
||||
return nil, fmt.Errorf("profile prefs require a config dir and profile ID")
|
||||
}
|
||||
pm := NewProfileManager(configDir)
|
||||
prefs, err := pm.serviceMgr.ProfilePrefs(profilemanager.ID(profileID), androidUsername)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolve profile prefs: %w", err)
|
||||
}
|
||||
return &profilePrefs{prefs: prefs}, nil
|
||||
}
|
||||
|
||||
func (p *profilePrefs) Get(namespace string, v any) (bool, error) {
|
||||
return p.prefs.Get(namespace, v)
|
||||
}
|
||||
|
||||
func (p *profilePrefs) Put(namespace string, v any) error {
|
||||
return p.prefs.Put(namespace, v)
|
||||
}
|
||||
@@ -1,649 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
gossh "golang.org/x/crypto/ssh"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal"
|
||||
"github.com/netbirdio/netbird/client/internal/auth"
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
nbssh "github.com/netbirdio/netbird/client/ssh"
|
||||
"github.com/netbirdio/netbird/client/ssh/detection"
|
||||
)
|
||||
|
||||
const (
|
||||
sshDialTimeout = 30 * time.Second
|
||||
sshDetectionTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
// PasswordRequiredMarker tells Java to prompt for a password and retry. It is
|
||||
// a string because gomobile flattens errors to their message, so a sentinel
|
||||
// value would not survive the binding.
|
||||
const PasswordRequiredMarker = "netbird-ssh-password-required"
|
||||
|
||||
// HostKeyUnknownMarker tells Java to show the fingerprint and, on confirmation,
|
||||
// retry with TrustHostKey set. The presented fingerprint is appended after the
|
||||
// marker so the prompt can display it and the retry can guard against a key
|
||||
// that changed between the two connects. Only regular (non-NetBird) servers
|
||||
// reach this: NetBird peers verify against the registry.
|
||||
const HostKeyUnknownMarker = "netbird-ssh-hostkey-unknown"
|
||||
|
||||
var (
|
||||
errPasswordRequired = errors.New(PasswordRequiredMarker)
|
||||
errClientClosed = errors.New("ssh client closed")
|
||||
)
|
||||
|
||||
// errHostKeyUnknown carries the presented fingerprint so Connect can build the
|
||||
// marker message the Java side parses.
|
||||
type errHostKeyUnknown struct {
|
||||
fingerprint string
|
||||
}
|
||||
|
||||
func (e *errHostKeyUnknown) Error() string {
|
||||
return HostKeyUnknownMarker + ":" + e.fingerprint
|
||||
}
|
||||
|
||||
// SSHTerminalListener receives SSH session events. It is implemented in Java.
|
||||
//
|
||||
// All callbacks are invoked from goroutines and may run concurrently with each
|
||||
// other; the implementation must be safe to call from any thread.
|
||||
type SSHTerminalListener interface {
|
||||
OnConnected()
|
||||
OnData(data []byte)
|
||||
OnClose(reason string)
|
||||
OnError(message string)
|
||||
}
|
||||
|
||||
// SSHClient is a NetBird-aware SSH client exposed to Java via gomobile.
|
||||
//
|
||||
// It dials through the running NetBird tunnel and runs a standard SSH session
|
||||
// on top with PTY enabled. Host-key verification uses the NetBird-provided
|
||||
// peer SSH host keys, identical to the desktop client.
|
||||
type SSHClient struct {
|
||||
nb *Client
|
||||
mu sync.Mutex
|
||||
listener SSHTerminalListener
|
||||
urlOpener URLOpener
|
||||
|
||||
sshClient *gossh.Client
|
||||
session *gossh.Session
|
||||
stdin io.WriteCloser
|
||||
closed bool
|
||||
|
||||
// gen identifies the current connection attempt. Connect and Close bump it,
|
||||
// so an in-flight dial or a reader left over from a previous connection
|
||||
// finds itself stale and stays silent instead of publishing OnConnected or
|
||||
// OnClose for a connection the caller already abandoned.
|
||||
gen uint64
|
||||
dialCancel context.CancelFunc
|
||||
|
||||
// knownHostsConfigDir and knownHostsProfile locate the TOFU store for
|
||||
// regular SSH servers in the profile's preferences. Java supplies them,
|
||||
// since an overlay IP is a different host under a different profile. Empty
|
||||
// until set: without them a regular server cannot be verified and Connect
|
||||
// refuses one.
|
||||
knownHostsConfigDir string
|
||||
knownHostsProfile string
|
||||
// trustHostKey carries the fingerprint the user confirmed on a previous
|
||||
// attempt, so the retry accepts exactly that key and persists it.
|
||||
trustHostKey string
|
||||
}
|
||||
|
||||
// NewSSHClient creates a new SSH client bound to the running NetBird Client.
|
||||
func NewSSHClient(c *Client) *SSHClient {
|
||||
return &SSHClient{nb: c}
|
||||
}
|
||||
|
||||
// SetListener registers the Java listener. Must be called before Connect to
|
||||
// receive any events.
|
||||
func (s *SSHClient) SetListener(l SSHTerminalListener) {
|
||||
s.mu.Lock()
|
||||
s.listener = l
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetURLOpener registers the Java URL opener used to display the device-code
|
||||
// authorization page in a Custom Tabs window when the target peer requires
|
||||
// JWT authentication. Must be set before Connect to be effective.
|
||||
func (s *SSHClient) SetURLOpener(opener URLOpener) {
|
||||
s.mu.Lock()
|
||||
s.urlOpener = opener
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetKnownHostsStore points the TOFU host-key store at a profile's preferences.
|
||||
// Must be set before connecting to a regular SSH server; without it such a
|
||||
// server cannot be verified and Connect refuses one.
|
||||
func (s *SSHClient) SetKnownHostsStore(configDir, profileID string) {
|
||||
s.mu.Lock()
|
||||
s.knownHostsConfigDir = configDir
|
||||
s.knownHostsProfile = profileID
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// TrustHostKey records the fingerprint the user confirmed for a regular server,
|
||||
// so the next Connect accepts that exact key and adds it to the known-hosts
|
||||
// store. Passing a fingerprint that no longer matches makes the connect fail
|
||||
// rather than trust a key that changed since the prompt.
|
||||
func (s *SSHClient) TrustHostKey(fingerprint string) {
|
||||
s.mu.Lock()
|
||||
s.trustHostKey = fingerprint
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Connect dials the SSH server through the NetBird tunnel and performs the
|
||||
// SSH handshake. It auto-detects the server type via SSH banner inspection
|
||||
// and selects the appropriate authentication path:
|
||||
//
|
||||
// - NetBird-SSH server requiring JWT: launches the OAuth 2.0 device-code
|
||||
// flow, opens the verification URL through the registered URLOpener, and
|
||||
// uses the resulting token as the SSH password. Host-key verification
|
||||
// uses the NetBird peer registry.
|
||||
// - NetBird-SSH server without JWT: authenticates with the NetBird SSH
|
||||
// private key. Host-key verification uses the NetBird peer registry.
|
||||
// - Regular SSH server (e.g. OpenSSH): authenticates with the NetBird key
|
||||
// first (so a user-installed NetBird public key works), then falls back
|
||||
// to the supplied password if non-empty. Host-key verification is
|
||||
// trust-on-first-use against the per-profile known-hosts store.
|
||||
//
|
||||
// The password parameter is only consulted for regular SSH servers.
|
||||
func (s *SSHClient) Connect(host string, port int, user, password string) error {
|
||||
if port < 1 || port > 65535 {
|
||||
return fmt.Errorf("invalid port: %d", port)
|
||||
}
|
||||
|
||||
cfg, cfgPath, cc := s.nb.authSnapshot()
|
||||
if cc == nil {
|
||||
return errors.New("netbird client not running")
|
||||
}
|
||||
if cfg == nil {
|
||||
return errors.New("netbird config not loaded")
|
||||
}
|
||||
engine := cc.Engine()
|
||||
if engine == nil {
|
||||
return errors.New("netbird engine not available")
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
s.gen++
|
||||
gen := s.gen
|
||||
s.mu.Unlock()
|
||||
|
||||
serverType := detectServerType(host, port)
|
||||
log.Debugf("SSH server type: %s", serverType)
|
||||
|
||||
authMethods, hostKeyCallback, err := s.buildAuth(cfg, cfgPath, engine, serverType, password)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
clientConfig := &gossh.ClientConfig{
|
||||
User: user,
|
||||
Auth: authMethods,
|
||||
HostKeyCallback: hostKeyCallback,
|
||||
Timeout: sshDialTimeout,
|
||||
}
|
||||
err = s.dialAndHandshake(gen, host, port, clientConfig)
|
||||
|
||||
// An unknown host key is a prompt, not a failure: return the marker intact
|
||||
// (rootCause would unwrap it) so Java can show the fingerprint and retry.
|
||||
var unknownHost *errHostKeyUnknown
|
||||
if errors.As(err, &unknownHost) {
|
||||
return errors.New(unknownHost.Error())
|
||||
}
|
||||
|
||||
// A regular server may still accept a password, so let the caller ask for
|
||||
// one instead of failing. NetBird servers never use a password, so a
|
||||
// failure there is genuine.
|
||||
if err != nil && serverType != detection.ServerTypeNetBirdJWT &&
|
||||
serverType != detection.ServerTypeNetBirdNoJWT && isAuthFailure(err) &&
|
||||
passwordCouldHelp(err, password != "") {
|
||||
return errPasswordRequired
|
||||
}
|
||||
if err != nil {
|
||||
return rootCause(err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StartSession requests a PTY and starts an interactive shell. Output from
|
||||
// the session is forwarded to the listener via OnData.
|
||||
func (s *SSHClient) StartSession(cols, rows int) error {
|
||||
err := s.startSession(cols, rows)
|
||||
if err != nil {
|
||||
log.Infof("SSH: start session failed: %v", err)
|
||||
return rootCause(err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Write sends data to the SSH session stdin.
|
||||
func (s *SSHClient) Write(data []byte) error {
|
||||
s.mu.Lock()
|
||||
stdin := s.stdin
|
||||
s.mu.Unlock()
|
||||
if stdin == nil {
|
||||
return errors.New("ssh session not started")
|
||||
}
|
||||
if _, err := stdin.Write(data); err != nil {
|
||||
return fmt.Errorf("write stdin: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resize updates the PTY window size.
|
||||
func (s *SSHClient) Resize(cols, rows int) error {
|
||||
s.mu.Lock()
|
||||
session := s.session
|
||||
s.mu.Unlock()
|
||||
if session == nil {
|
||||
return errors.New("ssh session not started")
|
||||
}
|
||||
return session.WindowChange(rows, cols)
|
||||
}
|
||||
|
||||
// Reset makes a closed client usable for another Connect: Close leaves the
|
||||
// one-shot guard set, and clearing it lets the same client back a reconnect.
|
||||
func (s *SSHClient) Reset() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.closed = false
|
||||
}
|
||||
|
||||
// Close terminates the SSH session and underlying connection. Safe to call
|
||||
// multiple times.
|
||||
func (s *SSHClient) Close() error {
|
||||
s.mu.Lock()
|
||||
s.gen++
|
||||
if s.dialCancel != nil {
|
||||
s.dialCancel()
|
||||
s.dialCancel = nil
|
||||
}
|
||||
sshClient := s.sshClient
|
||||
session := s.session
|
||||
stdin := s.stdin
|
||||
s.sshClient = nil
|
||||
s.session = nil
|
||||
s.stdin = nil
|
||||
notify := !s.closed
|
||||
s.closed = true
|
||||
listener := s.listener
|
||||
s.mu.Unlock()
|
||||
|
||||
if stdin != nil {
|
||||
if err := stdin.Close(); err != nil {
|
||||
log.Debugf("ssh: stdin close: %v", err)
|
||||
}
|
||||
}
|
||||
if session != nil {
|
||||
if err := session.Close(); err != nil && !errors.Is(err, io.EOF) {
|
||||
log.Debugf("ssh: session close: %v", err)
|
||||
}
|
||||
}
|
||||
var firstErr error
|
||||
if sshClient != nil {
|
||||
if err := sshClient.Close(); err != nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
if notify && listener != nil {
|
||||
listener.OnClose("closed by client")
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (s *SSHClient) startSession(cols, rows int) error {
|
||||
log.Debugf("SSH: starting session %dx%d", cols, rows)
|
||||
s.mu.Lock()
|
||||
sshClient := s.sshClient
|
||||
gen := s.gen
|
||||
s.mu.Unlock()
|
||||
|
||||
if sshClient == nil {
|
||||
return errors.New("ssh client not connected")
|
||||
}
|
||||
|
||||
pty, err := nbssh.StartPTYSession(sshClient, cols, rows)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if gen != s.gen {
|
||||
s.mu.Unlock()
|
||||
closeQuiet(pty.Session, "stale session")
|
||||
return errClientClosed
|
||||
}
|
||||
s.session = pty.Session
|
||||
s.stdin = pty.Stdin
|
||||
s.mu.Unlock()
|
||||
|
||||
readerDone := make(chan string, 2)
|
||||
go func() { readerDone <- s.readLoop(pty.Stdout, "stdout") }()
|
||||
go func() { readerDone <- s.readLoop(pty.Stderr, "stderr") }()
|
||||
go func() {
|
||||
reason := <-readerDone
|
||||
if second := <-readerDone; reason == "" {
|
||||
reason = second
|
||||
}
|
||||
s.notifyClose(gen, reason)
|
||||
}()
|
||||
log.Debug("SSH: session started, shell running")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SSHClient) buildAuth(cfg *profilemanager.Config, cfgPath string, engine *internal.Engine,
|
||||
serverType detection.ServerType, password string) ([]gossh.AuthMethod, gossh.HostKeyCallback, error) {
|
||||
|
||||
switch serverType {
|
||||
case detection.ServerTypeNetBirdJWT:
|
||||
token, err := s.requestJWTToken(cfg, cfgPath)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("jwt: %w", err)
|
||||
}
|
||||
auths := []gossh.AuthMethod{gossh.Password(token)}
|
||||
return auths, nbssh.CreateHostKeyCallback(nbssh.PeerKeyLookup(engine.GetPeerSSHKey)), nil
|
||||
|
||||
case detection.ServerTypeNetBirdNoJWT:
|
||||
if cfg.SSHKey == "" {
|
||||
return nil, nil, errors.New("no NetBird SSH key available")
|
||||
}
|
||||
signer, err := gossh.ParsePrivateKey([]byte(cfg.SSHKey))
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("parse netbird ssh key: %w", err)
|
||||
}
|
||||
auths := []gossh.AuthMethod{gossh.PublicKeys(signer)}
|
||||
return auths, nbssh.CreateHostKeyCallback(nbssh.PeerKeyLookup(engine.GetPeerSSHKey)), nil
|
||||
|
||||
case detection.ServerTypeRegular:
|
||||
var auths []gossh.AuthMethod
|
||||
if cfg.SSHKey != "" {
|
||||
if signer, err := gossh.ParsePrivateKey([]byte(cfg.SSHKey)); err == nil {
|
||||
auths = append(auths, gossh.PublicKeys(signer))
|
||||
} else {
|
||||
log.Debugf("ssh: parse netbird key for regular auth: %v", err)
|
||||
}
|
||||
}
|
||||
if password != "" {
|
||||
pw := password
|
||||
auths = append(auths, gossh.Password(pw))
|
||||
auths = append(auths, gossh.KeyboardInteractive(func(_, _ string, questions []string, _ []bool) ([]string, error) {
|
||||
answers := make([]string, len(questions))
|
||||
for i := range questions {
|
||||
answers[i] = pw
|
||||
}
|
||||
return answers, nil
|
||||
}))
|
||||
}
|
||||
if len(auths) == 0 {
|
||||
// Nothing to offer at all: ask for a password rather than failing,
|
||||
// so the caller can retry once the user supplies one.
|
||||
return nil, nil, errPasswordRequired
|
||||
}
|
||||
callback, err := s.tofuHostKeyCallback()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return auths, callback, nil
|
||||
|
||||
default:
|
||||
return nil, nil, fmt.Errorf("unsupported SSH server type: %v", serverType)
|
||||
}
|
||||
}
|
||||
|
||||
// tofuHostKeyCallback verifies a regular server's host key against the
|
||||
// per-profile known-hosts store. An unknown host returns errHostKeyUnknown so
|
||||
// Java can show the fingerprint and, once confirmed, retry with the key
|
||||
// trusted; a changed key is rejected outright, as OpenSSH does. When the user
|
||||
// has confirmed a fingerprint, the callback accepts exactly that key and
|
||||
// appends it to the store.
|
||||
func (s *SSHClient) tofuHostKeyCallback() (gossh.HostKeyCallback, error) {
|
||||
s.mu.Lock()
|
||||
configDir := s.knownHostsConfigDir
|
||||
profileID := s.knownHostsProfile
|
||||
trusted := s.trustHostKey
|
||||
s.mu.Unlock()
|
||||
|
||||
if configDir == "" || profileID == "" {
|
||||
return nil, errors.New("no known-hosts store configured for regular SSH")
|
||||
}
|
||||
|
||||
store, err := openKnownHostsStore(configDir, profileID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load known-hosts store: %w", err)
|
||||
}
|
||||
|
||||
return func(hostname string, remote net.Addr, key gossh.PublicKey) error {
|
||||
verdict, err := store.verify(hostname, remote, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if verdict == hostKeyMatched {
|
||||
return nil
|
||||
}
|
||||
if verdict == hostKeyChanged {
|
||||
return fmt.Errorf("SSH host key changed for %s (possible attack)", hostname)
|
||||
}
|
||||
|
||||
fingerprint := gossh.FingerprintSHA256(key)
|
||||
if trusted == "" {
|
||||
return &errHostKeyUnknown{fingerprint: fingerprint}
|
||||
}
|
||||
if trusted != fingerprint {
|
||||
return fmt.Errorf("SSH host key changed since it was confirmed for %s", hostname)
|
||||
}
|
||||
if err := store.append(hostname, remote, key); err != nil {
|
||||
return fmt.Errorf("persist trusted host key: %w", err)
|
||||
}
|
||||
// The confirmation is spent: now that the key is stored, a later
|
||||
// reconnect must verify against the file, not re-accept this fingerprint.
|
||||
s.mu.Lock()
|
||||
s.trustHostKey = ""
|
||||
s.mu.Unlock()
|
||||
return nil
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *SSHClient) requestJWTToken(cfg *profilemanager.Config, cfgPath string) (string, error) {
|
||||
s.mu.Lock()
|
||||
urlOpener := s.urlOpener
|
||||
s.mu.Unlock()
|
||||
if urlOpener == nil {
|
||||
return "", errors.New("URL opener not configured for JWT auth")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
flow, err := auth.NewOAuthFlow(ctx, cfg, false, true, profileLoginHint(cfgPath))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create oauth flow: %w", err)
|
||||
}
|
||||
|
||||
// The status callback covers the browser round-trip, which would
|
||||
// otherwise leave the terminal blank.
|
||||
tokenInfo, err := runOAuthFlow(ctx, flow, urlOpener, func() {
|
||||
s.notifyStatus("Waiting for browser authentication...")
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
token := tokenInfo.GetTokenToUse()
|
||||
if token == "" {
|
||||
return "", errors.New("empty token returned by IdP")
|
||||
}
|
||||
|
||||
// Tells the client the browser round-trip is over so it can dismiss the
|
||||
// surface it opened, the same way the login and session-extend flows do.
|
||||
// Without it the Custom Tab stays in front of the terminal even though the
|
||||
// token has already been collected.
|
||||
urlOpener.OnLoginSuccess()
|
||||
|
||||
return token, nil
|
||||
}
|
||||
|
||||
func (s *SSHClient) dialAndHandshake(gen uint64, host string, port int, clientConfig *gossh.ClientConfig) error {
|
||||
addr := net.JoinHostPort(host, strconv.Itoa(port))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), sshDialTimeout)
|
||||
defer cancel()
|
||||
|
||||
s.mu.Lock()
|
||||
if gen != s.gen {
|
||||
s.mu.Unlock()
|
||||
return errClientClosed
|
||||
}
|
||||
s.dialCancel = cancel
|
||||
s.mu.Unlock()
|
||||
|
||||
var dialer net.Dialer
|
||||
conn, err := dialer.DialContext(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("dial %s: %w", addr, err)
|
||||
}
|
||||
|
||||
client, err := nbssh.Handshake(ctx, conn, addr, clientConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if gen != s.gen {
|
||||
s.mu.Unlock()
|
||||
closeQuiet(client, "stale ssh client")
|
||||
return errClientClosed
|
||||
}
|
||||
s.sshClient = client
|
||||
listener := s.listener
|
||||
s.mu.Unlock()
|
||||
|
||||
if listener != nil {
|
||||
listener.OnConnected()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SSHClient) readLoop(r io.Reader, name string) string {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := r.Read(buf)
|
||||
if n > 0 {
|
||||
s.mu.Lock()
|
||||
listener := s.listener
|
||||
s.mu.Unlock()
|
||||
if listener != nil {
|
||||
chunk := make([]byte, n)
|
||||
copy(chunk, buf[:n])
|
||||
listener.OnData(chunk)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
// EOF is a normal shell exit, so report it without a reason.
|
||||
if errors.Is(err, io.EOF) {
|
||||
return ""
|
||||
}
|
||||
log.Debugf("ssh %s read: %v", name, err)
|
||||
return rootCause(err).Error()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// notifyStatus writes a progress line to the terminal through the normal
|
||||
// output path, so long steps are visible while nothing else is arriving.
|
||||
func (s *SSHClient) notifyStatus(text string) {
|
||||
s.mu.Lock()
|
||||
listener := s.listener
|
||||
s.mu.Unlock()
|
||||
if listener != nil {
|
||||
listener.OnData([]byte("\r\n\x1b[33m" + text + "\x1b[0m\r\n"))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SSHClient) notifyClose(gen uint64, reason string) {
|
||||
s.mu.Lock()
|
||||
if gen != s.gen || s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
listener := s.listener
|
||||
s.mu.Unlock()
|
||||
if listener != nil {
|
||||
listener.OnClose(reason)
|
||||
}
|
||||
}
|
||||
|
||||
func closeQuiet(c io.Closer, label string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if err := c.Close(); err != nil && !errors.Is(err, io.EOF) {
|
||||
log.Debugf("ssh: close %s: %v", label, err)
|
||||
}
|
||||
}
|
||||
|
||||
func detectServerType(host string, port int) detection.ServerType {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), sshDetectionTimeout)
|
||||
defer cancel()
|
||||
|
||||
dialer := &net.Dialer{}
|
||||
serverType, err := detection.DetectSSHServerType(ctx, dialer, host, port)
|
||||
if err != nil {
|
||||
log.Debugf("ssh: server detection failed: %v (assuming regular SSH)", err)
|
||||
return detection.ServerTypeRegular
|
||||
}
|
||||
return serverType
|
||||
}
|
||||
|
||||
// rootCause returns the innermost error of a %w chain, so the terminal shows
|
||||
// "i/o timeout" rather than every layer that added context on the way up.
|
||||
func rootCause(err error) error {
|
||||
for {
|
||||
// A joined error has no single root, so keep it as-is.
|
||||
if _, ok := err.(interface{ Unwrap() []error }); ok {
|
||||
return err
|
||||
}
|
||||
next := errors.Unwrap(err)
|
||||
if next == nil {
|
||||
return err
|
||||
}
|
||||
err = next
|
||||
}
|
||||
}
|
||||
|
||||
// isAuthFailure distinguishes credential rejection from dial, timeout and
|
||||
// host-key errors, which retrying with a password would not fix.
|
||||
func isAuthFailure(err error) bool {
|
||||
if errors.Is(err, errPasswordRequired) {
|
||||
return true
|
||||
}
|
||||
var partial *gossh.PartialSuccessError
|
||||
if errors.As(err, &partial) {
|
||||
return true
|
||||
}
|
||||
return strings.Contains(err.Error(), "unable to authenticate")
|
||||
}
|
||||
|
||||
// passwordCouldHelp reports whether prompting for a password again can change
|
||||
// the outcome. gossh lists a method under "attempted methods" only when the
|
||||
// server offered it, so a supplied password that was never attempted means the
|
||||
// server does not accept passwords and the real error should surface instead.
|
||||
func passwordCouldHelp(err error, passwordOffered bool) bool {
|
||||
if !passwordOffered {
|
||||
return true
|
||||
}
|
||||
msg := err.Error()
|
||||
return strings.Contains(msg, "password") || strings.Contains(msg, "keyboard-interactive")
|
||||
}
|
||||
@@ -1,168 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
gossh "golang.org/x/crypto/ssh"
|
||||
"golang.org/x/crypto/ssh/knownhosts"
|
||||
)
|
||||
|
||||
const knownHostsNamespace = "ssh"
|
||||
|
||||
const (
|
||||
hostKeyUnknown hostKeyVerdict = iota
|
||||
hostKeyMatched
|
||||
hostKeyChanged
|
||||
)
|
||||
|
||||
var knownHostsMu sync.Mutex
|
||||
|
||||
type hostKeyVerdict uint8
|
||||
|
||||
type knownHostsSection struct {
|
||||
KnownHosts []string `json:"knownHosts"`
|
||||
}
|
||||
|
||||
type knownHostsStore struct {
|
||||
prefs prefsStore
|
||||
}
|
||||
|
||||
// RemoveKnownHost deletes every known-hosts entry for host:port from the
|
||||
// profile's store, so a host trusted for a session that is being deleted does
|
||||
// not linger. Java calls this only once no session targets that host, so a
|
||||
// shared host stays trusted. A missing entry is not an error: the goal state
|
||||
// is "absent".
|
||||
func RemoveKnownHost(configDir, profileID, host string, port int) error {
|
||||
store, err := openKnownHostsStore(configDir, profileID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return store.removeHost(host, port)
|
||||
}
|
||||
|
||||
func openKnownHostsStore(configDir, profileID string) (*knownHostsStore, error) {
|
||||
prefs, err := newProfilePrefs(configDir, profileID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &knownHostsStore{prefs: prefs}, nil
|
||||
}
|
||||
|
||||
func (st *knownHostsStore) verify(hostname string, remote net.Addr, key gossh.PublicKey) (hostKeyVerdict, error) {
|
||||
lines, err := st.lines()
|
||||
if err != nil {
|
||||
return hostKeyUnknown, err
|
||||
}
|
||||
targets := knownHostsTargets(hostname, remote)
|
||||
|
||||
verdict := hostKeyUnknown
|
||||
for _, line := range lines {
|
||||
pubKey, ok := knownHostsLineKey(line, targets)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if pubKey.Type() == key.Type() && bytes.Equal(pubKey.Marshal(), key.Marshal()) {
|
||||
return hostKeyMatched, nil
|
||||
}
|
||||
verdict = hostKeyChanged
|
||||
}
|
||||
return verdict, nil
|
||||
}
|
||||
|
||||
func (st *knownHostsStore) append(hostname string, remote net.Addr, key gossh.PublicKey) error {
|
||||
line := knownhosts.Line(knownHostsTargets(hostname, remote), key)
|
||||
|
||||
knownHostsMu.Lock()
|
||||
defer knownHostsMu.Unlock()
|
||||
|
||||
lines, err := st.lines()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return st.prefs.Put(knownHostsNamespace, knownHostsSection{KnownHosts: append(lines, line)})
|
||||
}
|
||||
|
||||
func (st *knownHostsStore) removeHost(host string, port int) error {
|
||||
target := knownhosts.Normalize(net.JoinHostPort(host, strconv.Itoa(port)))
|
||||
|
||||
knownHostsMu.Lock()
|
||||
defer knownHostsMu.Unlock()
|
||||
|
||||
lines, err := st.lines()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
kept := make([]string, 0, len(lines))
|
||||
for _, line := range lines {
|
||||
if knownHostsLineMatches(line, target) {
|
||||
continue
|
||||
}
|
||||
kept = append(kept, line)
|
||||
}
|
||||
if len(kept) == len(lines) {
|
||||
return nil
|
||||
}
|
||||
return st.prefs.Put(knownHostsNamespace, knownHostsSection{KnownHosts: kept})
|
||||
}
|
||||
|
||||
func (st *knownHostsStore) lines() ([]string, error) {
|
||||
var section knownHostsSection
|
||||
if _, err := st.prefs.Get(knownHostsNamespace, §ion); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return section.KnownHosts, nil
|
||||
}
|
||||
|
||||
func knownHostsTargets(hostname string, remote net.Addr) []string {
|
||||
targets := []string{knownhosts.Normalize(hostname)}
|
||||
if remote != nil {
|
||||
if normalized := knownhosts.Normalize(remote.String()); normalized != targets[0] {
|
||||
targets = append(targets, normalized)
|
||||
}
|
||||
}
|
||||
return targets
|
||||
}
|
||||
|
||||
func knownHostsLineKey(line string, targets []string) (gossh.PublicKey, bool) {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if trimmed == "" || strings.HasPrefix(trimmed, "#") {
|
||||
return nil, false
|
||||
}
|
||||
_, hosts, pubKey, _, _, err := gossh.ParseKnownHosts([]byte(trimmed))
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
for _, host := range hosts {
|
||||
for _, target := range targets {
|
||||
if host == target {
|
||||
return pubKey, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// knownHostsLineMatches reports whether a known-hosts line's address list
|
||||
// contains the normalized target. Comment and blank lines never match.
|
||||
func knownHostsLineMatches(line, target string) bool {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if trimmed == "" || strings.HasPrefix(trimmed, "#") {
|
||||
return false
|
||||
}
|
||||
fields := strings.Fields(trimmed)
|
||||
if len(fields) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, addr := range strings.Split(fields[0], ",") {
|
||||
if addr == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -1,104 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
const (
|
||||
sshSessionsNamespace = "ssh-sessions"
|
||||
maxStoredSSHSessions = 50
|
||||
)
|
||||
|
||||
type sshSessionRecord struct {
|
||||
ID string `json:"id"`
|
||||
Host string `json:"host"`
|
||||
Port int `json:"port"`
|
||||
User string `json:"user"`
|
||||
}
|
||||
|
||||
type sshSessionsSection struct {
|
||||
Sessions []sshSessionRecord `json:"sessions"`
|
||||
}
|
||||
|
||||
// SSHSessionEntry is one stored SSH session, without any credential.
|
||||
type SSHSessionEntry struct {
|
||||
ID string
|
||||
Host string
|
||||
Port int
|
||||
User string
|
||||
}
|
||||
|
||||
// SSHSessionArray wraps stored SSH sessions for gomobile compatibility.
|
||||
type SSHSessionArray struct {
|
||||
items []*SSHSessionEntry
|
||||
}
|
||||
|
||||
// NewSSHSessionArray creates an empty session array to fill via Add.
|
||||
func NewSSHSessionArray() *SSHSessionArray {
|
||||
return &SSHSessionArray{}
|
||||
}
|
||||
|
||||
// Add appends a session entry, oldest first.
|
||||
func (a *SSHSessionArray) Add(id, host string, port int, user string) {
|
||||
a.items = append(a.items, &SSHSessionEntry{ID: id, Host: host, Port: port, User: user})
|
||||
}
|
||||
|
||||
// Length returns the number of entries.
|
||||
func (a *SSHSessionArray) Length() int {
|
||||
return len(a.items)
|
||||
}
|
||||
|
||||
// Get returns the entry at index i, or nil when out of range.
|
||||
func (a *SSHSessionArray) Get(i int) *SSHSessionEntry {
|
||||
if i < 0 || i >= len(a.items) {
|
||||
return nil
|
||||
}
|
||||
return a.items[i]
|
||||
}
|
||||
|
||||
// SSHSessionStore reads and writes a profile's stored SSH sessions.
|
||||
type SSHSessionStore struct {
|
||||
prefs prefsStore
|
||||
}
|
||||
|
||||
// NewSSHSessionStore opens the session store of the given profile.
|
||||
func NewSSHSessionStore(configDir, profileID string) (*SSHSessionStore, error) {
|
||||
prefs, err := newProfilePrefs(configDir, profileID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &SSHSessionStore{prefs: prefs}, nil
|
||||
}
|
||||
|
||||
// Load returns the stored sessions, oldest first.
|
||||
func (s *SSHSessionStore) Load() (*SSHSessionArray, error) {
|
||||
var section sshSessionsSection
|
||||
if _, err := s.prefs.Get(sshSessionsNamespace, §ion); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := NewSSHSessionArray()
|
||||
for _, record := range section.Sessions {
|
||||
if record.ID == "" || record.Host == "" {
|
||||
continue
|
||||
}
|
||||
out.Add(record.ID, record.Host, record.Port, record.User)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Save replaces the stored sessions, keeping only the newest entries when the
|
||||
// list exceeds the storage cap.
|
||||
func (s *SSHSessionStore) Save(sessions *SSHSessionArray) error {
|
||||
var items []*SSHSessionEntry
|
||||
if sessions != nil {
|
||||
items = sessions.items
|
||||
}
|
||||
if len(items) > maxStoredSSHSessions {
|
||||
items = items[len(items)-maxStoredSSHSessions:]
|
||||
}
|
||||
|
||||
records := make([]sshSessionRecord, 0, len(items))
|
||||
for _, item := range items {
|
||||
records = append(records, sshSessionRecord{ID: item.ID, Host: item.Host, Port: item.Port, User: item.User})
|
||||
}
|
||||
return s.prefs.Put(sshSessionsNamespace, sshSessionsSection{Sessions: records})
|
||||
}
|
||||
@@ -21,7 +21,7 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/auth"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
nbssh "github.com/netbirdio/netbird/client/ssh"
|
||||
sshcommon "github.com/netbirdio/netbird/client/ssh"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
mgmProto "github.com/netbirdio/netbird/shared/management/proto"
|
||||
@@ -521,7 +521,12 @@ func (c *Client) VerifySSHHostKey(peerAddress string, key []byte) error {
|
||||
return err
|
||||
}
|
||||
|
||||
return nbssh.PeerKeyLookup(engine.GetPeerSSHKey).VerifySSHHostKey(peerAddress, key)
|
||||
storedKey, found := engine.GetPeerSSHKey(peerAddress)
|
||||
if !found {
|
||||
return sshcommon.ErrPeerNotFound
|
||||
}
|
||||
|
||||
return sshcommon.VerifyHostKey(storedKey, key, peerAddress)
|
||||
}
|
||||
|
||||
// SetPerformance retunes a running Client. Only PreallocatedBuffersPerPool
|
||||
|
||||
@@ -16,47 +16,28 @@ import (
|
||||
"google.golang.org/grpc"
|
||||
|
||||
nbnet "github.com/netbirdio/netbird/client/net"
|
||||
"github.com/netbirdio/netbird/client/netsweep"
|
||||
)
|
||||
|
||||
func WithCustomDialer(_ bool, _ string) grpc.DialOption {
|
||||
return grpc.WithContextDialer(dialContext)
|
||||
}
|
||||
|
||||
// WithSweeper dials like WithCustomDialer but registers connections and
|
||||
// dials with the sweeper. Append it after WithCustomDialer: gRPC applies
|
||||
// dial options in order, so the later context dialer wins.
|
||||
func WithSweeper(sweeper *netsweep.Sweeper) grpc.DialOption {
|
||||
return grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
|
||||
dial := sweeper.StartDial(ctx)
|
||||
defer dial.Release()
|
||||
if runtime.GOOS == "linux" {
|
||||
currentUser, err := user.Current()
|
||||
if err != nil {
|
||||
return nil, status.Errorf(codes.FailedPrecondition, "failed to get current user: %v", err)
|
||||
}
|
||||
|
||||
conn, err := dialContext(dial.Ctx(), addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
// the custom dialer requires root permissions which are not required for use cases run as non-root
|
||||
if currentUser.Uid != "0" {
|
||||
log.Debug("Not running as root, using standard dialer")
|
||||
dialer := &net.Dialer{}
|
||||
return dialer.DialContext(ctx, "tcp", addr)
|
||||
}
|
||||
}
|
||||
return dial.WrapConn(conn)
|
||||
|
||||
conn, err := nbnet.NewDialer().DialContext(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("nbnet.NewDialer().DialContext: %w", err)
|
||||
}
|
||||
return conn, nil
|
||||
})
|
||||
}
|
||||
|
||||
func dialContext(ctx context.Context, addr string) (net.Conn, error) {
|
||||
if runtime.GOOS == "linux" {
|
||||
currentUser, err := user.Current()
|
||||
if err != nil {
|
||||
return nil, status.Errorf(codes.FailedPrecondition, "failed to get current user: %v", err)
|
||||
}
|
||||
|
||||
// the custom dialer requires root permissions which are not required for use cases run as non-root
|
||||
if currentUser.Uid != "0" {
|
||||
log.Debug("Not running as root, using standard dialer")
|
||||
dialer := &net.Dialer{}
|
||||
return dialer.DialContext(ctx, "tcp", addr)
|
||||
}
|
||||
}
|
||||
|
||||
conn, err := nbnet.NewDialer().DialContext(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("nbnet.NewDialer().DialContext: %w", err)
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ package grpc
|
||||
import (
|
||||
"google.golang.org/grpc"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netsweep"
|
||||
"github.com/netbirdio/netbird/util/wsproxy/client"
|
||||
)
|
||||
|
||||
@@ -12,8 +11,3 @@ import (
|
||||
func WithCustomDialer(tlsEnabled bool, component string) grpc.DialOption {
|
||||
return client.WithWebSocketDialer(tlsEnabled, component)
|
||||
}
|
||||
|
||||
// WithSweeper is a no-op on WASM/JS: there is no network change signal.
|
||||
func WithSweeper(_ *netsweep.Sweeper) grpc.DialOption {
|
||||
return grpc.EmptyDialOption{}
|
||||
}
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
// Retry mirrors backoff.Retry, but the sleep between attempts also wakes on
|
||||
// OS network availability transitions: an operation cut down by a network
|
||||
// change retries the moment the network settles instead of sleeping through
|
||||
// the recovery. A nil netState never fires, leaving plain backoff.Retry
|
||||
// behavior.
|
||||
func Retry(ctx context.Context, operation backoff.Operation, bo backoff.BackOff, netState *netstate.State) error {
|
||||
bo.Reset()
|
||||
for {
|
||||
err := operation()
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var permanent *backoff.PermanentError
|
||||
if errors.As(err, &permanent) {
|
||||
return permanent.Err
|
||||
}
|
||||
|
||||
next := bo.NextBackOff()
|
||||
if next == backoff.Stop {
|
||||
if cerr := ctx.Err(); cerr != nil {
|
||||
return cerr
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
timer := time.NewTimer(next)
|
||||
select {
|
||||
case <-timer.C:
|
||||
case <-netState.Changed():
|
||||
timer.Stop()
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,91 +0,0 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
func TestRetryWakesOnNetworkChange(t *testing.T) {
|
||||
ns := netstate.New()
|
||||
attempts := 0
|
||||
operation := func() error {
|
||||
attempts++
|
||||
if attempts == 1 {
|
||||
return errors.New("cut by network change")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
go func() {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
ns.Set(false)
|
||||
}()
|
||||
|
||||
start := time.Now()
|
||||
err := Retry(context.Background(), operation, backoff.NewConstantBackOff(time.Minute), ns)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, attempts, "network change must cause one immediate retry")
|
||||
assert.Less(t, time.Since(start), time.Second, "the transition must cut the minute-long sleep short")
|
||||
}
|
||||
|
||||
func TestRetryPermanentError(t *testing.T) {
|
||||
sentinel := errors.New("permission denied")
|
||||
operation := func() error {
|
||||
return backoff.Permanent(sentinel)
|
||||
}
|
||||
|
||||
err := Retry(context.Background(), operation, backoff.NewConstantBackOff(time.Millisecond), nil)
|
||||
assert.ErrorIs(t, err, sentinel, "permanent errors must stop retries")
|
||||
}
|
||||
|
||||
func TestRetryNilNetState(t *testing.T) {
|
||||
attempts := 0
|
||||
operation := func() error {
|
||||
attempts++
|
||||
if attempts < 3 {
|
||||
return errors.New("transient")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
err := Retry(context.Background(), operation, backoff.NewConstantBackOff(time.Millisecond), nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 3, attempts, "nil network state must preserve timed retries")
|
||||
}
|
||||
|
||||
func TestRetryStops(t *testing.T) {
|
||||
failure := errors.New("still failing")
|
||||
operation := func() error {
|
||||
return failure
|
||||
}
|
||||
|
||||
err := Retry(context.Background(), operation, &backoff.StopBackOff{}, nil)
|
||||
assert.ErrorIs(t, err, failure, "stop backoff must return the operation error")
|
||||
}
|
||||
|
||||
func TestRetryCtxCancelDuringSleep(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
operation := func() error {
|
||||
return errors.New("failing")
|
||||
}
|
||||
|
||||
go func() {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
start := time.Now()
|
||||
err := Retry(ctx, operation, backoff.NewConstantBackOff(time.Minute), netstate.New())
|
||||
|
||||
assert.ErrorIs(t, err, context.Canceled, "context cancellation must stop the retry loop")
|
||||
assert.Less(t, time.Since(start), time.Second, "context cancellation must interrupt backoff sleep")
|
||||
}
|
||||
@@ -138,37 +138,26 @@ func (a *Auth) IsSSOSupported(ctx context.Context) (bool, error) {
|
||||
|
||||
// GetOAuthFlow returns an OAuth flow (PKCE or Device) using the existing management connection
|
||||
// This avoids creating a new connection to the management server
|
||||
func (a *Auth) GetOAuthFlow(ctx context.Context, forceDeviceAuth bool, hint string) (OAuthFlow, error) {
|
||||
func (a *Auth) GetOAuthFlow(ctx context.Context, forceDeviceAuth bool) (OAuthFlow, error) {
|
||||
var flow OAuthFlow
|
||||
var err error
|
||||
|
||||
err := a.withRetry(ctx, func(client *mgm.GrpcClient) error {
|
||||
err = a.withRetry(ctx, func(client *mgm.GrpcClient) error {
|
||||
if forceDeviceAuth {
|
||||
deviceFlow, err := a.getDeviceFlow(client)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
deviceFlow.SetLoginHint(hint)
|
||||
flow = deviceFlow
|
||||
return nil
|
||||
flow, err = a.getDeviceFlow(client)
|
||||
return err
|
||||
}
|
||||
|
||||
// Try PKCE flow first
|
||||
pkceFlow, err := a.getPKCEFlow(client)
|
||||
flow, err = a.getPKCEFlow(client)
|
||||
if err != nil {
|
||||
// If PKCE not supported, try Device flow
|
||||
if s, ok := status.FromError(err); ok && (s.Code() == codes.NotFound || s.Code() == codes.Unimplemented) {
|
||||
deviceFlow, err := a.getDeviceFlow(client)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
deviceFlow.SetLoginHint(hint)
|
||||
flow = deviceFlow
|
||||
return nil
|
||||
flow, err = a.getDeviceFlow(client)
|
||||
return err
|
||||
}
|
||||
return err
|
||||
}
|
||||
pkceFlow.SetLoginHint(hint)
|
||||
flow = pkceFlow
|
||||
return nil
|
||||
})
|
||||
|
||||
|
||||
@@ -97,7 +97,9 @@ func authenticateWithPKCEFlow(ctx context.Context, config *profilemanager.Config
|
||||
return nil, fmt.Errorf("getting pkce authorization flow info failed with error: %v", err)
|
||||
}
|
||||
|
||||
pkceFlowInfo.SetLoginHint(hint)
|
||||
if hint != "" {
|
||||
pkceFlowInfo.SetLoginHint(hint)
|
||||
}
|
||||
|
||||
return pkceFlowInfo, nil
|
||||
}
|
||||
@@ -125,7 +127,9 @@ func authenticateWithDeviceCodeFlow(ctx context.Context, config *profilemanager.
|
||||
}
|
||||
}
|
||||
|
||||
deviceFlowInfo.SetLoginHint(hint)
|
||||
if hint != "" {
|
||||
deviceFlowInfo.SetLoginHint(hint)
|
||||
}
|
||||
|
||||
return deviceFlowInfo, nil
|
||||
}
|
||||
|
||||
@@ -38,8 +38,6 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/updater"
|
||||
"github.com/netbirdio/netbird/client/internal/updater/installer"
|
||||
nbnet "github.com/netbirdio/netbird/client/net"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
"github.com/netbirdio/netbird/client/netsweep"
|
||||
cProto "github.com/netbirdio/netbird/client/proto"
|
||||
"github.com/netbirdio/netbird/client/ssh"
|
||||
sshconfig "github.com/netbirdio/netbird/client/ssh/config"
|
||||
@@ -72,42 +70,18 @@ type ConnectClient struct {
|
||||
updateManager *updater.Manager
|
||||
|
||||
persistSyncResponse bool
|
||||
|
||||
// netState gates every reconnection loop on OS-reported network
|
||||
// availability. Nil (the default) disables gating; mobile platforms
|
||||
// inject it via WithNetworkState.
|
||||
netState *netstate.State
|
||||
|
||||
// sweeper cuts the management, signal and relay connections on network
|
||||
// change; nil disables it.
|
||||
sweeper *netsweep.Sweeper
|
||||
}
|
||||
|
||||
// ConnectClientOption configures optional ConnectClient behavior.
|
||||
type ConnectClientOption func(*ConnectClient)
|
||||
|
||||
// WithNetworkState injects the OS network availability state that gates every
|
||||
// reconnection loop; without it gating is disabled.
|
||||
func WithNetworkState(netState *netstate.State) ConnectClientOption {
|
||||
return func(c *ConnectClient) { c.netState = netState }
|
||||
}
|
||||
|
||||
// WithSweeper injects the network change sweeper.
|
||||
func WithSweeper(sweeper *netsweep.Sweeper) ConnectClientOption {
|
||||
return func(c *ConnectClient) { c.sweeper = sweeper }
|
||||
}
|
||||
|
||||
func NewConnectClient(
|
||||
ctx context.Context,
|
||||
config *profilemanager.Config,
|
||||
statusRecorder *peer.Status,
|
||||
opts ...ConnectClientOption,
|
||||
) *ConnectClient {
|
||||
// Derive the run context here so Stop owns the cancel that unblocks the run
|
||||
// loop. runCancel is set once at construction, so Stop can call it without
|
||||
// racing the run loop's startup. Callers therefore need not cancel before Stop.
|
||||
runCtx, runCancel := context.WithCancel(ctx)
|
||||
c := &ConnectClient{
|
||||
return &ConnectClient{
|
||||
ctx: runCtx,
|
||||
runCancel: runCancel,
|
||||
runExited: make(chan struct{}),
|
||||
@@ -115,10 +89,6 @@ func NewConnectClient(
|
||||
statusRecorder: statusRecorder,
|
||||
engineMutex: sync.Mutex{},
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(c)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *ConnectClient) SetUpdateManager(um *updater.Manager) {
|
||||
@@ -304,13 +274,6 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
return nil
|
||||
}
|
||||
|
||||
// suspend connection attempts while the OS reports no usable network
|
||||
if waited, err := c.netState.Wait(c.ctx); err != nil {
|
||||
return nil
|
||||
} else if waited {
|
||||
backOff.Reset()
|
||||
}
|
||||
|
||||
state.Set(StatusConnecting)
|
||||
|
||||
engineCtx, cancel := context.WithCancel(c.ctx)
|
||||
@@ -322,8 +285,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
}()
|
||||
|
||||
log.Debugf("connecting to the Management service %s", c.config.ManagementURL.Host)
|
||||
mgmClient, err := mgm.NewClient(engineCtx, c.config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled,
|
||||
mgm.WithNetworkState(c.netState), mgm.WithSweeper(c.sweeper))
|
||||
mgmClient, err := mgm.NewClient(engineCtx, c.config.ManagementURL.Host, myPrivateKey, mgmTlsEnabled)
|
||||
if err != nil {
|
||||
// On daemon shutdown / Down() the parent context is cancelled
|
||||
// and the dial fails with "context canceled". Wrapping that
|
||||
@@ -398,7 +360,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
}()
|
||||
|
||||
// with the global Netbird config in hand connect (just a connection, no stream yet) Signal
|
||||
signalClient, err := connectToSignal(engineCtx, loginResp.GetNetbirdConfig(), myPrivateKey, c.netState, c.sweeper)
|
||||
signalClient, err := connectToSignal(engineCtx, loginResp.GetNetbirdConfig(), myPrivateKey)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
return wrapErr(err)
|
||||
@@ -434,8 +396,7 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
engineConfig.StateDir = filepath.Dir(path)
|
||||
}
|
||||
|
||||
relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String(), engineConfig.MTU,
|
||||
relayClient.WithNetworkState(c.netState), relayClient.WithSweeper(c.sweeper))
|
||||
relayManager := relayClient.NewManager(engineCtx, relayURLs, myPrivateKey.PublicKey().String(), engineConfig.MTU)
|
||||
c.statusRecorder.SetRelayMgr(relayManager)
|
||||
if len(relayURLs) > 0 {
|
||||
if token != nil {
|
||||
@@ -463,7 +424,6 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
UpdateManager: c.updateManager,
|
||||
ClientMetrics: c.clientMetrics,
|
||||
MetricsCtx: c.ctx,
|
||||
NetState: c.netState,
|
||||
}, mobileDependency)
|
||||
engine.SetSyncResponsePersistence(c.persistSyncResponse)
|
||||
c.engine = engine
|
||||
@@ -520,16 +480,6 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
|
||||
// status stream stuck at Connecting.
|
||||
err = backoff.Retry(operation, backoff.WithContext(backOff, c.ctx))
|
||||
if err != nil {
|
||||
// Once the client context is cancelled backoff.WithContext surfaces the
|
||||
// bare context error, and any attempt torn down mid-flight reports the
|
||||
// same. That cancellation is the caller asking us to stop (Stop, Down or
|
||||
// an engine restart), so exit cleanly instead of handing back a failure
|
||||
// the caller would have to distinguish from a real one.
|
||||
if c.ctx.Err() != nil && errors.Is(err, context.Canceled) {
|
||||
log.Info("exiting client retry loop, context cancelled")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Debugf("exiting client retry loop due to unrecoverable error: %s", err)
|
||||
if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.PermissionDenied) {
|
||||
state.Set(StatusNeedsLogin)
|
||||
@@ -723,7 +673,7 @@ func selectMTU(localMTU uint16, peerMTU int32) uint16 {
|
||||
}
|
||||
|
||||
// connectToSignal creates Signal Service client and established a connection
|
||||
func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourPrivateKey wgtypes.Key, netState *netstate.State, sweeper *netsweep.Sweeper) (*signal.GrpcClient, error) {
|
||||
func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourPrivateKey wgtypes.Key) (*signal.GrpcClient, error) {
|
||||
var sigTLSEnabled bool
|
||||
if wtConfig.Signal.Protocol == mgmProto.HostConfig_HTTPS {
|
||||
sigTLSEnabled = true
|
||||
@@ -731,8 +681,7 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.NetbirdConfig, ourP
|
||||
sigTLSEnabled = false
|
||||
}
|
||||
|
||||
signalClient, err := signal.NewClient(ctx, wtConfig.Signal.Uri, ourPrivateKey, sigTLSEnabled,
|
||||
signal.WithNetworkState(netState), signal.WithSweeper(sweeper))
|
||||
signalClient, err := signal.NewClient(ctx, wtConfig.Signal.Uri, ourPrivateKey, sigTLSEnabled)
|
||||
if err != nil {
|
||||
log.Errorf("error while connecting to the Signal Exchange Service %s: %s", wtConfig.Signal.Uri, err)
|
||||
return nil, gstatus.Errorf(codes.FailedPrecondition, "failed connecting to Signal Service : %s", err)
|
||||
|
||||
@@ -35,8 +35,6 @@ var (
|
||||
// exported so a diagnostic reader reports the same locations that are written.
|
||||
const (
|
||||
// NRPTKeyPrefix starts the name of every NRPT rule key this client creates.
|
||||
// Older versions used different layouts under the same prefix: a single
|
||||
// unsuffixed key, then one key per domain, now one key per batch of domains.
|
||||
NRPTKeyPrefix = "NetBird-Match"
|
||||
|
||||
// DNSPolicyConfigRoot holds the NRPT rules of the local policy store.
|
||||
@@ -91,6 +89,7 @@ type registryConfigurator struct {
|
||||
guid string
|
||||
routingAll bool
|
||||
gpo bool
|
||||
nrptEntryCount int
|
||||
origNameservers []netip.Addr
|
||||
}
|
||||
|
||||
@@ -323,9 +322,14 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
|
||||
}
|
||||
|
||||
if len(matchDomains) != 0 {
|
||||
if err := r.addDNSMatchPolicy(matchDomains, config.ServerIP); err != nil {
|
||||
count, err := r.addDNSMatchPolicy(matchDomains, config.ServerIP)
|
||||
// Update count even on error to ensure cleanup covers partially created rules
|
||||
r.nrptEntryCount = count
|
||||
if err != nil {
|
||||
return fmt.Errorf("add dns match policy: %w", err)
|
||||
}
|
||||
} else {
|
||||
r.nrptEntryCount = 0
|
||||
}
|
||||
|
||||
r.updateState(stateManager)
|
||||
@@ -341,8 +345,9 @@ func (r *registryConfigurator) applyDNSConfig(config HostDNSConfig, stateManager
|
||||
|
||||
func (r *registryConfigurator) updateState(stateManager *statemanager.Manager) {
|
||||
if err := stateManager.UpdateState(&ShutdownState{
|
||||
Guid: r.guid,
|
||||
GPO: r.gpo,
|
||||
Guid: r.guid,
|
||||
GPO: r.gpo,
|
||||
NRPTEntryCount: r.nrptEntryCount,
|
||||
}); err != nil {
|
||||
log.Errorf("failed to update shutdown state: %s", err)
|
||||
}
|
||||
@@ -357,7 +362,7 @@ func (r *registryConfigurator) addDNSSetupForAll(ip netip.Addr) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip netip.Addr) error {
|
||||
func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip netip.Addr) (int, error) {
|
||||
// if the gpo key is present, we need to put our DNS settings there, otherwise our config might be ignored
|
||||
// see https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-gpnrpt/8cc31cb9-20cb-4140-9e85-3e08703b4745
|
||||
|
||||
@@ -374,17 +379,19 @@ func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip netip.Addr
|
||||
gpoPath := fmt.Sprintf("%s-%d", gpoDnsPolicyConfigMatchPath, ruleIndex)
|
||||
|
||||
if err := r.configureDNSPolicy(localPath, batchDomains, ip); err != nil {
|
||||
return fmt.Errorf("configure DNS Local policy for rule %d: %w", ruleIndex, err)
|
||||
return ruleIndex, fmt.Errorf("configure DNS Local policy for rule %d: %w", ruleIndex, err)
|
||||
}
|
||||
|
||||
// Increment immediately so the caller's cleanup path knows about this rule
|
||||
ruleIndex++
|
||||
|
||||
if r.gpo {
|
||||
if err := r.configureDNSPolicy(gpoPath, batchDomains, ip); err != nil {
|
||||
return fmt.Errorf("configure gpo DNS policy for rule %d: %w", ruleIndex, err)
|
||||
return ruleIndex, fmt.Errorf("configure gpo DNS policy for rule %d: %w", ruleIndex-1, err)
|
||||
}
|
||||
}
|
||||
|
||||
log.Debugf("added NRPT rule %d with %d domains", ruleIndex, len(batchDomains))
|
||||
ruleIndex++
|
||||
log.Debugf("added NRPT rule %d with %d domains", ruleIndex-1, len(batchDomains))
|
||||
}
|
||||
|
||||
if r.gpo {
|
||||
@@ -394,7 +401,7 @@ func (r *registryConfigurator) addDNSMatchPolicy(domains []string, ip netip.Addr
|
||||
}
|
||||
|
||||
log.Infof("added %d NRPT rules for %d domains", ruleIndex, len(domains))
|
||||
return nil
|
||||
return ruleIndex, nil
|
||||
}
|
||||
|
||||
func (r *registryConfigurator) configureDNSPolicy(policyPath string, domains []string, ip netip.Addr) error {
|
||||
@@ -527,28 +534,28 @@ func (r *registryConfigurator) restoreHostDNS() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// removeDNSMatchPolicies deletes every NRPT rule this client may have created,
|
||||
// from the local and the GPO policy store. The rules are found by enumerating
|
||||
// the registry, the only authoritative record of what was written. Cleanup must
|
||||
// not depend on a rule count: the in-memory one is scoped to a single
|
||||
// registryConfigurator and the persisted one is deleted on every clean
|
||||
// disconnect, and a rule left behind keeps resolving names over an interface
|
||||
// that is gone, until reboot discards the volatile key.
|
||||
func (r *registryConfigurator) removeDNSMatchPolicies() error {
|
||||
var merr *multierror.Error
|
||||
|
||||
for _, root := range []string{DNSPolicyConfigRoot, GPODNSPolicyConfigRoot} {
|
||||
names, err := listNRPTRuleKeys(root)
|
||||
if err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("list rule keys under %s: %w", root, err))
|
||||
continue
|
||||
// Try to remove the base entries (for backward compatibility)
|
||||
if err := removeRegistryKeyFromDNSPolicyConfig(dnsPolicyConfigMatchPath); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove local base entry: %w", err))
|
||||
}
|
||||
|
||||
if err := removeRegistryKeyFromDNSPolicyConfig(gpoDnsPolicyConfigMatchPath); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove GPO base entry: %w", err))
|
||||
}
|
||||
|
||||
for i := 0; i < r.nrptEntryCount; i++ {
|
||||
localPath := fmt.Sprintf("%s-%d", dnsPolicyConfigMatchPath, i)
|
||||
gpoPath := fmt.Sprintf("%s-%d", gpoDnsPolicyConfigMatchPath, i)
|
||||
|
||||
if err := removeRegistryKeyFromDNSPolicyConfig(localPath); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove local entry %d: %w", i, err))
|
||||
}
|
||||
|
||||
for _, name := range names {
|
||||
path := root + `\` + name
|
||||
if err := removeRegistryKeyFromDNSPolicyConfig(path); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove entry %s: %w", path, err))
|
||||
}
|
||||
if err := removeRegistryKeyFromDNSPolicyConfig(gpoPath); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove GPO entry %d: %w", i, err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -563,39 +570,6 @@ func (r *registryConfigurator) restoreUncleanShutdownDNS() error {
|
||||
return r.restoreHostDNS()
|
||||
}
|
||||
|
||||
// listNRPTRuleKeys returns the names of our NRPT rule keys under a policy store
|
||||
// root. An absent root holds nothing to clean up, which is the normal state of
|
||||
// the GPO store on a machine without DNS Client policy.
|
||||
func listNRPTRuleKeys(root string) ([]string, error) {
|
||||
k, err := registry.OpenKey(registry.LOCAL_MACHINE, root, registry.ENUMERATE_SUB_KEYS)
|
||||
switch {
|
||||
case errors.Is(err, registry.ErrNotExist), errors.Is(err, syscall.ERROR_PATH_NOT_FOUND):
|
||||
// the GPO store is absent on a machine without DNS client policy
|
||||
log.Debugf("HKEY_LOCAL_MACHINE\\%s does not exist", root)
|
||||
return nil, nil
|
||||
case err != nil:
|
||||
// any other failure has to reach the caller: reporting no rules would
|
||||
// report a successful cleanup while leaving the rules in place
|
||||
return nil, fmt.Errorf("open HKEY_LOCAL_MACHINE\\%s: %w", root, err)
|
||||
}
|
||||
defer closer(k)
|
||||
|
||||
names, err := k.ReadSubKeyNames(-1)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read subkey names: %w", err)
|
||||
}
|
||||
|
||||
var ruleKeys []string
|
||||
for _, name := range names {
|
||||
// registry key names are case insensitive
|
||||
if strings.HasPrefix(strings.ToLower(name), strings.ToLower(NRPTKeyPrefix)) {
|
||||
ruleKeys = append(ruleKeys, name)
|
||||
}
|
||||
}
|
||||
|
||||
return ruleKeys, nil
|
||||
}
|
||||
|
||||
func removeRegistryKeyFromDNSPolicyConfig(regKeyPath string) error {
|
||||
k, err := registry.OpenKey(registry.LOCAL_MACHINE, regKeyPath, registry.QUERY_VALUE)
|
||||
if err != nil {
|
||||
|
||||
@@ -25,7 +25,7 @@ func TestNRPTEntriesCleanupOnConfigChange(t *testing.T) {
|
||||
|
||||
// Create a test interface registry key so updateSearchDomains doesn't fail
|
||||
testGUID := "{12345678-1234-1234-1234-123456789ABC}"
|
||||
interfacePath := InterfaceConfigPath + `\` + testGUID
|
||||
interfacePath := `SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces\` + testGUID
|
||||
testKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, interfacePath, registry.SET_VALUE)
|
||||
require.NoError(t, err, "Should create test interface registry key")
|
||||
testKey.Close()
|
||||
@@ -56,7 +56,7 @@ func TestNRPTEntriesCleanupOnConfigChange(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify 3 NRPT rules exist
|
||||
assert.Equal(t, 3, countNRPTRuleKeys(t), "Should create 3 NRPT rules for 125 domains")
|
||||
assert.Equal(t, 3, cfg.nrptEntryCount, "Should create 3 NRPT rules for 125 domains")
|
||||
for i := 0; i < 3; i++ {
|
||||
exists, err := registryKeyExists(fmt.Sprintf("%s-%d", dnsPolicyConfigMatchPath, i))
|
||||
require.NoError(t, err)
|
||||
@@ -81,7 +81,7 @@ func TestNRPTEntriesCleanupOnConfigChange(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify first 2 NRPT rules exist
|
||||
assert.Equal(t, 2, countNRPTRuleKeys(t), "Should create 2 NRPT rules for 75 domains")
|
||||
assert.Equal(t, 2, cfg.nrptEntryCount, "Should create 2 NRPT rules for 75 domains")
|
||||
for i := 0; i < 2; i++ {
|
||||
exists, err := registryKeyExists(fmt.Sprintf("%s-%d", dnsPolicyConfigMatchPath, i))
|
||||
require.NoError(t, err)
|
||||
@@ -106,65 +106,9 @@ func registryKeyExists(path string) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// TestNRPTCleanupWithoutRuleCount verifies that rules written by a previous run
|
||||
// are removed by a configurator that has no record of how many there are: an
|
||||
// unclean exit loses the in-memory count and a clean disconnect deletes the
|
||||
// persisted one, so cleanup cannot depend on either.
|
||||
func TestNRPTCleanupWithoutRuleCount(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("skipping registry integration test in short mode")
|
||||
}
|
||||
|
||||
defer cleanupRegistryKeys(t)
|
||||
cleanupRegistryKeys(t)
|
||||
|
||||
testIP := netip.MustParseAddr("100.64.0.1")
|
||||
|
||||
// 75 domains produce two indexed rules, as the current layout does
|
||||
domains := make([]string, 75)
|
||||
for i := range domains {
|
||||
domains[i] = fmt.Sprintf(".domain%d.com", i+1)
|
||||
}
|
||||
|
||||
previousRun := ®istryConfigurator{}
|
||||
require.NoError(t, previousRun.addDNSMatchPolicy(domains, testIP))
|
||||
|
||||
// the unsuffixed key an older version would have written
|
||||
require.NoError(t, previousRun.configureDNSPolicy(dnsPolicyConfigMatchPath, []string{".legacy.example.com"}, testIP))
|
||||
|
||||
// a policy owned by someone else, which cleanup must not touch
|
||||
foreignPath := DNSPolicyConfigRoot + `\DnsPolicyConfigTestForeign`
|
||||
foreignKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, foreignPath, registry.SET_VALUE)
|
||||
require.NoError(t, err, "Should create foreign policy key")
|
||||
foreignKey.Close()
|
||||
defer func() {
|
||||
_ = registry.DeleteKey(registry.LOCAL_MACHINE, foreignPath)
|
||||
}()
|
||||
|
||||
require.Equal(t, 3, countNRPTRuleKeys(t), "Should have two indexed rules and the legacy one")
|
||||
|
||||
// a configurator that never applied a DNS config, as one built after a
|
||||
// restart or from a shutdown state without a count is
|
||||
freshRun := ®istryConfigurator{}
|
||||
require.NoError(t, freshRun.removeDNSMatchPolicies())
|
||||
|
||||
assert.Equal(t, 0, countNRPTRuleKeys(t), "Should remove every rule left by the previous run")
|
||||
|
||||
exists, err := registryKeyExists(foreignPath)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, exists, "Should not remove a policy that is not ours")
|
||||
}
|
||||
|
||||
func countNRPTRuleKeys(t *testing.T) int {
|
||||
t.Helper()
|
||||
|
||||
names, err := listNRPTRuleKeys(DNSPolicyConfigRoot)
|
||||
require.NoError(t, err, "Should list NRPT rule keys")
|
||||
return len(names)
|
||||
}
|
||||
|
||||
func cleanupRegistryKeys(*testing.T) {
|
||||
cfg := ®istryConfigurator{}
|
||||
// Clean up more entries to account for batching tests with many domains
|
||||
cfg := ®istryConfigurator{nrptEntryCount: 20}
|
||||
_ = cfg.removeDNSMatchPolicies()
|
||||
}
|
||||
|
||||
@@ -181,7 +125,7 @@ func TestNRPTDomainBatching(t *testing.T) {
|
||||
|
||||
// Create a test interface registry key so updateSearchDomains doesn't fail
|
||||
testGUID := "{12345678-1234-1234-1234-123456789ABC}"
|
||||
interfacePath := InterfaceConfigPath + `\` + testGUID
|
||||
interfacePath := `SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces\` + testGUID
|
||||
testKey, _, err := registry.CreateKey(registry.LOCAL_MACHINE, interfacePath, registry.SET_VALUE)
|
||||
require.NoError(t, err, "Should create test interface registry key")
|
||||
testKey.Close()
|
||||
@@ -249,7 +193,7 @@ func TestNRPTDomainBatching(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify that exactly expectedRuleCount rules were created
|
||||
assert.Equal(t, tc.expectedRuleCount, countNRPTRuleKeys(t),
|
||||
assert.Equal(t, tc.expectedRuleCount, cfg.nrptEntryCount,
|
||||
"Should create %d NRPT rules for %d domains", tc.expectedRuleCount, tc.domainCount)
|
||||
|
||||
// Verify all expected rules exist
|
||||
|
||||
@@ -5,8 +5,9 @@ import (
|
||||
)
|
||||
|
||||
type ShutdownState struct {
|
||||
Guid string
|
||||
GPO bool
|
||||
Guid string
|
||||
GPO bool
|
||||
NRPTEntryCount int
|
||||
}
|
||||
|
||||
func (s *ShutdownState) Name() string {
|
||||
@@ -15,8 +16,9 @@ func (s *ShutdownState) Name() string {
|
||||
|
||||
func (s *ShutdownState) Cleanup() error {
|
||||
manager := ®istryConfigurator{
|
||||
guid: s.Guid,
|
||||
gpo: s.GPO,
|
||||
guid: s.Guid,
|
||||
gpo: s.GPO,
|
||||
nrptEntryCount: s.NRPTEntryCount,
|
||||
}
|
||||
|
||||
if err := manager.restoreUncleanShutdownDNS(); err != nil {
|
||||
|
||||
@@ -2,21 +2,17 @@ package ebpf
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/cilium/ebpf/link"
|
||||
"github.com/cilium/ebpf/rlimit"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/sys/unix"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/ebpf/manager"
|
||||
)
|
||||
|
||||
const (
|
||||
xdpProgName = "nb_xdp_prog"
|
||||
|
||||
mapKeyFeatures uint32 = 0
|
||||
|
||||
featureFlagWGProxy = 0b00000001
|
||||
@@ -72,50 +68,21 @@ func (tf *GeneralManager) loadXdp() error {
|
||||
return err
|
||||
}
|
||||
|
||||
// lo has no native XDP, so the program runs in generic mode. Unless it
|
||||
// declares multi-buffer support the kernel must linearize every non-linear
|
||||
// skb before running it. Loopback packets are up to 64 KB, so that is a
|
||||
// contiguous GFP_ATOMIC allocation per packet, and when it fails the packet
|
||||
// is dropped before the program runs, stalling local TCP connections.
|
||||
// Multi-buffer XDP in generic mode requires kernel 6.3, so fall back to a
|
||||
// plain attach when the kernel rejects it.
|
||||
err = tf.attachXdp(iFace.Index, true)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
log.Debugf("failed to attach multi-buffer xdp program, retrying without it: %s", err)
|
||||
|
||||
return tf.attachXdp(iFace.Index, false)
|
||||
}
|
||||
|
||||
func (tf *GeneralManager) attachXdp(iFaceIndex int, multiBuffer bool) error {
|
||||
spec, err := loadBpf()
|
||||
// load pre-compiled programs into the kernel.
|
||||
err = loadBpfObjects(&tf.bpfObjs, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load bpf spec: %w", err)
|
||||
}
|
||||
|
||||
if multiBuffer {
|
||||
prog, ok := spec.Programs[xdpProgName]
|
||||
if !ok {
|
||||
return fmt.Errorf("program %s not found in bpf spec", xdpProgName)
|
||||
}
|
||||
prog.Flags |= unix.BPF_F_XDP_HAS_FRAGS
|
||||
}
|
||||
|
||||
if err := spec.LoadAndAssign(&tf.bpfObjs, nil); err != nil {
|
||||
return fmt.Errorf("load bpf objects: %w", err)
|
||||
return err
|
||||
}
|
||||
|
||||
tf.link, err = link.AttachXDP(link.XDPOptions{
|
||||
Program: tf.bpfObjs.NbXdpProg,
|
||||
Interface: iFaceIndex,
|
||||
Interface: iFace.Index,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
if closeErr := tf.bpfObjs.Close(); closeErr != nil {
|
||||
log.Debugf("failed to close bpf objects after xdp attach error: %s", closeErr)
|
||||
}
|
||||
_ = tf.bpfObjs.Close()
|
||||
tf.link = nil
|
||||
return fmt.Errorf("attach xdp: %w", err)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -59,7 +59,6 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/syncstore"
|
||||
"github.com/netbirdio/netbird/client/internal/updater"
|
||||
"github.com/netbirdio/netbird/client/jobexec"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
cProto "github.com/netbirdio/netbird/client/proto"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
@@ -182,9 +181,6 @@ type EngineServices struct {
|
||||
UpdateManager *updater.Manager
|
||||
ClientMetrics *metrics.ClientMetrics
|
||||
MetricsCtx context.Context
|
||||
// NetState gates the reconnection loops on OS-reported network
|
||||
// availability; nil disables gating.
|
||||
NetState *netstate.State
|
||||
}
|
||||
|
||||
// Engine is a mechanism responsible for reacting on Signal and Management stream events and managing connections to the remote peers.
|
||||
@@ -208,10 +204,6 @@ type Engine struct {
|
||||
config *EngineConfig
|
||||
mobileDep MobileDependency
|
||||
|
||||
// netState gates the peer reconnection guards on OS-reported network
|
||||
// availability; nil disables gating.
|
||||
netState *netstate.State
|
||||
|
||||
// STUNs is a list of STUN servers used by ICE
|
||||
STUNs []*stun.URI
|
||||
// TURNs is a list of STUN servers used by ICE
|
||||
@@ -345,7 +337,6 @@ func NewEngine(
|
||||
syncMsgMux: &sync.Mutex{},
|
||||
config: config,
|
||||
mobileDep: mobileDep,
|
||||
netState: services.NetState,
|
||||
STUNs: []*stun.URI{},
|
||||
TURNs: []*stun.URI{},
|
||||
networkSerial: 0,
|
||||
@@ -1902,8 +1893,7 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
|
||||
Addr: e.getRosenpassAddr(),
|
||||
PermissiveMode: e.config.RosenpassPermissive,
|
||||
},
|
||||
ICEConfig: e.createICEConfig(),
|
||||
NetworkState: e.netState,
|
||||
ICEConfig: e.createICEConfig(),
|
||||
}
|
||||
|
||||
serviceDependencies := peer.ServiceDependencies{
|
||||
|
||||
@@ -26,7 +26,6 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/portforward"
|
||||
"github.com/netbirdio/netbird/client/internal/rosenpass"
|
||||
"github.com/netbirdio/netbird/client/internal/stdnet"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
relayClient "github.com/netbirdio/netbird/shared/relay/client"
|
||||
)
|
||||
@@ -94,10 +93,6 @@ type ConnConfig struct {
|
||||
|
||||
// ICEConfig ICE protocol configuration
|
||||
ICEConfig icemaker.Config
|
||||
|
||||
// NetworkState gates the reconnection guard on OS-reported network
|
||||
// availability; nil disables gating.
|
||||
NetworkState *netstate.State
|
||||
}
|
||||
|
||||
type Conn struct {
|
||||
@@ -259,7 +254,7 @@ func (conn *Conn) open(engineCtx context.Context, firstPacket []byte) error {
|
||||
conn.handshaker.AddICEListener(conn.workerICE.OnNewOffer)
|
||||
}
|
||||
|
||||
conn.guard = guard.NewGuard(conn.Log, conn.isConnectedOnAllWay, conn.config.Timeout, conn.srWatcher, conn.config.NetworkState)
|
||||
conn.guard = guard.NewGuard(conn.Log, conn.isConnectedOnAllWay, conn.config.Timeout, conn.srWatcher)
|
||||
|
||||
conn.wg.Add(1)
|
||||
go func() {
|
||||
|
||||
@@ -6,8 +6,6 @@ import (
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
// ConnStatus represents the connection state as seen by the guard.
|
||||
@@ -33,26 +31,20 @@ type connStatusFunc func() ConnStatus
|
||||
// - Relayed connection disconnected
|
||||
// - ICE candidate changes
|
||||
type Guard struct {
|
||||
log *log.Entry
|
||||
isConnectedOnAllWay connStatusFunc
|
||||
timeout time.Duration
|
||||
srWatcher *SRWatcher
|
||||
// netState gates reconnect attempts on OS-reported network availability;
|
||||
// nil disables gating.
|
||||
netState *netstate.State
|
||||
log *log.Entry
|
||||
isConnectedOnAllWay connStatusFunc
|
||||
timeout time.Duration
|
||||
srWatcher *SRWatcher
|
||||
relayedConnDisconnected chan struct{}
|
||||
iCEConnDisconnected chan struct{}
|
||||
}
|
||||
|
||||
// NewGuard creates a reconnection guard for a peer connection. A nil netState
|
||||
// disables network availability gating.
|
||||
func NewGuard(log *log.Entry, isConnectedFn connStatusFunc, timeout time.Duration, srWatcher *SRWatcher, netState *netstate.State) *Guard {
|
||||
func NewGuard(log *log.Entry, isConnectedFn connStatusFunc, timeout time.Duration, srWatcher *SRWatcher) *Guard {
|
||||
return &Guard{
|
||||
log: log,
|
||||
isConnectedOnAllWay: isConnectedFn,
|
||||
timeout: timeout,
|
||||
srWatcher: srWatcher,
|
||||
netState: netState,
|
||||
relayedConnDisconnected: make(chan struct{}, 1),
|
||||
iCEConnDisconnected: make(chan struct{}, 1),
|
||||
}
|
||||
@@ -104,16 +96,9 @@ func (g *Guard) reconnectLoopWithRetry(ctx context.Context, callback func()) {
|
||||
iceState := &iceRetryState{log: g.log}
|
||||
defer iceState.reset()
|
||||
|
||||
netChanged := g.netState.Changed()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-tickerChannel:
|
||||
// skip attempts while the OS reports no usable network; the
|
||||
// netChanged case below resumes the loop once it returns
|
||||
if !g.netState.IsOnline() {
|
||||
continue
|
||||
}
|
||||
switch g.isConnectedOnAllWay() {
|
||||
case ConnStatusConnected:
|
||||
// all good, nothing to do
|
||||
@@ -150,23 +135,6 @@ func (g *Guard) reconnectLoopWithRetry(ctx context.Context, callback func()) {
|
||||
tickerChannel = ticker.C
|
||||
iceState.reset()
|
||||
|
||||
case <-netChanged:
|
||||
// Re-arm for the next transition before acting on this one.
|
||||
netChanged = g.netState.Changed()
|
||||
if !g.netState.IsOnline() {
|
||||
continue
|
||||
}
|
||||
// Ticks skipped while offline drove the backoff towards its
|
||||
// maximum without ever attempting, and left the ICE budget
|
||||
// frozen — possibly in hourly mode. Recover on our own so the
|
||||
// peer does not depend on a signal or relay event that never
|
||||
// comes when both stayed up across the outage.
|
||||
g.log.Debugf("network is back, reset reconnection ticker")
|
||||
ticker.Stop()
|
||||
ticker = g.newReconnectTicker(ctx)
|
||||
tickerChannel = ticker.C
|
||||
iceState.reset()
|
||||
|
||||
case <-ctx.Done():
|
||||
g.log.Debugf("context is done, stop reconnect loop")
|
||||
return
|
||||
|
||||
@@ -15,7 +15,7 @@ import (
|
||||
|
||||
func newTestGuard(status connStatusFunc) *Guard {
|
||||
srw := NewSRWatcher(nil, nil, nil, ice.Config{})
|
||||
return NewGuard(log.WithField("test", "guard"), status, 50*time.Millisecond, srw, nil)
|
||||
return NewGuard(log.WithField("test", "guard"), status, 50*time.Millisecond, srw)
|
||||
}
|
||||
|
||||
// countBackoffTickerGoroutines returns how many goroutines are currently sitting
|
||||
|
||||
@@ -1,107 +0,0 @@
|
||||
package guard
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/peer/ice"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
// newTestGuardWithNetState builds a guard with a realistic MaxInterval: the
|
||||
// backoff must be able to grow well past the outage, as it does in production
|
||||
// where the timeout is seconds to minutes.
|
||||
func newTestGuardWithNetState(status connStatusFunc, netState *netstate.State) *Guard {
|
||||
srw := NewSRWatcher(nil, nil, nil, ice.Config{})
|
||||
return NewGuard(log.WithField("test", "guard"), status, 30*time.Second, srw, netState)
|
||||
}
|
||||
|
||||
// TestGuard_RecoversAfterOfflineToOnline covers a peer that stays disconnected
|
||||
// across a network outage while neither signal nor relay reports an event —
|
||||
// both stayed up, as on a short airplane mode toggle over Wi-Fi.
|
||||
//
|
||||
// Every tick taken while offline is skipped, but it still advances the
|
||||
// exponential backoff, so by the time the network returns the next tick can be
|
||||
// tens of seconds away. Without an explicit reaction to the transition the
|
||||
// peer waits out that interval for a recovery that could start immediately.
|
||||
func TestGuard_RecoversAfterOfflineToOnline(t *testing.T) {
|
||||
netState := netstate.New()
|
||||
|
||||
var attempts atomic.Int32
|
||||
g := newTestGuardWithNetState(func() ConnStatus { return ConnStatusDisconnected }, netState)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
// Start from the reconnect ticker (800ms initial interval), the state a
|
||||
// peer is in after it loses its connection.
|
||||
go g.Start(ctx, func() { attempts.Add(1) })
|
||||
g.SetRelayedConnDisconnected()
|
||||
|
||||
// Let the backoff climb: 0.8s, 1.6s, 3.2s, 6.4s ... every tick is skipped
|
||||
// while offline, but each one doubles the wait for the next.
|
||||
netState.Set(false)
|
||||
time.Sleep(8 * time.Second)
|
||||
|
||||
offlineAttempts := attempts.Load()
|
||||
if offlineAttempts != 0 {
|
||||
t.Fatalf("callback ran %d times while offline, want 0", offlineAttempts)
|
||||
}
|
||||
|
||||
netState.Set(true)
|
||||
|
||||
// The next organic tick is now several seconds out, so anything within
|
||||
// this window can only come from reacting to the transition itself.
|
||||
pollCtx, stopPolling := context.WithTimeout(ctx, 2*time.Second)
|
||||
defer stopPolling()
|
||||
|
||||
select {
|
||||
case <-pollCtx.Done():
|
||||
t.Fatal("peer was not retried within 2s of the network coming back, " +
|
||||
"with neither a signal nor a relay event to fall back on")
|
||||
case <-pollUntil(pollCtx, func() bool { return attempts.Load() > 0 }):
|
||||
}
|
||||
}
|
||||
|
||||
// TestGuard_OfflineTransitionDoesNotRetry checks the other direction: going
|
||||
// offline must not itself trigger an attempt.
|
||||
func TestGuard_OfflineTransitionDoesNotRetry(t *testing.T) {
|
||||
netState := netstate.New()
|
||||
|
||||
var attempts atomic.Int32
|
||||
g := newTestGuardWithNetState(func() ConnStatus { return ConnStatusDisconnected }, netState)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go g.Start(ctx, func() { attempts.Add(1) })
|
||||
|
||||
netState.Set(false)
|
||||
time.Sleep(5 * time.Second)
|
||||
|
||||
if got := attempts.Load(); got != 0 {
|
||||
t.Fatalf("callback ran %d times after going offline, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
// pollUntil closes the returned channel once cond holds. It gives up when ctx
|
||||
// is done, so the polling goroutine never outlives the test that started it.
|
||||
func pollUntil(ctx context.Context, cond func() bool) <-chan struct{} {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
for {
|
||||
if cond() {
|
||||
close(done)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}()
|
||||
return done
|
||||
}
|
||||
@@ -1,40 +1,11 @@
|
||||
package peer
|
||||
|
||||
// ClientState identifies the client connection state delivered via
|
||||
// Listener.OnStateChanged.
|
||||
type ClientState int
|
||||
|
||||
// Client states. The numeric values cross the gomobile boundary (the mobile
|
||||
// bindings re-export them as integer constants), so they are a wire format:
|
||||
// append new states at the end, never reorder or insert.
|
||||
const (
|
||||
ClientStateDisconnected ClientState = iota
|
||||
ClientStateConnected
|
||||
ClientStateConnecting
|
||||
ClientStateDisconnecting
|
||||
// ClientStateNoNetwork is an overlay state: it is never stored as the
|
||||
// last notification, only derived from ClientStateConnecting while the
|
||||
// OS reports no usable network (see notifier.effectiveState).
|
||||
ClientStateNoNetwork
|
||||
)
|
||||
|
||||
// Listener is a callback type about the NetBird network connection state
|
||||
type Listener interface {
|
||||
// OnStateChanged reports every client state transition. New states are
|
||||
// delivered only through this callback; the per-state callbacks below
|
||||
// are kept for compatibility and will be removed once all consumers
|
||||
// have migrated.
|
||||
OnStateChanged(state ClientState)
|
||||
|
||||
// Deprecated: consume OnStateChanged instead.
|
||||
OnConnected()
|
||||
// Deprecated: consume OnStateChanged instead.
|
||||
OnDisconnected()
|
||||
// Deprecated: consume OnStateChanged instead.
|
||||
OnConnecting()
|
||||
// Deprecated: consume OnStateChanged instead.
|
||||
OnDisconnecting()
|
||||
|
||||
OnAddressChanged(string, string)
|
||||
OnPeersListChanged(int)
|
||||
}
|
||||
|
||||
@@ -4,64 +4,31 @@ import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
const (
|
||||
stateDisconnected = iota
|
||||
stateConnected
|
||||
stateConnecting
|
||||
stateDisconnecting
|
||||
)
|
||||
|
||||
type notifier struct {
|
||||
// publishLock orders state publication: it is held across computing the
|
||||
// effective state and handing it to the listener, so a transition cannot
|
||||
// overtake a newer one and leave the listener on a stale state.
|
||||
publishLock sync.Mutex
|
||||
serverStateLock sync.Mutex
|
||||
listenersLock sync.Mutex
|
||||
listener Listener
|
||||
currentClientState bool
|
||||
lastNotification ClientState
|
||||
lastNotification int
|
||||
lastNumberOfPeers int
|
||||
lastFqdnAddress string
|
||||
lastIPAddress string
|
||||
networkAvailable bool
|
||||
}
|
||||
|
||||
func newNotifier() *notifier {
|
||||
return ¬ifier{
|
||||
networkAvailable: true,
|
||||
}
|
||||
}
|
||||
|
||||
// effectiveState maps the computed state to what listeners should see:
|
||||
// while the OS reports no usable network, "Connecting" would be a lie —
|
||||
// connection attempts are suspended — so it is reported as NoNetwork.
|
||||
// Caller must hold serverStateLock.
|
||||
func (n *notifier) effectiveState(state ClientState) ClientState {
|
||||
if !n.networkAvailable && state == ClientStateConnecting {
|
||||
return ClientStateNoNetwork
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
// setNetworkAvailable records the OS network availability and re-notifies
|
||||
// the listener when the flag flips the effective state (Connecting <->
|
||||
// NoNetwork).
|
||||
func (n *notifier) setNetworkAvailable(available bool) {
|
||||
n.publishLock.Lock()
|
||||
defer n.publishLock.Unlock()
|
||||
|
||||
n.serverStateLock.Lock()
|
||||
if n.networkAvailable == available {
|
||||
n.serverStateLock.Unlock()
|
||||
return
|
||||
}
|
||||
previous := n.effectiveState(n.lastNotification)
|
||||
n.networkAvailable = available
|
||||
current := n.effectiveState(n.lastNotification)
|
||||
n.serverStateLock.Unlock()
|
||||
|
||||
if previous != current {
|
||||
n.notify(current)
|
||||
}
|
||||
return ¬ifier{}
|
||||
}
|
||||
|
||||
func (n *notifier) setListener(listener Listener) {
|
||||
n.serverStateLock.Lock()
|
||||
lastNotification := n.effectiveState(n.lastNotification)
|
||||
lastNotification := n.lastNotification
|
||||
numOfPeers := n.lastNumberOfPeers
|
||||
fqdnAddress := n.lastFqdnAddress
|
||||
address := n.lastIPAddress
|
||||
@@ -85,9 +52,6 @@ func (n *notifier) removeListener() {
|
||||
}
|
||||
|
||||
func (n *notifier) updateServerStates(mgmState bool, signalState bool) {
|
||||
n.publishLock.Lock()
|
||||
defer n.publishLock.Unlock()
|
||||
|
||||
n.serverStateLock.Lock()
|
||||
calculatedState := n.calculateState(mgmState, signalState)
|
||||
|
||||
@@ -97,54 +61,43 @@ func (n *notifier) updateServerStates(mgmState bool, signalState bool) {
|
||||
}
|
||||
|
||||
n.lastNotification = calculatedState
|
||||
effective := n.effectiveState(calculatedState)
|
||||
n.serverStateLock.Unlock()
|
||||
|
||||
n.notify(effective)
|
||||
n.notify(calculatedState)
|
||||
}
|
||||
|
||||
func (n *notifier) clientStart() {
|
||||
n.publishLock.Lock()
|
||||
defer n.publishLock.Unlock()
|
||||
|
||||
n.serverStateLock.Lock()
|
||||
n.currentClientState = true
|
||||
n.lastNotification = ClientStateConnecting
|
||||
effective := n.effectiveState(ClientStateConnecting)
|
||||
n.lastNotification = stateConnecting
|
||||
n.serverStateLock.Unlock()
|
||||
|
||||
n.notify(effective)
|
||||
n.notify(stateConnecting)
|
||||
}
|
||||
|
||||
func (n *notifier) clientStop() {
|
||||
n.publishLock.Lock()
|
||||
defer n.publishLock.Unlock()
|
||||
|
||||
n.serverStateLock.Lock()
|
||||
n.currentClientState = false
|
||||
n.lastNotification = ClientStateDisconnected
|
||||
n.lastNotification = stateDisconnected
|
||||
n.serverStateLock.Unlock()
|
||||
|
||||
n.notify(ClientStateDisconnected)
|
||||
n.notify(stateDisconnected)
|
||||
}
|
||||
|
||||
func (n *notifier) clientTearDown() {
|
||||
n.publishLock.Lock()
|
||||
defer n.publishLock.Unlock()
|
||||
|
||||
n.serverStateLock.Lock()
|
||||
n.currentClientState = false
|
||||
n.lastNotification = ClientStateDisconnecting
|
||||
n.lastNotification = stateDisconnecting
|
||||
n.serverStateLock.Unlock()
|
||||
|
||||
n.notify(ClientStateDisconnecting)
|
||||
n.notify(stateDisconnecting)
|
||||
}
|
||||
|
||||
func (n *notifier) isServerStateChanged(newState ClientState) bool {
|
||||
func (n *notifier) isServerStateChanged(newState int) bool {
|
||||
return n.lastNotification != newState
|
||||
}
|
||||
|
||||
func (n *notifier) notify(state ClientState) {
|
||||
func (n *notifier) notify(state int) {
|
||||
n.listenersLock.Lock()
|
||||
listener := n.listener
|
||||
n.listenersLock.Unlock()
|
||||
@@ -156,20 +109,20 @@ func (n *notifier) notify(state ClientState) {
|
||||
notifyListener(listener, state)
|
||||
}
|
||||
|
||||
func (n *notifier) calculateState(managementConn, signalConn bool) ClientState {
|
||||
func (n *notifier) calculateState(managementConn, signalConn bool) int {
|
||||
if managementConn && signalConn {
|
||||
return ClientStateConnected
|
||||
return stateConnected
|
||||
}
|
||||
|
||||
if !managementConn && !signalConn && !n.currentClientState {
|
||||
return ClientStateDisconnected
|
||||
return stateDisconnected
|
||||
}
|
||||
|
||||
if n.lastNotification == ClientStateDisconnecting {
|
||||
return ClientStateDisconnecting
|
||||
if n.lastNotification == stateDisconnecting {
|
||||
return stateDisconnecting
|
||||
}
|
||||
|
||||
return ClientStateConnecting
|
||||
return stateConnecting
|
||||
}
|
||||
|
||||
func (n *notifier) peerListChanged(numOfPeers int) {
|
||||
@@ -206,19 +159,15 @@ func (n *notifier) localAddressChanged(fqdn, address string) {
|
||||
listener.OnAddressChanged(fqdn, address)
|
||||
}
|
||||
|
||||
func notifyListener(l Listener, state ClientState) {
|
||||
// legacy per-state callbacks; NoNetwork is delivered only via
|
||||
// OnStateChanged below
|
||||
func notifyListener(l Listener, state int) {
|
||||
switch state {
|
||||
case ClientStateDisconnected:
|
||||
case stateDisconnected:
|
||||
l.OnDisconnected()
|
||||
case ClientStateConnected:
|
||||
case stateConnected:
|
||||
l.OnConnected()
|
||||
case ClientStateConnecting:
|
||||
case stateConnecting:
|
||||
l.OnConnecting()
|
||||
case ClientStateDisconnecting:
|
||||
case stateDisconnecting:
|
||||
l.OnDisconnecting()
|
||||
}
|
||||
|
||||
l.OnStateChanged(state)
|
||||
}
|
||||
|
||||
@@ -1,108 +0,0 @@
|
||||
package peer
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type recordingListener struct {
|
||||
mu sync.Mutex
|
||||
states []ClientState
|
||||
onState func(ClientState)
|
||||
}
|
||||
|
||||
func (l *recordingListener) OnStateChanged(state ClientState) {
|
||||
l.mu.Lock()
|
||||
l.states = append(l.states, state)
|
||||
hook := l.onState
|
||||
l.mu.Unlock()
|
||||
|
||||
if hook != nil {
|
||||
hook(state)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *recordingListener) last() (ClientState, bool) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
if len(l.states) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
return l.states[len(l.states)-1], true
|
||||
}
|
||||
|
||||
func (l *recordingListener) snapshot() []ClientState {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return append([]ClientState(nil), l.states...)
|
||||
}
|
||||
|
||||
func (l *recordingListener) OnConnected() {}
|
||||
func (l *recordingListener) OnDisconnected() {}
|
||||
func (l *recordingListener) OnConnecting() {}
|
||||
func (l *recordingListener) OnDisconnecting() {}
|
||||
func (l *recordingListener) OnAddressChanged(string, string) {}
|
||||
func (l *recordingListener) OnPeersListChanged(int) {}
|
||||
|
||||
// TestNotifier_ConcurrentAvailabilityFlipOrdersPublication holds the first
|
||||
// transition inside the listener callback and flips availability again from
|
||||
// another goroutine while it is parked. The second flip must not publish
|
||||
// ahead of the one in flight, otherwise the listener ends up on a state the
|
||||
// notifier already superseded.
|
||||
func TestNotifier_ConcurrentAvailabilityFlipOrdersPublication(t *testing.T) {
|
||||
n := newNotifier()
|
||||
n.currentClientState = true
|
||||
n.lastNotification = ClientStateConnecting
|
||||
|
||||
entered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
|
||||
l := &recordingListener{}
|
||||
l.onState = func(state ClientState) {
|
||||
if state != ClientStateNoNetwork {
|
||||
return
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.onState = nil
|
||||
l.mu.Unlock()
|
||||
close(entered)
|
||||
<-release
|
||||
}
|
||||
n.listener = l
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
n.setNetworkAvailable(false)
|
||||
}()
|
||||
|
||||
<-entered
|
||||
|
||||
flipped := make(chan struct{})
|
||||
go func() {
|
||||
defer close(flipped)
|
||||
n.setNetworkAvailable(true)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-flipped:
|
||||
t.Fatal("the online transition published while the offline one was " +
|
||||
"still in flight; publication is not serialized")
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
}
|
||||
|
||||
close(release)
|
||||
<-flipped
|
||||
wg.Wait()
|
||||
|
||||
got, ok := l.last()
|
||||
if !ok {
|
||||
t.Fatal("listener never observed a state")
|
||||
}
|
||||
if got != ClientStateConnecting {
|
||||
t.Fatalf("listener holds %v after the network came back, want Connecting; sequence: %v",
|
||||
got, l.snapshot())
|
||||
}
|
||||
}
|
||||
@@ -6,32 +6,29 @@ import (
|
||||
)
|
||||
|
||||
type mocListener struct {
|
||||
lastState ClientState
|
||||
lastState int
|
||||
wg sync.WaitGroup
|
||||
peersWg sync.WaitGroup
|
||||
peers int
|
||||
}
|
||||
|
||||
func (l *mocListener) OnConnected() {
|
||||
l.lastState = ClientStateConnected
|
||||
l.lastState = stateConnected
|
||||
l.wg.Done()
|
||||
}
|
||||
func (l *mocListener) OnDisconnected() {
|
||||
l.lastState = ClientStateDisconnected
|
||||
l.lastState = stateDisconnected
|
||||
l.wg.Done()
|
||||
}
|
||||
func (l *mocListener) OnConnecting() {
|
||||
l.lastState = ClientStateConnecting
|
||||
l.lastState = stateConnecting
|
||||
l.wg.Done()
|
||||
}
|
||||
func (l *mocListener) OnDisconnecting() {
|
||||
l.lastState = ClientStateDisconnecting
|
||||
l.lastState = stateDisconnecting
|
||||
l.wg.Done()
|
||||
}
|
||||
|
||||
func (l *mocListener) OnStateChanged(state ClientState) {
|
||||
|
||||
}
|
||||
func (l *mocListener) OnAddressChanged(host, addr string) {
|
||||
|
||||
}
|
||||
@@ -60,15 +57,15 @@ func Test_notifier_serverState(t *testing.T) {
|
||||
|
||||
type scenario struct {
|
||||
name string
|
||||
expected ClientState
|
||||
expected int
|
||||
mgmState bool
|
||||
signalState bool
|
||||
}
|
||||
scenarios := []scenario{
|
||||
{"connected", ClientStateConnected, true, true},
|
||||
{"mgm down", ClientStateConnecting, false, true},
|
||||
{"signal down", ClientStateConnecting, true, false},
|
||||
{"disconnected", ClientStateDisconnected, false, false},
|
||||
{"connected", stateConnected, true, true},
|
||||
{"mgm down", stateConnecting, false, true},
|
||||
{"signal down", stateConnecting, true, false},
|
||||
{"disconnected", stateDisconnected, false, false},
|
||||
}
|
||||
|
||||
for _, tt := range scenarios {
|
||||
@@ -88,7 +85,7 @@ func Test_notifier_SetListener(t *testing.T) {
|
||||
listener.setPeersWaiter()
|
||||
|
||||
n := newNotifier()
|
||||
n.lastNotification = ClientStateConnecting
|
||||
n.lastNotification = stateConnecting
|
||||
n.setListener(listener)
|
||||
listener.wait()
|
||||
listener.waitPeers()
|
||||
@@ -102,7 +99,7 @@ func Test_notifier_RemoveListener(t *testing.T) {
|
||||
listener.setWaiter()
|
||||
listener.setPeersWaiter()
|
||||
n := newNotifier()
|
||||
n.lastNotification = ClientStateConnecting
|
||||
n.lastNotification = stateConnecting
|
||||
n.setListener(listener)
|
||||
// setListener replays cached state on a goroutine; wait for both the state
|
||||
// and peers callbacks to finish so we don't race on listener.peers.
|
||||
|
||||
@@ -1211,12 +1211,6 @@ func (d *Status) ClientTeardown() {
|
||||
d.notifyStateChange()
|
||||
}
|
||||
|
||||
// SetNetworkAvailable records the OS-reported network availability; while
|
||||
// unavailable, listeners see NoNetwork instead of Connecting.
|
||||
func (d *Status) SetNetworkAvailable(available bool) {
|
||||
d.notifier.setNetworkAvailable(available)
|
||||
}
|
||||
|
||||
// SetConnectionListener set a listener to the notifier
|
||||
func (d *Status) SetConnectionListener(listener Listener) {
|
||||
d.notifier.setListener(listener)
|
||||
|
||||
@@ -1,130 +0,0 @@
|
||||
package profilemanager
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
const prefsFileSuffix = ".prefs.json"
|
||||
|
||||
var prefsMu sync.Mutex
|
||||
|
||||
// Prefs is a namespaced per-profile preference store backed by a single JSON
|
||||
// file next to the profile config; it is deleted together with the profile.
|
||||
type Prefs struct {
|
||||
path string
|
||||
}
|
||||
|
||||
// ProfilePrefs returns the preference store of the profile identified by id.
|
||||
func (s *ServiceManager) ProfilePrefs(id ID, username string) (*Prefs, error) {
|
||||
if !IsValidProfileFilenameStem(id) {
|
||||
return nil, fmt.Errorf("invalid profile ID: %q", id)
|
||||
}
|
||||
if id == defaultProfileName {
|
||||
return &Prefs{path: filepath.Join(filepath.Dir(DefaultConfigPath), id.String()+prefsFileSuffix)}, nil
|
||||
}
|
||||
configDir, err := s.getConfigDir(username)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get config directory for user %s: %w", username, err)
|
||||
}
|
||||
return &Prefs{path: filepath.Join(configDir, id.String()+prefsFileSuffix)}, nil
|
||||
}
|
||||
|
||||
// Get unmarshals the namespace section into v and reports whether it exists.
|
||||
func (p *Prefs) Get(namespace string, v any) (bool, error) {
|
||||
if namespace == "" {
|
||||
return false, fmt.Errorf("empty prefs namespace")
|
||||
}
|
||||
|
||||
prefsMu.Lock()
|
||||
defer prefsMu.Unlock()
|
||||
|
||||
sections, err := readPrefsFile(p.path)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
raw, ok := sections[namespace]
|
||||
if !ok {
|
||||
return false, nil
|
||||
}
|
||||
if err := json.Unmarshal(raw, v); err != nil {
|
||||
return false, fmt.Errorf("decode prefs namespace %q: %w", namespace, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// Put stores v as the namespace section, replacing any previous value.
|
||||
func (p *Prefs) Put(namespace string, v any) error {
|
||||
if namespace == "" {
|
||||
return fmt.Errorf("empty prefs namespace")
|
||||
}
|
||||
raw, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode prefs namespace %q: %w", namespace, err)
|
||||
}
|
||||
|
||||
prefsMu.Lock()
|
||||
defer prefsMu.Unlock()
|
||||
|
||||
sections, err := readPrefsFile(p.path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sections[namespace] = raw
|
||||
return writePrefsFile(p.path, sections)
|
||||
}
|
||||
|
||||
// Remove deletes the namespace section; a missing one is not an error.
|
||||
func (p *Prefs) Remove(namespace string) error {
|
||||
if namespace == "" {
|
||||
return fmt.Errorf("empty prefs namespace")
|
||||
}
|
||||
|
||||
prefsMu.Lock()
|
||||
defer prefsMu.Unlock()
|
||||
|
||||
sections, err := readPrefsFile(p.path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, ok := sections[namespace]; !ok {
|
||||
return nil
|
||||
}
|
||||
delete(sections, namespace)
|
||||
return writePrefsFile(p.path, sections)
|
||||
}
|
||||
|
||||
func removePrefsFile(path string) error {
|
||||
prefsMu.Lock()
|
||||
defer prefsMu.Unlock()
|
||||
return os.Remove(path)
|
||||
}
|
||||
|
||||
func readPrefsFile(path string) (map[string]json.RawMessage, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if os.IsNotExist(err) {
|
||||
return map[string]json.RawMessage{}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read prefs: %w", err)
|
||||
}
|
||||
|
||||
sections := map[string]json.RawMessage{}
|
||||
if err := json.Unmarshal(data, §ions); err != nil {
|
||||
return nil, fmt.Errorf("decode prefs: %w", err)
|
||||
}
|
||||
return sections, nil
|
||||
}
|
||||
|
||||
func writePrefsFile(path string, sections map[string]json.RawMessage) error {
|
||||
if err := util.WriteJsonWithRestrictedPermission(context.Background(), path, sections); err != nil {
|
||||
return fmt.Errorf("write prefs: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
package profilemanager
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type testPrefsSection struct {
|
||||
Mode uint8 `json:"mode"`
|
||||
Dest string `json:"dest"`
|
||||
}
|
||||
|
||||
func TestProfilePrefs_RoundTrip(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
created, err := sm.AddProfile("work", username)
|
||||
require.NoError(t, err)
|
||||
|
||||
prefs, err := sm.ProfilePrefs(created.ID, username)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, prefs.Put("filedrop", testPrefsSection{Mode: 2, Dest: "/tmp/x"}))
|
||||
require.NoError(t, prefs.Put("other", map[string]int{"n": 1}))
|
||||
|
||||
var got testPrefsSection
|
||||
found, err := prefs.Get("filedrop", &got)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, testPrefsSection{Mode: 2, Dest: "/tmp/x"}, got)
|
||||
|
||||
var other map[string]int
|
||||
found, err = prefs.Get("other", &other)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, map[string]int{"n": 1}, other)
|
||||
})
|
||||
}
|
||||
|
||||
func TestProfilePrefs_GetMissingNamespace(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
created, err := sm.AddProfile("work", username)
|
||||
require.NoError(t, err)
|
||||
|
||||
prefs, err := sm.ProfilePrefs(created.ID, username)
|
||||
require.NoError(t, err)
|
||||
|
||||
var got testPrefsSection
|
||||
found, err := prefs.Get("filedrop", &got)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, found)
|
||||
})
|
||||
}
|
||||
|
||||
func TestProfilePrefs_RemoveNamespace(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
created, err := sm.AddProfile("work", username)
|
||||
require.NoError(t, err)
|
||||
|
||||
prefs, err := sm.ProfilePrefs(created.ID, username)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, prefs.Put("filedrop", testPrefsSection{Mode: 1}))
|
||||
require.NoError(t, prefs.Put("other", map[string]int{"n": 1}))
|
||||
require.NoError(t, prefs.Remove("filedrop"))
|
||||
require.NoError(t, prefs.Remove("missing"))
|
||||
|
||||
var got testPrefsSection
|
||||
found, err := prefs.Get("filedrop", &got)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
var other map[string]int
|
||||
found, err = prefs.Get("other", &other)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, found)
|
||||
assert.Equal(t, map[string]int{"n": 1}, other)
|
||||
})
|
||||
}
|
||||
|
||||
func TestProfilePrefs_RejectsInvalidID(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
_, err := sm.ProfilePrefs("../escape", username)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestProfilePrefs_RejectsEmptyNamespace(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
created, err := sm.AddProfile("work", username)
|
||||
require.NoError(t, err)
|
||||
|
||||
prefs, err := sm.ProfilePrefs(created.ID, username)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = prefs.Get("", &testPrefsSection{})
|
||||
assert.Error(t, err)
|
||||
assert.Error(t, prefs.Put("", testPrefsSection{}))
|
||||
assert.Error(t, prefs.Remove(""))
|
||||
})
|
||||
}
|
||||
|
||||
func TestProfilePrefs_DefaultProfile(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
prefs, err := sm.ProfilePrefs(defaultProfileName, username)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, prefs.Put("filedrop", testPrefsSection{Mode: 1}))
|
||||
|
||||
expected := filepath.Join(filepath.Dir(DefaultConfigPath), "default"+prefsFileSuffix)
|
||||
_, err = os.Stat(expected)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestRemoveProfile_DeletesPrefsFile(t *testing.T) {
|
||||
withTestSM(t, func(sm *ServiceManager, username string) {
|
||||
created, err := sm.AddProfile("work", username)
|
||||
require.NoError(t, err)
|
||||
|
||||
prefs, err := sm.ProfilePrefs(created.ID, username)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, prefs.Put("filedrop", testPrefsSection{Mode: 2}))
|
||||
|
||||
configDir, err := sm.getConfigDir(username)
|
||||
require.NoError(t, err)
|
||||
prefsPath := filepath.Join(configDir, created.ID.String()+prefsFileSuffix)
|
||||
_, err = os.Stat(prefsPath)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, sm.RemoveProfile(created.ID, username))
|
||||
_, err = os.Stat(prefsPath)
|
||||
assert.True(t, errors.Is(err, os.ErrNotExist), "prefs file should be removed")
|
||||
})
|
||||
}
|
||||
@@ -420,11 +420,6 @@ func (s *ServiceManager) RemoveProfile(id ID, username string) error {
|
||||
log.Warnf("failed to remove profile state file %s: %v", stateFile, err)
|
||||
}
|
||||
|
||||
prefsFile := filepath.Join(filepath.Dir(target.Path), id.String()+prefsFileSuffix)
|
||||
if err := removePrefsFile(prefsFile); err != nil && !os.IsNotExist(err) {
|
||||
log.Warnf("failed to remove profile prefs file %s: %v", prefsFile, err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -37,32 +37,23 @@
|
||||
// Updater Process (Setup):
|
||||
//
|
||||
// 1. Receives parameters from service via command-line arguments
|
||||
// 2. Terminates the UI so the installer does not have to replace a locked image
|
||||
// file, which would otherwise leave the install needing a reboot
|
||||
// 3. Runs installer with appropriate silent/quiet flags:
|
||||
// 2. Runs installer with appropriate silent/quiet flags:
|
||||
// - Windows EXE: installer.exe /S
|
||||
// - Windows MSI: msiexec.exe /i installer.msi /qn /norestart REBOOT=ReallySuppress /l*v msi.log
|
||||
// - Windows MSI: msiexec.exe /i installer.msi /quiet /qn /l*v msi.log
|
||||
// - macOS PKG: installer -pkg installer.pkg -target /
|
||||
// - macOS Homebrew: brew upgrade netbirdio/tap/netbird
|
||||
// 4. Installer terminates the daemon
|
||||
// 5. Installer replaces binaries with new version
|
||||
// 6. Updater waits for installer to complete. On Windows, MSI exit codes 3010
|
||||
// (ERROR_SUCCESS_REBOOT_REQUIRED) and 1641 (ERROR_SUCCESS_REBOOT_INITIATED)
|
||||
// are a pending-reboot outcome, not a failure: the install succeeded, but
|
||||
// some files are only replaced on the next restart (the reboot itself is
|
||||
// suppressed via /norestart and REBOOT=ReallySuppress), and the flow
|
||||
// continues as on success
|
||||
// 7. Updater restarts daemon:
|
||||
// 3. Installer terminates daemon and UI processes
|
||||
// 4. Installer replaces binaries with new version
|
||||
// 5. Updater waits for installer to complete
|
||||
// 6. Updater restarts daemon:
|
||||
// - Windows: netbird.exe service start
|
||||
// - macOS/Linux: netbird service start
|
||||
// 8. Updater restarts UI:
|
||||
// - Windows: Launches netbird-ui.exe using CreateProcessAsUser in every
|
||||
// session it was terminated in, falling back to the active console session
|
||||
// 7. Updater restarts UI:
|
||||
// - Windows: Launches netbird-ui.exe as active console user using CreateProcessAsUser
|
||||
// - macOS: Uses launchctl asuser to launch NetBird.app for console user
|
||||
// - Linux: Not implemented (UI typically auto-starts)
|
||||
// 9. Updater writes result.json with success/error status (a pending reboot is
|
||||
// recorded as success)
|
||||
// 10. Updater process exits
|
||||
// 8. Updater writes result.json with success/error status
|
||||
// 9. Updater process exits
|
||||
//
|
||||
// # Result Communication
|
||||
//
|
||||
|
||||
@@ -42,9 +42,6 @@ func NewWithDir(tempDir string) *Installer {
|
||||
// This will run by the original service process
|
||||
func (u *Installer) RunInstallation(ctx context.Context, targetVersion string) (err error) {
|
||||
resultHandler := NewResultHandler(u.tempDir)
|
||||
if err := resultHandler.ClearStaleResult(); err != nil {
|
||||
log.Warnf("clear stale installer result: %v", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if err != nil {
|
||||
|
||||
@@ -2,7 +2,6 @@ package installer
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
@@ -23,12 +22,6 @@ const (
|
||||
|
||||
msiLogFile = "msi.log"
|
||||
|
||||
// ERROR_SUCCESS_REBOOT_REQUIRED and ERROR_SUCCESS_REBOOT_INITIATED
|
||||
msiRebootRequired = 3010
|
||||
msiRebootInitiated = 1641
|
||||
|
||||
processExitWait = 10 * time.Second
|
||||
|
||||
msiDownloadURL = "https://github.com/netbirdio/netbird/releases/download/v%version/netbird_installer_%version_windows_%arch.msi"
|
||||
exeDownloadURL = "https://github.com/netbirdio/netbird/releases/download/v%version/netbird_installer_%version_windows_%arch.exe"
|
||||
)
|
||||
@@ -45,8 +38,6 @@ var (
|
||||
func (u *Installer) Setup(ctx context.Context, dryRun bool, installerFile string, daemonFolder string) (resultErr error) {
|
||||
resultHandler := NewResultHandler(u.tempDir)
|
||||
|
||||
var uiSessions []uint32
|
||||
|
||||
// Always ensure daemon and UI are restarted after setup
|
||||
defer func() {
|
||||
log.Infof("starting daemon back")
|
||||
@@ -55,7 +46,7 @@ func (u *Installer) Setup(ctx context.Context, dryRun bool, installerFile string
|
||||
}
|
||||
|
||||
log.Infof("starting UI back")
|
||||
if err := u.startUI(daemonFolder, uiSessions); err != nil {
|
||||
if err := u.startUIAsUser(daemonFolder); err != nil {
|
||||
log.Errorf("failed to start UI: %v", err)
|
||||
}
|
||||
|
||||
@@ -84,14 +75,6 @@ func (u *Installer) Setup(ctx context.Context, dryRun bool, installerFile string
|
||||
return
|
||||
}
|
||||
|
||||
// The UI holds an open handle on its own image. Left running, Restart Manager
|
||||
// cannot shut it down (msiexec runs as LocalSystem here, the UI as the
|
||||
// interactive user), so the MSI falls back to replacing the file on reboot and
|
||||
// marks the install as restart-required. The deferred close-application action
|
||||
// in the package runs too late to prevent that, it happens after
|
||||
// InstallValidate has already registered the file as in use.
|
||||
uiSessions = killUI()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
switch installerType {
|
||||
case TypeExe:
|
||||
@@ -101,9 +84,7 @@ func (u *Installer) Setup(ctx context.Context, dryRun bool, installerFile string
|
||||
installerDir := filepath.Dir(installerFile)
|
||||
logPath := filepath.Join(installerDir, msiLogFile)
|
||||
log.Infof("run msi installer: %s", installerFile)
|
||||
// REBOOT=ReallySuppress: a silent install has no way to ask, so without it
|
||||
// msiexec reboots the machine on its own if it decides one is needed.
|
||||
cmd = exec.CommandContext(ctx, "msiexec.exe", "/i", filepath.Base(installerFile), "/qn", "/norestart", "REBOOT=ReallySuppress", "/l*v", logPath)
|
||||
cmd = exec.CommandContext(ctx, "msiexec.exe", "/i", filepath.Base(installerFile), "/quiet", "/qn", "/l*v", logPath)
|
||||
}
|
||||
|
||||
cmd.Dir = filepath.Dir(installerFile)
|
||||
@@ -114,13 +95,9 @@ func (u *Installer) Setup(ctx context.Context, dryRun bool, installerFile string
|
||||
}
|
||||
|
||||
log.Infof("installer started with PID %d", cmd.Process.Pid)
|
||||
if err := cmd.Wait(); err != nil {
|
||||
if !isRebootPending(err) {
|
||||
resultErr = err
|
||||
log.Errorf("installer process finished with error: %v", err)
|
||||
return
|
||||
}
|
||||
log.Warnf("installer completed but reported a pending reboot, some files will be replaced on the next restart")
|
||||
if resultErr = cmd.Wait(); resultErr != nil {
|
||||
log.Errorf("installer process finished with error: %v", resultErr)
|
||||
return
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -140,142 +117,16 @@ func (u *Installer) startDaemon(daemonFolder string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *Installer) startUI(daemonFolder string, sessionIDs []uint32) error {
|
||||
func (u *Installer) startUIAsUser(daemonFolder string) error {
|
||||
uiPath := filepath.Join(daemonFolder, uiName)
|
||||
log.Infof("starting netbird-ui: %s", uiPath)
|
||||
|
||||
if len(sessionIDs) == 0 {
|
||||
sessionID := windows.WTSGetActiveConsoleSessionId()
|
||||
if sessionID == 0xFFFFFFFF {
|
||||
return fmt.Errorf("no active user session found")
|
||||
}
|
||||
sessionIDs = []uint32{sessionID}
|
||||
// Get the active console session ID
|
||||
sessionID := windows.WTSGetActiveConsoleSessionId()
|
||||
if sessionID == 0xFFFFFFFF {
|
||||
return fmt.Errorf("no active user session found")
|
||||
}
|
||||
|
||||
var errs []error
|
||||
for _, sessionID := range sessionIDs {
|
||||
if err := startUIInSession(uiPath, sessionID); err != nil {
|
||||
errs = append(errs, fmt.Errorf("session %d: %w", sessionID, err))
|
||||
continue
|
||||
}
|
||||
log.Infof("netbird-ui started successfully in session %d", sessionID)
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
// isRebootPending reports whether the installer exit code means it succeeded but
|
||||
// left work for the next restart. The reboot itself is suppressed, so this is not
|
||||
// a failure.
|
||||
func isRebootPending(err error) bool {
|
||||
var exitErr *exec.ExitError
|
||||
if !errors.As(err, &exitErr) {
|
||||
return false
|
||||
}
|
||||
|
||||
switch exitErr.ExitCode() {
|
||||
case msiRebootRequired, msiRebootInitiated:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// killUI terminates any running netbird-ui process and returns the IDs of the
|
||||
// interactive sessions the terminated processes belonged to. Setup starts the
|
||||
// UI again in those sessions once the installer is done.
|
||||
func killUI() []uint32 {
|
||||
pids, err := processIDsByName(uiName)
|
||||
if err != nil {
|
||||
log.Warnf("failed to look up %s processes: %v", uiName, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
sessions := make(map[uint32]struct{})
|
||||
for _, pid := range pids {
|
||||
var sessionID uint32
|
||||
if err := windows.ProcessIdToSessionId(pid, &sessionID); err != nil {
|
||||
log.Warnf("failed to look up session of %s (PID %d): %v", uiName, pid, err)
|
||||
}
|
||||
|
||||
if err := terminateProcess(pid); err != nil {
|
||||
log.Warnf("failed to terminate %s (PID %d): %v", uiName, pid, err)
|
||||
continue
|
||||
}
|
||||
log.Infof("terminated %s (PID %d) in session %d", uiName, pid, sessionID)
|
||||
|
||||
if sessionID != 0 {
|
||||
sessions[sessionID] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
sessionIDs := make([]uint32, 0, len(sessions))
|
||||
for sessionID := range sessions {
|
||||
sessionIDs = append(sessionIDs, sessionID)
|
||||
}
|
||||
return sessionIDs
|
||||
}
|
||||
|
||||
func processIDsByName(name string) ([]uint32, error) {
|
||||
snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create process snapshot: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := windows.CloseHandle(snapshot); err != nil {
|
||||
log.Warnf("failed to close process snapshot: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
var entry windows.ProcessEntry32
|
||||
entry.Size = uint32(unsafe.Sizeof(entry))
|
||||
|
||||
var pids []uint32
|
||||
for err = windows.Process32First(snapshot, &entry); err == nil; err = windows.Process32Next(snapshot, &entry) {
|
||||
if strings.EqualFold(windows.UTF16ToString(entry.ExeFile[:]), name) {
|
||||
pids = append(pids, entry.ProcessID)
|
||||
}
|
||||
}
|
||||
if !errors.Is(err, windows.ERROR_NO_MORE_FILES) {
|
||||
return nil, fmt.Errorf("enumerate processes: %w", err)
|
||||
}
|
||||
|
||||
return pids, nil
|
||||
}
|
||||
|
||||
func terminateProcess(pid uint32) error {
|
||||
handle, err := windows.OpenProcess(windows.PROCESS_TERMINATE|windows.SYNCHRONIZE, false, pid)
|
||||
if err != nil {
|
||||
// The process may have exited between enumeration and now.
|
||||
if errors.Is(err, windows.ERROR_INVALID_PARAMETER) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("open process: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := windows.CloseHandle(handle); err != nil {
|
||||
log.Warnf("failed to close process handle: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
if err := windows.TerminateProcess(handle, 0); err != nil {
|
||||
return fmt.Errorf("terminate process: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the handle to signal so the image file is released before the
|
||||
// installer tries to overwrite it. A timeout is reported through the returned
|
||||
// event, not through err, which stays nil unless the wait itself failed.
|
||||
event, err := windows.WaitForSingleObject(handle, uint32(processExitWait.Milliseconds()))
|
||||
if err != nil {
|
||||
return fmt.Errorf("wait for process exit: %w", err)
|
||||
}
|
||||
if event != windows.WAIT_OBJECT_0 {
|
||||
return fmt.Errorf("wait for process exit: unexpected wait result %#x", event)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func startUIInSession(uiPath string, sessionID uint32) error {
|
||||
// Get the user token for that session
|
||||
var userToken windows.Token
|
||||
err := windows.WTSQueryUserToken(sessionID, &userToken)
|
||||
@@ -307,16 +158,6 @@ func startUIInSession(uiPath string, sessionID uint32) error {
|
||||
}
|
||||
}()
|
||||
|
||||
var env *uint16
|
||||
if err := windows.CreateEnvironmentBlock(&env, primaryToken, false); err != nil {
|
||||
return fmt.Errorf("create environment block: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := windows.DestroyEnvironmentBlock(env); err != nil {
|
||||
log.Warnf("failed to destroy environment block: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Prepare startup info
|
||||
var si windows.StartupInfo
|
||||
si.Cb = uint32(unsafe.Sizeof(si))
|
||||
@@ -339,7 +180,7 @@ func startUIInSession(uiPath string, sessionID uint32) error {
|
||||
nil,
|
||||
false,
|
||||
creationFlags,
|
||||
env,
|
||||
nil,
|
||||
nil,
|
||||
&si,
|
||||
&pi,
|
||||
@@ -356,6 +197,7 @@ func startUIInSession(uiPath string, sessionID uint32) error {
|
||||
log.Warnf("failed to close thread handle: %v", err)
|
||||
}
|
||||
|
||||
log.Infof("netbird-ui started successfully in session %d", sessionID)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -1,108 +0,0 @@
|
||||
package installer
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os/exec"
|
||||
"slices"
|
||||
"strconv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// exitErrorWithCode returns a real *exec.ExitError carrying the given exit code.
|
||||
func exitErrorWithCode(t *testing.T, code int) error {
|
||||
t.Helper()
|
||||
|
||||
err := exec.Command("cmd.exe", "/c", "exit "+strconv.Itoa(code)).Run()
|
||||
if err == nil {
|
||||
t.Fatalf("expected a non-zero exit for code %d", code)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func TestIsRebootPending(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
code int
|
||||
want bool
|
||||
}{
|
||||
{name: "reboot required", code: msiRebootRequired, want: true},
|
||||
{name: "reboot initiated", code: msiRebootInitiated, want: true},
|
||||
{name: "generic failure", code: 1603, want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := isRebootPending(exitErrorWithCode(t, tt.code)); got != tt.want {
|
||||
t.Errorf("isRebootPending(exit %d) = %v, want %v", tt.code, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestProcessIDsByNameAndTerminate spawns a long-running system process, finds it
|
||||
// by name and terminates it, covering the path the updater uses to release the UI
|
||||
// image file before the installer replaces it.
|
||||
func TestProcessIDsByNameAndTerminate(t *testing.T) {
|
||||
cmd := exec.Command("ping.exe", "-n", "60", "127.0.0.1")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start ping: %v", err)
|
||||
}
|
||||
|
||||
pid := uint32(cmd.Process.Pid)
|
||||
killed := false
|
||||
t.Cleanup(func() {
|
||||
if !killed {
|
||||
_ = cmd.Process.Kill()
|
||||
}
|
||||
_ = cmd.Wait()
|
||||
})
|
||||
|
||||
// Name matching must be case-insensitive: the snapshot reports PING.EXE.
|
||||
pids, err := processIDsByName("ping.exe")
|
||||
if err != nil {
|
||||
t.Fatalf("processIDsByName: %v", err)
|
||||
}
|
||||
|
||||
if !slices.Contains(pids, pid) {
|
||||
t.Fatalf("PID %d not among the ping.exe processes found: %v", pid, pids)
|
||||
}
|
||||
|
||||
if err := terminateProcess(pid); err != nil {
|
||||
t.Fatalf("terminateProcess: %v", err)
|
||||
}
|
||||
killed = true
|
||||
|
||||
// terminateProcess only returns once the handle has signalled, so the process
|
||||
// is already gone and Wait must not block. It exits with the code passed to
|
||||
// TerminateProcess, which is 0, so Wait reports no error.
|
||||
if err := cmd.Wait(); err != nil {
|
||||
t.Fatalf("wait for terminated ping: %v", err)
|
||||
}
|
||||
if !cmd.ProcessState.Exited() {
|
||||
t.Error("process did not exit after terminateProcess")
|
||||
}
|
||||
|
||||
remaining, err := processIDsByName("ping.exe")
|
||||
if err != nil {
|
||||
t.Fatalf("processIDsByName after terminate: %v", err)
|
||||
}
|
||||
if slices.Contains(remaining, pid) {
|
||||
t.Errorf("PID %d still listed after terminateProcess", pid)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessIDsByNameNoMatch(t *testing.T) {
|
||||
pids, err := processIDsByName("netbird-nonexistent-process.exe")
|
||||
if err != nil {
|
||||
t.Fatalf("processIDsByName: %v", err)
|
||||
}
|
||||
if len(pids) != 0 {
|
||||
t.Errorf("expected no matches, got %v", pids)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsRebootPendingNonExitError(t *testing.T) {
|
||||
if isRebootPending(errors.New("start installer: file not found")) {
|
||||
t.Error("a non-exit error must not be treated as a pending reboot")
|
||||
}
|
||||
}
|
||||
@@ -54,12 +54,6 @@ func (rh *ResultHandler) GetErrorResultReason() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// ClearStaleResult removes a result file left over from a previous installation
|
||||
// attempt so result watchers cannot read an outdated outcome for the current attempt.
|
||||
func (rh *ResultHandler) ClearStaleResult() error {
|
||||
return rh.cleanup()
|
||||
}
|
||||
|
||||
func (rh *ResultHandler) WriteSuccess() error {
|
||||
result := Result{
|
||||
Success: true,
|
||||
|
||||
@@ -22,8 +22,6 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/listener"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
"github.com/netbirdio/netbird/client/netsweep"
|
||||
"github.com/netbirdio/netbird/client/system"
|
||||
"github.com/netbirdio/netbird/formatter"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
@@ -38,6 +36,11 @@ const (
|
||||
AnonymizeLevelStrict = nbAnonymize.LevelStrictString
|
||||
)
|
||||
|
||||
// ConnectionListener export internal Listener for mobile
|
||||
type ConnectionListener interface {
|
||||
peer.Listener
|
||||
}
|
||||
|
||||
// RouteListener export internal RouteListener for mobile
|
||||
type NetworkChangeListener interface {
|
||||
listener.NetworkChangeListener
|
||||
@@ -84,12 +87,6 @@ type Client struct {
|
||||
onHostDnsFn func([]string)
|
||||
dnsManager dns.IosDnsManager
|
||||
loginComplete bool
|
||||
// netState outlives engine restarts: it mirrors the OS connectivity, not
|
||||
// the engine lifecycle. Run injects it into each new ConnectClient, which
|
||||
// distributes it to every reconnection loop.
|
||||
netState *netstate.State
|
||||
// sweeper also outlives engine restarts; NotifyNetworkChange sweeps it.
|
||||
sweeper *netsweep.Sweeper
|
||||
// preloadedConfig holds config loaded from JSON (used on tvOS where file writes are blocked)
|
||||
preloadedConfig *profilemanager.Config
|
||||
|
||||
@@ -112,8 +109,6 @@ func NewClient(cfgFile, stateFile, cacheDir, logFilePath, deviceName string, osV
|
||||
ctxCancelLock: &sync.Mutex{},
|
||||
networkChangeListener: networkChangeListener,
|
||||
dnsManager: dnsManager,
|
||||
netState: netstate.New(),
|
||||
sweeper: netsweep.New(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -189,8 +184,7 @@ func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
|
||||
c.onHostDnsFn = func([]string) {}
|
||||
cfg.WgIface = interfaceName
|
||||
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder,
|
||||
internal.WithNetworkState(c.netState), internal.WithSweeper(c.sweeper))
|
||||
connectClient := internal.NewConnectClient(ctx, cfg, c.recorder)
|
||||
c.setState(cfg, connectClient)
|
||||
// Persist the latest sync response so DebugBundle can include the network
|
||||
// map. On iOS this is backed by disk to keep it out of the constrained
|
||||
@@ -199,25 +193,6 @@ func (c *Client) Run(fd int32, interfaceName string, envList *EnvList) error {
|
||||
return connectClient.RunOniOS(fd, c.networkChangeListener, c.dnsManager, c.stateFile, c.cacheDir, c.logFilePath)
|
||||
}
|
||||
|
||||
// SetNetworkAvailable feeds OS-reported network availability into the client
|
||||
// (e.g. from NWPathMonitor). While unavailable, the internal reconnect loops
|
||||
// suspend their attempts and the connection listener reports NoNetwork
|
||||
// instead of Connecting; when availability returns, the loops resume
|
||||
// immediately with a fresh backoff.
|
||||
func (c *Client) SetNetworkAvailable(available bool) {
|
||||
c.netState.Set(available)
|
||||
c.recorder.SetNetworkAvailable(available)
|
||||
}
|
||||
|
||||
// NotifyNetworkChange marks the management, signal and relay connections
|
||||
// stale after the OS switched networks and schedules a sweep that cuts
|
||||
// whatever has not redialed on the new network by then. The engine and the
|
||||
// TUN device stay untouched.
|
||||
func (c *Client) NotifyNetworkChange() {
|
||||
c.sweeper.MarkNetworkChange()
|
||||
log.Infof("network change: connections marked stale")
|
||||
}
|
||||
|
||||
// Stop the internal client and free the resources
|
||||
func (c *Client) Stop() {
|
||||
c.ctxCancelLock.Lock()
|
||||
@@ -356,11 +331,7 @@ func (c *Client) GetStatusDetails() *StatusDetails {
|
||||
|
||||
// SetConnectionListener set the network connection listener
|
||||
func (c *Client) SetConnectionListener(listener ConnectionListener) {
|
||||
if listener == nil {
|
||||
c.recorder.RemoveConnectionListener()
|
||||
return
|
||||
}
|
||||
c.recorder.SetConnectionListener(connectionListenerAdapter{listener})
|
||||
c.recorder.SetConnectionListener(listener)
|
||||
}
|
||||
|
||||
// RemoveConnectionListener remove connection listener
|
||||
|
||||
@@ -1,43 +0,0 @@
|
||||
//go:build ios
|
||||
|
||||
package NetBirdSDK
|
||||
|
||||
import (
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
)
|
||||
|
||||
// Client state values, re-exported as basic constants so gomobile emits them
|
||||
// into the generated bindings. They mirror peer.ClientState*: append-only,
|
||||
// never reorder.
|
||||
const (
|
||||
ClientStateDisconnected = int(peer.ClientStateDisconnected)
|
||||
ClientStateConnected = int(peer.ClientStateConnected)
|
||||
ClientStateConnecting = int(peer.ClientStateConnecting)
|
||||
ClientStateDisconnecting = int(peer.ClientStateDisconnecting)
|
||||
ClientStateNoNetwork = int(peer.ClientStateNoNetwork)
|
||||
)
|
||||
|
||||
// ConnectionListener export internal Listener for mobile.
|
||||
//
|
||||
// It intentionally lacks OnStateChanged for now: adding a method to a gomobile
|
||||
// interface breaks every Swift implementation, so the iOS app keeps building
|
||||
// against the legacy per-state callbacks. A follow-up will extend it together
|
||||
// with the app.
|
||||
type ConnectionListener interface {
|
||||
OnConnected()
|
||||
OnDisconnected()
|
||||
OnConnecting()
|
||||
OnDisconnecting()
|
||||
OnAddressChanged(string, string)
|
||||
OnPeersListChanged(int)
|
||||
}
|
||||
|
||||
// connectionListenerAdapter adapts the gomobile-facing ConnectionListener to
|
||||
// peer.Listener.
|
||||
type connectionListenerAdapter struct {
|
||||
ConnectionListener
|
||||
}
|
||||
|
||||
// OnStateChanged is dropped on iOS until the app adopts the state callback;
|
||||
// the legacy per-state callbacks continue to fire.
|
||||
func (a connectionListenerAdapter) OnStateChanged(peer.ClientState) {}
|
||||
@@ -323,7 +323,7 @@ func (a *Auth) login(urlOpener URLOpener, forceDeviceAuth bool, deviceName strin
|
||||
const authInfoRequestTimeout = 30 * time.Second
|
||||
|
||||
func (a *Auth) foregroundGetTokenInfo(authClient *auth.Auth, urlOpener URLOpener, forceDeviceAuth bool) (*auth.TokenInfo, error) {
|
||||
oAuthFlow, err := authClient.GetOAuthFlow(a.ctx, forceDeviceAuth, "")
|
||||
oAuthFlow, err := authClient.GetOAuthFlow(a.ctx, forceDeviceAuth)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get OAuth flow: %v", err)
|
||||
}
|
||||
|
||||
@@ -1,110 +0,0 @@
|
||||
// Package netstate tracks OS-reported network availability for the client.
|
||||
//
|
||||
// A State instance is owned by the platform integration (e.g. the Android or
|
||||
// iOS bindings, fed from ConnectivityManager callbacks or NWPathMonitor) and
|
||||
// is injected into the connection retry loops (management, signal, relay,
|
||||
// peer guards and the top-level connect loop), which consult it to avoid
|
||||
// burning CPU and battery on reconnect attempts while the device has no
|
||||
// network at all (e.g. airplane mode), and to reset their backoff as soon as
|
||||
// the network returns.
|
||||
//
|
||||
// Consumers hold a *State that may be nil — every non-mobile platform leaves
|
||||
// it unset. The read methods are safe on a nil receiver: they report online
|
||||
// and never block, so consumers behave as if this package did not exist.
|
||||
package netstate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// State holds the OS-reported network availability. The zero value is not
|
||||
// usable; create instances with New.
|
||||
type State struct {
|
||||
mu sync.Mutex
|
||||
online bool
|
||||
changed chan struct{}
|
||||
}
|
||||
|
||||
// New creates a State that starts online. Platforms without network tracking
|
||||
// pass a nil *State instead: the read methods treat nil as always online and
|
||||
// never block, so consumers need no nil guards.
|
||||
func New() *State {
|
||||
return &State{
|
||||
online: true,
|
||||
changed: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Set records whether the OS reports any usable network. Transitions wake up
|
||||
// all Wait callers immediately. Unlike the read methods, Set is not nil-safe:
|
||||
// it is only for the platform owner that created the State with New.
|
||||
func (s *State) Set(online bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.online == online {
|
||||
return
|
||||
}
|
||||
s.online = online
|
||||
close(s.changed)
|
||||
s.changed = make(chan struct{})
|
||||
log.Infof("OS network availability changed: online=%t", online)
|
||||
}
|
||||
|
||||
// IsOnline reports whether the OS reports at least one usable network. On a
|
||||
// nil receiver — no State injected — it reports online.
|
||||
func (s *State) IsOnline() bool {
|
||||
if s == nil {
|
||||
return true
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.online
|
||||
}
|
||||
|
||||
// Changed returns a channel closed on the next availability transition, for
|
||||
// callers that already own a select loop and cannot block in Wait. Re-read it
|
||||
// after every fire: each transition installs a fresh channel. On a nil
|
||||
// receiver — no State injected — it returns nil, which blocks forever in a
|
||||
// select, so the caller simply never observes a transition.
|
||||
func (s *State) Changed() <-chan struct{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.changed
|
||||
}
|
||||
|
||||
// Wait blocks while the network is offline. It reports whether it had to
|
||||
// wait, so callers can reset their backoff after an outage. It returns early
|
||||
// with the context error when ctx is done. On a nil receiver — no State
|
||||
// injected — it returns immediately.
|
||||
func (s *State) Wait(ctx context.Context) (bool, error) {
|
||||
if s == nil {
|
||||
return false, nil
|
||||
}
|
||||
waited := false
|
||||
for {
|
||||
s.mu.Lock()
|
||||
if s.online {
|
||||
s.mu.Unlock()
|
||||
return waited, nil
|
||||
}
|
||||
ch := s.changed
|
||||
s.mu.Unlock()
|
||||
|
||||
if !waited {
|
||||
waited = true
|
||||
log.Debugf("network is offline, pausing connection attempts")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return waited, ctx.Err()
|
||||
case <-ch:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,170 +0,0 @@
|
||||
package netstate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNewStateIsOnline(t *testing.T) {
|
||||
assert.True(t, New().IsOnline(), "a fresh State should start online")
|
||||
}
|
||||
|
||||
func TestSetTogglesOnlineState(t *testing.T) {
|
||||
s := New()
|
||||
|
||||
s.Set(false)
|
||||
assert.False(t, s.IsOnline(), "state should be offline after Set(false)")
|
||||
|
||||
s.Set(true)
|
||||
assert.True(t, s.IsOnline(), "state should be online after Set(true)")
|
||||
}
|
||||
|
||||
func TestWaitReturnsImmediatelyWhenOnline(t *testing.T) {
|
||||
s := New()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
waited, err := s.Wait(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, waited, "Wait should not block when the network is online")
|
||||
}
|
||||
|
||||
func TestWaitBlocksUntilOnline(t *testing.T) {
|
||||
s := New()
|
||||
s.Set(false)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
result := make(chan bool, 1)
|
||||
go func() {
|
||||
waited, err := s.Wait(ctx)
|
||||
if err != nil {
|
||||
result <- false
|
||||
return
|
||||
}
|
||||
result <- waited
|
||||
}()
|
||||
|
||||
// Verify Wait is actually blocking while offline
|
||||
select {
|
||||
case <-result:
|
||||
t.Fatal("Wait should block while the network is offline")
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
|
||||
s.Set(true)
|
||||
|
||||
select {
|
||||
case waited := <-result:
|
||||
assert.True(t, waited, "Wait should report that it had to wait for the network")
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Wait should return promptly after the network becomes available")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitReturnsOnContextCancel(t *testing.T) {
|
||||
s := New()
|
||||
s.Set(false)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
result := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := s.Wait(ctx)
|
||||
result <- err
|
||||
}()
|
||||
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case err := <-result:
|
||||
assert.ErrorIs(t, err, context.Canceled)
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Wait should return promptly after context cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitWakesAllWaiters(t *testing.T) {
|
||||
s := New()
|
||||
s.Set(false)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
const waiters = 10
|
||||
var wg sync.WaitGroup
|
||||
results := make(chan bool, waiters)
|
||||
for i := 0; i < waiters; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
waited, err := s.Wait(ctx)
|
||||
if err != nil {
|
||||
results <- false
|
||||
return
|
||||
}
|
||||
results <- waited
|
||||
}()
|
||||
}
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
s.Set(true)
|
||||
wg.Wait()
|
||||
|
||||
close(results)
|
||||
count := 0
|
||||
for waited := range results {
|
||||
assert.True(t, waited, "every waiter should report that it waited")
|
||||
count++
|
||||
}
|
||||
assert.Equal(t, waiters, count, "all waiters should have returned")
|
||||
}
|
||||
|
||||
func TestNilStateReadsAreNoops(t *testing.T) {
|
||||
var s *State
|
||||
|
||||
assert.True(t, s.IsOnline(), "nil State should report online")
|
||||
|
||||
waited, err := s.Wait(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.False(t, waited, "nil State's Wait should not block")
|
||||
}
|
||||
|
||||
func TestConcurrentSetAndWait(t *testing.T) {
|
||||
s := New()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 4; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 100; j++ {
|
||||
s.Set(j%2 == 0)
|
||||
s.IsOnline()
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < 4; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 100; j++ {
|
||||
if _, err := s.Wait(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
@@ -1,267 +0,0 @@
|
||||
// Package netsweep cuts network-bound activity when the OS switches networks:
|
||||
// a sweep closes the registered connections and aborts the in-flight dials, so
|
||||
// their owners redial immediately instead of waiting for the old sockets to
|
||||
// time out.
|
||||
//
|
||||
// A nil *Sweeper disables everything: all methods are nil-safe no-ops.
|
||||
package netsweep
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
// DefaultSweepDelay absorbs network flapping while the OS settles on a
|
||||
// default network before the stale registrations are cut.
|
||||
const DefaultSweepDelay = 500 * time.Millisecond
|
||||
|
||||
const recentMarkWindow = 3 * time.Second
|
||||
|
||||
// Config customizes a Sweeper. The zero value applies the defaults.
|
||||
type Config struct {
|
||||
// SweepDelay overrides DefaultSweepDelay when positive.
|
||||
SweepDelay time.Duration
|
||||
}
|
||||
|
||||
// ErrSwept reports that a dial finished after a network change swept its
|
||||
// registration. The connection is already closed; the caller must treat it
|
||||
// as a failed dial and redial on the new network.
|
||||
var ErrSwept = errors.New("netsweep: connection swept by network change")
|
||||
|
||||
// sweepID identifies one registration in a sweeper. Connections and dials
|
||||
// draw from the same counter, so an id is unique across both registries.
|
||||
type sweepID uint64
|
||||
|
||||
type connEntry struct {
|
||||
conn net.Conn
|
||||
gen uint64
|
||||
}
|
||||
|
||||
// Dial tracks one dial from start to connection registration. It hands the
|
||||
// dialed connection to the sweeper atomically, so a sweep can never fall
|
||||
// between the dial finishing and the connection being registered.
|
||||
type Dial struct {
|
||||
sweeper *Sweeper
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
id sweepID
|
||||
done bool // set by a sweep, WrapConn or Release; guarded by sweeper.mu
|
||||
gen uint64
|
||||
}
|
||||
|
||||
// Ctx returns the dial's context. A sweep cancels it, so a dial started on the
|
||||
// old network aborts instead of waiting out its handshake timeout.
|
||||
func (d *Dial) Ctx() context.Context {
|
||||
return d.ctx
|
||||
}
|
||||
|
||||
// Release ends the dial's registration and cancels its context. It is
|
||||
// idempotent and safe after WrapConn, so callers can defer it.
|
||||
func (d *Dial) Release() {
|
||||
s := d.sweeper
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
d.done = true
|
||||
delete(s.dials, d.id)
|
||||
s.mu.Unlock()
|
||||
|
||||
d.cancel()
|
||||
}
|
||||
|
||||
// sweptConn deregisters itself from the sweeper when closed.
|
||||
type sweptConn struct {
|
||||
net.Conn
|
||||
sweeper *Sweeper
|
||||
id sweepID
|
||||
}
|
||||
|
||||
func (c *sweptConn) Close() error {
|
||||
c.sweeper.deregister(c.id)
|
||||
return c.Conn.Close()
|
||||
}
|
||||
|
||||
// Sweeper registers live connections and in-flight dials so the
|
||||
// network-change sweep can cut everything registered before the change.
|
||||
type Sweeper struct {
|
||||
mu sync.Mutex
|
||||
conns map[sweepID]connEntry
|
||||
dials map[sweepID]*Dial
|
||||
nextID sweepID
|
||||
gen uint64
|
||||
timer *time.Timer
|
||||
sweepDelay time.Duration
|
||||
lastMark time.Time
|
||||
}
|
||||
|
||||
// New creates an empty sweeper with the default configuration.
|
||||
func New() *Sweeper {
|
||||
return NewWithConfig(Config{})
|
||||
}
|
||||
|
||||
// NewWithConfig creates an empty sweeper customized by cfg.
|
||||
func NewWithConfig(cfg Config) *Sweeper {
|
||||
delay := cfg.SweepDelay
|
||||
if delay <= 0 {
|
||||
delay = DefaultSweepDelay
|
||||
}
|
||||
return &Sweeper{
|
||||
conns: make(map[sweepID]connEntry),
|
||||
dials: make(map[sweepID]*Dial),
|
||||
sweepDelay: delay,
|
||||
}
|
||||
}
|
||||
|
||||
// StartDial registers an in-flight dial. Dial with Ctx, hand the result to
|
||||
// WrapConn, and Release the dial when the attempt is over, typically deferred.
|
||||
func (s *Sweeper) StartDial(ctx context.Context) *Dial {
|
||||
if s == nil {
|
||||
return &Dial{ctx: ctx}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
d := &Dial{sweeper: s, ctx: ctx, cancel: cancel}
|
||||
|
||||
s.mu.Lock()
|
||||
d.id = s.nextID
|
||||
s.nextID++
|
||||
d.gen = s.gen
|
||||
s.dials[d.id] = d
|
||||
s.mu.Unlock()
|
||||
|
||||
return d
|
||||
}
|
||||
|
||||
// WrapConn hands conn over to the sweeper. If a sweep ran since StartDial,
|
||||
// the connection belongs to the old network: it is closed and ErrSwept is
|
||||
// returned. Otherwise conn is registered against the next sweep and returned
|
||||
// wrapped, deregistering itself on Close. Call it once, before Release.
|
||||
func (d *Dial) WrapConn(conn net.Conn) (net.Conn, error) {
|
||||
s := d.sweeper
|
||||
if s == nil {
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if d.done {
|
||||
s.mu.Unlock()
|
||||
if err := conn.Close(); err != nil {
|
||||
log.Debugf("swept dial close error: %v", err)
|
||||
}
|
||||
return nil, ErrSwept
|
||||
}
|
||||
d.done = true
|
||||
delete(s.dials, d.id)
|
||||
id := s.nextID
|
||||
s.nextID++
|
||||
// The conn inherits the dial's generation: the socket was bound to the
|
||||
// network that was default when the dial started, not when it finished.
|
||||
s.conns[id] = connEntry{conn: conn, gen: d.gen}
|
||||
s.mu.Unlock()
|
||||
|
||||
return &sweptConn{Conn: conn, sweeper: s, id: id}, nil
|
||||
}
|
||||
|
||||
// MarkNetworkChange records that the OS switched networks: everything
|
||||
// registered so far becomes stale, and a sweep is (re)scheduled after the
|
||||
// configured delay to cut whatever is still stale by then. Owners that
|
||||
// redialed in the meantime hold fresh-generation registrations and survive,
|
||||
// so no cancellation is needed around the sweep.
|
||||
func (s *Sweeper) MarkNetworkChange() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
s.gen++
|
||||
cutoff := s.gen
|
||||
s.lastMark = time.Now()
|
||||
if s.timer != nil {
|
||||
s.timer.Stop()
|
||||
}
|
||||
s.timer = time.AfterFunc(s.sweepDelay, func() {
|
||||
n := s.sweep(cutoff)
|
||||
log.Infof("network change sweep: closed %d stale connections", n)
|
||||
})
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// QuickRetryBackoff wraps bo so that after each Reset the first retry comes
|
||||
// quickly when the disconnect followed a recent network change and the
|
||||
// network is online. Any other failure keeps bo's spread, so the clients of
|
||||
// a restarted server still scatter their reconnects. A nil sweeper returns
|
||||
// bo unchanged.
|
||||
func (s *Sweeper) QuickRetryBackoff(ctx context.Context, bo backoff.BackOff, netState *netstate.State) backoff.BackOff {
|
||||
if s == nil {
|
||||
return bo
|
||||
}
|
||||
return backoff.WithContext(newQuickRetryBackoff(bo, s, netState), ctx)
|
||||
}
|
||||
|
||||
func (s *Sweeper) markedRecently() bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return !s.lastMark.IsZero() && time.Since(s.lastMark) < recentMarkWindow
|
||||
}
|
||||
|
||||
// sweep closes the registered connections and aborts the in-flight dials
|
||||
// older than cutoff, and returns how many connections it closed. A dial
|
||||
// whose connection was not yet handed to WrapConn is marked, so the late
|
||||
// WrapConn closes it instead of registering it.
|
||||
func (s *Sweeper) sweep(cutoff uint64) int {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
var conns []net.Conn
|
||||
for id, e := range s.conns {
|
||||
if e.gen < cutoff {
|
||||
delete(s.conns, id)
|
||||
conns = append(conns, e.conn)
|
||||
}
|
||||
}
|
||||
var dials []*Dial
|
||||
for id, d := range s.dials {
|
||||
if d.gen < cutoff {
|
||||
d.done = true
|
||||
delete(s.dials, id)
|
||||
dials = append(dials, d)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
if len(dials) > 0 {
|
||||
log.Debugf("aborting %d in-flight dials", len(dials))
|
||||
for _, d := range dials {
|
||||
d.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
for _, conn := range conns {
|
||||
log.Debugf("sweeping connection %s -> %s", conn.LocalAddr(), conn.RemoteAddr())
|
||||
if err := conn.Close(); err != nil {
|
||||
log.Debugf("swept connection close error: %v", err)
|
||||
}
|
||||
}
|
||||
return len(conns)
|
||||
}
|
||||
|
||||
func (s *Sweeper) deregister(id sweepID) {
|
||||
s.mu.Lock()
|
||||
delete(s.conns, id)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
@@ -1,241 +0,0 @@
|
||||
package netsweep
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSweepClosesRegisteredConns(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
c1 := wrap(t, sweeper, connPair(t))
|
||||
c2 := wrap(t, sweeper, connPair(t))
|
||||
|
||||
assert.Equal(t, 2, sweeper.sweepAll(), "both live connections should be closed")
|
||||
|
||||
// The wrappers must report closed now.
|
||||
buf := make([]byte, 1)
|
||||
_, err := c1.Read(buf)
|
||||
assert.Error(t, err, "first connection should be unusable after the sweep")
|
||||
_, err = c2.Read(buf)
|
||||
assert.Error(t, err, "second connection should be unusable after the sweep")
|
||||
|
||||
assert.Equal(t, 0, sweeper.sweepAll(), "second sweep should find nothing")
|
||||
}
|
||||
|
||||
func TestCloseDeregisters(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
conn := wrap(t, sweeper, connPair(t))
|
||||
require.NoError(t, conn.Close())
|
||||
|
||||
assert.Equal(t, 0, sweeper.sweepAll(), "closed connection must leave the registry")
|
||||
}
|
||||
|
||||
func TestCloseIsIdempotent(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
conn := wrap(t, sweeper, connPair(t))
|
||||
require.NoError(t, conn.Close())
|
||||
assert.Error(t, conn.Close(), "double close surfaces the underlying error but must not panic")
|
||||
}
|
||||
|
||||
func TestSweepOnlyAffectsOlderConns(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
_ = wrap(t, sweeper, connPair(t))
|
||||
assert.Equal(t, 1, sweeper.sweepAll())
|
||||
|
||||
// A connection dialed after the sweep must survive until the next one.
|
||||
_ = wrap(t, sweeper, connPair(t))
|
||||
assert.Equal(t, 1, sweeper.sweepAll(), "post-sweep connection belongs to the next sweep")
|
||||
}
|
||||
|
||||
func TestSweepAbortsInFlightDials(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
dial := sweeper.StartDial(context.Background())
|
||||
defer dial.Release()
|
||||
|
||||
sweeper.sweepAll()
|
||||
|
||||
assert.ErrorIs(t, dial.Ctx().Err(), context.Canceled, "sweep must cancel the in-flight dial context")
|
||||
}
|
||||
|
||||
func TestReleasedDialIsNotAborted(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
// Simulate a dial that finished before the sweep.
|
||||
released := sweeper.StartDial(context.Background())
|
||||
released.Release()
|
||||
|
||||
// A dial still in flight during the sweep.
|
||||
pending := sweeper.StartDial(context.Background())
|
||||
defer pending.Release()
|
||||
|
||||
sweeper.sweepAll()
|
||||
assert.ErrorIs(t, pending.Ctx().Err(), context.Canceled, "pending dial must be aborted")
|
||||
}
|
||||
|
||||
func TestSweepBetweenDialAndHandoffClosesConn(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
dial := sweeper.StartDial(context.Background())
|
||||
defer dial.Release()
|
||||
|
||||
// The dial succeeds on the old network, then the sweep lands before the
|
||||
// connection is handed over.
|
||||
conn := connPair(t)
|
||||
assert.Equal(t, 0, sweeper.sweepAll(), "the connection is not registered yet")
|
||||
|
||||
wrapped, err := dial.WrapConn(conn)
|
||||
require.ErrorIs(t, err, ErrSwept)
|
||||
require.Nil(t, wrapped)
|
||||
|
||||
buf := make([]byte, 1)
|
||||
_, err = conn.Read(buf)
|
||||
assert.Error(t, err, "the old-network connection must be closed, not leaked")
|
||||
|
||||
assert.Equal(t, 0, sweeper.sweepAll(), "nothing may leak into the next sweep")
|
||||
}
|
||||
|
||||
func TestMarkNetworkChangeSparesFreshConns(t *testing.T) {
|
||||
sweeper := NewWithConfig(Config{SweepDelay: 10 * time.Millisecond})
|
||||
|
||||
stale := wrap(t, sweeper, connPair(t))
|
||||
sweeper.MarkNetworkChange()
|
||||
_ = wrap(t, sweeper, connPair(t))
|
||||
|
||||
_ = stale.SetReadDeadline(time.Now().Add(time.Second))
|
||||
buf := make([]byte, 1)
|
||||
_, err := stale.Read(buf)
|
||||
require.ErrorIs(t, err, net.ErrClosed, "stale connection must be closed by the delayed sweep")
|
||||
|
||||
assert.Equal(t, 1, sweeper.sweepAll(), "the fresh connection must survive the stale sweep")
|
||||
}
|
||||
|
||||
func TestMarkNetworkChangeAbortsStaleDials(t *testing.T) {
|
||||
sweeper := NewWithConfig(Config{SweepDelay: 10 * time.Millisecond})
|
||||
|
||||
stale := sweeper.StartDial(context.Background())
|
||||
defer stale.Release()
|
||||
sweeper.MarkNetworkChange()
|
||||
fresh := sweeper.StartDial(context.Background())
|
||||
defer fresh.Release()
|
||||
|
||||
assert.Eventually(t, func() bool {
|
||||
return stale.Ctx().Err() != nil
|
||||
}, time.Second, 5*time.Millisecond, "stale dial must be aborted by the delayed sweep")
|
||||
assert.NoError(t, fresh.Ctx().Err(), "post-mark dial must not be aborted")
|
||||
}
|
||||
|
||||
func TestConnInheritsDialGeneration(t *testing.T) {
|
||||
sweeper := NewWithConfig(Config{SweepDelay: 20 * time.Millisecond})
|
||||
|
||||
// The dial starts before the network change but completes after it: the
|
||||
// socket is bound to the old network, so the sweep must still cut it.
|
||||
dial := sweeper.StartDial(context.Background())
|
||||
defer dial.Release()
|
||||
sweeper.MarkNetworkChange()
|
||||
|
||||
wrapped, err := dial.WrapConn(connPair(t))
|
||||
require.NoError(t, err)
|
||||
|
||||
_ = wrapped.SetReadDeadline(time.Now().Add(time.Second))
|
||||
buf := make([]byte, 1)
|
||||
_, err = wrapped.Read(buf)
|
||||
require.ErrorIs(t, err, net.ErrClosed, "old-generation connection must be swept")
|
||||
}
|
||||
|
||||
func TestRepeatedMarksCoalesce(t *testing.T) {
|
||||
sweeper := NewWithConfig(Config{SweepDelay: 20 * time.Millisecond})
|
||||
|
||||
first := wrap(t, sweeper, connPair(t))
|
||||
sweeper.MarkNetworkChange()
|
||||
second := wrap(t, sweeper, connPair(t))
|
||||
sweeper.MarkNetworkChange()
|
||||
_ = wrap(t, sweeper, connPair(t))
|
||||
|
||||
buf := make([]byte, 1)
|
||||
for _, conn := range []net.Conn{first, second} {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(time.Second))
|
||||
_, err := conn.Read(buf)
|
||||
require.ErrorIs(t, err, net.ErrClosed, "every pre-mark connection must be swept by the rescheduled sweep")
|
||||
}
|
||||
assert.Equal(t, 1, sweeper.sweepAll(), "only the newest-generation connection may remain")
|
||||
}
|
||||
|
||||
func TestNilSweeperIsNoop(t *testing.T) {
|
||||
var sweeper *Sweeper
|
||||
|
||||
conn := connPair(t)
|
||||
dial := sweeper.StartDial(context.Background())
|
||||
defer dial.Release()
|
||||
|
||||
wrapped, err := dial.WrapConn(conn)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, conn, wrapped, "nil sweeper must return the conn unchanged")
|
||||
assert.NoError(t, dial.Ctx().Err(), "nil sweeper must not cancel the dial context")
|
||||
assert.Equal(t, 0, sweeper.sweepAll(), "nil sweeper closes nothing")
|
||||
}
|
||||
|
||||
// wrap registers conn with the sweeper through a completed dial.
|
||||
func wrap(t *testing.T, sweeper *Sweeper, conn net.Conn) net.Conn {
|
||||
t.Helper()
|
||||
|
||||
dial := sweeper.StartDial(context.Background())
|
||||
defer dial.Release()
|
||||
|
||||
wrapped, err := dial.WrapConn(conn)
|
||||
require.NoError(t, err)
|
||||
return wrapped
|
||||
}
|
||||
|
||||
// connPair dials a loopback TCP connection and keeps the accepted peer open
|
||||
// until the test ends: a peer that closed early would make the connection
|
||||
// unreadable on its own, so a read error after the sweep would prove nothing.
|
||||
func connPair(t *testing.T) net.Conn {
|
||||
t.Helper()
|
||||
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
if err := l.Close(); err != nil {
|
||||
t.Logf("listener close error: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
accepted := make(chan net.Conn, 1)
|
||||
go func() {
|
||||
conn, err := l.Accept()
|
||||
if err != nil {
|
||||
close(accepted)
|
||||
return
|
||||
}
|
||||
accepted <- conn
|
||||
}()
|
||||
|
||||
conn, err := net.Dial("tcp", l.Addr().String())
|
||||
require.NoError(t, err)
|
||||
|
||||
peer, ok := <-accepted
|
||||
require.True(t, ok, "listener must accept the dialed connection")
|
||||
t.Cleanup(func() {
|
||||
if err := peer.Close(); err != nil {
|
||||
t.Logf("peer close error: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
return conn
|
||||
}
|
||||
|
||||
// sweepAll cuts every registration regardless of generation.
|
||||
func (s *Sweeper) sweepAll() int {
|
||||
return s.sweep(math.MaxUint64)
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
package netsweep
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
|
||||
"github.com/netbirdio/netbird/client/netstate"
|
||||
)
|
||||
|
||||
const quickRetryDelay = 200 * time.Millisecond
|
||||
|
||||
type quickRetryBackoff struct {
|
||||
backoff.BackOff
|
||||
sweeper *Sweeper
|
||||
netState *netstate.State
|
||||
used bool
|
||||
}
|
||||
|
||||
func newQuickRetryBackoff(bo backoff.BackOff, sweeper *Sweeper, netState *netstate.State) *quickRetryBackoff {
|
||||
return &quickRetryBackoff{
|
||||
BackOff: bo,
|
||||
sweeper: sweeper,
|
||||
netState: netState,
|
||||
}
|
||||
}
|
||||
|
||||
func (b *quickRetryBackoff) NextBackOff() time.Duration {
|
||||
if !b.used && b.sweeper.markedRecently() && b.netState.IsOnline() {
|
||||
b.used = true
|
||||
return quickRetryDelay
|
||||
}
|
||||
return b.BackOff.NextBackOff()
|
||||
}
|
||||
|
||||
func (b *quickRetryBackoff) Reset() {
|
||||
b.used = false
|
||||
b.BackOff.Reset()
|
||||
}
|
||||
@@ -1,58 +0,0 @@
|
||||
package netsweep
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff/v4"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestQuickRetryAfterRecentMark(t *testing.T) {
|
||||
sweeper := New()
|
||||
sweeper.MarkNetworkChange()
|
||||
|
||||
slow := backoff.NewConstantBackOff(5 * time.Second)
|
||||
bo := sweeper.QuickRetryBackoff(context.Background(), slow, nil)
|
||||
|
||||
assert.Equal(t, quickRetryDelay, bo.NextBackOff(), "first retry after a mark must be quick")
|
||||
assert.Equal(t, 5*time.Second, bo.NextBackOff(), "second retry must fall back to the wrapped backoff")
|
||||
|
||||
bo.Reset()
|
||||
assert.Equal(t, quickRetryDelay, bo.NextBackOff(), "reset must re-arm the quick retry")
|
||||
}
|
||||
|
||||
func TestQuickRetryWithoutMarkKeepsSpread(t *testing.T) {
|
||||
sweeper := New()
|
||||
|
||||
slow := backoff.NewConstantBackOff(5 * time.Second)
|
||||
bo := sweeper.QuickRetryBackoff(context.Background(), slow, nil)
|
||||
|
||||
assert.Equal(t, 5*time.Second, bo.NextBackOff(), "without a mark the wrapped backoff decides")
|
||||
|
||||
sweeper.mu.Lock()
|
||||
sweeper.lastMark = time.Now().Add(-recentMarkWindow)
|
||||
sweeper.mu.Unlock()
|
||||
assert.Equal(t, 5*time.Second, bo.NextBackOff(), "a stale mark must not trigger the quick retry")
|
||||
}
|
||||
|
||||
func TestQuickRetryNilSweeperPassthrough(t *testing.T) {
|
||||
var sweeper *Sweeper
|
||||
|
||||
slow := backoff.NewConstantBackOff(5 * time.Second)
|
||||
bo := sweeper.QuickRetryBackoff(context.Background(), slow, nil)
|
||||
|
||||
assert.Equal(t, backoff.BackOff(slow), bo, "nil sweeper must return the backoff unchanged")
|
||||
}
|
||||
|
||||
func TestQuickRetryHonorsContext(t *testing.T) {
|
||||
sweeper := New()
|
||||
sweeper.MarkNetworkChange()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
bo := sweeper.QuickRetryBackoff(ctx, backoff.NewConstantBackOff(time.Millisecond), nil)
|
||||
|
||||
assert.Equal(t, backoff.Stop, bo.NextBackOff(), "cancelled context must stop the retry loop")
|
||||
}
|
||||
@@ -313,23 +313,21 @@ func Dial(ctx context.Context, addr, user string, opts DialOptions) (*Client, er
|
||||
|
||||
// dialSSH establishes an SSH connection without JWT authentication
|
||||
func dialSSH(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*Client, error) {
|
||||
if config.Timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, config.Timeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
dialer := &net.Dialer{}
|
||||
conn, err := dialer.DialContext(ctx, network, addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dial %s: %w", addr, err)
|
||||
}
|
||||
|
||||
client, err := nbssh.Handshake(ctx, conn, addr, config)
|
||||
clientConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
if closeErr := conn.Close(); closeErr != nil {
|
||||
log.Debugf("connection close after handshake failure: %v", closeErr)
|
||||
}
|
||||
return nil, fmt.Errorf("ssh handshake: %w", err)
|
||||
}
|
||||
|
||||
client := ssh.NewClient(clientConn, chans, reqs)
|
||||
return &Client{
|
||||
client: client,
|
||||
}, nil
|
||||
|
||||
@@ -12,8 +12,6 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/ssh"
|
||||
"golang.org/x/term"
|
||||
|
||||
nbssh "github.com/netbirdio/netbird/client/ssh"
|
||||
)
|
||||
|
||||
func (c *Client) setupTerminalMode(ctx context.Context, session *ssh.Session) error {
|
||||
@@ -84,7 +82,37 @@ func (c *Client) setupTerminal(session *ssh.Session, fd int) error {
|
||||
return fmt.Errorf("get terminal size: %w", err)
|
||||
}
|
||||
|
||||
modes := nbssh.DefaultTerminalModes
|
||||
modes := ssh.TerminalModes{
|
||||
ssh.ECHO: 1,
|
||||
ssh.TTY_OP_ISPEED: 14400,
|
||||
ssh.TTY_OP_OSPEED: 14400,
|
||||
// Ctrl+C
|
||||
ssh.VINTR: 3,
|
||||
// Ctrl+\
|
||||
ssh.VQUIT: 28,
|
||||
// Backspace
|
||||
ssh.VERASE: 127,
|
||||
// Ctrl+U
|
||||
ssh.VKILL: 21,
|
||||
// Ctrl+D
|
||||
ssh.VEOF: 4,
|
||||
ssh.VEOL: 0,
|
||||
ssh.VEOL2: 0,
|
||||
// Ctrl+Q
|
||||
ssh.VSTART: 17,
|
||||
// Ctrl+S
|
||||
ssh.VSTOP: 19,
|
||||
// Ctrl+Z
|
||||
ssh.VSUSP: 26,
|
||||
// Ctrl+O
|
||||
ssh.VDISCARD: 15,
|
||||
// Ctrl+R
|
||||
ssh.VREPRINT: 18,
|
||||
// Ctrl+W
|
||||
ssh.VWERASE: 23,
|
||||
// Ctrl+V
|
||||
ssh.VLNEXT: 22,
|
||||
}
|
||||
|
||||
terminal := os.Getenv("TERM")
|
||||
if terminal == "" {
|
||||
|
||||
@@ -10,8 +10,6 @@ import (
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/ssh"
|
||||
|
||||
nbssh "github.com/netbirdio/netbird/client/ssh"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -82,14 +80,28 @@ func (c *Client) setupTerminalMode(_ context.Context, session *ssh.Session) erro
|
||||
w, h := c.getWindowsConsoleSize()
|
||||
|
||||
modes := ssh.TerminalModes{
|
||||
ssh.ICRNL: 1,
|
||||
ssh.OPOST: 1,
|
||||
ssh.ONLCR: 1,
|
||||
ssh.ISIG: 1,
|
||||
ssh.ICANON: 1,
|
||||
}
|
||||
for mode, value := range nbssh.DefaultTerminalModes {
|
||||
modes[mode] = value
|
||||
ssh.ECHO: 1,
|
||||
ssh.TTY_OP_ISPEED: 14400,
|
||||
ssh.TTY_OP_OSPEED: 14400,
|
||||
ssh.ICRNL: 1,
|
||||
ssh.OPOST: 1,
|
||||
ssh.ONLCR: 1,
|
||||
ssh.ISIG: 1,
|
||||
ssh.ICANON: 1,
|
||||
ssh.VINTR: 3, // Ctrl+C
|
||||
ssh.VQUIT: 28, // Ctrl+\
|
||||
ssh.VERASE: 127, // Backspace
|
||||
ssh.VKILL: 21, // Ctrl+U
|
||||
ssh.VEOF: 4, // Ctrl+D
|
||||
ssh.VEOL: 0,
|
||||
ssh.VEOL2: 0,
|
||||
ssh.VSTART: 17, // Ctrl+Q
|
||||
ssh.VSTOP: 19, // Ctrl+S
|
||||
ssh.VSUSP: 26, // Ctrl+Z
|
||||
ssh.VDISCARD: 15, // Ctrl+O
|
||||
ssh.VWERASE: 23, // Ctrl+W
|
||||
ssh.VLNEXT: 22, // Ctrl+V
|
||||
ssh.VREPRINT: 18, // Ctrl+R
|
||||
}
|
||||
|
||||
if err := session.RequestPty("xterm-256color", h, w, modes); err != nil {
|
||||
|
||||
@@ -35,19 +35,6 @@ type HostKeyVerifier interface {
|
||||
VerifySSHHostKey(peerAddress string, key []byte) error
|
||||
}
|
||||
|
||||
// PeerKeyLookup returns the stored SSH host key for a peer address.
|
||||
type PeerKeyLookup func(peerAddress string) ([]byte, bool)
|
||||
|
||||
// VerifySSHHostKey implements HostKeyVerifier by looking up the stored key
|
||||
// and comparing it against the presented key.
|
||||
func (l PeerKeyLookup) VerifySSHHostKey(peerAddress string, presentedKey []byte) error {
|
||||
storedKey, found := l(peerAddress)
|
||||
if !found {
|
||||
return ErrPeerNotFound
|
||||
}
|
||||
return VerifyHostKey(storedKey, presentedKey, peerAddress)
|
||||
}
|
||||
|
||||
// DaemonHostKeyVerifier implements HostKeyVerifier using the NetBird daemon
|
||||
type DaemonHostKeyVerifier struct {
|
||||
client proto.DaemonServiceClient
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
// Handshake runs the SSH client handshake on an already dialed conn and
|
||||
// returns the resulting client. Dialing bounds only the TCP establishment;
|
||||
// without a deadline on the socket a peer that accepts and then goes silent
|
||||
// blocks the handshake forever, so the context deadline is applied to conn
|
||||
// for the duration of the handshake. conn is closed on any error.
|
||||
func Handshake(ctx context.Context, conn net.Conn, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
if err := conn.SetDeadline(deadline); err != nil {
|
||||
closeHandshake(conn, "conn after deadline error")
|
||||
return nil, fmt.Errorf("set handshake deadline: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
|
||||
if err != nil {
|
||||
closeHandshake(conn, "conn after handshake error")
|
||||
return nil, fmt.Errorf("ssh handshake: %w", err)
|
||||
}
|
||||
|
||||
if err := conn.SetDeadline(time.Time{}); err != nil {
|
||||
closeHandshake(sshConn, "ssh conn after deadline clear error")
|
||||
return nil, fmt.Errorf("clear handshake deadline: %w", err)
|
||||
}
|
||||
|
||||
return ssh.NewClient(sshConn, chans, reqs), nil
|
||||
}
|
||||
|
||||
func closeHandshake(c io.Closer, label string) {
|
||||
if err := c.Close(); err != nil {
|
||||
log.Debugf("ssh: close %s: %v", label, err)
|
||||
}
|
||||
}
|
||||
@@ -610,10 +610,13 @@ func (p *SSHProxy) dialBackend(ctx context.Context, addr, user, jwtToken string)
|
||||
return nil, fmt.Errorf("connect to server: %w", err)
|
||||
}
|
||||
|
||||
handshakeCtx, cancel := context.WithTimeout(ctx, sshHandshakeTimeout)
|
||||
defer cancel()
|
||||
clientConn, chans, reqs, err := cryptossh.NewClientConn(conn, addr, config)
|
||||
if err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, fmt.Errorf("SSH handshake: %w", err)
|
||||
}
|
||||
|
||||
return nbssh.Handshake(handshakeCtx, conn, addr, config)
|
||||
return cryptossh.NewClient(clientConn, chans, reqs), nil
|
||||
}
|
||||
|
||||
func (p *SSHProxy) verifyHostKey(hostname string, remote net.Addr, key cryptossh.PublicKey) error {
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
// DefaultTerminalModes are the PTY modes used by the interactive terminal clients.
|
||||
var DefaultTerminalModes = ssh.TerminalModes{
|
||||
ssh.ECHO: 1,
|
||||
ssh.TTY_OP_ISPEED: 14400,
|
||||
ssh.TTY_OP_OSPEED: 14400,
|
||||
ssh.VINTR: 3, // Ctrl+C
|
||||
ssh.VQUIT: 28, // Ctrl+\
|
||||
ssh.VERASE: 127, // Backspace
|
||||
ssh.VKILL: 21, // Ctrl+U
|
||||
ssh.VEOF: 4, // Ctrl+D
|
||||
ssh.VEOL: 0,
|
||||
ssh.VEOL2: 0,
|
||||
ssh.VSTART: 17, // Ctrl+Q
|
||||
ssh.VSTOP: 19, // Ctrl+S
|
||||
ssh.VSUSP: 26, // Ctrl+Z
|
||||
ssh.VDISCARD: 15, // Ctrl+O
|
||||
ssh.VREPRINT: 18, // Ctrl+R
|
||||
ssh.VWERASE: 23, // Ctrl+W
|
||||
ssh.VLNEXT: 22, // Ctrl+V
|
||||
}
|
||||
|
||||
// PTYSession is an interactive shell session with a PTY and its I/O pipes.
|
||||
type PTYSession struct {
|
||||
Session *ssh.Session
|
||||
Stdin io.WriteCloser
|
||||
Stdout io.Reader
|
||||
Stderr io.Reader
|
||||
}
|
||||
|
||||
// StartPTYSession opens a session on the client, requests an xterm-256color PTY
|
||||
// with the default terminal modes, wires up the I/O pipes and starts a shell.
|
||||
// The session is closed on any error.
|
||||
func StartPTYSession(client *ssh.Client, cols, rows int) (*PTYSession, error) {
|
||||
session, err := client.NewSession()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("new session: %w", err)
|
||||
}
|
||||
|
||||
pty, err := setupPTYSession(session, cols, rows)
|
||||
if err != nil {
|
||||
if closeErr := session.Close(); closeErr != nil {
|
||||
log.Debugf("ssh: session close after setup error: %v", closeErr)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return pty, nil
|
||||
}
|
||||
|
||||
// setupPTYSession requests the PTY, opens the pipes and starts the shell on an
|
||||
// already created session.
|
||||
func setupPTYSession(session *ssh.Session, cols, rows int) (*PTYSession, error) {
|
||||
if err := session.RequestPty("xterm-256color", rows, cols, DefaultTerminalModes); err != nil {
|
||||
return nil, fmt.Errorf("request pty: %w", err)
|
||||
}
|
||||
|
||||
stdin, err := session.StdinPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stdin pipe: %w", err)
|
||||
}
|
||||
stdout, err := session.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stdout pipe: %w", err)
|
||||
}
|
||||
stderr, err := session.StderrPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stderr pipe: %w", err)
|
||||
}
|
||||
|
||||
if err := session.Shell(); err != nil {
|
||||
return nil, fmt.Errorf("start shell: %w", err)
|
||||
}
|
||||
|
||||
return &PTYSession{Session: session, Stdin: stdin, Stdout: stdout, Stderr: stderr}, nil
|
||||
}
|
||||
@@ -30,11 +30,6 @@ func GetInfo(ctx context.Context) *Info {
|
||||
kernelVersion = osInfo[2]
|
||||
}
|
||||
|
||||
addrs, err := networkAddresses()
|
||||
if err != nil {
|
||||
log.Warnf("discover network addresses: %s", err)
|
||||
}
|
||||
|
||||
gio := &Info{
|
||||
GoOS: runtime.GOOS,
|
||||
Kernel: kernel,
|
||||
@@ -46,7 +41,6 @@ func GetInfo(ctx context.Context) *Info {
|
||||
NetbirdVersion: version.NetbirdVersion(),
|
||||
UIVersion: extractUIVersion(ctx),
|
||||
KernelVersion: kernelVersion,
|
||||
NetworkAddresses: addrs,
|
||||
SystemSerialNumber: serial(),
|
||||
SystemProductName: productModel(),
|
||||
SystemManufacturer: productManufacturer(),
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build !ios && !android
|
||||
//go:build !ios
|
||||
|
||||
package system
|
||||
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var iFaceDiscover IFaceDiscover
|
||||
|
||||
type IFaceDiscover interface {
|
||||
IFaces() (string, error)
|
||||
}
|
||||
|
||||
// SetIFaceDiscover configures the Android interface discovery provider.
|
||||
func SetIFaceDiscover(discover IFaceDiscover) {
|
||||
iFaceDiscover = discover
|
||||
}
|
||||
|
||||
func networkAddresses() ([]NetworkAddress, error) {
|
||||
if iFaceDiscover == nil {
|
||||
return nil, nil
|
||||
}
|
||||
ifaces, err := iFaceDiscover.IFaces()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var netAddresses []NetworkAddress
|
||||
for _, line := range strings.Split(ifaces, "\n") {
|
||||
addresses, ok := interfaceAddresses(line)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, address := range addresses {
|
||||
netAddr, ok := toNetworkAddress(address)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if isDuplicated(netAddresses, netAddr) {
|
||||
continue
|
||||
}
|
||||
netAddresses = append(netAddresses, netAddr)
|
||||
}
|
||||
}
|
||||
return netAddresses, nil
|
||||
}
|
||||
|
||||
func interfaceAddresses(line string) ([]string, bool) {
|
||||
parts := strings.Split(line, "|")
|
||||
if len(parts) != 2 {
|
||||
return nil, false
|
||||
}
|
||||
flags := strings.Fields(parts[0])
|
||||
if len(flags) != 8 {
|
||||
return nil, false
|
||||
}
|
||||
up, loopback := flags[3], flags[5]
|
||||
if up != "true" || loopback == "true" {
|
||||
return nil, false
|
||||
}
|
||||
return strings.Fields(parts[1]), true
|
||||
}
|
||||
|
||||
func toNetworkAddress(address string) (NetworkAddress, bool) {
|
||||
prefix, err := netip.ParsePrefix(address)
|
||||
if err != nil {
|
||||
return NetworkAddress{}, false
|
||||
}
|
||||
if prefix.Addr().Is4In6() {
|
||||
if prefix.Bits() < 96 {
|
||||
return NetworkAddress{}, false
|
||||
}
|
||||
prefix = netip.PrefixFrom(prefix.Addr().Unmap(), prefix.Bits()-96)
|
||||
}
|
||||
ip := prefix.Addr()
|
||||
if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsMulticast() {
|
||||
return NetworkAddress{}, false
|
||||
}
|
||||
return NetworkAddress{NetIP: prefix}, true
|
||||
}
|
||||
|
||||
func isDuplicated(addresses []NetworkAddress, addr NetworkAddress) bool {
|
||||
for _, duplicated := range addresses {
|
||||
if duplicated.NetIP == addr.NetIP {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build !ios && !android
|
||||
//go:build !ios
|
||||
|
||||
package system
|
||||
|
||||
|
||||
@@ -80,12 +80,13 @@ func (c *Client) Connect(host string, port int, username, jwtToken string, ipVer
|
||||
return fmt.Errorf("dial %s: %w", addr, err)
|
||||
}
|
||||
|
||||
sshClient, err := nbssh.Handshake(ctx, conn, addr, config)
|
||||
sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
|
||||
if err != nil {
|
||||
return err
|
||||
closeWithLog(conn, "connection after handshake error")
|
||||
return fmt.Errorf("SSH handshake: %w", err)
|
||||
}
|
||||
|
||||
c.sshClient = sshClient
|
||||
c.sshClient = ssh.NewClient(sshConn, chans, reqs)
|
||||
logrus.Infof("SSH: Connected to %s", addr)
|
||||
|
||||
return nil
|
||||
@@ -118,26 +119,57 @@ func (c *Client) getAuthMethods(jwtToken string) ([]ssh.AuthMethod, error) {
|
||||
return []ssh.AuthMethod{ssh.PublicKeys(signer)}, nil
|
||||
}
|
||||
|
||||
// StartSession starts an SSH session with PTY. It holds the client lock for
|
||||
// the whole startup so Close cannot tear the client down mid-setup and the
|
||||
// new session cannot be installed into an already closed client.
|
||||
// StartSession starts an SSH session with PTY
|
||||
func (c *Client) StartSession(cols, rows int) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.sshClient == nil {
|
||||
return fmt.Errorf("SSH client not connected")
|
||||
}
|
||||
|
||||
pty, err := nbssh.StartPTYSession(c.sshClient, cols, rows)
|
||||
session, err := c.sshClient.NewSession()
|
||||
if err != nil {
|
||||
return err
|
||||
return fmt.Errorf("create session: %w", err)
|
||||
}
|
||||
|
||||
c.session = pty.Session
|
||||
c.stdin = pty.Stdin
|
||||
c.stdout = pty.Stdout
|
||||
c.stderr = pty.Stderr
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.session = session
|
||||
|
||||
modes := ssh.TerminalModes{
|
||||
ssh.ECHO: 1,
|
||||
ssh.TTY_OP_ISPEED: 14400,
|
||||
ssh.TTY_OP_OSPEED: 14400,
|
||||
ssh.VINTR: 3,
|
||||
ssh.VQUIT: 28,
|
||||
ssh.VERASE: 127,
|
||||
}
|
||||
|
||||
if err := session.RequestPty("xterm-256color", rows, cols, modes); err != nil {
|
||||
closeWithLog(session, "session after PTY error")
|
||||
return fmt.Errorf("PTY request: %w", err)
|
||||
}
|
||||
|
||||
c.stdin, err = session.StdinPipe()
|
||||
if err != nil {
|
||||
closeWithLog(session, "session after stdin error")
|
||||
return fmt.Errorf("get stdin: %w", err)
|
||||
}
|
||||
|
||||
c.stdout, err = session.StdoutPipe()
|
||||
if err != nil {
|
||||
closeWithLog(session, "session after stdout error")
|
||||
return fmt.Errorf("get stdout: %w", err)
|
||||
}
|
||||
|
||||
c.stderr, err = session.StderrPipe()
|
||||
if err != nil {
|
||||
closeWithLog(session, "session after stderr error")
|
||||
return fmt.Errorf("get stderr: %w", err)
|
||||
}
|
||||
|
||||
if err := session.Shell(); err != nil {
|
||||
closeWithLog(session, "session after shell error")
|
||||
return fmt.Errorf("start shell: %w", err)
|
||||
}
|
||||
|
||||
logrus.Info("SSH: Session started with PTY")
|
||||
return nil
|
||||
|
||||
@@ -115,7 +115,7 @@ sequenceDiagram
|
||||
Resp->>Resp: parse usage tokens, completion
|
||||
Note over Resp: capture_completion gates raw<br/>completion capture
|
||||
Resp->>Cost: tokens
|
||||
Cost->>Cost: lookup rates from config-delivered<br/>pricing table + compute cost
|
||||
Cost->>Cost: lookup pricing.yaml + compute cost
|
||||
Cost->>Rec: tokens + cost
|
||||
Rec->>MgmtGrpc: RecordLLMUsage(provider, model, prompt_t, completion_t, cost, groups, user)
|
||||
Rec-->>Log: emit access-log entry<br/>(if EnableLogCollection)
|
||||
|
||||
@@ -15,10 +15,6 @@ Inside the package: `manager.go` is the CRUD + permissions-gated facade; `synthe
|
||||
| ---- | ---- |
|
||||
| `agentnetwork/manager.go` | Manager interface + CRUD + permission gates + bootstrap-settings + reconcile trigger |
|
||||
| `agentnetwork/synthesizer.go` | Settings/policy → wire-format synthesis; sole writer of the proxy middleware chain |
|
||||
| `agentnetwork/synthesizer_pricing.go` | `buildCostMeterConfigJSON` — default table + per-provider prices → `cost_meter` config |
|
||||
| `agentnetwork/pricing/defaults.go` | Default pricing table derived from the catalog + supplementals; `DefaultTable`, `LookupDefault`, wire `Entry` |
|
||||
| `agentnetwork/pricing/override.go` | `LoadFile`/`StartReloader` for `AgentNetwork.PricingDefaultsFile` (mtime poll, merge over compiled-in base) |
|
||||
| `agentnetwork/pricing/{exampleyaml,gen}.go` | Generates `defaults_llm_pricing.example.yaml` from the compiled-in table (golden-tested) |
|
||||
| `agentnetwork/policyselect.go` | Per-request policy attribution + account-budget ceiling (min-wins) |
|
||||
| `agentnetwork/reconcile.go` | Per-account synth diff vs in-memory cache → Create/Update/Delete |
|
||||
| `agentnetwork/catalog/catalog.go` | Static provider catalogue (auth headers, identity-injection shapes) |
|
||||
@@ -52,8 +48,6 @@ flowchart TD
|
||||
I --> J[indexProviderGroups: providerID -> sorted source groups]
|
||||
J --> K[buildRouterConfigJSON drops orphan providers]
|
||||
J --> L[buildIdentityInjectConfigJSON per catalog entry]
|
||||
J --> K2[buildCostMeterConfigJSON: default table + per-provider prices]
|
||||
K2 --> P
|
||||
H --> M[mergeGuardrails: union allowlist, OR redact]
|
||||
M --> N[applyAccountCollectionControls account toggle = SOLE capture control]
|
||||
N --> O[marshalGuardrailConfig]
|
||||
@@ -66,84 +60,6 @@ flowchart TD
|
||||
R --> T[accountManager.UpdateAccountPeers — fans synth ACLs into network map]
|
||||
```
|
||||
|
||||
### LLM pricing (management is the sole authority)
|
||||
|
||||
**The proxy carries no price list.** Management synthesizes the entire pricing
|
||||
table and ships it inside `cost_meter`'s `ConfigJSON`, so a price change reaches
|
||||
the proxies as an ordinary mapping push — the chain rebuild installs a fresh
|
||||
table and there is nothing to reload on the proxy side.
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[catalog.All — PricingSurfaces x Models] --> B[buildDefaultTable + supplementalDefaults]
|
||||
B --> C{AgentNetwork.PricingDefaultsFile}
|
||||
C -- absent --> D[compiled-in table serves]
|
||||
C -- loaded --> E[LoadFile: merge file entries WHOLE over compiled base]
|
||||
E --> F[mergedTable atomic.Pointer]
|
||||
D --> G[DefaultTable]
|
||||
F --> G
|
||||
G --> H[buildCostMeterConfigJSON — pricing.defaults]
|
||||
I[types.Provider.Models operator prices] --> J[normalizePricingModelID<br/>bedrock ARN/region/version, vertex @version]
|
||||
J --> K[materializeEntry: default entry as base,<br/>operator input/output verbatim,<br/>cache pointers only when non-nil]
|
||||
K --> L[pricing.providers keyed by provider record ID]
|
||||
H --> M[cost_meter ConfigJSON]
|
||||
L --> M
|
||||
G --> N[GET /catalog — applyDefaultPricing prefills dashboard rows]
|
||||
O[StartReloader: mtime poll every ReloadInterval 1m] --> E
|
||||
```
|
||||
|
||||
**Two tiers, resolved per request on the proxy** (`synthesizer_pricing.go:22-35`):
|
||||
|
||||
- `pricing.defaults` — surface (`openai`/`anthropic`/`bedrock`) → normalized model
|
||||
id → rates. The **full** default table ships to every account: it is small
|
||||
(~10 KB) and it is what keeps gateway-style providers (which enumerate no
|
||||
models, so they claim every model) priced.
|
||||
- `pricing.providers` — provider **record** id → normalized model id → rates,
|
||||
matched against the `llm.resolved_provider_id` the router stamps. Entries are
|
||||
**fully materialized here**, at synth time: `materializeEntry` starts from the
|
||||
default entry for that model so cache rates the operator didn't state are
|
||||
inherited, overlays operator `input`/`output` verbatim (**including an explicit
|
||||
0**, which prices a self-hosted or internal endpoint as free rather than
|
||||
silently reverting to list price), and overlays cache-rate **pointers only when
|
||||
non-nil** — `nil` means "inherit the default", an explicit `0` means "no
|
||||
discount, bill this bucket at the input rate". The proxy therefore does two map
|
||||
lookups and no merging.
|
||||
|
||||
Same orphan rule as the router: a provider no enabled policy authorises is
|
||||
unreachable, so its prices aren't shipped. Model ids are normalized with the
|
||||
**same** functions the request parser uses (`NormalizeBedrockModel` /
|
||||
`NormalizeVertexModel`), which is what makes the per-record lookup key compare
|
||||
equal to the `llm.model` the proxy meters. Post-normalization duplicates resolve
|
||||
first-occurrence-wins, matching the routing dedup order.
|
||||
|
||||
**`AgentNetwork.PricingDefaultsFile`** (`config.go:190-207`) lets an operator
|
||||
replace default rates without a rebuild. Schema is `surface → model → rates`
|
||||
(`input_per_1k`, `output_per_1k`, and optional `cached_input_per_1k` /
|
||||
`cache_read_per_1k` / `cache_creation_per_1k`). Semantics:
|
||||
|
||||
- A **relative** path resolves against `<Datadir>`, so a bare filename lands
|
||||
alongside the store. Empty config probes `<Datadir>/defaults_llm_pricing.yaml`.
|
||||
- An **explicitly configured** path is *required to load*: a typo or malformed
|
||||
file fails startup, because the operator believes those rates are live. The
|
||||
conventional probe is optional — an absent file just serves compiled-in
|
||||
defaults, and the path stays watched in case it appears later.
|
||||
- File entries **replace** the compiled-in entry for the same (surface, model)
|
||||
**whole** — they are not field-merged, so an entry must repeat the cache rates
|
||||
it wants to keep. Everything the file doesn't mention keeps built-in rates.
|
||||
- Unknown YAML fields are rejected (`KnownFields(true)`) and every rate must be
|
||||
finite and non-negative — the same constraints the HTTP API enforces on
|
||||
operator per-provider prices.
|
||||
- Reload is an mtime poll (`ReloadInterval`, 1 min) and is **lenient at runtime**:
|
||||
a parse error keeps the previous table, a deleted file reverts to compiled-in
|
||||
defaults. A mid-edit save can never take pricing down.
|
||||
|
||||
The live table feeds **both** consumers, which is what keeps them consistent: the
|
||||
synthesizer (what proxies actually bill with) and `GET /api/agent-network/catalog`
|
||||
via `applyDefaultPricing` (what the dashboard's model-row prices prefill with).
|
||||
`defaults_llm_pricing.example.yaml` is generated from the compiled-in table
|
||||
(`go generate ./management/internals/modules/agentnetwork/pricing`) and
|
||||
golden-tested, so operators start from a file matching the built-in rates exactly.
|
||||
|
||||
### Budget rule resolution (min-wins, group+user bound)
|
||||
|
||||
```mermaid
|
||||
@@ -208,7 +124,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
|
||||
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
|
||||
| on_response | 6 | `cost_meter` | `{"pricing":{"defaults":{surface:{model:rates}},"providers"?:{providerRecordID:{model:rates}}}}` — rates are `{input_per_1k, output_per_1k, cached_input_per_1k?, cache_read_per_1k?, cache_creation_per_1k?}` | – |
|
||||
| on_response | 6 | `cost_meter` | `{}` | – |
|
||||
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |
|
||||
- **Synthesized service shape** (`synthesizer.go:739`): `Mode=HTTP`, `Private=true`, `Domain=<subdomain>.<cluster>`, `AccessGroups=unionSourceGroups(enabledPolicies)`, one `TargetTypeCluster` target with `Host=noop.invalid:443` (router rewrites per request), `Options.{DirectUpstream,AgentNetwork}=true`, `DisableAccessLog=!settings.EnableLogCollection`, `CaptureMax{Req,Resp}Bytes=1<<20`, `CaptureContentTypes=["application/json","text/event-stream"]`.
|
||||
|
||||
@@ -223,12 +139,6 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
- **Orphan providers (no enabled policy authorises them) NEVER reach the router** (`synthesizer.go:351-357`); skipped from `identity_inject` for symmetry.
|
||||
- **Provider creation refuses empty `api_key`** (`manager.go:175`); **deletion refuses while any policy still references it** (`manager.go:265-273`).
|
||||
- **Session keypair stability across provider edits** (`manager.go:226-228`) — server-managed, copied through every `UpdateProvider`, never API-surfaced.
|
||||
- **Management is the sole pricing authority.** The proxy has no embedded price list, so an account whose `cost_meter` config carries no `pricing` block bills **nothing** (`cost.skipped=unknown_model`, $0) rather than falling back to stale built-ins. The top-level `pricing` wrapper is also the feature-detection signal in both directions: an old proxy ignores it as an unknown field, and a new proxy reads its absence as "old management".
|
||||
- **Per-provider prices are materialized at synth time, not merged on the proxy** (`synthesizer_pricing.go:114-131`). A per-record entry is always complete, so the proxy's lookup is per-record-then-defaults with no field-level fallback between tiers.
|
||||
- **An explicit operator price of `0` prices the model as free** — it must not be treated as "unset" and reverted to list price (`synthesizer_pricing.go:49-54`). Only *cache*-rate fields distinguish unset from zero, via `*float64`.
|
||||
- **Pricing model ids are normalized with the same functions the request parser uses** (`normalizePricingModelID`). If the two ever diverge, per-record prices silently stop matching and every request falls through to surface defaults.
|
||||
- **The default table's coverage is structural, not curated.** It is derived from the catalog via each provider's `PricingSurfaces`; `TestDefaultTable_CoversEveryCatalogModel` fails on an unpriced catalog model and `TestDefaultTable_NoConflictingContributions` fails if two providers contribute the same (surface, model) at different rates.
|
||||
- **A pricing-defaults file failure is fatal only at startup, and only for an explicitly configured path.** Runtime reload failures keep the previous table; a deleted file reverts to compiled-in defaults (`pricing/override.go:62-81, 113-148`).
|
||||
|
||||
## Things to scrutinize
|
||||
|
||||
@@ -266,12 +176,10 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
- **Capture-pointer semantics (restated):** non-agent-network callers see no field → legacy nil-default emit, identical to pre-PR. Agent-network targets always carry an explicit `capture_*` value.
|
||||
- **`TestSynthesizeServices_HappyPath` was updated:** request-parser config moved from `{}` to `{"capture_prompt":false}` (`synthesizer_test.go:174`). External snapshot tests against synth output need updating.
|
||||
- **`MergedGuardrails` retains zeroed `TokenLimits`/`Budget`/`Retention`** even though `Policy.Limits` carries the real values now; `llm_limit_check` is the authoritative enforcement. Comment at `synthesizer.go:940-948` calls this out.
|
||||
- **`cost_meter`'s `pricing` block is version-skew-safe in both directions.** A proxy predating config-delivered pricing ignores the field as unknown JSON (it previously priced from its own embedded table, so it keeps billing — at its own rates, which is the skew to be aware of during a rolling upgrade). A current proxy paired with old management sees no `pricing` block, logs one warning at chain-build time, and records `cost.skipped=unknown_model` — token counting and cap enforcement are unaffected, only the USD annotation goes to $0.
|
||||
|
||||
### Performance
|
||||
|
||||
- **`SynthesizeServices` runs on every controller tick / mutation reconcile.** Cost: 4 store reads + optional per-provider keypair backfill. Sort + index + merge are O(N log N) / O(P × G); dominant cost is JSON marshalling. No nested loops escape these dimensions.
|
||||
- **The full default pricing table is marshalled into every account's `cost_meter` config on every synth** (~10 KB serialized). This is a deliberate trade: it keeps gateway-style providers priced for every catalog model, and it is the largest single contributor to the synth JSON. `DefaultTable()` itself is a pointer load (or a `sync.Once`-built map) — the cost is the marshal, not the build.
|
||||
- **`reconcile.diffMappings` is O(N + M)** with N=M=1 per account today — effectively constant.
|
||||
- **`SynthesizeServicesForCluster`** (`synthesizer.go:71`) walks every account on a cluster; per-account failures are **swallowed** (`synthesizer.go:91-93`) so a single misconfigured account doesn't drop the cluster. Runs per proxy reconnect.
|
||||
|
||||
@@ -280,7 +188,6 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
- **Activity codes:** `AgentNetwork{Provider,Policy,Guardrail,BudgetRule}{Created,Updated,Deleted}`; `AgentNetworkSettingsUpdated` with `log_collection/prompt_collection/redact_pii` payload (`manager.go:567-571`). **No activity code for `SelectPolicyForRequest` denies** — surfaced via proxy access log only (likely intentional given volume).
|
||||
- **Deny codes** namespaced: `llm_policy.{token,budget}_cap_exceeded`, `llm_account.{token,budget}_cap_exceeded` (`policyselect.go:18-26`).
|
||||
- **Reconcile failures are logged at warn and swallowed** (`reconcile.go:42-44`). Persistent synth failures (e.g. unknown catalog id) silently keep the proxy out of sync — consider a manager-level synth-health surface if this becomes a support burden.
|
||||
- **Pricing-file lifecycle logs at info** (load, reload, revert-to-built-ins) and **at warn** for a runtime reload failure; the mtime check itself is `Debugf`. There is no metric on reload failures, so an operator who breaks the file mid-flight keeps billing at the previous table with only a log line to show it (`pricing/override.go:113-148`).
|
||||
|
||||
## Test coverage
|
||||
|
||||
@@ -291,9 +198,6 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
| `synthesizer_guardrail_realstore_test.go` | `PromptCaptureAccountIsSoleControl`; `PromptCaptureFlowsWhenAccountOptsIn`; `AccountRedactWithoutGuardrailRedact`; `NoGuardrail_CaptureOff`. |
|
||||
| `synthesizer_log_collection_realstore_test.go` | `LogCollection{Off_SuppressesAccessLog,On_PermitsAccessLog}` — verifies `DisableAccessLog` propagation through `ToProtoMapping`. |
|
||||
| `synthesizer_parser_redact_realstore_test.go` | **Capture-pointer regression suite:** `ParserConfigsCarryRedactPii`; `ParserConfigsSuppressCaptureWhenLogCollectionOnly` (log=on/prompt=off ⇒ both capture flags false); `ParserConfigsOmitRedactPiiWhenOff`. |
|
||||
| `synthesizer_pricing_test.go` | `BuildCostMeterConfig_{BedrockModelNormalization,CacheRateNilVsZero,OrphanAndGatewayProviders}` — the per-record tier's three load-bearing rules: keys normalized like the parser's, `nil` cache pointer inherits vs explicit `0` bills at input rate, and orphan / gateway (empty `Models`) providers ship no per-record entry. |
|
||||
| `pricing/defaults_test.go` | `DefaultTable_{CoversEveryCatalogModel,NoConflictingContributions,AllRatesFiniteNonNegative,PinnedRates}`; `LookupDefault_SurfaceOrder`. Catalog-derived coverage + rate sanity are structural, not curated. |
|
||||
| `pricing/override_test.go` | `LoadFile_{MergesOverCompiledDefaults,MissingPath,RejectsInvalid}`; `Reload_LifeCycle` (mtime detect, parse error keeps previous, delete reverts to built-ins); `ExampleYAML_InSyncWithBuiltins` golden. |
|
||||
| `policyselect_test.go` | Mock-store: `NoApplicablePolicies`; `AllowWithLowestGroupAttribution`; `LargerPoolWinsAcrossUsageLevels`; `StaysOnLargerPoolAfterPartialDrain`; `FallsThroughToSmallerPoolWhenLargerExhausted`; `TiebreakBy{LargerGroupPool,CreatedAt}`; `DeniesWhenAllExhausted`; `UncappedPolicyAlwaysWinsAgainstCapped`; `DisabledPolicyIgnored`; `StoreErrorPropagates`; `RejectsEmptyAccount`; `SharesGroupCounterAcrossPolicies`; `AntiFallThroughOnLowestGroup`; `BudgetOnlyExhaustionDenies`; `BudgetTighterThanTokenWins`. |
|
||||
| `policyselect_realstore_test.go` | Real-sqlite regression guard: `NoApplicablePolicies`; `AllowAndLowestGroupAttribution`; `LargerPoolWins_FallsThroughWhenExhausted`; `BudgetCapDenies`; `GroupCounterSharedAcrossPolicies`; `DisabledPolicyIgnored`. |
|
||||
| `policyselect_account_realstore_test.go` | Account budget rules: `AccountCeilingBindsEvenWithUncappedPolicy` (min-wins); `AccountGroupCeiling`; `AccountTargetUsersBindsOnlyThatUser`; `AccountRuleRecordsToOwnWindow`. |
|
||||
|
||||
@@ -5,7 +5,7 @@ LLM request. The two highest-blast-radius areas are the **capture-pointer
|
||||
semantics** and the **limit_check ⇒ limit_record** record-once invariant.
|
||||
|
||||
Sibling module: [32-proxy-llm-parsers.md](./32-proxy-llm-parsers.md) — the SDK
|
||||
adapters + pricing table and cost formula this chain delegates to.
|
||||
adapters + pricing catalog this chain delegates to.
|
||||
|
||||
---
|
||||
|
||||
@@ -34,7 +34,7 @@ rewrites.
|
||||
| `llm_identity_inject` | OnRequest | `llm.{resolved_provider_id,authorising_groups}`, `Input.{UserEmail,UserID,UserGroups,UserGroupNames}` | none | header strip/inject + optional body rewrite |
|
||||
| `llm_guardrail` | OnRequest | `llm.{model,request_prompt_raw}` | `llm_policy.{decision,reason}`, `llm.request_prompt` | none (model allowlist deny) |
|
||||
| `llm_response_parser` | OnResponse | `llm.provider`, `Input.{RespHeaders,RespBody,Status}` | `llm.{input,output,total,cached_input,cache_creation}_tokens`, `llm.response_completion` | none |
|
||||
| `cost_meter` | OnResponse | `llm.{provider,model,resolved_provider_id}`, token buckets | `cost.usd_{input,cached_input,cache_creation,output,total,cache}` or `cost.skipped` | none (in-memory pricing lookup) |
|
||||
| `cost_meter` | OnResponse | `llm.{provider,model}`, token buckets | `cost.usd_total` or `cost.skipped` | pricing lookup |
|
||||
| `llm_limit_record` | OnResponse | `llm.{attribution_group_id,attribution_window_seconds,input_tokens,output_tokens}`, `cost.usd_total` | none | gRPC `RecordLLMUsage` |
|
||||
|
||||
[all_test.go:26–40](../../../proxy/internal/middleware/builtin/all_test.go)
|
||||
@@ -44,7 +44,7 @@ locks the ID set; adding or removing one is a conscious extension.
|
||||
|
||||
| File | LOC | Notes |
|
||||
|---|---:|---|
|
||||
| `builtin.go` | 90 | Registry + `FactoryContext` (ctx, meter, logger, mgmt client) |
|
||||
| `builtin.go` | 86 | Registry + `FactoryContext` (ctx, data dir, meter, logger, mgmt client) |
|
||||
| `all_test.go` | 41 | Locks the 8-ID registry surface |
|
||||
| `agentnetwork_chain_integration_test.go` | 319 | Live sqlite + real gRPC bufconn; gate→recorder wire path |
|
||||
| `llm_request_parser/*` | 162 / 66 / 356 | Provider detection, body parse, prompt extraction with capture-pointer gating |
|
||||
@@ -53,7 +53,7 @@ locks the ID set; adding or removing one is a conscious extension.
|
||||
| `llm_identity_inject/*` | 440 / 108 / 666 | HeaderPair (LiteLLM) + JSONMetadata (Portkey) + ExtraHeaders |
|
||||
| `llm_guardrail/*` | 176 / 82 / 75 / 219 / 217 | Model allowlist + optional prompt capture with PII redaction |
|
||||
| `llm_response_parser/*` | 258 / 222 / 43 / 433 / 169 / 111 | Buffered + SSE accumulation; AWS event-stream accumulator (`streaming_bedrock.go`) for Bedrock; capture-pointer gates completion emit |
|
||||
| `cost_meter/*` | 236 / 98 / 586 | Token → USD via `proxy/internal/llm/pricing`; both pricing tiers arrive in the middleware config |
|
||||
| `cost_meter/*` | 181 / 84 / 439 | Token → USD via `proxy/internal/llm/pricing` |
|
||||
| `llm_limit_record/*` | 144 / 35 / 191 | Post-flight `RecordLLMUsage` (5s, debug-on-error) |
|
||||
|
||||
## Per-middleware
|
||||
@@ -168,46 +168,12 @@ token schema.
|
||||
|
||||
### cost_meter
|
||||
|
||||
Reads `llm.provider` + `llm.model` + token buckets, looks up the per-1k rates,
|
||||
and emits the full `cost.usd_*` breakdown (four per-bucket values plus the
|
||||
`_total` and `_cache` aggregates) or a closed-set `cost.skipped` reason
|
||||
(`missing_provider/model/tokens`, `unparseable_tokens`, `zero_tokens`,
|
||||
`unknown_model`).
|
||||
|
||||
**Management owns pricing.** The proxy carries no embedded price list: the whole
|
||||
table arrives in this middleware's `ConfigJSON` as
|
||||
`{pricing: {defaults, providers}}`, synthesized by management from the catalog
|
||||
plus the operator's stored per-provider prices
|
||||
([factory.go:13–34](../../../proxy/internal/middleware/builtin/cost_meter/factory.go)).
|
||||
Both tiers are validated by `pricing.NewTable` / `pricing.NewEntries` at
|
||||
construction, so a non-finite or negative rate fails the chain build. A price
|
||||
change is an ordinary mapping push — the chain rebuild yields a fresh instance
|
||||
over a fresh immutable table, so there is no data dir, no pricing file, no
|
||||
reload goroutine, and nothing to invalidate.
|
||||
|
||||
**Two-tier lookup**
|
||||
([middleware.go:165–183](../../../proxy/internal/middleware/builtin/cost_meter/middleware.go)):
|
||||
|
||||
1. **Per-provider-record** — the operator's stored price for the route that
|
||||
actually served the request, keyed by the `llm.resolved_provider_id` that
|
||||
`llm_router` stamped on the allow path, then by normalized model id. Entries
|
||||
arrive fully materialized (management folds default cache rates in at synth
|
||||
time), so there is no merging here. Absent metadata — no router in the chain
|
||||
— skips this tier.
|
||||
2. **Surface defaults** — the catalog-derived table keyed by `llm.provider`
|
||||
(`openai`/`anthropic`/`bedrock`). This is also what prices gateway-style
|
||||
providers, which enumerate no models and therefore get no per-record entry.
|
||||
|
||||
**Backward compatibility:** a config with no `pricing` block means management
|
||||
predates config-delivered pricing. The factory logs one warning at build time
|
||||
and the instance records `cost.skipped=unknown_model` ($0) for every request
|
||||
rather than falling back to a stale built-in price list
|
||||
([factory.go:55–60](../../../proxy/internal/middleware/builtin/cost_meter/factory.go)).
|
||||
|
||||
**Key invariant:** the provider-shape switch lives in `pricing.EntryCosts`
|
||||
(sibling doc) and is selected by the **surface**, not by which tier the entry
|
||||
came from — `cost_meter` stays provider-agnostic, and a per-record override on
|
||||
an Anthropic route still bills its cache buckets additively.
|
||||
Reads `llm.provider` + `llm.model` + token buckets, looks up per-1k rate via
|
||||
`pricing.Loader`, emits `cost.usd_total` or a closed-set `cost.skipped`
|
||||
reason (`missing_provider/model/tokens`, `unparseable_tokens`, `zero_tokens`,
|
||||
`unknown_model`). Loader's hot-reload goroutine is bound to proxy-lifetime
|
||||
context via `startReloader`. **Key invariant:** provider-shape switch lives
|
||||
in `pricing.Table.Cost` (sibling doc) — `cost_meter` stays provider-agnostic.
|
||||
|
||||
### llm_limit_record
|
||||
|
||||
@@ -280,14 +246,12 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
|
||||
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
|
||||
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
|
||||
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
|
||||
| `cost_meter` | `{pricing: {defaults: {surface: {model: rates}}, providers: {providerRecordID: {model: rates}}}}` — rates are `{input_per_1k, output_per_1k, cached_input_per_1k?, cache_read_per_1k?, cache_creation_per_1k?}`. A missing `pricing` key means "management predates config-delivered pricing": every request records `cost.skipped=unknown_model` |
|
||||
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
|
||||
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |
|
||||
|
||||
All factories accept empty / null / `{}` / whitespace as zero-value config;
|
||||
only structurally invalid JSON is rejected so misconfig surfaces at chain
|
||||
build time. `cost_meter` adds a semantic check on top of that: a `pricing`
|
||||
block carrying a negative or non-finite rate fails the build too, rather than
|
||||
mispricing live traffic.
|
||||
build time.
|
||||
|
||||
## Invariants
|
||||
|
||||
@@ -356,11 +320,10 @@ non-object `metadata` field
|
||||
— header path still attributes, but body-level tag-budget enforcement
|
||||
doesn't run for that request.
|
||||
|
||||
**Concurrency.** `cost_meter`'s two pricing tables are built once from the
|
||||
middleware config and never mutated, so the lookup path needs no lock or atomic
|
||||
swap — a price change replaces the whole instance. Every middleware is
|
||||
otherwise a stateless value receiver. Integration test uses real bufconn gRPC —
|
||||
race detector is the meaningful bar.
|
||||
**Concurrency.** `cost_meter` shares a `pricing.Loader` via
|
||||
`atomic.Pointer[Table]`; readers always see a consistent table. Every
|
||||
middleware is a stateless value receiver. Integration test uses real bufconn
|
||||
gRPC — race detector is the meaningful bar.
|
||||
|
||||
**Perf.** Hot path is `lookupKV` linear scan over <10 KVs; `cost_meter.Cost`
|
||||
is O(1); SSE accumulation is single-pass. No map allocation per call.
|
||||
@@ -386,13 +349,13 @@ counter accuracy.
|
||||
| `llm_guardrail/redact_test.go` | 15 | Email, SSN, phone (E.164 + NA), bearer, IPv4; fixture-driven |
|
||||
| `llm_response_parser/middleware_test.go` | 18 | Buffered OAI+Anthro, capture-pointer, redact, truncation |
|
||||
| `llm_response_parser/streaming_test.go` | 7 | OAI usage frame, Anthro message_delta, truncated body best-effort |
|
||||
| `cost_meter/middleware_test.go` | 22 | Each skip reason, provider-shape formulas, config-delivered defaults, per-record-beats-defaults + miss-falls-back, per-record uses surface formula, nil-pricing skips everything, invalid-rate rejection |
|
||||
| `cost_meter/middleware_test.go` | 17 | Each skip reason, provider-shape, pricing loader integration |
|
||||
| `llm_limit_record/middleware_test.go` | 7 | Skip-on-no-signal, skip-on-missing-attribution, RPC failure swallowed |
|
||||
|
||||
## Cross-references
|
||||
|
||||
- Sibling: [32-proxy-llm-parsers.md](./32-proxy-llm-parsers.md) — SDK adapters
|
||||
+ SSE framer + pricing table and cost formula.
|
||||
+ SSE framer + pricing loader.
|
||||
- Path-routed providers (Vertex AI + Bedrock), `keyfile::` credential, GCP
|
||||
token minting, `/bedrock` prefix:
|
||||
[50-path-routed-providers.md](./50-path-routed-providers.md).
|
||||
|
||||
@@ -9,7 +9,7 @@ pricing table's per-provider cost formula is the highest-leverage place a
|
||||
small bug would silently mis-bill operators.
|
||||
|
||||
Sibling module: [31-proxy-middleware-builtin.md](./31-proxy-middleware-builtin.md)
|
||||
— the 8 middlewares that consume this package's parsers + pricing table.
|
||||
— the 8 middlewares that consume this package's parsers + pricing loader.
|
||||
|
||||
---
|
||||
|
||||
@@ -24,9 +24,8 @@ proxy-framework dependencies:
|
||||
- `openai.go` / `anthropic.go` / `bedrock.go` — per-provider `Parser` impls.
|
||||
- `sse.go` — SSE scanner (`Scanner`, `Event`, `NewScanner`).
|
||||
- `errors.go` — sentinels callers branch on with `errors.Is`.
|
||||
- `pricing/` — immutable pricing table + the per-surface cost formula. The
|
||||
rates themselves come from management inside `cost_meter`'s middleware
|
||||
config; this package holds no price list and reads no files.
|
||||
- `pricing/` — embedded-default + hot-reload override table with
|
||||
symlink-safe Unix loader (build-tagged stub elsewhere).
|
||||
- `fixtures/` — captured request/response/stream bodies the tests replay.
|
||||
|
||||
The package carries zero proxy-framework dependencies so the same parsers can
|
||||
@@ -48,9 +47,12 @@ be reused later by a WASM adapter
|
||||
| `sse_test.go` | 175 | 12 tests; fixture replay + multiline + size limits |
|
||||
| `parser_test.go` | 53 | `Parsers()`, `DetectParser`, provider enum values |
|
||||
| `errors.go` | 31 | 6 sentinels: `Err{Unknown,Unsupported}Provider/Model`, `Err{NotLLM,Malformed}Response`, `ErrStreamingUnsupported`, `ErrMalformedRequest` |
|
||||
| `pricing/pricing.go` | 234 | `Table`, `Entry`, `EntryJSON`, `Costs`; `NewTable`/`NewEntries` validation + `EntryCosts` formula. No I/O, no reload, no embedded rates |
|
||||
| `pricing/pricing_test.go` | 177 | 10 tests — provider-shape formulas, cached clamp, rate fallback, nil-safety, rate validation |
|
||||
| `fixtures/*` | 21–59 | OAI chat/responses/stream + Anthro messages/stream |
|
||||
| `pricing/pricing.go` | 421 | `Loader`, `Table`, `Entry`; embedded defaults + atomic swap + mtime reload |
|
||||
| `pricing/pricing_unix.go` | 69 | `O_NOFOLLOW` + fstat-from-FD + 1 MiB cap |
|
||||
| `pricing/pricing_other.go` | 21 | Stub returning "not supported on this platform" |
|
||||
| `pricing/pricing_test.go` | 432 | 21 tests — symlink rejection, reload race, path traversal, oversize |
|
||||
| `pricing/defaults_pricing.yaml` | 85 | go:embed source of truth |
|
||||
| `fixtures/*` | 21–59 | OAI chat/responses/stream + Anthro messages/stream + pricing starter |
|
||||
|
||||
## Request body → parser dispatch
|
||||
|
||||
@@ -186,11 +188,9 @@ response leg, covering both Bedrock body shapes:
|
||||
`totalTokens`). `firstNonZero` folds the two naming conventions into one
|
||||
`Usage`; when Converse omits `totalTokens` the parser sums the buckets.
|
||||
|
||||
`ProviderName()` returns `"bedrock"` — its own pricing surface in the table
|
||||
management ships, keyed by the **normalised** model id (region prefix + version
|
||||
suffix stripped by the request parser; management normalises its keys the same
|
||||
way at synth time so the two compare equal). `ParseResponse` returns
|
||||
`ErrStreamingUnsupported` for an
|
||||
`ProviderName()` returns `"bedrock"` — its own `defaults_pricing.yaml` block,
|
||||
keyed by the **normalised** model id (region prefix + version suffix stripped by
|
||||
the request parser). `ParseResponse` returns `ErrStreamingUnsupported` for an
|
||||
AWS binary event-stream content-type (`application/vnd.amazon.eventstream`,
|
||||
`isAWSEventStream`) so the caller routes to the streaming accumulator instead.
|
||||
|
||||
@@ -205,34 +205,11 @@ response body. Streaming accumulators live in the middleware package
|
||||
([llm_response_parser/streaming.go](../../../proxy/internal/middleware/builtin/llm_response_parser/streaming.go))
|
||||
but use `llm.NewScanner` so the framing contract stays here.
|
||||
|
||||
### Pricing table
|
||||
### Pricing catalog
|
||||
|
||||
**Management is the sole pricing authority.** The proxy carries no embedded
|
||||
price list and reads no pricing file: the whole table arrives inside
|
||||
`cost_meter`'s `ConfigJSON` on the ordinary mapping push, and a price change
|
||||
is just another push — the chain rebuild constructs a fresh `Table`, so there
|
||||
is nothing to reload
|
||||
([pricing.go:1–7](../../../proxy/internal/llm/pricing/pricing.go)). The
|
||||
management side of the contract (catalog defaults, the operator's stored
|
||||
per-provider prices, and `AgentNetwork.PricingDefaultsFile`) is covered in the
|
||||
management-side module guide; `cost_meter`'s wire shape is in
|
||||
[31-proxy-middleware-builtin.md](./31-proxy-middleware-builtin.md).
|
||||
|
||||
`EntryJSON`
|
||||
([pricing.go:36–45](../../../proxy/internal/llm/pricing/pricing.go)) is the
|
||||
management→proxy contract — five USD-per-1k rates under `input_per_1k`,
|
||||
`output_per_1k`, `cached_input_per_1k`, `cache_read_per_1k`,
|
||||
`cache_creation_per_1k`. Management's `pricing.Entry` marshals the identical
|
||||
names, and `EntryJSON`/`Entry` are field-identical so `NewEntries` converts by
|
||||
direct struct conversion rather than field-by-field copying (a new rate can't
|
||||
be silently dropped in transit).
|
||||
|
||||
`EntryCosts`
|
||||
([pricing.go:183–234](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
is the cost formula — most security-relevant math in this module. The
|
||||
**surface** (the `llm.provider` value the request parser stamped) selects the
|
||||
formula, never the tier the entry came from: a per-provider-record override on
|
||||
an Anthropic route still bills its cache buckets additively.
|
||||
`Table.Cost`
|
||||
([pricing.go:129–174](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
is the cost formula — most security-relevant math in this module:
|
||||
|
||||
| Provider | Formula |
|
||||
|---|---|
|
||||
@@ -241,7 +218,7 @@ an Anthropic route still bills its cache buckets additively.
|
||||
| default | `inTokens × InputPer1K + outTokens × OutputPer1K` |
|
||||
|
||||
`bedrock` shares the Anthropic additive-cache formula
|
||||
([pricing.go:214–229](../../../proxy/internal/llm/pricing/pricing.go)):
|
||||
([pricing.go:172-174](../../../proxy/internal/llm/pricing/pricing.go)):
|
||||
Anthropic-on-Bedrock reports the same additive cache buckets, while non-Anthropic
|
||||
Bedrock models (Nova, Llama) simply report zero in those buckets so cost reduces
|
||||
to `input + output`.
|
||||
@@ -249,12 +226,15 @@ to `input + output`.
|
||||
Each per-bucket rate falls back to `InputPer1K` when zero — operators opt in
|
||||
to discounts by setting the field.
|
||||
|
||||
`Costs`
|
||||
([pricing.go:143–163](../../../proxy/internal/llm/pricing/pricing.go)) is the
|
||||
per-request split. The four per-bucket fields are the base; `TotalUSD` and
|
||||
`CacheUSD` are **derived** in `newCosts` so the aggregates can never drift from
|
||||
the breakdown. `InputUSD` is always the non-cached input bucket on both
|
||||
provider shapes, so input and cached-input never double-count.
|
||||
`Loader`
|
||||
([pricing.go:212–268](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
overlays an optional `pricing.yaml` from data-dir on top of the go:embed
|
||||
defaults. Atomic pointer swap means readers never observe a partial update.
|
||||
The mtime-poll reloader (30s default cadence) keeps the previous table on
|
||||
parse failure so cost annotation never goes blank during a botched edit.
|
||||
|
||||
`defaults_pricing.yaml` is the source of truth for built-in pricing.
|
||||
Operator overrides only carry the entries they want to change.
|
||||
|
||||
## Public contracts
|
||||
|
||||
@@ -284,38 +264,29 @@ Order matters: `DetectFromURL` ties resolve by registration order.
|
||||
`ProviderBedrock = 3`. Numeric values are persisted in nothing today but treat
|
||||
them as wire-stable — new providers must take fresh numbers.
|
||||
|
||||
**`Pricing` construction + lookup**
|
||||
([pricing.go:60–130](../../../proxy/internal/llm/pricing/pricing.go)):
|
||||
**`Pricing` lookup**
|
||||
([pricing.go:129](../../../proxy/internal/llm/pricing/pricing.go)):
|
||||
|
||||
```go
|
||||
func NewEntries(raw map[string]map[string]EntryJSON) (map[string]map[string]Entry, error)
|
||||
func NewTable(raw map[string]map[string]EntryJSON) (*Table, error)
|
||||
|
||||
func (t *Table) Lookup(provider, model string) (Entry, bool)
|
||||
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool)
|
||||
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool)
|
||||
func EntryCosts(entry Entry, surface string, inTokens, outTokens, cachedInput, cacheCreation int64) Costs
|
||||
```
|
||||
|
||||
`NewTable` is the surface-keyed defaults table; `NewEntries` returns the raw
|
||||
two-level map `cost_meter` uses for the per-provider-record tier (it looks up an
|
||||
`Entry` directly and calls `EntryCosts`, so it needs no `Table` wrapper). Both
|
||||
reject any non-finite or negative rate, so a corrupt config fails the chain
|
||||
build rather than mispricing silently. Nil input yields an empty,
|
||||
never-matching table.
|
||||
|
||||
Nil-safe: `t.Cost`/`t.Lookup` on a nil receiver returns `ok=false`
|
||||
([pricing.go:96–99](../../../proxy/internal/llm/pricing/pricing.go)).
|
||||
`ok=false` means the surface or model is absent from the table management sent;
|
||||
the caller emits `cost.skipped=unknown_model`.
|
||||
Nil-safe: `t.Cost` on a nil receiver returns `(0, false)`
|
||||
([pricing.go:130–132](../../../proxy/internal/llm/pricing/pricing.go)).
|
||||
`ok=false` means provider or model is absent from the loaded table; the caller
|
||||
emits `cost.skipped=unknown_model`.
|
||||
|
||||
## Invariants
|
||||
|
||||
1. **The pricing package is pure and platform-independent.** No file I/O, no
|
||||
`//go:embed`, no goroutines, no build tags — the rates arrive as config, so
|
||||
there is nothing platform-specific left to port. Anything reintroducing a
|
||||
read-from-disk path here re-splits pricing authority between management and
|
||||
the proxy, which is exactly what this design removed.
|
||||
1. **Cross-platform pricing build.** `pricing_unix.go` carries the only
|
||||
functional `loadPricing` (uses `syscall.O_NOFOLLOW` and `f.Stat()` on an
|
||||
open descriptor — both Unix-only). `pricing_other.go` is a build-tag
|
||||
fallback that returns `"not supported on this platform"`
|
||||
([pricing_other.go:14–16](../../../proxy/internal/llm/pricing/pricing_other.go)).
|
||||
The proxy is Linux-only in production today; a Windows port needs an
|
||||
equivalent path-as-handle implementation. Reviewers building on Windows
|
||||
should expect this surface to return an error at startup if an override
|
||||
file is configured.
|
||||
|
||||
2. **SSE scanner handles partial chunks.** A buffered prefix that doesn't end
|
||||
in `\n\n` still yields its accumulated event before `io.EOF`
|
||||
@@ -327,45 +298,38 @@ the caller emits `cost.skipped=unknown_model`.
|
||||
usage rather than aborting
|
||||
([streaming.go:68–73, 144–150](../../../proxy/internal/middleware/builtin/llm_response_parser/streaming.go)).
|
||||
|
||||
3. **Management is the only source of rates.** `Table` has no constructor that
|
||||
invents prices: the only way in is `NewTable`/`NewEntries` over the wire map
|
||||
management sent. A missing or empty `pricing` block therefore means *no
|
||||
prices at all* (`cost_meter` records `cost.skipped=unknown_model`, $0) —
|
||||
never a stale built-in fallback that would silently bill list price.
|
||||
3. **`defaults_pricing.yaml` is the source of truth.** Compiled into the
|
||||
binary via `//go:embed`
|
||||
([pricing.go:29–30](../../../proxy/internal/llm/pricing/pricing.go)).
|
||||
`DefaultTable()` parses once and panics on parse failure
|
||||
([pricing.go:42–49](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
— by design: a broken embedded YAML must not ship to production.
|
||||
|
||||
4. **Tables are immutable once built.** `Table.entries` is written only in
|
||||
`NewEntries` and never mutated afterwards, and `cost_meter`'s `perRecord`
|
||||
map is likewise build-time-only
|
||||
([pricing.go:47–52](../../../proxy/internal/llm/pricing/pricing.go)). This
|
||||
is what makes the no-reload design safe: a price change arrives as a mapping
|
||||
push that builds a new middleware instance over a new table, so concurrent
|
||||
readers can't observe a half-updated price list and no atomic swap or lock
|
||||
is needed on the hot path.
|
||||
4. **Loader path validation.** `resolveMiddlewareDataPath`
|
||||
([pricing.go:370–394](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
rejects absolute paths, traversal segments, and basenames that fail
|
||||
`basenameRegex = ^[a-zA-Z0-9._-]+$`. The resolved path must remain
|
||||
inside `baseDir` even after `filepath.Clean`. Tests:
|
||||
`TestNewLoader_PathValidation`, `TestNewLoader_PathValidation_Extended`,
|
||||
`TestNewLoader_SymlinkOutsideBaseDirRejected`, `TestNewLoader_SymlinkRejected`.
|
||||
|
||||
5. **Rate validation happens at chain-build time, not per request.**
|
||||
`NewEntries` rejects negative, NaN, and ±Inf rates field by field
|
||||
([pricing.go:60–83](../../../proxy/internal/llm/pricing/pricing.go)), naming
|
||||
the offending surface/model/field in the error. Management enforces the same
|
||||
constraints at its API boundary and in its YAML parser, so this is
|
||||
defense-in-depth — but it means a corrupt push fails loudly at build instead
|
||||
of producing negative costs on live traffic. Test:
|
||||
`TestNewTable_ValidatesRates`.
|
||||
5. **Unix loader symlink safety.** `O_NOFOLLOW` on open, `f.Stat()` on the
|
||||
open descriptor (never re-stat by path), `info.Mode().IsRegular()` check,
|
||||
`io.LimitReader(f, maxPricingBytes+1)` with a final size assertion
|
||||
([pricing_unix.go:25–57](../../../proxy/internal/llm/pricing/pricing_unix.go)).
|
||||
A mid-read symlink swap is detected because the fstat is on the original
|
||||
fd. Test: `TestNewLoader_RejectsOversizedFile_FixesM4`.
|
||||
|
||||
6. **New rates must be added to `Entry`, `EntryJSON`, *and* management's
|
||||
`pricing.Entry` together.** `NewEntries` converts by direct struct
|
||||
conversion `Entry(e)`
|
||||
([pricing.go:76–78](../../../proxy/internal/llm/pricing/pricing.go)), which
|
||||
only compiles while the two structs stay field-identical — so the proxy half
|
||||
is compiler-enforced. The management half is not: a rate added there but not
|
||||
here unmarshals into nothing and prices that bucket at `InputPer1K`.
|
||||
6. **`yaml.NewDecoder(...).KnownFields(true)`**
|
||||
([pricing.go:397–398](../../../proxy/internal/llm/pricing/pricing.go))
|
||||
rejects YAML files that carry fields not in the schema. A typo in an
|
||||
operator override file fails loud instead of silently zeroing rates.
|
||||
|
||||
## Things to scrutinise
|
||||
|
||||
**Correctness.** Verify the OpenAI cached-prompt clamp at
|
||||
[pricing.go:203–206](../../../proxy/internal/llm/pricing/pricing.go)
|
||||
short-circuits before subtraction. Negative token counts are clamped to zero up
|
||||
front ([pricing.go:186–197](../../../proxy/internal/llm/pricing/pricing.go)) so
|
||||
no formula can yield a negative cost. `Anthropic.TotalTokens` sums all four
|
||||
**Correctness.** Verify OpenAI cached-prompt clamp at
|
||||
[pricing.go:147–149](../../../proxy/internal/llm/pricing/pricing.go)
|
||||
short-circuits before subtraction. `Anthropic.TotalTokens` sums all four
|
||||
buckets (in + out + cache_read + cache_creation) — downstream dashboards
|
||||
need to know this differs from `input + output`.
|
||||
`OpenAIParser.ExtractPrompt` falls through `messages → input → prompt`; a
|
||||
@@ -374,27 +338,22 @@ noting).
|
||||
|
||||
**Security.** `Scanner.maxLine = 1 MiB`; a 2 MiB single-line `data:` event
|
||||
errors from `Scanner.Next` and both accumulators stop with partial usage.
|
||||
Pricing is no longer file-backed, so the loader's path-traversal / symlink /
|
||||
oversize surface is gone entirely — the config channel (an authenticated
|
||||
mapping push from management) is now the only way rates enter the proxy, and
|
||||
`NewEntries` is the validation boundary on it. A new rate added to management's
|
||||
`pricing.Entry` but not to `EntryJSON` here is the remaining silent-mispricing
|
||||
path (see invariant 6).
|
||||
Pricing file 1 MiB cap is orders of magnitude larger than realistic. Confirm
|
||||
new schema additions are mirrored in both `pricingFile` and `Entry`;
|
||||
`KnownFields(true)` will reject silently-typo'd operator overrides
|
||||
otherwise.
|
||||
|
||||
**Concurrency.** Nothing in this package is shared mutable state: tables are
|
||||
built once and never written again, so `cost_meter`'s hot path is lock-free by
|
||||
construction rather than by atomic swap. Per-call `Scanner` instances mean no
|
||||
shared state across concurrent response-parser calls.
|
||||
**Concurrency.** `Loader.table` is `atomic.Pointer[Table]`; readers never
|
||||
block or see a torn table. `Loader.Reload` is one goroutine, cancelled via
|
||||
context (`TestLoader_ReloadBackgroundLoopCancellation`). `DefaultTable()`
|
||||
uses `sync.Once`. Per-call `Scanner` instances mean no shared state across
|
||||
concurrent response-parser calls.
|
||||
|
||||
**Perf.** `Table.Cost` is two map lookups + multiplications, O(1); the
|
||||
per-provider-record tier adds at most one more lookup. `Scanner.Next` is one
|
||||
`ReadString('\n')` per line. No background goroutines and no per-request
|
||||
allocation of pricing state.
|
||||
**Perf.** `Table.Cost` is two map lookups + multiplications, O(1).
|
||||
`Scanner.Next` is one `ReadString('\n')` per line. Pricing reload poll 30s.
|
||||
|
||||
**Observability.** A config carrying no `pricing` block logs one warning at
|
||||
chain-build time (`cost_meter` factory) and then records
|
||||
`cost.skipped=unknown_model` per request, so an old-management deployment is
|
||||
visible in both logs and the access log rather than quietly reporting $0.
|
||||
**Observability.** Reload failures count via `metric.Int64Counter` keyed
|
||||
`plugin`; warning log rate-limited at 5 min so a broken file doesn't flood.
|
||||
Parser errors return sentinels — middleware uses `errors.Is` to map to the
|
||||
right `cost.skipped` reason.
|
||||
|
||||
@@ -406,7 +365,7 @@ right `cost.skipped` reason.
|
||||
| `openai_test.go` | 11 | Chat Completions + Responses API + legacy `prompt`; cached-tokens subset for both naming conventions; fixture replays |
|
||||
| `anthropic_test.go` | 7 | Messages + legacy `/v1/complete`; streaming REJECTED on `ParseResponse` (must use scanner); fixture replays |
|
||||
| `sse_test.go` | 12 | Fixture replay both providers; multiline `data:`; CRLF; comment skip; trailing-event-without-blank-line; oversize rejection |
|
||||
| `pricing/pricing_test.go` | 10 | Provider-shape switch (surface selects the formula); cached-rate + cache-read/creation fallback to `InputPer1K`; cached-clamp; negative-token clamp; nil-receiver safety; rate validation (negative / NaN / Inf rejected); nil + empty table |
|
||||
| `pricing/pricing_test.go` | 21 | Provider-shape switch; cached-rate fallback; cached-clamp; symlink rejection (target outside basedir + symlink to file); path validation matrix; oversize rejection; reload-keeps-previous-on-parse-error; mtime change detection; goroutine cancellation |
|
||||
|
||||
**Fixtures** ([proxy/internal/llm/fixtures/](../../../proxy/internal/llm/fixtures/)):
|
||||
`openai_chat_completion.json` (chat.completions with usage),
|
||||
@@ -414,15 +373,14 @@ right `cost.skipped` reason.
|
||||
`openai_stream.txt` (3 deltas + usage + `[DONE]`),
|
||||
`anthropic_messages.json` (Messages API non-streaming),
|
||||
`anthropic_stream.txt` (full 7-event sequence: message_start →
|
||||
content_block_{start,delta×2,stop} → message_delta (usage) → message_stop).
|
||||
No pricing fixture: the table is config-delivered, so pricing tests construct
|
||||
it in-process from a wire-shape map.
|
||||
content_block_{start,delta×2,stop} → message_delta (usage) → message_stop),
|
||||
`pricing.yaml` (realistic-pricing starter for operator overrides).
|
||||
|
||||
## Cross-references
|
||||
|
||||
- Sibling: [31-proxy-middleware-builtin.md](./31-proxy-middleware-builtin.md)
|
||||
— the chain that calls `llm.Parsers()`, `llm.ParserByName`,
|
||||
`llm.NewScanner`, `pricing.NewTable` / `pricing.NewEntries`.
|
||||
`llm.NewScanner`, `pricing.NewLoader`.
|
||||
- Path-routed providers (Vertex AI + Bedrock), credential syntax, and the
|
||||
Bedrock AWS event-stream accumulator:
|
||||
[50-path-routed-providers.md](./50-path-routed-providers.md).
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# proxy/runtime — translate + serve + log
|
||||
|
||||
> **Risk level:** High — every config push from management is translated here, and the chain runs on every HTTP request to a synth target.
|
||||
> **Backward-compat impact:** Additive at the wire (`PathTargetOptions.middlewares`, `agent_network`, `disable_access_log`, capture caps) and on the proxy `Server` struct (`MiddlewareCaptureBudgetBytes`). Non-agent-network targets stay on the no-middleware fast path. Middleware config is entirely wire-delivered — no proxy-side data dir is involved, including for LLM pricing, which management ships inside `cost_meter`'s config.
|
||||
> **Backward-compat impact:** Additive at the wire (`PathTargetOptions.middlewares`, `agent_network`, `disable_access_log`, capture caps) and on the proxy `Server` struct (`MiddlewareDataDir`, `MiddlewareCaptureBudgetBytes`). Non-agent-network targets stay on the no-middleware fast path.
|
||||
|
||||
## Module boundary
|
||||
|
||||
@@ -114,7 +114,8 @@ At **request time** the access-log middleware stamps `CapturedData`; the auth ch
|
||||
|
||||
## Public contracts touched
|
||||
|
||||
- `proxy.Server.MiddlewareCaptureBudgetBytes` (int64) — process-wide capture cap; defaults to 256 MiB (server.go:249-253). There is no `MiddlewareDataDir`: no built-in middleware reads config from disk, so `builtin.FactoryContext` carries only the proxy-lifetime context, meter, logger, and management client.
|
||||
- `proxy.Server.MiddlewareDataDir` (string) — base dir for file-backed middleware config (server.go:238-241).
|
||||
- `proxy.Server.MiddlewareCaptureBudgetBytes` (int64) — process-wide capture cap; defaults to 256 MiB (server.go:248-250).
|
||||
- `proxy/internal/proxy.WithMiddlewareManager(*middleware.Manager) Option` — new option on `NewReverseProxy`; nil keeps the fast path (reverseproxy.go:48-56).
|
||||
- `proxy/internal/proxy.PathTarget` adds `Middlewares`, `CaptureConfig`, `AgentNetwork`, `DisableAccessLog` (servicemapping.go:27-51), all zero-default.
|
||||
- `proxy/internal/proxy.CapturedData` adds `agentNetwork`, `suppressAccessLog`, `userGroupNames` behind `sync.RWMutex`; slices deep-copied (context.go:47-66, 183-258).
|
||||
|
||||
@@ -87,9 +87,9 @@ strips the `@version` suffix from the model, and maps the publisher to a parser
|
||||
surface via `vertexPublisherVendor`:
|
||||
|
||||
- `anthropic` → `llm.provider="anthropic"` → metered through the Anthropic
|
||||
parser, priced under the **`anthropic`** surface of the pricing table
|
||||
management ships (the parser emits the standard Anthropic provider label, so
|
||||
Vertex Claude reuses first-party Anthropic prices).
|
||||
parser, priced under the **`anthropic`** block in `defaults_pricing.yaml`
|
||||
(the parser emits the standard Anthropic provider label, so Vertex Claude
|
||||
reuses first-party Anthropic prices).
|
||||
- `openai` → `llm.provider="openai"` (reserved; not in the catalog lineup
|
||||
today).
|
||||
- anything else (notably `google` / Gemini) → empty vendor → **no parser**.
|
||||
@@ -104,9 +104,8 @@ is omitted from the catalog.
|
||||
|
||||
> Caveat: cross-region inference profiles in `eu` / `apac` carry a ~10% price
|
||||
> premium that the base per-token rates do **not** model — cost annotations for
|
||||
> those regions read low. Operators who need exact regional billing set the
|
||||
> affected models' prices on the provider record, or replace the default entries
|
||||
> via management's `AgentNetwork.PricingDefaultsFile`.
|
||||
> those regions read low. Operators who need exact regional billing override
|
||||
> the affected entries in `pricing.yaml`.
|
||||
|
||||
## AWS Bedrock (`bedrock_api`)
|
||||
|
||||
@@ -212,19 +211,15 @@ so a model-listing call can't be rewritten onto an upstream that would 404 it.
|
||||
## Catalog ↔ pricing cross-check
|
||||
|
||||
Catalog prices and context windows are cross-checked against LiteLLM's
|
||||
`model_prices_and_context_window.json`. The **catalog is the source of default
|
||||
prices**: management's `pricing.DefaultTable` folds every catalog provider's
|
||||
models into the surfaces that provider declares (`PricingSurfaces`), so coverage
|
||||
is structural rather than maintained in a parallel file
|
||||
([pricing/defaults.go](../../../management/internals/modules/agentnetwork/pricing/defaults.go)).
|
||||
`TestDefaultTable_CoversEveryCatalogModel` fails if a catalog model ends up
|
||||
unpriced, and `TestDefaultTable_NoConflictingContributions` fails if two
|
||||
providers contribute the same (surface, model) at different rates. Bedrock
|
||||
entries are keyed by the **normalised** id the request parser emits (region
|
||||
prefix + version suffix stripped) — management applies the same normalisation to
|
||||
per-provider prices at synth time, so the two keys compare equal. Vertex Claude
|
||||
carries no Bedrock-style prefix, so it prices straight off the `anthropic`
|
||||
surface.
|
||||
`model_prices_and_context_window.json`. The proxy's embedded
|
||||
`defaults_pricing.yaml` covers **every metered first-party model** the catalog
|
||||
enumerates — guarded by
|
||||
`TestDefaultTable_FirstPartyModelCoverage`
|
||||
([pricing/defaults_coverage_test.go](../../../proxy/internal/llm/pricing/defaults_coverage_test.go)),
|
||||
which fails if a catalog model has no embedded price. Bedrock entries are keyed
|
||||
by the **normalised** id the request parser emits (region prefix + version
|
||||
suffix stripped). Vertex Claude carries no Bedrock-style prefix, so it prices
|
||||
straight off the `anthropic` block.
|
||||
|
||||
## Things to scrutinise
|
||||
|
||||
@@ -237,17 +232,16 @@ operator-misconfigured Vertex provider and unmetered Gemini traffic; verify
|
||||
publishers).
|
||||
|
||||
**Correctness.** `normalizeBedrockModel` is the join between the wire id and the
|
||||
pricing key — a model that normalises to something absent from the shipped
|
||||
pricing table meters at `cost.skipped=unknown_model` rather than failing the
|
||||
request. The
|
||||
pricing key — a model that normalises to something not in `defaults_pricing.yaml`
|
||||
meters at `cost.skipped=unknown_model` rather than failing the request. The
|
||||
`/bedrock` prefix strip must run on both the parser side (so the model is
|
||||
extracted) and the router side (so the upstream path is native); a regression in
|
||||
either silently breaks the other.
|
||||
|
||||
**Metering caveats.** eu/apac cross-region Bedrock + Vertex profiles carry a
|
||||
~10% premium not modelled by base pricing — flagged in the catalog comment.
|
||||
Operators needing exact regional billing set per-provider prices on the model
|
||||
rows (or replace the default entries via `AgentNetwork.PricingDefaultsFile`).
|
||||
~10% premium not modelled by base pricing — flagged in both the catalog comment
|
||||
and `defaults_pricing.yaml`. Operators needing exact regional billing override
|
||||
the relevant entries.
|
||||
|
||||
## Cross-references
|
||||
|
||||
|
||||
@@ -23,9 +23,10 @@ import (
|
||||
// model the client asks for. The proxy prices off the REQUEST model, not the
|
||||
// upstream response model, so a made-up model id billed at operator rates lets
|
||||
// these tests assert exact costs without a real vendor key.
|
||||
// Sourced from the harness so the counts can't drift from the mock's config.
|
||||
const (
|
||||
vllmPromptTokens = 11
|
||||
vllmCompletionTokens = 2
|
||||
vllmPromptTokens = harness.VLLMChatInputTokens
|
||||
vllmCompletionTokens = harness.VLLMChatOutputTokens
|
||||
)
|
||||
|
||||
// pricedEnv is a connected single-provider agent-network deployment pointed at
|
||||
@@ -169,23 +170,48 @@ func chatOnce(t *testing.T, ctx context.Context, env pricedEnv, model, sessionID
|
||||
return body
|
||||
}
|
||||
|
||||
// findAccessLogBySession polls the access-log page for the row carrying sessionID.
|
||||
func findAccessLogBySession(t *testing.T, ctx context.Context, sessionID string) api.AgentNetworkAccessLog {
|
||||
t.Helper()
|
||||
var row api.AgentNetworkAccessLog
|
||||
require.Eventually(t, func() bool {
|
||||
logs, lerr := srv.ListAccessLogs(ctx)
|
||||
if lerr != nil {
|
||||
return false
|
||||
}
|
||||
for _, r := range logs.Data {
|
||||
if r.SessionId != nil && *r.SessionId == sessionID {
|
||||
row = r
|
||||
return true
|
||||
// accessLogIngestWindow is how long a single request's access-log row is given
|
||||
// to appear before the caller gives up on it.
|
||||
// accessLogIngestWindow bounds how long a row may take to appear after its
|
||||
// request returned. The proxy streams each entry to management with a 10s send
|
||||
// timeout of its own, so a request whose send hits one full timeout and is
|
||||
// retried has not yet missed anything real — 30s left barely three send
|
||||
// attempts of headroom and lost the race on a loaded runner.
|
||||
const accessLogIngestWindow = 60 * time.Second
|
||||
|
||||
// lookupAccessLogBySession polls the access-log page for the row carrying
|
||||
// sessionID and reports whether it arrived within the window. It never fails
|
||||
// the test: callers that can recover — by firing a fresh request under a new
|
||||
// session — need to see the miss rather than die on it.
|
||||
func lookupAccessLogBySession(ctx context.Context, sessionID string, within time.Duration) (api.AgentNetworkAccessLog, bool) {
|
||||
deadline := time.Now().Add(within)
|
||||
for {
|
||||
if logs, lerr := srv.ListAccessLogs(ctx); lerr == nil {
|
||||
for _, r := range logs.Data {
|
||||
if r.SessionId != nil && *r.SessionId == sessionID {
|
||||
return r, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row", sessionID)
|
||||
if time.Now().After(deadline) {
|
||||
return api.AgentNetworkAccessLog{}, false
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return api.AgentNetworkAccessLog{}, false
|
||||
case <-time.After(2 * time.Second):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// findAccessLogBySession polls the access-log page for the row carrying
|
||||
// sessionID, failing the test if it never lands. Use it for a request whose row
|
||||
// must exist; where a missing row is a recoverable race, use
|
||||
// lookupAccessLogBySession and retry.
|
||||
func findAccessLogBySession(t *testing.T, ctx context.Context, sessionID string) api.AgentNetworkAccessLog {
|
||||
t.Helper()
|
||||
row, ok := lookupAccessLogBySession(ctx, sessionID, accessLogIngestWindow)
|
||||
require.True(t, ok, "session id %q must be recorded in an access-log row", sessionID)
|
||||
return row
|
||||
}
|
||||
|
||||
@@ -319,6 +345,11 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) {
|
||||
outRateA = 0.020
|
||||
inRateB = 0.050 // 5x / 4x the original, so a repriced row is unmistakable
|
||||
outRateB = 0.080
|
||||
// Per-attempt ingest wait, shorter than the default so a request that
|
||||
// produces no row costs one retry rather than most of the budget, and an
|
||||
// overall deadline long enough to hold several attempts.
|
||||
repriceIngestWindow = 20 * time.Second
|
||||
repriceDeadline = 180 * time.Second
|
||||
)
|
||||
|
||||
env := provisionPricedProvider(t, ctx, "reprice", []api.AgentNetworkProviderModel{
|
||||
@@ -353,10 +384,15 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) {
|
||||
// reading its cost, so an un-ingested row is never mistaken for "still rate A".
|
||||
// The expected new input cost is unmistakably higher than rate A, so a
|
||||
// lingering old-rate row can't satisfy the check.
|
||||
//
|
||||
// Every way an iteration can come up short — the request failing, its row not
|
||||
// landing, or the row still carrying rate A — is a symptom of the same
|
||||
// in-flight rebuild, so each one retries under a fresh session rather than
|
||||
// ending the test. Only the outer deadline is fatal.
|
||||
wantInputB := float64(vllmPromptTokens) / 1000 * inRateB
|
||||
var repriced api.AgentNetworkAccessLog
|
||||
var lastSession string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
deadline := time.Now().Add(repriceDeadline)
|
||||
for time.Now().Before(deadline) {
|
||||
lastSession = fmt.Sprintf("e2e-session-reprice-b-%d", time.Now().UnixNano())
|
||||
code, _, cerr := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, customModel, "Reply with exactly: pong", lastSession)
|
||||
@@ -364,7 +400,15 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) {
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
row := findAccessLogBySession(t, ctx, lastSession)
|
||||
row, ok := lookupAccessLogBySession(ctx, lastSession, repriceIngestWindow)
|
||||
if !ok {
|
||||
// No row for this request. The provider update rebuilds the proxy's
|
||||
// middleware chain, and a request served mid-rebuild can complete
|
||||
// without a resolved provider — 200 to the caller, nothing to
|
||||
// attribute, so no row is ever written for it. Fire another one.
|
||||
t.Logf("no access-log row for session %q within %s; retrying under a fresh session", lastSession, repriceIngestWindow)
|
||||
continue
|
||||
}
|
||||
if inDelta(row.InputCostUsd, wantInputB, 1e-6) {
|
||||
repriced = row
|
||||
break
|
||||
@@ -630,3 +674,47 @@ func inDelta(a, b, tol float64) bool {
|
||||
}
|
||||
return d <= tol
|
||||
}
|
||||
|
||||
// TestCustomDatedModelKeepsItsOwnPrice covers the review fix that anchored the
|
||||
// release-date fallback to Claude ids. Pricing looks every model up through
|
||||
// that helper, so while it matched a bare trailing date any operator id ending
|
||||
// in eight digits inherited the rate of its undated sibling — a silent
|
||||
// mis-bill on models NetBird knows nothing about.
|
||||
func TestCustomDatedModelKeepsItsOwnPrice(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
baseModel = "internal-llm"
|
||||
datedModel = "internal-llm-20250101"
|
||||
baseIn = 0.010
|
||||
baseOut = 0.020
|
||||
// An order of magnitude apart, so a row billed at the wrong entry is
|
||||
// unmistakable rather than a rounding argument.
|
||||
datedIn = 0.100
|
||||
datedOut = 0.200
|
||||
)
|
||||
|
||||
env := provisionPricedProvider(t, ctx, "customdated", []api.AgentNetworkProviderModel{
|
||||
{Id: baseModel, InputPer1k: baseIn, OutputPer1k: baseOut},
|
||||
{Id: datedModel, InputPer1k: datedIn, OutputPer1k: datedOut},
|
||||
})
|
||||
|
||||
t.Run("the undated id bills at its own rate", func(t *testing.T) {
|
||||
session := fmt.Sprintf("e2e-session-customdated-base-%d", time.Now().UnixNano())
|
||||
chatOnce(t, ctx, env, baseModel, session)
|
||||
assertOpenAICostAtRates(t, findAccessLogBySession(t, ctx, session), baseIn, baseOut)
|
||||
})
|
||||
|
||||
t.Run("the dated id keeps its own rate", func(t *testing.T) {
|
||||
session := fmt.Sprintf("e2e-session-customdated-dated-%d", time.Now().UnixNano())
|
||||
chatOnce(t, ctx, env, datedModel, session)
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
assertOpenAICostAtRates(t, row, datedIn, datedOut)
|
||||
|
||||
// Spelled out because it is the regression: inheriting the sibling's
|
||||
// rate would bill this request at a tenth of its price.
|
||||
assert.Greater(t, row.InputCostUsd, float64(vllmPromptTokens)/1000*baseIn*2,
|
||||
"a custom dated id must not inherit the undated entry's rate")
|
||||
})
|
||||
}
|
||||
|
||||
400
e2e/agentnetwork/discovery_live_test.go
Normal file
400
e2e/agentnetwork/discovery_live_test.go
Normal file
@@ -0,0 +1,400 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestLiveModelDiscovery drives model discovery against the REAL vendor
|
||||
// endpoints — OpenAI, Anthropic, Bedrock and Vertex — rather than the mock.
|
||||
//
|
||||
// The mock upstream proves the filter's mechanics: it advertises ids we chose,
|
||||
// so a listing narrowing to the ones we authorised is arithmetic we already
|
||||
// controlled both sides of. What it cannot prove is that the filter survives
|
||||
// contact with a real catalogue — ids we never enumerated, dated builds whose
|
||||
// suffix the vendor picks, surfaces that answer a listing request with
|
||||
// something other than a listing. That is what this covers, and it is the part
|
||||
// a QA engineer would otherwise have to walk through by hand.
|
||||
//
|
||||
// One proxy serves every case. Each provider gets its own group, policy and
|
||||
// client, because a model-less request matches exactly ONE route
|
||||
// (matchModelless): with two providers authorised for the same caller, the
|
||||
// listing would go to whichever won the tiebreak and the other would go
|
||||
// untested. Group-scoping the caller makes each provider the only candidate
|
||||
// for its own client.
|
||||
func TestLiveModelDiscovery(t *testing.T) {
|
||||
cases := liveDiscoveryCases()
|
||||
if len(cases) == 0 {
|
||||
t.Skip("no provider keys set; source ~/.llm-keys to run live model discovery")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
t.Logf("[discovery] live matrix: %s", strings.Join(caseNames(cases), ", "))
|
||||
|
||||
// Provision every provider, group and policy before the proxy starts: the
|
||||
// proxy takes a configuration snapshot at connect time and does not
|
||||
// reconcile provider changes made afterwards.
|
||||
keys := make(map[string]string, len(cases))
|
||||
for i := range cases {
|
||||
keys[cases[i].name] = provisionLiveDiscovery(t, ctx, &cases[i])
|
||||
}
|
||||
|
||||
endpoint, firstIP, firstClient, px := connectClient(t, ctx, "disc-live", keys[cases[0].name])
|
||||
clients := map[string]*harness.Client{cases[0].name: firstClient}
|
||||
ips := map[string]string{cases[0].name: firstIP}
|
||||
for _, tc := range cases[1:] {
|
||||
cl := joinClient(t, ctx, px, endpoint, keys[tc.name])
|
||||
ip, err := cl.ResolveProxyIP(ctx, endpoint)
|
||||
require.NoError(t, err, "resolve endpoint from the %s client", tc.name)
|
||||
clients[tc.name] = cl
|
||||
ips[tc.name] = ip
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
runLiveDiscoveryCase(t, ctx, tc, clients[tc.name], endpoint, ips[tc.name])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// discoveryOutcome is what a discovery request must produce end to end. The
|
||||
// three are genuinely different contracts, not degrees of success: only the
|
||||
// first puts a bounded listing in front of the caller.
|
||||
type discoveryOutcome int
|
||||
|
||||
const (
|
||||
// outcomeFiltered: the proxy routes the request and bounds the response to
|
||||
// what the caller may use.
|
||||
outcomeFiltered discoveryOutcome = iota
|
||||
// outcomeDenied: no provider of this shape can serve the surface, so the
|
||||
// proxy refuses rather than rewriting the request onto an upstream that
|
||||
// would 404 it. The caller gets a NetBird error, not a vendor one.
|
||||
outcomeDenied
|
||||
// outcomeUpstreamNoListing: the proxy routes the request to the configured
|
||||
// upstream, and the vendor does not implement the endpoint there. Proxy
|
||||
// side correct, product side a dead end — see the Bedrock case.
|
||||
outcomeUpstreamNoListing
|
||||
)
|
||||
|
||||
// liveDiscoveryCase is one provider's discovery surface and what the proxy
|
||||
// must make of it.
|
||||
type liveDiscoveryCase struct {
|
||||
name string
|
||||
catalogID string
|
||||
upstream string
|
||||
apiKey string
|
||||
|
||||
// path is the discovery endpoint the client calls. Not every surface uses
|
||||
// /v1/models: Bedrock lists inference profiles instead.
|
||||
path string
|
||||
// headers the vendor requires on a bare GET (Anthropic versions its API
|
||||
// through a header, and rejects a request without one).
|
||||
headers []string
|
||||
|
||||
// models the provider record enumerates. Empty models a gateway record,
|
||||
// which enumerates nothing and claims everything.
|
||||
models []string
|
||||
// allowlist, when non-empty, is a guardrail narrowing the policy below the
|
||||
// provider's own enumeration — the second of the two bounds discovery
|
||||
// applies, and the only one a provider record alone cannot demonstrate.
|
||||
allowlist []string
|
||||
|
||||
// outcome is what this surface must produce end to end.
|
||||
outcome discoveryOutcome
|
||||
|
||||
// permitted is every id allowed to survive filtering, in the form the
|
||||
// provider record registers it. A surviving id counts as permitted when it
|
||||
// matches one of these outright or after Anthropic date-normalisation.
|
||||
permitted []string
|
||||
// wantHidden are ids the upstream is known to advertise and the bound must
|
||||
// remove. Only set where we enumerate the model ourselves, so the
|
||||
// expectation cannot rot when a vendor changes its catalogue.
|
||||
wantHidden []string
|
||||
}
|
||||
|
||||
// liveDiscoveryCases builds the matrix from whichever provider credentials are
|
||||
// present, mirroring availableProviders' env-var gating so a partial key set
|
||||
// still yields partial coverage.
|
||||
func liveDiscoveryCases() []liveDiscoveryCase {
|
||||
var cases []liveDiscoveryCase
|
||||
|
||||
// OpenAI enumerates TWO real models and the policy permits one. That is
|
||||
// the only case here where both bounds are observable at once: the
|
||||
// upstream advertises dozens of ids, the provider record cuts them to two,
|
||||
// and the guardrail cuts those to one.
|
||||
if k := os.Getenv("OPENAI_TOKEN"); k != "" {
|
||||
cases = append(cases, liveDiscoveryCase{
|
||||
name: "openai", catalogID: "openai_api", upstream: "https://api.openai.com", apiKey: k,
|
||||
path: "/v1/models",
|
||||
models: []string{"gpt-4o-mini", "gpt-4o"},
|
||||
allowlist: []string{"gpt-4o-mini"},
|
||||
outcome: outcomeFiltered,
|
||||
permitted: []string{"gpt-4o-mini"},
|
||||
wantHidden: []string{"gpt-4o"},
|
||||
})
|
||||
}
|
||||
|
||||
// Anthropic is the surface Claude Code actually calls. Its listing returns
|
||||
// DATED build ids (claude-haiku-4-5-20251001) while the provider record
|
||||
// registers the undated id, so this is the case that proves the filter's
|
||||
// date-normalisation against ids the vendor chose rather than ids we wrote.
|
||||
if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" {
|
||||
cases = append(cases, liveDiscoveryCase{
|
||||
name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k,
|
||||
path: "/v1/models",
|
||||
headers: []string{"anthropic-version: 2023-06-01"},
|
||||
models: []string{"claude-haiku-4-5"},
|
||||
outcome: outcomeFiltered,
|
||||
permitted: []string{"claude-haiku-4-5"},
|
||||
})
|
||||
}
|
||||
|
||||
// Bedrock lists inference profiles, not models: matchModelless routes
|
||||
// /inference-profiles to a Bedrock route and refuses /v1/models for one.
|
||||
//
|
||||
// The request reaches AWS and AWS refuses it — bedrock-runtime answers
|
||||
// <UnknownOperationException/>, because ListInferenceProfiles is a CONTROL
|
||||
// PLANE operation served by bedrock.<region>.amazonaws.com, not the runtime
|
||||
// host. A provider record carries one upstream and it has to be the runtime
|
||||
// host for InvokeModel to work, so no Bedrock record can serve a listing as
|
||||
// the model stands today.
|
||||
//
|
||||
// The mock upstream hides this entirely: it answers /inference-profiles on
|
||||
// the same listener as everything else, so the routing test passes there
|
||||
// while the real endpoint 404s. That is the whole reason this file exists,
|
||||
// so the case is kept, asserting what actually happens.
|
||||
if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
|
||||
region := os.Getenv("AWS_REGION")
|
||||
if region == "" {
|
||||
region = "eu-central-1"
|
||||
}
|
||||
model := os.Getenv("AWS_BEDROCK_MODEL")
|
||||
if model == "" {
|
||||
model = "global.anthropic.claude-sonnet-4-6"
|
||||
}
|
||||
cases = append(cases, liveDiscoveryCase{
|
||||
name: "bedrock", catalogID: "bedrock_api",
|
||||
upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k,
|
||||
path: "/inference-profiles",
|
||||
models: []string{sharedllm.NormalizeAnthropicModel(strings.TrimPrefix(model, "global."))},
|
||||
outcome: outcomeUpstreamNoListing,
|
||||
})
|
||||
}
|
||||
|
||||
// Vertex carries the model in the rawPredict path and serves no listing
|
||||
// endpoint at all, so the proxy must refuse discovery rather than rewrite
|
||||
// it onto an upstream that would 404.
|
||||
if sa := os.Getenv("GOOGLE_VERTEX_SA_BASE64"); sa != "" {
|
||||
if project := os.Getenv("GOOGLE_VERTEX_PROJECT"); project != "" {
|
||||
region := os.Getenv("GOOGLE_VERTEX_REGION")
|
||||
if region == "" {
|
||||
region = "global"
|
||||
}
|
||||
host := "aiplatform.googleapis.com"
|
||||
if region != "global" {
|
||||
host = region + "-aiplatform.googleapis.com"
|
||||
}
|
||||
cases = append(cases, liveDiscoveryCase{
|
||||
name: "vertex", catalogID: "vertex_ai_api", upstream: "https://" + host,
|
||||
apiKey: "keyfile::" + sa,
|
||||
path: "/v1/models",
|
||||
outcome: outcomeDenied,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return cases
|
||||
}
|
||||
|
||||
// provisionLiveDiscovery creates the group, provider, optional guardrail and
|
||||
// policy for one case, and returns the setup key a client joins that group
|
||||
// with. Scoping each provider to its own group is what keeps it the only
|
||||
// candidate for its own client's model-less request.
|
||||
func provisionLiveDiscovery(t *testing.T, ctx context.Context, tc *liveDiscoveryCase) string {
|
||||
t.Helper()
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-live-" + tc.name})
|
||||
require.NoError(t, err, "create group for %s", tc.name)
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-disc-live-" + tc.name,
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key for %s", tc.name)
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext for %s", tc.name)
|
||||
|
||||
req := api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-disc-live-" + tc.name,
|
||||
ProviderId: tc.catalogID,
|
||||
UpstreamUrl: tc.upstream,
|
||||
ApiKey: &tc.apiKey,
|
||||
Enabled: ptr(true),
|
||||
}
|
||||
if len(tc.models) > 0 {
|
||||
models := make([]api.AgentNetworkProviderModel, 0, len(tc.models))
|
||||
for _, id := range tc.models {
|
||||
models = append(models, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.002})
|
||||
}
|
||||
req.Models = &models
|
||||
}
|
||||
prov, err := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, err, "create provider %s", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
polReq := api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-disc-live-" + tc.name,
|
||||
Enabled: ptr(true),
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
}
|
||||
if len(tc.allowlist) > 0 {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-disc-live-" + tc.name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = tc.allowlist
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail for %s", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
polReq.GuardrailIds = &[]string{g.Id}
|
||||
}
|
||||
pol, err := srv.CreatePolicy(ctx, polReq)
|
||||
require.NoError(t, err, "create policy for %s", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
return sk.Key
|
||||
}
|
||||
|
||||
// runLiveDiscoveryCase issues the discovery request and reports everything the
|
||||
// vendor said before asserting on any of it. The log is the point on the first
|
||||
// run: a live catalogue is the one input we do not control, so a failure has to
|
||||
// arrive with the response that caused it rather than just a count.
|
||||
func runLiveDiscoveryCase(t *testing.T, ctx context.Context, tc liveDiscoveryCase, cl *harness.Client, endpoint, proxyIP string) {
|
||||
t.Helper()
|
||||
|
||||
// A single request is enough for the two non-listing outcomes, and retrying
|
||||
// them would burn the retry window waiting for a status that is never
|
||||
// coming.
|
||||
if tc.outcome != outcomeFiltered {
|
||||
code, body, err := cl.Get(ctx, endpoint, proxyIP, tc.path, tc.headers)
|
||||
require.NoError(t, err, "request must reach the proxy")
|
||||
t.Logf("[discovery] %s GET %s -> %d; body: %s", tc.name, tc.path, code, truncate(body, 2000))
|
||||
assert.NotEqual(t, 200, code,
|
||||
"%s serves no bounded listing, so a 200 here would mean the caller was handed a picker nothing narrows; body: %s",
|
||||
tc.name, truncate(body, 2000))
|
||||
|
||||
// Which side refused is the whole distinction between these two
|
||||
// outcomes, and a NetBird error is the thing that tells them apart: the
|
||||
// middleware chain stamps its own name on anything it generates.
|
||||
if tc.outcome == outcomeDenied {
|
||||
assert.True(t, isProxyError(body),
|
||||
"%s serves no listing endpoint at all, so the proxy must refuse the request itself rather than forward it to an upstream that would answer for us; body: %s",
|
||||
tc.name, truncate(body, 2000))
|
||||
return
|
||||
}
|
||||
assert.False(t, isProxyError(body),
|
||||
"%s discovery must be routed to the configured upstream and refused by the vendor, not blocked by the proxy; body: %s",
|
||||
tc.name, truncate(body, 2000))
|
||||
return
|
||||
}
|
||||
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Get(ctx, endpoint, proxyIP, tc.path, tc.headers)
|
||||
}, 200)
|
||||
t.Logf("[discovery] %s GET %s -> %d; body: %s", tc.name, tc.path, code, truncate(body, 4000))
|
||||
require.Equal(t, 200, code, "%s discovery must be served; body: %s", tc.name, truncate(body, 2000))
|
||||
|
||||
ids, ok := listingIDs(body)
|
||||
require.Truef(t, ok,
|
||||
"%s answered discovery with something other than a {\"data\":[{\"id\":…}]} listing, which the filter forwards untouched — the caller would get an unbounded picker; body: %s",
|
||||
tc.name, truncate(body, 2000))
|
||||
sort.Strings(ids)
|
||||
t.Logf("[discovery] %s: %d ids after filtering: %s", tc.name, len(ids), strings.Join(ids, ", "))
|
||||
|
||||
require.NotEmpty(t, ids, "%s filtered the listing down to nothing; the caller would see an empty picker", tc.name)
|
||||
|
||||
permitted := make(map[string]struct{}, len(tc.permitted)*2)
|
||||
for _, id := range tc.permitted {
|
||||
permitted[id] = struct{}{}
|
||||
permitted[sharedllm.NormalizeAnthropicModel(id)] = struct{}{}
|
||||
}
|
||||
for _, id := range ids {
|
||||
_, direct := permitted[id]
|
||||
_, normalised := permitted[sharedllm.NormalizeAnthropicModel(id)]
|
||||
assert.Truef(t, direct || normalised,
|
||||
"%s offered %q, which no policy on this route permits — every entry the picker shows must be a request the guardrail would allow", tc.name, id)
|
||||
}
|
||||
for _, hidden := range tc.wantHidden {
|
||||
assert.NotContainsf(t, ids, hidden,
|
||||
"%s offered %q, which the provider enumerates but the policy does not permit", tc.name, hidden)
|
||||
}
|
||||
}
|
||||
|
||||
// isProxyError reports whether a response body was generated by the middleware
|
||||
// chain rather than forwarded from a vendor. Every chain-generated error names
|
||||
// the middleware that raised it, which no upstream's error body does — so this
|
||||
// separates "the proxy refused" from "the proxy routed it and the vendor
|
||||
// refused", the two failures that otherwise look alike from the client side.
|
||||
func isProxyError(body string) bool {
|
||||
return strings.Contains(body, `"middleware":`)
|
||||
}
|
||||
|
||||
// listingIDs pulls the model ids out of a listing response. ok is false when
|
||||
// the body is not the {"data":[{"id":…}]} shape the filter recognises.
|
||||
func listingIDs(body string) ([]string, bool) {
|
||||
var doc struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(body), &doc); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if doc.Data == nil {
|
||||
return nil, false
|
||||
}
|
||||
ids := make([]string, 0, len(doc.Data))
|
||||
for _, entry := range doc.Data {
|
||||
ids = append(ids, entry.ID)
|
||||
}
|
||||
return ids, true
|
||||
}
|
||||
|
||||
func caseNames(cases []liveDiscoveryCase) []string {
|
||||
names := make([]string, 0, len(cases))
|
||||
for _, c := range cases {
|
||||
names = append(names, c.name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// truncate bounds a logged response body. A live catalogue can run to tens of
|
||||
// kilobytes, and the useful part is the front.
|
||||
func truncate(s string, limit int) string {
|
||||
if len(s) <= limit {
|
||||
return s
|
||||
}
|
||||
return s[:limit] + "… (" + strconv.Itoa(len(s)-limit) + " more bytes)"
|
||||
}
|
||||
168
e2e/agentnetwork/discovery_multipolicy_test.go
Normal file
168
e2e/agentnetwork/discovery_multipolicy_test.go
Normal file
@@ -0,0 +1,168 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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"
|
||||
)
|
||||
|
||||
// TestDiscoveryBoundToCallersPolicies covers a model listing on a provider two
|
||||
// teams reach under different allowlists.
|
||||
//
|
||||
// Bounding the listing by the provider's enumerated models alone is not enough
|
||||
// once more than one policy is in play: the caller would be offered every model
|
||||
// any team may use, and each one outside their own policy is a request the
|
||||
// guardrail refuses a moment later — the empty-or-wrong picker this endpoint
|
||||
// exists to avoid, just moved one level up.
|
||||
//
|
||||
// The client joins the main group only. Both models are enumerated by the same
|
||||
// provider and both are advertised by the upstream, so a listing that leaked
|
||||
// the other team's model would visibly contain it.
|
||||
func TestDiscoveryBoundToCallersPolicies(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
mkKey := func(name, groupID string) string {
|
||||
sk, kerr := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: name,
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{groupID},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, kerr, "mint setup key %s", name)
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
return sk.Key
|
||||
}
|
||||
// One client per group. The second is what makes the first assertion mean
|
||||
// something: without a client that DOES see the other team's model, its
|
||||
// absence from the main client's listing could equally be a policy that
|
||||
// never propagated.
|
||||
keyMain := mkKey("e2e-disc-mp-main-client", grpMain.Id)
|
||||
keyOther := mkKey("e2e-disc-mp-other-client", grpOther.Id)
|
||||
|
||||
// One provider enumerating both models the upstream advertises, so the
|
||||
// listing is narrowed by policy rather than by what the provider serves.
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-disc-mp",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
{Id: harness.VLLMUnlistedModel, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-disc-mp-main", harness.VLLMModel)
|
||||
gOther := mkGuardrail("e2e-disc-mp-other", harness.VLLMUnlistedModel)
|
||||
|
||||
enabled := true
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-disc-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// The other team's policy, on the same provider, permitting the model the
|
||||
// client must never be offered.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-disc-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
endpoint, proxyIP, clMain, px := connectClient(t, ctx, "disc-mp", keyMain)
|
||||
clOther := joinClient(t, ctx, px, endpoint, keyOther)
|
||||
|
||||
listing := func(t *testing.T, cl *harness.Client, ip string) string {
|
||||
t.Helper()
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Get(ctx, endpoint, ip, "/v1/models?limit=1000", nil)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "discovery must be served; body: %s", body)
|
||||
return body
|
||||
}
|
||||
|
||||
otherIP, err := clOther.ResolveProxyIP(ctx, endpoint)
|
||||
require.NoError(t, err, "resolve endpoint from the other client")
|
||||
|
||||
// The other team's client first: seeing its own model proves polOther is
|
||||
// live, so the main client's listing is narrowed by policy scoping rather
|
||||
// than by the other policy having failed to apply at all.
|
||||
otherBody := listing(t, clOther, otherIP)
|
||||
assert.Contains(t, otherBody, harness.VLLMUnlistedModel,
|
||||
"the other group's policy must be in force, or this test proves nothing")
|
||||
assert.NotContains(t, otherBody, harness.VLLMModel,
|
||||
"and it must not be offered the main group's model either — isolation runs both ways")
|
||||
|
||||
mainBody := listing(t, clMain, proxyIP)
|
||||
assert.Contains(t, mainBody, harness.VLLMModel,
|
||||
"the model the caller's own policy permits must reach the picker")
|
||||
assert.NotContains(t, mainBody, harness.VLLMUnlistedModel,
|
||||
"a model only another group's policy permits must not be offered to this caller")
|
||||
}
|
||||
|
||||
// joinClient starts a second tunnel client against an already-running proxy, so
|
||||
// a test can drive the same endpoint as two different group memberships without
|
||||
// paying for a second proxy.
|
||||
func joinClient(t *testing.T, ctx context.Context, px *harness.Proxy, endpoint, setupKey string) *harness.Client {
|
||||
t.Helper()
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, setupKey)
|
||||
require.NoError(t, err, "start second client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "second client must connect to management")
|
||||
if _, err := cl.ResolveProxyIP(ctx, endpoint); err != nil {
|
||||
t.Fatalf("second client could not resolve the endpoint: %v", err)
|
||||
}
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("second client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
return cl
|
||||
}
|
||||
455
e2e/agentnetwork/gateway_protocol_test.go
Normal file
455
e2e/agentnetwork/gateway_protocol_test.go
Normal file
@@ -0,0 +1,455 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"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"
|
||||
)
|
||||
|
||||
// Models each catalog surface is registered with in the matrix below. They
|
||||
// differ per provider so the router's choice is unambiguous: a request that
|
||||
// lands on the wrong provider record fails the surface assertion instead of
|
||||
// passing by coincidence.
|
||||
const (
|
||||
matrixAnthropicModel = "claude-sonnet-5"
|
||||
matrixBedrockModel = "anthropic.claude-sonnet-5"
|
||||
// matrixBedrockPathModel is what a Bedrock SDK client puts in the URL: a
|
||||
// cross-region inference profile with a release date and version suffix.
|
||||
// The proxy must normalise it back to matrixBedrockModel to route and price.
|
||||
matrixBedrockPathModel = "us.anthropic.claude-sonnet-5-20250101-v1:0"
|
||||
// matrixVertexModel differs from the Anthropic record's model on purpose:
|
||||
// a shared id would leave two routes claiming it and make which one serves
|
||||
// /v1/messages depend on declaration order.
|
||||
matrixVertexModel = "claude-haiku-4-5"
|
||||
matrixVertexProject = "e2e-project"
|
||||
matrixVertexRegion = "us-east5"
|
||||
)
|
||||
|
||||
// gatewayEnv is a connected client plus a set of provider records, all pointed
|
||||
// at one mock upstream, so several wire shapes can be driven over a single
|
||||
// tunnel.
|
||||
type gatewayEnv struct {
|
||||
endpoint string
|
||||
proxyIP string
|
||||
client *harness.Client
|
||||
proxy *harness.Proxy
|
||||
vllm *harness.VLLM
|
||||
// providerIDs maps the catalog id to the created provider record id.
|
||||
providerIDs map[string]string
|
||||
}
|
||||
|
||||
// provisionGatewayMatrix brings up one mock upstream and one provider record
|
||||
// per catalog surface, all authorised for the same group by a single policy.
|
||||
// Sharing one proxy and client keeps the wire-shape cases to one tunnel setup;
|
||||
// each case still creates its own session id so its access-log row is findable.
|
||||
func provisionGatewayMatrix(t *testing.T, ctx context.Context) gatewayEnv {
|
||||
t.Helper()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-matrix"})
|
||||
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-gw-matrix-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// The mock ignores auth, so a dummy credential satisfies each catalog
|
||||
// entry's auth template. Vertex is the exception: its api_key is a GCP
|
||||
// service-account keyfile the proxy mints an OAuth token from, and a dummy
|
||||
// one cannot mint. That is deliberate — the Vertex case below asserts on
|
||||
// routing, which happens before the token mint.
|
||||
dummyKey := "sk-gw-e2e"
|
||||
dummyKeyfile := "keyfile::" + "e2e-not-a-real-service-account-key"
|
||||
|
||||
specs := []struct {
|
||||
name string
|
||||
catalogID string
|
||||
apiKey string
|
||||
models []api.AgentNetworkProviderModel
|
||||
}{
|
||||
{
|
||||
name: "openai", catalogID: "openai_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002}},
|
||||
},
|
||||
{
|
||||
name: "anthropic", catalogID: "anthropic_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixAnthropicModel, InputPer1k: 0.003, OutputPer1k: 0.015}},
|
||||
},
|
||||
{
|
||||
name: "bedrock", catalogID: "bedrock_api", apiKey: dummyKey,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixBedrockModel, InputPer1k: 0.003, OutputPer1k: 0.015}},
|
||||
},
|
||||
{
|
||||
name: "vertex", catalogID: "vertex_ai_api", apiKey: dummyKeyfile,
|
||||
models: []api.AgentNetworkProviderModel{{Id: matrixVertexModel, InputPer1k: 0.001, OutputPer1k: 0.005}},
|
||||
},
|
||||
}
|
||||
|
||||
providerIDs := make(map[string]string, len(specs))
|
||||
ids := make([]string, 0, len(specs))
|
||||
for _, spec := range specs {
|
||||
key := spec.apiKey
|
||||
models := spec.models
|
||||
prov, perr := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-gw-" + spec.name,
|
||||
ProviderId: spec.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &key,
|
||||
Enabled: ptr(true),
|
||||
Models: &models,
|
||||
})
|
||||
require.NoError(t, perr, "create %s provider", spec.name)
|
||||
id := prov.Id
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) })
|
||||
providerIDs[spec.catalogID] = id
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// Uncapped token limit: never blocks the handful of tokens driven here, but
|
||||
// switches on usage metering so consumption and cost land in the row.
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gw-matrix",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: ids,
|
||||
Limits: &api.AgentNetworkPolicyLimits{
|
||||
TokenLimit: api.AgentNetworkPolicyTokenLimit{
|
||||
Enabled: true,
|
||||
GroupCap: 10_000_000,
|
||||
UserCap: 10_000_000,
|
||||
WindowSeconds: 60,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-matrix", sk.Key)
|
||||
return gatewayEnv{
|
||||
endpoint: endpoint,
|
||||
proxyIP: proxyIP,
|
||||
client: cl,
|
||||
proxy: px,
|
||||
vllm: vllm,
|
||||
providerIDs: providerIDs,
|
||||
}
|
||||
}
|
||||
|
||||
// connectClient starts a proxy and a tunnel client for the shared account and
|
||||
// waits until the client can reach the proxy peer, returning the endpoint and
|
||||
// the proxy's tunnel IP to pin requests to.
|
||||
func connectClient(t *testing.T, ctx context.Context, name, setupKey string) (string, string, *harness.Client, *harness.Proxy) {
|
||||
t.Helper()
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-"+name+"-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
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, setupKey)
|
||||
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")
|
||||
// The probe resolves the endpoint and its first packet wakes the lazy proxy
|
||||
// peer, so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
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()))
|
||||
}
|
||||
return settings.Endpoint, proxyIP, cl, px
|
||||
}
|
||||
|
||||
// callUntil retries an HTTP call through the tunnel until it returns one of the
|
||||
// wanted statuses or the deadline passes, absorbing the DNS and tunnel jitter
|
||||
// the first call through a fresh tunnel can hit. The last status and body are
|
||||
// returned either way so the caller can assert with real detail.
|
||||
func callUntil(t *testing.T, call func() (int, string, error), want ...int) (int, string) {
|
||||
t.Helper()
|
||||
wanted := make(map[int]struct{}, len(want))
|
||||
for _, w := range want {
|
||||
wanted[w] = struct{}{}
|
||||
}
|
||||
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
c, b, err := call()
|
||||
if err == nil {
|
||||
code, body = c, b
|
||||
if _, ok := wanted[code]; ok {
|
||||
return code, body
|
||||
}
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
// TestGatewayProtocolProviderMatrix drives one request per wire shape over a
|
||||
// single tunnel, with a provider record per catalog surface behind it. It is
|
||||
// the regression net for the routing and parser-selection changes: each case
|
||||
// asserts the surface the request was metered under and the token counts that
|
||||
// surface's own usage block carries, so a request parsed by the wrong provider's
|
||||
// parser meters zero and fails rather than passing on a coincidence.
|
||||
func TestGatewayProtocolProviderMatrix(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionGatewayMatrix(t, ctx)
|
||||
diag := func() string {
|
||||
return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s",
|
||||
env.vllm.Logs(context.Background()), env.proxy.Logs(context.Background()))
|
||||
}
|
||||
|
||||
t.Run("openai chat completions", func(t *testing.T) {
|
||||
session := "e2e-gw-openai"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, harness.VLLMModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "openai chat must succeed; body: %s%s", body, diag())
|
||||
require.Contains(t, body, "chat.completion", "body must be an OpenAI completion; got: %s", body)
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "openai", *row.Provider, "the OpenAI chat path must meter under the openai surface")
|
||||
assert.Equal(t, int64(harness.VLLMChatInputTokens), row.InputTokens, "OpenAI usage block must be read")
|
||||
assert.Equal(t, int64(harness.VLLMChatOutputTokens), row.OutputTokens)
|
||||
})
|
||||
|
||||
t.Run("anthropic messages", func(t *testing.T) {
|
||||
session := "e2e-gw-anthropic"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, matrixAnthropicModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "anthropic messages must succeed; body: %s%s", body, diag())
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "anthropic", *row.Provider, "the /v1/messages path must meter under the anthropic surface")
|
||||
// These counts only appear if the Anthropic parser read the response:
|
||||
// its usage fields are named differently from the OpenAI block.
|
||||
assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens,
|
||||
"Anthropic input_tokens must be read; zero here means the wrong parser ran")
|
||||
assert.Equal(t, int64(harness.VLLMMessagesOutputTokens), row.OutputTokens)
|
||||
assert.Positive(t, row.CachedInputTokens, "the Anthropic cache-read bucket must be recorded")
|
||||
assert.Positive(t, row.CostUsd, "a metered request must carry a cost")
|
||||
require.NotNil(t, row.ResolvedProviderId)
|
||||
assert.Equal(t, env.providerIDs["anthropic_api"], *row.ResolvedProviderId,
|
||||
"a vendor-tagged request must not cross to another provider's record")
|
||||
})
|
||||
|
||||
t.Run("bedrock invoke normalises the path model", func(t *testing.T) {
|
||||
session := "e2e-gw-bedrock"
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Bedrock(ctx, env.endpoint, env.proxyIP, matrixBedrockPathModel, "ping", session)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "bedrock invoke must succeed; body: %s%s", body, diag())
|
||||
|
||||
row := findAccessLogBySession(t, ctx, session)
|
||||
require.NotNil(t, row.Provider)
|
||||
assert.Equal(t, "bedrock", *row.Provider, "a native Bedrock path must meter under the bedrock surface")
|
||||
require.NotNil(t, row.Model)
|
||||
assert.Equal(t, matrixBedrockModel, *row.Model,
|
||||
"the inference-profile prefix, release date and version suffix must be normalised away")
|
||||
assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens)
|
||||
})
|
||||
|
||||
t.Run("anthropic token counting", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/messages/count_tokens",
|
||||
fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"ping"}]}`, matrixAnthropicModel),
|
||||
[]string{"anthropic-version: 2023-06-01"})
|
||||
}, 200)
|
||||
assert.Equal(t, 200, code, "token counting must route rather than deny; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("bedrock token counting", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP,
|
||||
"/model/"+matrixBedrockPathModel+"/count-tokens",
|
||||
`{"input":{"converse":{"messages":[{"role":"user","content":[{"text":"ping"}]}]}}}`, nil)
|
||||
}, 200)
|
||||
assert.Equal(t, 200, code,
|
||||
"the Bedrock count-tokens action must route; denying it pushes counting onto the billable inference path; body: %s%s",
|
||||
body, diag())
|
||||
})
|
||||
|
||||
t.Run("vertex token counting reaches its provider", func(t *testing.T) {
|
||||
// The dummy service-account key cannot mint an OAuth token, so the
|
||||
// request stops at the upstream credential. Both outcomes render as
|
||||
// 403, so the deny code is what distinguishes them: upstream_auth_failed
|
||||
// means the path resolved to the Vertex route and only the credential
|
||||
// failed, while model_not_routable would mean the method segment was
|
||||
// swallowed into the model id and no route ever claimed it.
|
||||
path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s/count-tokens:rawPredict",
|
||||
matrixVertexProject, matrixVertexRegion, matrixVertexModel)
|
||||
_, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path,
|
||||
`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"ping"}]}`, nil)
|
||||
}, 403)
|
||||
assert.NotContains(t, body, "model_not_routable",
|
||||
"the count-tokens method segment must not be parsed as part of the model id; body: %s%s", body, diag())
|
||||
assert.Contains(t, body, "llm_policy.upstream_auth_failed",
|
||||
"the request must reach the Vertex route and fail only at the credential; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("connection warming probe", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Get(ctx, env.endpoint, env.proxyIP, "/api/hello", nil)
|
||||
}, 200)
|
||||
assert.NotEqual(t, 403, code,
|
||||
"the warm-up probe carries no model and must not be refused as unroutable; body: %s%s", body, diag())
|
||||
})
|
||||
|
||||
t.Run("unknown model denies in the caller's error shape", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages,
|
||||
"claude-not-a-real-model-9", "ping", "e2e-gw-unknown")
|
||||
}, 403)
|
||||
require.Equal(t, 403, code, "a model no provider claims must still be refused; body: %s%s", body, diag())
|
||||
|
||||
// The NetBird fields stay where they were for existing consumers.
|
||||
assert.Contains(t, body, "llm_policy.model_not_routable", "the deny code must be preserved")
|
||||
// And the vendor's own envelope rides alongside, so the client can show
|
||||
// the reason instead of an unexplained API error.
|
||||
assert.Contains(t, body, `"type":"error"`, "an Anthropic caller must get the Anthropic error envelope")
|
||||
assert.Contains(t, body, "permission_error", "403 must map to the vendor's permission error type")
|
||||
})
|
||||
}
|
||||
|
||||
// TestModelDiscoveryWithModelAllowlist covers gateway model discovery on an
|
||||
// account that restricts models, which is the configuration that broke: the
|
||||
// listing carries no model, and the per-model allowlist fails closed on an
|
||||
// undetermined one, so discovery denied for exactly the accounts using the
|
||||
// feature. It also asserts the allowlist still refuses a model outside it, so
|
||||
// the exemption cannot be read as a way around the gate.
|
||||
func TestModelDiscoveryWithModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-discovery"})
|
||||
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-gw-discovery-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// One provider enumerating a single model, while the upstream's own listing
|
||||
// advertises two. The proxy must serve the shorter list.
|
||||
dummyKey := "sk-discovery-e2e"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-gw-discovery",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &dummyKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// The model allowlist is what makes this a regression test: without a
|
||||
// guardrail enabled, discovery was never gated in the first place.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-gw-discovery-allowlist"
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel}
|
||||
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-gw-discovery",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-discovery", sk.Key)
|
||||
diag := func() string {
|
||||
return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s",
|
||||
vllm.Logs(context.Background()), px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
t.Run("listing is served and bounded by policy", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Get(ctx, endpoint, proxyIP, "/v1/models?limit=1000", nil)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code,
|
||||
"discovery must not be refused because the request carries no model; body: %s%s", body, diag())
|
||||
|
||||
assert.Contains(t, body, harness.VLLMModel, "the authorised model must reach the picker")
|
||||
assert.NotContains(t, body, harness.VLLMUnlistedModel,
|
||||
"a model the policy does not authorise must not be offered; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("allowlist still refuses a model outside it", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat,
|
||||
harness.VLLMUnlistedModel, "ping", "e2e-gw-discovery-blocked")
|
||||
}, 403)
|
||||
require.Equal(t, 403, code,
|
||||
"exempting model-less endpoints must not exempt inference; body: %s%s", body, diag())
|
||||
assert.True(t,
|
||||
strings.Contains(body, "llm_policy.model_blocked") || strings.Contains(body, "llm_policy.model_not_routable"),
|
||||
"the refusal must name a model policy code; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("allowlisted model still routes", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat,
|
||||
harness.VLLMModel, "ping", "e2e-gw-discovery-allowed")
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "the allowlisted model must still be served; body: %s%s", body, diag())
|
||||
})
|
||||
}
|
||||
242
e2e/agentnetwork/gateway_review_test.go
Normal file
242
e2e/agentnetwork/gateway_review_test.go
Normal file
@@ -0,0 +1,242 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"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"
|
||||
)
|
||||
|
||||
// The cases in this file cover behaviour that arrived from code review, after
|
||||
// the gateway-protocol end-to-end tests were written. Each had unit coverage
|
||||
// only; none needed a new harness capability, which is why they belong here
|
||||
// rather than on a manual checklist.
|
||||
|
||||
// TestNonInferenceEndpointsAreAuthorised covers the two review findings on the
|
||||
// endpoints that carry no body: the per-model lookup must be authorised
|
||||
// against the same allowlist that bounds the listing beside it, and only a read
|
||||
// method may claim the non-inference exemption that skips the token pre-flight.
|
||||
func TestNonInferenceEndpointsAreAuthorised(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionDiscoveryProvider(t, ctx)
|
||||
|
||||
t.Run("lookup of an authorised model succeeds", func(t *testing.T) {
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMModel, nil)
|
||||
}, 200)
|
||||
assert.Equal(t, 200, code, "an allowlisted model must remain reachable; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("lookup of an unauthorised model is refused", func(t *testing.T) {
|
||||
code, body, err := env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMUnlistedModel, nil)
|
||||
require.NoError(t, err, "request must reach the proxy")
|
||||
assert.Equal(t, 403, code,
|
||||
"a model the policy does not authorise must not be confirmed by the detail lookup; body: %s", body)
|
||||
})
|
||||
|
||||
// A write must not claim the exemption that lets the listing skip the token
|
||||
// pre-flight. The body names no model on purpose: that is what a request
|
||||
// probing for the exemption looks like, and it is the case the method gate
|
||||
// exists to refuse. (A POST that does name a model is a different thing —
|
||||
// it routes and meters as the inference request it is.)
|
||||
for _, path := range []string{"/v1/models", "/v1/models/" + harness.VLLMModel, "/api/hello"} {
|
||||
t.Run("write to "+path+" is refused", func(t *testing.T) {
|
||||
code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path,
|
||||
`{"messages":[{"role":"user","content":"hi"}]}`, nil)
|
||||
require.NoError(t, err, "request must reach the proxy")
|
||||
assert.NotEqual(t, 200, code,
|
||||
"a write to a non-inference path must not be served unmetered; body: %s", body)
|
||||
})
|
||||
}
|
||||
|
||||
// A request carrying the sub-agent attribution headers must still be served
|
||||
// and metered normally. Asserting the ids themselves is not possible yet:
|
||||
// the parser lifts them onto the request's metadata, but nothing persists
|
||||
// them, so they have no queryable surface to check against.
|
||||
t.Run("sub-agent headers do not disturb the request", func(t *testing.T) {
|
||||
sessionID := fmt.Sprintf("e2e-session-agentid-%d", time.Now().UnixNano())
|
||||
code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/chat/completions",
|
||||
fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"Reply with exactly: pong"}]}`, harness.VLLMModel),
|
||||
[]string{
|
||||
"x-session-id: " + sessionID,
|
||||
"x-claude-code-agent-id: agent-child-7",
|
||||
"x-claude-code-parent-agent-id: agent-root-1",
|
||||
})
|
||||
require.NoError(t, err, "request must reach the proxy")
|
||||
require.Equal(t, 200, code, "the request must succeed; body: %s", body)
|
||||
|
||||
row := findAccessLogBySession(t, ctx, sessionID)
|
||||
assert.Positive(t, row.InputTokens, "the request must still be metered normally")
|
||||
})
|
||||
}
|
||||
|
||||
// TestDatedModelIdRouting covers both halves of the dated-id rule that review
|
||||
// tightened: a dated id still reaches an undated registration, but a route
|
||||
// pinned to one dated build must never serve a different one.
|
||||
func TestDatedModelIdRouting(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
undated = "claude-sonnet-9"
|
||||
datedA = "claude-sonnet-9-20250101"
|
||||
datedB = "claude-sonnet-9-20250202"
|
||||
)
|
||||
|
||||
t.Run("a dated id reaches its undated registration", func(t *testing.T) {
|
||||
env := provisionModelProvider(t, ctx, "dated-undated", "anthropic_api", undated)
|
||||
|
||||
sessionID := fmt.Sprintf("e2e-session-dated-%d", time.Now().UnixNano())
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", sessionID)
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "a pinned release of a registered family must route; body: %s", body)
|
||||
|
||||
row := findAccessLogBySession(t, ctx, sessionID)
|
||||
assert.Positive(t, row.InputTokens, "the dated request must price at the registered rate, not zero")
|
||||
})
|
||||
|
||||
t.Run("a route pinned to one dated build refuses another", func(t *testing.T) {
|
||||
env := provisionModelProvider(t, ctx, "dated-pinned", "anthropic_api", datedA)
|
||||
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", "")
|
||||
}, 200)
|
||||
require.Equal(t, 200, code, "the exact dated id must still route; body: %s", body)
|
||||
|
||||
code, body, err := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedB, "Reply with exactly: pong", "")
|
||||
require.NoError(t, err, "request must reach the proxy")
|
||||
assert.Equal(t, 403, code,
|
||||
"a provider pinned to one dated build must not serve another; body: %s", body)
|
||||
})
|
||||
}
|
||||
|
||||
// TestBedrockInferenceProfilesReachTheUpstream covers the startup lookup a
|
||||
// Bedrock client makes. The proxy forwards it to the configured upstream rather
|
||||
// than denying it, so what comes back is the upstream's answer — never a
|
||||
// NetBird policy rejection.
|
||||
func TestBedrockInferenceProfilesReachTheUpstream(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionModelProvider(t, ctx, "infprofiles", "bedrock_api", "anthropic.claude-sonnet-5")
|
||||
|
||||
code, body := callUntil(t, func() (int, string, error) {
|
||||
return env.client.Get(ctx, env.endpoint, env.proxyIP, "/inference-profiles", nil)
|
||||
}, 200)
|
||||
|
||||
assert.Equal(t, 200, code, "the lookup must reach the upstream; body: %s", body)
|
||||
assert.NotContains(t, body, "llm_policy.",
|
||||
"the proxy must not answer a control-plane lookup with a policy denial")
|
||||
assert.Contains(t, body, "inferenceProfileSummaries",
|
||||
"the upstream's own answer must come back untouched")
|
||||
}
|
||||
|
||||
// provisionDiscoveryProvider brings up one mock-backed provider enumerating a
|
||||
// single model, with an allowlist guardrail in effect, plus a connected client.
|
||||
func provisionDiscoveryProvider(t *testing.T, ctx context.Context) pricedEnv {
|
||||
t.Helper()
|
||||
env := provisionModelProvider(t, ctx, "noninference", "openai_api", harness.VLLMModel)
|
||||
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-noninference-allowlist-" + fmt.Sprint(time.Now().UnixNano())
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
_, err = srv.UpdatePolicy(ctx, env.policyID, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-noninference",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{env.groupID},
|
||||
DestinationProviderIds: []string{env.providerID},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "attach guardrail to policy")
|
||||
return env
|
||||
}
|
||||
|
||||
// provisionModelProvider brings up the mock, one provider under the given
|
||||
// catalog id enumerating exactly one model, an authorising policy, and a
|
||||
// connected proxy + client.
|
||||
func provisionModelProvider(t *testing.T, ctx context.Context, name, catalogID, model string) pricedEnv {
|
||||
t.Helper()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
suffix := strings.ToLower(name)
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gwr-" + suffix})
|
||||
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-gwr-" + suffix + "-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
dummyKey := "sk-gwr-e2e"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-gwr-" + suffix,
|
||||
ProviderId: catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &dummyKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: model, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gwr-" + suffix,
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
Limits: &api.AgentNetworkPolicyLimits{
|
||||
TokenLimit: api.AgentNetworkPolicyTokenLimit{
|
||||
Enabled: true,
|
||||
GroupCap: 10_000_000,
|
||||
UserCap: 10_000_000,
|
||||
WindowSeconds: 60,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, "gwr-"+suffix, sk.Key)
|
||||
return pricedEnv{
|
||||
providerID: prov.Id,
|
||||
groupID: grp.Id,
|
||||
policyID: pol.Id,
|
||||
upstream: vllm.URL,
|
||||
endpoint: endpoint,
|
||||
proxyIP: proxyIP,
|
||||
client: cl,
|
||||
proxy: px,
|
||||
}
|
||||
}
|
||||
199
e2e/agentnetwork/streaming_test.go
Normal file
199
e2e/agentnetwork/streaming_test.go
Normal file
@@ -0,0 +1,199 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"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"
|
||||
)
|
||||
|
||||
// streamedModel is priced high enough that a mis-metered request is obvious in
|
||||
// the recorded cost, and named so it cannot collide with another test's route.
|
||||
const streamedModel = "e2e-streamed-model"
|
||||
|
||||
const (
|
||||
streamInRate = 0.010
|
||||
streamOutRate = 0.020
|
||||
// The cache-read bucket is priced separately from input, so a run that
|
||||
// folded the two together fails the per-bucket assertions below.
|
||||
streamCacheReadRate = 0.001
|
||||
)
|
||||
|
||||
// TestStreamingResponseMetersInputTokens is the end-to-end guard for the
|
||||
// metering bug this endpoint's gateway-protocol work fixed.
|
||||
//
|
||||
// On a streamed answer the input-token count exists only in the opening
|
||||
// message_start event; every later frame reports output. A response read with
|
||||
// the wrong vendor's parser — the shape a gateway record produces when it names
|
||||
// one API surface and serves another — never looks at that event, so input
|
||||
// metered as zero and the bulk of the bill silently vanished. Nothing in the
|
||||
// suite sent stream: true before this test, so the whole branch went unrun.
|
||||
//
|
||||
// The provider points at the mock's streaming listener, which answers every
|
||||
// request as SSE with token counts that differ from the buffered surface. That
|
||||
// difference is the point: passing these assertions is only possible if the
|
||||
// stream accumulator ran.
|
||||
func TestStreamingResponseMetersInputTokens(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionStreamingProvider(t, ctx, "anthropic_api")
|
||||
|
||||
sessionID := fmt.Sprintf("e2e-session-stream-%d", time.Now().UnixNano())
|
||||
code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID)
|
||||
require.Equal(t, 200, code, "streamed chat must succeed; body: %s", body)
|
||||
assert.Contains(t, body, "message_start",
|
||||
"the client must receive the event stream itself, not a buffered rewrite of it")
|
||||
|
||||
row := findAccessLogBySession(t, ctx, sessionID)
|
||||
|
||||
assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens),
|
||||
"input tokens live in message_start; zero here is the bug this test exists for")
|
||||
assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens),
|
||||
"output tokens ride message_delta and supersede the message_start seed")
|
||||
assert.Equal(t, harness.VLLMStreamCacheReadTokens, int(row.CachedInputTokens),
|
||||
"the Anthropic cache bucket rides message_start too, and only its own parser reads it")
|
||||
|
||||
// The Anthropic surface bills cache reads additively, so the input bucket
|
||||
// prices the full input count rather than a remainder.
|
||||
wantInput := float64(harness.VLLMStreamInputTokens) / 1000 * streamInRate
|
||||
wantOutput := float64(harness.VLLMStreamOutputTokens) / 1000 * streamOutRate
|
||||
assert.InDelta(t, wantInput, row.InputCostUsd, 1e-6, "input cost must price the streamed input tokens")
|
||||
assert.InDelta(t, wantOutput, row.OutputCostUsd, 1e-6, "output cost must price the streamed output tokens")
|
||||
assert.Greater(t, row.CostUsd, 0.0, "a streamed request must never record as free")
|
||||
}
|
||||
|
||||
// TestStreamingOnGatewayTypedProvider drives the same streamed Anthropic call
|
||||
// through a provider record whose catalog id names the OpenAI surface — the
|
||||
// exact misconfiguration that hid the bug, since gateway records commonly pin
|
||||
// one parser while the upstream serves another shape entirely.
|
||||
//
|
||||
// The router must choose the parser from the request path rather than the
|
||||
// record's provider id, or the Anthropic usage block goes unread and input
|
||||
// meters at zero all over again.
|
||||
func TestStreamingOnGatewayTypedProvider(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
env := provisionStreamingProvider(t, ctx, "openai_api")
|
||||
|
||||
sessionID := fmt.Sprintf("e2e-session-stream-gw-%d", time.Now().UnixNano())
|
||||
code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID)
|
||||
require.Equal(t, 200, code, "streamed chat through a gateway record must succeed; body: %s", body)
|
||||
|
||||
row := findAccessLogBySession(t, ctx, sessionID)
|
||||
|
||||
assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens),
|
||||
"a record typed openai_api must still read the Anthropic usage block it is actually serving")
|
||||
assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens),
|
||||
"output tokens must survive the surface mismatch too")
|
||||
assert.InDelta(t, float64(harness.VLLMStreamInputTokens)/1000*streamInRate, row.InputCostUsd, 1e-6,
|
||||
"the request must be priced on the surface it spoke, not the one the record names")
|
||||
}
|
||||
|
||||
// provisionStreamingProvider brings up the mock, one provider pointed at its
|
||||
// streaming listener under the given catalog id, a policy authorising it, and a
|
||||
// connected proxy + client.
|
||||
func provisionStreamingProvider(t *testing.T, ctx context.Context, catalogID string) pricedEnv {
|
||||
t.Helper()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
name := "stream-" + catalogID
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-" + name})
|
||||
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-" + name + "-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
dummyKey := "sk-stream-e2e"
|
||||
cacheRead := streamCacheReadRate
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: name,
|
||||
ProviderId: catalogID,
|
||||
UpstreamUrl: vllm.StreamURL,
|
||||
ApiKey: &dummyKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{{
|
||||
Id: streamedModel,
|
||||
InputPer1k: streamInRate,
|
||||
OutputPer1k: streamOutRate,
|
||||
CacheReadPer1k: &cacheRead,
|
||||
}},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-" + name,
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
Limits: &api.AgentNetworkPolicyLimits{
|
||||
TokenLimit: api.AgentNetworkPolicyTokenLimit{
|
||||
Enabled: true,
|
||||
GroupCap: 10_000_000,
|
||||
UserCap: 10_000_000,
|
||||
WindowSeconds: 60,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
endpoint, proxyIP, cl, px := connectClient(t, ctx, name, sk.Key)
|
||||
return pricedEnv{
|
||||
providerID: prov.Id,
|
||||
groupID: grp.Id,
|
||||
policyID: pol.Id,
|
||||
upstream: vllm.StreamURL,
|
||||
endpoint: endpoint,
|
||||
proxyIP: proxyIP,
|
||||
client: cl,
|
||||
proxy: px,
|
||||
}
|
||||
}
|
||||
|
||||
// chatStreamUntil drives one streamed chat, retrying to absorb the tunnel and
|
||||
// DNS jitter a first call through a fresh peer can hit.
|
||||
func chatStreamUntil(t *testing.T, ctx context.Context, env pricedEnv, kind, model, sessionID string) (int, string) {
|
||||
t.Helper()
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
c, b, cerr := env.client.ChatStream(ctx, env.endpoint, env.proxyIP, kind, model, "Reply with exactly: pong", sessionID)
|
||||
if cerr == nil {
|
||||
code, body = c, b
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
if code != 200 {
|
||||
t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background()))
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -199,12 +200,18 @@ func (cl *Client) pollStatus(ctx context.Context, timeout time.Duration, want st
|
||||
const (
|
||||
// curlExitCouldNotResolve is curl's exit code for a DNS resolution failure, distinct from connection-level failures.
|
||||
curlExitCouldNotResolve = 6
|
||||
// dnsProbeRetryWindow bounds DNS-failure retries: the synthesized zone lands a beat after management connects, so early NXDOMAIN is propagation; a zone still absent after this window is a real failure.
|
||||
dnsProbeRetryWindow = 30 * time.Second
|
||||
dnsProbeRetryInterval = 2 * time.Second
|
||||
// curlExitCouldNotConnect is curl's exit code for a connection that never
|
||||
// established. The probe exists to WAKE the lazy proxy peer, so the first
|
||||
// attempt legitimately arrives before WireGuard has brought the tunnel up
|
||||
// and fails here — which is propagation, exactly like an early NXDOMAIN,
|
||||
// and belongs inside the retry window rather than failing the test outright.
|
||||
curlExitCouldNotConnect = 7
|
||||
// endpointProbeRetryWindow bounds retries of the transient failures above: the synthesized zone and the tunnel both land a beat after management connects. Still failing after this window is a real failure.
|
||||
endpointProbeRetryWindow = 30 * time.Second
|
||||
endpointProbeRetryInterval = 2 * time.Second
|
||||
)
|
||||
|
||||
// ResolveProxyIP GETs https://<endpoint>/ from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; only DNS failures retry, within dnsProbeRetryWindow. Returns the connected IP for --resolve pinning.
|
||||
// ResolveProxyIP GETs https://<endpoint>/ from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; DNS and connect failures retry, within endpointProbeRetryWindow. Returns the connected IP for --resolve pinning.
|
||||
func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
|
||||
args := []string{
|
||||
"run", "--rm",
|
||||
@@ -215,7 +222,7 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string,
|
||||
"-w", "%{remote_ip}",
|
||||
"https://" + endpoint + "/",
|
||||
}
|
||||
deadline := time.Now().Add(dnsProbeRetryWindow)
|
||||
deadline := time.Now().Add(endpointProbeRetryWindow)
|
||||
for {
|
||||
cmd := exec.CommandContext(ctx, "docker", args...)
|
||||
var stdout, stderr strings.Builder
|
||||
@@ -231,21 +238,29 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string,
|
||||
}
|
||||
|
||||
var exitErr *exec.ExitError
|
||||
if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve {
|
||||
if !errors.As(err, &exitErr) || !isTransientProbeExit(exitErr.ExitCode()) {
|
||||
return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String()))
|
||||
}
|
||||
dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String()))
|
||||
if time.Until(deadline) < dnsProbeRetryInterval {
|
||||
return "", dnsErr
|
||||
probeErr := fmt.Errorf("endpoint %s not reachable yet: %s", endpoint, strings.TrimSpace(stderr.String()))
|
||||
if time.Until(deadline) < endpointProbeRetryInterval {
|
||||
return "", probeErr
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err())
|
||||
case <-time.After(dnsProbeRetryInterval):
|
||||
return "", fmt.Errorf("%w (%w)", probeErr, ctx.Err())
|
||||
case <-time.After(endpointProbeRetryInterval):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// isTransientProbeExit reports whether a curl exit code describes a state the
|
||||
// endpoint is expected to pass THROUGH on its way up, rather than a settled
|
||||
// failure. Anything else — TLS refusal, a protocol error, a bad argument —
|
||||
// would still be failing after the retry window, so it fails immediately.
|
||||
func isTransientProbeExit(code int) bool {
|
||||
return code == curlExitCouldNotResolve || code == curlExitCouldNotConnect
|
||||
}
|
||||
|
||||
// Wire shapes for Chat.
|
||||
const (
|
||||
// WireChat is the OpenAI-compatible /v1/chat/completions shape.
|
||||
@@ -292,6 +307,27 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi
|
||||
return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
|
||||
}
|
||||
|
||||
// ChatStream is Chat with "stream": true in the request body, so the proxy's
|
||||
// request parser marks the call as streaming and its response parser takes the
|
||||
// SSE accumulator rather than the buffered-body path. Pair it with a provider
|
||||
// pointed at VLLM.StreamURL, which answers every request as an event stream.
|
||||
func (cl *Client) ChatStream(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) {
|
||||
var path, body string
|
||||
var headers []string
|
||||
switch kind {
|
||||
case WireMessages:
|
||||
path = "/v1/messages"
|
||||
headers = []string{"anthropic-version: 2023-06-01"}
|
||||
body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"stream":true,"messages":[{"role":"user","content":%q}]}`, model, prompt)
|
||||
default:
|
||||
path = "/v1/chat/completions"
|
||||
// include_usage is what makes a real OpenAI stream emit its final usage
|
||||
// frame; without it the last chunk carries no tokens at all.
|
||||
body = fmt.Sprintf(`{"model":%q,"stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":%q}]}`, model, prompt)
|
||||
}
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID))
|
||||
}
|
||||
|
||||
// Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
|
||||
// Chat, the model is carried in the request path (project/region/model), so the
|
||||
// proxy routes by path and mints the service-account OAuth token; the body uses
|
||||
@@ -322,10 +358,29 @@ func withSessionID(headers []string, sessionID string) []string {
|
||||
return append(headers, "x-session-id: "+sessionID)
|
||||
}
|
||||
|
||||
// post runs curl in a throwaway container sharing the client's network
|
||||
// namespace so the request traverses the WireGuard tunnel, pinning the endpoint
|
||||
// to the proxy IP. It returns the HTTP status and response body.
|
||||
// Get issues a GET to the agent-network endpoint over the client's tunnel.
|
||||
// Model discovery and the connection-warming probe are read-only endpoints
|
||||
// that carry no body, so they can't go through the chat helpers.
|
||||
func (cl *Client) Get(ctx context.Context, endpoint, proxyIP, path string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodGet, endpoint, proxyIP, path, "", extraHeaders)
|
||||
}
|
||||
|
||||
// PostJSON issues an arbitrary JSON POST over the client's tunnel, for wire
|
||||
// shapes the typed helpers don't cover (token counting, say).
|
||||
func (cl *Client) PostJSON(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders)
|
||||
}
|
||||
|
||||
// post issues a JSON POST. Retained as the shorthand the chat helpers use.
|
||||
func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders)
|
||||
}
|
||||
|
||||
// do runs curl in a throwaway container sharing the client's network
|
||||
// namespace so the request traverses the WireGuard tunnel, pinning the endpoint
|
||||
// to the proxy IP. It returns the HTTP status and response body. An empty body
|
||||
// sends no payload, which is what a GET needs.
|
||||
func (cl *Client) do(ctx context.Context, method, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
|
||||
url := "https://" + endpoint + path
|
||||
args := []string{
|
||||
"run", "--rm",
|
||||
@@ -334,13 +389,15 @@ func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string
|
||||
"-sk", "--connect-timeout", "5", "--max-time", "90",
|
||||
"--resolve", endpoint + ":443:" + proxyIP,
|
||||
"-o", "/dev/stderr", "-w", "%{http_code}",
|
||||
"-X", "POST", url,
|
||||
"-X", method, url,
|
||||
"-H", "Content-Type: application/json",
|
||||
}
|
||||
for _, h := range extraHeaders {
|
||||
args = append(args, "-H", h)
|
||||
}
|
||||
args = append(args, "--data", body)
|
||||
if body != "" {
|
||||
args = append(args, "--data", body)
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, "docker", args...)
|
||||
// -w writes the status code to stdout; -o /dev/stderr writes the body to
|
||||
// stderr so we can capture both separately.
|
||||
|
||||
@@ -18,18 +18,63 @@ const (
|
||||
vllmImage = "nginx:alpine"
|
||||
vllmAlias = "vllm"
|
||||
vllmPort = "8000/tcp"
|
||||
// vllmStreamPort serves the same wire shapes as an SSE stream. See the
|
||||
// nginx config for why streaming lives on its own listener.
|
||||
vllmStreamPort = "8001/tcp"
|
||||
// VLLMModel is the served model id the mock advertises and echoes back. It
|
||||
// matches a real small model commonly served by vLLM so the provider's
|
||||
// enumerated model and the client's request line up.
|
||||
VLLMModel = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
// VLLMUnlistedModel is a second id the mock's model listing advertises but
|
||||
// no test provider enumerates, so a filtered listing is observably shorter
|
||||
// than the upstream's own.
|
||||
VLLMUnlistedModel = "Qwen/Qwen2.5-7B-Instruct"
|
||||
)
|
||||
|
||||
// Token counts the mock reports per wire shape. Tests assert on these rather
|
||||
// than on "> 0" so a response parsed with the wrong provider's parser (which
|
||||
// would read a different field, or none) fails loudly instead of passing on
|
||||
// a coincidental non-zero.
|
||||
const (
|
||||
// VLLMChatInputTokens / VLLMChatOutputTokens ride the OpenAI usage block.
|
||||
VLLMChatInputTokens = 11
|
||||
VLLMChatOutputTokens = 2
|
||||
// VLLMMessagesInputTokens / VLLMMessagesOutputTokens ride the Anthropic
|
||||
// usage block, whose field names the OpenAI parser cannot read.
|
||||
VLLMMessagesInputTokens = 17
|
||||
VLLMMessagesOutputTokens = 3
|
||||
)
|
||||
|
||||
// Token counts the streaming surface reports. They differ from the
|
||||
// non-streaming ones on purpose: a test that asserts these numbers proves the
|
||||
// SSE accumulator ran, rather than a buffered JSON body having been parsed.
|
||||
//
|
||||
// Input and cache-read arrive on message_start; output arrives on
|
||||
// message_delta and supersedes the seed value message_start carries. Any
|
||||
// parser that cannot read message_start reports zero input tokens — which is
|
||||
// exactly the bug these counts exist to catch.
|
||||
const (
|
||||
VLLMStreamInputTokens = 29
|
||||
VLLMStreamOutputTokens = 5
|
||||
VLLMStreamCacheReadTokens = 7
|
||||
)
|
||||
|
||||
// vllmNginxConf emulates a vLLM OpenAI-compatible server over plain HTTP (vLLM's
|
||||
// default: no TLS, port 8000). It answers /v1/models with a one-model list and
|
||||
// any chat/completions path with a canned OpenAI-shaped chat completion carrying
|
||||
// a non-zero usage block, so the proxy's OpenAI parser records real token
|
||||
// consumption. Running actual vLLM in CI is infeasible (GPU + multi-GB model
|
||||
// default: no TLS, port 8000), and additionally answers the wire shapes the
|
||||
// other catalog surfaces speak so one mock can stand in for every provider the
|
||||
// proxy routes to. Running actual vLLM in CI is infeasible (GPU + multi-GB model
|
||||
// download), so this stands in for the wire contract the proxy depends on.
|
||||
//
|
||||
// Each shape answers with its own vendor's usage block, so a response parsed
|
||||
// under the wrong surface meters zero rather than passing by accident:
|
||||
//
|
||||
// - /v1/chat/completions (and any unmatched path): OpenAI chat completion.
|
||||
// - /v1/messages: Anthropic Messages, snake_case usage plus a cache bucket.
|
||||
// - /model/{id}/invoke: Bedrock InvokeModel, which carries the Anthropic body.
|
||||
// - the token-counting endpoints: a count, with no usage block at all.
|
||||
//
|
||||
// The model listing advertises two models so a policy that authorises one
|
||||
// produces an observably shorter list than the upstream's own.
|
||||
const vllmNginxConf = `pid /tmp/nginx.pid;
|
||||
events {}
|
||||
http {
|
||||
@@ -37,13 +82,75 @@ http {
|
||||
listen 8000;
|
||||
location = /v1/models {
|
||||
default_type application/json;
|
||||
return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"}]}';
|
||||
return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"},{"id":"Qwen/Qwen2.5-7B-Instruct","object":"model","owned_by":"vllm"}]}';
|
||||
}
|
||||
location = /v1/messages {
|
||||
default_type application/json;
|
||||
return 200 '{"id":"msg_e2e","type":"message","role":"assistant","model":"claude-sonnet-5","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}';
|
||||
}
|
||||
location = /v1/messages/count_tokens {
|
||||
default_type application/json;
|
||||
return 200 '{"input_tokens":7}';
|
||||
}
|
||||
location ~ ^/model/.+/invoke$ {
|
||||
default_type application/json;
|
||||
return 200 '{"id":"msg_e2e_bedrock","type":"message","role":"assistant","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}';
|
||||
}
|
||||
location ~ ^/model/.+/count-tokens$ {
|
||||
default_type application/json;
|
||||
return 200 '{"inputTokens":9}';
|
||||
}
|
||||
location = /api/hello {
|
||||
return 200;
|
||||
}
|
||||
location = /inference-profiles {
|
||||
default_type application/json;
|
||||
return 200 '{"inferenceProfileSummaries":[{"inferenceProfileId":"us.anthropic.claude-sonnet-5","status":"ACTIVE"}]}';
|
||||
}
|
||||
location / {
|
||||
default_type application/json;
|
||||
return 200 '{"id":"chatcmpl-e2e-vllm","object":"chat.completion","created":1700000000,"model":"Qwen/Qwen2.5-0.5B-Instruct","choices":[{"index":0,"message":{"role":"assistant","content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":2,"total_tokens":13}}';
|
||||
}
|
||||
}
|
||||
|
||||
# The streaming surface, on its own port so the response content type is a
|
||||
# property of the listener rather than of a per-request branch: nginx sets
|
||||
# Content-Type from default_type, which cannot be varied inside an "if", and
|
||||
# a second Content-Type via add_header would leave the proxy reading the
|
||||
# wrong one. A provider record pointed at this port streams every answer.
|
||||
#
|
||||
# Input and cache-read tokens ride message_start, output rides message_delta
|
||||
# — the split that makes a stream different from a buffered body, and the
|
||||
# reason a parser that ignores message_start meters input as zero.
|
||||
server {
|
||||
listen 8001;
|
||||
location = /v1/messages {
|
||||
default_type text/event-stream;
|
||||
return 200 'event: message_start
|
||||
data: {"type":"message_start","message":{"id":"msg_e2e_stream","type":"message","role":"assistant","model":"claude-sonnet-5","content":[],"usage":{"input_tokens":29,"output_tokens":1,"cache_read_input_tokens":7}}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}}
|
||||
|
||||
event: message_delta
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}
|
||||
|
||||
event: message_stop
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
';
|
||||
}
|
||||
location / {
|
||||
default_type text/event-stream;
|
||||
return 200 'data: {"choices":[{"delta":{"content":"pong"}}]}
|
||||
|
||||
data: {"choices":[],"usage":{"prompt_tokens":29,"completion_tokens":5,"total_tokens":34}}
|
||||
|
||||
data: [DONE]
|
||||
|
||||
';
|
||||
}
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
@@ -55,6 +162,10 @@ type VLLM struct {
|
||||
workDir string
|
||||
// URL is the upstream URL the vllm provider points at (http://<alias>:8000).
|
||||
URL string
|
||||
// StreamURL is the same mock's streaming listener. A provider pointed here
|
||||
// answers every request as SSE, so the proxy's streaming accumulator runs
|
||||
// instead of its buffered-body parser.
|
||||
StreamURL string
|
||||
}
|
||||
|
||||
// StartVLLM runs the mock vLLM server on the shared network over plain HTTP.
|
||||
@@ -73,14 +184,17 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) {
|
||||
|
||||
req := testcontainers.ContainerRequest{
|
||||
Image: vllmImage,
|
||||
ExposedPorts: []string{vllmPort},
|
||||
ExposedPorts: []string{vllmPort, vllmStreamPort},
|
||||
Networks: []string{c.network.Name},
|
||||
NetworkAliases: map[string][]string{c.network.Name: {vllmAlias}},
|
||||
Cmd: []string{"nginx", "-c", "/conf/nginx.conf", "-g", "daemon off;"},
|
||||
HostConfigModifier: func(hc *container.HostConfig) {
|
||||
hc.Binds = append(hc.Binds, workDir+":/conf:ro")
|
||||
},
|
||||
WaitingFor: wait.ForListeningPort(vllmPort).WithStartupTimeout(60 * time.Second),
|
||||
WaitingFor: wait.ForAll(
|
||||
wait.ForListeningPort(vllmPort),
|
||||
wait.ForListeningPort(vllmStreamPort),
|
||||
).WithStartupTimeout(60 * time.Second),
|
||||
}
|
||||
|
||||
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
@@ -92,7 +206,12 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) {
|
||||
return nil, fmt.Errorf("start vllm container: %w", err)
|
||||
}
|
||||
|
||||
return &VLLM{container: ctr, workDir: workDir, URL: "http://" + vllmAlias + ":8000"}, nil
|
||||
return &VLLM{
|
||||
container: ctr,
|
||||
workDir: workDir,
|
||||
URL: "http://" + vllmAlias + ":8000",
|
||||
StreamURL: "http://" + vllmAlias + ":8001",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Logs returns the vLLM container logs, for diagnostics on failure.
|
||||
|
||||
12
funding.json
12
funding.json
@@ -6,7 +6,7 @@
|
||||
"name": "NetBird GmbH",
|
||||
"email": "hello@netbird.io",
|
||||
"phone": "",
|
||||
"description": "NetBird GmbH is a Berlin-based software company specializing in the development of open source network security solutions. Network security is utterly complex and expensive, accessible only to companies with multi-million dollar IT budgets. In contrast, there are millions of companies left behind. Our mission is to create an advanced network and cybersecurity platform that is both easy-to-use and affordable for teams of all sizes and budgets. By leveraging the open source strategy and technological advancements, NetBird aims to set the industry standard for connecting and securing IT infrastructure.",
|
||||
"description": "NetBird GmbH is a Berlin-based software company specializing in the development of open-source network security solutions. Network security is utterly complex and expensive, accessible only to companies with multi-million dollar IT budgets. In contrast, there are millions of companies left behind. Our mission is to create an advanced network and cybersecurity platform that is both easy-to-use and affordable for teams of all sizes and budgets. By leveraging the open-source strategy and technological advancements, NetBird aims to set the industry standard for connecting and securing IT infrastructure.",
|
||||
"webpageUrl": {
|
||||
"url": "https://github.com/netbirdio"
|
||||
}
|
||||
@@ -15,7 +15,7 @@
|
||||
{
|
||||
"guid": "netbird",
|
||||
"name": "NetBird",
|
||||
"description": "NetBird is a configuration-free peer-to-peer private network and a centralized access control system combined in a single open source platform. It makes it easy to create secure WireGuard-based private networks for your organization or home.",
|
||||
"description": "NetBird is a configuration-free peer-to-peer private network and a centralized access control system combined in a single open-source platform. It makes it easy to create secure WireGuard-based private networks for your organization or home.",
|
||||
"webpageUrl": {
|
||||
"url": "https://github.com/netbirdio/netbird"
|
||||
},
|
||||
@@ -59,7 +59,7 @@
|
||||
"guid": "support-yearly",
|
||||
"status": "active",
|
||||
"name": "Support Open Source Development and Maintenance - Yearly",
|
||||
"description": "This will help us partially cover the yearly cost of maintaining the open source NetBird project.",
|
||||
"description": "This will help us partially cover the yearly cost of maintaining the open-source NetBird project.",
|
||||
"amount": 100000,
|
||||
"currency": "USD",
|
||||
"frequency": "yearly",
|
||||
@@ -72,7 +72,7 @@
|
||||
"guid": "support-one-time-year",
|
||||
"status": "active",
|
||||
"name": "Support Open Source Development and Maintenance - One Year",
|
||||
"description": "This will help us partially cover the yearly cost of maintaining the open source NetBird project.",
|
||||
"description": "This will help us partially cover the yearly cost of maintaining the open-source NetBird project.",
|
||||
"amount": 100000,
|
||||
"currency": "USD",
|
||||
"frequency": "one-time",
|
||||
@@ -85,7 +85,7 @@
|
||||
"guid": "support-one-time-monthly",
|
||||
"status": "active",
|
||||
"name": "Support Open Source Development and Maintenance - Monthly",
|
||||
"description": "This will help us partially cover the monthly cost of maintaining the open source NetBird project.",
|
||||
"description": "This will help us partially cover the monthly cost of maintaining the open-source NetBird project.",
|
||||
"amount": 10000,
|
||||
"currency": "USD",
|
||||
"frequency": "monthly",
|
||||
@@ -98,7 +98,7 @@
|
||||
"guid": "support-monthly",
|
||||
"status": "active",
|
||||
"name": "Support Open Source Development and Maintenance - One Month",
|
||||
"description": "This will help us partially cover the monthly cost of maintaining the open source NetBird project.",
|
||||
"description": "This will help us partially cover the monthly cost of maintaining the open-source NetBird project.",
|
||||
"amount": 10000,
|
||||
"currency": "USD",
|
||||
"frequency": "monthly",
|
||||
|
||||
12
go.mod
12
go.mod
@@ -19,7 +19,7 @@ require (
|
||||
github.com/spf13/cobra v1.10.2
|
||||
github.com/spf13/pflag v1.0.10
|
||||
github.com/vishvananda/netlink v1.3.1
|
||||
golang.org/x/crypto v0.55.0
|
||||
golang.org/x/crypto v0.54.0
|
||||
golang.org/x/sys v0.47.0
|
||||
golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
|
||||
golang.zx2c4.com/wireguard/wgctrl v0.0.0-20241231184526-a9ab2273dd10
|
||||
@@ -127,9 +127,9 @@ require (
|
||||
go.uber.org/zap v1.27.0
|
||||
goauthentik.io/api/v3 v3.2023051.3
|
||||
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f
|
||||
golang.org/x/mobile v0.0.0-20260816165457-f98cc9b3c733
|
||||
golang.org/x/mod v0.39.0
|
||||
golang.org/x/net v0.58.0
|
||||
golang.org/x/mobile v0.0.0-20251113184115-a159579294ab
|
||||
golang.org/x/mod v0.37.0
|
||||
golang.org/x/net v0.56.0
|
||||
golang.org/x/oauth2 v0.36.0
|
||||
golang.org/x/sync v0.22.0
|
||||
golang.org/x/term v0.45.0
|
||||
@@ -313,8 +313,8 @@ require (
|
||||
go.opentelemetry.io/otel/trace v1.43.0 // indirect
|
||||
go.uber.org/multierr v1.11.0 // indirect
|
||||
go.yaml.in/yaml/v2 v2.4.3 // indirect
|
||||
golang.org/x/text v0.41.0 // indirect
|
||||
golang.org/x/tools v0.49.0 // indirect
|
||||
golang.org/x/text v0.40.0 // indirect
|
||||
golang.org/x/tools v0.47.0 // indirect
|
||||
golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7 // indirect
|
||||
gopkg.in/square/go-jose.v2 v2.6.0 // indirect
|
||||
|
||||
24
go.sum
24
go.sum
@@ -728,13 +728,13 @@ golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1m
|
||||
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
|
||||
golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M=
|
||||
golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis=
|
||||
golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
|
||||
golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
|
||||
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f h1:W3F4c+6OLc6H2lb//N1q4WpJkhzJCK5J6kUi1NTVXfM=
|
||||
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f/go.mod h1:J1xhfL/vlindoeF/aINzNzt2Bket5bjo9sdOYzOsU80=
|
||||
golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
|
||||
golang.org/x/mobile v0.0.0-20260816165457-f98cc9b3c733 h1:XKMObIaAElmkdO+4SQh1iCfzwciZHJi1OblnX9BED9k=
|
||||
golang.org/x/mobile v0.0.0-20260816165457-f98cc9b3c733/go.mod h1:jMwjxoDSx9jqhNaZqPnr6nnKzb7cs+Dy1Czk7wdX+R8=
|
||||
golang.org/x/mobile v0.0.0-20251113184115-a159579294ab h1:Iqyc+2zr7aGyLuEadIm0KRJP0Wwt+fhlXLa51Fxf1+Q=
|
||||
golang.org/x/mobile v0.0.0-20251113184115-a159579294ab/go.mod h1:Eq3Nh/5pFSWug2ohiudJ1iyU59SO78QFuh4qTTN++I0=
|
||||
golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
|
||||
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
|
||||
@@ -744,8 +744,8 @@ golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.39.0 h1:UF5zwQdCRRUpHfyPwr7d4UrGiVeldIsogtzWVnczL74=
|
||||
golang.org/x/mod v0.39.0/go.mod h1:bvIbwjQ0HUFFf5AKukeeYQG4ZBUG9yxQbR9aEweIwYY=
|
||||
golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ=
|
||||
golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0=
|
||||
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
|
||||
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
|
||||
golang.org/x/net v0.0.0-20190603091049-60506f45cf65/go.mod h1:HSz+uSET+XFnRR8LxR5pz3Of3rY3CfYBVs4xY44aLks=
|
||||
@@ -764,8 +764,8 @@ golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
|
||||
golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
|
||||
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To=
|
||||
golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU=
|
||||
golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o=
|
||||
golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec=
|
||||
golang.org/x/oauth2 v0.8.0/go.mod h1:yr7u4HXZRm1R1kBWqr/xKNqewf0plRYoB7sla+BCIXE=
|
||||
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
||||
@@ -843,8 +843,8 @@ golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
|
||||
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
|
||||
golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
|
||||
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
|
||||
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
|
||||
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
|
||||
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
@@ -858,8 +858,8 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
|
||||
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
|
||||
golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI=
|
||||
golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo=
|
||||
golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q=
|
||||
golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
||||
@@ -221,6 +221,9 @@ render_override() {
|
||||
# Remove this file (and config.yaml.enterprise if present) to revert.
|
||||
|
||||
services:
|
||||
${DASHBOARD_SERVICE}:
|
||||
image: \${NETBIRD_DASHBOARD_IMAGE:-ghcr.io/netbirdio/dashboard-cloud:latest}
|
||||
|
||||
${COMBINED_SERVICE}:
|
||||
image: \${NETBIRD_SERVER_IMAGE:-ghcr.io/netbirdio/netbird-server-cloud:latest}
|
||||
environment:
|
||||
|
||||
@@ -113,8 +113,61 @@ type Provider struct {
|
||||
// upstream provider + credentials on Portkey's hosted side).
|
||||
ExtraHeaders []ExtraHeader
|
||||
Models []Model
|
||||
// Discovery, when non-nil, describes how to ask this vendor which
|
||||
// models the operator's own credential can actually reach, so the
|
||||
// provider form can offer a live list instead of only the hand-curated
|
||||
// Models above. Nil for entries with no listing endpoint (gateways
|
||||
// vary too much) — those keep free-text entry.
|
||||
Discovery *Discovery
|
||||
}
|
||||
|
||||
// ListingShape names the response envelope a vendor returns its model
|
||||
// listing in. Every vendor invented its own, and none of them can be
|
||||
// guessed from the request, so the catalog states it.
|
||||
type ListingShape string
|
||||
|
||||
const (
|
||||
// ShapeOpenAIData is {"data":[{"id":…}]} — OpenAI, and Anthropic, which
|
||||
// adopted the same envelope.
|
||||
ShapeOpenAIData ListingShape = "openai_data"
|
||||
// ShapeBedrockInferenceProfiles is
|
||||
// {"inferenceProfileSummaries":[{"inferenceProfileId":…}]}. The ids carry
|
||||
// the region prefix that makes them invocable, which is exactly what an
|
||||
// operator cannot reconstruct by hand.
|
||||
ShapeBedrockInferenceProfiles ListingShape = "bedrock_inference_profiles"
|
||||
// ShapeVertexPublisherModels is {"publisherModels":[{"name":…}]}, where
|
||||
// name is a resource path and the invocable id is its last segment joined
|
||||
// to a separate versionId field.
|
||||
ShapeVertexPublisherModels ListingShape = "vertex_publisher_models"
|
||||
)
|
||||
|
||||
// Discovery describes one vendor's model-listing endpoint.
|
||||
//
|
||||
// Host is deliberately separate from the provider record's upstream URL:
|
||||
// Bedrock serves listings from the control plane (bedrock.<region>) while
|
||||
// inference must go to the runtime host (bedrock-runtime.<region>), so the
|
||||
// two cannot be the same value. Empty Host means "use the record's own
|
||||
// upstream", which is right for every vendor that serves both from one host.
|
||||
//
|
||||
// The regionPlaceholder in Host is substituted from the provider record's
|
||||
// region. Deriving the discovery host from the catalog rather than accepting
|
||||
// one from the caller is also what keeps this from being an open proxy: the
|
||||
// only hosts management will dial are the ones written here.
|
||||
type Discovery struct {
|
||||
Host string
|
||||
Path string
|
||||
Query string
|
||||
Shape ListingShape
|
||||
// Headers are static headers the vendor requires beyond the credential
|
||||
// (Anthropic versions its API through one and rejects a request without
|
||||
// it). The auth header itself comes from AuthHeaderName/Template.
|
||||
Headers map[string]string
|
||||
}
|
||||
|
||||
// RegionPlaceholder is replaced in Discovery.Host by the provider record's
|
||||
// configured region.
|
||||
const RegionPlaceholder = "<region>"
|
||||
|
||||
// ExtraHeader names a single optional per-provider routing/config
|
||||
// header. Catalog declares N of these per provider type; the operator
|
||||
// fills any subset on the provider record (see Provider.ExtraValues).
|
||||
@@ -245,8 +298,12 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#10A37F",
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1/models",
|
||||
Shape: ShapeOpenAIData,
|
||||
},
|
||||
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
|
||||
@@ -284,8 +341,18 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#D97757",
|
||||
ParserID: "anthropic",
|
||||
PricingSurfaces: []string{"anthropic"},
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1/models",
|
||||
// The default page is short and a picker wants the whole
|
||||
// catalogue in one call.
|
||||
Query: "limit=1000",
|
||||
Shape: ShapeOpenAIData,
|
||||
// Anthropic versions its API through a header and refuses a
|
||||
// request that omits it, listing included.
|
||||
Headers: map[string]string{"anthropic-version": "2023-06-01"},
|
||||
},
|
||||
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
|
||||
@@ -296,6 +363,8 @@ var providers = []Provider{
|
||||
// account to be on >= 30-day data retention or all requests
|
||||
// 400.
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Label: "Claude Opus 5", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-5", Label: "Claude Sonnet 5", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{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},
|
||||
@@ -343,6 +412,22 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#FF9900",
|
||||
// Listings come from the CONTROL PLANE, not the runtime host in
|
||||
// DefaultHost above: ListInferenceProfiles is not an operation
|
||||
// bedrock-runtime implements, and answers <UnknownOperationException/>
|
||||
// there. Inference has to go to the runtime host, so the two hosts
|
||||
// genuinely differ and Discovery.Host carries the difference.
|
||||
//
|
||||
// Inference profiles rather than foundation models because the profile
|
||||
// id is the invocable one: it carries the region prefix (eu., us.,
|
||||
// global.) that AWS requires and that cannot be derived from the
|
||||
// configured region — an eu-central-1 account legitimately holds
|
||||
// global.* profiles.
|
||||
Discovery: &Discovery{
|
||||
Host: "bedrock." + RegionPlaceholder + ".amazonaws.com",
|
||||
Path: "/inference-profiles",
|
||||
Shape: ShapeBedrockInferenceProfiles,
|
||||
},
|
||||
// ParserID stays empty (path-style dispatch via IsBedrockPathStyle);
|
||||
// the request parser meters these under the "bedrock" surface.
|
||||
PricingSurfaces: []string{"bedrock"},
|
||||
@@ -355,6 +440,8 @@ 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-5", Label: "Claude Opus 5 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "anthropic.claude-sonnet-5", Label: "Claude Sonnet 5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{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},
|
||||
@@ -391,6 +478,15 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#4285F4",
|
||||
// Only the v1beta1 publisher listing answers: the v1 form and the
|
||||
// project-scoped form under BOTH versions return 404. That means the
|
||||
// list is publisher-global — it cannot say which models this project
|
||||
// has enabled — so it is offered as a suggestion beside the catalog
|
||||
// rather than replacing it. See the discovery e2e for the probes.
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1beta1/publishers/anthropic/models",
|
||||
Shape: ShapeVertexPublisherModels,
|
||||
},
|
||||
// ParserID stays empty (path-style dispatch via IsVertexPathStyle);
|
||||
// Anthropic-on-Vertex requests are metered under the "anthropic"
|
||||
// surface with the bare, unversioned model id.
|
||||
@@ -406,6 +502,8 @@ var providers = []Provider{
|
||||
// exists — the router denies unmeterable publishers rather than forward
|
||||
// them uncounted.
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Label: "Claude Opus 5 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000},
|
||||
{ID: "claude-sonnet-5", Label: "Claude Sonnet 5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000},
|
||||
{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},
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestClaudeLineupSelectable pins the models Claude Code resolves to by
|
||||
// default. A model absent from the lineup can't be ticked on a provider
|
||||
// record, so llm_router denies it as not-routable and the operator has no
|
||||
// way to authorise the client's own default.
|
||||
func TestClaudeLineupSelectable(t *testing.T) {
|
||||
for providerID, wanted := range map[string][]string{
|
||||
"anthropic_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"},
|
||||
"bedrock_api": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5", "anthropic.claude-haiku-4-5"},
|
||||
"vertex_ai_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"},
|
||||
} {
|
||||
provider, ok := Lookup(providerID)
|
||||
require.True(t, ok, "catalog must define %s", providerID)
|
||||
|
||||
selectable := make(map[string]Model, len(provider.Models))
|
||||
for _, m := range provider.Models {
|
||||
selectable[m.ID] = m
|
||||
}
|
||||
for _, id := range wanted {
|
||||
model, found := selectable[id]
|
||||
require.True(t, found, "%s must offer %s", providerID, id)
|
||||
assert.NotEmpty(t, model.Label, "%s/%s needs a label for the picker", providerID, id)
|
||||
assert.Positive(t, model.InputPer1k, "%s/%s needs an input rate", providerID, id)
|
||||
assert.Positive(t, model.OutputPer1k, "%s/%s needs an output rate", providerID, id)
|
||||
assert.Positive(t, model.ContextWindow, "%s/%s needs a context window", providerID, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/modeldiscovery"
|
||||
nbcontext "github.com/netbirdio/netbird/management/server/context"
|
||||
"github.com/netbirdio/netbird/shared/auth"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// discoveryManagerStub records what the handler asked for and returns a canned
|
||||
// answer. The Manager interface is embedded rather than implemented: only the
|
||||
// one method is reachable from this handler, and a call to any other should
|
||||
// fail loudly rather than silently return a zero value.
|
||||
type discoveryManagerStub struct {
|
||||
agentnetwork.Manager
|
||||
|
||||
gotReq modeldiscovery.Request
|
||||
gotRecordID string
|
||||
models []modeldiscovery.Model
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *discoveryManagerStub) DiscoverProviderModels(
|
||||
_ context.Context, _, _ string, req modeldiscovery.Request, recordID string,
|
||||
) ([]modeldiscovery.Model, error) {
|
||||
s.gotReq = req
|
||||
s.gotRecordID = recordID
|
||||
return s.models, s.err
|
||||
}
|
||||
|
||||
// postDiscovery drives the handler with an authenticated request.
|
||||
func postDiscovery(t *testing.T, stub *discoveryManagerStub, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
h := &handler{manager: stub}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/agent-network/catalog/providers/models", strings.NewReader(body))
|
||||
req = req.WithContext(nbcontext.SetUserAuthInContext(req.Context(), auth.UserAuth{
|
||||
AccountId: "acc-1",
|
||||
UserId: "user-1",
|
||||
}))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.discoverProviderModels(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestDiscoverModelsReturnsTheVendorList(t *testing.T) {
|
||||
stub := &discoveryManagerStub{models: []modeldiscovery.Model{
|
||||
{ID: "eu.anthropic.claude-haiku-4-5-20251001-v1:0", Label: "EU Claude Haiku 4.5", PricingKnown: true},
|
||||
{ID: "global.cohere.embed-v4:0", Label: "Global Cohere Embed v4"},
|
||||
}}
|
||||
|
||||
rec := postDiscovery(t, stub, `{
|
||||
"catalog_provider_id":"bedrock_api",
|
||||
"upstream_url":"https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
"api_key":"aws-bearer"
|
||||
}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String())
|
||||
|
||||
var out api.AgentNetworkModelDiscoveryResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &out))
|
||||
require.Len(t, out.Models, 2)
|
||||
|
||||
assert.Equal(t, "eu.anthropic.claude-haiku-4-5-20251001-v1:0", out.Models[0].Id)
|
||||
assert.True(t, out.Models[0].PricingKnown)
|
||||
// An unpriced model must say so rather than arriving indistinguishable
|
||||
// from a priced one: registering it silently would meter at zero.
|
||||
assert.False(t, out.Models[1].PricingKnown)
|
||||
|
||||
assert.Equal(t, "bedrock_api", stub.gotReq.CatalogID)
|
||||
assert.Equal(t, "aws-bearer", stub.gotReq.APIKey)
|
||||
assert.Empty(t, stub.gotRecordID)
|
||||
}
|
||||
|
||||
func TestDiscoverModelsUsesAStoredRecordWithoutAKey(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
|
||||
rec := postDiscovery(t, stub, `{"catalog_provider_id":"openai_api","provider_id":"prov-42"}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String())
|
||||
|
||||
// The dashboard refreshes a saved provider's list without ever holding
|
||||
// the credential, so the record id has to reach the manager.
|
||||
assert.Equal(t, "prov-42", stub.gotRecordID)
|
||||
assert.Empty(t, stub.gotReq.APIKey)
|
||||
}
|
||||
|
||||
// TestDiscoverModelsRefusesMixedCredentials covers the case where a caller
|
||||
// names a saved provider AND supplies a key. Accepting it would run an
|
||||
// arbitrary credential under the identity of a record the caller may only be
|
||||
// permitted to read.
|
||||
func TestDiscoverModelsRefusesMixedCredentials(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
|
||||
rec := postDiscovery(t, stub, `{
|
||||
"catalog_provider_id":"openai_api",
|
||||
"provider_id":"prov-42",
|
||||
"api_key":"sk-attacker"
|
||||
}`)
|
||||
assert.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
assert.Empty(t, stub.gotRecordID, "the request must be refused before it reaches the manager")
|
||||
}
|
||||
|
||||
// TestDiscoverModelsReportsNoDiscoveryDistinctly matters because the caller
|
||||
// falls back to the catalog's own model list on this outcome. Collapsing it
|
||||
// into a generic 500 would turn "this provider has no listing endpoint" into
|
||||
// "something went wrong", and the form would show an error instead of a list.
|
||||
func TestDiscoverModelsReportsNoDiscoveryDistinctly(t *testing.T) {
|
||||
stub := &discoveryManagerStub{err: modeldiscovery.ErrNoDiscovery}
|
||||
|
||||
rec := postDiscovery(t, stub, `{"catalog_provider_id":"litellm_proxy","upstream_url":"https://gw.example.com","api_key":"sk"}`)
|
||||
assert.Equal(t, http.StatusUnprocessableEntity, rec.Code)
|
||||
}
|
||||
|
||||
func TestDiscoverModelsRejectsMalformedRequests(t *testing.T) {
|
||||
for name, body := range map[string]string{
|
||||
"not json": `{`,
|
||||
"no catalog provider": `{"api_key":"sk"}`,
|
||||
"blank catalog provide": `{"catalog_provider_id":" ","api_key":"sk"}`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
rec := postDiscovery(t, stub, body)
|
||||
assert.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
@@ -16,6 +17,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/modeldiscovery"
|
||||
"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"
|
||||
@@ -32,6 +34,7 @@ type handler struct {
|
||||
func RegisterEndpoints(manager agentnetwork.Manager, router *mux.Router) {
|
||||
h := &handler{manager: manager}
|
||||
router.HandleFunc("/agent-network/catalog/providers", h.getCatalogProviders).Methods("GET", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/catalog/providers/models", h.discoverProviderModels).Methods("POST", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers", h.getAllProviders).Methods("GET", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers", h.createProvider).Methods("POST", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET", "OPTIONS")
|
||||
@@ -61,6 +64,73 @@ func (h *handler) getCatalogProviders(w http.ResponseWriter, r *http.Request) {
|
||||
util.WriteJSONObject(r.Context(), w, out)
|
||||
}
|
||||
|
||||
// discoverProviderModels asks the vendor which models the operator's own
|
||||
// credential can reach, so the provider form can offer a live list rather than
|
||||
// only the static catalog.
|
||||
func (h *handler) discoverProviderModels(w http.ResponseWriter, r *http.Request) {
|
||||
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
|
||||
if err != nil {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
|
||||
var body api.AgentNetworkModelDiscoveryRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
util.WriteErrorResponse("invalid json", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(body.CatalogProviderId) == "" {
|
||||
util.WriteErrorResponse("catalog_provider_id is required", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
|
||||
recordID := strValue(body.ProviderId)
|
||||
req := modeldiscovery.Request{
|
||||
CatalogID: body.CatalogProviderId,
|
||||
UpstreamURL: strValue(body.UpstreamUrl),
|
||||
APIKey: strValue(body.ApiKey),
|
||||
}
|
||||
// One source of credential or the other, never a mix: taking a key from
|
||||
// the request while addressing a saved record would let a caller run an
|
||||
// arbitrary credential against a provider they can only read.
|
||||
if recordID != "" && req.APIKey != "" {
|
||||
util.WriteErrorResponse("provide either provider_id or api_key, not both", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
|
||||
models, err := h.manager.DiscoverProviderModels(r.Context(), userAuth.AccountId, userAuth.UserId, req, recordID)
|
||||
if err != nil {
|
||||
// A provider with no listing endpoint is a fact about the catalog
|
||||
// entry, not a failure: the caller falls back to the catalog's own
|
||||
// models, so it must be able to tell the two apart.
|
||||
if errors.Is(err, modeldiscovery.ErrNoDiscovery) {
|
||||
util.WriteErrorResponse(err.Error(), http.StatusUnprocessableEntity, w)
|
||||
return
|
||||
}
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
|
||||
out := api.AgentNetworkModelDiscoveryResponse{Models: make([]api.AgentNetworkDiscoveredModel, 0, len(models))}
|
||||
for _, m := range models {
|
||||
entry := api.AgentNetworkDiscoveredModel{Id: m.ID, PricingKnown: m.PricingKnown}
|
||||
if m.Label != "" {
|
||||
label := m.Label
|
||||
entry.Label = &label
|
||||
}
|
||||
out.Models = append(out.Models, entry)
|
||||
}
|
||||
util.WriteJSONObject(r.Context(), w, out)
|
||||
}
|
||||
|
||||
// strValue reads an optional string field, treating absent as empty.
|
||||
func strValue(v *string) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(*v)
|
||||
}
|
||||
|
||||
// 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.
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/labelgen"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/modeldiscovery"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
|
||||
@@ -50,6 +51,7 @@ type Manager interface {
|
||||
CreateProvider(ctx context.Context, userID string, provider *types.Provider) (*types.Provider, error)
|
||||
UpdateProvider(ctx context.Context, userID string, provider *types.Provider) (*types.Provider, error)
|
||||
DeleteProvider(ctx context.Context, accountID, userID, providerID string) error
|
||||
DiscoverProviderModels(ctx context.Context, accountID, userID string, req modeldiscovery.Request, recordID string) ([]modeldiscovery.Model, error)
|
||||
|
||||
GetAllPolicies(ctx context.Context, accountID, userID string) ([]*types.Policy, error)
|
||||
GetPolicy(ctx context.Context, accountID, userID, policyID string) (*types.Policy, error)
|
||||
@@ -123,6 +125,11 @@ type managerImpl struct {
|
||||
permissionsManager permissions.Manager
|
||||
proxyController proxy.Controller
|
||||
|
||||
// modelDiscovery queries vendors for the models a credential can reach.
|
||||
// A field rather than a package call so tests can drive it without
|
||||
// reaching the network.
|
||||
modelDiscovery *modeldiscovery.Client
|
||||
|
||||
// reconcileCache holds the last set of synthesised proxy mappings
|
||||
// per account, each paired with the proxy that served it, so a change
|
||||
// of serving proxy can be diffed without re-deriving it.
|
||||
@@ -151,6 +158,7 @@ func NewManager(
|
||||
accountManager: accountManager,
|
||||
permissionsManager: permissionsManager,
|
||||
proxyController: proxyController,
|
||||
modelDiscovery: &modeldiscovery.Client{},
|
||||
reconcileCache: make(map[string]map[string]syntheticMapping),
|
||||
labelRng: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
}
|
||||
@@ -170,6 +178,37 @@ func (m *managerImpl) GetProvider(ctx context.Context, accountID, userID, provid
|
||||
return m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID)
|
||||
}
|
||||
|
||||
// DiscoverProviderModels asks the vendor which models a credential can reach.
|
||||
//
|
||||
// recordID, when set, names an existing provider whose stored credential and
|
||||
// upstream are used instead of the ones in req — so the dashboard can refresh
|
||||
// the list without ever holding the key. Reading a stored credential is a read
|
||||
// of that provider, and is permission-checked as one.
|
||||
//
|
||||
// Gated on Create rather than Read: this spends the operator's credential
|
||||
// against a third party, which is not something a read-only role should be
|
||||
// able to make the server do.
|
||||
func (m *managerImpl) DiscoverProviderModels(ctx context.Context, accountID, userID string, req modeldiscovery.Request, recordID string) ([]modeldiscovery.Model, error) {
|
||||
if err := m.requirePermission(ctx, accountID, userID, modules.AgentNetworkProviders, operations.Create); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if recordID != "" {
|
||||
record, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, recordID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// The catalog id comes from the stored record too: letting the caller
|
||||
// name a different one would run a provider's credential against
|
||||
// whichever vendor endpoint they picked.
|
||||
req.CatalogID = record.ProviderID
|
||||
req.UpstreamURL = record.UpstreamURL
|
||||
req.APIKey = record.APIKey
|
||||
}
|
||||
|
||||
return m.modelDiscovery.Fetch(ctx, req)
|
||||
}
|
||||
|
||||
// CreateProvider persists a new provider for the account. Providers have no
|
||||
// settings side effects: the account's endpoint is bootstrapped separately and
|
||||
// explicitly via CreateSettings, and every provider in the account routes
|
||||
@@ -1017,6 +1056,10 @@ func (*mockManager) GetAllProviders(_ context.Context, _, _ string) ([]*types.Pr
|
||||
return []*types.Provider{}, nil
|
||||
}
|
||||
|
||||
func (*mockManager) DiscoverProviderModels(_ context.Context, _, _ string, _ modeldiscovery.Request, _ string) ([]modeldiscovery.Model, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (*mockManager) GetProvider(_ context.Context, _, _, _ string) (*types.Provider, error) {
|
||||
return &types.Provider{}, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
// Package modeldiscovery asks a vendor which models an operator's own
|
||||
// credential can reach, so the provider form can offer a live list instead of
|
||||
// only the catalog's hand-curated one.
|
||||
//
|
||||
// The catalog cannot know two things that matter. It goes stale — its entries
|
||||
// carry comments tracking which models a vendor retired on which date — and it
|
||||
// cannot see an account: which OpenAI models an org is entitled to, which
|
||||
// Bedrock inference profiles a given account and region hold, which Vertex
|
||||
// models a project has enabled. Those are exactly the facts an operator needs
|
||||
// when filling in a provider record, and only the vendor has them.
|
||||
//
|
||||
// The vendor is authoritative for the model ID. The catalog remains
|
||||
// authoritative for pricing, and a discovered model the catalog cannot price
|
||||
// is reported as such rather than silently registered at a rate of zero.
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/oauth2/google"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
const (
|
||||
// fetchTimeout bounds one vendor call end to end. A listing is a single
|
||||
// small GET; anything slower is a vendor problem and the operator is
|
||||
// waiting on a form.
|
||||
fetchTimeout = 8 * time.Second
|
||||
// maxListingBytes bounds the response we will buffer. The largest real
|
||||
// listing observed is Bedrock's foundation-model catalogue at ~70KB, so
|
||||
// this is a wide margin over anything legitimate.
|
||||
maxListingBytes = 2 << 20
|
||||
// gcpScope matches the scope llm_router mints Vertex tokens under, so a
|
||||
// credential that works for discovery works for inference too.
|
||||
gcpScope = "https://www.googleapis.com/auth/cloud-platform"
|
||||
// vertexKeyfilePrefix marks an api_key that is a base64 service-account
|
||||
// JSON key rather than a bearer token.
|
||||
vertexKeyfilePrefix = "keyfile::"
|
||||
)
|
||||
|
||||
// ErrNoDiscovery is returned for a catalog entry that declares no listing
|
||||
// endpoint. Gateways vary too much to have one, and the caller should fall
|
||||
// back to the catalog list plus free-text entry rather than treating this as
|
||||
// a failure.
|
||||
var ErrNoDiscovery = errors.New("provider has no model-discovery endpoint")
|
||||
|
||||
// Model is one discovered model.
|
||||
type Model struct {
|
||||
// ID is the identifier to register on the provider record, in the form the
|
||||
// vendor issues it. For Bedrock that is the region-prefixed inference
|
||||
// profile id, which is the only form AWS accepts at invoke time.
|
||||
ID string
|
||||
// Label is the vendor's display name where it supplies one.
|
||||
Label string
|
||||
// PricingKnown reports whether the shipped pricing table can price this
|
||||
// model. False means the operator must set rates, or the request would
|
||||
// meter at zero.
|
||||
PricingKnown bool
|
||||
}
|
||||
|
||||
// Request identifies which vendor to ask and with what credential.
|
||||
type Request struct {
|
||||
// CatalogID selects the catalog entry, which supplies the endpoint, the
|
||||
// auth header and the response shape. The caller never supplies those.
|
||||
CatalogID string
|
||||
// UpstreamURL is the provider record's configured upstream. It is used
|
||||
// only when the catalog entry declares no discovery host of its own.
|
||||
UpstreamURL string
|
||||
// Region substitutes the catalog host's <region> placeholder.
|
||||
Region string
|
||||
// APIKey is the operator's credential, exactly as stored on the record.
|
||||
APIKey string
|
||||
}
|
||||
|
||||
// Client fetches model listings. The zero value is usable; Resolver and
|
||||
// HTTPClient exist so tests can drive it against a local server.
|
||||
type Client struct {
|
||||
HTTPClient *http.Client
|
||||
// Resolver looks up the host for the SSRF check. Nil uses the default.
|
||||
Resolver *net.Resolver
|
||||
// AllowPrivateHosts disables the private-address guard. Only tests set it:
|
||||
// their server is on loopback, which is precisely what the guard blocks.
|
||||
AllowPrivateHosts bool
|
||||
}
|
||||
|
||||
// Fetch returns the models the credential can reach.
|
||||
func (c *Client) Fetch(ctx context.Context, req Request) ([]Model, error) {
|
||||
entry, ok := catalog.Lookup(req.CatalogID)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown catalog provider %q", req.CatalogID)
|
||||
}
|
||||
if entry.Discovery == nil {
|
||||
return nil, ErrNoDiscovery
|
||||
}
|
||||
|
||||
endpoint, err := c.discoveryURL(entry, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, fetchTimeout)
|
||||
defer cancel()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build discovery request: %w", err)
|
||||
}
|
||||
if err := applyAuth(httpReq, entry, req.APIKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for name, value := range entry.Discovery.Headers {
|
||||
httpReq.Header.Set(name, value)
|
||||
}
|
||||
httpReq.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient().Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reach %s: %w", entry.Name, err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxListingBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read %s listing: %w", entry.Name, err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// Surface the vendor's own status. An operator whose key lacks a scope
|
||||
// needs to see 403 rather than a generic failure.
|
||||
return nil, fmt.Errorf("%s returned %d for its model listing", entry.Name, resp.StatusCode)
|
||||
}
|
||||
|
||||
ids, err := parseListing(entry.Discovery.Shape, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return decorate(entry, ids), nil
|
||||
}
|
||||
|
||||
// discoveryURL builds the listing URL and refuses one that does not point at a
|
||||
// public host.
|
||||
//
|
||||
// The path, query and (for Bedrock) the host all come from the catalog rather
|
||||
// than from the caller, so the only operator-controlled part is the host of an
|
||||
// entry whose listing lives on its own upstream. That still has to be checked:
|
||||
// management holds credentials for every provider, and an upstream pointed at
|
||||
// an internal address would turn this endpoint into a probe of the management
|
||||
// server's own network.
|
||||
func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, error) {
|
||||
host := entry.Discovery.Host
|
||||
if host == "" {
|
||||
parsed, err := url.Parse(strings.TrimSpace(req.UpstreamURL))
|
||||
if err != nil || parsed.Host == "" {
|
||||
return "", fmt.Errorf("provider upstream %q is not a usable URL", req.UpstreamURL)
|
||||
}
|
||||
host = parsed.Host
|
||||
}
|
||||
if strings.Contains(host, catalog.RegionPlaceholder) {
|
||||
region := strings.TrimSpace(req.Region)
|
||||
if region == "" {
|
||||
// A provider record carries no region field: the region lives
|
||||
// inside the upstream host the operator already configured, so
|
||||
// read it back out rather than asking them for it twice.
|
||||
region = regionFromUpstream(entry, req.UpstreamURL)
|
||||
}
|
||||
if region == "" {
|
||||
return "", fmt.Errorf("%s discovery needs a region, and none could be read from the provider upstream", entry.Name)
|
||||
}
|
||||
host = strings.ReplaceAll(host, catalog.RegionPlaceholder, region)
|
||||
}
|
||||
|
||||
target := &url.URL{Scheme: "https", Host: host, Path: entry.Discovery.Path, RawQuery: entry.Discovery.Query}
|
||||
if err := c.checkPublicHost(target.Hostname()); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return target.String(), nil
|
||||
}
|
||||
|
||||
// regionFromUpstream recovers the region an operator embedded in the provider
|
||||
// upstream, by matching it against the catalog's own host template. Bedrock's
|
||||
// template is "bedrock-runtime.<region>.amazonaws.com" and Vertex's is
|
||||
// "<region>-aiplatform.googleapis.com", so the region is whatever sits between
|
||||
// the fixed halves. Returns empty when the upstream does not match the
|
||||
// template, which is the case for a custom or proxied endpoint.
|
||||
func regionFromUpstream(entry catalog.Provider, upstreamURL string) string {
|
||||
prefix, suffix, found := strings.Cut(entry.DefaultHost, catalog.RegionPlaceholder)
|
||||
if !found {
|
||||
return ""
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimSpace(upstreamURL))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
host := parsed.Hostname()
|
||||
if host == "" {
|
||||
// A bare host with no scheme parses as a path, not a host.
|
||||
host = strings.TrimSpace(upstreamURL)
|
||||
}
|
||||
if !strings.HasPrefix(host, prefix) || !strings.HasSuffix(host, suffix) {
|
||||
return ""
|
||||
}
|
||||
region := host[len(prefix) : len(host)-len(suffix)]
|
||||
if region == "" || strings.Contains(region, ".") {
|
||||
return ""
|
||||
}
|
||||
return region
|
||||
}
|
||||
|
||||
// checkPublicHost refuses hosts that resolve to an address the management
|
||||
// server should never be asked to reach on an operator's behalf.
|
||||
func (c *Client) checkPublicHost(host string) error {
|
||||
if c.AllowPrivateHosts {
|
||||
return nil
|
||||
}
|
||||
if host == "" {
|
||||
return errors.New("discovery host is empty")
|
||||
}
|
||||
resolver := c.Resolver
|
||||
if resolver == nil {
|
||||
resolver = net.DefaultResolver
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), fetchTimeout)
|
||||
defer cancel()
|
||||
|
||||
addrs, err := resolver.LookupNetIP(ctx, "ip", host)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolve discovery host %q: %w", host, err)
|
||||
}
|
||||
// Every address must be public: a name that resolves to one public and one
|
||||
// loopback address is still a way to reach loopback.
|
||||
for _, addr := range addrs {
|
||||
if !isPublic(addr) {
|
||||
return fmt.Errorf("discovery host %q resolves to a non-public address", host)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// isPublic reports whether an address is one we are willing to dial.
|
||||
func isPublic(addr netip.Addr) bool {
|
||||
addr = addr.Unmap()
|
||||
switch {
|
||||
case !addr.IsValid(),
|
||||
addr.IsLoopback(),
|
||||
addr.IsPrivate(),
|
||||
addr.IsLinkLocalUnicast(),
|
||||
addr.IsLinkLocalMulticast(),
|
||||
addr.IsInterfaceLocalMulticast(),
|
||||
addr.IsMulticast(),
|
||||
addr.IsUnspecified():
|
||||
return false
|
||||
}
|
||||
// 100.64.0.0/10 (carrier NAT) is where NetBird's own overlay addresses
|
||||
// live, so it is emphatically not somewhere to send a provider credential.
|
||||
if addr.Is4() {
|
||||
b := addr.As4()
|
||||
if b[0] == 100 && b[1] >= 64 && b[1] <= 127 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// applyAuth sets the credential header the catalog entry declares. A Vertex
|
||||
// service-account key is exchanged for an OAuth token first, the same way the
|
||||
// proxy does at request time.
|
||||
func applyAuth(req *http.Request, entry catalog.Provider, apiKey string) error {
|
||||
key := strings.TrimSpace(apiKey)
|
||||
if key == "" {
|
||||
return fmt.Errorf("%s discovery needs an API key", entry.Name)
|
||||
}
|
||||
if rest, ok := strings.CutPrefix(key, vertexKeyfilePrefix); ok {
|
||||
token, err := mintGCPToken(req.Context(), rest)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key = token
|
||||
}
|
||||
name := entry.AuthHeaderName
|
||||
if name == "" {
|
||||
name = "Authorization"
|
||||
}
|
||||
template := entry.AuthHeaderTemplate
|
||||
if template == "" {
|
||||
template = "${API_KEY}"
|
||||
}
|
||||
req.Header.Set(name, strings.ReplaceAll(template, "${API_KEY}", key))
|
||||
return nil
|
||||
}
|
||||
|
||||
// mintGCPToken exchanges a base64 service-account key for an access token.
|
||||
func mintGCPToken(ctx context.Context, saKeyB64 string) (string, error) {
|
||||
jsonKey, err := base64.StdEncoding.DecodeString(strings.TrimSpace(saKeyB64))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("decode service-account key: %w", err)
|
||||
}
|
||||
conf, err := google.JWTConfigFromJSON(jsonKey, gcpScope)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parse service-account key: %w", err)
|
||||
}
|
||||
tok, err := conf.TokenSource(ctx).Token()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("mint gcp token: %w", err)
|
||||
}
|
||||
return tok.AccessToken, nil
|
||||
}
|
||||
|
||||
// decorate turns raw vendor ids into the models the caller renders, marking
|
||||
// each with whether the shipped pricing table can price it.
|
||||
func decorate(entry catalog.Provider, ids []listedModel) []Model {
|
||||
priced := make(map[string]struct{}, len(entry.Models))
|
||||
for _, m := range entry.Models {
|
||||
priced[m.ID] = struct{}{}
|
||||
}
|
||||
|
||||
out := make([]Model, 0, len(ids))
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
for _, listed := range ids {
|
||||
if listed.id == "" {
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[listed.id]; dup {
|
||||
continue
|
||||
}
|
||||
seen[listed.id] = struct{}{}
|
||||
|
||||
// The catalog keys pricing by the normalised id while the vendor
|
||||
// issues the wire form, so normalise before asking whether we can
|
||||
// price it — otherwise every Bedrock profile would report unpriced.
|
||||
_, known := priced[normalizeForPricing(entry.ID, listed.id)]
|
||||
out = append(out, Model{ID: listed.id, Label: listed.label, PricingKnown: known})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *Client) httpClient() *http.Client {
|
||||
if c.HTTPClient != nil {
|
||||
return c.HTTPClient
|
||||
}
|
||||
return &http.Client{
|
||||
Timeout: fetchTimeout,
|
||||
// A redirect is a way to move the request to a host the guard above
|
||||
// never checked, so none are followed.
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
// stubTransport answers every request with one canned response and records the
|
||||
// request it was given, so a test can assert on the URL and headers the client
|
||||
// built without a network round trip.
|
||||
type stubTransport struct {
|
||||
status int
|
||||
body string
|
||||
got *http.Request
|
||||
}
|
||||
|
||||
func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
s.got = req
|
||||
status := s.status
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Body: io.NopCloser(strings.NewReader(s.body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Request: req,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// newStubClient returns a client that never leaves the process. The host guard
|
||||
// is disabled because it would otherwise resolve the vendor's real name, which
|
||||
// would make these tests depend on DNS.
|
||||
func newStubClient(status int, body string) (*Client, *stubTransport) {
|
||||
tr := &stubTransport{status: status, body: body}
|
||||
return &Client{
|
||||
HTTPClient: &http.Client{Transport: tr},
|
||||
AllowPrivateHosts: true,
|
||||
}, tr
|
||||
}
|
||||
|
||||
// The payloads below are trimmed from what the vendors actually returned in
|
||||
// the discovery e2e, rather than invented, so a parser that only works against
|
||||
// an idealised shape fails here.
|
||||
|
||||
const openAIListing = `{"object":"list","data":[
|
||||
{"id":"gpt-4o-mini","object":"model","created":1721172741,"owned_by":"system"},
|
||||
{"id":"gpt-4o","object":"model","created":1715367049,"owned_by":"system"}
|
||||
]}`
|
||||
|
||||
const anthropicListing = `{"data":[
|
||||
{"type":"model","id":"claude-haiku-4-5-20251001","display_name":"Claude Haiku 4.5"},
|
||||
{"type":"model","id":"claude-sonnet-4-6","display_name":"Claude Sonnet 4.6"}
|
||||
],"has_more":false}`
|
||||
|
||||
const bedrockListing = `{"inferenceProfileSummaries":[
|
||||
{"inferenceProfileId":"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"inferenceProfileName":"EU Anthropic Claude Haiku 4.5","status":"ACTIVE","type":"SYSTEM_DEFINED"},
|
||||
{"inferenceProfileId":"global.cohere.embed-v4:0",
|
||||
"inferenceProfileName":"Global Cohere Embed v4","status":"ACTIVE","type":"SYSTEM_DEFINED"},
|
||||
{"inferenceProfileId":"eu.meta.llama3-2-1b-instruct-v1:0",
|
||||
"inferenceProfileName":"EU Meta Llama 3.2 1B","status":"INACTIVE","type":"SYSTEM_DEFINED"}
|
||||
]}`
|
||||
|
||||
const vertexListing = `{"publisherModels":[
|
||||
{"name":"publishers/anthropic/models/claude-3-opus","versionId":"20240229","launchStage":"GA"},
|
||||
{"name":"publishers/anthropic/models/claude-sonnet-4-5","versionId":"20250929","launchStage":"GA"}
|
||||
]}`
|
||||
|
||||
func TestFetchOpenAIListing(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "https://api.openai.com/v1/models", tr.got.URL.String())
|
||||
assert.Equal(t, "Bearer sk-test", tr.got.Header.Get("Authorization"),
|
||||
"the credential must be injected through the catalog's auth template")
|
||||
assert.Equal(t, []string{"gpt-4o-mini", "gpt-4o"}, ids(models))
|
||||
for _, m := range models {
|
||||
assert.True(t, m.PricingKnown, "both models are in the shipped catalog: %s", m.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAnthropicSendsTheVersionHeader(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, anthropicListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "anthropic_api",
|
||||
UpstreamURL: "https://api.anthropic.com",
|
||||
APIKey: "sk-ant-test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Anthropic rejects a request without the version header, so a listing
|
||||
// that reached us at all proves it was sent — but assert it, because the
|
||||
// failure mode otherwise only shows up against the live API.
|
||||
assert.Equal(t, "2023-06-01", tr.got.Header.Get("anthropic-version"))
|
||||
assert.Equal(t, "sk-ant-test", tr.got.Header.Get("x-api-key"),
|
||||
"Anthropic takes a bare key under its own header, not a Bearer token")
|
||||
assert.Equal(t, "limit=1000", tr.got.URL.RawQuery)
|
||||
|
||||
assert.Equal(t, []string{"claude-haiku-4-5-20251001", "claude-sonnet-4-6"}, ids(models))
|
||||
assert.Equal(t, "Claude Haiku 4.5", models[0].Label)
|
||||
}
|
||||
|
||||
func TestFetchBedrockUsesTheControlPlaneAndKeepsWireIDs(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, bedrockListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "bedrock_api",
|
||||
// The record's upstream is the RUNTIME host, which does not serve
|
||||
// listings. The catalog's own discovery host must win over it.
|
||||
UpstreamURL: "https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
Region: "eu-central-1",
|
||||
APIKey: "aws-bearer",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "https://bedrock.eu-central-1.amazonaws.com/inference-profiles",
|
||||
tr.got.URL.String(), "listings come from the control plane, not the runtime host")
|
||||
|
||||
// Region-prefixed ids verbatim: the prefix is what makes them invocable
|
||||
// and it cannot be reconstructed — global.* alongside eu.* is exactly the
|
||||
// case that defeats deriving it from the configured region.
|
||||
assert.Equal(t, []string{
|
||||
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"global.cohere.embed-v4:0",
|
||||
}, ids(models), "an INACTIVE profile must not be offered")
|
||||
|
||||
assert.True(t, models[0].PricingKnown,
|
||||
"the catalog prices anthropic.claude-haiku-4-5, which this id normalises to")
|
||||
assert.False(t, models[1].PricingKnown,
|
||||
"cohere embed is not in the shipped Bedrock catalog, so the operator must price it")
|
||||
}
|
||||
|
||||
func TestFetchVertexJoinsNameAndVersion(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, vertexListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "vertex_ai_api",
|
||||
UpstreamURL: "https://us-east5-aiplatform.googleapis.com",
|
||||
Region: "us-east5",
|
||||
APIKey: "ya29.test-token",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Vertex addresses a model as "<id>@<version>" on rawPredict, and splits
|
||||
// those across two fields in the listing.
|
||||
assert.Equal(t, []string{"claude-3-opus@20240229", "claude-sonnet-4-5@20250929"}, ids(models))
|
||||
assert.Equal(t, "claude-3-opus", models[0].Label)
|
||||
}
|
||||
|
||||
func TestFetchSurfacesTheVendorStatus(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusForbidden, `{"error":{"message":"no access"}}`)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "403",
|
||||
"an operator whose key lacks access needs to see which status the vendor returned")
|
||||
}
|
||||
|
||||
func TestFetchRejectsAProviderWithoutDiscovery(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "litellm_proxy",
|
||||
UpstreamURL: "https://gateway.example.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
assert.ErrorIs(t, err, ErrNoDiscovery,
|
||||
"a gateway with no listing endpoint must be distinguishable from a failure, so the caller can fall back")
|
||||
}
|
||||
|
||||
func TestFetchRequiresACredential(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "API key")
|
||||
}
|
||||
|
||||
func TestDiscoveryURLNeedsARegionWhenTheHostTemplatesOne(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, bedrockListing)
|
||||
|
||||
// An upstream that matches no catalog template — a proxy in front of
|
||||
// Bedrock, say — leaves nothing to read the region from. Refusing beats
|
||||
// guessing: an unsubstituted placeholder would dial a host that does not
|
||||
// exist, and a guessed region would dial the wrong account's endpoint.
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "bedrock_api",
|
||||
UpstreamURL: "https://bedrock.internal-proxy.example.com",
|
||||
APIKey: "aws-bearer",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "region")
|
||||
}
|
||||
|
||||
// TestHostGuardRejectsNonPublicAddresses is the SSRF guard. Management holds a
|
||||
// credential for every provider, so an upstream pointed at an internal address
|
||||
// would turn discovery into a way to probe — and hand a token to — the
|
||||
// management server's own network.
|
||||
func TestHostGuardRejectsNonPublicAddresses(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
addr string
|
||||
want bool
|
||||
}{
|
||||
{"loopback v4", "127.0.0.1", false},
|
||||
{"loopback v6", "::1", false},
|
||||
{"private 10/8", "10.0.0.5", false},
|
||||
{"private 172.16/12", "172.16.4.1", false},
|
||||
{"private 192.168/16", "192.168.1.1", false},
|
||||
{"link-local", "169.254.169.254", false}, // cloud metadata
|
||||
{"unspecified", "0.0.0.0", false},
|
||||
{"multicast", "224.0.0.1", false},
|
||||
{"netbird overlay 100.64/10", "100.90.1.2", false},
|
||||
{"v4-mapped loopback", "::ffff:127.0.0.1", false},
|
||||
{"public v4", "1.1.1.1", true},
|
||||
{"public v6", "2606:4700:4700::1111", true},
|
||||
{"just outside CGNAT", "100.128.0.1", true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
addr, err := netip.ParseAddr(tc.addr)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, isPublic(addr))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostGuardResolvesAndRejectsLocalhost(t *testing.T) {
|
||||
cl := &Client{}
|
||||
err := cl.checkPublicHost("localhost")
|
||||
require.Error(t, err, "a name resolving to loopback must be refused, not just a literal address")
|
||||
assert.Contains(t, err.Error(), "non-public")
|
||||
}
|
||||
|
||||
// TestEveryDiscoveryEntryHasAParser keeps the catalog and the parser table from
|
||||
// drifting: adding a Discovery block with a shape nothing parses would fail
|
||||
// only at runtime, in front of an operator.
|
||||
func TestEveryDiscoveryEntryHasAParser(t *testing.T) {
|
||||
for _, entry := range catalog.All() {
|
||||
if entry.Discovery == nil {
|
||||
continue
|
||||
}
|
||||
t.Run(entry.ID, func(t *testing.T) {
|
||||
assert.NotEmpty(t, entry.Discovery.Path, "a discovery entry needs a path")
|
||||
_, err := parseListing(entry.Discovery.Shape, []byte(`{}`))
|
||||
assert.NoError(t, err, "shape %q has no parser", entry.Discovery.Shape)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func ids(models []Model) []string {
|
||||
out := make([]string, 0, len(models))
|
||||
for _, m := range models {
|
||||
out = append(out, m.ID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestRegionIsReadBackFromTheUpstream covers the reason the API takes no
|
||||
// region field: a provider record has none, and the operator already encoded
|
||||
// it in the upstream host when they configured inference.
|
||||
func TestRegionIsReadBackFromTheUpstream(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, bedrockListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "bedrock_api",
|
||||
UpstreamURL: "https://bedrock-runtime.us-west-2.amazonaws.com",
|
||||
APIKey: "aws-bearer",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "bedrock.us-west-2.amazonaws.com", tr.got.URL.Host)
|
||||
}
|
||||
|
||||
func TestRegionFromUpstream(t *testing.T) {
|
||||
bedrock, ok := catalog.Lookup("bedrock_api")
|
||||
require.True(t, ok)
|
||||
vertex, ok := catalog.Lookup("vertex_ai_api")
|
||||
require.True(t, ok)
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
entry catalog.Provider
|
||||
upstream string
|
||||
want string
|
||||
}{
|
||||
{"bedrock runtime host", bedrock, "https://bedrock-runtime.eu-central-1.amazonaws.com", "eu-central-1"},
|
||||
{"bedrock without scheme", bedrock, "bedrock-runtime.ap-south-1.amazonaws.com", "ap-south-1"},
|
||||
{"vertex regional host", vertex, "https://us-east5-aiplatform.googleapis.com", "us-east5"},
|
||||
// A proxied or self-hosted upstream matches no template, and guessing
|
||||
// a region from it would build a URL pointing somewhere arbitrary.
|
||||
{"unrelated upstream", bedrock, "https://llm.internal.example.com", ""},
|
||||
{"vertex global host has no region segment", vertex, "https://aiplatform.googleapis.com", ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
assert.Equal(t, tc.want, regionFromUpstream(tc.entry, tc.upstream))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// listedModel is one entry lifted out of a vendor listing before the catalog
|
||||
// is consulted about it.
|
||||
type listedModel struct {
|
||||
id string
|
||||
label string
|
||||
}
|
||||
|
||||
// parseListing extracts model ids from a vendor listing. Each vendor invented
|
||||
// its own envelope, and the shape is declared by the catalog rather than
|
||||
// sniffed, so a vendor that changes shape fails loudly instead of silently
|
||||
// returning nothing.
|
||||
func parseListing(shape catalog.ListingShape, body []byte) ([]listedModel, error) {
|
||||
switch shape {
|
||||
case catalog.ShapeOpenAIData:
|
||||
return parseOpenAIData(body)
|
||||
case catalog.ShapeBedrockInferenceProfiles:
|
||||
return parseBedrockInferenceProfiles(body)
|
||||
case catalog.ShapeVertexPublisherModels:
|
||||
return parseVertexPublisherModels(body)
|
||||
default:
|
||||
return nil, fmt.Errorf("no parser for listing shape %q", shape)
|
||||
}
|
||||
}
|
||||
|
||||
// parseOpenAIData reads {"data":[{"id":…}]}, which OpenAI defined and
|
||||
// Anthropic adopted. Anthropic additionally supplies display_name.
|
||||
func parseOpenAIData(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
DisplayName string `json:"display_name"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode model listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Data))
|
||||
for _, entry := range doc.Data {
|
||||
out = append(out, listedModel{id: entry.ID, label: entry.DisplayName})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// parseBedrockInferenceProfiles reads
|
||||
// {"inferenceProfileSummaries":[{"inferenceProfileId":…}]}.
|
||||
//
|
||||
// The profile id is taken verbatim because its region prefix (eu., us.,
|
||||
// global.) is what makes it invocable, and it is not derivable from the
|
||||
// configured region — an account in one region legitimately holds global.*
|
||||
// profiles alongside its regional ones.
|
||||
//
|
||||
// Only ACTIVE profiles are offered: AWS reports others, and registering one
|
||||
// would produce a model that routes inside NetBird and fails at AWS.
|
||||
func parseBedrockInferenceProfiles(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Summaries []struct {
|
||||
ID string `json:"inferenceProfileId"`
|
||||
Name string `json:"inferenceProfileName"`
|
||||
Status string `json:"status"`
|
||||
} `json:"inferenceProfileSummaries"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode inference-profile listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Summaries))
|
||||
for _, entry := range doc.Summaries {
|
||||
if entry.Status != "" && !strings.EqualFold(entry.Status, "ACTIVE") {
|
||||
continue
|
||||
}
|
||||
out = append(out, listedModel{id: entry.ID, label: entry.Name})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// parseVertexPublisherModels reads {"publisherModels":[{"name":…}]}, where
|
||||
// name is a resource path ("publishers/anthropic/models/claude-3-opus") and
|
||||
// the version lives in a separate field.
|
||||
//
|
||||
// Vertex addresses a model as "<id>@<version>" on the rawPredict path, so the
|
||||
// two are joined here: reporting the bare name would hand the operator an id
|
||||
// that looks usable and is not.
|
||||
func parseVertexPublisherModels(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
VersionID string `json:"versionId"`
|
||||
} `json:"publisherModels"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode publisher-model listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Models))
|
||||
for _, entry := range doc.Models {
|
||||
id := entry.Name
|
||||
if slash := strings.LastIndex(id, "/"); slash >= 0 {
|
||||
id = id[slash+1:]
|
||||
}
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
label := id
|
||||
if entry.VersionID != "" {
|
||||
id += "@" + entry.VersionID
|
||||
}
|
||||
out = append(out, listedModel{id: id, label: label})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// normalizeForPricing maps a vendor's wire id onto the key the catalog prices
|
||||
// it under. It mirrors the synthesiser's normalizePricingModelID: the two must
|
||||
// agree, or a model reported here as priced would meter at the default rate
|
||||
// instead of the operator's.
|
||||
func normalizeForPricing(catalogProviderID, modelID string) string {
|
||||
switch {
|
||||
case catalog.IsBedrockPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
case catalog.IsVertexPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeVertexModel(modelID)
|
||||
default:
|
||||
return modelID
|
||||
}
|
||||
}
|
||||
@@ -47,17 +47,11 @@ var supplementalDefaults = map[string]map[string]Entry{
|
||||
"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 (
|
||||
|
||||
@@ -82,6 +82,11 @@ anthropic:
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-sonnet-5:
|
||||
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
|
||||
@@ -145,6 +150,11 @@ bedrock:
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-sonnet-5:
|
||||
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
|
||||
|
||||
@@ -116,11 +116,13 @@ func TestDefaultTable_PinnedRates(t *testing.T) {
|
||||
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.
|
||||
// Every id below must stay priced whichever source provides it: the
|
||||
// catalog lineup for the current Claude 5 family, supplementalDefaults
|
||||
// for the ids the dashboard deliberately doesn't offer.
|
||||
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"},
|
||||
"anthropic": {"claude-opus-5", "claude-sonnet-5", "kimi-k3[1m]", "kimi-k3"},
|
||||
"bedrock": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5"},
|
||||
} {
|
||||
for _, id := range ids {
|
||||
_, ok := table[surface][id]
|
||||
|
||||
@@ -211,7 +211,19 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
groupIndex := indexProviderGroups(enabledPolicies)
|
||||
|
||||
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex)
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in that map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
|
||||
// Discovery gets the finer view: per policy rather than flattened per
|
||||
// provider, so a listing can be bounded to what the calling groups may
|
||||
// actually use instead of the union across everyone who reaches the
|
||||
// provider.
|
||||
modelPolicies := buildModelPolicies(enabledPolicies, guardrailsByID)
|
||||
|
||||
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex, modelPolicies)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -228,11 +240,6 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
// 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
|
||||
@@ -351,6 +358,11 @@ type routerProviderRoute struct {
|
||||
AuthHeaderName string `json:"auth_header_name"`
|
||||
AuthHeaderValue string `json:"auth_header_value"`
|
||||
AllowedGroupIDs []string `json:"allowed_group_ids,omitempty"`
|
||||
// ModelPolicies is one entry per enabled policy authorising this provider,
|
||||
// carrying that policy's source groups and the models it permits. The
|
||||
// router bounds a model listing with it, so a provider two groups reach
|
||||
// under different allowlists offers each only its own.
|
||||
ModelPolicies []routerModelPolicy `json:"model_policies,omitempty"`
|
||||
// Vertex marks a Google Vertex AI provider, whose requests carry the
|
||||
// model in the URL path. The router selects it by path, bypassing the
|
||||
// model/vendor table.
|
||||
@@ -422,7 +434,7 @@ func indexProviderGroups(policies []*types.Policy) map[string][]string {
|
||||
// path-prefix tiebreak. Providers no enabled policy authorises
|
||||
// (orphans) are intentionally OMITTED so the router never observes a
|
||||
// route with an empty ACL.
|
||||
func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) {
|
||||
func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string, modelPolicies map[string][]routerModelPolicy) ([]byte, error) {
|
||||
cfg := routerConfig{Providers: make([]routerProviderRoute, 0, len(providers))}
|
||||
for _, p := range providers {
|
||||
groups, hasPolicy := groupIndex[p.ID]
|
||||
@@ -449,6 +461,7 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]
|
||||
AuthHeaderName: headerName,
|
||||
AuthHeaderValue: headerValue,
|
||||
AllowedGroupIDs: groups,
|
||||
ModelPolicies: modelPolicies[p.ID],
|
||||
Vertex: catalog.IsVertexPathStyle(p.ProviderID),
|
||||
Bedrock: catalog.IsBedrockPathStyle(p.ProviderID),
|
||||
GCPServiceAccountKeyB64: gcpSAKeyB64,
|
||||
@@ -1098,3 +1111,46 @@ func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// routerModelPolicy mirrors the router's ModelPolicyRule: one authorising
|
||||
// policy's source groups plus the models it permits. Models is nil for a
|
||||
// policy that sets no model allowlist, which lifts the restriction for the
|
||||
// groups it binds — so nil and empty must survive the round trip distinctly.
|
||||
type routerModelPolicy struct {
|
||||
GroupIDs []string `json:"group_ids"`
|
||||
Models []string `json:"models"`
|
||||
}
|
||||
|
||||
// buildModelPolicies indexes, per provider, one rule for each enabled policy
|
||||
// authorising it: the policy's source groups and the models its guardrail
|
||||
// permits.
|
||||
//
|
||||
// This is deliberately finer than buildProviderAllowlists, which flattens the
|
||||
// same inputs into one list per provider for the proxy's fail-closed guardrail.
|
||||
// A flattened list cannot answer "what may THIS caller see", so a provider two
|
||||
// teams reach under different allowlists would offer each team the other's
|
||||
// models — a picker full of entries the next request refuses. Keeping the
|
||||
// source groups alongside the models lets the router answer it at request time,
|
||||
// where it knows the caller's groups.
|
||||
func buildModelPolicies(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]routerModelPolicy {
|
||||
out := make(map[string][]routerModelPolicy)
|
||||
for _, p := range policies {
|
||||
if p == nil || len(p.SourceGroups) == 0 {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
rule := routerModelPolicy{GroupIDs: append([]string(nil), p.SourceGroups...)}
|
||||
if restricted {
|
||||
// Never nil when restricted: an allowlist permitting nothing must
|
||||
// stay distinguishable from no allowlist at all.
|
||||
rule.Models = append([]string{}, models...)
|
||||
}
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
out[providerID] = append(out[providerID], rule)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
@@ -93,3 +94,75 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
|
||||
})
|
||||
}
|
||||
|
||||
// policyForGroups builds an enabled policy binding the given source groups to
|
||||
// the given providers under an optional guardrail.
|
||||
func policyForGroups(id string, groups []string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
SourceGroups: groups,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildModelPolicies covers the finer index discovery needs. Where
|
||||
// buildProviderAllowlists flattens every authorising policy into one list per
|
||||
// provider — enough for a fail-closed backstop, but blind to who is asking —
|
||||
// this keeps each policy's source groups beside its models so the router can
|
||||
// bound a listing to the calling groups.
|
||||
func TestBuildModelPolicies(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("each policy keeps its own groups and models", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-sales"}, []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Equal(t, []routerModelPolicy{
|
||||
{GroupIDs: []string{"grp-eng"}, Models: []string{"gpt-4o"}},
|
||||
{GroupIDs: []string{"grp-sales"}, Models: []string{"claude-opus-4"}},
|
||||
}, got["prov-x"],
|
||||
"the two policies must stay separable so neither group is offered the other's models")
|
||||
})
|
||||
|
||||
t.Run("an unrestricted policy carries nil models", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-admin"}, nil, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Nil(t, got["prov-x"][1].Models,
|
||||
"no allowlist must reach the router as nil, which lifts the restriction for its groups")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist is not a restriction", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-disabled"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
assert.Nil(t, got["prov-x"][0].Models,
|
||||
"a guardrail with the allowlist check off restricts nothing")
|
||||
})
|
||||
|
||||
t.Run("an enabled allowlist with no models permits nothing", func(t *testing.T) {
|
||||
byIDEmpty := map[string]*types.Guardrail{
|
||||
"g-empty": {ID: "g-empty", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true}}},
|
||||
}
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-empty"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byIDEmpty)
|
||||
require.NotNil(t, got["prov-x"][0].Models,
|
||||
"an empty allowlist must not arrive as nil — that would read as unrestricted")
|
||||
assert.Empty(t, got["prov-x"][0].Models)
|
||||
})
|
||||
|
||||
t.Run("a policy binding no groups is skipped", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", nil, []string{"g-4o"}, "prov-x")}
|
||||
assert.Empty(t, buildModelPolicies(policies, byID),
|
||||
"a policy with no source groups authorises nobody, so it bounds nobody's listing")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -64,19 +64,6 @@ func (h *handler) createSetupKey(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// A one-off key can be used once, and GenerateSetupKey pins its usage limit
|
||||
// at 1 whatever the request says. Silently overriding a caller that asked
|
||||
// for a different number leaves them holding a key that does not do what
|
||||
// they configured, and no way to find out except by using it. Only values
|
||||
// above 1 are refused: usage_limit is a required field with no null, so 0
|
||||
// cannot be told apart from a caller that has nothing to say about it.
|
||||
if types.SetupKeyType(req.Type) == types.SetupKeyOneOff && req.UsageLimit > 1 {
|
||||
util.WriteError(r.Context(), status.Errorf(status.InvalidArgument,
|
||||
"usage_limit %d is not valid for a one-off setup key, which can be used once; use type reusable for a key that can be used more than once",
|
||||
req.UsageLimit), w)
|
||||
return
|
||||
}
|
||||
|
||||
expiresIn := time.Duration(req.ExpiresIn) * time.Second
|
||||
|
||||
if expiresIn < 0 {
|
||||
|
||||
@@ -134,40 +134,6 @@ func TestSetupKeysHandlers(t *testing.T) {
|
||||
expectedBody: true,
|
||||
expectedSetupKey: expectedNewKey,
|
||||
},
|
||||
{
|
||||
// A one-off key is used once. Asking for more used to be accepted
|
||||
// and then quietly reduced to 1.
|
||||
name: "Create One-Off Setup Key With Conflicting Usage Limit",
|
||||
requestType: http.MethodPost,
|
||||
requestPath: "/api/setup-keys",
|
||||
requestBody: bytes.NewBuffer(
|
||||
[]byte(fmt.Sprintf("{\"name\":\"%s\",\"type\":\"one-off\",\"expires_in\":86400,\"usage_limit\":5}", newSetupKeyName))),
|
||||
expectedStatus: http.StatusUnprocessableEntity,
|
||||
expectedBody: false,
|
||||
},
|
||||
{
|
||||
// 0 is what a caller sends when it has nothing to say about the
|
||||
// usage limit, since the field is required and has no null, so it
|
||||
// has to keep working.
|
||||
name: "Create One-Off Setup Key Without Usage Limit",
|
||||
requestType: http.MethodPost,
|
||||
requestPath: "/api/setup-keys",
|
||||
requestBody: bytes.NewBuffer(
|
||||
[]byte(fmt.Sprintf("{\"name\":\"%s\",\"type\":\"one-off\",\"expires_in\":86400,\"usage_limit\":0}", newSetupKeyName))),
|
||||
expectedStatus: http.StatusOK,
|
||||
expectedBody: false,
|
||||
},
|
||||
{
|
||||
// Only one-off keys are constrained; a reusable key means what it
|
||||
// says.
|
||||
name: "Create Reusable Setup Key With Usage Limit",
|
||||
requestType: http.MethodPost,
|
||||
requestPath: "/api/setup-keys",
|
||||
requestBody: bytes.NewBuffer(
|
||||
[]byte(fmt.Sprintf("{\"name\":\"%s\",\"type\":\"reusable\",\"expires_in\":86400,\"usage_limit\":5}", newSetupKeyName))),
|
||||
expectedStatus: http.StatusOK,
|
||||
expectedBody: false,
|
||||
},
|
||||
{
|
||||
name: "Update Setup Key",
|
||||
requestType: http.MethodPut,
|
||||
|
||||
@@ -136,10 +136,7 @@ func Test_SetupKeys_Create(t *testing.T) {
|
||||
},
|
||||
},
|
||||
{
|
||||
// The key used to be created anyway, with its usage limit quietly
|
||||
// reduced to 1, so the caller was told a key they had not asked for
|
||||
// was what they asked for.
|
||||
name: "Create Setup Key as one-off with more than one usage",
|
||||
name: "Create Setup Key as on-off with more than one usage",
|
||||
requestType: http.MethodPost,
|
||||
requestPath: "/api/setup-keys",
|
||||
requestBody: &api.CreateSetupKeyRequest{
|
||||
@@ -149,7 +146,23 @@ func Test_SetupKeys_Create(t *testing.T) {
|
||||
Type: "one-off",
|
||||
UsageLimit: 3,
|
||||
},
|
||||
expectedStatus: http.StatusUnprocessableEntity,
|
||||
expectedStatus: http.StatusOK,
|
||||
expectedResponse: &api.SetupKey{
|
||||
AutoGroups: []string{},
|
||||
Ephemeral: false,
|
||||
Expires: time.Time{},
|
||||
Id: "",
|
||||
Key: "",
|
||||
LastUsed: time.Time{},
|
||||
Name: testing_tools.NewKeyName,
|
||||
Revoked: false,
|
||||
State: "valid",
|
||||
Type: "one-off",
|
||||
UpdatedAt: time.Now(),
|
||||
UsageLimit: 1,
|
||||
UsedTimes: 0,
|
||||
Valid: true,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Create Setup Key with expiration in the past",
|
||||
|
||||
@@ -13,6 +13,14 @@ func NormalizeBedrockModel(modelID string) string {
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeAnthropicModel strips the trailing "-YYYYMMDD" release-date suffix
|
||||
// from an Anthropic model id so a dated id a client pins matches the undated
|
||||
// one the operator registered. Thin delegate to shared/llm for the same
|
||||
// contract reason as the two below.
|
||||
func NormalizeAnthropicModel(modelID string) string {
|
||||
return sharedllm.NormalizeAnthropicModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
|
||||
// so it matches the catalog/pricing key. Thin delegate to shared/llm, kept
|
||||
// beside NormalizeBedrockModel for the same contract reason.
|
||||
|
||||
@@ -10,6 +10,8 @@ package pricing
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// Entry is a single model's input and output pricing, expressed in USD per
|
||||
@@ -92,7 +94,10 @@ func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) {
|
||||
return &Table{entries: entries}, nil
|
||||
}
|
||||
|
||||
// Lookup returns the entry for the given provider surface and model.
|
||||
// Lookup returns the entry for the given provider surface and model. A
|
||||
// dated Anthropic id falls back to its undated form, so a client pinning
|
||||
// "claude-sonnet-4-5-20250929" bills at the registered "claude-sonnet-4-5"
|
||||
// rate instead of recording no cost at all.
|
||||
func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if t == nil {
|
||||
return Entry{}, false
|
||||
@@ -101,7 +106,14 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if !ok {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[model]
|
||||
if e, found := byModel[model]; found {
|
||||
return e, true
|
||||
}
|
||||
undated := sharedllm.NormalizeAnthropicModel(model)
|
||||
if undated == model {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[undated]
|
||||
return e, ok
|
||||
}
|
||||
|
||||
|
||||
@@ -175,3 +175,22 @@ func TestNewTable_NilAndEmpty(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map")
|
||||
}
|
||||
|
||||
// TestLookup_DatedAnthropicIDFallsBackToUndated covers a client pinning a
|
||||
// release date on a model priced under its undated id. Without the
|
||||
// fallback the request records no cost at all.
|
||||
func TestLookup_DatedAnthropicIDFallsBackToUndated(t *testing.T) {
|
||||
table, err := NewTable(map[string]map[string]EntryJSON{
|
||||
"anthropic": {
|
||||
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "table must build from a valid defaults map")
|
||||
|
||||
entry, ok := table.Lookup("anthropic", "claude-sonnet-4-5-20250929")
|
||||
require.True(t, ok, "a dated id must resolve to the undated entry")
|
||||
assert.InDelta(t, 0.003, entry.InputPer1K, 1e-9, "dated id must bill at the registered rate")
|
||||
|
||||
_, ok = table.Lookup("anthropic", "claude-sonnet-9-9-20250929")
|
||||
assert.False(t, ok, "an unknown family must stay unpriced")
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user