Compare commits

..

3 Commits

Author SHA1 Message Date
pascal
6edcee6a94 Merge branch 'main' into feature/header-auth-on-proxy 2026-08-20 17:25:44 +02:00
pascal
68ccc7e0b3 skip session cookie for header auth 2026-08-20 14:42:02 +02:00
pascal
80bfa33f71 validate header auth on proxy 2026-08-20 14:11:32 +02:00
92 changed files with 1616 additions and 7927 deletions

View File

@@ -12,13 +12,6 @@ on:
AWS issues it. Leave empty for the Sonnet 4.6 default.
required: false
default: ""
test_pattern:
description: >-
Package pattern to run. Defaults to the whole suite; narrow it to one
package (e.g. ./e2e/agentnetwork/...) when a run only needs that
package's answer and not the sixteen minutes the container suite costs.
required: false
default: "./e2e/..."
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
@@ -84,8 +77,4 @@ jobs:
GOOGLE_VERTEX_PROJECT: ${{ secrets.E2E_GOOGLE_VERTEX_PROJECT }}
GOOGLE_VERTEX_REGION: ${{ secrets.E2E_GOOGLE_VERTEX_REGION }}
GOOGLE_VERTEX_MODEL: ${{ secrets.E2E_GOOGLE_VERTEX_MODEL }}
# Read through an env var rather than interpolated into the run
# script: a dispatch input reaching a shell command directly is a
# script-injection seam, however trusted the dispatcher.
TEST_PATTERN: ${{ inputs.test_pattern || './e2e/...' }}
run: go test -tags e2e -timeout 40m -v "$TEST_PATTERN"
run: go test -tags e2e -timeout 40m -v ./e2e/...

View File

@@ -43,19 +43,8 @@ jobs:
run: /usr/local/lib/android/sdk/cmdline-tools/7.0/bin/sdkmanager --install "ndk;23.1.7779620"
- name: install gomobile
run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20251113184115-a159579294ab
# `gomobile init` re-installs gobind from golang.org/x/mobile@latest
# regardless of the pin above (cmd/gomobile/init.go: "Make sure gobind is
# up to date"), so this step resolves a version nobody chose, on every run.
#
# setup-go sets GOTOOLCHAIN=local, so that install fails outright once
# x/mobile@latest declares a newer Go than go.mod does — which it did on
# 2026-08-21, breaking both jobs on every branch at once. GOTOOLCHAIN=auto
# lets this one install fetch the toolchain it asks for. Scoped to the
# step: the repo's own Go version, and every build below, is unaffected.
- name: gomobile init
run: gomobile init
env:
GOTOOLCHAIN: auto
- name: build android netbird lib
run: PATH=$PATH:$(go env GOPATH) gomobile bind -o $GITHUB_WORKSPACE/netbird.aar -javapkg=io.netbird.gomobile -ldflags="-checklinkname=0 -X golang.zx2c4.com/wireguard/ipc.socketDirectory=/data/data/io.netbird.client/cache/wireguard -X github.com/netbirdio/netbird/version.version=buildtest" $GITHUB_WORKSPACE/client/android
env:
@@ -75,13 +64,8 @@ jobs:
go-version-file: "go.mod"
- name: install gomobile
run: go install golang.org/x/mobile/cmd/gomobile@v0.0.0-20251113184115-a159579294ab
# See the Android job: `gomobile init` re-installs gobind from
# golang.org/x/mobile@latest regardless of the pin above, and needs a
# toolchain it may pick newer than go.mod's.
- name: gomobile init
run: gomobile init
env:
GOTOOLCHAIN: auto
- name: build iOS netbird lib
run: PATH=$PATH:$(go env GOPATH) gomobile bind -target=ios -bundleid=io.netbird.framework -ldflags="-X github.com/netbirdio/netbird/version.version=buildtest" -o ./NetBirdSDK.xcframework ./client/ios/NetBirdSDK
env:

View File

@@ -1,78 +0,0 @@
name: No New Replace Directives
on:
pull_request:
paths:
- "go.mod"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || github.actor_id }}
cancel-in-progress: true
jobs:
check-replace-directives:
name: check-replace-directives
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
fetch-depth: 0
- name: Install Go
uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0
with:
go-version-file: go.mod
- name: Compare replace directives against the base branch
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
set -euo pipefail
# A replace directive only applies when this module is the main
# module. Anything importing netbird as a library, the embedded
# clients among them, resolves the replaced path upstream instead and
# fails to build against whatever the replacement provides. Requiring
# a fork under its own module path avoids that; a replace does not.
#
# go.mod is parsed rather than diffed so that reordering, comments and
# single-line versus block syntax do not register as changes.
#
# Versions are part of the key because a replace can be scoped to one
# version of a module. Keyed on paths alone, retargeting such a
# directive at a different version would read as unchanged.
list_replaces() {
go mod edit -json "$1" \
| jq -r '
def ref: .Path + (if (.Version // "") == "" then "" else " " + .Version end);
(.Replace // [])[] | "\(.Old | ref) => \(.New | ref)"
' \
| sort
}
git show "${BASE_SHA}:go.mod" > /tmp/base-go.mod
list_replaces /tmp/base-go.mod > /tmp/base-replaces
list_replaces go.mod > /tmp/head-replaces
added=$(comm -13 /tmp/base-replaces /tmp/head-replaces)
if [ -n "$added" ]; then
echo "::error::This PR adds a replace directive to go.mod:"
echo "$added" | sed 's/^/ /'
echo ""
echo "A replace directive applies only to the main module, so it does not"
echo "reach anything that imports netbird as a library. Require the module"
echo "under a path you control instead, as done for github.com/netbirdio/go-nat."
exit 1
fi
removed=$(comm -23 /tmp/base-replaces /tmp/head-replaces)
if [ -n "$removed" ]; then
echo "This PR removes replace directives:"
echo "$removed" | sed 's/^/ /'
fi
echo "No new replace directives."

View File

@@ -40,35 +40,6 @@ 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:

View File

@@ -91,13 +91,6 @@ type Options struct {
// when the embedded client must never act as a stepping stone into
// the host's local network (e.g. the proxy's overlay peer).
BlockLANAccess bool
// LazyConnectionEnabled is a tri-state local override for lazy connections,
// mirroring the NB_LAZY_CONN env var. Nil defers to the management feature
// flag; a set value overrides it in both directions. A short-lived client
// that reaches only a few known peers can set this to false, so its peers
// connect eagerly and the first request does not wait for the connection to
// be established.
LazyConnectionEnabled *bool
// WireguardPort is the port for the tunnel interface. Use 0 for a random port.
WireguardPort *int
// MTU is the MTU for the tunnel interface.
@@ -227,15 +220,6 @@ func New(opts Options) (*Client, error) {
config.PrivateKey = opts.PrivateKey
}
if opts.LazyConnectionEnabled != nil {
// Runtime-only override, read back through lazyconn.ParseState; a set value
// wins over the management feature flag in both directions.
config.LazyConnection = "off"
if *opts.LazyConnectionEnabled {
config.LazyConnection = "on"
}
}
if opts.Performance.PreallocatedBuffersPerPool != nil {
wgdevice.SetPreallocatedBuffersPerPool(*opts.Performance.PreallocatedBuffersPerPool)
}

View File

@@ -389,17 +389,6 @@ func (w *WorkerICE) injectPortForwardedCandidate(srflxCandidate ice.Candidate) {
return
}
// A forwarded candidate only makes sense for an IPv4 mapping, which
// translates a port on the gateway's address. An IPv6 pinhole translates
// nothing: it unblocks the address ICE already gathers as a host candidate,
// so there is no second address to advertise. Injecting one here would also
// paste an IPv6 address onto whichever server-reflexive candidate arrived
// first, which is usually IPv4.
if mapping.ExternalIP != nil && mapping.ExternalIP.To4() == nil {
w.log.Debugf("skipping port-forwarded candidate: %s mapping is IPv6-only", mapping.NATType)
return
}
w.muxAgent.Lock()
if w.portForwardAttempted {
w.muxAgent.Unlock()

View File

@@ -10,8 +10,10 @@ import (
"sync"
"time"
"github.com/netbirdio/go-nat"
"github.com/libp2p/go-nat"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/client/internal/portforward/pcp"
)
const (
@@ -166,11 +168,6 @@ func (m *Manager) setup(ctx context.Context) (nat.NAT, *Mapping, error) {
if err != nil {
return nil, nil, fmt.Errorf("create port mapping: %w", err)
}
// Only meaningful once a mapping has been attempted: that is what opens the
// pinhole and records its outcome.
logIPv6Pinhole(gateway)
return gateway, mapping, nil
}
@@ -268,9 +265,7 @@ func (m *Manager) checkHealthAndRecreate(ctx context.Context, gateway nat.NAT) b
return false
}
// Assert on the interface, not on a concrete type: a dual-stack gateway is
// a wrapper around the IPv4 NAT, so a type assertion misses it.
checker, ok := gateway.(nat.HealthChecker)
pcpNAT, ok := gateway.(*pcp.NAT)
if !ok {
return false
}
@@ -278,7 +273,7 @@ func (m *Manager) checkHealthAndRecreate(ctx context.Context, gateway nat.NAT) b
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
epoch, serverRestarted, err := checker.CheckServerHealth(ctx)
epoch, serverRestarted, err := pcpNAT.CheckServerHealth(ctx)
if err != nil {
log.Debugf("PCP health check failed: %v", err)
return false
@@ -345,18 +340,3 @@ func (m *Manager) startTearDown(ctx context.Context) {
func isPermanentLeaseRequired(err error) bool {
return err != nil && upnpErrPermanentLeaseOnly.MatchString(err.Error())
}
// logIPv6Pinhole reports the outcome of the IPv6 pinhole. Pinholes are best
// effort and never fail a mapping on their own, so this is the only way to see
// whether one was actually opened.
func logIPv6Pinhole(gateway nat.NAT) {
reporter, ok := gateway.(nat.IPv6PinholeReporter)
if !ok {
return
}
if err := reporter.IPv6PinholeError(); err != nil {
log.Warnf("IPv6 pinhole: %v", err)
return
}
log.Infof("IPv6 pinhole open")
}

View File

@@ -0,0 +1,408 @@
package pcp
import (
"context"
"crypto/rand"
"errors"
"fmt"
"net"
"net/netip"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
const (
defaultTimeout = 3 * time.Second
responseBufferSize = 128
// RFC 6887 Section 8.1.1 retry timing
initialRetryDelay = 3 * time.Second
maxRetryDelay = 1024 * time.Second
maxRetries = 4 // 3s + 6s + 12s + 24s = 45s total worst case
)
// Client is a PCP protocol client.
// All methods are safe for concurrent use.
type Client struct {
gateway netip.Addr
timeout time.Duration
mu sync.Mutex
// localIP caches the resolved local IP address.
localIP netip.Addr
// lastEpoch is the last observed server epoch value.
lastEpoch uint32
// epochTime tracks when lastEpoch was received for state loss detection.
epochTime time.Time
// externalIP caches the external IP from the last successful MAP response.
externalIP netip.Addr
// epochStateLost is set when epoch indicates server restart.
epochStateLost bool
}
// NewClient creates a new PCP client for the gateway at the given IP.
func NewClient(gateway net.IP) *Client {
addr, ok := netip.AddrFromSlice(gateway)
if !ok {
log.Debugf("invalid gateway IP: %v", gateway)
}
return &Client{
gateway: addr.Unmap(),
timeout: defaultTimeout,
}
}
// NewClientWithTimeout creates a new PCP client with a custom timeout.
func NewClientWithTimeout(gateway net.IP, timeout time.Duration) *Client {
addr, ok := netip.AddrFromSlice(gateway)
if !ok {
log.Debugf("invalid gateway IP: %v", gateway)
}
return &Client{
gateway: addr.Unmap(),
timeout: timeout,
}
}
// SetLocalIP sets the local IP address to use in PCP requests.
func (c *Client) SetLocalIP(ip net.IP) {
addr, ok := netip.AddrFromSlice(ip)
if !ok {
log.Debugf("invalid local IP: %v", ip)
}
c.mu.Lock()
c.localIP = addr.Unmap()
c.mu.Unlock()
}
// Gateway returns the gateway IP address.
func (c *Client) Gateway() net.IP {
return c.gateway.AsSlice()
}
// Announce sends a PCP ANNOUNCE request to discover PCP support.
// Returns the server's epoch time on success.
func (c *Client) Announce(ctx context.Context) (epoch uint32, err error) {
localIP, err := c.getLocalIP()
if err != nil {
return 0, fmt.Errorf("get local IP: %w", err)
}
req := buildAnnounceRequest(localIP)
resp, err := c.sendRequest(ctx, req)
if err != nil {
return 0, fmt.Errorf("send announce: %w", err)
}
parsed, err := parseResponse(resp)
if err != nil {
return 0, fmt.Errorf("parse announce response: %w", err)
}
if parsed.ResultCode != ResultSuccess {
return 0, fmt.Errorf("PCP ANNOUNCE failed: %s", ResultCodeString(parsed.ResultCode))
}
c.mu.Lock()
if c.updateEpochLocked(parsed.Epoch) {
log.Warnf("PCP server epoch indicates state loss - mappings may need refresh")
}
c.mu.Unlock()
return parsed.Epoch, nil
}
// AddPortMapping requests a port mapping from the PCP server.
func (c *Client) AddPortMapping(ctx context.Context, protocol string, internalPort int, lifetime time.Duration) (*MapResponse, error) {
return c.addPortMappingWithHint(ctx, protocol, internalPort, internalPort, netip.Addr{}, lifetime)
}
// AddPortMappingWithHint requests a port mapping with suggested external port and IP.
// Use lifetime <= 0 to delete a mapping.
func (c *Client) AddPortMappingWithHint(ctx context.Context, protocol string, internalPort, suggestedExtPort int, suggestedExtIP net.IP, lifetime time.Duration) (*MapResponse, error) {
var extIP netip.Addr
if suggestedExtIP != nil {
var ok bool
extIP, ok = netip.AddrFromSlice(suggestedExtIP)
if !ok {
log.Debugf("invalid suggested external IP: %v", suggestedExtIP)
}
extIP = extIP.Unmap()
}
return c.addPortMappingWithHint(ctx, protocol, internalPort, suggestedExtPort, extIP, lifetime)
}
func (c *Client) addPortMappingWithHint(ctx context.Context, protocol string, internalPort, suggestedExtPort int, suggestedExtIP netip.Addr, lifetime time.Duration) (*MapResponse, error) {
localIP, err := c.getLocalIP()
if err != nil {
return nil, fmt.Errorf("get local IP: %w", err)
}
proto, err := protocolNumber(protocol)
if err != nil {
return nil, fmt.Errorf("parse protocol: %w", err)
}
var nonce [12]byte
if _, err := rand.Read(nonce[:]); err != nil {
return nil, fmt.Errorf("generate nonce: %w", err)
}
// Convert lifetime to seconds. Lifetime 0 means delete, so only apply
// default for positive durations that round to 0 seconds.
var lifetimeSec uint32
if lifetime > 0 {
lifetimeSec = uint32(lifetime.Seconds())
if lifetimeSec == 0 {
lifetimeSec = DefaultLifetime
}
}
req := buildMapRequest(localIP, nonce, proto, uint16(internalPort), uint16(suggestedExtPort), suggestedExtIP, lifetimeSec)
resp, err := c.sendRequest(ctx, req)
if err != nil {
return nil, fmt.Errorf("send map request: %w", err)
}
mapResp, err := parseMapResponse(resp)
if err != nil {
return nil, fmt.Errorf("parse map response: %w", err)
}
if mapResp.Nonce != nonce {
return nil, fmt.Errorf("nonce mismatch in response")
}
if mapResp.Protocol != proto {
return nil, fmt.Errorf("protocol mismatch: requested %d, got %d", proto, mapResp.Protocol)
}
if mapResp.InternalPort != uint16(internalPort) {
return nil, fmt.Errorf("internal port mismatch: requested %d, got %d", internalPort, mapResp.InternalPort)
}
if mapResp.ResultCode != ResultSuccess {
return nil, &Error{
Code: mapResp.ResultCode,
Message: ResultCodeString(mapResp.ResultCode),
}
}
c.mu.Lock()
if c.updateEpochLocked(mapResp.Epoch) {
log.Warnf("PCP server epoch indicates state loss - mappings may need refresh")
}
c.cacheExternalIPLocked(mapResp.ExternalIP)
c.mu.Unlock()
return mapResp, nil
}
// DeletePortMapping removes a port mapping by requesting zero lifetime.
func (c *Client) DeletePortMapping(ctx context.Context, protocol string, internalPort int) error {
if _, err := c.addPortMappingWithHint(ctx, protocol, internalPort, 0, netip.Addr{}, 0); err != nil {
var pcpErr *Error
if errors.As(err, &pcpErr) && pcpErr.Code == ResultNotAuthorized {
return nil
}
return fmt.Errorf("delete mapping: %w", err)
}
return nil
}
// GetExternalAddress returns the external IP address.
// First checks for a cached value from previous MAP responses.
// If not cached, creates a short-lived mapping to discover the external IP.
func (c *Client) GetExternalAddress(ctx context.Context) (net.IP, error) {
c.mu.Lock()
if c.externalIP.IsValid() {
ip := c.externalIP.AsSlice()
c.mu.Unlock()
return ip, nil
}
c.mu.Unlock()
// Use an ephemeral port in the dynamic range (49152-65535).
// Port 0 is not valid with UDP/TCP protocols per RFC 6887.
ephemeralPort := 49152 + int(uint16(time.Now().UnixNano()))%(65535-49152)
// Use minimal lifetime (1 second) for discovery.
resp, err := c.AddPortMapping(ctx, "udp", ephemeralPort, time.Second)
if err != nil {
return nil, fmt.Errorf("create temporary mapping: %w", err)
}
if err := c.DeletePortMapping(ctx, "udp", ephemeralPort); err != nil {
log.Debugf("cleanup temporary PCP mapping: %v", err)
}
return resp.ExternalIP.AsSlice(), nil
}
// LastEpoch returns the last observed server epoch value.
// A decrease in epoch indicates the server may have restarted and mappings may be lost.
func (c *Client) LastEpoch() uint32 {
c.mu.Lock()
defer c.mu.Unlock()
return c.lastEpoch
}
// EpochStateLost returns true if epoch state loss was detected and clears the flag.
func (c *Client) EpochStateLost() bool {
c.mu.Lock()
defer c.mu.Unlock()
lost := c.epochStateLost
c.epochStateLost = false
return lost
}
// updateEpoch updates the epoch tracking and detects potential state loss.
// Returns true if state loss was detected (server likely restarted).
// Caller must hold c.mu.
func (c *Client) updateEpochLocked(newEpoch uint32) bool {
now := time.Now()
stateLost := false
// RFC 6887 Section 8.5: Detect invalid epoch indicating server state loss.
// client_delta = time since last response
// server_delta = epoch change since last response
// Invalid if: client_delta+2 < server_delta - server_delta/16
// OR: server_delta+2 < client_delta - client_delta/16
// The +2 handles quantization, /16 (6.25%) handles clock drift.
if !c.epochTime.IsZero() && c.lastEpoch > 0 {
clientDelta := uint32(now.Sub(c.epochTime).Seconds())
serverDelta := newEpoch - c.lastEpoch
// Check for epoch going backwards or jumping unexpectedly.
// Subtraction is safe: serverDelta/16 is always <= serverDelta.
if clientDelta+2 < serverDelta-(serverDelta/16) ||
serverDelta+2 < clientDelta-(clientDelta/16) {
stateLost = true
c.epochStateLost = true
}
}
c.lastEpoch = newEpoch
c.epochTime = now
return stateLost
}
// cacheExternalIP stores the external IP from a successful MAP response.
// Caller must hold c.mu.
func (c *Client) cacheExternalIPLocked(ip netip.Addr) {
if ip.IsValid() && !ip.IsUnspecified() {
c.externalIP = ip
}
}
// sendRequest sends a PCP request with retries per RFC 6887 Section 8.1.1.
func (c *Client) sendRequest(ctx context.Context, req []byte) ([]byte, error) {
addr := &net.UDPAddr{IP: c.gateway.AsSlice(), Port: Port}
var lastErr error
delay := initialRetryDelay
for range maxRetries {
resp, err := c.sendOnce(ctx, addr, req)
if err == nil {
return resp, nil
}
lastErr = err
if ctx.Err() != nil {
return nil, ctx.Err()
}
// RFC 6887 Section 8.1.1: RT = (1 + RAND) * MIN(2 * RTprev, MRT)
// RAND is random between -0.1 and +0.1
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(retryDelayWithJitter(delay)):
}
delay = min(delay*2, maxRetryDelay)
}
return nil, fmt.Errorf("PCP request failed after %d retries: %w", maxRetries, lastErr)
}
// retryDelayWithJitter applies RFC 6887 jitter: multiply by (1 + RAND) where RAND is [-0.1, +0.1].
func retryDelayWithJitter(d time.Duration) time.Duration {
var b [1]byte
_, _ = rand.Read(b[:])
// Convert byte to range [-0.1, +0.1]: (b/255 * 0.2) - 0.1
jitter := (float64(b[0])/255.0)*0.2 - 0.1
return time.Duration(float64(d) * (1 + jitter))
}
func (c *Client) sendOnce(ctx context.Context, addr *net.UDPAddr, req []byte) ([]byte, error) {
// Use ListenUDP instead of DialUDP to validate response source address per RFC 6887 §8.3.
conn, err := net.ListenUDP("udp", nil)
if err != nil {
return nil, fmt.Errorf("listen: %w", err)
}
defer func() {
if err := conn.Close(); err != nil {
log.Debugf("close UDP connection: %v", err)
}
}()
timeout := c.timeout
if deadline, ok := ctx.Deadline(); ok {
if remaining := time.Until(deadline); remaining < timeout {
timeout = remaining
}
}
if err := conn.SetDeadline(time.Now().Add(timeout)); err != nil {
return nil, fmt.Errorf("set deadline: %w", err)
}
if _, err := conn.WriteToUDP(req, addr); err != nil {
return nil, fmt.Errorf("write: %w", err)
}
resp := make([]byte, responseBufferSize)
n, from, err := conn.ReadFromUDP(resp)
if err != nil {
return nil, fmt.Errorf("read: %w", err)
}
// RFC 6887 §8.3: Validate response came from expected PCP server.
if !from.IP.Equal(addr.IP) {
return nil, fmt.Errorf("response from unexpected source %s (expected %s)", from.IP, addr.IP)
}
return resp[:n], nil
}
func (c *Client) getLocalIP() (netip.Addr, error) {
c.mu.Lock()
defer c.mu.Unlock()
if !c.localIP.IsValid() {
return netip.Addr{}, fmt.Errorf("local IP not set for gateway %s", c.gateway)
}
return c.localIP, nil
}
func protocolNumber(protocol string) (uint8, error) {
switch protocol {
case "udp", "UDP":
return ProtoUDP, nil
case "tcp", "TCP":
return ProtoTCP, nil
default:
return 0, fmt.Errorf("unsupported protocol: %s", protocol)
}
}
// Error represents a PCP error response.
type Error struct {
Code uint8
Message string
}
func (e *Error) Error() string {
return fmt.Sprintf("PCP error: %s (%d)", e.Message, e.Code)
}

View File

@@ -0,0 +1,187 @@
package pcp
import (
"context"
"net"
"net/netip"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAddrConversion(t *testing.T) {
tests := []struct {
name string
addr netip.Addr
}{
{"IPv4", netip.MustParseAddr("192.168.1.100")},
{"IPv4 loopback", netip.MustParseAddr("127.0.0.1")},
{"IPv6", netip.MustParseAddr("2001:db8::1")},
{"IPv6 loopback", netip.MustParseAddr("::1")},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
b16 := addrTo16(tt.addr)
recovered := addrFrom16(b16)
assert.Equal(t, tt.addr, recovered, "address should round-trip")
})
}
}
func TestBuildAnnounceRequest(t *testing.T) {
clientIP := netip.MustParseAddr("192.168.1.100")
req := buildAnnounceRequest(clientIP)
require.Len(t, req, headerSize)
assert.Equal(t, byte(Version), req[0], "version")
assert.Equal(t, byte(OpAnnounce), req[1], "opcode")
// Check client IP is properly encoded as IPv4-mapped IPv6
assert.Equal(t, byte(0xff), req[18], "IPv4-mapped prefix byte 10")
assert.Equal(t, byte(0xff), req[19], "IPv4-mapped prefix byte 11")
assert.Equal(t, byte(192), req[20], "IP octet 1")
assert.Equal(t, byte(168), req[21], "IP octet 2")
assert.Equal(t, byte(1), req[22], "IP octet 3")
assert.Equal(t, byte(100), req[23], "IP octet 4")
}
func TestBuildMapRequest(t *testing.T) {
clientIP := netip.MustParseAddr("192.168.1.100")
nonce := [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
req := buildMapRequest(clientIP, nonce, ProtoUDP, 51820, 51820, netip.Addr{}, 3600)
require.Len(t, req, mapRequestSize)
assert.Equal(t, byte(Version), req[0], "version")
assert.Equal(t, byte(OpMap), req[1], "opcode")
// Lifetime at bytes 4-7
assert.Equal(t, uint32(3600), (uint32(req[4])<<24)|(uint32(req[5])<<16)|(uint32(req[6])<<8)|uint32(req[7]), "lifetime")
// Nonce at bytes 24-35
assert.Equal(t, nonce[:], req[24:36], "nonce")
// Protocol at byte 36
assert.Equal(t, byte(ProtoUDP), req[36], "protocol")
// Internal port at bytes 40-41
assert.Equal(t, uint16(51820), (uint16(req[40])<<8)|uint16(req[41]), "internal port")
// External port at bytes 42-43
assert.Equal(t, uint16(51820), (uint16(req[42])<<8)|uint16(req[43]), "external port")
}
func TestParseResponse(t *testing.T) {
// Construct a valid ANNOUNCE response
resp := make([]byte, headerSize)
resp[0] = Version
resp[1] = OpAnnounce | OpReply
// Result code = 0 (success)
// Lifetime = 0
// Epoch = 12345
resp[8] = 0
resp[9] = 0
resp[10] = 0x30
resp[11] = 0x39
parsed, err := parseResponse(resp)
require.NoError(t, err)
assert.Equal(t, uint8(Version), parsed.Version)
assert.Equal(t, uint8(OpAnnounce|OpReply), parsed.Opcode)
assert.Equal(t, uint8(ResultSuccess), parsed.ResultCode)
assert.Equal(t, uint32(12345), parsed.Epoch)
}
func TestParseResponseErrors(t *testing.T) {
t.Run("too short", func(t *testing.T) {
_, err := parseResponse([]byte{1, 2, 3})
assert.Error(t, err)
})
t.Run("wrong version", func(t *testing.T) {
resp := make([]byte, headerSize)
resp[0] = 1 // Wrong version
resp[1] = OpReply
_, err := parseResponse(resp)
assert.Error(t, err)
})
t.Run("missing reply bit", func(t *testing.T) {
resp := make([]byte, headerSize)
resp[0] = Version
resp[1] = OpAnnounce // Missing OpReply bit
_, err := parseResponse(resp)
assert.Error(t, err)
})
}
func TestResultCodeString(t *testing.T) {
assert.Equal(t, "SUCCESS", ResultCodeString(ResultSuccess))
assert.Equal(t, "NOT_AUTHORIZED", ResultCodeString(ResultNotAuthorized))
assert.Equal(t, "ADDRESS_MISMATCH", ResultCodeString(ResultAddressMismatch))
assert.Contains(t, ResultCodeString(255), "UNKNOWN")
}
func TestProtocolNumber(t *testing.T) {
proto, err := protocolNumber("udp")
require.NoError(t, err)
assert.Equal(t, uint8(ProtoUDP), proto)
proto, err = protocolNumber("tcp")
require.NoError(t, err)
assert.Equal(t, uint8(ProtoTCP), proto)
proto, err = protocolNumber("UDP")
require.NoError(t, err)
assert.Equal(t, uint8(ProtoUDP), proto)
_, err = protocolNumber("icmp")
assert.Error(t, err)
}
func TestClientCreation(t *testing.T) {
gateway := netip.MustParseAddr("192.168.1.1").AsSlice()
client := NewClient(gateway)
assert.Equal(t, net.IP(gateway), client.Gateway())
assert.Equal(t, defaultTimeout, client.timeout)
clientWithTimeout := NewClientWithTimeout(gateway, 5*time.Second)
assert.Equal(t, 5*time.Second, clientWithTimeout.timeout)
}
func TestNATType(t *testing.T) {
n := NewNAT(netip.MustParseAddr("192.168.1.1").AsSlice(), netip.MustParseAddr("192.168.1.100").AsSlice())
assert.Equal(t, "PCP", n.Type())
}
// Integration test - skipped unless PCP_TEST_GATEWAY env is set
func TestClientIntegration(t *testing.T) {
t.Skip("Integration test - run manually with PCP_TEST_GATEWAY=<gateway-ip>")
gateway := netip.MustParseAddr("10.0.1.1").AsSlice() // Change to your test gateway
localIP := netip.MustParseAddr("10.0.1.100").AsSlice() // Change to your local IP
client := NewClient(gateway)
client.SetLocalIP(localIP)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// Test ANNOUNCE
epoch, err := client.Announce(ctx)
require.NoError(t, err)
t.Logf("Server epoch: %d", epoch)
// Test MAP
resp, err := client.AddPortMapping(ctx, "udp", 51820, 1*time.Hour)
require.NoError(t, err)
t.Logf("Mapping: internal=%d external=%d externalIP=%s",
resp.InternalPort, resp.ExternalPort, resp.ExternalIP)
// Cleanup
err = client.DeletePortMapping(ctx, "udp", 51820)
require.NoError(t, err)
}

View File

@@ -0,0 +1,222 @@
package pcp
import (
"context"
"fmt"
"net"
"net/netip"
"runtime"
"sync"
"time"
log "github.com/sirupsen/logrus"
"github.com/libp2p/go-nat"
"github.com/libp2p/go-netroute"
)
var _ nat.NAT = (*NAT)(nil)
// NAT implements the go-nat NAT interface using PCP.
// Supports dual-stack (IPv4 and IPv6) when available.
// All methods are safe for concurrent use.
//
// TODO: IPv6 pinholes use the local IPv6 address. If the address changes
// (e.g., due to SLAAC rotation or network change), the pinhole becomes stale
// and needs to be recreated with the new address.
type NAT struct {
client *Client
mu sync.RWMutex
// client6 is the IPv6 PCP client, nil if IPv6 is unavailable.
client6 *Client
// localIP6 caches the local IPv6 address used for PCP requests.
localIP6 netip.Addr
}
// NewNAT creates a new NAT instance backed by PCP.
func NewNAT(gateway, localIP net.IP) *NAT {
client := NewClient(gateway)
client.SetLocalIP(localIP)
return &NAT{
client: client,
}
}
// Type returns "PCP" as the NAT type.
func (n *NAT) Type() string {
return "PCP"
}
// GetDeviceAddress returns the gateway IP address.
func (n *NAT) GetDeviceAddress() (net.IP, error) {
return n.client.Gateway(), nil
}
// GetExternalAddress returns the external IP address.
func (n *NAT) GetExternalAddress() (net.IP, error) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
return n.client.GetExternalAddress(ctx)
}
// GetInternalAddress returns the local IP address used to communicate with the gateway.
func (n *NAT) GetInternalAddress() (net.IP, error) {
addr, err := n.client.getLocalIP()
if err != nil {
return nil, err
}
return addr.AsSlice(), nil
}
// AddPortMapping creates a port mapping on both IPv4 and IPv6 (if available).
func (n *NAT) AddPortMapping(ctx context.Context, protocol string, internalPort int, _ string, timeout time.Duration) (int, error) {
resp, err := n.client.AddPortMapping(ctx, protocol, internalPort, timeout)
if err != nil {
return 0, fmt.Errorf("add mapping: %w", err)
}
n.mu.RLock()
client6 := n.client6
localIP6 := n.localIP6
n.mu.RUnlock()
if client6 == nil {
return int(resp.ExternalPort), nil
}
if _, err := client6.AddPortMapping(ctx, protocol, internalPort, timeout); err != nil {
log.Warnf("IPv6 PCP mapping failed (continuing with IPv4): %v", err)
return int(resp.ExternalPort), nil
}
log.Infof("created IPv6 PCP pinhole: %s:%d", localIP6, internalPort)
return int(resp.ExternalPort), nil
}
// DeletePortMapping removes a port mapping from both IPv4 and IPv6.
func (n *NAT) DeletePortMapping(ctx context.Context, protocol string, internalPort int) error {
err := n.client.DeletePortMapping(ctx, protocol, internalPort)
n.mu.RLock()
client6 := n.client6
n.mu.RUnlock()
if client6 != nil {
if err6 := client6.DeletePortMapping(ctx, protocol, internalPort); err6 != nil {
log.Warnf("IPv6 PCP delete mapping failed: %v", err6)
}
}
if err != nil {
return fmt.Errorf("delete mapping: %w", err)
}
return nil
}
// CheckServerHealth sends an ANNOUNCE to verify the server is still responsive.
// Returns the current epoch and whether the server may have restarted (epoch state loss detected).
func (n *NAT) CheckServerHealth(ctx context.Context) (epoch uint32, serverRestarted bool, err error) {
epoch, err = n.client.Announce(ctx)
if err != nil {
return 0, false, fmt.Errorf("announce: %w", err)
}
return epoch, n.client.EpochStateLost(), nil
}
// DiscoverPCP attempts to discover a PCP-capable gateway.
// Returns a NAT interface if PCP is supported, or an error otherwise.
// Discovers both IPv4 and IPv6 gateways when available.
func DiscoverPCP(ctx context.Context) (nat.NAT, error) {
gateway, localIP, err := getDefaultGateway()
if err != nil {
return nil, fmt.Errorf("get default gateway: %w", err)
}
client := NewClient(gateway)
client.SetLocalIP(localIP)
if _, err := client.Announce(ctx); err != nil {
return nil, fmt.Errorf("PCP announce: %w", err)
}
result := &NAT{client: client}
discoverIPv6(ctx, result)
return result, nil
}
func discoverIPv6(ctx context.Context, result *NAT) {
gateway6, localIP6, err := getDefaultGateway6()
if err != nil {
log.Debugf("IPv6 gateway discovery failed: %v", err)
return
}
client6 := NewClient(gateway6)
client6.SetLocalIP(localIP6)
if _, err := client6.Announce(ctx); err != nil {
log.Debugf("PCP IPv6 announce failed: %v", err)
return
}
addr, ok := netip.AddrFromSlice(localIP6)
if !ok {
log.Debugf("invalid IPv6 local IP: %v", localIP6)
return
}
result.mu.Lock()
result.client6 = client6
result.localIP6 = addr
result.mu.Unlock()
log.Debugf("PCP IPv6 gateway discovered: %s (local: %s)", gateway6, localIP6)
}
// getDefaultGateway returns the default IPv4 gateway and local IP using the system routing table.
func getDefaultGateway() (gateway net.IP, localIP net.IP, err error) {
router, err := netroute.New()
if err != nil {
return nil, nil, err
}
dst := net.IPv4zero
if runtime.GOOS == "linux" || runtime.GOOS == "android" {
// go-netroute v0.4.0 rejects unspecified destinations client-side on Linux/Android.
// TODO: on android/ios, use platform APIs (ConnectivityManager.getLinkProperties /
// NWPathMonitor) when netlink-based lookup is restricted or unavailable.
dst = net.IPv4(0, 0, 0, 1)
}
_, gateway, localIP, err = router.Route(dst)
if err != nil {
return nil, nil, err
}
if gateway == nil {
return nil, nil, nat.ErrNoNATFound
}
return gateway, localIP, nil
}
// getDefaultGateway6 returns the default IPv6 gateway IP address using the system routing table.
func getDefaultGateway6() (gateway net.IP, localIP net.IP, err error) {
router, err := netroute.New()
if err != nil {
return nil, nil, err
}
dst := net.IPv6zero
if runtime.GOOS == "linux" || runtime.GOOS == "android" {
// ::2
dst = net.IP{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}
}
_, gateway, localIP, err = router.Route(dst)
if err != nil {
return nil, nil, err
}
if gateway == nil {
return nil, nil, nat.ErrNoNATFound
}
return gateway, localIP, nil
}

View File

@@ -0,0 +1,225 @@
// Package pcp implements the Port Control Protocol (RFC 6887).
//
// # Implemented Features
//
// - ANNOUNCE opcode: Discovers PCP server support
// - MAP opcode: Creates/deletes port mappings (IPv4 NAT) and firewall pinholes (IPv6)
// - Dual-stack: Simultaneous IPv4 and IPv6 support via separate clients
// - Nonce validation: Prevents response spoofing
// - Epoch tracking: Detects server restarts per Section 8.5
// - RFC-compliant retry timing: 3s initial, exponential backoff to 1024s max (Section 8.1.1)
//
// # Not Implemented
//
// - PEER opcode: For outbound peer connections (not needed for inbound NAT traversal)
// - THIRD_PARTY option: For managing mappings on behalf of other devices
// - PREFER_FAILURE option: Requires exact external port or fail (IPv4 NAT only, not needed for IPv6 pinholing)
// - FILTER option: To restrict remote peer addresses
//
// These optional features are omitted because the primary use case is simple
// port forwarding for WireGuard, which only requires MAP with default behavior.
package pcp
import (
"encoding/binary"
"fmt"
"net/netip"
)
const (
// Version is the PCP protocol version (RFC 6887).
Version = 2
// Port is the standard PCP server port.
Port = 5351
// DefaultLifetime is the default requested mapping lifetime in seconds.
DefaultLifetime = 7200 // 2 hours
// Header sizes
headerSize = 24
mapPayloadSize = 36
mapRequestSize = headerSize + mapPayloadSize // 60 bytes
)
// Opcodes
const (
OpAnnounce = 0
OpMap = 1
OpPeer = 2
OpReply = 0x80 // OR'd with opcode in responses
)
// Protocol numbers for MAP requests
const (
ProtoUDP = 17
ProtoTCP = 6
)
// Result codes (RFC 6887 Section 7.4)
const (
ResultSuccess = 0
ResultUnsuppVersion = 1
ResultNotAuthorized = 2
ResultMalformedRequest = 3
ResultUnsuppOpcode = 4
ResultUnsuppOption = 5
ResultMalformedOption = 6
ResultNetworkFailure = 7
ResultNoResources = 8
ResultUnsuppProtocol = 9
ResultUserExQuota = 10
ResultCannotProvideExt = 11
ResultAddressMismatch = 12
ResultExcessiveRemotePeers = 13
)
// ResultCodeString returns a human-readable string for a result code.
func ResultCodeString(code uint8) string {
switch code {
case ResultSuccess:
return "SUCCESS"
case ResultUnsuppVersion:
return "UNSUPP_VERSION"
case ResultNotAuthorized:
return "NOT_AUTHORIZED"
case ResultMalformedRequest:
return "MALFORMED_REQUEST"
case ResultUnsuppOpcode:
return "UNSUPP_OPCODE"
case ResultUnsuppOption:
return "UNSUPP_OPTION"
case ResultMalformedOption:
return "MALFORMED_OPTION"
case ResultNetworkFailure:
return "NETWORK_FAILURE"
case ResultNoResources:
return "NO_RESOURCES"
case ResultUnsuppProtocol:
return "UNSUPP_PROTOCOL"
case ResultUserExQuota:
return "USER_EX_QUOTA"
case ResultCannotProvideExt:
return "CANNOT_PROVIDE_EXTERNAL"
case ResultAddressMismatch:
return "ADDRESS_MISMATCH"
case ResultExcessiveRemotePeers:
return "EXCESSIVE_REMOTE_PEERS"
default:
return fmt.Sprintf("UNKNOWN(%d)", code)
}
}
// Response represents a parsed PCP response header.
type Response struct {
Version uint8
Opcode uint8
ResultCode uint8
Lifetime uint32
Epoch uint32
}
// MapResponse contains the full response to a MAP request.
type MapResponse struct {
Response
Nonce [12]byte
Protocol uint8
InternalPort uint16
ExternalPort uint16
ExternalIP netip.Addr
}
// addrTo16 converts an address to its 16-byte IPv4-mapped IPv6 representation.
func addrTo16(addr netip.Addr) [16]byte {
if addr.Is4() {
return netip.AddrFrom4(addr.As4()).As16()
}
return addr.As16()
}
// addrFrom16 extracts an address from a 16-byte representation, unmapping IPv4.
func addrFrom16(b [16]byte) netip.Addr {
return netip.AddrFrom16(b).Unmap()
}
// buildAnnounceRequest creates a PCP ANNOUNCE request packet.
func buildAnnounceRequest(clientIP netip.Addr) []byte {
req := make([]byte, headerSize)
req[0] = Version
req[1] = OpAnnounce
mapped := addrTo16(clientIP)
copy(req[8:24], mapped[:])
return req
}
// buildMapRequest creates a PCP MAP request packet.
func buildMapRequest(clientIP netip.Addr, nonce [12]byte, protocol uint8, internalPort, suggestedExtPort uint16, suggestedExtIP netip.Addr, lifetime uint32) []byte {
req := make([]byte, mapRequestSize)
// Header
req[0] = Version
req[1] = OpMap
binary.BigEndian.PutUint32(req[4:8], lifetime)
mapped := addrTo16(clientIP)
copy(req[8:24], mapped[:])
// MAP payload
copy(req[24:36], nonce[:])
req[36] = protocol
binary.BigEndian.PutUint16(req[40:42], internalPort)
binary.BigEndian.PutUint16(req[42:44], suggestedExtPort)
if suggestedExtIP.IsValid() {
extMapped := addrTo16(suggestedExtIP)
copy(req[44:60], extMapped[:])
}
return req
}
// parseResponse parses the common PCP response header.
func parseResponse(data []byte) (*Response, error) {
if len(data) < headerSize {
return nil, fmt.Errorf("response too short: %d bytes", len(data))
}
resp := &Response{
Version: data[0],
Opcode: data[1],
ResultCode: data[3], // Byte 2 is reserved, byte 3 is result code (RFC 6887 §7.2)
Lifetime: binary.BigEndian.Uint32(data[4:8]),
Epoch: binary.BigEndian.Uint32(data[8:12]),
}
if resp.Version != Version {
return nil, fmt.Errorf("unsupported PCP version: %d", resp.Version)
}
if resp.Opcode&OpReply == 0 {
return nil, fmt.Errorf("response missing reply bit: opcode=0x%02x", resp.Opcode)
}
return resp, nil
}
// parseMapResponse parses a complete MAP response.
func parseMapResponse(data []byte) (*MapResponse, error) {
if len(data) < mapRequestSize {
return nil, fmt.Errorf("MAP response too short: %d bytes", len(data))
}
resp, err := parseResponse(data)
if err != nil {
return nil, fmt.Errorf("parse header: %w", err)
}
mapResp := &MapResponse{
Response: *resp,
Protocol: data[36],
InternalPort: binary.BigEndian.Uint16(data[40:42]),
ExternalPort: binary.BigEndian.Uint16(data[42:44]),
ExternalIP: addrFrom16([16]byte(data[44:60])),
}
copy(mapResp.Nonce[:], data[24:36])
return mapResp, nil
}

View File

@@ -1,116 +0,0 @@
//go:build !js
package portforward
import (
"context"
"errors"
"strings"
"testing"
"github.com/netbirdio/go-nat"
log "github.com/sirupsen/logrus"
"github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// mockPinholeNAT is a gateway that also reports an IPv6 pinhole outcome, the
// shape a dual-stack gateway has.
type mockPinholeNAT struct {
*mockNAT
pinholeErr error
}
func (m *mockPinholeNAT) IPv6PinholeError() error {
return m.pinholeErr
}
func TestSetupLogsPinholeOutcome(t *testing.T) {
pinholeErr := errors.New("pcp ipv6: NOT_AUTHORIZED")
tests := []struct {
name string
pinholeErr error
mappingErr error
wantLevel log.Level
wantText string
}{
{
name: "an open pinhole is reported",
wantLevel: log.InfoLevel,
wantText: "IPv6 pinhole open",
},
{
name: "a failed pinhole is reported without failing the mapping",
// The IPv4 mapping is what the caller asked for, so the pinhole
// failure surfaces only in the log.
pinholeErr: pinholeErr,
wantLevel: log.WarnLevel,
wantText: pinholeErr.Error(),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gateway := &mockPinholeNAT{mockNAT: newMockNAT(), pinholeErr: tt.pinholeErr}
hook := stubGatewayDiscovery(t, gateway)
m := NewManager()
m.wgPort = 51820
_, mapping, err := m.setup(context.Background())
require.NoError(t, err)
require.NotNil(t, mapping)
entry := findEntry(hook, tt.wantText)
require.NotNil(t, entry, "no log entry mentioning %q", tt.wantText)
assert.Equal(t, tt.wantLevel, entry.Level)
})
}
t.Run("a failed mapping reports no pinhole outcome", func(t *testing.T) {
// Nothing opened the pinhole, so whatever it currently reports says
// nothing about this attempt.
gateway := &mockPinholeNAT{mockNAT: newMockNAT()}
gateway.addMappingErr = errors.New("gateway refused")
hook := stubGatewayDiscovery(t, gateway)
m := NewManager()
m.wgPort = 51820
_, _, err := m.setup(context.Background())
require.Error(t, err)
assert.Nil(t, findEntry(hook, "IPv6 pinhole"))
})
}
// stubGatewayDiscovery makes discovery return gateway and captures log output.
func stubGatewayDiscovery(t *testing.T, gateway nat.NAT) *test.Hook {
t.Helper()
orig := discoverGateway
discoverGateway = func(context.Context) (nat.NAT, error) { return gateway, nil }
t.Cleanup(func() { discoverGateway = orig })
hook := test.NewGlobal()
origLevel := log.GetLevel()
log.SetLevel(log.DebugLevel)
t.Cleanup(func() {
hook.Reset()
log.SetLevel(origLevel)
})
return hook
}
func findEntry(hook *test.Hook, substr string) *log.Entry {
for _, entry := range hook.AllEntries() {
if strings.Contains(entry.Message, substr) {
return entry
}
}
return nil
}

View File

@@ -4,94 +4,27 @@ package portforward
import (
"context"
"errors"
"fmt"
"time"
"github.com/netbirdio/go-nat"
"github.com/netbirdio/go-nat/pcp"
"github.com/libp2p/go-nat"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/client/internal/portforward/pcp"
)
// discoverGateway is the function used for NAT gateway discovery.
// It can be replaced in tests to avoid real network operations.
// Tries PCP first, then falls back to NAT-PMP/UPnP.
var discoverGateway = defaultDiscoverGateway
// pinholeDiscoveryTimeout is the slice of the discovery budget held back for
// the IPv6 pinhole probe.
//
// Sizing it is coarser than it looks: PCP retransmits on a 3s socket timeout
// and a 3s first backoff, so a second attempt needs about 9s. Anything from
// roughly 1s to 8s therefore buys exactly one attempt, and this only sets how
// long that attempt waits. A PCP server sits on the local link and answers in
// milliseconds, so 3s is margin rather than need, and the rest is left to
// gateway discovery, whose multicast SSDP search alone takes 5s. A probe lost
// to a dropped packet is retried by the next discovery round.
//
// It is a variable so tests can shorten it.
var pinholeDiscoveryTimeout = 3 * time.Second
// Discovery entry points, as variables so tests can drive the fallback without
// touching the network.
var (
discoverNATGateway = nat.DiscoverGateway
discoverPCPPinhole = func(ctx context.Context) (nat.NAT, error) {
pinhole, err := pcp.DiscoverPCP(ctx)
if err != nil {
return nil, err
}
return pinhole, nil
}
)
// defaultDiscoverGateway finds a gateway that can make the WireGuard port
// reachable. DiscoverGateway prefers PCP for IPv4, races UPnP and NAT-PMP
// behind it, and attaches an IPv6 pinhole independently of which IPv4 protocol
// wins.
//
// It reports no gateway on a network offering only IPv6, having no IPv4 mapping
// to attach a pinhole to. Such a network still needs one: there is no
// translation to traverse, but the router drops inbound IPv6 until something
// opens it. Fall back to PCP alone, which yields a gateway holding just the
// pinhole.
func defaultDiscoverGateway(ctx context.Context) (nat.NAT, error) {
gatewayCtx, cancel := reserveForPinhole(ctx)
defer cancel()
gateway, err := discoverNATGateway(gatewayCtx)
pcpGateway, err := pcp.DiscoverPCP(ctx)
if err == nil {
return gateway, nil
}
if !errors.Is(err, nat.ErrNoNATFound) {
return nil, err
return pcpGateway, nil
}
log.Debugf("PCP discovery failed: %v, trying NAT-PMP/UPnP", err)
pinhole, pinholeErr := discoverPCPPinhole(ctx)
if pinholeErr != nil {
log.Debugf("no IPv6 pinhole after %v: %v", err, pinholeErr)
return nil, err
}
log.Infof("no IPv4 gateway, continuing with an IPv6 pinhole only")
return pinhole, nil
}
// reserveForPinhole shortens ctx so that a pinhole probe still has time to run
// afterwards. Finding nothing takes gateway discovery everything it is given,
// so on the unshortened context the probe would start already expired. A budget
// too small to divide is left to gateway discovery, which is the likelier win.
func reserveForPinhole(ctx context.Context) (context.Context, context.CancelFunc) {
deadline, ok := ctx.Deadline()
if !ok {
return context.WithCancel(ctx)
}
remaining := time.Until(deadline)
if remaining <= pinholeDiscoveryTimeout {
return context.WithCancel(ctx)
}
return context.WithTimeout(ctx, remaining-pinholeDiscoveryTimeout)
return nat.DiscoverGateway(ctx)
}
// State is persisted only for crash recovery cleanup

View File

@@ -1,140 +0,0 @@
//go:build !js
package portforward
import (
"context"
"errors"
"testing"
"time"
"github.com/netbirdio/go-nat"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// stubDiscovery replaces both discovery entry points for the duration of a
// test. gatewayDelay simulates gateway discovery spending everything it is
// given before reporting that it found nothing.
func stubDiscovery(t *testing.T, gateway nat.NAT, gatewayErr error, gatewayDelay time.Duration, pinhole nat.NAT, pinholeErr error) {
t.Helper()
origGateway, origPinhole := discoverNATGateway, discoverPCPPinhole
discoverNATGateway = func(ctx context.Context) (nat.NAT, error) {
if gatewayDelay > 0 {
select {
case <-time.After(gatewayDelay):
case <-ctx.Done():
}
}
return gateway, gatewayErr
}
discoverPCPPinhole = func(ctx context.Context) (nat.NAT, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
return pinhole, pinholeErr
}
t.Cleanup(func() { discoverNATGateway, discoverPCPPinhole = origGateway, origPinhole })
}
func TestDefaultDiscoverGateway(t *testing.T) {
ipv4Gateway := &mockNAT{natType: "PCP+PCPv6"}
ipv6Pinhole := &mockNAT{natType: "PCP"}
otherErr := errors.New("routing table unavailable")
t.Run("an IPv4 gateway is used as is", func(t *testing.T) {
stubDiscovery(t, ipv4Gateway, nil, 0, ipv6Pinhole, nil)
got, err := defaultDiscoverGateway(context.Background())
require.NoError(t, err)
assert.Same(t, ipv4Gateway, got)
})
t.Run("no IPv4 gateway still opens an IPv6 pinhole", func(t *testing.T) {
stubDiscovery(t, nil, nat.ErrNoNATFound, 0, ipv6Pinhole, nil)
got, err := defaultDiscoverGateway(context.Background())
require.NoError(t, err)
assert.Same(t, ipv6Pinhole, got)
})
t.Run("no gateway and no pinhole reports the original failure", func(t *testing.T) {
stubDiscovery(t, nil, nat.ErrNoNATFound, 0, nil, errors.New("no IPv6 route"))
got, err := defaultDiscoverGateway(context.Background())
assert.Nil(t, got)
assert.ErrorIs(t, err, nat.ErrNoNATFound, "the pinhole failure must not mask why no gateway was found")
})
t.Run("a failure other than no-gateway is reported as is", func(t *testing.T) {
stubDiscovery(t, nil, otherErr, 0, ipv6Pinhole, nil)
got, err := defaultDiscoverGateway(context.Background())
assert.Nil(t, got)
assert.ErrorIs(t, err, otherErr)
})
t.Run("the pinhole survives gateway discovery using its whole budget", func(t *testing.T) {
// On one shared context the probe would start already expired, which is
// how this failed against a real gateway.
reserve := 50 * time.Millisecond
origReserve := pinholeDiscoveryTimeout
pinholeDiscoveryTimeout = reserve
t.Cleanup(func() { pinholeDiscoveryTimeout = origReserve })
budget := 4 * reserve
ctx, cancel := context.WithTimeout(context.Background(), budget)
defer cancel()
stubDiscovery(t, nil, nat.ErrNoNATFound, budget, ipv6Pinhole, nil)
got, err := defaultDiscoverGateway(ctx)
require.NoError(t, err)
assert.Same(t, ipv6Pinhole, got)
})
}
func TestReserveForPinhole(t *testing.T) {
origReserve := pinholeDiscoveryTimeout
pinholeDiscoveryTimeout = time.Second
t.Cleanup(func() { pinholeDiscoveryTimeout = origReserve })
t.Run("a budget is divided", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
gatewayCtx, cancelGateway := reserveForPinhole(ctx)
defer cancelGateway()
deadline, ok := gatewayCtx.Deadline()
require.True(t, ok)
assert.InDelta(t, 9*time.Second, time.Until(deadline), float64(500*time.Millisecond))
})
t.Run("a budget too small to divide is left whole", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
defer cancel()
gatewayCtx, cancelGateway := reserveForPinhole(ctx)
defer cancelGateway()
deadline, ok := gatewayCtx.Deadline()
require.True(t, ok)
assert.InDelta(t, 500*time.Millisecond, time.Until(deadline), float64(100*time.Millisecond))
})
t.Run("no deadline stays unbounded", func(t *testing.T) {
gatewayCtx, cancelGateway := reserveForPinhole(context.Background())
defer cancelGateway()
_, ok := gatewayCtx.Deadline()
assert.False(t, ok)
})
}

View File

@@ -15,7 +15,7 @@ func UpdateStaticInfoAsync() {
}
// GetInfo retrieves system information for WASM environment
func GetInfo(ctx context.Context) *Info {
func GetInfo(_ context.Context) *Info {
info := &Info{
GoOS: runtime.GOOS,
Kernel: runtime.GOARCH,
@@ -30,13 +30,6 @@ func GetInfo(ctx context.Context) *Info {
collectBrowserInfo(info)
collectLocationInfo(info)
collectSystemInfo(info)
// A caller-provided device name wins, as on the other platforms. A peer
// registered over an API keeps reporting the name it was registered with,
// so its meta does not change on the first sync.
if name := extractDeviceName(ctx, info.Hostname); name != "" {
info.Hostname = name
}
return info
}

View File

@@ -1,27 +0,0 @@
//go:build js
package system
import (
"context"
"testing"
)
// TestGetInfoHonorsDeviceName covers a caller-provided device name reaching the
// reported hostname, so a peer registered over an API keeps reporting the name
// it was registered with instead of renaming itself on its first sync.
func TestGetInfoHonorsDeviceName(t *testing.T) {
ctx := context.WithValue(context.Background(), DeviceNameCtxKey, "session-name")
if got := GetInfo(ctx).Hostname; got != "session-name" {
t.Errorf("hostname should carry the caller's device name, got %q", got)
}
}
// TestGetInfoWithoutDeviceNameKeepsFallback covers the embed layer's habit of
// always setting the context value: an empty name must not blank the hostname.
func TestGetInfoWithoutDeviceNameKeepsFallback(t *testing.T) {
ctx := context.WithValue(context.Background(), DeviceNameCtxKey, "")
if got := GetInfo(ctx).Hostname; got == "" {
t.Error("an empty device name must not blank the hostname")
}
}

View File

@@ -1,5 +1,3 @@
//go:build windows || (linux && !android) || (darwin && !ios) || freebsd
package system
import (

View File

@@ -764,19 +764,7 @@
"message": "Sensible Informationen anonymisieren"
},
"settings.troubleshooting.anonymize.help": {
"message": "Verbirgt IP-Adressen, Domains und andere sensible Werte."
},
"settings.troubleshooting.anonymize.info": {
"message": "Der Standardmodus lässt interne IPv4-Adressen und Peer-Namen für den Support lesbar. Der strikte Modus anonymisiert zusätzlich private (RFC 1918), CGNAT- und Link-Local-IP-Adressen, Peer-Namen und öffentliche WireGuard-Schlüssel. Wiederkehrende Werte erhalten denselben Platzhalter, sodass Peers unterscheidbar bleiben. Verwenden Sie den strikten Modus, wenn Sie das Debug-Paket außerhalb Ihrer Organisation weitergeben."
},
"settings.troubleshooting.anonymize.none": {
"message": "Keine"
},
"settings.troubleshooting.anonymize.default": {
"message": "Standard"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Strikt"
"message": "Versteckt öffentliche IP-Adressen und nicht-NetBird-Domains in Logs."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Systeminformationen einschließen"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "Vorgang fehlgeschlagen."
},
"settings.ssh.privilege.hint": {
"message": "Erfordert {actor}. Führen Sie stattdessen dies aus:"
},
"settings.ssh.privilege.oneWay": {
"message": "Sie können dies deaktivieren, aber zum erneuten Aktivieren sind {actor} erforderlich:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Sie können dies aktivieren, aber zum erneuten Deaktivieren sind {actor} erforderlich:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Anonimizar información sensible"
},
"settings.troubleshooting.anonymize.help": {
"message": "Oculta direcciones IP, dominios y otros valores sensibles."
},
"settings.troubleshooting.anonymize.info": {
"message": "El modo predeterminado mantiene legibles las direcciones IPv4 internas y los nombres de los peers para el soporte. El modo estricto anonimiza además las direcciones IP privadas (RFC 1918), CGNAT y de enlace local, los nombres de los peers y las claves públicas de WireGuard. Los valores recurrentes se asignan al mismo marcador de posición, por lo que los peers siguen siendo distinguibles. Use el modo estricto cuando comparta el paquete de diagnóstico fuera de su organización."
},
"settings.troubleshooting.anonymize.none": {
"message": "Ninguno"
},
"settings.troubleshooting.anonymize.default": {
"message": "Predeterminado"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Estricto"
"message": "Oculta las direcciones IP públicas y los dominios ajenos a NetBird de los registros."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Incluir información del sistema"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "La operación falló."
},
"settings.ssh.privilege.hint": {
"message": "Requiere {actor}. Ejecute esto en su lugar:"
},
"settings.ssh.privilege.oneWay": {
"message": "Puede desactivarlo, pero volver a activarlo requiere {actor}:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Puede activarlo, pero volver a desactivarlo requiere {actor}:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Anonymiser les informations sensibles"
},
"settings.troubleshooting.anonymize.help": {
"message": "Masque les adresses IP, les domaines et d'autres valeurs sensibles."
},
"settings.troubleshooting.anonymize.info": {
"message": "Le mode par défaut garde les adresses IPv4 internes et les noms des pairs lisibles pour le support. Le mode strict anonymise en plus les adresses IP privées (RFC 1918), CGNAT et de lien local, les noms des pairs et les clés publiques WireGuard. Les valeurs récurrentes reçoivent le même espace réservé, les pairs restent donc distinguables. Utilisez le mode strict lorsque vous partagez le lot de diagnostic en dehors de votre organisation."
},
"settings.troubleshooting.anonymize.none": {
"message": "Aucune"
},
"settings.troubleshooting.anonymize.default": {
"message": "Par défaut"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Strict"
"message": "Masque les adresses IP publiques et les domaines non-NetBird dans les journaux."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Inclure les informations système"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "Lopération a échoué."
},
"settings.ssh.privilege.hint": {
"message": "Nécessite {actor}. Exécutez plutôt ceci :"
},
"settings.ssh.privilege.oneWay": {
"message": "Vous pouvez le désactiver, mais le réactiver nécessite {actor} :"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Vous pouvez lactiver, mais le désactiver de nouveau nécessite {actor} :"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Érzékeny információk anonimizálása"
},
"settings.troubleshooting.anonymize.help": {
"message": "Elrejti az IP-címeket, a tartományokat és más érzékeny értékeket."
},
"settings.troubleshooting.anonymize.info": {
"message": "Az Alapértelmezett szint a belső IPv4-címeket és a peer-neveket olvashatóan hagyja a támogatás számára. A Szigorú ezen felül anonimizálja a privát (RFC 1918), CGNAT és link-local IP-címeket, a peer-neveket és a WireGuard nyilvános kulcsokat. Az ismétlődő értékek ugyanazt a helyettesítőt kapják, így a peerek megkülönböztethetők maradnak. Használja a Szigorú szintet, ha a hibakeresési csomagot a szervezetén kívül osztja meg."
},
"settings.troubleshooting.anonymize.none": {
"message": "Nincs"
},
"settings.troubleshooting.anonymize.default": {
"message": "Alapértelmezett"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Szigorú"
"message": "Elrejti a nyilvános IP-címeket és a nem-NetBird tartományokat a naplókban."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Rendszerinformációk beillesztése"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "A művelet meghiúsult."
},
"settings.ssh.privilege.hint": {
"message": "{actor} szükséges hozzá. Futtassa inkább ezt:"
},
"settings.ssh.privilege.oneWay": {
"message": "Kikapcsolhatja, de a visszakapcsolásához {actor} szükséges:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Bekapcsolhatja, de az ismételt kikapcsolásához {actor} szükséges:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Anonimizza informazioni sensibili"
},
"settings.troubleshooting.anonymize.help": {
"message": "Nasconde indirizzi IP, domini e altri valori sensibili."
},
"settings.troubleshooting.anonymize.info": {
"message": "La modalità predefinita mantiene leggibili gli indirizzi IPv4 interni e i nomi dei peer per il supporto. La modalità rigorosa anonimizza inoltre gli indirizzi IP privati (RFC 1918), CGNAT e link-local, i nomi dei peer e le chiavi pubbliche WireGuard. I valori ricorrenti vengono associati allo stesso segnaposto, quindi i peer restano distinguibili. Usa la modalità rigorosa quando condividi il pacchetto di debug al di fuori della tua organizzazione."
},
"settings.troubleshooting.anonymize.none": {
"message": "Nessuna"
},
"settings.troubleshooting.anonymize.default": {
"message": "Predefinito"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Rigoroso"
"message": "Nasconde gli indirizzi IP pubblici e i domini non NetBird dai log."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Includi informazioni di sistema"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "Operazione non riuscita."
},
"settings.ssh.privilege.hint": {
"message": "Richiede {actor}. Esegua invece questo:"
},
"settings.ssh.privilege.oneWay": {
"message": "Può disabilitarlo, ma riabilitarlo richiede {actor}:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Può abilitarlo, ma disabilitarlo di nuovo richiede {actor}:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "機密情報を匿名化"
},
"settings.troubleshooting.anonymize.help": {
"message": "IP アドレス、ドメイン、その他の機密性の高い値を隠します。"
},
"settings.troubleshooting.anonymize.info": {
"message": "「デフォルト」では、サポートのために内部 IPv4 アドレスとピア名は読める状態のまま残ります。「厳格」では、さらにプライベート (RFC 1918)、CGNAT、リンクローカルの IP アドレス、ピア名、WireGuard 公開鍵も匿名化されます。繰り返し現れる値は同じプレースホルダーに置き換えられるため、ピアは区別できます。デバッグバンドルを組織外に共有する場合は「厳格」を使用してください。"
},
"settings.troubleshooting.anonymize.none": {
"message": "なし"
},
"settings.troubleshooting.anonymize.default": {
"message": "デフォルト"
},
"settings.troubleshooting.anonymize.strict": {
"message": "厳格"
"message": "ログからパブリック IP アドレスと NetBird 以外のドメインを隠します。"
},
"settings.troubleshooting.systemInfo.label": {
"message": "システム情報を含める"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "操作に失敗しました。"
},
"settings.ssh.privilege.hint": {
"message": "{actor}が必要です。代わりに次のコマンドを実行してください:"
},
"settings.ssh.privilege.oneWay": {
"message": "無効にはできますが、再度有効にするには{actor}が必要です:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "有効にはできますが、再度無効にするには{actor}が必要です:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Anonimizar informações sensíveis"
},
"settings.troubleshooting.anonymize.help": {
"message": "Oculta endereços IP, domínios e outros valores sensíveis."
},
"settings.troubleshooting.anonymize.info": {
"message": "O modo padrão mantém os endereços IPv4 internos e os nomes dos peers legíveis para o suporte. O modo estrito anonimiza também os endereços IP privados (RFC 1918), CGNAT e link-local, os nomes dos peers e as chaves públicas do WireGuard. Valores recorrentes recebem o mesmo marcador, então os peers continuam distinguíveis. Use o modo estrito ao compartilhar o pacote de depuração fora da sua organização."
},
"settings.troubleshooting.anonymize.none": {
"message": "Nenhum"
},
"settings.troubleshooting.anonymize.default": {
"message": "Padrão"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Estrito"
"message": "Oculta endereços IP públicos e domínios que não são do NetBird nos logs."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Incluir informações do sistema"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "A operação falhou."
},
"settings.ssh.privilege.hint": {
"message": "Requer {actor}. Execute isto em vez disso:"
},
"settings.ssh.privilege.oneWay": {
"message": "Você pode desativar isto, mas ativar novamente requer {actor}:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Você pode ativar isto, mas desativar novamente requer {actor}:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "Анонимизировать конфиденциальную информацию"
},
"settings.troubleshooting.anonymize.help": {
"message": "Скрывает IP-адреса, домены и другие конфиденциальные значения."
},
"settings.troubleshooting.anonymize.info": {
"message": "Режим «По умолчанию» оставляет внутренние IPv4-адреса и имена пиров читаемыми для поддержки. Режим «Строгий» дополнительно анонимизирует частные (RFC 1918), CGNAT и link-local IP-адреса, имена пиров и публичные ключи WireGuard. Повторяющиеся значения заменяются одним и тем же заполнителем, поэтому пиры остаются различимыми. Используйте режим «Строгий», когда передаёте отладочный пакет за пределы вашей организации."
},
"settings.troubleshooting.anonymize.none": {
"message": "Нет"
},
"settings.troubleshooting.anonymize.default": {
"message": "По умолчанию"
},
"settings.troubleshooting.anonymize.strict": {
"message": "Строгий"
"message": "Скрывает публичные IP-адреса и сторонние (не относящиеся к NetBird) домены в журналах."
},
"settings.troubleshooting.systemInfo.label": {
"message": "Включить сведения о системе"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "Не удалось выполнить операцию."
},
"settings.ssh.privilege.hint": {
"message": "Требуются {actor}. Выполните вместо этого:"
},
"settings.ssh.privilege.oneWay": {
"message": "Отключить можно, но чтобы включить снова, нужны {actor}:"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "Включить можно, но чтобы отключить снова, нужны {actor}:"
}
}

View File

@@ -764,19 +764,7 @@
"message": "匿名化敏感信息"
},
"settings.troubleshooting.anonymize.help": {
"message": "隐藏 IP 地址、域名和其他敏感值。"
},
"settings.troubleshooting.anonymize.info": {
"message": "默认级别保留内部 IPv4 地址和对等节点名称,便于支持人员阅读。严格级别还会匿名化私有 (RFC 1918)、CGNAT 和链路本地 IP 地址、对等节点名称以及 WireGuard 公钥。相同的值会映射到相同的占位符,因此对等节点仍可区分。向组织外部分享调试包时请使用严格级别。"
},
"settings.troubleshooting.anonymize.none": {
"message": "无"
},
"settings.troubleshooting.anonymize.default": {
"message": "默认"
},
"settings.troubleshooting.anonymize.strict": {
"message": "严格"
"message": "从日志中隐藏公共 IP 地址和非 NetBird 域名。"
},
"settings.troubleshooting.systemInfo.label": {
"message": "包含系统信息"
@@ -1350,14 +1338,5 @@
},
"error.unknown": {
"message": "操作失败。"
},
"settings.ssh.privilege.hint": {
"message": "需要{actor}。请改为运行:"
},
"settings.ssh.privilege.oneWay": {
"message": "您可以关闭此项,但重新开启需要{actor}"
},
"settings.ssh.privilege.oneWayInverted": {
"message": "您可以开启此项,但再次关闭需要{actor}"
}
}

View File

@@ -56,7 +56,8 @@ func startClient(ctx context.Context, nbClient *netbird.Client) error {
// parseClientOptions extracts NetBird options from JavaScript object
func parseClientOptions(jsOptions js.Value) (netbird.Options, error) {
options := netbird.Options{
LogLevel: defaultLogLevel,
DeviceName: "dashboard-client",
LogLevel: defaultLogLevel,
}
if jwtToken := jsOptions.Get("jwtToken"); !jwtToken.IsNull() && !jwtToken.IsUndefined() {
@@ -86,41 +87,13 @@ func parseClientOptions(jsOptions js.Value) (netbird.Options, error) {
options.DeviceName = deviceName.String()
}
disableIPv6, err := boolOption(jsOptions, "disableIPv6")
if err != nil {
return options, err
if disableIPv6 := jsOptions.Get("disableIPv6"); !disableIPv6.IsNull() && !disableIPv6.IsUndefined() {
options.DisableIPv6 = disableIPv6.Bool()
}
if disableIPv6 != nil {
options.DisableIPv6 = *disableIPv6
}
// The caller decides whether this client uses lazy connections; left unset it
// defers to the management feature flag. A short-lived, interactive caller
// turns it off so its sessions reach the few peers their grant covers eagerly,
// instead of the first request waiting for the connection to be established.
lazyConnectionEnabled, err := boolOption(jsOptions, "lazyConnectionEnabled")
if err != nil {
return options, err
}
options.LazyConnectionEnabled = lazyConnectionEnabled
return options, nil
}
// boolOption reads a boolean option, returning nil when the caller left it out.
// js.Value.Bool panics on any other type, so a wrong type is reported instead.
func boolOption(jsOptions js.Value, name string) (*bool, error) {
v := jsOptions.Get(name)
if v.IsNull() || v.IsUndefined() {
return nil, nil
}
if v.Type() != js.TypeBoolean {
return nil, fmt.Errorf("option %s must be a boolean, got %s", name, v.Type())
}
b := v.Bool()
return &b, nil
}
// createStartMethod creates the start method for the client
func createStartMethod(client *netbird.Client) js.Func {
return js.FuncOf(func(this js.Value, args []js.Value) any {

View File

@@ -1,64 +0,0 @@
//go:build js
package main
import (
"syscall/js"
"testing"
)
// TestParseClientOptionsBooleans covers the boolean options against the value
// kinds a JS caller can pass: js.Value.Bool panics on anything but a boolean,
// so a wrong type has to be rejected before it reaches the client.
func TestParseClientOptionsBooleans(t *testing.T) {
t.Run("unset leaves the lazy override empty", func(t *testing.T) {
options, err := parseClientOptions(js.Global().Get("Object").New())
if err != nil {
t.Fatalf("parse options: %v", err)
}
if options.LazyConnectionEnabled != nil {
t.Errorf("lazy override should stay unset, got %v", *options.LazyConnectionEnabled)
}
if options.DisableIPv6 {
t.Error("disableIPv6 should default to false")
}
})
t.Run("null defers to the management flag", func(t *testing.T) {
jsOptions := js.Global().Get("Object").New()
jsOptions.Set("lazyConnectionEnabled", js.Null())
options, err := parseClientOptions(jsOptions)
if err != nil {
t.Fatalf("parse options: %v", err)
}
if options.LazyConnectionEnabled != nil {
t.Errorf("lazy override should stay unset, got %v", *options.LazyConnectionEnabled)
}
})
t.Run("booleans are carried through", func(t *testing.T) {
jsOptions := js.Global().Get("Object").New()
jsOptions.Set("lazyConnectionEnabled", false)
jsOptions.Set("disableIPv6", true)
options, err := parseClientOptions(jsOptions)
if err != nil {
t.Fatalf("parse options: %v", err)
}
if options.LazyConnectionEnabled == nil || *options.LazyConnectionEnabled {
t.Errorf("lazy override should be false, got %v", options.LazyConnectionEnabled)
}
if !options.DisableIPv6 {
t.Error("disableIPv6 should be true")
}
})
t.Run("a non-boolean is rejected", func(t *testing.T) {
for _, value := range []any{"true", 1, js.Global().Get("Object").New()} {
jsOptions := js.Global().Get("Object").New()
jsOptions.Set("lazyConnectionEnabled", value)
if _, err := parseClientOptions(jsOptions); err == nil {
t.Errorf("value %v should be rejected", value)
}
}
})
}

View File

@@ -1,254 +0,0 @@
//go:build e2e
package agentnetwork
import (
"context"
"encoding/json"
"io"
"net/http"
"net/url"
"os"
"sort"
"strings"
"testing"
"time"
"github.com/stretchr/testify/require"
"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"
sharedllm "github.com/netbirdio/netbird/shared/llm"
)
// TestDebugBedrockProfileCount investigates a listing that reports 100+ models
// for an account whose console shows 38 in the same region. It asserts almost
// nothing — it prints what the production path throws away.
//
// parseListing keeps an id, a name and a status and discards the rest of every
// summary, so the type, the geography and whether a nextToken came back never
// reach a log line. Fetch then works from a list that has already been
// filtered. Neither can answer where the surplus comes from.
//
// It reads the listing three ways:
//
// [1] one GET with no query parameters — byte for byte what Fetch issues,
// which shows how much of the account a single page carries
// [2] the same call followed through nextToken, for the real total
// [3] Fetch itself, for what reaches the dashboard
//
// then decomposes the full set by status, type, geography and vendor, and
// counts distinct models after normalization. Two outcomes need opposite
// fixes and look identical in the dashboard:
//
// - distinct-after-normalization lands near the console's count → the
// surplus is one model offered once per geography, and the question is
// what to offer rather than what broke
// - it does not → we are being handed profiles the console does not show,
// and the filter is what to look at
//
// Uses the same credential as the rest of the live suite:
//
// go test -tags e2e ./e2e/agentnetwork/ -run TestDebugBedrockProfileCount -v
func TestDebugBedrockProfileCount(t *testing.T) {
token := os.Getenv("AWS_BEARER_TOKEN_BEDROCK")
if token == "" {
t.Skip("AWS_BEARER_TOKEN_BEDROCK not set; source ~/.llm-keys to run the Bedrock count debug")
}
region := os.Getenv("AWS_REGION")
if region == "" {
region = "eu-central-1"
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
host := "bedrock." + region + ".amazonaws.com"
t.Logf("=== region %s, control plane %s ===", region, host)
// [1] Exactly what Fetch asks for: no maxResults, no type filter.
first, firstRaw := listInferenceProfiles(t, ctx, host, token, nil)
t.Logf("[1] production-shaped call: %d summaries, %d bytes, nextToken present: %t",
len(first.Summaries), len(firstRaw), first.NextToken != "")
// [2] Followed to exhaustion, so the total is not just a page size.
all := append([]bedrockProfileSummary(nil), first.Summaries...)
next, pages := first.NextToken, 1
for next != "" && pages < 20 {
page, _ := listInferenceProfiles(t, ctx, host, token, map[string]string{"nextToken": next})
all = append(all, page.Summaries...)
next, pages = page.NextToken, pages+1
}
t.Logf("[2] paginated: %d summaries across %d page(s)", len(all), pages)
if next != "" {
t.Logf(" WARNING: stopped at the page cap with a nextToken still outstanding")
}
// [3] The path the Load models button drives, with its ACTIVE filter and
// its dedup.
var cl modeldiscovery.Client
fetched, err := cl.Fetch(ctx, modeldiscovery.Request{
CatalogID: "bedrock_api",
UpstreamURL: "https://bedrock-runtime." + region + ".amazonaws.com",
APIKey: token,
})
require.NoError(t, err, "Fetch must reach the control plane")
t.Logf("[3] Fetch returned %d models (this is what the dashboard renders)", len(fetched))
if len(first.Summaries) == len(all) && len(fetched) > len(all) {
t.Logf(" NOTE: Fetch returned more than the raw listing — the surplus is ours, not AWS's")
}
byStatus, byType, byGeo, byVendor := map[string]int{}, map[string]int{}, map[string]int{}, map[string]int{}
normalized := map[string]struct{}{}
perModel := map[string][]string{}
active := 0
for _, s := range all {
byStatus[orAbsent(s.Status)]++
byType[orAbsent(s.Type)]++
geo, vendor := splitBedrockProfileID(s.ID)
byGeo[geo]++
byVendor[vendor]++
if s.Status != "" && !strings.EqualFold(s.Status, "ACTIVE") {
continue
}
active++
key := sharedllm.NormalizeBedrockModel(s.ID)
normalized[key] = struct{}{}
perModel[key] = append(perModel[key], s.ID)
}
t.Logf("--- ACTIVE summaries: %d of %d", active, len(all))
t.Logf("--- distinct models after normalization: %d <<< compare with the console", len(normalized))
logProfileCounts(t, "by status", byStatus)
logProfileCounts(t, "by type", byType)
logProfileCounts(t, "by geography", byGeo)
logProfileCounts(t, "by vendor", byVendor)
var repeated []string
for key, ids := range perModel {
if len(ids) > 1 {
sort.Strings(ids)
repeated = append(repeated, key+" ("+strings.Join(ids, ", ")+")")
}
}
sort.Strings(repeated)
t.Logf("--- models offered under more than one geography: %d", len(repeated))
for _, line := range repeated {
t.Logf(" %s", line)
}
// A model the catalog cannot price is a catalog gap, not a normalization
// failure. Both render as $0 with a yellow border and need opposite fixes.
entry, ok := catalog.Lookup("bedrock_api")
require.True(t, ok)
var priced, unpriced []string
for id := range normalized {
if _, known := pricing.LookupDefault(entry.PricingSurfaces, id); known {
priced = append(priced, id)
continue
}
unpriced = append(unpriced, id)
}
sort.Strings(priced)
sort.Strings(unpriced)
t.Logf("--- priced by the catalog: %d", len(priced))
for _, id := range priced {
t.Logf(" + %s", id)
}
t.Logf("--- NOT priced by the catalog: %d (catalog coverage, not normalization)", len(unpriced))
for _, id := range unpriced {
t.Logf(" - %s", id)
}
}
type bedrockProfileSummary struct {
ID string `json:"inferenceProfileId"`
Name string `json:"inferenceProfileName"`
Status string `json:"status"`
Type string `json:"type"`
ARN string `json:"inferenceProfileArn"`
}
type bedrockProfilePage struct {
Summaries []bedrockProfileSummary `json:"inferenceProfileSummaries"`
NextToken string `json:"nextToken"`
}
// listInferenceProfiles calls the control plane directly so the whole summary
// is visible, rather than the three fields parseListing keeps.
func listInferenceProfiles(t *testing.T, ctx context.Context, host, token string, query map[string]string) (bedrockProfilePage, []byte) {
t.Helper()
target := url.URL{Scheme: "https", Host: host, Path: "/inference-profiles"}
if len(query) > 0 {
q := target.Query()
for k, v := range query {
q.Set(k, v)
}
target.RawQuery = q.Encode()
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil)
require.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Accept", "application/json")
resp, err := http.DefaultClient.Do(req)
require.NoError(t, err, "reach the Bedrock control plane")
defer func() { _ = resp.Body.Close() }()
raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
require.NoError(t, err)
if resp.StatusCode != http.StatusOK {
// The body is the point of a failure here: an IAM denial names the
// action it refused, which is a different fix from a bad token.
t.Logf("control plane answered %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
t.Logf(" x-amzn-errortype: %s", resp.Header.Get("x-amzn-errortype"))
}
require.Equal(t, http.StatusOK, resp.StatusCode, "control plane must answer the listing")
var page bedrockProfilePage
require.NoError(t, json.Unmarshal(raw, &page), "listing must parse")
return page, raw
}
// splitBedrockProfileID reports the geography and vendor segments of a
// profile id.
func splitBedrockProfileID(id string) (geo, vendor string) {
parts := strings.SplitN(id, ".", 3)
switch len(parts) {
case 3:
return parts[0], parts[1]
case 2:
return "(none)", parts[0]
default:
return "(none)", "(none)"
}
}
func orAbsent(s string) string {
if s == "" {
return "(absent)"
}
return s
}
func logProfileCounts(t *testing.T, label string, counts map[string]int) {
t.Helper()
keys := make([]string, 0, len(counts))
for k := range counts {
keys = append(keys, k)
}
sort.Slice(keys, func(i, j int) bool {
if counts[keys[i]] != counts[keys[j]] {
return counts[keys[i]] > counts[keys[j]]
}
return keys[i] < keys[j]
})
t.Logf("--- %s:", label)
for _, k := range keys {
t.Logf(" %-30s %d", k, counts[k])
}
}

View File

@@ -23,10 +23,9 @@ 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 = harness.VLLMChatInputTokens
vllmCompletionTokens = harness.VLLMChatOutputTokens
vllmPromptTokens = 11
vllmCompletionTokens = 2
)
// pricedEnv is a connected single-provider agent-network deployment pointed at
@@ -163,90 +162,30 @@ func chatOnce(t *testing.T, ctx context.Context, env pricedEnv, model, sessionID
break
}
}
if !waitBeforeRetry(ctx, 5*time.Second) {
break
}
time.Sleep(5 * time.Second)
}
require.Equal(t, 200, code,
"chat for %s must return 200; body: %s\n=== proxy logs ===\n%s", model, body, env.proxy.Logs(context.Background()))
return body
}
// 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
// accessLogPollInterval is how long the lookup waits between pages. Ingest is
// asynchronous, so the row lands somewhere inside the window rather than on
// any particular poll.
const accessLogPollInterval = 2 * 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 {
// Each poll is bounded by what is left of the window rather than by the
// caller's context: a single stalled request would otherwise hold the
// loop open long past the ingest window it is meant to enforce, and the
// caller would read the delay as a missing row.
if logs, lerr := listAccessLogsBy(ctx, deadline); lerr == nil {
for _, r := range logs.Data {
if r.SessionId != nil && *r.SessionId == sessionID {
return r, true
}
}
}
// The wait is bounded by the window as well, so the answer arrives when
// the caller's budget runs out rather than a poll interval later: a
// full interval slept past the deadline reports "no row" up to two
// seconds late, which reads as a slower lookup than the one asked for.
wait := time.Until(deadline)
if wait > accessLogPollInterval {
wait = accessLogPollInterval
}
if wait <= 0 {
return api.AgentNetworkAccessLog{}, false
}
timer := time.NewTimer(wait)
select {
case <-ctx.Done():
timer.Stop()
return api.AgentNetworkAccessLog{}, false
case <-timer.C:
}
// Checked after the wait rather than before the request: a poll issued
// past the deadline carries no budget and would fail on arrival.
if !time.Now().Before(deadline) {
return api.AgentNetworkAccessLog{}, false
}
}
}
// listAccessLogsBy fetches one access-log page under a context that expires at
// deadline, so no single call can outlive the window its caller is polling
// within. The parent's cancellation still applies: the child inherits it.
func listAccessLogsBy(ctx context.Context, deadline time.Time) (api.AgentNetworkAccessLogsResponse, error) {
reqCtx, cancel := context.WithDeadline(ctx, deadline)
defer cancel()
return srv.ListAccessLogs(reqCtx)
}
// 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.
// 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()
row, ok := lookupAccessLogBySession(ctx, sessionID, accessLogIngestWindow)
require.True(t, ok, "session id %q must be recorded in an access-log row", sessionID)
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
}
}
return false
}, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row", sessionID)
return row
}
@@ -380,11 +319,6 @@ 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{
@@ -419,61 +353,27 @@ 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
// The cost last read, kept separately: repriced is the zero value on every
// path that gives up, so reporting its cost would say "$0.000000" whether
// the rows were still at rate A or no row was ever read.
var lastCost float64
var sawRow bool
deadline := time.Now().Add(repriceDeadline)
// Everything inside the loop runs under the deadline rather than the
// test's own context. An attempt started just before it would otherwise
// run well past it: the chat container is capped at 90s of its own and the
// row lookup at another 20s, so the loop could report a repricing failure
// nearly two minutes after the window it was given had closed.
repriceCtx, cancelReprice := context.WithDeadline(ctx, deadline)
defer cancelReprice()
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
lastSession = fmt.Sprintf("e2e-session-reprice-b-%d", time.Now().UnixNano())
code, _, cerr := env.client.Chat(repriceCtx, env.endpoint, env.proxyIP, harness.WireChat, customModel, "Reply with exactly: pong", lastSession)
code, _, cerr := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, customModel, "Reply with exactly: pong", lastSession)
if cerr != nil || code != 200 {
if !waitBeforeRetry(repriceCtx, 5*time.Second) {
break
}
continue
}
row, ok := lookupAccessLogBySession(repriceCtx, lastSession, repriceIngestWindow)
if !ok {
// No row for this request. The proxy now publishes a rebuilt chain
// before the route that reaches it, so a request can no longer be
// served unattributed mid-update; this retry covers the ingest
// window alone. Fire another one under a fresh session.
t.Logf("no access-log row for session %q within %s; retrying under a fresh session", lastSession, repriceIngestWindow)
time.Sleep(5 * time.Second)
continue
}
row := findAccessLogBySession(t, ctx, lastSession)
if inDelta(row.InputCostUsd, wantInputB, 1e-6) {
repriced = row
break
}
// Still priced at the old rate — the push hasn't landed yet; retry.
lastCost, sawRow = row.InputCostUsd, true
if !waitBeforeRetry(repriceCtx, 5*time.Second) {
break
}
time.Sleep(5 * time.Second)
}
lastSeen := "no row was ever read"
if sawRow {
lastSeen = fmt.Sprintf("last input_cost_usd=$%.6f", lastCost)
}
require.NotEmpty(t, repriced.Id, "a request after the price change must be priced at the new rate B; %s, wanted $%.6f\n=== proxy logs ===\n%s",
lastSeen, wantInputB, env.proxy.Logs(context.Background()))
require.NotEmpty(t, repriced.Id, "a request after the price change must be priced at the new rate B; last input_cost_usd=$%.6f, wanted $%.6f\n=== proxy logs ===\n%s",
repriced.InputCostUsd, wantInputB, env.proxy.Logs(context.Background()))
assertOpenAICostAtRates(t, repriced, inRateB, outRateB)
verifyUsageRowForSession(t, lastSession, inRateB, outRateB)
@@ -730,47 +630,3 @@ 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")
})
}

View File

@@ -1,447 +0,0 @@
//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 listing is served by the CONTROL PLANE (bedrock.<region>), not the
// runtime host a provider record must point at for InvokeModel — the
// runtime host answers <UnknownOperationException/>. The router now sends
// the listing, and only the listing, to the control plane, so this case
// asserts a real filtered listing rather than the 404 it used to get.
//
// The mock upstream cannot show any of this: it answers
// /inference-profiles on the same listener as everything else, so a
// mock-based test passes whichever host the request went to.
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",
// Registered verbatim, as an operator would copy it from AWS: the
// region prefix is what makes the id invocable, and the listing
// returns ids in exactly this form.
models: []string{model},
outcome: outcomeFiltered,
permitted: []string{model},
})
}
// 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)
// Status only, not the body. A Bedrock listing embeds inference-profile
// ARNs carrying the 12-digit AWS account id, and these job logs are
// readable by anyone who can see the run. The ids line below is the finding
// anyway. The failure paths below are the same log: a listing that fails to
// arrive is an AWS refusal naming the resource it refused, and that name is
// an ARN carrying the same account id.
t.Logf("[discovery] %s GET %s -> %d", tc.name, tc.path, code)
require.Equal(t, 200, code, "%s discovery must be served; response was %s", tc.name, bodyShape(body))
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; response was %s",
tc.name, bodyShape(body))
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]
_, dated := permitted[sharedllm.NormalizeAnthropicModel(id)]
// Bedrock ids carry a region prefix and version suffix the record may
// not repeat; the proxy's filter tries the same forms.
_, bedrock := permitted[sharedllm.NormalizeBedrockModel(id)]
assert.Truef(t, direct || dated || bedrock,
"%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 neither envelope the proxy's filter recognises — the two must
// stay in step, or this test reports "not a listing" for a response the proxy
// filtered perfectly well.
func listingIDs(body string) ([]string, bool) {
var doc struct {
// OpenAI's shape, which Anthropic adopted.
Data []struct {
ID string `json:"id"`
} `json:"data"`
// Bedrock returns inference-profile summaries under a key of its own,
// with the id under a field of its own.
Summaries []struct {
ID string `json:"inferenceProfileId"`
} `json:"inferenceProfileSummaries"`
}
if err := json.Unmarshal([]byte(body), &doc); err != nil {
return nil, false
}
switch {
case doc.Data != nil:
ids := make([]string, 0, len(doc.Data))
for _, entry := range doc.Data {
ids = append(ids, entry.ID)
}
return ids, true
case doc.Summaries != nil:
ids := make([]string, 0, len(doc.Summaries))
for _, entry := range doc.Summaries {
ids = append(ids, entry.ID)
}
return ids, true
}
return nil, false
}
func caseNames(cases []liveDiscoveryCase) []string {
names := make([]string, 0, len(cases))
for _, c := range cases {
names = append(names, c.name)
}
return names
}
// bodyShape describes a response without quoting any of it: its size and the
// top-level keys it arrived under. That is what a discovery failure is
// diagnosed from — which envelope the vendor answered with — and it is all
// that may go in a message rendered into a public job log, because the values
// underneath can carry an ARN and its account id.
func bodyShape(body string) string {
var doc map[string]json.RawMessage
if err := json.Unmarshal([]byte(body), &doc); err != nil {
return strconv.Itoa(len(body)) + " bytes, not a JSON object"
}
keys := make([]string, 0, len(doc))
for key := range doc {
keys = append(keys, key)
}
sort.Strings(keys)
if len(keys) == 0 {
return strconv.Itoa(len(body)) + " bytes, an empty JSON object"
}
return strconv.Itoa(len(body)) + " bytes, keyed by: " + strings.Join(keys, ", ")
}
// 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)"
}

View File

@@ -1,170 +0,0 @@
//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")
_, err = cl.ResolveProxyIP(ctx, endpoint)
require.NoError(t, err, "second client could not resolve the endpoint")
// Guarded rather than passed straight to require: px.Logs pulls the whole
// proxy container log, which is only worth fetching when the wait failed.
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
require.NoError(t, err, "second client did not see the proxy peer\n=== proxy logs ===\n%s",
px.Logs(context.Background()))
}
return cl
}

View File

@@ -1,455 +0,0 @@
//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())
})
}

View File

@@ -1,242 +0,0 @@
//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,
}
}

View File

@@ -54,19 +54,3 @@ func run(m *testing.M) int {
return m.Run()
}
// waitBeforeRetry pauses between attempts of a polling loop and reports
// whether the caller should keep going. A cancelled context ends the loop
// where a plain sleep would keep retrying against it: every call fails
// instantly once ctx is done, so the loop would spend its whole remaining
// window sleeping between failures nobody is waiting for any more.
func waitBeforeRetry(ctx context.Context, d time.Duration) bool {
timer := time.NewTimer(d)
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}

View File

@@ -1,209 +0,0 @@
//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
wantCacheRead := float64(harness.VLLMStreamCacheReadTokens) / 1000 * streamCacheReadRate
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")
// The total, not merely a positive number: input and output alone are
// positive, so a cache bucket parsed and then never billed would pass any
// weaker assertion. The gap is 7e-6, well outside the delta.
assert.InDelta(t, wantInput+wantOutput+wantCacheRead, row.CostUsd, 1e-6,
"the recorded cost must be every bucket the surface bills, cache reads included")
}
// 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")
// Deleting the group does not delete the key it auto-joins, so the key
// needs a cleanup of its own.
t.Cleanup(func() { _ = srv.API().SetupKeys.Delete(context.Background(), sk.Id) })
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
}
}
if !waitBeforeRetry(ctx, 5*time.Second) {
break
}
}
if code != 200 {
t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background()))
}
return code, body
}

View File

@@ -7,7 +7,6 @@ import (
"errors"
"fmt"
"io"
"net/http"
"os/exec"
"strconv"
"strings"
@@ -200,18 +199,12 @@ 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
// 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
// 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
)
// 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.
// 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.
func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
args := []string{
"run", "--rm",
@@ -222,7 +215,7 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string,
"-w", "%{remote_ip}",
"https://" + endpoint + "/",
}
deadline := time.Now().Add(endpointProbeRetryWindow)
deadline := time.Now().Add(dnsProbeRetryWindow)
for {
cmd := exec.CommandContext(ctx, "docker", args...)
var stdout, stderr strings.Builder
@@ -238,29 +231,21 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string,
}
var exitErr *exec.ExitError
if !errors.As(err, &exitErr) || !isTransientProbeExit(exitErr.ExitCode()) {
if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve {
return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String()))
}
probeErr := fmt.Errorf("endpoint %s not reachable yet: %s", endpoint, strings.TrimSpace(stderr.String()))
if time.Until(deadline) < endpointProbeRetryInterval {
return "", probeErr
dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String()))
if time.Until(deadline) < dnsProbeRetryInterval {
return "", dnsErr
}
select {
case <-ctx.Done():
return "", fmt.Errorf("%w (%w)", probeErr, ctx.Err())
case <-time.After(endpointProbeRetryInterval):
return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err())
case <-time.After(dnsProbeRetryInterval):
}
}
}
// 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.
@@ -307,27 +292,6 @@ 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
@@ -358,29 +322,10 @@ func withSessionID(headers []string, sessionID string) []string {
return append(headers, "x-session-id: "+sessionID)
}
// 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
// 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. 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) {
// to the proxy IP. It returns the HTTP status and response body.
func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) {
url := "https://" + endpoint + path
args := []string{
"run", "--rm",
@@ -389,15 +334,13 @@ func (cl *Client) do(ctx context.Context, method, endpoint, proxyIP, path, body
"-sk", "--connect-timeout", "5", "--max-time", "90",
"--resolve", endpoint + ":443:" + proxyIP,
"-o", "/dev/stderr", "-w", "%{http_code}",
"-X", method, url,
"-X", "POST", url,
"-H", "Content-Type: application/json",
}
for _, h := range extraHeaders {
args = append(args, "-H", h)
}
if body != "" {
args = append(args, "--data", 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.

View File

@@ -18,63 +18,18 @@ 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), 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
// 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
// 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 {
@@ -82,75 +37,13 @@ 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"},{"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"}]}';
return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"}]}';
}
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]
';
}
}
}
`
@@ -162,10 +55,6 @@ 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.
@@ -184,17 +73,14 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) {
req := testcontainers.ContainerRequest{
Image: vllmImage,
ExposedPorts: []string{vllmPort, vllmStreamPort},
ExposedPorts: []string{vllmPort},
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.ForAll(
wait.ForListeningPort(vllmPort),
wait.ForListeningPort(vllmStreamPort),
).WithStartupTimeout(60 * time.Second),
WaitingFor: wait.ForListeningPort(vllmPort).WithStartupTimeout(60 * time.Second),
}
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
@@ -206,12 +92,7 @@ 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",
StreamURL: "http://" + vllmAlias + ":8001",
}, nil
return &VLLM{container: ctr, workDir: workDir, URL: "http://" + vllmAlias + ":8000"}, nil
}
// Logs returns the vLLM container logs, for diagnostics on failure.

2
go.mod
View File

@@ -73,6 +73,7 @@ require (
github.com/hashicorp/go-version v1.7.0
github.com/jackc/pgx/v5 v5.5.5
github.com/libdns/route53 v1.5.0
github.com/libp2p/go-nat v0.2.0
github.com/libp2p/go-netroute v0.4.0
github.com/lrh3321/ipset-go v0.0.0-20250619021614-54a0a98ace81
github.com/mdlayher/socket v0.5.1
@@ -80,7 +81,6 @@ require (
github.com/miekg/dns v1.1.72
github.com/mitchellh/hashstructure/v2 v2.0.2
github.com/moby/moby/api v1.54.1
github.com/netbirdio/go-nat v0.0.0-20260821095157-6b2c8c5c74e8
github.com/netbirdio/management-integrations/integrations v0.0.0-20260416123949-2355d972be42
github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45
github.com/oapi-codegen/runtime v1.1.2

4
go.sum
View File

@@ -407,6 +407,8 @@ github.com/libdns/libdns v0.2.2 h1:O6ws7bAfRPaBsgAYt8MDe2HcNBGC29hkZ9MX2eUSX3s=
github.com/libdns/libdns v0.2.2/go.mod h1:4Bj9+5CQiNMVGf87wjX4CY3HQJypUHRuLvlsfsZqLWQ=
github.com/libdns/route53 v1.5.0 h1:2SKdpPFl/qgWsXQvsLNJJAoX7rSxlk7zgoL4jnWdXVA=
github.com/libdns/route53 v1.5.0/go.mod h1:joT4hKmaTNKHEwb7GmZ65eoDz1whTu7KKYPS8ZqIh6Q=
github.com/libp2p/go-nat v0.2.0 h1:Tyz+bUFAYqGyJ/ppPPymMGbIgNRH+WqC5QrT5fKrrGk=
github.com/libp2p/go-nat v0.2.0/go.mod h1:3MJr+GRpRkyT65EpVPBstXLvOlAPzUVlG6Pwg9ohLJk=
github.com/libp2p/go-netroute v0.4.0 h1:sZZx9hyANYUx9PZyqcgE/E1GUG3iEtTZHUEvdtXT7/Q=
github.com/libp2p/go-netroute v0.4.0/go.mod h1:Nkd5ShYgSMS5MUKy/MU2T57xFoOKvvLR92Lic48LEyA=
github.com/lrh3321/ipset-go v0.0.0-20250619021614-54a0a98ace81 h1:J56rFEfUTFT9j9CiRXhi1r8lUJ4W5idG3CiaBZGojNU=
@@ -478,8 +480,6 @@ github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1 h1:neE7z+FPUk
github.com/netbirdio/dex/api/v2 v2.0.0-20260512110716-8d70ad8647c1/go.mod h1:awuTyT29CYALpEyET0S307EgNlPWrc7fFKRAyhsO45M=
github.com/netbirdio/easyjson v0.9.0 h1:6Nw2lghSVuy8RSkAYDhDv1thBVEmfVbKZnV7T7Z6Aus=
github.com/netbirdio/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=
github.com/netbirdio/go-nat v0.0.0-20260821095157-6b2c8c5c74e8 h1:pBxXEsxcsO3qVUND//5j1kelYlO57x5IrRviNF0+0iA=
github.com/netbirdio/go-nat v0.0.0-20260821095157-6b2c8c5c74e8/go.mod h1:mFViabv4PpnoDw9w7W21a7xux6APA4q7KQZRsv4BCl8=
github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51 h1:Ov4qdafATOgGMB1wbSuh+0aAHcwz9hdvB6VZjh1mVMI=
github.com/netbirdio/ice/v4 v4.0.0-20250908184934-6202be846b51/go.mod h1:ZSIbPdBn5hePO8CpF1PekH2SfpTxg1PDhEwtbqZS7R8=
github.com/netbirdio/management-integrations/integrations v0.0.0-20260416123949-2355d972be42 h1:F3zS5fT9xzD1OFLfcdAE+3FfyiwjGukF1hvj0jErgs8=

View File

@@ -15,12 +15,6 @@ set -o pipefail
# 2. Postgres migration — add Postgres, migrate SQLite data via migrate-store.
# 3. Traffic flow — add NATS + flow-enricher + flow-receiver.
#
# Step 2 is skipped when the deployment already runs on Postgres
# (server.store.engine: postgres in config.yaml). Nothing is provisioned or
# migrated in that case and the store config is left exactly as the operator
# wrote it — the enterprise image reads the same Postgres the community image
# did. Such a deployment gets the image swap, and can still opt into step 3.
#
# If any step fails once the stack has been touched, the script rolls itself
# back automatically: generated files are removed, the Postgres volume this run
# created is dropped, and the original deployment is started again.
@@ -44,18 +38,6 @@ ENV_BACKUP=""
PG_VOLUME_NAME=""
BACKUP_DIR=""
# Store state. STORE_ENGINE is what the deployment runs on today; when it is
# already postgres, MIGRATE_POSTGRES stays "no" and nothing is provisioned.
# POSTGRES_SERVICE is empty when Postgres lives outside this compose project.
STORE_ENGINE=""
EXISTING_POSTGRES="no"
POSTGRES_DSN=""
POSTGRES_SERVICE=""
POSTGRES_DEPENDS_CONDITION="service_healthy"
# Whether this run needs to generate config.yaml.enterprise at all. A pure
# image swap does not.
ENTERPRISE_CONFIG="no"
NETBIRD_EULA_URL="https://netbird.io/self-hosted-EULA"
check_docker_compose() {
@@ -210,85 +192,6 @@ detect_exposed_address() {
yq eval '.server.exposedAddress // ""' "$CONFIG_YAML_HOST"
}
# The engine is a config.yaml-only setting — there is no env override for it
# (combined/cmd/root.go reads it from YAML and derives the env vars), so
# config.yaml is authoritative. Absent means the sqlite default.
detect_store_engine() {
local engine
engine=$(yq eval '.server.store.engine // ""' "$CONFIG_YAML_HOST")
if [[ -z "$engine" ]] || [[ "$engine" == "null" ]]; then
engine="sqlite"
fi
echo "$engine" | tr '[:upper:]' '[:lower:]'
}
detect_store_dsn() {
yq eval '.server.store.dsn // ""' "$CONFIG_YAML_HOST"
}
# config.yaml is where a combined deployment carries its DSN; this only covers
# hand-rolled installs that keep it in the environment instead.
detect_store_dsn_from_compose() {
# `compose config` re-escapes a literal $ as $$ on the way out, so undo that
# to get the value the container actually receives.
$DOCKER_COMPOSE_COMMAND config 2>/dev/null | yq eval "
.services[\"$COMBINED_SERVICE\"].environment.NB_STORE_ENGINE_POSTGRES_DSN //
.services[\"$COMBINED_SERVICE\"].environment.NETBIRD_STORE_ENGINE_POSTGRES_DSN // \"\"
" - 2>/dev/null | sed 's/\$\$/$/g'
}
# Reads either DSN form: "host=db ..." or "postgres://user:pass@db:5432/name".
dsn_host() {
local dsn="$1"
case "$dsn" in
*://*) printf '%s' "$dsn" | sed -n 's,^[a-zA-Z+]*://\([^/?]*\).*,\1,p' | sed -e 's,.*@,,' -e 's,:.*,,' ;;
*) printf '%s' "$dsn" | sed -n 's/.*[[:space:]]*host=\([^[:space:]]*\).*/\1/p' ;;
esac
}
# flow-enricher is its own container, so a loopback host or a socket path would
# reach the enricher rather than Postgres. Only flag hosts we can positively
# identify — an unparseable DSN must not leave the operator with no way forward.
dsn_host_reachable() {
local dsn="$1"
case "$(dsn_host "$dsn")" in
localhost | 127.* | ::1 | 0.0.0.0 | /*) return 1 ;;
*) return 0 ;;
esac
}
# Names the compose service running this deployment's Postgres, for depends_on.
# Empty means external — the DSN host matched no service. A DSN with no readable
# host falls back to matching on image.
detect_postgres_service() {
local host
host=$(dsn_host "$POSTGRES_DSN")
if [[ -n "$host" ]]; then
if [[ "$(host="$host" yq eval '.services | has(env(host))' "$COMPOSE_FILE" 2>/dev/null)" == "true" ]]; then
echo "$host"
fi
return
fi
yq eval '.services | to_entries | map(select(.value.image // "" | test("(^|/)(postgres|postgis|pgvector|timescaledb)(:|@|$)"))) | .[0].key // ""' "$COMPOSE_FILE"
}
# depends_on: service_healthy is only legal if the service defines a healthcheck.
detect_postgres_depends_condition() {
local tag
tag=$(yq eval ".services[\"$POSTGRES_SERVICE\"].healthcheck | tag" "$COMPOSE_FILE" 2>/dev/null)
if [[ "$tag" == "!!map" ]]; then
echo "service_healthy"
else
echo "service_started"
fi
}
env_value() {
local value="$1"
value=$(printf '%s' "$value" | sed -e 's/\\/\\\\/g' -e 's/"/\\"/g' -e 's/\$/$$/g')
printf '"%s"' "$value"
}
detect_compose_network() {
local tag
tag=$(yq eval ".services[\"$COMBINED_SERVICE\"].networks | tag" "$COMPOSE_FILE" 2>/dev/null)
@@ -325,30 +228,16 @@ services:
NETBIRD_LICENSE_SERVER_BASE_URL: \${NETBIRD_LICENSE_SERVER_BASE_URL}
EOF
# An existing Postgres is already wired up by the operator's own compose file,
# so only a Postgres this run creates needs a depends_on.
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
cat <<EOF
depends_on:
${POSTGRES_SERVICE}:
condition: ${POSTGRES_DEPENDS_CONDITION}
EOF
fi
# The server is only pointed at a different config file when this run
# generates one. A pure image swap leaves it on its original config.yaml.
if [[ "$ENTERPRISE_CONFIG" == "yes" ]]; then
cat <<EOF
postgres:
condition: service_healthy
volumes:
- ./${ENTERPRISE_CONFIG_FILE}:/etc/netbird/config.yaml.enterprise:ro
command: ["--config", "/etc/netbird/config.yaml.enterprise"]
EOF
fi
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
cat <<EOF
${POSTGRES_SERVICE}:
postgres:
image: postgres:17
container_name: netbird-postgres
restart: unless-stopped
@@ -368,14 +257,6 @@ EOF
fi
if [[ "$ENABLE_FLOW" == "yes" ]]; then
# Nothing to wait on when Postgres is managed outside this compose project.
local enricher_depends=""
if [[ -n "$POSTGRES_SERVICE" ]]; then
enricher_depends="
${POSTGRES_SERVICE}:
condition: ${POSTGRES_DEPENDS_CONDITION}"
fi
cat <<EOF
nats:
@@ -392,7 +273,9 @@ EOF
container_name: netbird-flow-enricher
restart: unless-stopped
networks: [${COMPOSE_NETWORK}]
depends_on:${enricher_depends}
depends_on:
postgres:
condition: service_healthy
nats:
condition: service_started
environment:
@@ -400,10 +283,10 @@ EOF
NETBIRD_LICENSE_SERVER_BASE_URL: \${NETBIRD_LICENSE_SERVER_BASE_URL}
NB_DATADIR: /var/lib/netbird
NB_MANAGEMENT_STORE_ENGINE: postgres
NB_MANAGEMENT_POSTGRES_DSN: \${NB_ENTERPRISE_POSTGRES_DSN}
NB_STORE_ENGINE_POSTGRES_DSN: \${NB_ENTERPRISE_POSTGRES_DSN}
NB_MANAGEMENT_POSTGRES_DSN: "host=postgres user=netbird password=\${POSTGRES_PASSWORD} dbname=netbird port=5432 sslmode=disable"
NB_STORE_ENGINE_POSTGRES_DSN: "host=postgres user=netbird password=\${POSTGRES_PASSWORD} dbname=netbird port=5432 sslmode=disable"
NB_TRAFFIC_EVENT_STORE_ENGINE: postgres
NB_TRAFFIC_EVENT_POSTGRES_DSN: \${NB_ENTERPRISE_POSTGRES_DSN}
NB_TRAFFIC_EVENT_POSTGRES_DSN: "host=postgres user=netbird password=\${POSTGRES_PASSWORD} dbname=netbird port=5432 sslmode=disable"
NB_MANAGEMENT_STORE_KEY: \${NETBIRD_ENCRYPTION_KEY}
NB_FLOW_ADAPTER_TYPE: nats
NB_FLOW_NATS_ENDPOINTS: nats://nats:4222
@@ -460,41 +343,27 @@ EOF
fi
}
# Build config.yaml.enterprise from the operator's existing config.yaml. We
# don't touch the original file. Values go through strenv() so a DSN carrying
# quotes, backslashes or $ cannot break out of the expression.
# Build config.yaml.enterprise by yq-editing the operator's existing
# config.yaml. We don't touch the original file.
render_enterprise_config() {
{
echo "# Generated by migrate-to-enterprise.sh from ${CONFIG_YAML_HOST}."
echo "# The enterprise server is started with --config pointing at this file,"
echo "# so later edits to ${CONFIG_YAML_HOST} have no effect until copied here."
cat "$CONFIG_YAML_HOST"
} > "$ENTERPRISE_CONFIG_FILE"
local pg_dsn="host=postgres user=netbird password=${POSTGRES_PASSWORD} dbname=netbird port=5432 sslmode=disable"
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
# Fresh Postgres: point every store section at it. migrate-store carries the
# SQLite contents across.
POSTGRES_DSN="$POSTGRES_DSN" yq eval -i '
.server.store.engine = "postgres" |
.server.store.dsn = strenv(POSTGRES_DSN) |
.server.activityStore.engine = "postgres" |
.server.activityStore.dsn = strenv(POSTGRES_DSN) |
.server.authStore.engine = "postgres" |
.server.authStore.dsn = strenv(POSTGRES_DSN)
' "$ENTERPRISE_CONFIG_FILE"
fi
# Otherwise the store config is the operator's and stays untouched.
# activityStore and authStore do not inherit from server.store — each falls
# back to its own SQLite file under dataDir — so repointing them at Postgres
# here would silently strand the existing audit log and the embedded IdP's
# users, with no migrate-store run to carry them over.
yq eval "
.server.store.engine = \"postgres\" |
.server.store.dsn = \"$pg_dsn\" |
.server.activityStore.engine = \"postgres\" |
.server.activityStore.dsn = \"$pg_dsn\" |
.server.authStore.engine = \"postgres\" |
.server.authStore.dsn = \"$pg_dsn\"
" "$CONFIG_YAML_HOST" > "$ENTERPRISE_CONFIG_FILE"
if [[ "$ENABLE_FLOW" == "yes" ]]; then
NETBIRD_DOMAIN="$NETBIRD_DOMAIN" yq eval -i '
local flow_addr="${NETBIRD_DOMAIN}"
yq eval -i "
.server.trafficFlow.enabled = true |
.server.trafficFlow.address = strenv(NETBIRD_DOMAIN) |
.server.trafficFlow.interval = "60s"
' "$ENTERPRISE_CONFIG_FILE"
.server.trafficFlow.address = \"$flow_addr\" |
.server.trafficFlow.interval = \"60s\"
" "$ENTERPRISE_CONFIG_FILE"
fi
}
@@ -761,91 +630,6 @@ on_exit() {
# Main
# ---------------------------------------------------------------------------
# Already on Postgres: there is nothing to provision and nothing to migrate.
# The enterprise image reads the very same store config the community image
# did, so step 2 collapses to a no-op and the run is a plain image swap.
configure_existing_postgres() {
EXISTING_POSTGRES="yes"
MIGRATE_POSTGRES="no"
# DSN first — detect_postgres_service prefers the host it names.
POSTGRES_DSN=$(detect_store_dsn)
if [[ -z "$POSTGRES_DSN" ]] || [[ "$POSTGRES_DSN" == "null" ]]; then
POSTGRES_DSN=$(detect_store_dsn_from_compose)
fi
if [[ "$POSTGRES_DSN" == "null" ]]; then
POSTGRES_DSN=""
fi
POSTGRES_SERVICE=$(detect_postgres_service)
if [[ -n "$POSTGRES_SERVICE" ]]; then
POSTGRES_DEPENDS_CONDITION=$(detect_postgres_depends_condition)
fi
echo "Step 2: Postgres migration not needed — this deployment already runs on"
echo " Postgres. Its store configuration is reused as-is and left"
echo " untouched; no database is created and no data is moved."
if [[ -n "$POSTGRES_SERVICE" ]]; then
echo " Postgres service: $POSTGRES_SERVICE (in $COMPOSE_FILE)"
else
echo " Postgres service: managed outside $COMPOSE_FILE"
fi
}
configure_sqlite_store() {
MIGRATE_POSTGRES=$(read_yes_no "Step 2: Migrate storage from SQLite to Postgres? (recommended)" "n")
[[ "$MIGRATE_POSTGRES" == "yes" ]] || return 0
# The override would otherwise merge into a service of the same name and
# quietly rewrite its image and credentials.
local existing
existing=$(yq eval '.services | has("postgres")' "$COMPOSE_FILE")
if [[ "$existing" == "true" ]]; then
echo "" > /dev/stderr
echo "$COMPOSE_FILE already defines a service named 'postgres', but config.yaml" > /dev/stderr
echo "still has server.store.engine: sqlite. This script would add its own" > /dev/stderr
echo "'postgres' service and Compose would merge the two." > /dev/stderr
echo "" > /dev/stderr
echo "Point server.store.engine at that Postgres yourself, or rename the service," > /dev/stderr
echo "then re-run." > /dev/stderr
exit 1
fi
echo ""
echo " ⚠ Data will be migrated from SQLite to Postgres. The SQLite store"
echo " will be backed up automatically. To fully revert later, restore"
echo " that backup and delete docker-compose.override.yml +"
echo " config.yaml.enterprise."
local confirm
confirm=$(read_yes_no " Continue?" "y")
if [[ "$confirm" != "yes" ]]; then
MIGRATE_POSTGRES="no"
echo " Skipping Postgres migration."
return 0
fi
POSTGRES_PASSWORD=$(rand_password)
POSTGRES_SERVICE="postgres"
POSTGRES_DEPENDS_CONDITION="service_healthy"
POSTGRES_DSN="host=postgres user=netbird password=${POSTGRES_PASSWORD} dbname=netbird port=5432 sslmode=disable"
}
# mysql, or something this script has never seen. Swapping the images is still
# valid; touching the store is not.
configure_unsupported_store() {
MIGRATE_POSTGRES="no"
echo " ⚠ server.store.engine is '$STORE_ENGINE'. This script only migrates"
echo " SQLite to Postgres, and traffic flow requires Postgres, so both are"
echo " unavailable here. The store configuration will be left untouched."
echo ""
local proceed
proceed=$(read_yes_no "Step 2 skipped. Continue with the image swap only?" "n")
if [[ "$proceed" != "yes" ]]; then
echo "Aborted."
exit 0
fi
}
init_migration() {
DOCKER_COMPOSE_COMMAND=$(check_docker_compose)
check_yq
@@ -895,15 +679,12 @@ init_migration() {
exit 1
fi
STORE_ENGINE=$(detect_store_engine)
echo "Detected existing deployment:"
echo " Combined service: $COMBINED_SERVICE"
echo " Dashboard: $DASHBOARD_SERVICE"
echo " config.yaml: $CONFIG_YAML_HOST"
echo " Data volume: $DATA_VOLUME"
echo " Network: $COMPOSE_NETWORK"
echo " Store engine: $STORE_ENGINE"
echo ""
require_eula_acceptance
@@ -922,17 +703,28 @@ init_migration() {
echo "Step 1: Image swap (community → Enterprise). License key required."
NB_LICENSE_KEY=$(read_secret " License key")
# Step 2 — what this does depends on what the deployment already stores in.
# Step 2 — optional
echo ""
case "$STORE_ENGINE" in
postgres) configure_existing_postgres ;;
sqlite) configure_sqlite_store ;;
*) configure_unsupported_store ;;
esac
MIGRATE_POSTGRES=$(read_yes_no "Step 2: Migrate storage from SQLite to Postgres? (recommended)" "n")
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
echo ""
echo " ⚠ Data will be migrated from SQLite to Postgres. The SQLite store"
echo " will be backed up automatically. To fully revert later, restore"
echo " that backup and delete docker-compose.override.yml +"
echo " config.yaml.enterprise."
local confirm
confirm=$(read_yes_no " Continue?" "y")
if [[ "$confirm" != "yes" ]]; then
MIGRATE_POSTGRES="no"
echo " Skipping Postgres migration."
else
POSTGRES_PASSWORD=$(rand_password)
fi
fi
# Step 3 — optional, only if Postgres is on (flow requires Postgres)
echo ""
if [[ "$MIGRATE_POSTGRES" == "yes" ]] || [[ "$EXISTING_POSTGRES" == "yes" ]]; then
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
ENABLE_FLOW=$(read_yes_no "Step 3: Enable traffic flow? (requires Postgres)" "n")
if [[ "$ENABLE_FLOW" == "yes" ]]; then
# Auth secret MUST match server.authSecret from config.yaml
@@ -956,46 +748,12 @@ init_migration() {
echo "Could not read server.store.encryptionKey from $CONFIG_YAML_HOST." > /dev/stderr
exit 1
fi
# flow-enricher talks to Postgres directly, so this is the one place an
# existing deployment's DSN is actually needed — and the one place a host
# that only works from inside the server container shows up.
while :; do
local dsn_problem=""
if [[ -z "$POSTGRES_DSN" ]]; then
dsn_problem="No DSN could be read from $CONFIG_YAML_HOST or from the $COMBINED_SERVICE environment."
elif ! dsn_host_reachable "$POSTGRES_DSN"; then
dsn_problem="Its host '$(dsn_host "$POSTGRES_DSN")' only resolves inside the server container."
fi
[[ -n "$dsn_problem" ]] || break
echo ""
echo " The flow enricher reaches Postgres from a container of its own."
echo " $dsn_problem"
echo " Enter a DSN reachable from other containers, or press Ctrl-C to abort."
POSTGRES_DSN=$(read_required " Postgres DSN (host=… user=… password=… dbname=… port=5432 sslmode=disable)")
done
# Only where the operator owns Postgres: a DSN entered above may name a
# different host. The sqlite path creates its own service, nothing to find.
if [[ "$EXISTING_POSTGRES" == "yes" ]]; then
POSTGRES_SERVICE=$(detect_postgres_service)
if [[ -n "$POSTGRES_SERVICE" ]]; then
POSTGRES_DEPENDS_CONDITION=$(detect_postgres_depends_condition)
fi
fi
fi
else
ENABLE_FLOW="no"
echo "Step 3 (traffic flow) skipped — requires Postgres."
fi
# config.yaml.enterprise only exists to hold changes; without any there is
# nothing to generate and the server keeps running on its own config.yaml.
if [[ "$MIGRATE_POSTGRES" == "yes" ]] || [[ "$ENABLE_FLOW" == "yes" ]]; then
ENTERPRISE_CONFIG="yes"
fi
check_data_directory
check_stale_postgres_volume
}
@@ -1013,7 +771,7 @@ apply_changes() {
sed -i.bak '/NETBIRD_LICENSE_SERVER_BASE_URL/d' "$OVERRIDE_FILE" && rm -f "$OVERRIDE_FILE.bak"
fi
if [[ "$ENTERPRISE_CONFIG" == "yes" ]]; then
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
echo "Writing $ENTERPRISE_CONFIG_FILE ..."
install -m 600 /dev/null "$ENTERPRISE_CONFIG_FILE"
render_enterprise_config
@@ -1049,9 +807,6 @@ apply_changes() {
echo "POSTGRES_PASSWORD=${POSTGRES_PASSWORD}"
fi
if [[ "$ENABLE_FLOW" == "yes" ]]; then
# Own variable name rather than NB_STORE_ENGINE_POSTGRES_DSN so that a
# deployment already setting that one keeps its own value.
echo "NB_ENTERPRISE_POSTGRES_DSN=$(env_value "$POSTGRES_DSN")"
echo "NB_FLOW_AUTH_SECRET=${NB_FLOW_AUTH_SECRET}"
echo "NETBIRD_ENCRYPTION_KEY=${NETBIRD_ENCRYPTION_KEY}"
fi
@@ -1113,19 +868,14 @@ print_summary() {
echo " Summary"
echo "──────────────────────────────────────────────────────────────────────"
echo " Images: swapped to enterprise"
if [[ "$MIGRATE_POSTGRES" == "yes" ]]; then
echo " Storage: Postgres (data migrated from SQLite)"
elif [[ "$EXISTING_POSTGRES" == "yes" ]]; then
echo " Storage: Postgres (pre-existing, configuration unchanged)"
else
echo " Storage: $STORE_ENGINE (unchanged)"
fi
[[ "$MIGRATE_POSTGRES" == "yes" ]] && echo " Storage: Postgres (data migrated from SQLite)"
[[ "$MIGRATE_POSTGRES" != "yes" ]] && echo " Storage: SQLite (unchanged)"
[[ "$ENABLE_FLOW" == "yes" ]] && echo " Traffic flow: enabled"
[[ "$ENABLE_FLOW" != "yes" ]] && echo " Traffic flow: disabled"
echo ""
echo " Generated files (next to your docker-compose.yml):"
echo " $OVERRIDE_FILE"
[[ "$ENTERPRISE_CONFIG" == "yes" ]] && echo " $ENTERPRISE_CONFIG_FILE"
[[ "$MIGRATE_POSTGRES" == "yes" ]] && echo " $ENTERPRISE_CONFIG_FILE"
echo " .env (license key + secrets, mode 600)"
[[ "$ENV_EXISTED" == "yes" ]] && [[ -f "$ENV_BACKUP" ]] && echo " $ENV_BACKUP (.env as it was before this run)"
[[ "$MIGRATE_POSTGRES" == "yes" ]] && echo " backups/sqlite-pre-enterprise-*/ (SQLite backup)"
@@ -1149,11 +899,7 @@ print_summary() {
else
echo " $DOCKER_COMPOSE_COMMAND down"
fi
if [[ "$ENTERPRISE_CONFIG" == "yes" ]]; then
echo " rm -f $OVERRIDE_FILE $ENTERPRISE_CONFIG_FILE"
else
echo " rm -f $OVERRIDE_FILE"
fi
echo " rm -f $OVERRIDE_FILE $ENTERPRISE_CONFIG_FILE"
if [[ "$ENV_EXISTED" == "yes" ]] && [[ -f "$ENV_BACKUP" ]]; then
echo " mv $ENV_BACKUP .env # restores .env as it was before this run"
elif [[ "$ENV_EXISTED" == "no" ]]; then

View File

@@ -113,61 +113,8 @@ 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).
@@ -298,12 +245,8 @@ var providers = []Provider{
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#10A37F",
Discovery: &Discovery{
Path: "/v1/models",
Shape: ShapeOpenAIData,
},
ParserID: "openai",
PricingSurfaces: []string{"openai"},
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
@@ -341,18 +284,8 @@ var providers = []Provider{
AuthHeaderTemplate: "${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#D97757",
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"},
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
@@ -363,8 +296,6 @@ 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},
@@ -412,22 +343,6 @@ 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"},
@@ -440,8 +355,6 @@ 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},
@@ -478,15 +391,6 @@ 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.
@@ -502,8 +406,6 @@ 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},

View File

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

View File

@@ -1,178 +0,0 @@
package handlers
import (
"context"
"encoding/json"
"fmt"
"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"},
// A vendor that supplies no display name at all. Bedrock does for
// every profile, but the OpenAI listing carries none.
{ID: "gpt-4o-mini", PricingKnown: true},
}}
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, 3)
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)
require.NotNil(t, out.Models[0].Label, "the vendor supplied a display name")
assert.Equal(t, "EU Claude Haiku 4.5", *out.Models[0].Label)
// A vendor that supplies no name must omit the key rather than send an
// empty string: the dashboard falls back to the id on absence, and would
// render a blank row for "".
assert.Nil(t, out.Models[2].Label, "an absent label must not serialize")
assert.NotContains(t, rec.Body.String(), `"label":""`)
assert.Equal(t, "bedrock_api", stub.gotReq.CatalogID)
assert.Equal(t, "aws-bearer", stub.gotReq.APIKey)
// The upstream is what the region is read back out of for Bedrock, so
// losing it here would break discovery for every regional provider.
assert.Equal(t, "https://bedrock-runtime.eu-central-1.amazonaws.com", stub.gotReq.UpstreamURL)
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)
}
// TestDiscoverModelsTrimsTheCatalogID pins that the id the emptiness check
// accepts is the id the manager receives. A padded value that clears the check
// but reaches the catalog untrimmed misses the lookup, and the operator is told
// their provider does not exist.
func TestDiscoverModelsTrimsTheCatalogID(t *testing.T) {
stub := &discoveryManagerStub{}
rec := postDiscovery(t, stub, `{"catalog_provider_id":" openai_api ","api_key":"sk"}`)
require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String())
assert.Equal(t, "openai_api", stub.gotReq.CatalogID)
}
// TestDiscoverModelsReportsCallerInputAsBadRequest covers the other half of the
// error mapping. These failures are all reachable from a well-formed request
// with a bad field value, so answering 500 both misinforms the operator and
// puts their typo into the server's error rate.
func TestDiscoverModelsReportsCallerInputAsBadRequest(t *testing.T) {
stub := &discoveryManagerStub{
err: fmt.Errorf("%w: unknown catalog provider %q", modeldiscovery.ErrInvalidRequest, "nope"),
}
rec := postDiscovery(t, stub, `{"catalog_provider_id":"nope","api_key":"sk"}`)
assert.Equal(t, http.StatusBadRequest, rec.Code)
assert.Contains(t, rec.Body.String(), "unknown catalog provider")
}
func TestDiscoverModelsRejectsMalformedRequests(t *testing.T) {
for name, body := range map[string]string{
"not json": `{`,
"no catalog provider": `{"api_key":"sk"}`,
"blank catalog provider": `{"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)
})
}
}

View File

@@ -7,7 +7,6 @@ package handlers
import (
"encoding/json"
"errors"
"math"
"net/http"
"net/url"
@@ -17,7 +16,6 @@ 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"
@@ -34,7 +32,6 @@ 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")
@@ -64,98 +61,6 @@ 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
}
// Trimmed once and carried, not trimmed for the emptiness test and then
// discarded: a padded " openai_api " would clear the check here and miss
// the catalog lookup, reporting the provider as unknown.
catalogID := strings.TrimSpace(body.CatalogProviderId)
if catalogID == "" {
util.WriteErrorResponse("catalog_provider_id is required", http.StatusBadRequest, w)
return
}
recordID := strValue(body.ProviderId)
req := modeldiscovery.Request{
CatalogID: catalogID,
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
}
// An unknown provider, an unusable upstream, a missing region or a
// missing key are all things the caller sent, reachable from a
// well-formed request. Reporting them as 500 tells the operator the
// server broke and buries genuine faults in the error rate.
if errors.Is(err, modeldiscovery.ErrInvalidRequest) {
util.WriteErrorResponse(err.Error(), http.StatusBadRequest, 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,
// Sent even when zero: the form prefills every discovered model as
// an editable row, and an unpriced one is shown at zero and flagged
// rather than left out.
InputPer1k: m.InputPer1k,
OutputPer1k: m.OutputPer1k,
// Cache rates stay absent when unset, matching the catalog
// response — a zero would read as "free", not "not applicable".
CachedInputPer1k: positiveRatePtr(m.CachedInputPer1k),
CacheReadPer1k: positiveRatePtr(m.CacheReadPer1k),
CacheCreationPer1k: positiveRatePtr(m.CacheCreationPer1k),
}
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.

View File

@@ -13,7 +13,6 @@ 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"
@@ -51,7 +50,6 @@ 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)
@@ -125,15 +123,6 @@ 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.
//
// One instance serves every request for the process's lifetime, so its
// fields must stay read-only after construction: lazy initialisation
// inside Fetch or httpClient would race across request goroutines.
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.
@@ -162,7 +151,6 @@ 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())),
}
@@ -182,38 +170,6 @@ 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.
//
// 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. That one check also covers reading the stored
// record — Create is strictly stronger than Read here, and the lookup is
// scoped to accountID, so another account's record is never reachable.
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
@@ -1061,10 +1017,6 @@ 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
}

View File

@@ -1,469 +0,0 @@
// 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"
"syscall"
"time"
"golang.org/x/oauth2/google"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
)
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")
// ErrInvalidRequest marks a discovery failure caused by the caller's own input
// rather than by the vendor or by this server. Every one of these is reachable
// from a well-formed request carrying a bad field value, so the handler owes
// the caller a 400 — a 500 would both misinform them and bury real server
// faults in the error rate.
var ErrInvalidRequest = errors.New("invalid discovery request")
// 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
// The rates below are the defaults for this model, taken from the same
// table the proxy bills with, so the form prefills exactly what a request
// would cost. All zero when PricingKnown is false — an unpriced model is
// offered at zero and flagged, rather than withheld: the vendor says the
// credential can reach it, and refusing to show it would hide a model the
// operator genuinely has.
InputPer1k float64
OutputPer1k float64
CachedInputPer1k float64
CacheReadPer1k float64
CacheCreationPer1k float64
}
// 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("%w: unknown catalog provider %q", ErrInvalidRequest, 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("%w: provider upstream %q is not a usable URL", ErrInvalidRequest, 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("%w: %s discovery needs a region, and none could be read from the provider upstream",
ErrInvalidRequest, 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)
}
// The two halves must not overlap. "bedrock-runtime.amazonaws.com" carries
// both of Bedrock's — it is the regionless endpoint — and satisfies both
// checks above while leaving nothing between them, so slicing it would
// panic on an inverted range rather than report "no region here".
if !strings.HasPrefix(host, prefix) || !strings.HasSuffix(host, suffix) ||
len(host) < len(prefix)+len(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("%w: discovery host %q resolves to a non-public address", ErrInvalidRequest, 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("%w: %s discovery needs an API key", ErrInvalidRequest, 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, attaching
// the rates the request would actually be billed at.
//
// Rates come from the live default pricing table rather than the compiled-in
// catalog, because that is the table the synthesiser ships to the proxy: an
// operator running a defaults_llm_pricing.yaml would otherwise be shown one
// price in the form and charged another. It is also the same lookup the catalog
// endpoint prefills from, so a model reached by either route prices identically.
func decorate(entry catalog.Provider, ids []listedModel) []Model {
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 table keys pricing by the normalised id while the vendor issues
// the wire form, so normalise before looking it up — otherwise every
// Bedrock profile would report unpriced.
model := Model{ID: listed.id, Label: listed.label}
if rate, known := pricing.LookupDefault(entry.PricingSurfaces, normalizeForPricing(entry.ID, listed.id)); known {
model.PricingKnown = true
model.InputPer1k = rate.InputPer1k
model.OutputPer1k = rate.OutputPer1k
model.CachedInputPer1k = rate.CachedInputPer1k
model.CacheReadPer1k = rate.CacheReadPer1k
model.CacheCreationPer1k = rate.CacheCreationPer1k
}
out = append(out, model)
}
return out
}
// refuseRedirect is the redirect policy every discovery request runs under. A
// redirect is a way to move the request to a host checkPublicHost never saw,
// so none are followed.
func refuseRedirect(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
}
func (c *Client) httpClient() *http.Client {
if c.HTTPClient != nil {
if c.HTTPClient.CheckRedirect != nil {
return c.HTTPClient
}
// An injected client that states no policy still gets ours: the
// no-redirect guarantee should not depend on the caller remembering it.
//
// Copied rather than assigned into: one Client is shared by every
// request for the process's lifetime, so writing to its fields here
// would race across request goroutines. The copy shares the Transport,
// which is safe for concurrent use by design.
clone := *c.HTTPClient
clone.CheckRedirect = refuseRedirect
return &clone
}
transport := guardedTransport
if c.AllowPrivateHosts {
transport = http.DefaultTransport
}
return &http.Client{
Timeout: fetchTimeout,
Transport: transport,
CheckRedirect: refuseRedirect,
}
}
// guardedTransport dials only addresses isPublic accepts.
//
// checkPublicHost resolves the host itself, and the transport then resolves it
// again when it dials — two lookups of a name whose owner chooses the answers.
// A record that returns a public address to the first and 127.0.0.1 to the
// second passes the guard and reaches loopback anyway, which is the whole of
// DNS rebinding. Re-checking at the socket closes that window: whatever the
// second lookup returned is what Control is handed, and an address the guard
// refuses never gets connected.
//
// Shared package-wide rather than built per Fetch so connections and their
// pool survive between calls; the guard holds no state.
var guardedTransport = newGuardedTransport()
func newGuardedTransport() http.RoundTripper {
base, ok := http.DefaultTransport.(*http.Transport)
if !ok {
// Something replaced the default transport. Fall back to it rather
// than dropping its behaviour, and rely on checkPublicHost alone.
return http.DefaultTransport
}
// Cloned so proxy settings, TLS defaults and timeouts come from the
// standard transport rather than being restated here.
transport := base.Clone()
dialer := &net.Dialer{
Timeout: fetchTimeout,
KeepAlive: 30 * time.Second,
Control: func(_, address string, _ syscall.RawConn) error {
return guardDialAddress(address)
},
}
transport.DialContext = dialer.DialContext
return transport
}
// guardDialAddress refuses a resolved socket address the discovery client has
// no business connecting to. Control hands it over post-resolution and
// pre-connect, once per address the dialer tries, so a name with several A
// records is checked at each one.
func guardDialAddress(address string) error {
host, _, err := net.SplitHostPort(address)
if err != nil {
return fmt.Errorf("discovery dial address %q is unreadable", address)
}
addr, err := netip.ParseAddr(host)
if err != nil {
// Control is documented to receive a resolved address; anything else
// is a state we cannot vet, so it does not get dialled.
return fmt.Errorf("discovery dial address %q is not an IP", host)
}
if !isPublic(addr) {
return fmt.Errorf("discovery refused to dial non-public address %s", addr)
}
return nil
}

View File

@@ -1,532 +0,0 @@
package modeldiscovery
import (
"context"
"io"
"net/http"
"net/http/httptest"
"net/netip"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
)
// 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")
// The rates travel with the model, so the form can prefill an editable row
// rather than making the operator look every price up by hand.
assert.Positive(t, models[0].InputPer1k, "a priced model must carry its input rate")
assert.Positive(t, models[0].OutputPer1k, "a priced model must carry its output rate")
// An unpriced model is offered at zero and flagged, not withheld: the
// vendor says the credential can reach it.
assert.Zero(t, models[1].InputPer1k)
assert.Zero(t, models[1].OutputPer1k)
}
// TestDiscoveredRatesMatchTheCatalogEndpoint pins the two prefill paths to one
// table. The provider form fills a model row either from the catalog response
// or from a discovery response, and an operator who switches between them must
// not see the price change — both must equal what the proxy will bill.
func TestDiscoveredRatesMatchTheCatalogEndpoint(t *testing.T) {
cl, _ := 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)
require.NotEmpty(t, models)
entry, ok := catalog.Lookup("openai_api")
require.True(t, ok)
for _, m := range models {
want, known := pricing.LookupDefault(entry.PricingSurfaces, m.ID)
require.True(t, known, "%s should be priced by the default table", m.ID)
assert.Equal(t, want.InputPer1k, m.InputPer1k, "input rate for %s", m.ID)
assert.Equal(t, want.OutputPer1k, m.OutputPer1k, "output rate for %s", m.ID)
assert.Equal(t, want.CachedInputPer1k, m.CachedInputPer1k, "cached-input rate for %s", m.ID)
}
}
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")
}
// TestRedirectsAreNotFollowed covers a gap the other tests leave open: they all
// inject an HTTPClient, which bypasses httpClient() and therefore the redirect
// policy entirely. The policy is a security control — a 302 moves the request
// to a host checkPublicHost never resolved — so it needs a test that goes
// through the constructor the manager actually uses.
func TestRedirectsAreNotFollowed(t *testing.T) {
var hits int
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits++
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
}))
t.Cleanup(srv.Close)
for name, cl := range map[string]*Client{
// The production shape: no injected client at all.
"default client": {AllowPrivateHosts: true},
// An injected client that states no policy must inherit ours rather
// than silently chasing the redirect.
"injected client with no policy": {
AllowPrivateHosts: true,
HTTPClient: &http.Client{},
},
} {
t.Run(name, func(t *testing.T) {
hits = 0
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
require.NoError(t, err)
resp, err := cl.httpClient().Do(req)
require.NoError(t, err)
t.Cleanup(func() { _ = resp.Body.Close() })
assert.Equal(t, http.StatusFound, resp.StatusCode,
"the redirect must be surfaced, not followed to an unchecked host")
assert.Equal(t, 1, hits, "exactly one request must leave the client")
})
}
}
// TestInjectedClientKeepsItsOwnRedirectPolicy pins that the default above is a
// default, not an override, and that supplying it does not mutate the caller's
// client — one Client is shared across every request, so a write here would
// race.
func TestInjectedClientKeepsItsOwnRedirectPolicy(t *testing.T) {
own := func(*http.Request, []*http.Request) error { return nil }
injected := &http.Client{CheckRedirect: own}
cl := &Client{HTTPClient: injected}
assert.Same(t, injected, cl.httpClient(),
"a client that states a policy must be handed back untouched")
bare := &http.Client{}
cl = &Client{HTTPClient: bare}
require.NotSame(t, bare, cl.httpClient(), "the policy must be applied to a copy")
assert.Nil(t, bare.CheckRedirect, "the caller's client must not be written to")
}
// TestDialGuardRejectsRebindingToANonPublicAddress covers the window between
// the two DNS lookups. checkPublicHost resolves the host, then the transport
// resolves it again to dial; a name whose owner answers the first with a public
// address and the second with 127.0.0.1 would otherwise pass the guard and
// still reach loopback. The dial-time check sees whatever the second lookup
// actually returned.
func TestDialGuardRejectsRebindingToANonPublicAddress(t *testing.T) {
for _, tc := range []struct {
name string
address string
wantErr string
}{
{"loopback", "127.0.0.1:443", "non-public"},
{"cloud metadata", "169.254.169.254:80", "non-public"},
{"rfc1918", "10.1.2.3:443", "non-public"},
{"netbird overlay", "100.90.1.2:443", "non-public"},
{"loopback v6", "[::1]:443", "non-public"},
{"unresolved name", "evil.example.com:443", "not an IP"},
{"no port", "1.1.1.1", "unreadable"},
} {
t.Run(tc.name, func(t *testing.T) {
err := guardDialAddress(tc.address)
require.Error(t, err)
assert.Contains(t, err.Error(), tc.wantErr)
})
}
assert.NoError(t, guardDialAddress("1.1.1.1:443"), "a public address must still be dialled")
assert.NoError(t, guardDialAddress("[2606:4700:4700::1111]:443"))
}
// TestDialGuardIsInstalledOnTheDefaultClient pins the wiring rather than the
// guard: a correct guard nothing calls protects nothing.
func TestDialGuardIsInstalledOnTheDefaultClient(t *testing.T) {
cl := &Client{}
transport, ok := cl.httpClient().Transport.(*http.Transport)
require.True(t, ok, "the default discovery client must carry the guarded transport")
require.NotNil(t, transport.DialContext, "the guarded transport must dial through the guard")
_, err := transport.DialContext(context.Background(), "tcp", "127.0.0.1:9")
require.Error(t, err, "the guard must refuse loopback even when the caller dials it directly")
assert.Contains(t, err.Error(), "non-public")
// Tests point the client at a loopback server on purpose, so the opt-out
// has to reach the dialer too.
relaxed := &Client{AllowPrivateHosts: true}
assert.Equal(t, http.DefaultTransport, relaxed.httpClient().Transport)
}
// TestCallerInputFailuresAreMarkedInvalid keeps the handler's 400 mapping
// honest: it branches on this sentinel, so an unmarked caller-input failure
// silently becomes a 500.
func TestCallerInputFailuresAreMarkedInvalid(t *testing.T) {
for _, tc := range []struct {
name string
req Request
}{
{"unknown provider", Request{CatalogID: "not_a_provider", APIKey: "k"}},
{"unusable upstream", Request{CatalogID: "openai_api", UpstreamURL: "://", APIKey: "k"}},
{"missing api key", Request{CatalogID: "openai_api", UpstreamURL: "https://api.openai.com"}},
{"no region to read", Request{
CatalogID: "bedrock_api",
UpstreamURL: "https://bedrock-runtime.amazonaws.com",
APIKey: "aws-bearer",
}},
} {
t.Run(tc.name, func(t *testing.T) {
cl, _ := newStubClient(http.StatusOK, openAIListing)
_, err := cl.Fetch(context.Background(), tc.req)
require.Error(t, err)
assert.ErrorIs(t, err, ErrInvalidRequest)
})
}
}
// 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", ""},
// Bedrock's regionless endpoint carries both halves of the template at
// once, with nothing between them. It has to read as "no region here"
// rather than as an inverted slice range.
{"bedrock regionless endpoint", bedrock, "https://bedrock-runtime.amazonaws.com", ""},
{"bedrock regionless without scheme", bedrock, "bedrock-runtime.amazonaws.com", ""},
} {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, RegionFromUpstream(tc.entry, tc.upstream))
})
}
}
// bedrockGeoListing carries profiles from geographies the original prefix list
// did not name. Every one reduces to a catalog key, so every one must arrive
// priced — an unstripped geography is what made a real account's listing come
// back almost entirely at zero.
const bedrockGeoListing = `{"inferenceProfileSummaries":[
{"inferenceProfileId":"jp.anthropic.claude-sonnet-5-20260514-v1:0",
"inferenceProfileName":"JP Anthropic Claude Sonnet 5","status":"ACTIVE","type":"SYSTEM_DEFINED"},
{"inferenceProfileId":"au.anthropic.claude-haiku-4-5-20251001-v1:0",
"inferenceProfileName":"AU Anthropic Claude Haiku 4.5","status":"ACTIVE","type":"SYSTEM_DEFINED"},
{"inferenceProfileId":"us-gov.anthropic.claude-sonnet-5-20260514-v1:0",
"inferenceProfileName":"GovCloud Anthropic Claude Sonnet 5","status":"ACTIVE","type":"SYSTEM_DEFINED"}
]}`
func TestBedrockProfilesFromAnyGeographyArrivePriced(t *testing.T) {
cl, _ := newStubClient(http.StatusOK, bedrockGeoListing)
models, err := cl.Fetch(context.Background(), Request{
CatalogID: "bedrock_api",
UpstreamURL: "https://bedrock-runtime.eu-central-1.amazonaws.com",
APIKey: "aws-token",
})
require.NoError(t, err)
require.Len(t, models, 3)
for _, m := range models {
assert.True(t, m.PricingKnown, "%s must resolve to a catalog rate", m.ID)
assert.Greater(t, m.InputPer1k, 0.0, "input rate for %s", m.ID)
assert.Greater(t, m.OutputPer1k, 0.0, "output rate for %s", m.ID)
assert.Greater(t, m.CacheReadPer1k, 0.0, "cache-read rate for %s", m.ID)
}
// The wire id is preserved whatever the pricing key reduced to: it is the
// only form that works at invoke time.
assert.Equal(t, "jp.anthropic.claude-sonnet-5-20260514-v1:0", models[0].ID)
}

View File

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

View File

@@ -47,11 +47,17 @@ 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 (

View File

@@ -82,11 +82,6 @@ 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
@@ -150,11 +145,6 @@ 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

View File

@@ -116,13 +116,11 @@ 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")
// 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.
// Supplementals present on their surfaces.
for surface, ids := range map[string][]string{
"openai": {"gpt-5", "gpt-5-mini", "gpt-5-nano"},
"anthropic": {"claude-opus-5", "claude-sonnet-5", "kimi-k3[1m]", "kimi-k3"},
"bedrock": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5"},
"anthropic": {"claude-opus-5", "kimi-k3[1m]", "kimi-k3"},
"bedrock": {"anthropic.claude-opus-5"},
} {
for _, id := range ids {
_, ok := table[surface][id]

View File

@@ -10,7 +10,6 @@ import (
"strings"
"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/types"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
@@ -212,19 +211,7 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
groupIndex := indexProviderGroups(enabledPolicies)
// 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)
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex)
if err != nil {
return nil, err
}
@@ -241,6 +228,11 @@ 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
@@ -359,11 +351,6 @@ 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.
@@ -381,9 +368,6 @@ type routerProviderRoute struct {
// proxy dials this provider's upstream. For self-hosted / internal gateways
// behind a private or self-signed certificate.
SkipTLSVerify bool `json:"skip_tls_verify,omitempty"`
// DiscoveryHost, when set, is the host serving this provider's model
// listing, for a vendor that does not serve it from the inference host.
DiscoveryHost string `json:"discovery_host,omitempty"`
}
// indexProviderGroups walks the enabled policies and returns, per
@@ -438,7 +422,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, modelPolicies map[string][]routerModelPolicy) ([]byte, error) {
func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) {
cfg := routerConfig{Providers: make([]routerProviderRoute, 0, len(providers))}
for _, p := range providers {
groups, hasPolicy := groupIndex[p.ID]
@@ -451,9 +435,6 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]
if err != nil {
return nil, fmt.Errorf("router config for provider %s: %w", p.ID, err)
}
// Lookup rather than assume: an unknown provider id yields the zero
// entry, which declares no discovery and so contributes nothing.
catalogEntry, _ := catalog.Lookup(p.ProviderID)
headerName, headerValue, gcpSAKeyB64, err := providerAuthHeader(p)
if err != nil {
return nil, err
@@ -468,12 +449,10 @@ 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,
SkipTLSVerify: p.SkipTLSVerification,
DiscoveryHost: discoveryHost(catalogEntry, p.UpstreamURL),
})
}
out, err := json.Marshal(cfg)
@@ -483,33 +462,6 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]
return out, nil
}
// discoveryHost returns the host serving this provider's model listing when it
// differs from the inference host, and empty when the two are the same — which
// is true of every vendor but Bedrock, whose ListInferenceProfiles is a control
// plane operation on bedrock.<region> while InvokeModel must go to
// bedrock-runtime.<region>. One provider record therefore needs two hosts.
//
// The catalog declares the listing host; the region is recovered from the
// upstream the operator configured, since a provider record carries no region
// field. An upstream matching no catalog template yields empty rather than a
// guess: a proxied or self-hosted Bedrock endpoint may serve both from one
// place, and inventing a host would send the credential somewhere the operator
// never configured.
func discoveryHost(entry catalog.Provider, upstreamURL string) string {
if entry.Discovery == nil || entry.Discovery.Host == "" {
return ""
}
host := entry.Discovery.Host
if !strings.Contains(host, catalog.RegionPlaceholder) {
return host
}
region := modeldiscovery.RegionFromUpstream(entry, upstreamURL)
if region == "" {
return ""
}
return strings.ReplaceAll(host, catalog.RegionPlaceholder, region)
}
// providerVendor returns the parser surface ("openai", "anthropic", …)
// the provider speaks, sourced from its catalog entry's ParserID. The
// router uses it to keep a request the parser tagged with a vendor on a
@@ -1146,46 +1098,3 @@ 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
}

View File

@@ -103,37 +103,3 @@ func TestBuildCostMeterConfig_OrphanAndGatewayProviders(t *testing.T) {
assert.NotContains(t, cfg.Pricing.Providers, "prov-litellm", "empty-models gateway needs no per-record entry")
assert.NotEmpty(t, cfg.Pricing.Defaults["openai"], "defaults still ship so the gateway's catalog-model traffic is priced")
}
// TestBuildCostMeterConfig_BedrockGeographyOutsideTheOriginalFour is the
// accounting half of the geography bug. The docs tell operators to register a
// Bedrock id exactly as AWS issues it, region prefix included, and the cost
// meter keys its table by the normalized form. While the geography was matched
// against a list of four, a profile issued anywhere else kept its prefix,
// missed the catalog entry it was meant to inherit from, and billed with a
// zero entry underneath the operator's own rates — so every cache bucket
// metered free and a model priced only by catalog defaults metered at nothing
// at all.
func TestBuildCostMeterConfig_BedrockGeographyOutsideTheOriginalFour(t *testing.T) {
for _, geo := range []string{"jp", "au", "ca", "sa", "us-gov"} {
t.Run(geo, func(t *testing.T) {
bedrock := &types.Provider{
ID: "prov-bedrock",
ProviderID: "bedrock_api",
Enabled: true,
Models: []types.ProviderModel{
{ID: geo + ".anthropic.claude-sonnet-5-20260514-v1:0", InputPer1k: 0.003, OutputPer1k: 0.015},
},
}
raw, err := buildCostMeterConfigJSON([]*types.Provider{bedrock}, map[string][]string{"prov-bedrock": {"grp"}})
require.NoError(t, err)
cfg := decodeCostMeterConfig(t, raw)
e, ok := cfg.Pricing.Providers["prov-bedrock"]["anthropic.claude-sonnet-5"]
require.True(t, ok, "a %s profile must key by the same normalized id the parser emits", geo)
assert.InDelta(t, 0.0003, e.CacheReadPer1k, 1e-9,
"cache read must be inherited from the bedrock default entry, not left at zero")
assert.InDelta(t, 0.00375, e.CacheCreationPer1k, 1e-9,
"cache creation must be inherited from the bedrock default entry, not left at zero")
})
}
}

View File

@@ -4,7 +4,6 @@ import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
@@ -94,75 +93,3 @@ 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")
})
}

View File

@@ -10,7 +10,6 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/server/store"
@@ -1246,57 +1245,3 @@ func TestSynthesizeServices_EmptyAPIKey_FailsClosed(t *testing.T) {
require.Error(t, err, "synthesis must refuse a provider with no api key")
assert.Contains(t, err.Error(), "no api key", "error must surface the missing credential")
}
// TestDiscoveryHost pins which providers get a separate listing host. Getting
// this wrong in either direction is costly: a missing host leaves Bedrock
// discovery 404ing at AWS, and a host on the wrong provider would send that
// provider's listing — and its credential — somewhere the operator never
// configured.
func TestDiscoveryHost(t *testing.T) {
entry := func(id string) catalog.Provider {
p, ok := catalog.Lookup(id)
require.True(t, ok, "catalog entry %s must exist", id)
return p
}
for _, tc := range []struct {
name string
entry catalog.Provider
upstream string
want string
}{
{
// ListInferenceProfiles is a control-plane operation; the runtime
// host answers <UnknownOperationException/> for it.
name: "bedrock splits the listing off the runtime host",
entry: entry("bedrock_api"), upstream: "https://bedrock-runtime.eu-central-1.amazonaws.com",
want: "bedrock.eu-central-1.amazonaws.com",
},
{
name: "bedrock in another region",
entry: entry("bedrock_api"), upstream: "https://bedrock-runtime.us-west-2.amazonaws.com",
want: "bedrock.us-west-2.amazonaws.com",
},
{
// A proxied Bedrock endpoint may well serve both from one place,
// and there is no region to read back out of it.
name: "proxied bedrock upstream yields no discovery host",
entry: entry("bedrock_api"), upstream: "https://bedrock.internal.example.com",
want: "",
},
{
name: "openai serves its listing from the same host",
entry: entry("openai_api"), upstream: "https://api.openai.com",
want: "",
},
{
name: "vertex serves its listing from the same host",
entry: entry("vertex_ai_api"), upstream: "https://us-east5-aiplatform.googleapis.com",
want: "",
},
} {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, discoveryHost(tc.entry, tc.upstream))
})
}
}

View File

@@ -1311,7 +1311,7 @@ func (s *ProxyServiceServer) authenticateHeader(ctx context.Context, serviceID s
lastErr = err
continue
}
return true, "header-user", proxyauth.MethodHeader
return true, proxyauth.HeaderUserID, proxyauth.MethodHeader
}
if lastErr != nil {

View File

@@ -30,6 +30,12 @@ const (
SessionJWTIssuer = "netbird-management"
)
// HeaderUserID is the synthetic user id recorded for header-authenticated
// requests. Header auth validates a per-service secret and resolves no user
// record, so proxy access logs and management-minted session tokens both
// attribute the request to this id.
const HeaderUserID = "header-user"
// ResolveProto determines the protocol scheme based on the forwarded proto
// configuration. When set to "http" or "https" the value is used directly.
// Otherwise TLS state is used: if conn is non-nil "https" is returned, else "http".

View File

@@ -1,36 +1,32 @@
package auth
import (
"errors"
"fmt"
"crypto/sha256"
"net/http"
"sync"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/hash/argon2id"
)
// ErrHeaderAuthFailed indicates that the header was present but the
// credential did not validate. Callers should return 401 instead of
// falling through to other auth schemes.
var ErrHeaderAuthFailed = errors.New("header authentication failed")
// Header implements header-based authentication. The proxy checks for the
// configured header in each request and validates its value via gRPC.
// Header implements header-based authentication. The service mapping carries
// the argon2id hash of every value accepted for the header, so the proxy
// verifies the credential locally rather than round-tripping to management.
type Header struct {
id types.ServiceID
accountId types.AccountID
headerName string
client authenticator
hashes []string
verified *verifiedValues
}
// NewHeader creates a Header authentication scheme for the given header name.
func NewHeader(client authenticator, id types.ServiceID, accountId types.AccountID, headerName string) Header {
// NewHeader creates a Header authentication scheme accepting any value whose
// argon2id hash appears in hashes. An empty hashes slice rejects every request
// carrying the header, so a mapping that arrived without its hashes fails
// closed instead of leaving the service unprotected.
func NewHeader(headerName string, hashes []string) Header {
return Header{
id: id,
accountId: accountId,
headerName: headerName,
client: client,
headerName: http.CanonicalHeaderKey(headerName),
hashes: hashes,
verified: &verifiedValues{seen: make(map[[32]byte]struct{}, len(hashes))},
}
}
@@ -39,31 +35,55 @@ func (Header) Type() auth.Method {
return auth.MethodHeader
}
// Authenticate checks for the configured header in the request. If absent,
// returns empty (unauthenticated). If present, validates via gRPC.
func (h Header) Authenticate(r *http.Request) (string, string, error) {
// Authenticate satisfies Scheme. Header credentials are resolved by Verify
// before the scheme loop runs, so a request that reaches here never carries
// the header and there is no credential to prompt for.
func (Header) Authenticate(*http.Request) (string, string, error) {
return "", "", nil
}
// Verify reports whether the request carries the configured header and, when
// it does, whether the value matches one of the service's hashes.
func (h Header) Verify(r *http.Request) (present, matched bool) {
value := r.Header.Get(h.headerName)
if value == "" {
return "", "", nil
return false, false
}
res, err := h.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
Id: string(h.id),
AccountId: string(h.accountId),
Request: &proto.AuthenticateRequest_HeaderAuth{
HeaderAuth: &proto.HeaderAuthRequest{
HeaderValue: value,
HeaderName: h.headerName,
},
},
})
if err != nil {
return "", "", fmt.Errorf("authenticate header: %w", err)
digest := sha256.Sum256([]byte(value))
if h.verified.has(digest) {
return true, true
}
if res.GetSuccess() {
return res.GetSessionToken(), "", nil
for _, hash := range h.hashes {
if argon2id.Verify(value, hash) == nil {
h.verified.add(digest)
return true, true
}
}
return "", "", ErrHeaderAuthFailed
return true, false
}
// verifiedValues remembers which header values already passed argon2id
// verification. argon2id is deliberately expensive (19 MiB, two passes) and
// header credentials repeat on every request, so re-deriving per request would
// dominate the hot path. The set cannot outgrow the number of configured
// hashes, and a mapping update builds a fresh scheme with an empty set.
// Values are keyed by digest so the plaintext credential is not retained.
type verifiedValues struct {
mu sync.Mutex
seen map[[32]byte]struct{}
}
func (v *verifiedValues) has(digest [32]byte) bool {
v.mu.Lock()
defer v.mu.Unlock()
_, ok := v.seen[digest]
return ok
}
func (v *verifiedValues) add(digest [32]byte) {
v.mu.Lock()
defer v.mu.Unlock()
v.seen[digest] = struct{}{}
}

View File

@@ -146,7 +146,7 @@ func (mw *Middleware) Protect(next http.Handler) http.Handler {
return
}
if mw.forwardWithHeaderAuth(w, r, host, config, next) {
if mw.forwardWithHeaderAuth(w, r, config, next) {
return
}
@@ -325,6 +325,16 @@ func (mw *Middleware) forwardWithSessionCookie(w http.ResponseWriter, r *http.Re
if err != nil {
return false
}
// Header auth is checked per request against the mapping's hashes and mints
// no session, so a header-method token can only predate that. Honouring it
// would keep a rotated credential working until the token expired.
if method == auth.MethodHeader.String() {
mw.logger.WithField("host", host).
Debug("ignoring header-auth session cookie; the header is required on every request")
return false
}
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(userID)
cd.SetUserEmail(email)
@@ -436,14 +446,14 @@ func isTunnelSourceIP(ip netip.Addr) bool {
// forwardWithHeaderAuth checks for a Header auth scheme. If the header validates,
// the request is forwarded directly (no redirect), which is important for API clients.
func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Request, config DomainConfig, next http.Handler) bool {
for _, scheme := range config.Schemes {
hdr, ok := scheme.(Header)
if !ok {
continue
}
handled := mw.tryHeaderScheme(w, r, host, config, hdr, next)
handled := mw.tryHeaderScheme(w, r, hdr, next)
if handled {
return true
}
@@ -451,40 +461,27 @@ func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Reque
return false
}
func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, hdr Header, next http.Handler) bool {
token, _, err := hdr.Authenticate(r)
if err != nil {
return mw.handleHeaderAuthError(w, r, err)
}
if token == "" {
// tryHeaderScheme verifies the credential against the hashes the service
// mapping carries. No session token is issued: the credential travels on
// every request, so there is nothing for a cookie to save.
func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, hdr Header, next http.Handler) bool {
present, matched := hdr.Verify(r)
if !present {
return false
}
result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
if err != nil {
if !matched {
mw.logger.WithFields(log.Fields{
"host": r.Host,
"header": hdr.headerName,
}).Debug("header auth rejected: value does not match any configured hash")
setHeaderCapturedData(r.Context(), "", "", nil, nil)
status := http.StatusBadRequest
msg := "invalid session token"
if errors.Is(err, errValidationUnavailable) {
status = http.StatusBadGateway
msg = "authentication service unavailable"
}
http.Error(w, msg, status)
return true
}
if !result.Valid {
setHeaderCapturedData(r.Context(), result.UserID, result.UserEmail, result.Groups, result.GroupNames)
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return true
}
setSessionCookie(w, token, config.SessionExpiration)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(result.UserID)
cd.SetUserEmail(result.UserEmail)
cd.SetUserGroups(result.Groups)
cd.SetUserGroupNames(result.GroupNames)
cd.SetUserID(auth.HeaderUserID)
cd.SetAuthMethod(auth.MethodHeader.String())
}
@@ -492,20 +489,6 @@ func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, ho
return true
}
func (mw *Middleware) handleHeaderAuthError(w http.ResponseWriter, r *http.Request, err error) bool {
if errors.Is(err, ErrHeaderAuthFailed) {
setHeaderCapturedData(r.Context(), "", "", nil, nil)
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return true
}
mw.logger.WithField("scheme", "header").Warnf("header auth infrastructure error: %v", err)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
}
http.Error(w, "authentication service unavailable", http.StatusBadGateway)
return true
}
func setHeaderCapturedData(ctx context.Context, userID, userEmail string, groups, groupNames []string) {
cd := proxy.CapturedDataFromContext(ctx)
if cd == nil {

View File

@@ -25,6 +25,7 @@ import (
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/restrict"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/hash/argon2id"
"github.com/netbirdio/netbird/shared/management/proto"
)
@@ -1023,38 +1024,24 @@ func TestProtect_OIDCWithOtherMethodShowsLoginPage(t *testing.T) {
assert.Equal(t, http.StatusUnauthorized, rec.Code, "should show login page when multiple methods exist")
}
// mockAuthenticator is a minimal mock for the authenticator gRPC interface
// used by the Header scheme.
type mockAuthenticator struct {
fn func(ctx context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error)
}
func (m *mockAuthenticator) Authenticate(ctx context.Context, in *proto.AuthenticateRequest, _ ...grpc.CallOption) (*proto.AuthenticateResponse, error) {
return m.fn(ctx, in)
}
// newHeaderSchemeWithToken creates a Header scheme backed by a mock that
// returns a signed session token when the expected header value is provided.
func newHeaderSchemeWithToken(t *testing.T, kp *sessionkey.KeyPair, headerName, expectedValue string) Header {
// newHeaderScheme creates a Header scheme accepting each of the given values,
// hashed the way management hashes them before putting them on the mapping.
func newHeaderScheme(t *testing.T, headerName string, acceptedValues ...string) Header {
t.Helper()
token, err := sessionkey.SignToken(kp.PrivateKey, "header-user", "", "example.com", auth.MethodHeader, nil, nil, time.Hour)
require.NoError(t, err)
mock := &mockAuthenticator{fn: func(_ context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
ha := req.GetHeaderAuth()
if ha != nil && ha.GetHeaderValue() == expectedValue {
return &proto.AuthenticateResponse{Success: true, SessionToken: token}, nil
}
return &proto.AuthenticateResponse{Success: false}, nil
}}
return NewHeader(mock, "svc1", "acc1", headerName)
hashes := make([]string, 0, len(acceptedValues))
for _, v := range acceptedValues {
hash, err := argon2id.Hash(v)
require.NoError(t, err, "hashing an accepted header value must succeed")
hashes = append(hashes, hash)
}
return NewHeader(headerName, hashes)
}
func TestProtect_HeaderAuth_ForwardsOnSuccess(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
hdr := newHeaderScheme(t, "X-API-Key", "secret-key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
var backendCalled bool
@@ -1075,19 +1062,12 @@ func TestProtect_HeaderAuth_ForwardsOnSuccess(t *testing.T) {
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, "ok", rec.Body.String())
// Session cookie should be set.
var sessionCookie *http.Cookie
// The credential rides on every request, so no session cookie is issued.
for _, c := range rec.Result().Cookies() {
if c.Name == auth.SessionCookieName {
sessionCookie = c
break
}
assert.NotEqual(t, auth.SessionCookieName, c.Name, "header auth must not issue a session cookie")
}
require.NotNil(t, sessionCookie, "session cookie should be set after successful header auth")
assert.True(t, sessionCookie.HttpOnly)
assert.True(t, sessionCookie.Secure)
assert.Equal(t, "header-user", capturedData.GetUserID())
assert.Equal(t, auth.HeaderUserID, capturedData.GetUserID())
assert.Equal(t, "header", capturedData.GetAuthMethod())
}
@@ -1095,7 +1075,7 @@ func TestProtect_HeaderAuth_MissingHeaderFallsThrough(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
hdr := newHeaderScheme(t, "X-API-Key", "secret-key")
// Also add a PIN scheme so we can verify fallthrough behavior.
pinScheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr, pinScheme}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
@@ -1114,10 +1094,7 @@ func TestProtect_HeaderAuth_WrongValueReturns401(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
mock := &mockAuthenticator{fn: func(_ context.Context, _ *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
return &proto.AuthenticateResponse{Success: false}, nil
}}
hdr := NewHeader(mock, "svc1", "acc1", "X-API-Key")
hdr := newHeaderScheme(t, "X-API-Key", "secret-key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
capturedData := proxy.NewCapturedData("")
@@ -1131,93 +1108,157 @@ func TestProtect_HeaderAuth_WrongValueReturns401(t *testing.T) {
assert.Equal(t, http.StatusUnauthorized, rec.Code)
assert.Equal(t, "header", capturedData.GetAuthMethod())
assert.Empty(t, hdr.verified.seen, "a rejected value must not be memoized")
}
func TestProtect_HeaderAuth_InfraErrorReturns502(t *testing.T) {
// TestProtect_HeaderAuth_NoHashesFailsClosed covers a mapping that names a
// header but carries no hash for it: the check cannot be evaluated, so the
// request must be denied rather than let through unauthenticated.
func TestProtect_HeaderAuth_NoHashesFailsClosed(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
mock := &mockAuthenticator{fn: func(_ context.Context, _ *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
return nil, errors.New("gRPC unavailable")
}}
hdr := NewHeader(mock, "svc1", "acc1", "X-API-Key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
handler := mw.Protect(newPassthroughHandler())
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.Header.Set("X-API-Key", "some-key")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusBadGateway, rec.Code)
}
func TestProtect_HeaderAuth_SubsequentRequestUsesSessionCookie(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderSchemeWithToken(t, kp, "X-API-Key", "secret-key")
hdr := NewHeader("X-API-Key", nil)
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
var backendCalled bool
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
backendCalled = true
w.WriteHeader(http.StatusOK)
}))
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.Header.Set("X-API-Key", "any-key")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusUnauthorized, rec.Code)
assert.False(t, backendCalled, "a header auth with no hashes must not admit the request")
}
// TestProtect_HeaderAuth_SubsequentRequestRequiresHeader verifies that header
// auth grants no ambient session: a follow-up request that drops the header is
// treated as unauthenticated.
func TestProtect_HeaderAuth_SubsequentRequestRequiresHeader(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderScheme(t, "X-API-Key", "secret-key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
var backendCalls int
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
backendCalls++
w.WriteHeader(http.StatusOK)
}))
// First request with header auth.
req1 := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req1.Header.Set("X-API-Key", "secret-key")
req1 = req1.WithContext(proxy.WithCapturedData(req1.Context(), proxy.NewCapturedData("")))
rec1 := httptest.NewRecorder()
handler.ServeHTTP(rec1, req1)
require.Equal(t, http.StatusOK, rec1.Code)
require.Equal(t, 1, backendCalls)
// Extract session cookie.
var sessionCookie *http.Cookie
for _, c := range rec1.Result().Cookies() {
if c.Name == auth.SessionCookieName {
sessionCookie = c
break
}
}
require.NotNil(t, sessionCookie)
// Second request with only the session cookie (no header).
capturedData2 := proxy.NewCapturedData("")
// Same client, second request, header omitted: no cookie was handed out, so
// there is nothing to carry the earlier success forward.
req2 := httptest.NewRequest(http.MethodGet, "http://example.com/other", nil)
req2.AddCookie(sessionCookie)
req2 = req2.WithContext(proxy.WithCapturedData(req2.Context(), capturedData2))
for _, c := range rec1.Result().Cookies() {
req2.AddCookie(c)
}
rec2 := httptest.NewRecorder()
handler.ServeHTTP(rec2, req2)
assert.Equal(t, http.StatusOK, rec2.Code)
assert.Equal(t, "header-user", capturedData2.GetUserID())
assert.Equal(t, "header", capturedData2.GetAuthMethod())
assert.Equal(t, http.StatusUnauthorized, rec2.Code, "dropping the header must revoke access")
assert.Equal(t, 1, backendCalls, "backend must not be reached without the header")
}
// TestProtect_HeaderAuth_MultipleValuesSameHeader verifies that the proxy
// correctly handles multiple valid credentials for the same header name.
// In production, the mgmt gRPC authenticateHeader iterates all configured
// header auths and accepts if any hash matches (OR semantics). The proxy
// creates one Header scheme per entry, but a single gRPC call checks all.
// TestProtect_HeaderAuth_LegacySessionCookieIsIgnored covers the upgrade
// window. Header auth used to mint a session token, so cookies with
// method=header survive a proxy upgrade and stay signature-valid for their full
// lifetime. They must not stand in for the header, or a credential rotated
// right after the upgrade would keep working until every such token expired.
func TestProtect_HeaderAuth_LegacySessionCookieIsIgnored(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderScheme(t, "X-API-Key", "secret-key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
// A token management would have minted for header auth before the upgrade.
legacyToken, err := sessionkey.SignToken(kp.PrivateKey, auth.HeaderUserID, "", "example.com", auth.MethodHeader, nil, nil, time.Hour)
require.NoError(t, err)
var backendCalls int
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
backendCalls++
w.WriteHeader(http.StatusOK)
}))
t.Run("cookie alone is rejected", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.AddCookie(&http.Cookie{Name: auth.SessionCookieName, Value: legacyToken})
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusUnauthorized, rec.Code, "a header-auth cookie must not authenticate on its own")
assert.Equal(t, 0, backendCalls, "backend must not be reached without the header")
})
t.Run("cookie does not block the header path", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.AddCookie(&http.Cookie{Name: auth.SessionCookieName, Value: legacyToken})
req.Header.Set("X-API-Key", "secret-key")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusOK, rec.Code, "a client sending both must still be admitted by the header")
assert.Equal(t, 1, backendCalls)
})
}
// TestProtect_HeaderAuth_RepeatedValueIsMemoized verifies the KDF is run once
// per distinct accepted value. argon2id is deliberately expensive, so a
// credential that repeats on every request must not be re-derived each time.
func TestProtect_HeaderAuth_RepeatedValueIsMemoized(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
hdr := newHeaderScheme(t, "X-API-Key", "key-a", "key-b")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
}))
get := func(value string) int {
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.Header.Set("X-API-Key", value)
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
return rec.Code
}
require.Equal(t, http.StatusOK, get("key-a"))
require.Equal(t, http.StatusOK, get("key-a"))
assert.Len(t, hdr.verified.seen, 1, "the same value must be memoized once")
require.Equal(t, http.StatusOK, get("key-b"))
assert.Len(t, hdr.verified.seen, 2, "each accepted value gets its own entry")
require.Equal(t, http.StatusUnauthorized, get("key-c"))
assert.Len(t, hdr.verified.seen, 2, "rejected values must not grow the set")
}
// TestProtect_HeaderAuth_MultipleValuesSameHeader verifies that a service with
// several accepted credentials for one header name accepts any of them.
// Management applied these OR semantics while it still validated the value; the
// proxy preserves them by carrying every hash for a name on one scheme.
func TestProtect_HeaderAuth_MultipleValuesSameHeader(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
// Mock simulates mgmt behavior: accepts either token-a or token-b.
accepted := map[string]bool{"Bearer token-a": true, "Bearer token-b": true}
mock := &mockAuthenticator{fn: func(_ context.Context, req *proto.AuthenticateRequest) (*proto.AuthenticateResponse, error) {
ha := req.GetHeaderAuth()
if ha != nil && accepted[ha.GetHeaderValue()] {
token, err := sessionkey.SignToken(kp.PrivateKey, "header-user", "", "example.com", auth.MethodHeader, nil, nil, time.Hour)
require.NoError(t, err)
return &proto.AuthenticateResponse{Success: true, SessionToken: token}, nil
}
return &proto.AuthenticateResponse{Success: false}, nil
}}
// Single Header scheme (as if one entry existed), but the mock checks both values.
hdr := NewHeader(mock, "svc1", "acc1", "Authorization")
hdr := newHeaderScheme(t, "Authorization", "Bearer token-a", "Bearer token-b")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
var backendCalled bool

View File

@@ -13,14 +13,6 @@ 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.

View File

@@ -10,8 +10,6 @@ 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
@@ -94,10 +92,7 @@ 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. 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.
// Lookup returns the entry for the given provider surface and model.
func (t *Table) Lookup(provider, model string) (Entry, bool) {
if t == nil {
return Entry{}, false
@@ -106,14 +101,7 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) {
if !ok {
return Entry{}, false
}
if e, found := byModel[model]; found {
return e, true
}
undated := sharedllm.NormalizeAnthropicModel(model)
if undated == model {
return Entry{}, false
}
e, ok := byModel[undated]
e, ok := byModel[model]
return e, ok
}

View File

@@ -175,22 +175,3 @@ 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")
}

View File

@@ -11,7 +11,6 @@ import (
"fmt"
"strconv"
"github.com/netbirdio/netbird/proxy/internal/llm"
"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
@@ -176,28 +175,13 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
// Anthropic route still bills its cache buckets additively.
func (m *Middleware) lookupCosts(md []middleware.KV, surface, model string, inTokens, outTokens, cachedTokens, cacheCreationTokens int64) (pricing.Costs, bool) {
if recordID := lookupKV(md, middleware.KeyLLMResolvedProviderID); recordID != "" {
if entry, ok := perRecordEntry(m.perRecord[recordID], model); ok {
if entry, ok := m.perRecord[recordID][model]; ok {
return pricing.EntryCosts(entry, surface, inTokens, outTokens, cachedTokens, cacheCreationTokens), true
}
}
return m.defaults.Costs(surface, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
}
// perRecordEntry resolves the operator's stored price for a model on one
// provider record, falling back to the undated form of a dated Anthropic id
// so a client that pins a release date still bills at the registered rate.
func perRecordEntry(byModel map[string]pricing.Entry, model string) (pricing.Entry, bool) {
if entry, ok := byModel[model]; ok {
return entry, true
}
undated := llm.NormalizeAnthropicModel(model)
if undated == model {
return pricing.Entry{}, false
}
entry, ok := byModel[undated]
return entry, ok
}
// usd renders a cost as the fixed-precision string every cost.usd_* key
// carries, so the per-bucket values and the aggregates round identically.
//

View File

@@ -84,10 +84,8 @@ func (m *Middleware) MutationsSupported() bool { return false }
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
nonInference, _ := lookupMetadata(in.Metadata, middleware.KeyLLMNonInference)
if denial := m.evaluateAllowlist(providerID, surface, model, modelPresent, nonInference == "true"); denial != nil {
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
return denial, nil
}
@@ -116,7 +114,7 @@ func (m *Middleware) Close() error { return nil }
// evaluateAllowlist denies when the resolved provider's allowlist rejects the
// model; nil means proceed. Scoped to the provider llm_router resolved, so an
// unrestricted provider (absent from config) is never caught by another's list.
func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelPresent, nonInference bool) *middleware.Output {
func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
if len(m.cfg.ProviderAllowlists) == 0 {
return nil
}
@@ -124,7 +122,7 @@ func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelP
// if this request targets a restricted provider — fail closed. llm_router
// normally stamps the provider first, so this is a defensive guard.
if providerID == "" {
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
if !restricted {
@@ -135,29 +133,18 @@ func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelP
// Fail closed: with an allowlist in effect for this provider, a request whose
// model the parser couldn't extract (absent/empty) is denied. This enforces
// the allowlist for path-routed providers (Bedrock, Vertex) with no body model.
//
// The exception is a non-inference endpoint the router already authorised.
// The model listing and the connection-warming probe name no model
// anywhere — not in a body, not in the path — so failing closed here
// rejected model discovery for exactly the accounts that configured an
// allowlist, which is the outage this endpoint is meant to avoid. The
// per-model lookup does name one (the router stamps it from the path), so
// it still falls through to the allowlist check below.
if !modelPresent || normaliseModel(model) == "" {
if nonInference {
return nil
}
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
if modelInAllowlist(allowlist, model) {
return nil
}
return denyModel(surface, model, denyCodeModel, denyMessageModel, denyReasonModel)
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
}
// denyModel builds a 403 deny Output for a model-allowlist rejection. model is
// included in the details only when non-empty.
func denyModel(surface, model, code, message, reason string) *middleware.Output {
func denyModel(model, code, message, reason string) *middleware.Output {
details := map[string]string{}
if model != "" {
details["model"] = model
@@ -169,7 +156,6 @@ func denyModel(surface, model, code, message, reason string) *middleware.Output
Code: code,
Message: message,
Details: details,
Surface: surface,
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},

View File

@@ -343,52 +343,3 @@ func TestFactoryNormalisesAllowlist(t *testing.T) {
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out2.Decision, "trimmed entry must still match")
}
// TestAllowlistSkipsNonInferenceWithoutModel covers the reported regression:
// GET /v1/models carries no model anywhere, so the fail-closed rule above
// denied model discovery for exactly the accounts that configured a provider
// allowlist — the clients that read a 403 here render an empty model picker.
// The router authorises those endpoints by path before the guardrail sees
// them, so an absent model there is expected rather than undeterminable.
func TestAllowlistSkipsNonInferenceWithoutModel(t *testing.T) {
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"model discovery must not be refused because it names no model")
}
// TestAllowlistStillAppliesToNonInferenceWithModel pins that the exemption is
// scoped to requests that genuinely name nothing. The per-model lookup
// (GET /v1/models/{id}) is non-inference too, but the router stamps the model
// from its path, so the allowlist must still decide it — otherwise the
// exemption becomes a way to confirm a model the policy blocks.
func TestAllowlistStillAppliesToNonInferenceWithModel(t *testing.T) {
mw := New(providerCfg("gpt-4o"))
t.Run("model in the allowlist", func(t *testing.T) {
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"an allowlisted model must stay reachable")
})
t.Run("model outside the allowlist", func(t *testing.T) {
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"},
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-5"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"non-inference must not become a way past the allowlist")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code,
"a named but blocked model is blocked, not unknown")
})
}

View File

@@ -217,32 +217,6 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
return mutations
}
// bodyInjectableSurfaces are the request-body dialects that accept the
// OpenAI-standard identity fields this middleware writes. A surface
// outside this set gets header-only stamping: "user" and "metadata.tags"
// are not part of the Anthropic Messages schema, which rejects unknown
// top-level fields and permits only "user_id" under metadata, so writing
// them into an Anthropic-shaped body turns a working request into a 400.
// Claude Code speaks that shape through gateway records pinned to the
// OpenAI parser, so the check keys on the detected surface rather than
// on the provider record.
var bodyInjectableSurfaces = map[string]struct{}{
"openai": {},
// An empty surface means no parser claimed the path (a custom gateway
// base). Those upstreams are OpenAI-compatible by convention, so keep
// the long-standing behaviour rather than silently dropping identity.
"": {},
}
// bodyAcceptsOpenAIIdentity reports whether the request body may carry the
// OpenAI-standard identity fields, read from the surface llm_request_parser
// resolved from the request path.
func bodyAcceptsOpenAIIdentity(in *middleware.Input) bool {
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
_, ok := bodyInjectableSurfaces[surface]
return ok
}
// injectIntoBody parses the request body and writes the supplied
// identity dimensions into it. Tags land at metadata.tags (creating
// the metadata object when absent); the user identity lands at the
@@ -251,8 +225,6 @@ func bodyAcceptsOpenAIIdentity(in *middleware.Input) bool {
// was written. Returns ok=false (no mutation) when:
//
// - both inputs are empty (nothing to write);
// - the body speaks a dialect without these fields (see
// bodyInjectableSurfaces);
// - the body is empty or truncated (we don't have the full document
// to safely round-trip);
// - the body isn't a JSON object (skip silently — this middleware
@@ -273,9 +245,6 @@ func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte,
if in == nil || len(in.Body) == 0 || in.BodyTruncated {
return nil, false
}
if !bodyAcceptsOpenAIIdentity(in) {
return nil, false
}
var doc map[string]any
if err := json.Unmarshal(in.Body, &doc); err != nil {
return nil, false

View File

@@ -704,57 +704,3 @@ func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) {
"empty extra value must not be stamped")
}
}
// TestInject_AnthropicBodyIsNotRewritten pins the shape gate. Claude Code
// reaches a LiteLLM record on /v1/messages, where "user" is not a
// permitted top-level field and metadata accepts only "user_id", so
// writing the OpenAI-standard fields would turn a working request into a
// 400 naming a field the client never sent. Header stamping still runs, so
// spend tracking and per-end-user budgets keep working.
func TestInject_AnthropicBodyIsNotRewritten(t *testing.T) {
rule := liteLLMRuleWithBody()
rule.HeaderPair.EndUserIDInBody = true
mw := New(Config{Providers: []ProviderInjection{rule}})
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
in.UserEmail = "alice@example.com"
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
in.Body = []byte(`{"model":"claude-sonnet-5","messages":[]}`)
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations)
assert.Empty(t, out.Mutations.BodyReplace,
"an Anthropic-shaped body must reach the upstream unmodified")
var endUser string
for _, kv := range out.Mutations.HeadersAdd {
if kv.Key == "x-litellm-end-user-id" {
endUser = kv.Value
}
}
assert.Equal(t, "alice@example.com", endUser,
"header stamping must still carry identity when body inject is skipped")
}
// TestInject_OpenAIBodyStillRewritten guards the gate against
// over-reaching: the OpenAI surface must keep its body-level identity,
// which is the only path LiteLLM's tag-budget check reads.
func TestInject_OpenAIBodyStillRewritten(t *testing.T) {
mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
in := newInput(litellmProvider, "alice", []string{"grp-eng"})
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "openai"})
in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations)
require.NotEmpty(t, out.Mutations.BodyReplace, "the OpenAI surface still gets body tags")
var doc map[string]any
require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
meta, ok := doc["metadata"].(map[string]any)
require.True(t, ok, "metadata must be an object")
assert.NotEmpty(t, meta["tags"], "metadata.tags must still be written")
}

View File

@@ -84,15 +84,6 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew
return allowNoAttribution(), nil
}
// Model-listing and other non-inference endpoints carry no model, and
// management's per-model allowlist fails closed on an empty one. The
// router has already authorised the route against the caller's groups
// and the request consumes no tokens, so gating it on a model that
// cannot exist would only break gateway model discovery.
if lookupKV(in.Metadata, middleware.KeyLLMNonInference) == "true" {
return allowNoAttribution(), nil
}
providerID := lookupKV(in.Metadata, middleware.KeyLLMResolvedProviderID)
if providerID == "" {
// llm_router didn't emit a resolved provider id — usually
@@ -126,7 +117,7 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew
}
if resp.GetDecision() == "deny" {
return denyFromManagement(resp, lookupKV(in.Metadata, middleware.KeyLLMProvider)), nil
return denyFromManagement(resp), nil
}
return allowFromManagement(resp), nil
}
@@ -170,7 +161,7 @@ func allowFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.O
// envelope. The deny code surfaces verbatim through the framework's
// fixed JSON template; arbitrary middleware bytes can't reach the
// wire.
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse, surface string) *middleware.Output {
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
code := resp.GetDenyCode()
if code == "" {
code = "llm_policy.cap_exceeded"
@@ -185,7 +176,6 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse, surface string
DenyReason: &middleware.DenyReason{
Code: code,
Message: denyMessageForCode(code),
Surface: surface,
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},

View File

@@ -224,35 +224,3 @@ func TestMetadataKeys_Allowlist(t *testing.T) {
}
assert.ElementsMatch(t, want, keys)
}
// TestInvoke_NonInferenceSkipsPreflight covers gateway model discovery:
// GET /v1/models carries no model, and management's per-model allowlist
// fails closed on an empty one, so a pre-flight would deny discovery for
// exactly the accounts that use the model allowlist. The router marks the
// request non-inference after authorising the route, and the gate must
// then allow without calling management at all.
func TestInvoke_NonInferenceSkipsPreflight(t *testing.T) {
mgmt := &fakeMgmt{
checkResp: &proto.CheckLLMPolicyLimitsResponse{
Decision: "deny",
DenyCode: "llm_policy.model_blocked",
},
}
m := New(mgmt, nil)
out := runInvoke(t, m, &middleware.Input{
AccountID: "acc-1",
UserID: "user-bob",
UserGroups: []string{"grp-engineers"},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
{Key: middleware.KeyLLMNonInference, Value: "true"},
},
})
assert.Equal(t, middleware.DecisionAllow, out.Decision, "model-less endpoints must not be gated on a model")
assert.Nil(t, mgmt.checkReq, "no pre-flight may be sent for a non-inference request")
assert.Empty(t, lookupKV(out.Metadata, middleware.KeyLLMSelectedPolicyID),
"no policy is attributed when nothing was metered")
}

View File

@@ -1,13 +1,9 @@
package llm_request_parser
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
func TestParseBedrockPath(t *testing.T) {
@@ -40,25 +36,3 @@ func TestParseBedrockPath(t *testing.T) {
}
}
}
// TestInvoke_BedrockCountTokens covers the dedicated token-counting
// endpoint. Denying it does not break the client, it just pushes context
// counting back onto the inference endpoint, which is billable.
func TestInvoke_BedrockCountTokens(t *testing.T) {
mw := newMiddleware(t)
out, err := mw.Invoke(context.Background(), &middleware.Input{
URL: "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens",
Body: []byte(`{"input":{"converse":{"messages":[]}}}`),
})
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision)
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
require.True(t, ok, "count-tokens carries a model in the path and must emit it")
assert.Equal(t, "anthropic.claude-sonnet-4-5", model, "model must be normalized like any other action")
stream, _ := metaValue(t, out.Metadata, middleware.KeyLLMStream)
assert.Equal(t, "false", stream, "count-tokens never streams")
}

View File

@@ -61,8 +61,6 @@ func (middlewareImpl) MetadataKeys() []string {
middleware.KeyLLMRequestPromptRaw,
middleware.KeyLLMCaptureTruncated,
middleware.KeyLLMSessionID,
middleware.KeyLLMAgentID,
middleware.KeyLLMParentAgentID,
}
}
@@ -74,9 +72,9 @@ func (middlewareImpl) Close() error { return nil }
// Invoke detects the LLM provider, parses request facts, and emits
// metadata. Always returns DecisionAllow; never errors. Provider
// selection prefers the request path, falling back to the configured
// providerID (synthesiser-stamped on agent-network targets) so requests
// routed to a custom upstream URL still resolve.
// selection prefers the configured providerID (synthesiser-stamped on
// agent-network targets) so requests routed to a custom upstream URL
// still resolve. Falls back to URL sniffing when no providerID is set.
func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
out := &middleware.Output{Decision: middleware.DecisionAllow}
if in == nil {
@@ -94,14 +92,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
return m.invokeBedrock(in, br), nil
}
// A path that names an API surface wins over the configured providerID:
// a gateway record pinned to "openai" still serves Claude Code on
// /v1/messages, and reading that body with the OpenAI parser loses the
// Anthropic usage block and prices the request on the wrong surface.
// providerID stays the fallback for upstreams whose path says nothing.
parser, ok := llm.DetectParser(extractPath(in.URL))
parser, ok := llm.ParserByName(m.providerID)
if !ok {
parser, ok = llm.ParserByName(m.providerID)
parser, ok = llm.DetectParser(extractPath(in.URL))
}
if !ok {
return out, nil
@@ -123,9 +116,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
}
appendSessionID := func(md []middleware.KV) []middleware.KV {
if sessionID != "" {
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
}
return appendAgentIDs(md, in.Headers)
return md
}
facts, err := parser.ParseRequest(in.Body)
@@ -167,41 +160,6 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
return out, nil
}
// agentIDHeader and parentAgentIDHeader carry sub-agent attribution: a
// coding agent that spawns helpers stamps the spawned agent's id, plus the
// spawning agent's when that helper is itself nested. Both are opaque
// identifiers rather than content, so they're emitted regardless of the
// prompt-collection toggle, the same way the session id is.
const (
agentIDHeader = "x-claude-code-agent-id"
parentAgentIDHeader = "x-claude-code-parent-agent-id"
)
// appendAgentIDs stamps the sub-agent attribution headers onto the metadata
// bag, skipping either one the request doesn't carry.
func appendAgentIDs(md []middleware.KV, headers []middleware.KV) []middleware.KV {
for _, pair := range []struct{ key, header string }{
{middleware.KeyLLMAgentID, agentIDHeader},
{middleware.KeyLLMParentAgentID, parentAgentIDHeader},
} {
if v := headerValue(headers, pair.header); v != "" {
md = append(md, middleware.KV{Key: pair.key, Value: v})
}
}
return md
}
// headerValue returns the first non-empty value for the named header.
// Headers arrive in canonical form, so the match is case-insensitive.
func headerValue(headers []middleware.KV, want string) string {
for _, kv := range headers {
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
return kv.Value
}
}
return ""
}
// sessionIDHeaders are request header names that may carry a client
// session identifier, checked in order, case-insensitively. Matching is
// against Go's canonical header form, so use the hyphenated names the
@@ -215,8 +173,10 @@ var sessionIDHeaders = []string{"x-claude-code-session-id", "session-id", "x-ses
// canonical form, so the match is case-insensitive.
func sessionIDFromHeaders(headers []middleware.KV) string {
for _, want := range sessionIDHeaders {
if v := headerValue(headers, want); v != "" {
return v
for _, kv := range headers {
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
return kv.Value
}
}
}
return ""
@@ -292,12 +252,6 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) {
if c := strings.LastIndex(rest, ":"); c >= 0 {
model, action = rest[:c], rest[c+1:]
}
// Token counting hangs off the model as its own path segment
// (".../models/{model}/count-tokens:rawPredict"), so anything past the
// first "/" belongs to the method rather than the model id.
if slash := strings.Index(model, "/"); slash >= 0 {
model = model[:slash]
}
model = llm.NormalizeVertexModel(model)
if model == "" {
return vertexRequest{}, false
@@ -344,7 +298,6 @@ func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *mi
if sessionID != "" {
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
}
md = appendAgentIDs(md, in.Headers)
promptTruncated := false
if parser != nil && m.capturePrompt {
@@ -392,9 +345,7 @@ func trimBedrockNamespace(reqPath string) string {
//
// /model/{modelId}/{action}
//
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream,
// count-tokens}. Token counting carries a model and no usage, so it routes
// like any other action and meters to zero.
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}.
// The modelId may be URL-encoded and may carry a cross-region inference-profile
// prefix and a version suffix; normalizeBedrockModel strips both so the model
// matches catalog pricing.
@@ -418,7 +369,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
return bedrockRequest{}, false
}
switch action {
case "invoke", "converse", "count-tokens":
case "invoke", "converse":
return bedrockRequest{model: model}, true
case "invoke-with-response-stream", "converse-stream":
return bedrockRequest{model: model, stream: true}, true
@@ -446,7 +397,6 @@ func (m middlewareImpl) invokeBedrock(in *middleware.Input, br bedrockRequest) *
if sessionID != "" {
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
}
md = appendAgentIDs(md, in.Headers)
promptTruncated := false
if parser != nil && m.capturePrompt {

View File

@@ -45,8 +45,6 @@ func TestMiddleware_StaticSurface(t *testing.T) {
middleware.KeyLLMRequestPromptRaw,
middleware.KeyLLMCaptureTruncated,
middleware.KeyLLMSessionID,
middleware.KeyLLMAgentID,
middleware.KeyLLMParentAgentID,
}
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
}
@@ -232,31 +230,6 @@ func TestInvoke_ProviderIDConfigBypassesURLSniff(t *testing.T) {
assert.Equal(t, "gpt-4o-mini", model)
}
func TestInvoke_PathSurfaceBeatsProviderIDConfig(t *testing.T) {
// Gateway records (LiteLLM, Portkey, OpenRouter) pin provider_id
// "openai", but the same record serves Claude Code on /v1/messages.
// Parsing that body as OpenAI reads no usage off the Anthropic
// response and prices the request on a surface where no claude-*
// model exists, so the path has to win.
mw, err := Factory{}.New([]byte(`{"provider_id":"openai"}`))
require.NoError(t, err, "factory must accept provider_id config")
out, err := mw.Invoke(context.Background(), &middleware.Input{
URL: "/v1/messages",
Body: []byte(`{"model":"claude-sonnet-5","stream":true,"messages":[{"role":"user","content":"Hi"}]}`),
})
require.NoError(t, err)
require.NotNil(t, out)
provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider)
require.True(t, ok, "provider must be emitted")
assert.Equal(t, "anthropic", provider, "the /v1/messages path selects the Anthropic surface")
model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel)
require.True(t, ok, "model must be extracted")
assert.Equal(t, "claude-sonnet-5", model)
}
func TestInvoke_UnknownProviderIDFallsBackToURL(t *testing.T) {
mw, err := Factory{}.New([]byte(`{"provider_id":"not-a-real-parser"}`))
require.NoError(t, err, "factory must accept any provider_id string")
@@ -443,81 +416,3 @@ func TestInvoke_NilInputAllows(t *testing.T) {
assert.Equal(t, middleware.DecisionAllow, out.Decision, "nil input still allows")
assert.Empty(t, out.Metadata, "nil input emits no metadata")
}
// TestParseVertexPath_CountTokensKeepsModel covers Vertex token counting,
// where the method hangs off the model as its own path segment. Splitting
// only on the final colon swallowed "/count-tokens" into the model id, so
// the router saw a model no route could claim.
func TestParseVertexPath_CountTokensKeepsModel(t *testing.T) {
cases := map[string]struct {
model string
stream bool
}{
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:rawPredict": {model: "claude-sonnet-5"},
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:streamRawPredict": {model: "claude-sonnet-5", stream: true},
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5/count-tokens:rawPredict": {model: "claude-sonnet-5"},
"/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5@20250929/count-tokens:rawPredict": {model: "claude-sonnet-5"},
}
for path, want := range cases {
vx, ok := parseVertexPath(path)
require.True(t, ok, "must parse %q", path)
assert.Equal(t, want.model, vx.model, "model for %q", path)
assert.Equal(t, want.stream, vx.stream, "stream flag for %q", path)
assert.Equal(t, "anthropic", vx.publisher, "publisher for %q", path)
}
}
// TestInvoke_EmitsAgentIDs covers sub-agent attribution: several agents run
// in parallel inside one session, and without their ids every request in
// the session attributes to the session alone.
func TestInvoke_EmitsAgentIDs(t *testing.T) {
mw := newMiddleware(t)
t.Run("spawned agent", func(t *testing.T) {
out, err := mw.Invoke(context.Background(), &middleware.Input{
URL: "/v1/messages",
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
Headers: []middleware.KV{
{Key: "X-Claude-Code-Session-Id", Value: "sess-1"},
{Key: "X-Claude-Code-Agent-Id", Value: "agent-7"},
},
})
require.NoError(t, err)
agent, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
require.True(t, ok, "the spawned agent's id must be emitted")
assert.Equal(t, "agent-7", agent)
_, ok = metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID)
assert.False(t, ok, "a top-level agent has no parent to emit")
})
t.Run("nested agent", func(t *testing.T) {
out, err := mw.Invoke(context.Background(), &middleware.Input{
URL: "/v1/messages",
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
Headers: []middleware.KV{
{Key: "X-Claude-Code-Agent-Id", Value: "agent-9"},
{Key: "X-Claude-Code-Parent-Agent-Id", Value: "agent-7"},
},
})
require.NoError(t, err)
agent, _ := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
assert.Equal(t, "agent-9", agent)
parent, ok := metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID)
require.True(t, ok, "a nested agent must carry the spawning agent's id")
assert.Equal(t, "agent-7", parent)
})
t.Run("absent on a plain request", func(t *testing.T) {
out, err := mw.Invoke(context.Background(), &middleware.Input{
URL: "/v1/messages",
Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`),
})
require.NoError(t, err)
_, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID)
assert.False(t, ok, "no key is emitted when the client sends no agent id")
})
}

View File

@@ -1,175 +0,0 @@
package llm_router
import (
"context"
"net/http"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
// bedrockRoute is a Bedrock provider whose listing lives on the control plane
// while inference goes to the runtime host — the split this file is about.
func bedrockRoute(models []string, policies []ModelPolicyRule) ProviderRoute {
return ProviderRoute{
ID: "prov-bedrock",
Bedrock: true,
Models: models,
ModelPolicies: policies,
UpstreamScheme: "https",
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
DiscoveryHost: "bedrock.eu-central-1.amazonaws.com",
AuthHeaderName: "Authorization",
AuthHeaderValue: "Bearer aws-token",
AllowedGroupIDs: []string{defaultTestGroup},
}
}
func getInput(path string) *middleware.Input {
return &middleware.Input{
Slot: middleware.SlotOnRequest,
Method: http.MethodGet,
URL: "https://endpoint.netbird.local" + path,
UserGroups: []string{defaultTestGroup},
}
}
// TestBedrockListingGoesToTheControlPlane is the whole point of DiscoveryHost.
// ListInferenceProfiles is not an operation bedrock-runtime implements — it
// answers <UnknownOperationException/> — so a listing forwarded to the
// inference upstream can only 404, however well it is routed.
func TestBedrockListingGoesToTheControlPlane(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{bedrockRoute(nil, nil)}})
out, err := mw.Invoke(context.Background(), getInput("/inference-profiles"))
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "bedrock.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host)
}
// TestBedrockInferenceStillGoesToTheRuntimeHost is the other half: the
// redirect must apply to the listing alone. Sending an InvokeModel call to the
// control plane would break every Bedrock request in the account.
func TestBedrockInferenceStillGoesToTheRuntimeHost(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{bedrockRoute(nil, nil)}})
in := newInputWithModelAndURL("anthropic.claude-haiku-4-5",
"https://endpoint.netbird.local/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/invoke")
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host)
}
// TestIsListingPath guards the narrower reading of "model-less". Both the
// upstream redirect and the policy bound key on this, and the warming probe
// must be excluded from both: it carries no listing to filter, and pointing it
// at the control plane would warm a pool the inference requests never use.
func TestIsListingPath(t *testing.T) {
for path, want := range map[string]bool{
"/v1/models": true,
"/inference-profiles": true,
"/bedrock/inference-profiles": true,
"/api/hello": false,
"/v1/models/gpt-4o": false, // the per-model lookup, routed elsewhere
"/v1/chat/completions": false,
} {
t.Run(path, func(t *testing.T) {
assert.Equal(t, want, isListingPath(path))
})
}
}
// TestBedrockListingIsBoundByPolicy covers the case that was previously
// unreachable: filtering keyed on /v1/models alone, so a Bedrock listing was
// routed but never narrowed to what the caller may use.
func TestBedrockListingIsBoundByPolicy(t *testing.T) {
route := bedrockRoute(
[]string{"eu.anthropic.claude-haiku-4-5-20251001-v1:0", "eu.anthropic.claude-sonnet-4-6"},
[]ModelPolicyRule{{
GroupIDs: []string{defaultTestGroup},
// A guardrail allowlist names the catalog key, which is the form an
// operator picks in the UI — not the region-prefixed wire id the
// record registers.
Models: []string{"anthropic.claude-haiku-4-5"},
}},
)
mw := New(Config{Providers: []ProviderRoute{route}})
out, err := mw.Invoke(context.Background(), getInput("/inference-profiles"))
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
// Exact-string intersection would find nothing here and bound the listing
// to empty, handing the caller a picker with no models on a provider that
// works perfectly well.
assert.Equal(t, []string{"eu.anthropic.claude-haiku-4-5-20251001-v1:0"},
out.Mutations.RewriteUpstream.DiscoveryModels)
}
// TestBedrockListingWithoutADiscoveryHostFallsThrough keeps a proxied or
// self-hosted Bedrock endpoint working: the synthesiser emits no discovery
// host for one, and the listing must then go to the configured upstream rather
// than nowhere.
func TestBedrockListingWithoutADiscoveryHostFallsThrough(t *testing.T) {
route := bedrockRoute(nil, nil)
route.UpstreamHost = "bedrock.internal.example.com"
route.DiscoveryHost = ""
mw := New(Config{Providers: []ProviderRoute{route}})
out, err := mw.Invoke(context.Background(), getInput("/inference-profiles"))
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "bedrock.internal.example.com", out.Mutations.RewriteUpstream.Host)
}
// TestBedrockProfileDetailHonoursTheModelTable covers GetInferenceProfile,
// which the listing filter cannot help with: it answers for one profile with a
// single object, not a set, so nothing narrows it on the way back. Authorising
// it by provider type alone would let any caller with a Bedrock route read the
// full configuration of every profile in the account.
//
// Both registration spellings are exercised, because a record may carry the
// raw profile id AWS issues or the catalog key it reduces to.
func TestBedrockProfileDetailHonoursTheModelTable(t *testing.T) {
const permitted = "eu.anthropic.claude-sonnet-5-20260514-v1:0"
for _, registered := range []string{permitted, "anthropic.claude-sonnet-5"} {
t.Run(registered, func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{bedrockRoute([]string{registered}, nil)}})
out, err := mw.Invoke(context.Background(), getInput("/inference-profiles/"+permitted))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"a profile the record registers must still resolve")
denied, err := mw.Invoke(context.Background(),
getInput("/inference-profiles/eu.anthropic.claude-opus-5-20260514-v1:0"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, denied.Decision,
"a profile outside the record's models must not be readable")
})
}
}
// TestBedrockProfileListingStaysModelLess pins the other half: the listing
// names no profile, so it must not be judged against the model table. It is
// bounded by DiscoveryModels in the response instead, and denying it here
// would take model discovery away from exactly the records that enumerate
// their models.
func TestBedrockProfileListingStaysModelLess(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{bedrockRoute([]string{"anthropic.claude-sonnet-5"}, nil)}})
out, err := mw.Invoke(context.Background(), getInput("/inference-profiles"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision)
}

View File

@@ -1,13 +1,9 @@
package llm_router
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
// TestRouteClaimsModel_BedrockNormalizesCandidate guards the fix for the native
@@ -32,86 +28,3 @@ func TestRouteClaimsModel_BedrockNormalizesCandidate(t *testing.T) {
assert.False(t, routeClaimsModel(openai, "us.gpt-4o"),
"non-Bedrock routes must not strip a us. prefix")
}
// TestRouter_BedrockCountTokensRoutes pins that the token-counting action
// reaches the Bedrock route instead of denying as not-routable.
func TestRouter_BedrockCountTokensRoutes(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{{
ID: "bedrock-prod",
Bedrock: true,
Models: []string{"anthropic.claude-sonnet-4-5"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
}}})
in := newInputWithModelAndURL("anthropic.claude-sonnet-4-5",
"/model/anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens")
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "bedrock"})
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "count-tokens must route, not deny")
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host)
}
// TestRouter_BedrockInferenceProfilesRoutes covers the startup lookups a
// client makes to resolve a configured inference profile. They carry no
// model, so before they were recognised they denied and wrote a policy
// rejection into the access log on every session start.
func TestRouter_BedrockInferenceProfilesRoutes(t *testing.T) {
bedrock := ProviderRoute{
ID: "bedrock-prod",
Bedrock: true,
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
}
openai := ProviderRoute{
ID: "openai-prod",
Models: []string{"gpt-4o"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "api.openai.com",
}
mw := New(Config{Providers: []ProviderRoute{openai, bedrock}})
for _, path := range []string{
"/inference-profiles?type=SYSTEM_DEFINED",
"/inference-profiles/us.anthropic.claude-sonnet-5",
} {
out, err := mw.Invoke(context.Background(), newModellessInput(path))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "%s must route", path)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host,
"%s must reach the Bedrock provider, not the first authorised one", path)
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
assert.Equal(t, "true", nonInference, "%s carries no model to gate on", path)
}
}
// TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix pins that the
// optional gateway namespace is removed before the request goes upstream.
func TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{{
ID: "bedrock-prod",
Bedrock: true,
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com",
}}})
out, err := mw.Invoke(context.Background(), newModellessInput("/bedrock/inference-profiles"))
require.NoError(t, err)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "/bedrock", out.Mutations.RewriteUpstream.StripPathPrefix,
"the namespace prefix must not reach the real Bedrock endpoint")
}

View File

@@ -44,19 +44,6 @@ type ProviderRoute struct {
AuthHeaderName string `json:"auth_header_name"`
AuthHeaderValue string `json:"auth_header_value"`
AllowedGroupIDs []string `json:"allowed_group_ids"`
// ModelPolicies carries, per authorising policy, the source groups it
// binds and the models it permits. The router uses it to bound a model
// listing to what THIS caller may use: a provider reachable by two groups
// under different allowlists must not offer either group the other's
// models. Empty means no policy restricts models on this route.
ModelPolicies []ModelPolicyRule `json:"model_policies,omitempty"`
// DiscoveryHost, when set, is the host that serves this provider's model
// listing, for a vendor that does not serve it from the same host as
// inference. Bedrock is why it exists: ListInferenceProfiles is a control
// plane operation on bedrock.<region>, while InvokeModel must go to
// bedrock-runtime.<region>, so one record genuinely needs two hosts.
// Empty means the listing is served from UpstreamHost like everything else.
DiscoveryHost string `json:"discovery_host,omitempty"`
// Vertex marks a Google Vertex AI provider. Vertex requests carry the
// model in the URL path, so the router selects this route by path
// (isVertexPath) and bypasses the model/vendor table entirely.
@@ -78,18 +65,6 @@ type ProviderRoute struct {
SkipTLSVerify bool `json:"skip_tls_verify,omitempty"`
}
// ModelPolicyRule is one authorising policy's contribution to what a caller
// may use on a route: the source groups it binds, and the models it permits.
//
// Models is nil when the policy sets no model allowlist — an unrestricted
// policy, which lifts the restriction for the groups it binds. That is why
// nil and empty must stay distinct: an empty list is a guardrail that permits
// nothing, and collapsing the two would let a listing fail open.
type ModelPolicyRule struct {
GroupIDs []string `json:"group_ids"`
Models []string `json:"models"`
}
// Config is the on-wire configuration accepted by the factory. An
// empty Providers slice yields a router that denies every request as
// not-routable; the synthesiser is responsible for stamping the

View File

@@ -109,10 +109,6 @@ func (m *Middleware) MetadataKeys() []string {
middleware.KeyLLMAuthorisingGroups,
middleware.KeyLLMPolicyDecision,
middleware.KeyLLMPolicyReason,
middleware.KeyLLMNonInference,
// Emitted only for the per-model lookup, whose model lives in the path
// rather than a body the parser could read.
middleware.KeyLLMModel,
}
}
@@ -141,26 +137,29 @@ const (
// known to a provider that no policy authorises for the caller deny
// with no_authorised_provider.
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
reqPath := requestPath(in.URL)
// The caller's API dialect, used to mirror a denial in the vendor's own
// error shape so the client can explain it to the user.
surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
// Vertex AI carries the model in the URL path, not the body, and is
// selected by path rather than by the model/vendor table. Route it before
// the model lookup so a model the parser extracted from the path can't be
// claimed by a same-vendor direct provider (e.g. claude-* on api.anthropic.com).
reqPath := requestPath(in.URL)
if isVertexPath(reqPath) {
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
// The request parser emits no llm.provider for a Vertex publisher it
// can't parse (e.g. google/gemini). Forwarding such a request would
// bypass token/budget metering, so deny it rather than serve it
// unmetered.
if surface == "" {
return denyUnmeterable(surface), nil
if vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider); vendor == "" {
return denyUnmeterable(), nil
}
route, outcome := m.matchVertex(reqPath, model, in.UserGroups)
return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil
switch outcome {
case matchOutcomeFound:
return m.allowWithRoute(route, in.UserGroups), nil
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(model), nil
default:
return denyUnknownModel(model), nil
}
}
// Bedrock likewise carries the model in the URL path (/model/{id}/{action}),
@@ -168,231 +167,52 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
// before the model lookup; when the prefix is present, strip it from the
// forwarded path so the real Bedrock endpoint receives its native path.
if isBedrockPath(reqPath) {
model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
native, hadPrefix := splitBedrockNamespace(reqPath)
route, outcome := m.matchBedrock(native, model, in.UserGroups)
return m.decide(route, outcome, surface, model, in.UserGroups, func(out *middleware.Output) {
if hadPrefix {
stripBedrockNamespace(out)
switch outcome {
case matchOutcomeFound:
out := m.allowWithRoute(route, in.UserGroups)
if hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
}
}), nil
return out, nil
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(model), nil
default:
return denyUnknownModel(model), nil
}
}
// GET /v1/models/{id} carries no body, so no model reaches the router in
// metadata — but the path names one, and answering it confirms a model
// exists and is reachable. Authorise it against the model table like any
// other per-model request, then mark it non-inference so it still skips
// the token pre-flight it would otherwise charge nothing against.
if detail, isDetail := modelDetailID(reqPath); isDetail && isNonInferenceMethod(in.Method) {
route, outcome := m.matchRoute(detail, surface, reqPath, in.UserGroups)
return m.decide(route, outcome, surface, detail, in.UserGroups, func(out *middleware.Output) {
markNonInference(out)
// The parser reads models from JSON bodies only, and this request
// has none, so stamp the one the path names. Without it the
// guardrail's own allowlist — a separate, possibly narrower list
// than the route's — never sees a model to check.
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMModel, Value: detail})
}), nil
model, ok := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
if !ok || model == "" {
// Non-inference endpoints (model listing) carry no model but still
// need rewriting from the synth placeholder to a real upstream;
// clients such as Codex call GET /v1/models at startup to enumerate
// availability and read a 403 as "model unavailable".
route, outcome := m.matchModelless(requestPath(in.URL), in.UserGroups)
switch outcome {
case matchOutcomeFound:
return m.allowWithRoute(route, in.UserGroups), nil
case matchOutcomeUnauthorised:
// A recognised model-less endpoint exists but no provider
// authorises the caller — deny as an authorisation failure
// rather than masking it as a missing model.
return denyNoAuthorisedRoute(model), nil
default:
return denyMissingModel(), nil
}
}
if model == "" {
return m.routeModelless(reqPath, surface, in.Method, in.UserGroups), nil
}
route, outcome := m.matchRoute(model, surface, reqPath, in.UserGroups)
return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil
}
// decide turns a per-model match result into the middleware's decision. Every
// surface that routes by model shares the same two denial arms — a model no
// route claims is not routable, one that some route claims but none authorises
// for this caller is an authorisation failure — so they live here once.
// decorate, when non-nil, adjusts the allow with whatever that surface needs.
func (m *Middleware) decide(
route ProviderRoute,
outcome matchOutcome,
surface, model string,
userGroups []string,
decorate func(*middleware.Output),
) *middleware.Output {
vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
route, outcome := m.matchRoute(model, vendor, requestPath(in.URL), in.UserGroups)
switch outcome {
case matchOutcomeFound:
out := m.allowWithRoute(route, surface, userGroups)
if decorate != nil {
decorate(out)
}
return out
return m.allowWithRoute(route, in.UserGroups), nil
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(surface, model)
return denyNoAuthorisedRoute(model), nil
default:
return denyUnknownModel(surface, model)
}
}
// routeModelless serves the endpoints that name no model at all: the model
// listing, the connection-warming probe, and the Bedrock inference-profile
// lookup. They still need rewriting from the synth placeholder to a real
// upstream — clients such as Codex call GET /v1/models at startup to enumerate
// availability and read a 403 as "model unavailable".
func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups []string) *middleware.Output {
route, outcome := m.matchModelless(reqPath, method, userGroups)
switch outcome {
case matchOutcomeFound:
out := m.allowWithRoute(route, surface, userGroups)
markNonInference(out)
if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix {
stripBedrockNamespace(out)
}
if isListingPath(reqPath) && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
// A vendor that serves its listing from somewhere other than its
// inference upstream is redirected here, and only for the listing
// — every other request still goes to the configured upstream.
if route.DiscoveryHost != "" {
out.Mutations.RewriteUpstream.Host = route.DiscoveryHost
}
// What the caller may actually use bounds what the picker may
// offer: every entry outside it is a request the chain will deny a
// moment later.
if models, bounded := discoverableModels(route, userGroups); bounded {
out.Mutations.RewriteUpstream.DiscoveryModels = models
}
}
return out
case matchOutcomeUnauthorised:
// A recognised model-less endpoint exists but no provider authorises
// the caller — deny as an authorisation failure rather than masking it
// as a missing model.
return denyNoAuthorisedRoute(surface, "")
default:
return denyMissingModel(surface)
}
}
// isNonInferenceMethod reports whether a request method is one the
// non-inference endpoints actually use: the listing and the per-model lookup
// are GET, the connection-warming probe is HEAD or GET. The method is the only
// thing separating "GET /v1/models/{id}" from a POST to the same path carrying
// an inference body, and the non-inference mark exempts a request from the
// token pre-flight — so anything else falls through to normal per-model
// routing, which denies when the request names no model.
func isNonInferenceMethod(method string) bool {
return method == http.MethodGet || method == http.MethodHead
}
// discoverableModels returns the model ids a caller in userGroups may actually
// use on this route, and whether the listing should be bounded to them at all.
//
// Two things narrow a listing, and both must apply or the picker offers models
// the very next request refuses:
//
// - the provider's own enumerated models, when it lists any (a gateway record
// enumerates nothing and claims everything);
// - the model allowlists of the policies that authorise THIS caller. A
// provider reachable by two groups under different allowlists must not
// offer either group the other's models, which is why the rules carry their
// source groups rather than arriving pre-flattened.
//
// A policy that sets no allowlist lifts the restriction for the groups it
// binds, so a caller holding one unrestricted policy sees the provider's full
// list. bounded is false when nothing narrows the listing — an unrestricted
// caller on a route that enumerates nothing — in which case the upstream's own
// answer passes through untouched.
func discoverableModels(route ProviderRoute, userGroups []string) ([]string, bool) {
permitted, restricted := policyPermittedModels(route, userGroups)
switch {
case !restricted && len(route.Models) == 0:
return nil, false
case !restricted:
return append([]string(nil), route.Models...), true
case len(route.Models) == 0:
// A gateway record enumerates nothing, so the allowlist is the whole
// bound — previously such a record offered the upstream's entire
// catalogue however narrow the policy was.
return sortedModels(permitted), true
}
// Both bound: only what the provider serves and the policy permits.
intersection := make(map[string]struct{}, len(route.Models))
for _, m := range route.Models {
if _, ok := permitted[m]; ok {
intersection[m] = struct{}{}
continue
}
// The two sides are not always written the same way. A Bedrock record
// may register the raw inference-profile id an operator copied from
// AWS while a guardrail allowlist names the catalog key, and comparing
// those verbatim finds nothing — which would bound a correctly
// configured provider's listing down to empty. routeClaimsModel
// already normalises the candidate for exactly this reason, and the
// listing bound has to agree with it or the picker disagrees with what
// the guardrail will actually allow.
if route.Bedrock {
if _, ok := permitted[llm.NormalizeBedrockModel(m)]; ok {
intersection[m] = struct{}{}
}
}
}
return sortedModels(intersection), true
}
// policyPermittedModels folds the rules whose groups intersect the caller's
// into the set of models they permit. restricted is false when the caller
// holds at least one authorising policy that sets no allowlist, or when no
// rule binds them at all.
func policyPermittedModels(route ProviderRoute, userGroups []string) (map[string]struct{}, bool) {
permitted := make(map[string]struct{})
restricted := false
for _, rule := range route.ModelPolicies {
if !groupsIntersect(rule.GroupIDs, userGroups) {
continue
}
if rule.Models == nil {
// An unrestricted policy the caller holds lifts the restriction
// entirely, whatever the others say.
return nil, false
}
restricted = true
for _, m := range rule.Models {
permitted[m] = struct{}{}
}
}
return permitted, restricted
}
// groupsIntersect reports whether the two group-id sets share a member.
func groupsIntersect(a, b []string) bool {
for _, x := range a {
for _, y := range b {
if x == y {
return true
}
}
}
return false
}
// sortedModels flattens a model set into a stable slice so the bound the proxy
// applies — and any test asserting on it — does not depend on map order.
func sortedModels(set map[string]struct{}) []string {
out := make([]string, 0, len(set))
for m := range set {
out = append(out, m)
}
sort.Strings(out)
return out
}
// markNonInference tags an allow as a request that spends no tokens, so the
// limit check skips the management pre-flight it would charge nothing against.
func markNonInference(out *middleware.Output) {
out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"})
}
// stripBedrockNamespace tells the rewrite to drop the optional "/bedrock"
// gateway namespace so the upstream receives its native Bedrock path.
func stripBedrockNamespace(out *middleware.Output) {
if out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
return denyUnknownModel(model), nil
}
}
@@ -480,91 +300,12 @@ func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []stri
return best, matchOutcomeFound
}
// connectionWarmPath is the probe Anthropic clients send before their first
// inference request to open the upstream connection early. Forwarding it
// warms the connection the request will actually use; denying it only fills
// the access log with rejections at every session start.
const connectionWarmPath = "/api/hello"
// modelListingPath is the endpoint clients read at startup to populate
// their model picker. Its response is a list the proxy can bound; the
// per-model "/v1/models/{id}" lookup returns a single object and is left
// alone.
const modelListingPath = "/v1/models"
// isListingPath reports whether reqPath asks for a MODEL LISTING, as opposed
// to the other model-less endpoints. Only a listing gets an upstream redirect
// and a policy bound: the connection-warming probe carries no model list to
// filter, and rewriting its host would send the warm-up to the wrong pool.
func isListingPath(reqPath string) bool {
return reqPath == modelListingPath || isBedrockModelLessPath(reqPath)
}
// isModelLessPath reports whether reqPath is a known non-inference endpoint
// that legitimately carries no model at all: the model listing and the
// connection-warming probe. These must route to an upstream rather than
// deny, so model enumeration works end to end. The per-model
// "/v1/models/{id}" lookup is deliberately excluded — it names a model, so
// it is authorised against the model table instead (see modelDetailID).
// isModelLessPath reports whether reqPath is a known OpenAI-shaped
// non-inference endpoint that legitimately carries no model in its
// request (the model-listing endpoints). These must route to an upstream
// rather than deny, so model enumeration works end to end.
func isModelLessPath(reqPath string) bool {
return reqPath == modelListingPath || reqPath == connectionWarmPath
}
// modelDetailID returns the model id named by a "/v1/models/{id}" lookup.
// reqPath comes from url.URL.Path, which is already percent-decoded, so an
// id carrying a "/" (a self-hosted "Qwen/Qwen2.5-0.5B-Instruct" sent as
// "Qwen%2FQwen2.5-...") arrives whole and everything after the prefix is the
// id, separators included.
func modelDetailID(reqPath string) (string, bool) {
if !strings.HasPrefix(reqPath, modelListingPath+"/") {
return "", false
}
id := strings.TrimPrefix(reqPath, modelListingPath+"/")
if id == "" {
return "", false
}
return id, true
}
// isBedrockModelLessPath reports whether reqPath is a Bedrock
// inference-profile lookup, optionally behind the "/bedrock" gateway
// namespace. Clients read these at startup to resolve a configured profile
// to its underlying model. They carry no model of their own, so they route
// by path to a Bedrock provider rather than through the model table.
//
// On native AWS these live on the control plane ("bedrock.<region>") while a
// provider's upstream is normally the runtime host ("bedrock-runtime.<region>"),
// so forwarding yields a 404 there. That is deliberate: a client has one base
// URL, so pointing it straight at the runtime host 404s identically, and
// forwarding keeps the proxy transparent instead of inventing a policy denial
// the client would never otherwise see. Operators whose Bedrock upstream is a
// gateway that does serve the lookup get a working answer.
func isBedrockModelLessPath(reqPath string) bool {
native, _ := splitBedrockNamespace(reqPath)
return native == "/inference-profiles" || strings.HasPrefix(native, bedrockProfileDetailPrefix)
}
// bedrockProfileDetailPrefix precedes the identifier in a GetInferenceProfile
// lookup, once any gateway namespace is off the front.
const bedrockProfileDetailPrefix = "/inference-profiles/"
// bedrockProfileID returns the inference profile a "/inference-profiles/{id}"
// lookup names. The listing beside it names none, which is what separates the
// two: a listing is a set the response filter can bound, while this answers
// for one profile with a single object no filter inspects.
//
// The id arrives as AWS issues it — region prefix and version suffix included
// — because that is the only form that works at invoke time.
func bedrockProfileID(reqPath string) (string, bool) {
native, _ := splitBedrockNamespace(reqPath)
if !strings.HasPrefix(native, bedrockProfileDetailPrefix) {
return "", false
}
id := strings.TrimPrefix(native, bedrockProfileDetailPrefix)
if id == "" {
return "", false
}
return id, true
return reqPath == "/v1/models" || strings.HasPrefix(reqPath, "/v1/models/")
}
// isVertexPath reports whether reqPath is a Google Vertex AI publisher
@@ -591,33 +332,20 @@ func splitBedrockNamespace(reqPath string) (string, bool) {
return reqPath, false
}
// bedrockActions are the runtime actions that follow the model id in a
// Bedrock path. count-tokens is here so a client can price its context
// against the dedicated endpoint; denying it pushes that work back onto
// the inference endpoint, which bills for it.
var bedrockActions = []string{
"/invoke",
"/invoke-with-response-stream",
"/converse",
"/converse-stream",
"/count-tokens",
}
// isBedrockPath reports whether reqPath is an AWS Bedrock runtime model
// endpoint: /model/{modelId}/{action} — optionally behind a "/bedrock"
// gateway-namespace prefix. The model lives in the path, so these requests
// are routed by path to the Bedrock provider.
// endpoint: /model/{modelId}/{action} where action is invoke,
// invoke-with-response-stream, converse, or converse-stream — optionally behind
// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these
// requests are routed by path to the Bedrock provider.
func isBedrockPath(reqPath string) bool {
native, _ := splitBedrockNamespace(reqPath)
if !strings.HasPrefix(native, "/model/") {
return false
}
for _, action := range bedrockActions {
if strings.HasSuffix(native, action) {
return true
}
}
return false
return strings.HasSuffix(native, "/invoke") ||
strings.HasSuffix(native, "/invoke-with-response-stream") ||
strings.HasSuffix(native, "/converse") ||
strings.HasSuffix(native, "/converse-stream")
}
// matchVertex selects the Vertex provider authorised for the caller's groups
@@ -697,42 +425,19 @@ func (m *Middleware) matchPathRoute(reqPath, model string, userGroups []string,
// declaration order), matchOutcomeUnauthorised when no provider authorises
// the caller, or matchOutcomeUnknownModel when the path isn't a recognised
// model-less endpoint.
func (m *Middleware) matchModelless(reqPath, method string, userGroups []string) (ProviderRoute, matchOutcome) {
if !isNonInferenceMethod(method) {
func (m *Middleware) matchModelless(reqPath string, userGroups []string) (ProviderRoute, matchOutcome) {
if !isModelLessPath(reqPath) {
return ProviderRoute{}, matchOutcomeUnknownModel
}
var eligible func(ProviderRoute) bool
switch {
case isBedrockModelLessPath(reqPath):
if profile, isDetail := bedrockProfileID(reqPath); isDetail {
// A detail lookup names one profile, so it is authorised like any
// other per-model request rather than by provider type alone. The
// listing beside it is bounded by DiscoveryModels on the way back,
// but this answers with a single object no filter inspects — so
// without the check here, a caller reads the full configuration of
// every profile in the account, including the ones its policy
// never named.
//
// The id is normalised first: a record may register the raw
// profile id or the catalog key it reduces to, and routeClaimsModel
// expects the normalised form an inference request would carry.
wanted := llm.NormalizeBedrockModel(profile)
eligible = func(r ProviderRoute) bool { return r.Bedrock && routeClaimsModel(r, wanted) }
} else {
eligible = func(r ProviderRoute) bool { return r.Bedrock }
}
case isModelLessPath(reqPath):
var candidates []ProviderRoute
for _, route := range m.cfg.Providers {
// Vertex/Bedrock are path-routed and don't serve OpenAI-style
// model-listing endpoints; including them here could rewrite a
// GET /v1/models to an upstream that 404s it.
eligible = func(r ProviderRoute) bool { return !r.Vertex && !r.Bedrock }
default:
return ProviderRoute{}, matchOutcomeUnknownModel
}
var candidates []ProviderRoute
for _, route := range m.cfg.Providers {
if eligible(route) && routeAuthorisesGroups(route, userGroups) {
if route.Vertex || route.Bedrock {
continue
}
if routeAuthorisesGroups(route, userGroups) {
candidates = append(candidates, route)
}
}
@@ -859,16 +564,6 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
return true
}
// A client may pin a dated Anthropic id ("claude-sonnet-4-5-20250929")
// where the operator registered the undated one. Only an undated
// registration absorbs a dated request: normalising both sides would
// let a route pinned to one dated release claim a different one, so an
// operator who deliberately pinned a build would silently serve
// another — and with several such routes, ordering would decide which.
if candidate == llm.NormalizeAnthropicModel(candidate) &&
candidate == llm.NormalizeAnthropicModel(model) {
return true
}
}
return false
}
@@ -917,7 +612,7 @@ func requestPath(raw string) string {
// provider id so identity-stamping middlewares (llm_identity_inject)
// tag the request with ONLY the groups that authorised this specific
// route — not every group the peer happens to be in.
func (m *Middleware) allowWithRoute(route ProviderRoute, surface string, userGroups []string) *middleware.Output {
func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *middleware.Output {
rewrite := &middleware.UpstreamRewrite{
Scheme: route.UpstreamScheme,
Host: route.UpstreamHost,
@@ -939,7 +634,7 @@ func (m *Middleware) allowWithRoute(route ProviderRoute, surface string, userGro
// request time (cached + auto-refreshed) instead of a static value.
bearer, err := m.gcpBearer(route.GCPServiceAccountKeyB64)
if err != nil {
return denyUpstreamAuth(surface)
return denyUpstreamAuth()
}
authValue = bearer
}
@@ -1009,12 +704,11 @@ func (m *Middleware) gcpTokenSource(saKeyB64 string) (oauth2.TokenSource, error)
// denyUpstreamAuth is returned when the router cannot obtain the upstream
// credential (e.g. a malformed service-account key or an unreachable token
// endpoint). It surfaces as a 502 — an upstream problem, not a policy denial.
func denyUpstreamAuth(surface string) *middleware.Output {
func denyUpstreamAuth() *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 502,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeUpstreamAuth,
Message: "could not obtain upstream credential",
},
@@ -1028,12 +722,11 @@ func denyUpstreamAuth(surface string) *middleware.Output {
// denyUnmeterable returns the deny envelope for a path-routed request whose
// publisher has no parser surface, so its usage can't be metered. Serving it
// would bypass token/budget caps, so it is rejected with a 403.
func denyUnmeterable(surface string) *middleware.Output {
func denyUnmeterable() *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeUnmeterable,
Message: "request publisher is not supported for metering",
},
@@ -1046,12 +739,11 @@ func denyUnmeterable(surface string) *middleware.Output {
// denyMissingModel returns the deny envelope for a request whose
// envelope has no llm.model metadata.
func denyMissingModel(surface string) *middleware.Output {
func denyMissingModel() *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeNotRoutable,
Message: "missing llm.model on request envelope",
},
@@ -1064,12 +756,11 @@ func denyMissingModel(surface string) *middleware.Output {
// denyUnknownModel returns the deny envelope for a model that no
// configured provider claims.
func denyUnknownModel(surface, model string) *middleware.Output {
func denyUnknownModel(model string) *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeNotRoutable,
Message: fmt.Sprintf("no provider configured for model %s", model),
Details: map[string]string{"model": model},
@@ -1084,12 +775,11 @@ func denyUnknownModel(surface, model string) *middleware.Output {
// denyNoAuthorisedRoute returns the deny envelope for a model that one
// or more providers claim, but where no policy authorises the caller's
// groups for any of those providers.
func denyNoAuthorisedRoute(surface, model string) *middleware.Output {
func denyNoAuthorisedRoute(model string) *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeNoAuthorisedRoute,
Message: fmt.Sprintf("no policy authorises model %s for the caller's groups", model),
Details: map[string]string{"model": model},

View File

@@ -2,7 +2,6 @@ package llm_router
import (
"context"
"net/http"
"testing"
"github.com/stretchr/testify/assert"
@@ -61,8 +60,6 @@ func TestMiddlewareIdentity(t *testing.T) {
[]string{
middleware.KeyLLMResolvedProviderID,
middleware.KeyLLMAuthorisingGroups,
middleware.KeyLLMNonInference,
middleware.KeyLLMModel,
middleware.KeyLLMPolicyDecision,
middleware.KeyLLMPolicyReason,
},
@@ -174,12 +171,8 @@ func TestRouter_MissingModel(t *testing.T) {
// from which a model could be parsed). UserGroups matches defaultTestGroup.
func newModellessInput(reqURL string) *middleware.Input {
return &middleware.Input{
Slot: middleware.SlotOnRequest,
URL: reqURL,
// The non-inference endpoints are read requests; the method is what
// separates them from an inference body posted to the same path, so
// state it rather than leaning on the zero value.
Method: http.MethodGet,
Slot: middleware.SlotOnRequest,
URL: reqURL,
UserGroups: []string{defaultTestGroup},
}
}
@@ -204,12 +197,6 @@ func TestRouter_ModelLessPath_RoutesToAuthorisedProvider(t *testing.T) {
provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID)
assert.Equal(t, "openai-prod", provider, "resolved provider must be the authorised route")
// The limits gate reads this to tell "no model applies here" from
// "the model could not be determined", which fails closed.
nonInference, ok := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
require.True(t, ok, "model-less allow must mark the request non-inference")
assert.Equal(t, "true", nonInference)
}
func TestRouter_ModelLessPath_MultiProviderDeclarationOrder(t *testing.T) {
@@ -886,403 +873,3 @@ func TestRouter_EmptyModelsClaimsAnyModel(t *testing.T) {
resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID)
assert.Equal(t, "litellm", resolved)
}
// TestRouter_DatedAnthropicModelRoutes covers a client pinning a release
// date on a model the operator registered undated. Exact matches still win,
// so an operator who registers both dated releases keeps them distinct.
func TestRouter_DatedAnthropicModelRoutes(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{{
ID: "anthropic-prod",
Vendor: "anthropic",
Models: []string{"claude-sonnet-4-5"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "api.anthropic.com",
}}})
in := newInputWithModelAndURL("claude-sonnet-4-5-20250929", "/v1/messages")
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "a dated id must route to the undated registration")
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
}
// TestRouter_ConnectionWarmProbeRoutes covers the HEAD /api/hello probe an
// Anthropic client sends before its first request. Forwarding it warms the
// connection that request will use; denying it only wrote a rejection into
// the access log at every session start.
func TestRouter_ConnectionWarmProbeRoutes(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{{
ID: "anthropic-prod",
Vendor: "anthropic",
Models: []string{"claude-sonnet-5"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "api.anthropic.com",
}}})
in := newModellessInput("/api/hello")
in.Method = http.MethodHead
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "the warm-up probe must reach the upstream")
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
assert.Equal(t, "true", nonInference, "the probe carries no model to gate on")
}
// TestRouter_ModelListingCarriesAuthorisedModels pins the list the proxy
// bounds the discovery response with. A catch-all route enumerates nothing,
// so it must not bound the upstream's list at all.
func TestRouter_ModelListingCarriesAuthorisedModels(t *testing.T) {
enumerated := ProviderRoute{
ID: "anthropic-prod",
Models: []string{"claude-sonnet-5", "claude-haiku-4-5"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "api.anthropic.com",
}
t.Run("enumerated route bounds the listing", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{enumerated}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models"))
require.NoError(t, err)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"},
out.Mutations.RewriteUpstream.DiscoveryModels,
"the picker must be bounded by what the route authorises")
})
t.Run("catch-all route leaves the listing alone", func(t *testing.T) {
catchAll := enumerated
catchAll.Models = nil
mw := New(Config{Providers: []ProviderRoute{catchAll}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models"))
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
"a route that claims every model cannot bound the upstream's list")
})
t.Run("per-model lookup is not a listing", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{enumerated}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5"))
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
"the single-object lookup has no data array to filter")
})
}
// TestRouter_ModelDetailHonoursAllowlist pins that GET /v1/models/{id} is
// authorised against the model table. It carries no body model, so treating
// it as a model-less endpoint would let a caller confirm a model the route
// does not list — the listing itself is bounded to the allowlist, so the
// detail lookup must be too.
func TestRouter_ModelDetailHonoursAllowlist(t *testing.T) {
enumerated := ProviderRoute{
ID: "anthropic-prod",
Models: []string{"claude-sonnet-5"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "api.anthropic.com",
}
t.Run("allowlisted model routes and skips metering", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{enumerated}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host)
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
assert.Equal(t, "true", nonInference, "a detail lookup spends no tokens")
})
t.Run("model outside the allowlist denies", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{enumerated}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-opus-5"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"a model no route lists must not be confirmed by the detail lookup")
})
t.Run("dated id matches its undated registration", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{enumerated}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5-20250929"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"a pinned release of an allowlisted family stays reachable")
})
t.Run("catch-all route still answers every lookup", func(t *testing.T) {
catchAll := enumerated
catchAll.Models = nil
mw := New(Config{Providers: []ProviderRoute{catchAll}})
out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/anything-at-all"))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"a gateway that enumerates nothing cannot refuse a lookup")
})
}
// TestRouter_NonInferenceRequiresReadMethod pins that the non-inference mark —
// which exempts a request from the token pre-flight — is reachable only by the
// read methods these endpoints actually use. A POST to the same path could
// carry an inference body, so it must not buy the exemption; it falls through
// to normal per-model routing instead, which denies when no model is named.
func TestRouter_NonInferenceRequiresReadMethod(t *testing.T) {
route := ProviderRoute{
ID: "gateway",
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "gateway.example.com",
}
for _, path := range []string{"/v1/models", "/v1/models/claude-sonnet-5", "/api/hello"} {
t.Run("POST "+path, func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(path)
in.Method = http.MethodPost
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"a write to a non-inference path must not route unmetered")
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
assert.NotEqual(t, "true", nonInference,
"only a read method may skip the token pre-flight")
})
}
t.Run("HEAD keeps the warm probe working", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(connectionWarmPath)
in.Method = http.MethodHead
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"the HEAD warm probe must still reach the upstream")
nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference)
assert.Equal(t, "true", nonInference,
"the HEAD warm probe carries no model to meter")
})
}
// TestRouter_PinnedDatedModelStaysDistinct pins that a route registered
// against one dated Anthropic release does not claim another. Normalising
// both sides of the comparison made every dated build of a family
// interchangeable, so an operator who deliberately pinned a build would have
// served a different one — and with several such routes, declaration or path
// order would have decided which.
func TestRouter_PinnedDatedModelStaysDistinct(t *testing.T) {
pinned := ProviderRoute{
ID: "anthropic-pinned",
Vendor: "anthropic",
Models: []string{"claude-sonnet-4-5-20250101"},
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "pinned.example.com",
}
t.Run("a different dated release is not claimed", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{pinned}})
in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages")
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"a route pinned to one dated build must not serve another")
})
t.Run("its own dated release still routes", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{pinned}})
in := newInputWithModelAndURL("claude-sonnet-4-5-20250101", "/v1/messages")
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "the exact match must still route")
})
t.Run("two pinned builds each route to their own provider", func(t *testing.T) {
other := pinned
other.ID = "anthropic-pinned-newer"
other.Models = []string{"claude-sonnet-4-5-20250202"}
other.UpstreamHost = "newer.example.com"
mw := New(Config{Providers: []ProviderRoute{pinned, other}})
in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages")
in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"})
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "newer.example.com", out.Mutations.RewriteUpstream.Host,
"declaration order must not decide between two deliberately pinned builds")
})
}
// TestRouter_DiscoveryBoundToCallersPolicies pins that a model listing is
// bounded by the policies that authorise the caller, not by the union across
// everyone who can reach the provider. Two teams sharing one provider record
// under different allowlists is the case that makes the difference visible: a
// flattened per-provider list would offer each team the other's models, and
// every one of those entries is a request the guardrail then refuses.
func TestRouter_DiscoveryBoundToCallersPolicies(t *testing.T) {
const (
eng = "grp-eng"
sales = "grp-sales"
)
route := ProviderRoute{
ID: "shared-gateway",
Models: []string{"claude-sonnet-5", "claude-haiku-4-5", "gpt-4o"},
AllowedGroupIDs: []string{eng, sales},
UpstreamScheme: "https",
UpstreamHost: "gateway.example.com",
ModelPolicies: []ModelPolicyRule{
{GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}},
{GroupIDs: []string{sales}, Models: []string{"gpt-4o"}},
},
}
listingFor := func(t *testing.T, group string) []string {
t.Helper()
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(modelListingPath)
in.UserGroups = []string{group}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.Equal(t, middleware.DecisionAllow, out.Decision)
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
return out.Mutations.RewriteUpstream.DiscoveryModels
}
t.Run("each group sees only its own policy's models", func(t *testing.T) {
assert.Equal(t, []string{"claude-sonnet-5"}, listingFor(t, eng),
"engineering must not be offered the model only sales may use")
assert.Equal(t, []string{"gpt-4o"}, listingFor(t, sales),
"sales must not be offered the model only engineering may use")
})
t.Run("a model no policy allows is offered to nobody", func(t *testing.T) {
for _, group := range []string{eng, sales} {
assert.NotContains(t, listingFor(t, group), "claude-haiku-4-5",
"the provider serves it, but no policy permits it")
}
})
}
// TestRouter_DiscoveryUnrestrictedPolicy covers the lifting rule: a caller
// holding one policy without a model allowlist sees everything the provider
// enumerates, whatever the other policies say.
func TestRouter_DiscoveryUnrestrictedPolicy(t *testing.T) {
const (
eng = "grp-eng"
admin = "grp-admin"
)
route := ProviderRoute{
ID: "shared-gateway",
Models: []string{"claude-sonnet-5", "gpt-4o"},
AllowedGroupIDs: []string{eng, admin},
UpstreamScheme: "https",
UpstreamHost: "gateway.example.com",
ModelPolicies: []ModelPolicyRule{
{GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}},
// nil Models: a policy that sets no allowlist at all.
{GroupIDs: []string{admin}},
},
}
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(modelListingPath)
in.UserGroups = []string{eng, admin}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.ElementsMatch(t, []string{"claude-sonnet-5", "gpt-4o"},
out.Mutations.RewriteUpstream.DiscoveryModels,
"an unrestricted policy the caller holds lifts the restriction")
}
// TestRouter_DiscoveryOnGatewayRecord covers a record that enumerates no
// models. It previously offered the upstream's whole catalogue however narrow
// the policy was, because there was nothing to intersect against; the policy
// allowlist is now the bound on its own.
func TestRouter_DiscoveryOnGatewayRecord(t *testing.T) {
const eng = "grp-eng"
base := ProviderRoute{
ID: "litellm",
AllowedGroupIDs: []string{eng},
UpstreamScheme: "https",
UpstreamHost: "litellm.internal",
}
t.Run("a policy allowlist bounds it", func(t *testing.T) {
route := base
route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{"gpt-4o"}}}
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(modelListingPath)
in.UserGroups = []string{eng}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, []string{"gpt-4o"}, out.Mutations.RewriteUpstream.DiscoveryModels,
"a catch-all record must still be bounded by what policy permits")
})
t.Run("an allowlist permitting nothing offers nothing", func(t *testing.T) {
route := base
route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{}}}
mw := New(Config{Providers: []ProviderRoute{route}})
in := newModellessInput(modelListingPath)
in.UserGroups = []string{eng}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels,
"an empty allowlist permits nothing, and must not be read as unrestricted")
})
t.Run("no policy restriction leaves the listing alone", func(t *testing.T) {
mw := New(Config{Providers: []ProviderRoute{base}})
in := newModellessInput(modelListingPath)
in.UserGroups = []string{eng}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Nil(t, out.Mutations.RewriteUpstream.DiscoveryModels,
"nothing narrows the listing, so the upstream's own answer passes through")
})
}

View File

@@ -11,78 +11,11 @@ var codeRegex = regexp.MustCompile(`^[a-z][a-z0-9._-]{0,63}$`)
// denyResponse is the on-wire shape rendered by RenderDenyResponse.
// Keeping this as a typed struct ensures we never leak
// middleware-supplied bytes outside known fields.
//
// Type and Error mirror the denial in the vendor's own error shape when
// the request reached a known LLM surface. LLM clients only parse their
// provider's envelope, so without the mirror a budget stop reaches the
// user as an unexplained API error. The NetBird fields stay where they
// were, so the body is a superset and existing consumers are unaffected.
type denyResponse struct {
Code string `json:"code"`
Message string `json:"message,omitempty"`
Details map[string]string `json:"details,omitempty"`
Middleware string `json:"middleware,omitempty"`
Type string `json:"type,omitempty"`
Error *providerError `json:"error,omitempty"`
}
// providerError is the nested error object both vendor envelopes carry.
type providerError struct {
Type string `json:"type"`
Message string `json:"message,omitempty"`
Code string `json:"code,omitempty"`
}
// Vendor error types keyed by HTTP status, per each provider's published
// error reference.
const (
anthropicErrInvalidRequest = "invalid_request_error"
anthropicErrPermission = "permission_error"
anthropicErrRateLimit = "rate_limit_error"
anthropicErrAPI = "api_error"
openAIErrInvalidRequest = "invalid_request_error"
openAIErrRateLimit = "rate_limit_error"
)
// providerEnvelope returns the vendor-shaped mirror for a denial on the
// given surface, or nil when the surface has no envelope we can speak.
// message is the already-redacted public message.
func providerEnvelope(surface, code, message string, status int) (string, *providerError) {
switch surface {
case "anthropic":
return "error", &providerError{
Type: anthropicErrorType(status),
Message: message,
}
case "openai":
return "", &providerError{
Type: openAIErrorType(status),
Message: message,
Code: code,
}
default:
return "", nil
}
}
func anthropicErrorType(status int) string {
switch status {
case http.StatusForbidden:
return anthropicErrPermission
case http.StatusTooManyRequests:
return anthropicErrRateLimit
case http.StatusBadRequest:
return anthropicErrInvalidRequest
default:
return anthropicErrAPI
}
}
func openAIErrorType(status int) string {
if status == http.StatusTooManyRequests {
return openAIErrRateLimit
}
return openAIErrInvalidRequest
}
// RenderDenyResponse writes a structured JSON deny body. Status is
@@ -103,7 +36,6 @@ func RenderDenyResponse(w http.ResponseWriter, middlewareID string, reason *Deny
Message: truncate(Scan(reason.Message), 256),
Middleware: truncate(Scan(middlewareID), 64),
}
resp.Type, resp.Error = providerEnvelope(reason.Surface, resp.Code, resp.Message, status)
if n := len(reason.Details); n > 0 {
resp.Details = make(map[string]string, min(n, 8))
for k, v := range reason.Details {

View File

@@ -1,92 +0,0 @@
package middleware
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// decodeDeny renders a denial and returns the parsed body plus the status.
func decodeDeny(t *testing.T, reason *DenyReason, status int) (map[string]any, int) {
t.Helper()
rec := httptest.NewRecorder()
RenderDenyResponse(rec, "llm_limit_check", reason, status)
var body map[string]any
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body), "deny body must be valid JSON")
return body, rec.Code
}
// TestRenderDeny_AnthropicSurfaceMirrorsVendorShape covers a budget stop
// reaching Claude Code. The client only parses the Anthropic envelope, so
// without the mirror the user sees an unexplained API error instead of the
// reason their request was refused.
func TestRenderDeny_AnthropicSurfaceMirrorsVendorShape(t *testing.T) {
body, status := decodeDeny(t, &DenyReason{
Code: "llm_policy.budget_cap_exceeded",
Message: "LLM policy limit exceeded",
Surface: "anthropic",
}, http.StatusForbidden)
assert.Equal(t, http.StatusForbidden, status)
assert.Equal(t, "error", body["type"], "Anthropic errors carry type=error at the top level")
errObj, ok := body["error"].(map[string]any)
require.True(t, ok, "error must be an object")
assert.Equal(t, "permission_error", errObj["type"], "403 maps to permission_error")
assert.Equal(t, "LLM policy limit exceeded", errObj["message"])
// The NetBird fields stay put so existing consumers keep working.
assert.Equal(t, "llm_policy.budget_cap_exceeded", body["code"])
assert.Equal(t, "LLM policy limit exceeded", body["message"])
assert.Equal(t, "llm_limit_check", body["middleware"])
}
// TestRenderDeny_OpenAISurfaceMirrorsVendorShape pins the OpenAI envelope,
// which nests the code and carries no top-level type.
func TestRenderDeny_OpenAISurfaceMirrorsVendorShape(t *testing.T) {
body, _ := decodeDeny(t, &DenyReason{
Code: "llm_policy.model_blocked",
Message: "model is not in the policy allowlist",
Surface: "openai",
}, http.StatusForbidden)
assert.NotContains(t, body, "type", "OpenAI errors have no top-level type")
errObj, ok := body["error"].(map[string]any)
require.True(t, ok, "error must be an object")
assert.Equal(t, "invalid_request_error", errObj["type"])
assert.Equal(t, "llm_policy.model_blocked", errObj["code"], "the NetBird code rides in the vendor code field")
assert.Equal(t, "model is not in the policy allowlist", errObj["message"])
}
// TestRenderDeny_RateLimitStatusMapsToVendorRateLimit pins the mapping a
// client's backoff keys on.
func TestRenderDeny_RateLimitStatusMapsToVendorRateLimit(t *testing.T) {
body, status := decodeDeny(t, &DenyReason{
Code: "llm_policy.token_cap_exceeded",
Message: "LLM policy limit exceeded",
Surface: "anthropic",
}, http.StatusTooManyRequests)
assert.Equal(t, http.StatusTooManyRequests, status, "429 must survive the status clamp")
errObj := body["error"].(map[string]any)
assert.Equal(t, "rate_limit_error", errObj["type"])
}
// TestRenderDeny_NoSurfaceKeepsLegacyShape guards non-LLM middlewares and
// denials raised before a surface is known.
func TestRenderDeny_NoSurfaceKeepsLegacyShape(t *testing.T) {
body, _ := decodeDeny(t, &DenyReason{
Code: "llm_policy.model_not_routable",
Message: "no provider configured for model x",
}, http.StatusForbidden)
assert.NotContains(t, body, "type", "no surface means no vendor mirror")
assert.NotContains(t, body, "error", "no surface means no vendor mirror")
assert.Equal(t, "llm_policy.model_not_routable", body["code"])
}

View File

@@ -22,15 +22,6 @@ const (
// body. Empty for clients that don't send one.
KeyLLMSessionID = "llm.session_id"
// Sub-agent attribution (emitted by llm_request_parser from the
// client's request headers). A coding agent that spawns helpers
// stamps the spawned agent's id, and the spawning agent's id when
// the helper is itself nested, so cost within one session can be
// split across the agents that ran in parallel. These identify an
// agent, not a person or a device: never treat them as a user id.
KeyLLMAgentID = "llm.agent_id"
KeyLLMParentAgentID = "llm.parent_agent_id"
// LLM response-side metadata (emitted by llm_response_parser).
//nolint:gosec // metadata key name, not a credential
KeyLLMInputTokens = "llm.input_tokens"
@@ -75,14 +66,6 @@ const (
// downstream gateways' spend logs.
KeyLLMAuthorisingGroups = "llm.authorising_groups"
// LLM non-inference marker (emitted by llm_router on the allow path
// for endpoints that legitimately carry no model, such as model
// listing). The router still authorises these against the caller's
// groups; the marker only tells the limits gate that a per-model
// allowlist has nothing to evaluate, so an empty model must not be
// read as an undetermined one. Never derived from client input.
KeyLLMNonInference = "llm.non_inference"
// LLM policy attribution (emitted by llm_limit_check on the allow
// path). Names the policy that paid for this request and the
// dimension counters the post-flight llm_limit_record middleware

View File

@@ -179,12 +179,6 @@ type DenyReason struct {
Code string
Message string
Details map[string]string
// Surface names the LLM API dialect the caller speaks (the
// llm.provider value), so the rendered body can mirror the denial in
// that vendor's error shape alongside the NetBird fields. Empty for
// non-LLM middlewares and for denials raised before a surface was
// resolved; the body then carries the NetBird fields alone.
Surface string
}
// Output is the value each middleware returns to the dispatcher. The
@@ -253,12 +247,6 @@ type UpstreamRewrite struct {
// without verifying its TLS certificate. Set by llm_router from the
// provider's skip_tls_verification for self-hosted / internal gateways.
SkipTLSVerify bool
// DiscoveryModels, when non-empty, is the set of model ids the resolved
// route authorises, and the proxy drops everything else from the
// model-listing response. Empty leaves the upstream's list untouched,
// which is what a route that claims every model wants. Set by
// llm_router on a model-listing request only.
DiscoveryModels []string
}
// AuthHeader is a single name/value pair the proxy injects on the

View File

@@ -1,237 +0,0 @@
package proxy
import (
"bytes"
"encoding/json"
"io"
"net/http"
"strconv"
"strings"
sharedllm "github.com/netbirdio/netbird/shared/llm"
)
// maxDiscoveryBodyBytes bounds the model-listing response the filter will
// buffer. A listing is a few kilobytes of ids; anything larger is not a
// listing we recognise, and buffering it to rewrite would cost more than
// the filtering is worth.
const maxDiscoveryBodyBytes = 1 << 20
// modelDiscoveryFilter returns a ModifyResponse hook that drops models the
// caller's policy does not authorise from a model-listing response, then
// delegates to next (which may be nil).
//
// Clients populate their model picker from this endpoint, so an unfiltered
// list offers models the very next request denies. The filter is
// best-effort: a response it cannot safely rewrite passes through
// untouched rather than reaching the client corrupted.
func modelDiscoveryFilter(allowed []string, next func(*http.Response) error) func(*http.Response) error {
permitted := make(map[string]struct{}, len(allowed)*2)
for _, id := range allowed {
permitted[id] = struct{}{}
permitted[sharedllm.NormalizeAnthropicModel(id)] = struct{}{}
}
return func(resp *http.Response) error {
if err := filterModelListing(resp, permitted); err != nil {
return err
}
if next == nil {
return nil
}
return next(resp)
}
}
// filterModelListing rewrites the response body in place, keeping only the
// entries whose id the policy authorises. Responses that are not a plain
// JSON listing are left alone.
func filterModelListing(resp *http.Response, permitted map[string]struct{}) error {
if !isPlainJSONListing(resp) {
return nil
}
// One byte past the cap, so an oversized body is detectable without
// buffering all of it.
body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryBodyBytes+1))
if err != nil {
_ = resp.Body.Close()
return err
}
if len(body) > maxDiscoveryBodyBytes {
// Too large to filter. Put the bytes already read back in front of the
// unread remainder and forward the response exactly as the upstream
// sent it, headers included. Buffering what was read and closing here
// would truncate the body at the cap and hand the client a short,
// invalid listing — worse than not filtering at all.
resp.Body = spliceBody(body, resp.Body)
return nil
}
if err := resp.Body.Close(); err != nil {
return err
}
filtered, ok := filterListingBody(body, permitted)
if !ok {
restoreBody(resp, body)
return nil
}
restoreBody(resp, filtered)
return nil
}
// isPlainJSONListing reports whether the response is a JSON body the filter
// can parse. A content-encoded body is skipped: the transport only
// transparently decompresses what it negotiated itself, and the client
// negotiates its own encoding on this request.
func isPlainJSONListing(resp *http.Response) bool {
if resp == nil || resp.Body == nil {
return false
}
if resp.StatusCode != http.StatusOK {
return false
}
if enc := resp.Header.Get("Content-Encoding"); enc != "" && !strings.EqualFold(enc, "identity") {
return false
}
return strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json")
}
// listingEnvelopes maps a listing's wrapper key to the field naming the model
// id inside it. Vendors did not converge on one shape: OpenAI's is what
// Anthropic adopted, while Bedrock returns inference-profile summaries under a
// key of its own. A body matching none of these is forwarded untouched.
var listingEnvelopes = []struct {
key string
idField string
}{
{"data", "id"},
{"inferenceProfileSummaries", "inferenceProfileId"},
}
// filterListingBody returns the listing with unauthorised entries removed.
// ok is false when the body is not a listing shape, in which case the
// caller must forward the original bytes.
func filterListingBody(body []byte, permitted map[string]struct{}) ([]byte, bool) {
var doc map[string]json.RawMessage
if err := json.Unmarshal(body, &doc); err != nil {
return nil, false
}
for _, envelope := range listingEnvelopes {
raw, present := doc[envelope.key]
if !present {
continue
}
var entries []map[string]json.RawMessage
if err := json.Unmarshal(raw, &entries); err != nil {
return nil, false
}
kept := make([]map[string]json.RawMessage, 0, len(entries))
for _, entry := range entries {
if entryPermitted(entry, envelope.idField, permitted) {
kept = append(kept, entry)
}
}
encoded, err := json.Marshal(kept)
if err != nil {
return nil, false
}
doc[envelope.key] = encoded
out, err := json.Marshal(doc)
if err != nil {
return nil, false
}
return out, true
}
return nil, false
}
// entryPermitted reports whether a listing entry names a model the policy
// authorises, trying every form the same model is written in.
func entryPermitted(entry map[string]json.RawMessage, idField string, permitted map[string]struct{}) bool {
raw, ok := entry[idField]
if !ok {
return false
}
var id string
if err := json.Unmarshal(raw, &id); err != nil {
return false
}
for _, candidate := range modelIDForms(id) {
if _, ok := permitted[candidate]; ok {
return true
}
}
return false
}
// gatewayNamespaces are the provider prefixes a gateway prepends to a model
// it re-exports: LiteLLM lists a Bedrock model the operator registered as
// "anthropic.claude-opus-5" under "bedrock/anthropic.claude-opus-5". Only
// these are stripped before matching.
//
// A slash is not by itself a namespace separator. Self-hosted backends ship
// ids that carry one ("Qwen/Qwen2.5-0.5B-Instruct"), and an upstream is free
// to scope ids per tenant ("tenant-b/claude-sonnet-5"). Treating every slash
// as a prefix let any such id match an allowed model by its tail, so the
// picker offered models the policy never named.
var gatewayNamespaces = map[string]struct{}{
"anthropic": {},
"azure": {},
"bedrock": {},
"mistral": {},
"openai": {},
"vertex_ai": {},
}
// modelIDForms returns the forms a single model id may be written in: the id
// itself, its undated form, and — when the id is namespaced by a gateway we
// recognise — the same two with that namespace removed
// ("vertex_ai/claude-sonnet-5"). The bare id is always tried first.
//
// The namespace is what precedes the FIRST slash: it is a prefix the gateway
// put in front of the whole id, and everything after it is the id the
// operator would have registered, separators included.
func modelIDForms(id string) []string {
if id == "" {
return nil
}
forms := []string{id, sharedllm.NormalizeAnthropicModel(id)}
// A Bedrock listing returns region-prefixed, version-suffixed profile ids
// ("eu.anthropic.claude-haiku-4-5-20251001-v1:0") while the record may
// register the catalog key. Stripping to the key is a no-op for ids that
// carry neither, so this costs nothing on the other surfaces.
if bedrock := sharedllm.NormalizeBedrockModel(id); bedrock != id {
forms = append(forms, bedrock)
}
if slash := strings.Index(id, "/"); slash > 0 {
if _, ok := gatewayNamespaces[id[:slash]]; ok {
tail := id[slash+1:]
forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail))
}
}
return forms
}
// restoreBody puts body back on the response and fixes the length headers
// so the client reads exactly what is there.
// spliceBody returns a ReadCloser that yields prefix followed by whatever is
// left in rest, closing rest when closed. It lets the filter put back bytes it
// consumed while deciding, without owning the rest of the stream.
func spliceBody(prefix []byte, rest io.ReadCloser) io.ReadCloser {
return struct {
io.Reader
io.Closer
}{
Reader: io.MultiReader(bytes.NewReader(prefix), rest),
Closer: rest,
}
}
func restoreBody(resp *http.Response, body []byte) {
resp.Body = io.NopCloser(bytes.NewReader(body))
resp.ContentLength = int64(len(body))
resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
}

View File

@@ -1,272 +0,0 @@
package proxy
import (
"bytes"
"encoding/json"
"io"
"net/http"
"strconv"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// jsonListingResponse builds a 200 model-listing response with the given
// body, as an upstream would return it.
func jsonListingResponse(body string) *http.Response {
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{},
Body: io.NopCloser(strings.NewReader(body)),
ContentLength: int64(len(body)),
}
resp.Header.Set("Content-Type", "application/json")
return resp
}
// listedIDs runs the filter and returns the ids left in the response.
func listedIDs(t *testing.T, allowed []string, body string) []string {
t.Helper()
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body, replaced by the filter
require.NoError(t, modelDiscoveryFilter(allowed, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
raw, err := io.ReadAll(resp.Body)
require.NoError(t, err)
var doc struct {
Data []struct {
ID string `json:"id"`
} `json:"data"`
}
require.NoError(t, json.Unmarshal(raw, &doc), "filtered body must stay valid JSON")
ids := make([]string, 0, len(doc.Data))
for _, entry := range doc.Data {
ids = append(ids, entry.ID)
}
return ids
}
// TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels covers the picker a
// developer sees: an unfiltered upstream list offers every model the shared
// key can reach, and each one the policy excludes is a request the chain
// denies a moment later.
func TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels(t *testing.T) {
ids := listedIDs(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, `{
"data": [
{"id": "claude-opus-5", "display_name": "Claude Opus 5"},
{"id": "claude-sonnet-5", "display_name": "Claude Sonnet 5"},
{"id": "claude-haiku-4-5"}
],
"has_more": false
}`)
assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, ids,
"only the models the route authorises may reach the picker")
}
// TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs pins the two id forms
// a gateway returns for a model the operator registered plainly.
func TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs(t *testing.T) {
ids := listedIDs(t, []string{"claude-sonnet-4-5", "anthropic.claude-opus-5"}, `{
"data": [
{"id": "claude-sonnet-4-5-20250929"},
{"id": "bedrock/anthropic.claude-opus-5"},
{"id": "gpt-4o"}
]
}`)
assert.Equal(t, []string{"claude-sonnet-4-5-20250929", "bedrock/anthropic.claude-opus-5"}, ids,
"a dated or provider-prefixed id must match its registered form")
}
// TestModelDiscoveryFilter_PreservesEnvelopeFields guards the rest of the
// document: clients read paging fields alongside data.
func TestModelDiscoveryFilter_PreservesEnvelopeFields(t *testing.T) {
resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}],"has_more":true,"first_id":"x"}`) //nolint:bodyclose // in-memory body, replaced by the filter
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
raw, err := io.ReadAll(resp.Body)
require.NoError(t, err)
var doc map[string]any
require.NoError(t, json.Unmarshal(raw, &doc))
assert.Equal(t, true, doc["has_more"], "paging fields must survive the rewrite")
assert.Equal(t, "x", doc["first_id"])
assert.Equal(t, strconv.Itoa(len(raw)), resp.Header.Get("Content-Length"),
"Content-Length must match the rewritten body")
}
// TestModelDiscoveryFilter_PassesThroughUnfilterable covers the responses
// the filter must not touch: a compressed body it cannot parse, a non-JSON
// body, an error status, and a document with no data array.
func TestModelDiscoveryFilter_PassesThroughUnfilterable(t *testing.T) {
cases := map[string]func() *http.Response{
"compressed": func() *http.Response {
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
resp.Header.Set("Content-Encoding", "gzip")
return resp
},
"not json": func() *http.Response {
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
resp.Header.Set("Content-Type", "text/html")
return resp
},
"error status": func() *http.Response {
resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`)
resp.StatusCode = http.StatusInternalServerError
return resp
},
"no data array": func() *http.Response {
return jsonListingResponse(`{"object":"list"}`)
},
}
for name, build := range cases {
t.Run(name, func(t *testing.T) {
resp := build() //nolint:bodyclose // in-memory body, replaced by the filter
original, err := io.ReadAll(resp.Body)
require.NoError(t, err)
resp.Body = io.NopCloser(bytes.NewReader(original))
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
got, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.Equal(t, string(original), string(got), "an unfilterable response must reach the client unchanged")
})
}
}
// TestModelDiscoveryFilter_RunsNextHook pins that an existing
// ModifyResponse hook still runs after filtering.
func TestModelDiscoveryFilter_RunsNextHook(t *testing.T) {
called := false
next := func(*http.Response) error {
called = true
return nil
}
resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}]}`) //nolint:bodyclose // in-memory body, replaced by the filter
require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, next)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter
assert.True(t, called, "the chained hook must still run")
}
// TestModelDiscoveryFilter_KeepsSlashBearingIDs covers self-hosted backends
// whose model ids carry a slash of their own. Treating the slash as a
// gateway prefix and keeping only the tail dropped every such model from
// the picker even though the policy named it exactly.
func TestModelDiscoveryFilter_KeepsSlashBearingIDs(t *testing.T) {
ids := listedIDs(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, `{
"object": "list",
"data": [
{"id": "Qwen/Qwen2.5-0.5B-Instruct"},
{"id": "Qwen/Qwen2.5-7B-Instruct"}
]
}`)
assert.Equal(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, ids,
"a slash inside the model id is part of the id, not a provider prefix")
}
// TestModelDiscoveryFilter_RejectsTailMatchOnUnknownNamespace covers the id
// an upstream scopes with a prefix of its own. "tenant-b/claude-sonnet-5"
// ends in a model the policy permits, but it is a different model on a
// different tenant, and the guardrail denies that string outright — so
// offering it hands the picker an entry the next request refuses.
func TestModelDiscoveryFilter_RejectsTailMatchOnUnknownNamespace(t *testing.T) {
ids := listedIDs(t, []string{"claude-sonnet-5"}, `{
"data": [
{"id": "claude-sonnet-5"},
{"id": "tenant-b/claude-sonnet-5"},
{"id": "Qwen/claude-sonnet-5"}
]
}`)
assert.Equal(t, []string{"claude-sonnet-5"}, ids,
"only a namespace a gateway is known to prepend may be stripped before matching")
}
// TestModelDiscoveryFilter_ForwardsOversizedBodyIntact covers a listing past
// the buffering cap. The filter reads one byte beyond the cap to detect the
// size; forwarding only what it read would hand the client a body truncated
// at exactly 1 MiB — valid-looking, short, and unparseable as JSON. The bytes
// already read must be spliced back in front of the unread remainder so the
// response reaches the client exactly as the upstream sent it.
func TestModelDiscoveryFilter_ForwardsOversizedBodyIntact(t *testing.T) {
// A well-formed listing whose single entry pads the body past the cap.
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
require.Greater(t, len(body), maxDiscoveryBodyBytes+1,
"the fixture must exceed the cap by more than the one-byte probe")
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
got, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.Equal(t, len(body), len(got),
"an oversized listing must reach the client whole, not truncated at the cap")
assert.Equal(t, body, string(got), "the forwarded bytes must be the upstream's own")
var doc map[string]json.RawMessage
assert.NoError(t, json.Unmarshal(got, &doc),
"the forwarded body must still parse as JSON")
}
// TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders pins that the
// oversized path leaves the response metadata alone. Rewriting Content-Length
// to the truncated prefix is what made the corruption invisible to the client
// until it tried to parse.
func TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders(t *testing.T) {
padding := strings.Repeat("x", maxDiscoveryBodyBytes)
body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}`
resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body
resp.Header.Set("Content-Length", strconv.Itoa(len(body)))
require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body
assert.Equal(t, int64(len(body)), resp.ContentLength,
"ContentLength must keep describing the body the client receives")
assert.Equal(t, strconv.Itoa(len(body)), resp.Header.Get("Content-Length"),
"the Content-Length header must not be rewritten to the truncated prefix")
}
// TestFilterBedrockInferenceProfiles covers the second listing envelope. AWS
// returns inference-profile summaries under a key of its own with an id field
// of its own, so a filter that only knew OpenAI's shape forwarded a Bedrock
// listing whole — offering every profile in the account regardless of policy.
func TestFilterBedrockInferenceProfiles(t *testing.T) {
body := []byte(`{"inferenceProfileSummaries":[
{"inferenceProfileId":"eu.anthropic.claude-haiku-4-5-20251001-v1:0","status":"ACTIVE"},
{"inferenceProfileId":"eu.anthropic.claude-sonnet-4-6","status":"ACTIVE"},
{"inferenceProfileId":"global.cohere.embed-v4:0","status":"ACTIVE"}
]}`)
// The permitted set holds what the record registers. Here that is the
// catalog key, while the vendor answers with region-prefixed wire ids —
// the two must still line up.
permitted := map[string]struct{}{"anthropic.claude-haiku-4-5": {}}
out, ok := filterListingBody(body, permitted)
require.True(t, ok, "a Bedrock listing must be recognised as filterable")
var doc struct {
Summaries []struct {
ID string `json:"inferenceProfileId"`
} `json:"inferenceProfileSummaries"`
}
require.NoError(t, json.Unmarshal(out, &doc))
require.Len(t, doc.Summaries, 1)
assert.Equal(t, "eu.anthropic.claude-haiku-4-5-20251001-v1:0", doc.Summaries[0].ID)
}
// TestFilterLeavesUnknownEnvelopesAlone keeps the best-effort contract: a body
// the filter cannot parse must reach the client exactly as the upstream sent
// it, rather than being rewritten into something shorter and wrong.
func TestFilterLeavesUnknownEnvelopesAlone(t *testing.T) {
_, ok := filterListingBody([]byte(`{"models":[{"name":"something"}]}`), map[string]struct{}{})
assert.False(t, ok)
}

View File

@@ -363,9 +363,6 @@ func (p *ReverseProxy) forwardUpstream(respWriter http.ResponseWriter, r *http.R
if result.rewriteRedirects {
rp.ModifyResponse = p.rewriteLocationFunc(effectiveURL, rewriteMatchedPath, r) //nolint:bodyclose
}
if upstreamRewrite != nil && len(upstreamRewrite.DiscoveryModels) > 0 {
rp.ModifyResponse = modelDiscoveryFilter(upstreamRewrite.DiscoveryModels, rp.ModifyResponse) //nolint:bodyclose // the hook replaces the body and closes the original
}
rp.ServeHTTP(respWriter, r.WithContext(ctx))
}

View File

@@ -20,6 +20,7 @@ import (
"net/url"
"path/filepath"
"reflect"
"slices"
"sync"
"time"
@@ -2062,9 +2063,7 @@ func (s *Server) updateMapping(ctx context.Context, mapping *proto.ProxyMapping)
if mapping.GetAuth().GetOidc() {
schemes = append(schemes, auth.NewOIDC(s.mgmtClient, svcID, accountID, s.ForwardedProto))
}
for _, ha := range mapping.GetAuth().GetHeaderAuths() {
schemes = append(schemes, auth.NewHeader(s.mgmtClient, svcID, accountID, ha.GetHeader()))
}
schemes = append(schemes, headerAuthSchemes(mapping.GetAuth().GetHeaderAuths())...)
ipRestrictions := s.parseRestrictions(mapping)
s.warnIfGeoUnavailable(mapping.GetDomain(), mapping.GetAccessRestrictions())
@@ -2074,20 +2073,40 @@ func (s *Server) updateMapping(ctx context.Context, mapping *proto.ProxyMapping)
return fmt.Errorf("auth setup for domain %s: %w", mapping.GetDomain(), err)
}
m := s.protoToMapping(ctx, mapping)
// The chain is published before the route that leads to it. A request
// arriving at a target whose chain has not been rebuilt yet is served
// straight through, so a provider update that added the route first left a
// window in which an inference could complete unrouted and unmetered.
// Rebuilding first inverts that: the worst a request in the window meets is
// the new chain in front of the previous target, which is still counted.
if err := s.rebuildMiddlewareChains(svcID, m); err != nil {
return err
}
s.meter.AddMapping(m)
s.proxy.AddMapping(m)
s.meter.AddMapping(m)
s.rebuildMiddlewareChains(svcID, m)
return nil
}
// headerAuthSchemes builds one scheme per canonical header name, carrying every
// hash configured for that name so any of them is accepted — the OR semantics
// management applied while it still validated the credential itself. A name
// whose entries arrive without a hash yields a scheme with none, which rejects
// the header rather than leaving the service unprotected.
func headerAuthSchemes(headerAuths []*proto.HeaderAuth) []auth.Scheme {
names := make([]string, 0, len(headerAuths))
hashes := make(map[string][]string, len(headerAuths))
for _, ha := range headerAuths {
name := http.CanonicalHeaderKey(ha.GetHeader())
if name == "" {
continue
}
if !slices.Contains(names, name) {
names = append(names, name)
}
if hash := ha.GetHashedValue(); hash != "" {
hashes[name] = append(hashes[name], hash)
}
}
schemes := make([]auth.Scheme, 0, len(names))
for _, name := range names {
schemes = append(schemes, auth.NewHeader(name, hashes[name]))
}
return schemes
}
// initMiddlewareManager wires the middleware subsystem at boot. It configures
// the per-process FactoryContext concrete middlewares consult, installs the
// live-service check, and binds the resolver to the registry concrete
@@ -2122,21 +2141,15 @@ func (s *Server) initMiddlewareManager(ctx context.Context) error {
}
// rebuildMiddlewareChains converts m into per-path bindings and calls
// Manager.Rebuild. Short-circuits when the middleware manager is unset, which
// is a deployment without middleware rather than a failure to install it.
//
// A rebuild that fails is reported rather than logged: the caller publishes
// the route once this returns, and a route published over chains that were
// not installed serves requests with no policy enforcement and no metering.
func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) error {
// Manager.Rebuild. Short-circuits when the middleware manager is unset.
func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) {
if s.middlewareManager == nil {
return nil
return
}
bindings := buildMiddlewareBindings(svcID, m)
if err := s.middlewareManager.Rebuild(string(svcID), bindings); err != nil {
return fmt.Errorf("rebuild middleware chains for service %s: %w", svcID, err)
s.Logger.WithError(err).WithField("service_id", svcID).Error("failed to rebuild middleware chains")
}
return nil
}
// isLiveService reports whether svcID is currently present in the live

View File

@@ -6,6 +6,8 @@ import (
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"testing"
"time"
@@ -15,8 +17,10 @@ import (
"go.opentelemetry.io/otel/metric/noop"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/auth"
proxymetrics "github.com/netbirdio/netbird/proxy/internal/metrics"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/hash/argon2id"
"github.com/netbirdio/netbird/shared/management/proto"
)
@@ -209,6 +213,50 @@ func TestRedactMappingForLog_HandlesEmptyOrNilFields(t *testing.T) {
assert.Empty(t, redacted.Path, "empty Path must remain empty")
}
// headerSchemeAccepts reports whether the scheme admits value for headerName.
func headerSchemeAccepts(t *testing.T, scheme auth.Scheme, headerName, value string) bool {
t.Helper()
hdr, ok := scheme.(auth.Header)
require.True(t, ok, "header auths must produce Header schemes")
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.Header.Set(headerName, value)
_, matched := hdr.Verify(req)
return matched
}
func TestHeaderAuthSchemes_GroupsValuesByCanonicalHeaderName(t *testing.T) {
hashOf := func(v string) string {
hash, err := argon2id.Hash(v)
require.NoError(t, err)
return hash
}
schemes := headerAuthSchemes([]*proto.HeaderAuth{
{Header: "Authorization", HashedValue: hashOf("Bearer a")},
{Header: "authorization", HashedValue: hashOf("Bearer b")},
{Header: "X-Api-Key", HashedValue: hashOf("key-1")},
})
require.Len(t, schemes, 2, "entries differing only in header-name case must collapse into one scheme")
assert.True(t, headerSchemeAccepts(t, schemes[0], "Authorization", "Bearer a"), "first value for the header must be accepted")
assert.True(t, headerSchemeAccepts(t, schemes[0], "Authorization", "Bearer b"), "second value for the same header must be accepted")
assert.False(t, headerSchemeAccepts(t, schemes[0], "Authorization", "Bearer c"), "unconfigured value must be rejected")
assert.True(t, headerSchemeAccepts(t, schemes[1], "X-Api-Key", "key-1"), "a second header name keeps its own scheme")
}
// TestHeaderAuthSchemes_MissingHashFailsClosed covers a mapping that names a
// header but carries no hash for it. Dropping the scheme would leave a service
// whose only auth is that header wide open, so the scheme is kept and denies.
func TestHeaderAuthSchemes_MissingHashFailsClosed(t *testing.T) {
schemes := headerAuthSchemes([]*proto.HeaderAuth{{Header: "X-Api-Key"}})
require.Len(t, schemes, 1, "a header without a hash must still register a scheme")
assert.False(t, headerSchemeAccepts(t, schemes[0], "X-Api-Key", "anything"),
"a header auth without a hash must reject every value")
}
type statusUpdateOnlyClient struct {
proto.ProxyServiceClient
}

View File

@@ -10,59 +10,9 @@ import (
"strings"
)
// bedrockVendorNamespaces are the vendor segments a Bedrock model id is
// published under. They identify the geography in front of a cross-region
// inference profile without enumerating geographies: in
// "eu.anthropic.claude-...", what makes "eu" a geography is that "anthropic"
// follows it.
//
// Listing geographies instead is what this replaced, and it aged badly — the
// list held us, eu, apac and global, so every profile issued under jp, au, ca,
// sa or us-gov carried its prefix into the pricing key, matched no catalog
// entry, and reported the model unpriced.
//
// A vendor missing from this map fails safe: its id keeps the geography, which
// is exactly the behaviour of the list this replaced. Over-stripping is the
// dangerous direction, because the result also decides which route may claim a
// model.
var bedrockVendorNamespaces = map[string]struct{}{
"ai21": {},
"amazon": {},
"anthropic": {},
"cohere": {},
"deepseek": {},
"luma": {},
"meta": {},
"mistral": {},
"openai": {},
"qwen": {},
"stability": {},
"twelvelabs": {},
"writer": {},
}
// stripBedrockGeography removes the cross-region inference-profile geography
// from a Bedrock model id, leaving the "<vendor>.<model>" form the catalog and
// the pricing table key on.
//
// A leading segment counts as a geography only when a known vendor follows it.
// "amazon.nova-pro" is a vendor and a model, not a geography and a model, and
// cutting its first segment would strip the vendor away.
func stripBedrockGeography(modelID string) string {
dot := strings.IndexByte(modelID, '.')
if dot <= 0 {
return modelID
}
rest := modelID[dot+1:]
vendor, _, found := strings.Cut(rest, ".")
if !found {
return modelID
}
if _, ok := bedrockVendorNamespaces[vendor]; !ok {
return modelID
}
return rest
}
// bedrockRegionPrefixes are the cross-region inference-profile prefixes that
// front a Bedrock model id (e.g. "eu.anthropic.claude-...").
var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
// bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]"
// version/throughput suffix of a Bedrock model id.
@@ -87,31 +37,15 @@ func NormalizeBedrockModel(modelID string) string {
m = m[i+1:]
}
}
m = stripBedrockGeography(m)
for _, p := range bedrockRegionPrefixes {
if strings.HasPrefix(m, p) {
m = m[len(p):]
break
}
}
return bedrockVersionSuffix.ReplaceAllString(m, "")
}
// anthropicDatedModel matches a Claude model id carrying the trailing
// "-YYYYMMDD" release-date suffix Anthropic appends to a pinned release,
// capturing the id without it. The "claude" anchor is load-bearing: pricing
// looks every model up through this helper regardless of surface, and an
// operator may register a custom id with any shape at all, so an unanchored
// "-\d{8}$" would let "internal-llm-20250101" silently inherit the rate
// registered for "internal-llm". The anchor also covers the vendor-prefixed
// forms ("anthropic.claude-...", "us.anthropic.claude-...").
var anthropicDatedModel = regexp.MustCompile(`(?i)^(.*claude.*)-\d{8}$`)
// NormalizeAnthropicModel strips the trailing release-date suffix from a
// Claude model id, e.g. "claude-sonnet-4-5-20250929" -> "claude-sonnet-4-5",
// so a dated id a client pins matches the undated one the operator
// registered. Ids that are not Claude-family are returned untouched.
// Callers try the verbatim id first and fall back to this, so two dated
// releases of the same family stay distinct wherever both are registered
// explicitly.
func NormalizeAnthropicModel(modelID string) string {
return anthropicDatedModel.ReplaceAllString(modelID, "$1")
}
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
// (e.g. "claude-sonnet-4-5@20250929" -> "claude-sonnet-4-5") so it matches
// the catalog/pricing key. Vertex publisher models are priced under their

View File

@@ -34,63 +34,3 @@ func TestNormalizeVertexModel(t *testing.T) {
require.Equal(t, want, NormalizeVertexModel(in), "normalize %q", in)
}
}
func TestNormalizeAnthropicModel(t *testing.T) {
cases := map[string]string{
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5",
"claude-3-5-haiku-20241022": "claude-3-5-haiku",
"claude-sonnet-5": "claude-sonnet-5",
"claude-opus-4-8": "claude-opus-4-8",
"anthropic.claude-haiku-4-5": "anthropic.claude-haiku-4-5",
"anthropic.claude-sonnet-4-5-20250929": "anthropic.claude-sonnet-4-5",
"us.anthropic.claude-opus-4-8-20250101": "us.anthropic.claude-opus-4-8",
// Non-Claude ids must survive untouched even when they end in eight
// consecutive digits: an operator can register a custom model under
// any id, and pricing looks every one of them up through this helper.
"gpt-4o": "gpt-4o",
"gpt-4o-2024-08-06": "gpt-4o-2024-08-06",
"gpt-4o-20240806": "gpt-4o-20240806",
"internal-llm-20250101": "internal-llm-20250101",
"deepseek-r1-20250120": "deepseek-r1-20250120",
"Qwen/Qwen2.5-20250101": "Qwen/Qwen2.5-20250101",
"gemini-2-5-pro-20250101": "gemini-2-5-pro-20250101",
"": "",
}
for in, want := range cases {
require.Equal(t, want, NormalizeAnthropicModel(in), "normalize %q", in)
}
}
// TestNormalizeBedrockModel_GeographiesBeyondTheOriginalFour covers the bug
// that made this vendor-anchored: the geography used to be matched against a
// list of four, so a profile issued anywhere else kept its prefix, missed the
// catalog key it was supposed to match, and reported the model unpriced.
func TestNormalizeBedrockModel_GeographiesBeyondTheOriginalFour(t *testing.T) {
for _, geo := range []string{"us", "eu", "apac", "global", "jp", "au", "ca", "sa", "us-gov", "il", "mx"} {
t.Run(geo, func(t *testing.T) {
got := NormalizeBedrockModel(geo + ".anthropic.claude-sonnet-5-20260514-v1:0")
require.Equal(t, "anthropic.claude-sonnet-5", got,
"a cross-region profile must reduce to the catalog key whatever geography issued it")
})
}
}
// TestNormalizeBedrockModel_KeepsAVendorItCannotMistakeForAGeography pins the
// direction that must never break: a plain "<vendor>.<model>" id has no
// geography, and cutting its first segment would strip the vendor away and
// hand the id to whichever route claims the bare model name.
func TestNormalizeBedrockModel_KeepsAVendorItCannotMistakeForAGeography(t *testing.T) {
cases := map[string]string{
"amazon.nova-pro-v1:0": "amazon.nova-pro",
"anthropic.claude-sonnet-5-v1:0": "anthropic.claude-sonnet-5",
"meta.llama3-3-70b-instruct-v1:0": "meta.llama3-3-70b-instruct",
"cohere.command-r-plus-v1:0": "cohere.command-r-plus",
"eu.unknownvendor.some-model-v1:0": "eu.unknownvendor.some-model",
"Qwen/Qwen2.5-0.5B-Instruct": "Qwen/Qwen2.5-0.5B-Instruct",
}
for in, want := range cases {
t.Run(in, func(t *testing.T) {
require.Equal(t, want, NormalizeBedrockModel(in))
})
}
}

View File

@@ -5335,84 +5335,6 @@ components:
- input_per_1k
- output_per_1k
- context_window
AgentNetworkModelDiscoveryRequest:
type: object
properties:
catalog_provider_id:
type: string
description: Catalog provider to query (AgentNetworkCatalogProvider.id). Determines the listing endpoint, the auth header and the response shape.
example: "bedrock_api"
upstream_url:
type: string
description: |
The upstream being configured. Used to reach vendors that serve their listing from the same host as inference, and to read back the region for those whose host embeds one. Ignored when provider_id is supplied.
example: "https://bedrock-runtime.eu-central-1.amazonaws.com"
api_key:
type: string
description: Credential to query the vendor with, for a provider that has not been saved yet. Mutually exclusive with provider_id.
example: "sk-..."
provider_id:
type: string
description: Existing Agent Network provider record whose stored credential and upstream should be used. Lets the form refresh the list without the client holding the key.
example: "ch8i4ug6lnn4g9hqv7m0"
required:
- catalog_provider_id
AgentNetworkModelDiscoveryResponse:
type: object
properties:
models:
type: array
description: Models the credential can reach, in the order the vendor returned them.
items:
$ref: '#/components/schemas/AgentNetworkDiscoveredModel'
required:
- models
AgentNetworkDiscoveredModel:
type: object
properties:
id:
type: string
description: |
Identifier to register on the provider record, in the form the vendor issues it. For Bedrock this is the region-prefixed inference-profile id, which is the only form AWS accepts at invoke time.
example: "eu.anthropic.claude-haiku-4-5-20251001-v1:0"
label:
type: string
description: Vendor-supplied display name, where the vendor supplies one.
example: "EU Anthropic Claude Haiku 4.5"
pricing_known:
type: boolean
description: Whether NetBird's shipped pricing table can price this model. When false the rates below are all zero and the operator must set them, or requests to this model would record a cost of zero.
example: true
input_per_1k:
type: number
format: double
description: Default input token price per 1k tokens, in USD, from the same table the proxy bills with. Zero when pricing_known is false.
example: 0.005
output_per_1k:
type: number
format: double
description: Default output token price per 1k tokens, in USD. Zero when pricing_known is false.
example: 0.015
cached_input_per_1k:
type: number
format: double
description: OpenAI-shape cache rate — default cost per 1k cached prompt tokens (a subset of input tokens), in USD. Absent when the model has no cached-input discount.
example: 0.000075
cache_read_per_1k:
type: number
format: double
description: Anthropic-shape cache rate — default cost per 1k cache-read tokens (additive to input tokens), in USD. Absent when the model has no cache-read rate.
example: 0.0003
cache_creation_per_1k:
type: number
format: double
description: Anthropic-shape cache rate — default cost per 1k cache-creation tokens (additive to input tokens), in USD. Absent when the model has no cache-creation rate.
example: 0.00375
required:
- id
- pricing_known
- input_per_1k
- output_per_1k
AgentNetworkCatalogProvider:
type: object
properties:
@@ -14082,42 +14004,6 @@ paths:
"$ref": "#/components/responses/forbidden"
'500':
"$ref": "#/components/responses/internal_error"
/api/agent-network/catalog/providers/models:
post:
summary: Discover the models a provider credential can reach
description: |
Asks the vendor which models the supplied credential can actually use, so the provider form can offer a live list instead of only the static catalog. The endpoint, auth header and response shape are taken from the catalog entry, never from the request.
Supply either an api_key together with the upstream_url being configured (before the provider is saved), or a provider_id of an existing record to reuse its stored credential.
Returns 422 for a catalog provider that has no listing endpoint (most gateways); the caller should fall back to the catalog's own model list. A model whose price the shipped table does not know is returned with pricing_known false, and the operator must set rates for it.
tags: [ Agent Network ]
security:
- BearerAuth: [ ]
- TokenAuth: [ ]
requestBody:
required: true
content:
application/json:
schema:
$ref: '#/components/schemas/AgentNetworkModelDiscoveryRequest'
responses:
'200':
description: The models the credential can reach
content:
application/json:
schema:
$ref: '#/components/schemas/AgentNetworkModelDiscoveryResponse'
'400':
"$ref": "#/components/responses/bad_request"
'401':
"$ref": "#/components/responses/requires_authentication"
'403':
"$ref": "#/components/responses/forbidden"
'422':
"$ref": "#/components/responses/validation_failed_simple"
'500':
"$ref": "#/components/responses/internal_error"
/api/agent-network/providers:
get:
summary: List all Agent Network Providers

View File

@@ -2120,33 +2120,6 @@ type AgentNetworkConsumption struct {
// AgentNetworkConsumptionDimensionKind Whether this row counts a single end user or a single source group across every member.
type AgentNetworkConsumptionDimensionKind string
// AgentNetworkDiscoveredModel defines model for AgentNetworkDiscoveredModel.
type AgentNetworkDiscoveredModel struct {
// CacheCreationPer1k Anthropic-shape cache rate — default cost per 1k cache-creation tokens (additive to input tokens), in USD. Absent when the model has no cache-creation rate.
CacheCreationPer1k *float64 `json:"cache_creation_per_1k,omitempty"`
// CacheReadPer1k Anthropic-shape cache rate — default cost per 1k cache-read tokens (additive to input tokens), in USD. Absent when the model has no cache-read rate.
CacheReadPer1k *float64 `json:"cache_read_per_1k,omitempty"`
// CachedInputPer1k OpenAI-shape cache rate — default cost per 1k cached prompt tokens (a subset of input tokens), in USD. Absent when the model has no cached-input discount.
CachedInputPer1k *float64 `json:"cached_input_per_1k,omitempty"`
// Id Identifier to register on the provider record, in the form the vendor issues it. For Bedrock this is the region-prefixed inference-profile id, which is the only form AWS accepts at invoke time.
Id string `json:"id"`
// InputPer1k Default input token price per 1k tokens, in USD, from the same table the proxy bills with. Zero when pricing_known is false.
InputPer1k float64 `json:"input_per_1k"`
// Label Vendor-supplied display name, where the vendor supplies one.
Label *string `json:"label,omitempty"`
// OutputPer1k Default output token price per 1k tokens, in USD. Zero when pricing_known is false.
OutputPer1k float64 `json:"output_per_1k"`
// PricingKnown Whether NetBird's shipped pricing table can price this model. When false the rates below are all zero and the operator must set them, or requests to this model would record a cost of zero.
PricingKnown bool `json:"pricing_known"`
}
// AgentNetworkGuardrail defines model for AgentNetworkGuardrail.
type AgentNetworkGuardrail struct {
// Checks Guardrail check parameters. Each entry has an `enabled` flag plus per-check configuration; disabled entries are inert.
@@ -2194,27 +2167,6 @@ type AgentNetworkGuardrailRequest struct {
Name string `json:"name"`
}
// AgentNetworkModelDiscoveryRequest defines model for AgentNetworkModelDiscoveryRequest.
type AgentNetworkModelDiscoveryRequest struct {
// ApiKey Credential to query the vendor with, for a provider that has not been saved yet. Mutually exclusive with provider_id.
ApiKey *string `json:"api_key,omitempty"`
// CatalogProviderId Catalog provider to query (AgentNetworkCatalogProvider.id). Determines the listing endpoint, the auth header and the response shape.
CatalogProviderId string `json:"catalog_provider_id"`
// ProviderId Existing Agent Network provider record whose stored credential and upstream should be used. Lets the form refresh the list without the client holding the key.
ProviderId *string `json:"provider_id,omitempty"`
// UpstreamUrl The upstream being configured. Used to reach vendors that serve their listing from the same host as inference, and to read back the region for those whose host embeds one. Ignored when provider_id is supplied.
UpstreamUrl *string `json:"upstream_url,omitempty"`
}
// AgentNetworkModelDiscoveryResponse defines model for AgentNetworkModelDiscoveryResponse.
type AgentNetworkModelDiscoveryResponse struct {
// Models Models the credential can reach, in the order the vendor returned them.
Models []AgentNetworkDiscoveredModel `json:"models"`
}
// AgentNetworkPolicy defines model for AgentNetworkPolicy.
type AgentNetworkPolicy struct {
// CreatedAt Timestamp when the policy was created.
@@ -6227,9 +6179,6 @@ type PostApiAgentNetworkBudgetRulesJSONRequestBody = AgentNetworkBudgetRuleReque
// PutApiAgentNetworkBudgetRulesRuleIdJSONRequestBody defines body for PutApiAgentNetworkBudgetRulesRuleId for application/json ContentType.
type PutApiAgentNetworkBudgetRulesRuleIdJSONRequestBody = AgentNetworkBudgetRuleRequest
// PostApiAgentNetworkCatalogProvidersModelsJSONRequestBody defines body for PostApiAgentNetworkCatalogProvidersModels for application/json ContentType.
type PostApiAgentNetworkCatalogProvidersModelsJSONRequestBody = AgentNetworkModelDiscoveryRequest
// PostApiAgentNetworkGuardrailsJSONRequestBody defines body for PostApiAgentNetworkGuardrails for application/json ContentType.
type PostApiAgentNetworkGuardrailsJSONRequestBody = AgentNetworkGuardrailRequest