Merge remote-tracking branch 'origin/main' into HEAD

# Conflicts:
#	shared/management/proto/proxy_service.pb.go
This commit is contained in:
Viktor Liu
2026-09-23 07:46:09 +02:00
1091 changed files with 82476 additions and 19116 deletions
+1 -1
View File
@@ -1,4 +1,4 @@
FROM golang:1.25-alpine AS builder
FROM golang:1.26.7-alpine AS builder
WORKDIR /app
RUN echo "netbird:x:1000:1000:netbird:/var/lib/netbird:/sbin/nologin" > /tmp/passwd && \
+1 -1
View File
@@ -1,4 +1,4 @@
FROM golang:1.25-alpine AS builder
FROM golang:1.26.7-alpine AS builder
WORKDIR /app
COPY go.mod go.sum ./
+6
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".
+12 -4
View File
@@ -2,12 +2,14 @@ package acme
import (
"context"
"fmt"
"path/filepath"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/flock"
"github.com/netbirdio/netbird/proxy/internal/k8s"
"github.com/netbirdio/netbird/shared/management/domain"
)
// certLocker provides distributed mutual exclusion for certificate operations.
@@ -74,9 +76,15 @@ func newFlockLocker(certDir string, logger *log.Logger) *flockLocker {
return &flockLocker{certDir: certDir, logger: logger}
}
// Lock acquires an advisory file lock for the given domain.
func (l *flockLocker) Lock(ctx context.Context, domain string) (func(), error) {
lockPath := filepath.Join(l.certDir, domain+".lock")
// Lock acquires an advisory file lock for the given domain. The domain must
// be a valid hostname so the lock file always resolves to a direct child of
// certDir; anything else is rejected before touching the filesystem.
func (l *flockLocker) Lock(ctx context.Context, name string) (func(), error) {
if !domain.IsValidDomainNoWildcard(name) {
return nil, fmt.Errorf("invalid domain %q for lock file", name)
}
lockPath := filepath.Join(l.certDir, name+".lock")
lockFile, err := flock.Lock(ctx, lockPath)
if err != nil {
return nil, err
@@ -89,7 +97,7 @@ func (l *flockLocker) Lock(ctx context.Context, domain string) (func(), error) {
return func() {
if err := flock.Unlock(lockFile); err != nil {
l.logger.Debugf("release cert lock for domain %q: %v", domain, err)
l.logger.Debugf("release cert lock for domain %q: %v", name, err)
}
}, nil
}
+30
View File
@@ -63,3 +63,33 @@ func TestNewCertLockerK8sFallsBackToFlock(t *testing.T) {
_, ok := locker.(*flockLocker)
assert.True(t, ok, "k8s-lease without SA should fall back to flockLocker")
}
func TestFlockLockerRejectsUnsafeDomain(t *testing.T) {
root := t.TempDir()
certDir := filepath.Join(root, "certs")
require.NoError(t, os.Mkdir(certDir, 0o700))
locker := newFlockLocker(certDir, nil)
for _, d := range []string{
"",
".",
"..",
"../escape",
"../../etc/cron.d/attacker",
"sub/dir.example.com",
`back\slash.example.com`,
"*.example.com",
} {
unlock, err := locker.Lock(context.Background(), d)
assert.Error(t, err, "domain %q", d)
assert.Nil(t, unlock, "domain %q", d)
}
assert.NoFileExists(t, filepath.Join(root, "escape.lock"))
certEntries, err := os.ReadDir(certDir)
require.NoError(t, err)
assert.Empty(t, certEntries)
rootEntries, err := os.ReadDir(root)
require.NoError(t, err)
assert.Len(t, rootEntries, 1)
}
+26
View File
@@ -0,0 +1,26 @@
# PIN and password authentication limits
PIN and password credentials are accepted only in a POST form body. Query-string
credentials and credentials on other HTTP methods are ignored.
The proxy permits a burst of five credential checks per account and service,
then replenishes one check every six seconds (ten per minute). PIN and password
checks share the same budget. Five failed checks from one client IP in a
rolling five-minute window block that source for fifteen minutes. In-flight checks
reserve failure slots; blocked requests do not extend the cooldown. Successful
authentication clears that source's failure history. Infrastructure failures
consume the service budget without counting as incorrect credentials.
Throttled requests return HTTP 429 with a `Retry-After` delay in seconds. The
login page displays that delay. Existing authenticated sessions and other
authentication methods do not consume these credential budgets.
The client IP comes from the existing trusted-proxy resolution. Deployments
behind a load balancer must configure trusted proxies correctly; otherwise
visitors share the load balancer's source budget. Visitors behind the same NAT
also share a source budget for a service.
State is held in memory per proxy process and resets on restart. Multiple
replicas have independent budgets. State is bounded to 16,384 source entries and
4,096 service entries; when capacity is exhausted, new checks are denied until
idle entries expire. Active blocks are never evicted to admit a new source.
+100
View File
@@ -0,0 +1,100 @@
package auth
import (
"errors"
"math"
"net/http"
"strconv"
"time"
"google.golang.org/genproto/googleapis/rpc/errdetails"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/proxy"
)
var errCredentialClientIP = errors.New("invalid client address")
type credentialLimitError struct {
retryAfter time.Duration
}
func (e *credentialLimitError) Error() string {
return "too many authentication attempts"
}
func credentialFormValue(r *http.Request, field string) string {
if r.Method != http.MethodPost {
return ""
}
return r.PostFormValue(field)
}
func (mw *Middleware) authenticateScheme(r *http.Request, config DomainConfig, scheme Scheme) (string, string, error) {
method := scheme.Type()
if (method != auth.MethodPIN && method != auth.MethodPassword) || !wasCredentialSubmitted(r, method) {
return scheme.Authenticate(r)
}
ip := mw.resolveClientIP(r).Unmap()
if !ip.IsValid() {
return "", "", errCredentialClientIP
}
source, retry := mw.credentials.begin(credentialSourceKey{
service: credentialServiceKey{accountID: config.AccountID, serviceID: config.ServiceID},
ip: ip,
})
if retry > 0 {
return "", "", &credentialLimitError{retryAfter: retry}
}
token, prompt, err := scheme.Authenticate(r)
outcome := credentialUnavailable
if err == nil {
outcome = credentialRejected
if token != "" {
outcome = credentialAccepted
}
}
mw.credentials.finish(source, outcome)
return token, prompt, err
}
func credentialRetryAfter(err error) time.Duration {
var limitErr *credentialLimitError
if errors.As(err, &limitErr) {
return limitErr.retryAfter
}
s := status.Convert(err)
if s.Code() != codes.ResourceExhausted {
return 0
}
for _, detail := range s.Details() {
if info, ok := detail.(*errdetails.RetryInfo); ok && info.RetryDelay != nil && info.RetryDelay.CheckValid() == nil {
if delay := info.RetryDelay.AsDuration(); delay > 0 {
return delay
}
}
}
return credentialCheckInterval
}
func (mw *Middleware) writeAuthenticationError(w http.ResponseWriter, r *http.Request, method auth.Method, err error) {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
cd.SetAuthMethod(method.String())
}
if retry := credentialRetryAfter(err); retry > 0 {
// RFC 6585 section 4 forbids caching 429 responses.
w.Header().Set("Cache-Control", "no-store")
w.Header().Set("Retry-After", strconv.FormatInt(int64(math.Ceil(retry.Seconds())), 10))
http.Error(w, "too many authentication attempts; try again later", http.StatusTooManyRequests)
return
}
if errors.Is(err, errCredentialClientIP) {
http.Error(w, "invalid client address", http.StatusBadRequest)
return
}
mw.logger.WithField("scheme", method.String()).Warnf("authentication infrastructure error: %v", err)
http.Error(w, "authentication service unavailable", http.StatusBadGateway)
}
+169
View File
@@ -0,0 +1,169 @@
package auth
import (
"net/netip"
"sync"
"time"
"golang.org/x/time/rate"
"github.com/netbirdio/netbird/proxy/internal/types"
)
const (
credentialFailureLimit = 5
credentialFailureWindow = 5 * time.Minute
credentialBlockDuration = 15 * time.Minute
credentialCheckInterval = 6 * time.Second
credentialCheckBurst = 5
credentialMaxSources = 16384
credentialMaxServices = 4096
credentialCleanupInterval = time.Minute
)
type credentialServiceKey struct {
accountID types.AccountID
serviceID types.ServiceID
}
type credentialSourceKey struct {
service credentialServiceKey
ip netip.Addr
}
type credentialSource struct {
failures []time.Time
pending int
expiresAt time.Time
blockedUntil time.Time
}
type credentialService struct {
limiter *rate.Limiter
lastUsed time.Time
}
type credentialOutcome string
const (
credentialUnavailable credentialOutcome = "unavailable"
credentialRejected credentialOutcome = "rejected"
credentialAccepted credentialOutcome = "accepted"
)
// State is local to this proxy process. Active blocks are never evicted to
// make room for a new source; exhausting capacity denies new checks.
type credentialLimiter struct {
mu sync.Mutex
now func() time.Time
sources map[credentialSourceKey]*credentialSource
services map[credentialServiceKey]*credentialService
nextCleanup time.Time
}
func newCredentialLimiter() *credentialLimiter {
return &credentialLimiter{
now: time.Now,
sources: make(map[credentialSourceKey]*credentialSource),
services: make(map[credentialServiceKey]*credentialService),
}
}
func (l *credentialLimiter) begin(key credentialSourceKey) (*credentialSource, time.Duration) {
l.mu.Lock()
defer l.mu.Unlock()
now := l.now()
l.cleanup(now)
source := l.sources[key]
if source != nil {
if now.Before(source.blockedUntil) {
return nil, source.blockedUntil.Sub(now)
}
if source.pending == 0 && !now.Before(source.expiresAt) {
*source = credentialSource{}
}
source.expireFailures(now)
// Reserve the failure budget before verification so concurrent guesses
// cannot all pass a check against the same completed failure count.
if len(source.failures)+source.pending >= credentialFailureLimit {
return nil, time.Second
}
} else if len(l.sources) >= credentialMaxSources {
return nil, credentialCleanupInterval
}
if retry := l.allowService(key.service, now); retry > 0 {
return nil, retry
}
if source == nil {
source = &credentialSource{}
l.sources[key] = source
}
if source.expiresAt.IsZero() {
source.expiresAt = now.Add(credentialFailureWindow)
}
source.pending++
return source, 0
}
func (l *credentialLimiter) allowService(key credentialServiceKey, now time.Time) time.Duration {
service := l.services[key]
if service == nil {
if len(l.services) >= credentialMaxServices {
return credentialCleanupInterval
}
service = &credentialService{limiter: rate.NewLimiter(rate.Every(credentialCheckInterval), credentialCheckBurst)}
l.services[key] = service
}
service.lastUsed = now
if service.limiter.AllowN(now, 1) {
return 0
}
return max(time.Nanosecond, time.Duration((1-service.limiter.TokensAt(now))*float64(credentialCheckInterval)))
}
func (l *credentialLimiter) finish(source *credentialSource, outcome credentialOutcome) {
l.mu.Lock()
defer l.mu.Unlock()
source.pending--
now := l.now()
source.expireFailures(now)
switch outcome {
case credentialRejected:
source.failures = append(source.failures, now)
source.expiresAt = now.Add(credentialFailureWindow)
if len(source.failures) >= credentialFailureLimit && source.blockedUntil.IsZero() {
source.blockedUntil = now.Add(credentialBlockDuration)
source.expiresAt = source.blockedUntil
}
case credentialAccepted:
if !now.Before(source.blockedUntil) {
source.failures = nil
source.expiresAt = now.Add(credentialFailureWindow)
}
case credentialUnavailable:
// Transport failures consume the service budget, but are not bad guesses.
}
}
func (s *credentialSource) expireFailures(now time.Time) {
for len(s.failures) > 0 && !now.Before(s.failures[0].Add(credentialFailureWindow)) {
s.failures = s.failures[1:]
}
}
func (l *credentialLimiter) cleanup(now time.Time) {
if now.Before(l.nextCleanup) {
return
}
l.nextCleanup = now.Add(credentialCleanupInterval)
for key, source := range l.sources {
if source.pending == 0 && !now.Before(source.expiresAt) {
delete(l.sources, key)
}
}
for key, service := range l.services {
if now.Sub(service.lastUsed) >= credentialBlockDuration {
delete(l.services, key)
}
}
}
@@ -0,0 +1,190 @@
package auth
import (
"net/netip"
"strconv"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/types"
)
func TestCredentialLimiterCooldown(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
for range credentialFailureLimit {
attempt, retry := l.begin(key)
require.Zero(t, retry, "initial guesses must reach verification")
l.finish(attempt, credentialRejected)
}
_, retry := l.begin(key)
assert.Equal(t, credentialBlockDuration, retry, "five failures must start a fifteen-minute block")
now = now.Add(credentialBlockDuration - time.Second)
_, retry = l.begin(key)
assert.Equal(t, time.Second, retry, "blocked requests must not extend the deadline")
now = now.Add(time.Second)
attempt, retry := l.begin(key)
require.Zero(t, retry, "the source must recover when its block expires")
l.finish(attempt, credentialAccepted)
}
func TestCredentialLimiterFailureWindowAndSuccess(t *testing.T) {
for _, outcome := range []credentialOutcome{credentialAccepted, credentialUnavailable} {
t.Run(map[credentialOutcome]string{credentialAccepted: "success", credentialUnavailable: "infrastructure error"}[outcome], func(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
for range 4 {
attempt, retry := l.begin(key)
require.Zero(t, retry, "four failures must fit the budget")
l.finish(attempt, credentialRejected)
}
attempt, retry := l.begin(key)
require.Zero(t, retry, "fifth check must be allowed")
l.finish(attempt, outcome)
now = now.Add(credentialCheckInterval)
attempt, retry = l.begin(key)
require.Zero(t, retry, "success or infrastructure error must not start a block")
l.finish(attempt, credentialRejected)
now = now.Add(credentialCheckInterval)
attempt, retry = l.begin(key)
if outcome == credentialUnavailable {
assert.Greater(t, retry, time.Duration(0), "infrastructure errors must preserve earlier failures")
return
}
require.Zero(t, retry, "success must clear earlier failures")
l.finish(attempt, credentialRejected)
now = now.Add(credentialFailureWindow)
for range credentialFailureLimit {
attempt, retry = l.begin(key)
require.Zero(t, retry, "old failures must expire")
l.finish(attempt, credentialRejected)
}
})
}
}
func TestCredentialLimiterRollingWindow(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
attempt, retry := l.begin(key)
require.Zero(t, retry, "the first failure starts the history")
l.finish(attempt, credentialRejected)
now = now.Add(4 * time.Minute)
for range 3 {
attempt, retry = l.begin(key)
require.Zero(t, retry, "three more failures must fit the budget")
l.finish(attempt, credentialRejected)
}
now = now.Add(time.Minute + time.Second)
for range 2 {
attempt, retry = l.begin(key)
require.Zero(t, retry, "only the oldest failure must have expired")
l.finish(attempt, credentialRejected)
}
_, retry = l.begin(key)
assert.Equal(t, credentialBlockDuration, retry, "five recent failures must block even across the first window boundary")
}
func TestCredentialLimiterServiceBudget(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
for range credentialCheckBurst {
attempt, retry := l.begin(key)
require.Zero(t, retry, "initial checks must fit the service burst")
l.finish(attempt, credentialAccepted)
key.ip = key.ip.Next()
}
_, retry := l.begin(key)
assert.Equal(t, credentialCheckInterval, retry, "changing IP must not bypass the service budget")
other := key
other.service.accountID = "another-account"
attempt, retry := l.begin(other)
require.Zero(t, retry, "accounts must have separate budgets")
l.finish(attempt, credentialAccepted)
other = key
other.service.serviceID = "another-service"
attempt, retry = l.begin(other)
require.Zero(t, retry, "services must have separate budgets")
l.finish(attempt, credentialAccepted)
now = now.Add(credentialCheckInterval)
attempt, retry = l.begin(key)
require.Zero(t, retry, "one check must refill every six seconds")
l.finish(attempt, credentialAccepted)
_, retry = l.begin(key)
assert.Equal(t, credentialCheckInterval, retry, "refill must only grant one new check")
}
func TestCredentialLimiterConcurrentReservations(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
var attempts []*credentialSource
for range credentialFailureLimit {
attempt, retry := l.begin(key)
require.Zero(t, retry, "initial requests must reserve the failure budget")
attempts = append(attempts, attempt)
}
// Refill the service budget while earlier verification calls are still running.
now = now.Add(time.Minute)
var admitted atomic.Int32
var wg sync.WaitGroup
for range 100 {
wg.Go(func() {
attempt, retry := l.begin(key)
if retry == 0 {
admitted.Add(1)
l.finish(attempt, credentialRejected)
}
})
}
wg.Wait()
assert.Zero(t, admitted.Load(), "in-flight guesses must reserve the failure budget despite a refilled service budget")
for _, attempt := range attempts {
wg.Go(func() { l.finish(attempt, credentialRejected) })
}
wg.Wait()
_, retry := l.begin(key)
assert.Equal(t, credentialBlockDuration, retry, "concurrent failures must activate the block")
}
func TestCredentialLimiterCapacityAndCleanup(t *testing.T) {
for _, fullSources := range []bool{true, false} {
t.Run(map[bool]string{true: "sources", false: "services"}[fullSources], func(t *testing.T) {
l := newCredentialLimiter()
now := time.Now()
l.now = func() time.Time { return now }
key := credentialSourceKey{service: credentialServiceKey{"account", "service"}, ip: netip.MustParseAddr("192.0.2.1")}
if fullSources {
ip := netip.MustParseAddr("198.18.0.1")
for range credentialMaxSources {
l.sources[credentialSourceKey{service: key.service, ip: ip}] = &credentialSource{expiresAt: now.Add(credentialBlockDuration), blockedUntil: now.Add(credentialBlockDuration)}
ip = ip.Next()
}
} else {
for i := range credentialMaxServices {
l.services[credentialServiceKey{serviceID: key.service.serviceID, accountID: types.AccountID(strconv.Itoa(i))}] = &credentialService{lastUsed: now}
}
}
_, retry := l.begin(key)
assert.Positive(t, retry, "full state must deny new checks without evicting active entries")
now = now.Add(credentialBlockDuration)
attempt, retry := l.begin(key)
require.Zero(t, retry, "expired state must release capacity")
l.finish(attempt, credentialAccepted)
})
}
}
+196
View File
@@ -0,0 +1,196 @@
package auth
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"net/netip"
"net/url"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/genproto/googleapis/rpc/errdetails"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/types/known/durationpb"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
servicemanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service/manager"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/store"
mgmttypes "github.com/netbirdio/netbird/management/server/types"
proxyauth "github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/shared/management/proto"
)
// localCredentialClient replaces the transport while keeping the real service
// store, credential verification, and session signing.
type localCredentialClient struct {
server *nbgrpc.ProxyServiceServer
}
func (c localCredentialClient) Authenticate(ctx context.Context, req *proto.AuthenticateRequest, _ ...grpc.CallOption) (*proto.AuthenticateResponse, error) {
return c.server.Authenticate(ctx, req)
}
func credentialHandler(t *testing.T, field string) (*Middleware, http.Handler) {
t.Helper()
ctx := context.Background()
s, err := store.NewStore(ctx, mgmttypes.SqliteStoreEngine, t.TempDir(), nil, false)
require.NoError(t, err)
t.Cleanup(func() { assert.NoError(t, s.Close(ctx)) })
require.NoError(t, s.SaveAccount(ctx, &mgmttypes.Account{Id: "account"}))
keys := generateTestKeyPair(t)
svc := &service.Service{
ID: "service", AccountID: "account", Name: "test", Domain: "example.com",
Enabled: true, SessionPrivateKey: keys.PrivateKey, SessionPublicKey: keys.PublicKey,
Auth: service.AuthConfig{
PinAuth: &service.PINAuthConfig{Enabled: true, Pin: "842716"},
PasswordAuth: &service.PasswordAuthConfig{Enabled: true, Password: "842716"},
},
}
require.NoError(t, svc.Auth.HashSecrets())
require.NoError(t, s.CreateService(ctx, svc))
server := nbgrpc.NewProxyServiceServer(nil, nil, nil, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil)
t.Cleanup(server.Close)
server.SetServiceManager(servicemanager.NewManager(s, nil, nil, nil, nil, nil))
client := localCredentialClient{server: server}
var scheme Scheme = NewPin(client, "service", "account")
if field == "password" {
scheme = NewPassword(client, "service", "account")
}
mw := NewMiddleware(nil, nil, nil)
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, keys.PublicKey, time.Hour, "account", "service", nil, false, nil))
return mw, mw.Protect(newPassthroughHandler())
}
func credentialRequest(method, field, value string) *http.Request {
r := httptest.NewRequest(method, "https://example.com/", strings.NewReader(url.Values{field: {value}}.Encode()))
r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
r.RemoteAddr = "198.51.100.25:12345"
return r
}
func TestCredentialAuthPOSTOnly(t *testing.T) {
for _, field := range []string{"pin", "password"} {
t.Run(field, func(t *testing.T) {
_, handler := credentialHandler(t, field)
for _, method := range []string{http.MethodGet, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodPost} {
r := credentialRequest(method, field, "")
r.URL.RawQuery = url.Values{field: {"842716"}}.Encode()
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, r)
assert.Equal(t, http.StatusUnauthorized, resp.Code, "%s query credentials must not authenticate", method)
assert.Empty(t, resp.Result().Cookies(), "query credentials must not issue a session")
}
for _, method := range []string{http.MethodGet, http.MethodPut, http.MethodPatch, http.MethodDelete} {
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(method, field, "842716"))
assert.Equal(t, http.StatusUnauthorized, resp.Code, "%s body credentials must not authenticate", method)
}
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(http.MethodPost, field, "842716"))
assert.Equal(t, http.StatusSeeOther, resp.Code, "POST body credentials must authenticate")
})
}
}
func TestCredentialAuthThrottling(t *testing.T) {
for _, field := range []string{"pin", "password"} {
t.Run(field, func(t *testing.T) {
_, handler := credentialHandler(t, field)
for range 5 {
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(http.MethodPost, field, "000000"))
require.Equal(t, http.StatusUnauthorized, resp.Code, "initial wrong credentials must be rejected")
}
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(http.MethodPost, field, "842716"))
assert.Equal(t, http.StatusTooManyRequests, resp.Code, "even correct credentials must wait for the block to expire")
assert.Equal(t, "900", resp.Header().Get("Retry-After"), "five failures must block the source for fifteen minutes")
assert.Empty(t, resp.Result().Cookies(), "blocked credentials must not issue a session")
})
}
}
func TestCredentialAuthSessionAndClientIP(t *testing.T) {
keys := generateTestKeyPair(t)
token, err := sessionkey.SignToken(keys.PrivateKey, "pin-user", "", "example.com", proxyauth.MethodPIN, nil, nil, time.Hour)
require.NoError(t, err)
mw := NewMiddleware(nil, nil, nil)
now := time.Now()
mw.credentials.now = func() time.Time { return now }
scheme := &stubScheme{method: proxyauth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, keys.PublicKey, 0, "account", "service", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
for range credentialFailureLimit {
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(http.MethodPost, "pin", "000000"))
require.Equal(t, http.StatusUnauthorized, resp.Code, "bad PIN must consume the failure budget")
}
now = now.Add(credentialCheckInterval)
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, keys.PublicKey, 0, "account", "service", nil, false, nil))
r := credentialRequest(http.MethodPost, "pin", "000000")
r.RemoteAddr = "[::ffff:198.51.100.25]:45678"
r.Header.Set("X-Forwarded-For", "192.0.2.5")
r.Header.Set("X-Real-IP", "192.0.2.6")
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, r)
assert.Equal(t, http.StatusTooManyRequests, resp.Code, "mapped addresses and untrusted forwarding headers must not bypass the source block")
assert.Equal(t, "no-store", resp.Header().Get("Cache-Control"), "rate limits must not be cached")
r.AddCookie(&http.Cookie{Name: proxyauth.SessionCookieName, Value: token})
resp = httptest.NewRecorder()
handler.ServeHTTP(resp, r)
assert.Equal(t, http.StatusOK, resp.Code, "an existing session must pass even with credentials in the request")
assert.Equal(t, "backend", resp.Body.String(), "the authenticated request must reach the application")
r = credentialRequest(http.MethodPost, "pin", "000000")
cd := proxy.NewCapturedData("test")
cd.SetClientIP(netip.MustParseAddr("192.0.2.9"))
r = r.WithContext(proxy.WithCapturedData(r.Context(), cd))
resp = httptest.NewRecorder()
handler.ServeHTTP(resp, r)
assert.Equal(t, http.StatusUnauthorized, resp.Code, "a client resolved by the trusted-proxy middleware must get its own source budget")
r = credentialRequest(http.MethodPost, "pin", "000000")
r.RemoteAddr = "invalid"
resp = httptest.NewRecorder()
handler.ServeHTTP(resp, r)
assert.Equal(t, http.StatusBadRequest, resp.Code, "an unresolvable client address must fail closed")
now = now.Add(credentialBlockDuration)
scheme.token = token
resp = httptest.NewRecorder()
handler.ServeHTTP(resp, credentialRequest(http.MethodPost, "pin", "842716"))
assert.Equal(t, http.StatusSeeOther, resp.Code, "credentials must work again after cooldown")
}
func TestCredentialAuthManagementThrottling(t *testing.T) {
s, err := status.New(codes.ResourceExhausted, "rate limited").WithDetails(&errdetails.RetryInfo{RetryDelay: durationpb.New(2500 * time.Millisecond)})
require.NoError(t, err)
for _, tc := range []struct {
name string
err error
code int
retry string
}{
{"retry info", fmt.Errorf("authenticate PIN: %w", s.Err()), http.StatusTooManyRequests, "3"},
{"missing retry info", status.Error(codes.ResourceExhausted, "rate limited"), http.StatusTooManyRequests, "6"},
{"unavailable", status.Error(codes.Unavailable, "unavailable"), http.StatusBadGateway, ""},
} {
t.Run(tc.name, func(t *testing.T) {
keys := generateTestKeyPair(t)
mw := NewMiddleware(nil, nil, nil)
scheme := &stubScheme{method: proxyauth.MethodPIN, authFn: func(*http.Request) (string, string, error) { return "", "", tc.err }}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, keys.PublicKey, 0, "account", "service", nil, false, nil))
resp := httptest.NewRecorder()
mw.Protect(newPassthroughHandler()).ServeHTTP(resp, credentialRequest(http.MethodPost, "pin", "000000"))
assert.Equal(t, tc.code, resp.Code, "management errors must keep their HTTP meaning")
assert.Equal(t, tc.retry, resp.Header().Get("Retry-After"), "retry hints must round up to whole seconds")
})
}
}
+68 -38
View File
@@ -1,36 +1,33 @@
package auth
import (
"crypto/sha256"
"errors"
"fmt"
"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 +36,64 @@ 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.
//
// A non-nil unusable is a diagnostic rather than a request error: a stored hash
// could not be decoded, so no credential can ever match it and the header stays
// unauthenticatable until the service is saved again. Folding that into an
// ordinary mismatch would hide the misconfiguration behind a permanent 401.
func (h Header) Verify(r *http.Request) (present, matched bool, unusable error) {
value := r.Header.Get(h.headerName)
if value == "" {
return "", "", nil
return false, false, nil
}
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, nil
}
if res.GetSuccess() {
return res.GetSessionToken(), "", nil
for _, hash := range h.hashes {
err := argon2id.Verify(value, hash)
if err == nil {
h.verified.add(digest)
return true, true, nil
}
if !errors.Is(err, argon2id.ErrMismatchedHashAndPassword) {
unusable = err
}
}
return true, false, unusable
}
return "", "", ErrHeaderAuthFailed
// 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{}{}
}
+129 -67
View File
@@ -59,6 +59,11 @@ type DomainConfig struct {
IPRestrictions *restrict.Filter
// Private routes the domain through ValidateTunnelPeer; failure → 403.
Private bool
// AllowedGroups holds the group ids that may reach the service through an
// OIDC identity. When non-empty, a session cookie is honoured only if its
// groups claim intersects this set. Empty means group membership does not
// restrict access.
AllowedGroups map[string]struct{}
}
type validationResult struct {
@@ -82,6 +87,7 @@ type Middleware struct {
sessionValidator SessionValidator
geo restrict.GeoResolver
tunnelCache *tunnelValidationCache
credentials *credentialLimiter
}
// NewMiddleware creates a new authentication middleware. The sessionValidator is
@@ -96,6 +102,7 @@ func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator, geo re
sessionValidator: sessionValidator,
geo: geo,
tunnelCache: newTunnelValidationCache(),
credentials: newCredentialLimiter(),
}
}
@@ -128,7 +135,7 @@ func (mw *Middleware) Protect(next http.Handler) http.Handler {
if mw.forwardWithTunnelPeer(w, r, host, config, next) {
return
}
http.Error(w, "Forbidden", http.StatusForbidden)
denyPrivate(w)
return
}
@@ -146,7 +153,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
}
@@ -223,7 +230,7 @@ func (mw *Middleware) checkIPRestrictions(w http.ResponseWriter, r *http.Request
clientIP := mw.resolveClientIP(r)
if !clientIP.IsValid() {
mw.logger.Debugf("IP restriction: cannot resolve client address for %q, denying", r.RemoteAddr)
http.Error(w, "Forbidden", http.StatusForbidden)
denyForbidden(w, config)
return false
}
@@ -258,10 +265,30 @@ func (mw *Middleware) checkIPRestrictions(w http.ResponseWriter, r *http.Request
reason := verdict.String()
mw.blockIPRestriction(r, reason)
http.Error(w, "Forbidden", http.StatusForbidden)
denyForbidden(w, config)
return false
}
// denyForbidden writes a 403, dropping the client connection when the
// domain is private so a later retry cannot reuse it.
func denyForbidden(w http.ResponseWriter, config DomainConfig) {
if config.Private {
denyPrivate(w)
return
}
http.Error(w, "Forbidden", http.StatusForbidden)
}
// denyPrivate writes a 403 and closes the connection, so a client refused
// before joining the overlay cannot keep retrying on the same warm socket.
// Go's HTTP/2 server turns the exact lowercase "close" token into a GOAWAY.
func denyPrivate(w http.ResponseWriter) {
h := w.Header()
h.Set("Connection", "close")
h.Set("Cache-Control", "no-store")
http.Error(w, "Forbidden", http.StatusForbidden)
}
// resolveClientIP extracts the real client IP from CapturedData, falling back to r.RemoteAddr.
func (mw *Middleware) resolveClientIP(r *http.Request) netip.Addr {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
@@ -316,6 +343,9 @@ func (mw *Middleware) handleOAuthCallbackError(w http.ResponseWriter, r *http.Re
// forwardWithSessionCookie checks for a valid session cookie and, if found,
// sets the user identity on the request context and forwards to the next handler.
// A signature-valid cookie is not on its own a grant: an OIDC session must also
// carry a group the service allows, so a token cannot be replayed past the
// group check that gated the login it came from.
func (mw *Middleware) forwardWithSessionCookie(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
cookie, err := r.Cookie(auth.SessionCookieName)
if err != nil {
@@ -325,6 +355,24 @@ 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 !sessionGroupsAllowed(config.AllowedGroups, auth.Method(method), groups) {
mw.logger.WithFields(log.Fields{
"host": host,
"user_id": userID,
}).Debug("session cookie rejected: groups claim does not intersect the service's allowed groups")
return false
}
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(userID)
cd.SetUserEmail(email)
@@ -436,73 +484,44 @@ 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 {
var presented []string
for _, scheme := range config.Schemes {
hdr, ok := scheme.(Header)
if !ok {
continue
}
handled := mw.tryHeaderScheme(w, r, host, config, hdr, next)
if handled {
present, matched, unusable := hdr.Verify(r)
if matched {
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetUserID(auth.HeaderUserID)
cd.SetAuthMethod(auth.MethodHeader.String())
}
next.ServeHTTP(w, r)
return true
}
if unusable != nil {
mw.logger.WithFields(log.Fields{
"host": r.Host,
"header": hdr.headerName,
}).WithError(unusable).Error("header auth: a configured hash cannot be decoded, so this header can never authenticate; re-save the service")
}
if present {
presented = append(presented, hdr.headerName)
}
}
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 == "" {
if len(presented) == 0 {
return false
}
result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
if err != nil {
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.SetAuthMethod(auth.MethodHeader.String())
}
next.ServeHTTP(w, r)
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)
mw.logger.WithFields(log.Fields{
"host": r.Host,
"headers": presented,
}).Debug("header auth rejected: no presented header matched a configured hash")
setHeaderCapturedData(r.Context(), "", "", nil, nil)
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return true
}
@@ -526,13 +545,9 @@ func (mw *Middleware) authenticateWithSchemes(w http.ResponseWriter, r *http.Req
var attemptedMethod string
for _, scheme := range config.Schemes {
token, promptData, err := scheme.Authenticate(r)
token, promptData, err := mw.authenticateScheme(r, config, scheme)
if err != nil {
mw.logger.WithField("scheme", scheme.Type().String()).Warnf("authentication infrastructure error: %v", err)
if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
cd.SetOrigin(proxy.OriginAuth)
}
http.Error(w, "authentication service unavailable", http.StatusBadGateway)
mw.writeAuthenticationError(w, r, scheme.Type(), err)
return
}
@@ -633,9 +648,9 @@ func setSessionCookie(w http.ResponseWriter, token string, expiration time.Durat
func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
switch method {
case auth.MethodPIN:
return r.FormValue("pin") != ""
return credentialFormValue(r, pinFormId) != ""
case auth.MethodPassword:
return r.FormValue("password") != ""
return credentialFormValue(r, passwordFormId) != ""
case auth.MethodOIDC:
return r.URL.Query().Get("session_token") != ""
}
@@ -644,7 +659,8 @@ func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
// AddDomain registers authentication schemes for the given domain. With schemes a valid session public key is required.
// private=true forces ValidateTunnelPeer enforcement (403 on failure) regardless of the schemes list.
func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter, private bool) error {
// allowedGroups restricts OIDC sessions to the given group ids; empty means unrestricted.
func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter, private bool, allowedGroups []string) error {
if len(schemes) == 0 {
mw.domainsMux.Lock()
defer mw.domainsMux.Unlock()
@@ -653,6 +669,7 @@ func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 st
ServiceID: serviceID,
IPRestrictions: ipRestrictions,
Private: private,
AllowedGroups: groupSet(allowedGroups),
}
return nil
}
@@ -675,6 +692,7 @@ func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 st
ServiceID: serviceID,
IPRestrictions: ipRestrictions,
Private: private,
AllowedGroups: groupSet(allowedGroups),
}
return nil
}
@@ -726,6 +744,50 @@ func (mw *Middleware) validateSessionToken(ctx context.Context, host, token stri
return &validationResult{UserID: userID, UserEmail: email, Valid: true, Groups: groups, GroupNames: groupNames}, nil
}
// groupSet builds the lookup set the cookie path consults, returning nil for an
// empty list so callers can test membership restriction with len().
func groupSet(groups []string) map[string]struct{} {
if len(groups) == 0 {
return nil
}
set := make(map[string]struct{}, len(groups))
for _, g := range groups {
if g != "" {
set[g] = struct{}{}
}
}
if len(set) == 0 {
return nil
}
return set
}
// sessionGroupsAllowed reports whether a session token's groups claim satisfies
// the service's allowed groups. Only OIDC sessions are gated: password, PIN and
// header credentials carry no group identity and are authorised by the secret
// itself, which mirrors how management validates them. A token minted before the
// groups claim existed carries none and is therefore denied on a group-restricted
// service, which sends the user back through login for a fresh decision. A method
// this build doesn't know carries no such argument, so it is denied.
func sessionGroupsAllowed(allowed map[string]struct{}, method auth.Method, groups []string) bool {
if len(allowed) == 0 {
return true
}
switch method {
case auth.MethodPassword, auth.MethodPIN, auth.MethodHeader:
return true
case auth.MethodOIDC:
for _, g := range groups {
if _, ok := allowed[g]; ok {
return true
}
}
return false
default:
return false
}
}
// stripSessionTokenParam returns the request URI with the session_token query
// parameter removed so it doesn't linger in the browser's address bar or history.
func stripSessionTokenParam(u *url.URL) string {
+311 -144
View File
@@ -16,6 +16,7 @@ import (
"time"
log "github.com/sirupsen/logrus"
logtest "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
@@ -25,6 +26,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"
)
@@ -64,7 +66,7 @@ func TestAddDomain_ValidKey(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
err := mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil)
require.NoError(t, err)
mw.domainsMux.RLock()
@@ -81,7 +83,7 @@ func TestAddDomain_EmptyKey(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
err := mw.AddDomain("example.com", []Scheme{scheme}, "", time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", []Scheme{scheme}, "", time.Hour, "", "", nil, false, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "invalid session public key size")
@@ -95,7 +97,7 @@ func TestAddDomain_InvalidBase64(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
err := mw.AddDomain("example.com", []Scheme{scheme}, "not-valid-base64!!!", time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", []Scheme{scheme}, "not-valid-base64!!!", time.Hour, "", "", nil, false, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "decode session public key")
@@ -110,7 +112,7 @@ func TestAddDomain_WrongKeySize(t *testing.T) {
shortKey := base64.StdEncoding.EncodeToString([]byte("tooshort"))
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
err := mw.AddDomain("example.com", []Scheme{scheme}, shortKey, time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", []Scheme{scheme}, shortKey, time.Hour, "", "", nil, false, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "invalid session public key size")
@@ -123,7 +125,7 @@ func TestAddDomain_WrongKeySize(t *testing.T) {
func TestAddDomain_NoSchemes_NoKeyRequired(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
err := mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil, false, nil)
require.NoError(t, err, "domains with no auth schemes should not require a key")
mw.domainsMux.RLock()
@@ -139,8 +141,8 @@ func TestAddDomain_OverwritesPreviousConfig(t *testing.T) {
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp2.PublicKey, 2*time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil, false, nil))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp2.PublicKey, 2*time.Hour, "", "", nil, false, nil))
mw.domainsMux.RLock()
config := mw.domains["example.com"]
@@ -156,7 +158,7 @@ func TestRemoveDomain(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
mw.RemoveDomain("example.com")
@@ -180,7 +182,7 @@ func TestProtect_UnknownDomainPassesThrough(t *testing.T) {
func TestProtect_DomainWithNoSchemesPassesThrough(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
require.NoError(t, mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", nil, "", time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -197,7 +199,7 @@ func TestProtect_UnauthenticatedRequestIsBlocked(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
var backendCalled bool
backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -218,7 +220,7 @@ func TestProtect_HostWithPortIsMatched(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
var backendCalled bool
backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -239,7 +241,7 @@ func TestProtect_ValidSessionCookiePassesThrough(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "", "example.com", auth.MethodPIN, nil, nil, time.Hour)
require.NoError(t, err)
@@ -272,7 +274,7 @@ func TestProtect_SessionCookieGroupsPropagate(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
groups := []string{"engineering", "sre"}
token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "", "example.com", auth.MethodPIN, groups, nil, time.Hour)
@@ -337,7 +339,7 @@ func TestProtect_PrivateService_TunnelPeerGroupsPropagate(t *testing.T) {
kp := generateTestKeyPair(t)
// Private service: no operator schemes — auth gates solely on the tunnel peer.
require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true))
require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true, nil))
cd := proxy.NewCapturedData("")
cd.SetClientIP(netip.MustParseAddr("100.90.1.14")) // CGNAT tunnel source
@@ -377,7 +379,7 @@ func TestProtect_PrivateService_TunnelPeerDenied(t *testing.T) {
}}
mw := NewMiddleware(log.StandardLogger(), validator, nil)
kp := generateTestKeyPair(t)
require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true))
require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true, nil))
cd := proxy.NewCapturedData("")
cd.SetClientIP(netip.MustParseAddr("100.90.1.14"))
@@ -405,7 +407,7 @@ func TestProtect_ExpiredSessionCookieIsRejected(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
// Sign a token that expired 1 second ago.
token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "", "example.com", auth.MethodPIN, nil, nil, -time.Second)
@@ -431,7 +433,7 @@ func TestProtect_WrongDomainCookieIsRejected(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
// Token signed for a different domain audience.
token, err := sessionkey.SignToken(kp.PrivateKey, "test-user", "", "other.com", auth.MethodPIN, nil, nil, time.Hour)
@@ -458,7 +460,7 @@ func TestProtect_WrongKeyCookieIsRejected(t *testing.T) {
kp2 := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp1.PublicKey, time.Hour, "", "", nil, false, nil))
// Token signed with a different private key.
token, err := sessionkey.SignToken(kp2.PrivateKey, "test-user", "", "example.com", auth.MethodPIN, nil, nil, time.Hour)
@@ -495,7 +497,7 @@ func TestProtect_SchemeAuthRedirectsWithCookie(t *testing.T) {
return "", "pin", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
var backendCalled bool
backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -548,7 +550,7 @@ func TestProtect_FailedAuthDoesNotSetCookie(t *testing.T) {
return "", "pin", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -584,7 +586,7 @@ func TestProtect_MultipleSchemes(t *testing.T) {
return "", "password", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{pinScheme, passwordScheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{pinScheme, passwordScheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
var backendCalled bool
backend := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -614,7 +616,7 @@ func TestProtect_InvalidTokenFromSchemeReturns400(t *testing.T) {
return "invalid-jwt-token", "", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -638,7 +640,7 @@ func TestAddDomain_RandomBytes32NotEd25519(t *testing.T) {
key := base64.StdEncoding.EncodeToString(randomBytes)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
err = mw.AddDomain("example.com", []Scheme{scheme}, key, time.Hour, "", "", nil, false)
err = mw.AddDomain("example.com", []Scheme{scheme}, key, time.Hour, "", "", nil, false, nil)
require.NoError(t, err, "any 32-byte key should be accepted at registration time")
}
@@ -647,10 +649,10 @@ func TestAddDomain_InvalidKeyDoesNotCorruptExistingConfig(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
// Attempt to overwrite with an invalid key.
err := mw.AddDomain("example.com", []Scheme{scheme}, "bad", time.Hour, "", "", nil, false)
err := mw.AddDomain("example.com", []Scheme{scheme}, "bad", time.Hour, "", "", nil, false, nil)
require.Error(t, err)
// The original valid config should still be intact.
@@ -674,7 +676,7 @@ func TestProtect_FailedPinAuthCapturesAuthMethod(t *testing.T) {
return "", "pin", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
capturedData := proxy.NewCapturedData("")
handler := mw.Protect(newPassthroughHandler())
@@ -701,7 +703,7 @@ func TestProtect_FailedPasswordAuthCapturesAuthMethod(t *testing.T) {
return "", "password", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
capturedData := proxy.NewCapturedData("")
handler := mw.Protect(newPassthroughHandler())
@@ -728,7 +730,7 @@ func TestProtect_NoCredentialsDoesNotCaptureAuthMethod(t *testing.T) {
return "", "pin", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
capturedData := proxy.NewCapturedData("")
handler := mw.Protect(newPassthroughHandler())
@@ -816,7 +818,7 @@ func TestCheckIPRestrictions_UnparseableAddress(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"10.0.0.0/8"}}), false)
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"10.0.0.0/8"}}), false, nil)
require.NoError(t, err)
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -852,7 +854,7 @@ func TestCheckIPRestrictions_UsesCapturedDataClientIP(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"203.0.113.0/24"}}), false)
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"203.0.113.0/24"}}), false, nil)
require.NoError(t, err)
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -893,7 +895,7 @@ func TestCheckIPRestrictions_NilGeoWithCountryRules(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
restrict.ParseFilter(restrict.FilterConfig{AllowedCountries: []string{"US"}}), false)
restrict.ParseFilter(restrict.FilterConfig{AllowedCountries: []string{"US"}}), false, nil)
require.NoError(t, err)
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -920,7 +922,7 @@ func TestCheckIPRestrictions_OverlayOriginSkipsCountryRules(t *testing.T) {
restrict.ParseFilter(restrict.FilterConfig{
AllowedCIDRs: []string{"100.64.0.0/10"},
AllowedCountries: []string{"US"},
}), false)
}), false, nil)
require.NoError(t, err)
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -954,7 +956,7 @@ func TestCheckIPRestrictions_OverlayOriginRespectsCIDR(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
err := mw.AddDomain("example.com", nil, "", 0, "acc1", "svc1",
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"100.64.0.0/16"}}), false)
restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"100.64.0.0/16"}}), false, nil)
require.NoError(t, err)
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -982,7 +984,7 @@ func TestProtect_OIDCOnlyRedirectsDirectly(t *testing.T) {
return "", oidcURL, nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1011,7 +1013,7 @@ func TestProtect_OIDCWithOtherMethodShowsLoginPage(t *testing.T) {
return "", "pin", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{oidcScheme, pinScheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{oidcScheme, pinScheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1023,39 +1025,25 @@ 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")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
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, nil))
var backendCalled bool
capturedData := proxy.NewCapturedData("")
@@ -1075,19 +1063,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,10 +1076,10 @@ 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))
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr, pinScheme}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1114,11 +1095,8 @@ 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")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
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, nil))
capturedData := proxy.NewCapturedData("")
handler := mw.Protect(newPassthroughHandler())
@@ -1131,94 +1109,283 @@ 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_MatchesAnyConfiguredHeader covers a client that carries
// a valid credential on one configured header while also sending an unrelated
// value on another — an app-level Authorization alongside an API key, say.
// Schemes OR across header names, so the valid credential admits the request no
// matter which order the mapping happened to list the headers in.
func TestProtect_HeaderAuth_MatchesAnyConfiguredHeader(t *testing.T) {
tests := []struct {
name string
matchedLast bool
}{
{name: "unmatched header listed first", matchedLast: true},
{name: "matched header listed first", matchedLast: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
authz := newHeaderScheme(t, "Authorization", "Bearer proxy-secret")
apiKey := newHeaderScheme(t, "X-Api-Key", "secret-key")
schemes := []Scheme{apiKey, authz}
if tt.matchedLast {
schemes = []Scheme{authz, apiKey}
}
require.NoError(t, mw.AddDomain("example.com", schemes, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false, nil))
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", "secret-key")
req.Header.Set("Authorization", "Bearer app-level-token")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.True(t, backendCalled, "a valid credential on one header must admit the request")
assert.Equal(t, http.StatusOK, rec.Code)
})
}
}
// TestProtect_HeaderAuth_RejectsWhenEveryPresentedHeaderFails is the other half
// of the OR: trying all schemes before rejecting must not turn into admitting a
// request that satisfied none of them.
func TestProtect_HeaderAuth_RejectsWhenEveryPresentedHeaderFails(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")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
authz := newHeaderScheme(t, "Authorization", "Bearer proxy-secret")
apiKey := newHeaderScheme(t, "X-Api-Key", "secret-key")
require.NoError(t, mw.AddDomain("example.com", []Scheme{authz, apiKey}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false, nil))
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", "wrong-key")
req.Header.Set("Authorization", "Bearer wrong-token")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.False(t, backendCalled)
assert.Equal(t, http.StatusUnauthorized, rec.Code)
}
// TestProtect_HeaderAuth_ReportsUndecodableHash covers a stored hash the proxy
// cannot decode. No credential can ever match it, so the header is permanently
// unauthenticatable — an operator fault that has to surface loudly instead of
// hiding behind the same quiet 401 a wrong credential earns.
func TestProtect_HeaderAuth_ReportsUndecodableHash(t *testing.T) {
validHash, err := argon2id.Hash("secret-key")
require.NoError(t, err)
tests := []struct {
name string
hashes []string
wantErrLog bool
}{
{name: "stored hash cannot be decoded", hashes: []string{"$argon2id$v=19$garbage"}, wantErrLog: true},
{name: "wrong credential against a good hash", hashes: []string{validHash}, wantErrLog: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
logger, hook := logtest.NewNullLogger()
logger.SetLevel(log.DebugLevel)
mw := NewMiddleware(logger, nil, nil)
kp := generateTestKeyPair(t)
require.NoError(t, mw.AddDomain("example.com", []Scheme{NewHeader("X-Api-Key", tt.hashes)},
kp.PublicKey, time.Hour, "acc1", "svc1", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
req.Header.Set("X-Api-Key", "wrong-key")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
require.Equal(t, http.StatusUnauthorized, rec.Code, "either way the request is denied")
var errored []string
for _, entry := range hook.AllEntries() {
if entry.Level == log.ErrorLevel {
errored = append(errored, entry.Message)
}
}
if !tt.wantErrLog {
assert.Empty(t, errored, "a wrong credential is not an operator fault")
return
}
require.Len(t, errored, 1, "an undecodable hash must be reported once")
assert.Contains(t, errored[0], "cannot be decoded")
})
}
}
// 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)
hdr := NewHeader("X-API-Key", nil)
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false, nil))
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, nil))
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, nil))
// 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, nil))
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")
require.NoError(t, mw.AddDomain("example.com", []Scheme{hdr}, kp.PublicKey, time.Hour, "acc1", "svc1", nil, false))
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, nil))
var backendCalled bool
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -1276,7 +1443,7 @@ func TestProtect_OIDCOnPlainHTTP_BlockedWith400(t *testing.T) {
return "", "https://idp.example.com/authorize", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1300,7 +1467,7 @@ func TestProtect_OIDCOverTLS_NotBlocked(t *testing.T) {
return "", "https://idp.example.com/authorize", nil
},
}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1320,7 +1487,7 @@ func TestProtect_NonOIDCSchemes_PlainHTTP_NotBlocked(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1350,7 +1517,7 @@ func TestProtect_TunnelPeerFastPath_RequiresInboundMarker(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
@@ -1385,7 +1552,7 @@ func TestProtect_TunnelPeerFastPath_TakesPathWithInboundMarker(t *testing.T) {
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false))
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "", "", nil, false, nil))
handler := mw.Protect(newPassthroughHandler())
+1 -1
View File
@@ -35,7 +35,7 @@ func (Password) Type() auth.Method {
// so that it can be injected into a request from the UI so that
// authentication may be successful.
func (p Password) Authenticate(r *http.Request) (string, string, error) {
password := r.FormValue(passwordFormId)
password := credentialFormValue(r, passwordFormId)
if password == "" {
// No password submitted; return the form ID so the UI can prompt the user.
+1 -1
View File
@@ -35,7 +35,7 @@ func (Pin) Type() auth.Method {
// so that it can be injected into a request from the UI so that
// authentication may be successful.
func (p Pin) Authenticate(r *http.Request) (string, string, error) {
pin := r.FormValue(pinFormId)
pin := credentialFormValue(r, pinFormId)
if pin == "" {
// No PIN submitted; return the form ID so the UI can prompt the user.
+272
View File
@@ -0,0 +1,272 @@
package auth
import (
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"net/http/httptrace"
"net/netip"
"sync"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/restrict"
"github.com/netbirdio/netbird/shared/management/proto"
)
// switchableTunnelValidator flips the ValidateTunnelPeer verdict between requests.
type switchableTunnelValidator struct {
mu sync.Mutex
valid bool
}
func (s *switchableTunnelValidator) setValid(v bool) {
s.mu.Lock()
defer s.mu.Unlock()
s.valid = v
}
func (s *switchableTunnelValidator) ValidateSession(context.Context, *proto.ValidateSessionRequest, ...grpc.CallOption) (*proto.ValidateSessionResponse, error) {
return nil, errors.New("not used in this test")
}
func (s *switchableTunnelValidator) ValidateTunnelPeer(context.Context, *proto.ValidateTunnelPeerRequest, ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) {
s.mu.Lock()
defer s.mu.Unlock()
if !s.valid {
return &proto.ValidateTunnelPeerResponse{Valid: false, DeniedReason: "not_in_group"}, nil
}
return &proto.ValidateTunnelPeerResponse{
Valid: true,
UserId: "user-1",
SessionToken: "tunnel-session-token",
}, nil
}
// testServerHost is the domain key Protect derives from the httptest listener.
const testServerHost = "127.0.0.1"
var testTunnelIP = netip.MustParseAddr("100.90.1.14")
// startProtectedServer serves mw.Protect and stamps requests as overlay traffic.
func startProtectedServer(t *testing.T, mw *Middleware, clientIP netip.Addr, lookup TunnelLookupFunc, h2 bool) *httptest.Server {
t.Helper()
protected := mw.Protect(newPassthroughHandler())
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
cd := proxy.NewCapturedData("")
cd.SetClientIP(clientIP)
ctx := proxy.WithCapturedData(r.Context(), cd)
ctx = WithTunnelLookup(ctx, lookup)
protected.ServeHTTP(w, r.WithContext(ctx))
})
srv := httptest.NewUnstartedServer(handler)
if h2 {
srv.EnableHTTP2 = true
srv.StartTLS()
} else {
srv.Start()
}
t.Cleanup(srv.Close)
return srv
}
// tracedResponse is what a test observes from one client round trip.
type tracedResponse struct {
status int
protoMajor int
close bool
connection string
cacheControl string
reused bool
}
// doTraced GETs url and reports whether the connection that served it was reused.
func doTraced(t *testing.T, client *http.Client, url string) tracedResponse {
t.Helper()
var reused bool
trace := &httptrace.ClientTrace{
GotConn: func(info httptrace.GotConnInfo) { reused = info.Reused },
}
req, err := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, url, nil)
require.NoError(t, err)
resp, err := client.Do(req)
require.NoError(t, err)
defer func() { require.NoError(t, resp.Body.Close()) }()
_, err = io.Copy(io.Discard, resp.Body)
require.NoError(t, err)
return tracedResponse{
status: resp.StatusCode,
protoMajor: resp.ProtoMajor,
close: resp.Close,
connection: resp.Header.Get("Connection"),
cacheControl: resp.Header.Get("Cache-Control"),
reused: reused,
}
}
func acceptAllLookup(_ netip.Addr) (PeerIdentity, bool) {
return PeerIdentity{TunnelIP: testTunnelIP}, true
}
func newPrivateMiddleware(t *testing.T, validator SessionValidator, ipRestrictions *restrict.Filter) *Middleware {
t.Helper()
mw := NewMiddleware(log.StandardLogger(), validator, nil)
kp := generateTestKeyPair(t)
require.NoError(t, mw.AddDomain(testServerHost, nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", ipRestrictions, true, nil))
return mw
}
// A rejected tunnel peer must emit the exact lowercase "close" token h2 matches on.
func TestProtect_PrivateService_DeniedSetsCloseHeaders(t *testing.T) {
mw := newPrivateMiddleware(t, &switchableTunnelValidator{}, nil)
handler := mw.Protect(newPassthroughHandler())
cd := proxy.NewCapturedData("")
cd.SetClientIP(testTunnelIP)
req := httptest.NewRequest(http.MethodGet, "http://"+testServerHost+"/", nil)
req.RemoteAddr = testTunnelIP.String() + ":5000"
req = req.WithContext(WithTunnelLookup(proxy.WithCapturedData(req.Context(), cd), acceptAllLookup))
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusForbidden, rec.Code)
assert.Equal(t, "close", rec.Header().Get("Connection"), "private denial must ask the client to drop the connection")
assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"), "private denial must not be cacheable")
}
// A denied client must not keep reusing the warm socket after joining the overlay.
func TestPrivateDeny_HTTP1_ClosesConnection(t *testing.T) {
validator := &switchableTunnelValidator{}
mw := newPrivateMiddleware(t, validator, nil)
srv := startProtectedServer(t, mw, testTunnelIP, acceptAllLookup, false)
client := srv.Client()
resp := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusForbidden, resp.status)
assert.Equal(t, 1, resp.protoMajor, "plain httptest server must speak HTTP/1.1")
// The Go client folds "Connection: close" into resp.close and drops the header.
assert.True(t, resp.close, "private denial must make the client mark the connection as not reusable")
assert.Equal(t, "no-store", resp.cacheControl, "private denial must not be cacheable")
validator.setValid(true)
resp2 := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusOK, resp2.status, "the retry must reach the upstream once the peer is valid")
assert.False(t, resp2.reused, "the retry must open a new connection")
}
// On HTTP/2 the header becomes a GOAWAY and the retry must use a new connection.
func TestPrivateDeny_HTTP2_SendsGoAway(t *testing.T) {
validator := &switchableTunnelValidator{}
mw := newPrivateMiddleware(t, validator, nil)
srv := startProtectedServer(t, mw, testTunnelIP, acceptAllLookup, true)
client := srv.Client()
resp := doTraced(t, client, srv.URL)
require.Equal(t, 2, resp.protoMajor, "test client must negotiate HTTP/2")
assert.Equal(t, http.StatusForbidden, resp.status)
assert.Empty(t, resp.connection, "HTTP/2 must not carry a Connection header on the wire")
assert.Equal(t, "no-store", resp.cacheControl)
validator.setValid(true)
resp2 := doTraced(t, client, srv.URL)
assert.Equal(t, 2, resp2.protoMajor)
assert.Equal(t, http.StatusOK, resp2.status, "the retry must reach the upstream once the peer is valid")
assert.False(t, resp2.reused, "GOAWAY must retire the connection so the retry opens a new one")
}
// Legitimate private traffic keeps its keep-alive connection.
func TestPrivateAllow_KeepsConnection(t *testing.T) {
validator := &switchableTunnelValidator{valid: true}
mw := newPrivateMiddleware(t, validator, nil)
srv := startProtectedServer(t, mw, testTunnelIP, acceptAllLookup, false)
client := srv.Client()
resp := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusOK, resp.status)
assert.Empty(t, resp.connection, "an allowed private request must not close the connection")
resp2 := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusOK, resp2.status)
assert.True(t, resp2.reused, "allowed private traffic must keep reusing the connection")
}
// Public denials keep the connection open; only private services change.
func TestPublicDeny_KeepsConnection(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
filter := restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"10.0.0.0/8"}})
require.NoError(t, mw.AddDomain(testServerHost, nil, "", 0, "acct-1", "svc-1", filter, false, nil))
srv := startProtectedServer(t, mw, netip.MustParseAddr("192.168.1.1"), nil, false)
client := srv.Client()
resp := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusForbidden, resp.status)
assert.Empty(t, resp.connection, "public denial must not close the connection")
assert.Empty(t, resp.cacheControl, "public denial must not gain cache headers")
resp2 := doTraced(t, client, srv.URL)
assert.Equal(t, http.StatusForbidden, resp2.status)
assert.True(t, resp2.reused, "public denials must keep reusing the connection")
}
// IP restriction denials on a private service must close the connection too.
func TestCheckIPRestrictions_PrivateDenialClosesConnection(t *testing.T) {
filter := restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"100.64.0.0/16"}})
mw := newPrivateMiddleware(t, &switchableTunnelValidator{valid: true}, filter)
handler := mw.Protect(newPassthroughHandler())
tests := []struct {
name string
remoteAddr string
}{
{"denied by CIDR", "100.65.5.6:5000"},
{"unresolvable client address", "not-an-ip:1234"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "http://"+testServerHost+"/", nil)
req.RemoteAddr = tt.remoteAddr
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusForbidden, rec.Code)
assert.Equal(t, "close", rec.Header().Get("Connection"), "private IP-restriction denial must close the connection")
assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"), "private IP-restriction denial must not be cacheable")
})
}
}
func TestCheckIPRestrictions_PublicDenialKeepsConnection(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
filter := restrict.ParseFilter(restrict.FilterConfig{AllowedCIDRs: []string{"10.0.0.0/8"}})
require.NoError(t, mw.AddDomain(testServerHost, nil, "", 0, "acct-1", "svc-1", filter, false, nil))
handler := mw.Protect(newPassthroughHandler())
tests := []struct {
name string
remoteAddr string
}{
{"denied by CIDR", "192.168.1.1:5000"},
{"unresolvable client address", "not-an-ip:1234"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "http://"+testServerHost+"/", nil)
req.RemoteAddr = tt.remoteAddr
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
assert.Equal(t, http.StatusForbidden, rec.Code)
assert.Empty(t, rec.Header().Get("Connection"), "public IP-restriction denial must not close the connection")
assert.Empty(t, rec.Header().Get("Cache-Control"), "public IP-restriction denial must not gain cache headers")
})
}
}
+167
View File
@@ -0,0 +1,167 @@
package auth
import (
"context"
"crypto/tls"
"net/http"
"net/http/httptest"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
"github.com/netbirdio/netbird/proxy/auth"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/shared/management/proto"
)
// denyingSessionValidator mimics management for a user who completed OIDC login
// but is outside the service's distribution groups: ValidateSession denies.
type denyingSessionValidator struct {
calls int
}
func (d *denyingSessionValidator) ValidateSession(context.Context, *proto.ValidateSessionRequest, ...grpc.CallOption) (*proto.ValidateSessionResponse, error) {
d.calls++
return &proto.ValidateSessionResponse{Valid: false, UserId: "user-1", DeniedReason: "not_in_group"}, nil
}
func (d *denyingSessionValidator) ValidateTunnelPeer(context.Context, *proto.ValidateTunnelPeerRequest, ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) {
return &proto.ValidateTunnelPeerResponse{Valid: false}, nil
}
// TestProtect_SelfInstalledCookieCannotBypassGroupCheck is the regression guard
// for the group-authorisation bypass: a user denied at login still holds the raw
// session token from the ?session_token= redirect, so pasting it into the
// nb_session cookie must not buy access. The cookie path validated only the JWT
// signature, which turned the token management had already refused into a bearer
// credential for the service.
func TestProtect_SelfInstalledCookieCannotBypassGroupCheck(t *testing.T) {
validator := &denyingSessionValidator{}
mw := NewMiddleware(log.StandardLogger(), validator, nil)
kp := generateTestKeyPair(t)
oidc := &stubScheme{method: auth.MethodOIDC, authFn: func(r *http.Request) (string, string, error) {
return r.URL.Query().Get("session_token"), "https://idp.example/authorize", nil
}}
require.NoError(t, mw.AddDomain("example.com", []Scheme{oidc}, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, false, []string{"grp-allowed"}))
// The token a denied user gets to see: validly signed for this service and
// domain, but carrying no group the service allows.
token, err := sessionkey.SignToken(kp.PrivateKey, "user-1", "john.doe@example.com", "example.com", auth.MethodOIDC, nil, nil, time.Hour)
require.NoError(t, err)
backendHits := 0
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
backendHits++
w.WriteHeader(http.StatusOK)
}))
t.Run("token in the callback URL is denied", func(t *testing.T) {
rec := serveWithCookie(t, handler, "https://example.com/?session_token="+token, nil)
assert.Equal(t, http.StatusForbidden, rec.Code, "group check must deny the login")
assert.Empty(t, rec.Result().Cookies(), "a denied login must not install a session cookie")
})
t.Run("same token pasted into the session cookie is denied", func(t *testing.T) {
rec := serveWithCookie(t, handler, "https://example.com/", &http.Cookie{Name: auth.SessionCookieName, Value: token})
assert.NotEqual(t, http.StatusOK, rec.Code, "a self-installed cookie must not reach the backend")
assert.Equal(t, 0, backendHits, "backend must never be reached without an allowed group")
})
}
// TestProtect_SessionCookieWithAllowedGroupPassesThrough is the positive half of
// the group gate: a member of an allowed group keeps the cookie fast-path, with
// no management round-trip.
func TestProtect_SessionCookieWithAllowedGroupPassesThrough(t *testing.T) {
validator := &denyingSessionValidator{}
mw := NewMiddleware(log.StandardLogger(), validator, nil)
kp := generateTestKeyPair(t)
oidc := &stubScheme{method: auth.MethodOIDC}
require.NoError(t, mw.AddDomain("example.com", []Scheme{oidc}, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, false, []string{"grp-other", "grp-allowed"}))
token, err := sessionkey.SignToken(kp.PrivateKey, "user-2", "jane@example.com", "example.com", auth.MethodOIDC,
[]string{"grp-unrelated", "grp-allowed"}, []string{"Unrelated", "Allowed"}, time.Hour)
require.NoError(t, err)
handler := mw.Protect(newPassthroughHandler())
rec := serveWithCookie(t, handler, "https://example.com/", &http.Cookie{Name: auth.SessionCookieName, Value: token})
assert.Equal(t, http.StatusOK, rec.Code, "a cookie carrying an allowed group must pass through")
assert.Equal(t, 0, validator.calls, "the cookie fast-path must not call management")
}
// TestProtect_NonOIDCSessionCookieIgnoresGroupRestriction locks the scope of the
// gate: PIN, password and header credentials carry no group identity and are
// authorised by the secret itself, exactly as management validates them.
func TestProtect_NonOIDCSessionCookieIgnoresGroupRestriction(t *testing.T) {
mw := NewMiddleware(log.StandardLogger(), nil, nil)
kp := generateTestKeyPair(t)
scheme := &stubScheme{method: auth.MethodPIN, promptID: "pin"}
require.NoError(t, mw.AddDomain("example.com", []Scheme{scheme}, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, false, []string{"grp-allowed"}))
token, err := sessionkey.SignToken(kp.PrivateKey, "pin-user", "", "example.com", auth.MethodPIN, nil, nil, time.Hour)
require.NoError(t, err)
handler := mw.Protect(newPassthroughHandler())
rec := serveWithCookie(t, handler, "https://example.com/", &http.Cookie{Name: auth.SessionCookieName, Value: token})
assert.Equal(t, http.StatusOK, rec.Code, "a PIN session must not be gated on OIDC group membership")
}
func TestSessionGroupsAllowed(t *testing.T) {
allowed := groupSet([]string{"a", "b"})
tests := []struct {
name string
allowed map[string]struct{}
method auth.Method
groups []string
want bool
}{
{"unrestricted service allows a groupless token", nil, auth.MethodOIDC, nil, true},
{"restricted service allows an intersecting token", allowed, auth.MethodOIDC, []string{"c", "b"}, true},
{"restricted service denies a disjoint token", allowed, auth.MethodOIDC, []string{"c"}, false},
{"restricted service denies a groupless token", allowed, auth.MethodOIDC, nil, false},
{"restricted service ignores a pin token", allowed, auth.MethodPIN, nil, true},
{"restricted service ignores a password token", allowed, auth.MethodPassword, nil, true},
{"restricted service ignores a header token", allowed, auth.MethodHeader, nil, true},
{"restricted service denies an unknown method", allowed, auth.Method("totp"), []string{"a"}, false},
{"restricted service denies a token with no method", allowed, auth.Method(""), []string{"a"}, false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, sessionGroupsAllowed(tc.allowed, tc.method, tc.groups))
})
}
}
func TestGroupSetDropsEmptyEntries(t *testing.T) {
assert.Nil(t, groupSet(nil), "no groups means unrestricted")
assert.Nil(t, groupSet([]string{"", ""}), "blank ids must not restrict access to nothing reachable")
assert.Equal(t, map[string]struct{}{"a": {}}, groupSet([]string{"a", ""}))
}
// serveWithCookie drives the middleware over TLS with captured data attached,
// optionally carrying a session cookie.
func serveWithCookie(t *testing.T, handler http.Handler, url string, cookie *http.Cookie) *httptest.ResponseRecorder {
t.Helper()
req := httptest.NewRequest(http.MethodGet, url, nil)
req.TLS = &tls.ConnectionState{}
if cookie != nil {
req.AddCookie(cookie)
}
req = req.WithContext(proxy.WithCapturedData(req.Context(), proxy.NewCapturedData("")))
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
return rec
}
+5 -5
View File
@@ -44,7 +44,7 @@ func (s *stubSessionValidator) ValidateTunnelPeer(_ context.Context, in *proto.V
func newTunnelMiddleware(t *testing.T, validator SessionValidator) *Middleware {
t.Helper()
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("svc.example", nil, "", 0, "acct-1", "svc-1", nil, false))
require.NoError(t, mw.AddDomain("svc.example", nil, "", 0, "acct-1", "svc-1", nil, false, nil))
return mw
}
@@ -235,8 +235,8 @@ func TestForwardWithTunnelPeer_RoutesAccountIDIntoCacheKey(t *testing.T) {
}
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("svc-a.example", nil, "", 0, "acct-a", "svc-a", nil, false))
require.NoError(t, mw.AddDomain("svc-b.example", nil, "", 0, "acct-b", "svc-b", nil, false))
require.NoError(t, mw.AddDomain("svc-a.example", nil, "", 0, "acct-a", "svc-a", nil, false, nil))
require.NoError(t, mw.AddDomain("svc-b.example", nil, "", 0, "acct-b", "svc-b", nil, false, nil))
// The fast-path requires the inbound-listener marker on the context.
// The peerstore lookup itself is account-agnostic at this level
@@ -299,7 +299,7 @@ func TestForwardWithTunnelPeer_LocalLookupShortCircuitDoesNotPopulateCache(t *te
func TestPrivateService_FailsClosedOnTunnelPeerFailure(t *testing.T) {
mw := NewMiddleware(log.New(), nil, nil)
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true, nil))
called := false
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -328,7 +328,7 @@ func TestPrivateService_ForwardsOnTunnelPeerSuccess(t *testing.T) {
},
}
mw := NewMiddleware(log.New(), validator, nil)
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true, nil))
called := false
handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+176 -13
View File
@@ -105,6 +105,20 @@ type Handler struct {
startTime time.Time
templates *template.Template
templateMu sync.RWMutex
// setPerformance applies a buffer cap to one client. Held as a field so
// tests can drive applyBufferCap without a live embedded client.
setPerformance func(*nbembed.Client, uint32) error
perfMu sync.Mutex
perfInflight map[types.AccountID]*perfWorker
}
// perfWorker is the single in-flight retune for one account. err is valid once
// done is closed.
type perfWorker struct {
done chan struct{}
err error
}
// NewHandler creates a new debug handler.
@@ -113,10 +127,11 @@ func NewHandler(provider clientProvider, healthChecker healthChecker, logger *lo
logger = log.StandardLogger()
}
h := &Handler{
provider: provider,
health: healthChecker,
logger: logger,
startTime: time.Now(),
provider: provider,
health: healthChecker,
logger: logger,
startTime: time.Now(),
setPerformance: setClientPerformance,
}
if err := h.loadTemplates(); err != nil {
logger.Errorf("failed to load embedded templates: %v", err)
@@ -716,15 +731,7 @@ func (h *Handler) handlePerf(w http.ResponseWriter, r *http.Request) {
}
capN := uint32(n)
applied := 0
failed := map[string]string{}
for accountID, client := range h.provider.ListClientsForStartup() {
if err := client.SetPerformance(nbembed.Performance{PreallocatedBuffersPerPool: &capN}); err != nil {
failed[string(accountID)] = err.Error()
continue
}
applied++
}
applied, failed, inFlight := h.applyBufferCap(capN)
resp := map[string]any{
"success": true,
@@ -734,9 +741,165 @@ func (h *Handler) handlePerf(w http.ResponseWriter, r *http.Request) {
if len(failed) > 0 {
resp["failed"] = failed
}
if len(inFlight) > 0 {
resp["in_flight"] = inFlight
}
h.writeJSON(w, resp)
}
// perfApplyTimeout bounds the whole apply, however many clients are registered.
// A var, not a const, so tests can shorten the wait.
var perfApplyTimeout = 5 * time.Second
type perfResult struct {
accountID types.AccountID
err error
}
// setClientPerformance is the production implementation behind Handler.setPerformance.
func setClientPerformance(client *nbembed.Client, capN uint32) error {
return client.SetPerformance(nbembed.Performance{PreallocatedBuffersPerPool: &capN})
}
// collectBuffered takes every result already sitting in the channel, removing
// those accounts from pending, and returns how many of them succeeded. It is
// called when the deadline fires: select picks at random among ready cases, so
// a result that landed in time would otherwise be reported as a timeout.
func collectBuffered(results <-chan perfResult, pending map[types.AccountID]*perfWorker, failed map[string]string) int {
applied := 0
for {
select {
case res := <-results:
delete(pending, res.accountID)
if res.err != nil {
failed[string(res.accountID)] = res.err.Error()
continue
}
applied++
default:
return applied
}
}
}
// resolvePending closes out the accounts still pending when the deadline fires.
// A worker whose done channel is closed has finished, whatever the results
// channel has managed to deliver, so its own error is the truth; the rest are
// genuinely still running and are reported as timed out. Returns how many of
// them had in fact succeeded.
func resolvePending(pending map[types.AccountID]*perfWorker, failed map[string]string) int {
applied := 0
for accountID, w := range pending {
select {
case <-w.done:
if w.err != nil {
failed[string(accountID)] = w.err.Error()
continue
}
applied++
default:
failed[string(accountID)] = fmt.Sprintf("timed out after %s waiting for the client", perfApplyTimeout)
}
}
return applied
}
// startPerfWorker returns the in-flight retune for the account, starting one if
// there is none. The bool reports whether this call started it.
//
// At most one retune runs per account at a time. A client wedged inside its own
// lock never returns, so without this a caller could add one permanently blocked
// goroutine per request just by retrying the endpoint.
func (h *Handler) startPerfWorker(accountID types.AccountID, client *nbembed.Client, capN uint32, results chan<- perfResult) (*perfWorker, bool) {
h.perfMu.Lock()
defer h.perfMu.Unlock()
if w, ok := h.perfInflight[accountID]; ok {
return w, false
}
w := &perfWorker{done: make(chan struct{})}
if h.perfInflight == nil {
h.perfInflight = make(map[types.AccountID]*perfWorker)
}
h.perfInflight[accountID] = w
go func() {
err := h.setPerformance(client, capN)
w.err = err
close(w.done)
// Publish before touching the registry: perfMu is taken once per
// account by every caller walking the fleet, so a finishing worker
// can queue behind a long apply and miss its own deadline.
results <- perfResult{accountID: accountID, err: err}
h.perfMu.Lock()
delete(h.perfInflight, accountID)
h.perfMu.Unlock()
}()
return w, true
}
// applyBufferCap sets the WireGuard buffer pool cap on every registered client
// and reports how many took it, a per-account error for those that did not, and
// the accounts whose earlier retune has not come back yet.
//
// Clients are handled concurrently and the wait is bounded: SetPerformance goes
// through the embedded client's lock, which Start and Stop hold for as long as
// they take - and on a wedged client Stop never returns. This endpoint is the
// recovery path for exactly that fleet, so one stuck account must neither delay
// the others nor accumulate goroutines across retries.
func (h *Handler) applyBufferCap(capN uint32) (int, map[string]string, []string) {
clients := h.provider.ListClientsForStartup()
results := make(chan perfResult, len(clients))
applied := 0
failed := map[string]string{}
var inFlight []string
pending := make(map[types.AccountID]*perfWorker, len(clients))
for accountID, client := range clients {
w, started := h.startPerfWorker(accountID, client, capN, results)
if started {
pending[accountID] = w
continue
}
// Another request owns this account's retune. Take its result if it
// has already landed, otherwise report it as still running instead of
// waiting on it again.
select {
case <-w.done:
if w.err != nil {
failed[string(accountID)] = w.err.Error()
continue
}
applied++
default:
inFlight = append(inFlight, string(accountID))
}
}
deadline := time.After(perfApplyTimeout)
for range len(pending) {
select {
case res := <-results:
delete(pending, res.accountID)
if res.err != nil {
failed[string(res.accountID)] = res.err.Error()
continue
}
applied++
case <-deadline:
applied += collectBuffered(results, pending, failed)
applied += resolvePending(pending, failed)
return applied, failed, inFlight
}
}
return applied, failed, inFlight
}
// handleRuntime returns cheap runtime and process stats. Safe to hit on a
// running proxy; does not read pprof profiles.
func (h *Handler) handleRuntime(w http.ResponseWriter, _ *http.Request) {
+158
View File
@@ -0,0 +1,158 @@
package debug
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
nbembed "github.com/netbirdio/netbird/client/embed"
"github.com/netbirdio/netbird/proxy/internal/health"
"github.com/netbirdio/netbird/proxy/internal/roundtrip"
"github.com/netbirdio/netbird/proxy/internal/types"
)
// perfProvider serves a fixed set of accounts. The clients are nil: the tests
// drive Handler.setPerformance, which never dereferences them.
type perfProvider struct {
accounts []types.AccountID
}
func (p *perfProvider) GetClient(types.AccountID) (*nbembed.Client, bool) { return nil, false }
func (p *perfProvider) ListClientsForDebug() map[types.AccountID]roundtrip.ClientDebugInfo {
return nil
}
func (p *perfProvider) ListClientsForStartup() map[types.AccountID]*nbembed.Client {
out := make(map[types.AccountID]*nbembed.Client, len(p.accounts))
for _, id := range p.accounts {
out[id] = nil
}
return out
}
type stubHealth struct{}
func (stubHealth) ReadinessProbe() bool { return true }
func (stubHealth) StartupProbe(context.Context) bool { return true }
func (stubHealth) CheckClientsConnected(context.Context) (bool, map[types.AccountID]health.ClientHealth) {
return true, nil
}
func shortenPerfTimeout(t *testing.T, d time.Duration) {
t.Helper()
prev := perfApplyTimeout
perfApplyTimeout = d
t.Cleanup(func() { perfApplyTimeout = prev })
}
// TestCollectBufferedCountsResultsReadyAtTheDeadline covers the select-ordering
// trap: when the deadline fires, results already buffered must be counted, not
// reported as timeouts. Driving collectBuffered directly keeps it deterministic
// - through applyBufferCap the two select cases race by construction.
func TestCollectBufferedCountsResultsReadyAtTheDeadline(t *testing.T) {
results := make(chan perfResult, 3)
results <- perfResult{accountID: "ok"}
results <- perfResult{accountID: "broken", err: errors.New("boom")}
pending := map[types.AccountID]*perfWorker{
"ok": {done: make(chan struct{})},
"broken": {done: make(chan struct{})},
"wedged": {done: make(chan struct{})},
}
failed := map[string]string{}
applied := collectBuffered(results, pending, failed)
if applied != 1 {
t.Fatalf("applied = %d, want 1", applied)
}
if failed["broken"] != "boom" {
t.Fatalf("failed = %v, want the error recorded for \"broken\"", failed)
}
if _, ok := pending["wedged"]; !ok || len(pending) != 1 {
t.Fatalf("pending = %v, want only the account that never answered", pending)
}
}
// TestApplyBufferCapSingleFlightPerAccount covers the goroutine accumulation
// reported on PR #7452: repeated calls against a client stuck in its own lock
// must not start a second attempt for the same account.
func TestApplyBufferCapSingleFlightPerAccount(t *testing.T) {
shortenPerfTimeout(t, 50*time.Millisecond)
release := make(chan struct{})
t.Cleanup(func() { close(release) })
var calls atomic.Int32
h := &Handler{
provider: &perfProvider{accounts: []types.AccountID{"wedged"}},
health: stubHealth{},
setPerformance: func(_ *nbembed.Client, _ uint32) error {
calls.Add(1)
<-release
return nil
},
}
for i := range 5 {
applied, failed, inFlight := h.applyBufferCap(4096)
if applied != 0 {
t.Fatalf("call %d: applied = %d, want 0", i, applied)
}
if i == 0 {
if len(failed) != 1 {
t.Fatalf("first call: failed = %v, want the account reported as timed out", failed)
}
continue
}
if len(inFlight) != 1 {
t.Fatalf("call %d: inFlight = %v, want the account reported as still running", i, inFlight)
}
if len(failed) != 0 {
t.Fatalf("call %d: failed = %v, want empty while the retune is in flight", i, failed)
}
}
if got := calls.Load(); got != 1 {
t.Fatalf("setPerformance called %d times, want 1: each retry started another blocked worker", got)
}
}
// TestResolvePendingTrustsFinishedWorkers covers the reporting race cubic
// flagged on PR #7452: a retune that finished just before the deadline must be
// reported by its outcome, not as a timeout, whatever the results channel has
// delivered so far.
func TestResolvePendingTrustsFinishedWorkers(t *testing.T) {
ok := &perfWorker{done: make(chan struct{})}
close(ok.done)
broken := &perfWorker{done: make(chan struct{}), err: errors.New("boom")}
close(broken.done)
stillRunning := &perfWorker{done: make(chan struct{})}
pending := map[types.AccountID]*perfWorker{
"ok": ok,
"broken": broken,
"running": stillRunning,
}
failed := map[string]string{}
applied := resolvePending(pending, failed)
if applied != 1 {
t.Fatalf("applied = %d, want 1", applied)
}
if failed["broken"] != "boom" {
t.Fatalf("failed[broken] = %q, want the worker's own error", failed["broken"])
}
if _, ok := failed["ok"]; ok {
t.Fatalf("failed = %v, want no entry for the account that succeeded", failed)
}
if got := failed["running"]; got == "" || got == "boom" {
t.Fatalf("failed[running] = %q, want the timeout message", got)
}
}
+8
View File
@@ -13,6 +13,14 @@ func NormalizeBedrockModel(modelID string) string {
return sharedllm.NormalizeBedrockModel(modelID)
}
// NormalizeAnthropicModel strips the trailing "-YYYYMMDD" release-date suffix
// from an Anthropic model id so a dated id a client pins matches the undated
// one the operator registered. Thin delegate to shared/llm for the same
// contract reason as the two below.
func NormalizeAnthropicModel(modelID string) string {
return sharedllm.NormalizeAnthropicModel(modelID)
}
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
// so it matches the catalog/pricing key. Thin delegate to shared/llm, kept
// beside NormalizeBedrockModel for the same contract reason.
+14 -2
View File
@@ -10,6 +10,8 @@ package pricing
import (
"fmt"
"math"
sharedllm "github.com/netbirdio/netbird/shared/llm"
)
// Entry is a single model's input and output pricing, expressed in USD per
@@ -92,7 +94,10 @@ func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) {
return &Table{entries: entries}, nil
}
// Lookup returns the entry for the given provider surface and model.
// Lookup returns the entry for the given provider surface and model. A
// dated Anthropic id falls back to its undated form, so a client pinning
// "claude-sonnet-4-5-20250929" bills at the registered "claude-sonnet-4-5"
// rate instead of recording no cost at all.
func (t *Table) Lookup(provider, model string) (Entry, bool) {
if t == nil {
return Entry{}, false
@@ -101,7 +106,14 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) {
if !ok {
return Entry{}, false
}
e, ok := byModel[model]
if e, found := byModel[model]; found {
return e, true
}
undated := sharedllm.NormalizeAnthropicModel(model)
if undated == model {
return Entry{}, false
}
e, ok := byModel[undated]
return e, ok
}
@@ -175,3 +175,22 @@ func TestNewTable_NilAndEmpty(t *testing.T) {
require.NoError(t, err)
assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map")
}
// TestLookup_DatedAnthropicIDFallsBackToUndated covers a client pinning a
// release date on a model priced under its undated id. Without the
// fallback the request records no cost at all.
func TestLookup_DatedAnthropicIDFallsBackToUndated(t *testing.T) {
table, err := NewTable(map[string]map[string]EntryJSON{
"anthropic": {
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015},
},
})
require.NoError(t, err, "table must build from a valid defaults map")
entry, ok := table.Lookup("anthropic", "claude-sonnet-4-5-20250929")
require.True(t, ok, "a dated id must resolve to the undated entry")
assert.InDelta(t, 0.003, entry.InputPer1K, 1e-9, "dated id must bill at the registered rate")
_, ok = table.Lookup("anthropic", "claude-sonnet-9-9-20250929")
assert.False(t, ok, "an unknown family must stay unpriced")
}
@@ -11,6 +11,7 @@ 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"
)
@@ -175,13 +176,28 @@ 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 := m.perRecord[recordID][model]; ok {
if entry, ok := perRecordEntry(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.
//
@@ -84,8 +84,10 @@ 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, model, modelPresent); denial != nil {
if denial := m.evaluateAllowlist(providerID, surface, model, modelPresent, nonInference == "true"); denial != nil {
return denial, nil
}
@@ -114,7 +116,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, model string, modelPresent bool) *middleware.Output {
func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelPresent, nonInference bool) *middleware.Output {
if len(m.cfg.ProviderAllowlists) == 0 {
return nil
}
@@ -122,7 +124,7 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo
// 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("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
if !restricted {
@@ -133,18 +135,29 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo
// 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) == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
if nonInference {
return nil
}
return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
if modelInAllowlist(allowlist, model) {
return nil
}
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
return denyModel(surface, model, denyCodeModel, denyMessageModel, denyReasonModel)
}
// denyModel builds a 403 deny Output for a model-allowlist rejection. model is
// included in the details only when non-empty.
func denyModel(model, code, message, reason string) *middleware.Output {
func denyModel(surface, model, code, message, reason string) *middleware.Output {
details := map[string]string{}
if model != "" {
details["model"] = model
@@ -156,6 +169,7 @@ func denyModel(model, code, message, reason string) *middleware.Output {
Code: code,
Message: message,
Details: details,
Surface: surface,
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
@@ -343,3 +343,52 @@ 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")
})
}
@@ -217,6 +217,32 @@ 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
@@ -225,6 +251,8 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut
// 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
@@ -245,6 +273,9 @@ 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
@@ -704,3 +704,57 @@ 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")
}
@@ -84,6 +84,15 @@ 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
@@ -117,7 +126,7 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew
}
if resp.GetDecision() == "deny" {
return denyFromManagement(resp), nil
return denyFromManagement(resp, lookupKV(in.Metadata, middleware.KeyLLMProvider)), nil
}
return allowFromManagement(resp), nil
}
@@ -161,7 +170,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) *middleware.Output {
func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse, surface string) *middleware.Output {
code := resp.GetDenyCode()
if code == "" {
code = "llm_policy.cap_exceeded"
@@ -176,6 +185,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
DenyReason: &middleware.DenyReason{
Code: code,
Message: denyMessageForCode(code),
Surface: surface,
},
Metadata: []middleware.KV{
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
@@ -224,3 +224,35 @@ 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")
}
@@ -1,9 +1,13 @@
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) {
@@ -36,3 +40,25 @@ 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")
}
@@ -61,6 +61,8 @@ func (middlewareImpl) MetadataKeys() []string {
middleware.KeyLLMRequestPromptRaw,
middleware.KeyLLMCaptureTruncated,
middleware.KeyLLMSessionID,
middleware.KeyLLMAgentID,
middleware.KeyLLMParentAgentID,
}
}
@@ -72,9 +74,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 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.
// 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.
func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
out := &middleware.Output{Decision: middleware.DecisionAllow}
if in == nil {
@@ -92,9 +94,14 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
return m.invokeBedrock(in, br), nil
}
parser, ok := llm.ParserByName(m.providerID)
// 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))
if !ok {
parser, ok = llm.DetectParser(extractPath(in.URL))
parser, ok = llm.ParserByName(m.providerID)
}
if !ok {
return out, nil
@@ -116,9 +123,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle
}
appendSessionID := func(md []middleware.KV) []middleware.KV {
if sessionID != "" {
return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID})
}
return md
return appendAgentIDs(md, in.Headers)
}
facts, err := parser.ParseRequest(in.Body)
@@ -160,6 +167,41 @@ 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
@@ -173,10 +215,8 @@ 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 {
for _, kv := range headers {
if strings.EqualFold(kv.Key, want) && kv.Value != "" {
return kv.Value
}
if v := headerValue(headers, want); v != "" {
return v
}
}
return ""
@@ -252,6 +292,12 @@ 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
@@ -298,6 +344,7 @@ 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 {
@@ -345,7 +392,9 @@ func trimBedrockNamespace(reqPath string) string {
//
// /model/{modelId}/{action}
//
// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}.
// 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.
// 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.
@@ -369,7 +418,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
return bedrockRequest{}, false
}
switch action {
case "invoke", "converse":
case "invoke", "converse", "count-tokens":
return bedrockRequest{model: model}, true
case "invoke-with-response-stream", "converse-stream":
return bedrockRequest{model: model, stream: true}, true
@@ -397,6 +446,7 @@ 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 {
@@ -45,6 +45,8 @@ 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")
}
@@ -230,6 +232,31 @@ 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")
@@ -416,3 +443,81 @@ 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")
})
}
@@ -0,0 +1,175 @@
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)
}
@@ -1,9 +1,13 @@
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
@@ -28,3 +32,103 @@ 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")
}
// TestRouteClaimsModel_VertexNormalizesCandidate is the Vertex counterpart of
// the Bedrock case above: the parser strips the "@version" suffix from the
// path model, so a provider registered with the versioned form must still
// match the normalized request model.
func TestRouteClaimsModel_VertexNormalizesCandidate(t *testing.T) {
route := ProviderRoute{Vertex: true, Models: []string{"claude-sonnet-4-5@20250929"}}
assert.True(t, routeClaimsModel(route, "claude-sonnet-4-5"),
"raw @version Vertex model must match the normalized request model")
assert.False(t, routeClaimsModel(route, "claude-opus-4-8"),
"a model outside the provider's list must not match")
// Non-Vertex routes keep exact matching (no @version stripping).
openai := ProviderRoute{Models: []string{"gpt-4o@2024"}}
assert.False(t, routeClaimsModel(openai, "gpt-4o"),
"non-Vertex routes must not strip an @version suffix")
}
@@ -36,7 +36,10 @@ type ProviderRoute struct {
// request on a same-vendor route so catch-all gateways of a different
// vendor can't swallow it. Empty disables vendor filtering for this
// route.
Vendor string `json:"vendor,omitempty"`
Vendor string `json:"vendor,omitempty"`
// Vendors lists every parser surface a multi-surface gateway accepts.
// Vendor remains supported for existing single-surface configurations.
Vendors []string `json:"vendors,omitempty"`
Models []string `json:"models"`
UpstreamScheme string `json:"upstream_scheme"`
UpstreamHost string `json:"upstream_host"`
@@ -44,6 +47,19 @@ 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.
@@ -65,6 +81,18 @@ 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
@@ -109,6 +109,10 @@ 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,
}
}
@@ -137,29 +141,26 @@ 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 vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider); vendor == "" {
return denyUnmeterable(), nil
if surface == "" {
return denyUnmeterable(surface), nil
}
route, outcome := m.matchVertex(reqPath, model, in.UserGroups)
switch outcome {
case matchOutcomeFound:
return m.allowWithRoute(route, in.UserGroups), nil
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(model), nil
default:
return denyUnknownModel(model), nil
}
return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil
}
// Bedrock likewise carries the model in the URL path (/model/{id}/{action}),
@@ -167,52 +168,236 @@ 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)
switch outcome {
case matchOutcomeFound:
out := m.allowWithRoute(route, in.UserGroups)
if hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil {
out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix
return m.decide(route, outcome, surface, model, in.UserGroups, func(out *middleware.Output) {
if hadPrefix {
stripBedrockNamespace(out)
}
return out, nil
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(model), nil
default:
return denyUnknownModel(model), nil
}
}), 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
}
// 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
}
vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider)
route, outcome := m.matchRoute(model, vendor, requestPath(in.URL), in.UserGroups)
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 {
switch outcome {
case matchOutcomeFound:
return m.allowWithRoute(route, in.UserGroups), nil
out := m.allowWithRoute(route, surface, userGroups)
if decorate != nil {
decorate(out)
}
return out
case matchOutcomeUnauthorised:
return denyNoAuthorisedRoute(model), nil
return denyNoAuthorisedRoute(surface, model)
default:
return denyUnknownModel(model), nil
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{}{}
}
}
if route.Vertex {
if _, ok := permitted[llm.NormalizeVertexModel(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
}
}
@@ -229,7 +414,7 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
// peer, return matchOutcomeUnauthorised so the caller can emit
// the dedicated no_authorised_provider deny code.
// 3. Vendor precedence: when the request carries a detected vendor
// (llm.provider) and at least one candidate is the same vendor,
// (llm.provider) and at least one candidate declares that vendor,
// drop the rest — a vendor-tagged request must never cross to
// another vendor's route (e.g. an Anthropic call landing on an
// OpenAI-compatible gateway that also claims the model).
@@ -252,9 +437,9 @@ func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []stri
// Vendor pinning runs BEFORE the group filter so a request the parser
// tagged with a vendor can never cross to another vendor's route — not
// even an authorised one. Narrow to same-vendor routes when any
// model-matched route declares that vendor; setups with no vendor tag on
// any route fall through unchanged. After narrowing, if no same-vendor
// even an authorised one. Narrow to supporting routes when any
// model-matched route declares that vendor; setups with no matching vendor
// declaration fall through unchanged. After narrowing, if no supporting
// route authorises the caller, that's matchOutcomeUnauthorised (no
// cross-vendor fallback).
if vendor != "" {
@@ -300,12 +485,91 @@ func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []stri
return best, matchOutcomeFound
}
// 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.
// 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).
func isModelLessPath(reqPath string) bool {
return reqPath == "/v1/models" || strings.HasPrefix(reqPath, "/v1/models/")
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
}
// isVertexPath reports whether reqPath is a Google Vertex AI publisher
@@ -332,20 +596,33 @@ 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} 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.
// 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.
func isBedrockPath(reqPath string) bool {
native, _ := splitBedrockNamespace(reqPath)
if !strings.HasPrefix(native, "/model/") {
return false
}
return strings.HasSuffix(native, "/invoke") ||
strings.HasSuffix(native, "/invoke-with-response-stream") ||
strings.HasSuffix(native, "/converse") ||
strings.HasSuffix(native, "/converse-stream")
for _, action := range bedrockActions {
if strings.HasSuffix(native, action) {
return true
}
}
return false
}
// matchVertex selects the Vertex provider authorised for the caller's groups
@@ -425,19 +702,42 @@ 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 string, userGroups []string) (ProviderRoute, matchOutcome) {
if !isModelLessPath(reqPath) {
func (m *Middleware) matchModelless(reqPath, method string, userGroups []string) (ProviderRoute, matchOutcome) {
if !isNonInferenceMethod(method) {
return ProviderRoute{}, matchOutcomeUnknownModel
}
var candidates []ProviderRoute
for _, route := range m.cfg.Providers {
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):
// 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.
if route.Vertex || route.Bedrock {
continue
}
if routeAuthorisesGroups(route, userGroups) {
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) {
candidates = append(candidates, route)
}
}
@@ -510,21 +810,31 @@ func authorisingGroupsCSV(routeGroups, userGroups []string) string {
return strings.Join(out, ",")
}
// matchingVendor returns the subset of routes whose Vendor equals the
// request's detected vendor. Routes with an empty Vendor never match — an
// untagged route can't be asserted to speak the request's surface, so it
// stays out of the vendor-filtered set (but remains eligible via the
// fall-through when no route matches the vendor at all).
// matchingVendor returns the routes that declare the request's detected
// vendor through either the legacy singular field or the multi-vendor field.
// Untagged routes remain eligible only when no route declares the vendor.
func matchingVendor(routes []ProviderRoute, vendor string) []ProviderRoute {
var out []ProviderRoute
for _, r := range routes {
if r.Vendor == vendor {
if routeSupportsVendor(r, vendor) {
out = append(out, r)
}
}
return out
}
func routeSupportsVendor(route ProviderRoute, vendor string) bool {
if route.Vendor == vendor {
return true
}
for _, candidate := range route.Vendors {
if candidate == vendor {
return true
}
}
return false
}
// explicitlyClaiming returns the subset of routes whose Models list
// names the model exactly. Catch-all routes (empty Models) are excluded,
// so callers can prefer a provider that genuinely declares the model over
@@ -564,6 +874,21 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
return true
}
// Vertex likewise: the parser strips the "@version" suffix from the
// path model, while the operator may register the versioned form.
if route.Vertex && llm.NormalizeVertexModel(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
}
@@ -612,7 +937,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, userGroups []string) *middleware.Output {
func (m *Middleware) allowWithRoute(route ProviderRoute, surface string, userGroups []string) *middleware.Output {
rewrite := &middleware.UpstreamRewrite{
Scheme: route.UpstreamScheme,
Host: route.UpstreamHost,
@@ -634,7 +959,7 @@ func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *m
// request time (cached + auto-refreshed) instead of a static value.
bearer, err := m.gcpBearer(route.GCPServiceAccountKeyB64)
if err != nil {
return denyUpstreamAuth()
return denyUpstreamAuth(surface)
}
authValue = bearer
}
@@ -704,11 +1029,12 @@ 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() *middleware.Output {
func denyUpstreamAuth(surface string) *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 502,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeUpstreamAuth,
Message: "could not obtain upstream credential",
},
@@ -722,11 +1048,12 @@ func denyUpstreamAuth() *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() *middleware.Output {
func denyUnmeterable(surface string) *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",
},
@@ -739,11 +1066,12 @@ func denyUnmeterable() *middleware.Output {
// denyMissingModel returns the deny envelope for a request whose
// envelope has no llm.model metadata.
func denyMissingModel() *middleware.Output {
func denyMissingModel(surface string) *middleware.Output {
return &middleware.Output{
Decision: middleware.DecisionDeny,
DenyStatus: 403,
DenyReason: &middleware.DenyReason{
Surface: surface,
Code: denyCodeNotRoutable,
Message: "missing llm.model on request envelope",
},
@@ -756,11 +1084,12 @@ func denyMissingModel() *middleware.Output {
// denyUnknownModel returns the deny envelope for a model that no
// configured provider claims.
func denyUnknownModel(model string) *middleware.Output {
func denyUnknownModel(surface, 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},
@@ -775,11 +1104,12 @@ func denyUnknownModel(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(model string) *middleware.Output {
func denyNoAuthorisedRoute(surface, 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},
@@ -2,6 +2,7 @@ package llm_router
import (
"context"
"net/http"
"testing"
"github.com/stretchr/testify/assert"
@@ -60,6 +61,8 @@ func TestMiddlewareIdentity(t *testing.T) {
[]string{
middleware.KeyLLMResolvedProviderID,
middleware.KeyLLMAuthorisingGroups,
middleware.KeyLLMNonInference,
middleware.KeyLLMModel,
middleware.KeyLLMPolicyDecision,
middleware.KeyLLMPolicyReason,
},
@@ -171,8 +174,12 @@ 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,
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,
UserGroups: []string{defaultTestGroup},
}
}
@@ -197,6 +204,12 @@ 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) {
@@ -399,6 +412,50 @@ func TestRouter_VendorKeepsOpenAIOffAnthropic(t *testing.T) {
assert.Equal(t, "api.openai.com", out.Mutations.RewriteUpstream.Host, "openai vendor must pin to the openai route despite anthropic being declared first")
}
func TestRouter_MultiVendorGatewayAcceptsBothSurfaces(t *testing.T) {
gateway := ProviderRoute{
ID: "agentgateway",
Vendors: []string{"openai", "anthropic"},
Models: nil,
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "gateway.example.com",
}
other := ProviderRoute{
ID: "other-vendor",
Vendor: "mistral",
Models: nil,
AllowedGroupIDs: []string{defaultTestGroup},
UpstreamScheme: "https",
UpstreamHost: "mistral.example.com",
}
mw := New(Config{Providers: []ProviderRoute{other, gateway}})
for _, tc := range []struct {
name string
vendor string
model string
path string
}{
{name: "OpenAI", vendor: "openai", model: "gpt-4o-mini", path: "/v1/chat/completions"},
{name: "Anthropic", vendor: "anthropic", model: "claude-sonnet-4-5", path: "/v1/messages"},
} {
t.Run(tc.name, func(t *testing.T) {
out, err := mw.Invoke(context.Background(), newInputVendorModelURL(tc.vendor, tc.model, tc.path))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"supported vendor must route through the multi-surface gateway")
require.NotNil(t, out.Mutations)
require.NotNil(t, out.Mutations.RewriteUpstream)
assert.Equal(t, "gateway.example.com", out.Mutations.RewriteUpstream.Host)
provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID)
assert.Equal(t, "agentgateway", provider)
})
}
}
// TestRouter_VendorAbsentFallsBackToModelPath confirms vendor filtering is
// inert when the request carries no detected vendor: routing then relies on
// model/path as before.
@@ -679,6 +736,23 @@ func TestRouter_FactoryRejectsBadJSON(t *testing.T) {
require.Error(t, err, "malformed JSON config must be rejected at chain build time")
}
func TestRouter_FactoryDecodesLegacyAndMultiVendorFields(t *testing.T) {
raw := []byte(`{"providers":[` +
`{"id":"legacy","vendor":"openai","models":[],"upstream_scheme":"https","upstream_host":"openai.example.com","auth_header_name":"Authorization","auth_header_value":"Bearer legacy","allowed_group_ids":["group"]},` +
`{"id":"multi","vendors":["openai","anthropic"],"models":[],"upstream_scheme":"https","upstream_host":"gateway.example.com","auth_header_name":"Authorization","auth_header_value":"Bearer multi","allowed_group_ids":["group"]}` +
`]}`)
resolved, err := Factory{}.New(raw)
require.NoError(t, err)
router, ok := resolved.(*Middleware)
require.True(t, ok, "factory must return the concrete router middleware")
require.Len(t, router.cfg.Providers, 2)
assert.Equal(t, "openai", router.cfg.Providers[0].Vendor,
"the legacy singular field must keep decoding")
assert.Equal(t, []string{"openai", "anthropic"}, router.cfg.Providers[1].Vendors,
"the multi-vendor field must decode both supported surfaces")
}
func TestRouter_FactoryAcceptsEmptyShapes(t *testing.T) {
cases := [][]byte{nil, []byte(""), []byte(" "), []byte("null"), []byte("{}"), []byte("[]")}
for _, raw := range cases {
@@ -873,3 +947,403 @@ 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")
})
}
+1 -1
View File
@@ -264,7 +264,7 @@ func applyMutations(ctx context.Context, d *Dispatcher, spec Spec, r *http.Reque
if m == nil {
return
}
add, remove, blocked := FilterHeaderMutations(m)
add, remove, blocked := filterHeaderMutations(m, spec.ID)
for _, h := range blocked {
d.metrics.IncHeaderMutationBlocked(ctx, spec.ID, h)
}
+59
View File
@@ -2,6 +2,7 @@ package middleware
import (
"context"
"net/http"
"strconv"
"testing"
@@ -278,6 +279,64 @@ func TestChain_ApplyMutations_RewriteGatedOnCanMutate(t *testing.T) {
assert.Nil(t, rewrite, "rewrite must be filtered when CanMutate=false")
}
func TestChain_IdentityInjectReplacesReservedNetBirdHeaders(t *testing.T) {
mw := &fakeMiddleware{
id: "llm_identity_inject",
slot: SlotOnRequest,
mutationsSupported: true,
canMutate: true,
mutations: &Mutations{
HeadersRemove: []string{"x-netbird-user-id", "x-netbird-groups"},
HeadersAdd: []KV{
{Key: "x-netbird-user-id", Value: "trusted-user"},
{Key: "x-netbird-groups", Value: "trusted-group"},
},
},
}
c := chainFor(t, mw)
req, err := http.NewRequest(http.MethodGet, "https://gateway.example.com/v1/models", nil)
require.NoError(t, err)
req.Header.Set("x-netbird-user-id", "spoofed-user")
req.Header.Set("x-netbird-groups", "spoofed-group")
denied, _, _, err := c.RunRequest(context.Background(), req, &Input{}, NewAccumulator(0))
require.NoError(t, err)
assert.Nil(t, denied, "identity injection must not deny the request")
assert.Equal(t, "trusted-user", req.Header.Get("x-netbird-user-id"),
"the built-in identity middleware must replace a spoofed user header")
assert.Equal(t, "trusted-group", req.Header.Get("x-netbird-groups"),
"the built-in identity middleware must replace spoofed groups")
}
func TestChain_OtherMiddlewareCannotReplaceReservedNetBirdHeaders(t *testing.T) {
mw := &fakeMiddleware{
id: "untrusted-middleware",
slot: SlotOnRequest,
mutationsSupported: true,
canMutate: true,
mutations: &Mutations{
HeadersRemove: []string{"x-netbird-user-id", "x-netbird-groups"},
HeadersAdd: []KV{
{Key: "x-netbird-user-id", Value: "replacement-user"},
{Key: "x-netbird-groups", Value: "replacement-group"},
},
},
}
c := chainFor(t, mw)
req, err := http.NewRequest(http.MethodGet, "https://gateway.example.com/v1/models", nil)
require.NoError(t, err)
req.Header.Set("x-netbird-user-id", "original-user")
req.Header.Set("x-netbird-groups", "original-group")
denied, _, _, err := c.RunRequest(context.Background(), req, &Input{}, NewAccumulator(0))
require.NoError(t, err)
assert.Nil(t, denied, "blocked mutations must not deny the request")
assert.Equal(t, "original-user", req.Header.Get("x-netbird-user-id"),
"other middleware must remain unable to mutate reserved identity headers")
assert.Equal(t, "original-group", req.Header.Get("x-netbird-groups"),
"other middleware must remain unable to mutate reserved identity headers")
}
// TestChain_RunRequest_PropagatesUserGroups asserts the chain forwards
// Input.UserGroups verbatim through cloneInputFor so policy-aware
// middlewares (e.g. llm_policy_check) can authorise without an extra
+68
View File
@@ -11,11 +11,78 @@ 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
@@ -36,6 +103,7 @@ 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 {
@@ -0,0 +1,92 @@
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"])
}
+16 -2
View File
@@ -2,6 +2,8 @@ package middleware
import "strings"
const trustedIdentityMiddlewareID = "llm_identity_inject"
var denyHeaders = []string{
"Authorization",
"Connection",
@@ -78,18 +80,22 @@ func isHeaderFieldName(name string) bool {
// header names so the dispatcher can increment the blocked-header
// metric.
func FilterHeaderMutations(m *Mutations) (filteredAdd []KV, filteredRemove []string, blocked []string) {
return filterHeaderMutations(m, "")
}
func filterHeaderMutations(m *Mutations, middlewareID string) (filteredAdd []KV, filteredRemove []string, blocked []string) {
if m == nil {
return nil, nil, nil
}
for _, kv := range m.HeadersAdd {
if IsHeaderMutable(kv.Key) {
if IsHeaderMutable(kv.Key) || isTrustedIdentityHeader(middlewareID, kv.Key) {
filteredAdd = append(filteredAdd, kv)
continue
}
blocked = append(blocked, kv.Key)
}
for _, name := range m.HeadersRemove {
if IsHeaderMutable(name) {
if IsHeaderMutable(name) || isTrustedIdentityHeader(middlewareID, name) {
filteredRemove = append(filteredRemove, name)
continue
}
@@ -97,3 +103,11 @@ func FilterHeaderMutations(m *Mutations) (filteredAdd []KV, filteredRemove []str
}
return filteredAdd, filteredRemove, blocked
}
func isTrustedIdentityHeader(middlewareID, name string) bool {
if middlewareID != trustedIdentityMiddlewareID {
return false
}
return strings.EqualFold(name, "x-netbird-user-id") ||
strings.EqualFold(name, "x-netbird-groups")
}
@@ -0,0 +1,26 @@
package middleware
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestFilterHeaderMutationsDoesNotTrustReservedHeaders(t *testing.T) {
mutations := &Mutations{
HeadersAdd: []KV{
{Key: "x-request-label", Value: "allowed"},
{Key: "x-netbird-user-id", Value: "spoofed-user"},
},
HeadersRemove: []string{"x-request-label", "x-netbird-groups"},
}
filteredAdd, filteredRemove, blocked := FilterHeaderMutations(mutations)
assert.Equal(t, []KV{{Key: "x-request-label", Value: "allowed"}}, filteredAdd,
"the public filter should retain mutable additions")
assert.Equal(t, []string{"x-request-label"}, filteredRemove,
"the public filter should retain mutable removals")
assert.ElementsMatch(t, []string{"x-netbird-user-id", "x-netbird-groups"}, blocked,
"the public filter must not grant the identity middleware exception")
}
+17
View File
@@ -22,6 +22,15 @@ 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"
@@ -66,6 +75,14 @@ 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
+12
View File
@@ -179,6 +179,12 @@ 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
@@ -247,6 +253,12 @@ 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
+237
View File
@@ -0,0 +1,237 @@
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)))
}
@@ -0,0 +1,272 @@
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)
}
+3
View File
@@ -363,6 +363,9 @@ 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))
}
+4 -5
View File
@@ -26,8 +26,8 @@ import (
// branch at all), construct the MultiTransport via NewDirectOnly.
type MultiTransport struct {
embedded http.RoundTripper
direct *http.Transport
insecure *http.Transport
direct *upstreamTransport
insecure *upstreamTransport
}
// errNoEmbeddedTransport is returned when a request reaches the
@@ -53,7 +53,6 @@ func NewMultiTransport(embedded http.RoundTripper, logger *log.Logger) *MultiTra
}
direct := &http.Transport{
DialContext: dialWithTimeout(dialer.DialContext),
ForceAttemptHTTP2: true,
MaxIdleConns: cfg.maxIdleConns,
MaxIdleConnsPerHost: cfg.maxIdleConnsPerHost,
MaxConnsPerHost: cfg.maxConnsPerHost,
@@ -70,8 +69,8 @@ func NewMultiTransport(embedded http.RoundTripper, logger *log.Logger) *MultiTra
return &MultiTransport{
embedded: embedded,
direct: direct,
insecure: insecure,
direct: newUpstreamTransport(direct, cfg.upstreamHTTPVersion, logger),
insecure: newUpstreamTransport(insecure, cfg.upstreamHTTPVersion, logger),
}
}
+60 -4
View File
@@ -5,6 +5,7 @@ import (
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
@@ -75,16 +76,71 @@ func TestMultiTransport_AppliesEnvOverridesToDirect(t *testing.T) {
mt := NewMultiTransport(&stubRoundTripper{body: "embedded"}, nil)
assert.Equal(t, 42, mt.direct.MaxIdleConns,
assert.Equal(t, 42, mt.direct.primary.MaxIdleConns,
"NB_PROXY_MAX_IDLE_CONNS must propagate to the direct transport")
assert.Equal(t, 11*time.Second, mt.direct.IdleConnTimeout,
assert.Equal(t, 11*time.Second, mt.direct.primary.IdleConnTimeout,
"NB_PROXY_IDLE_CONN_TIMEOUT must propagate to the direct transport")
assert.Equal(t, 7*time.Second, mt.direct.TLSHandshakeTimeout,
assert.Equal(t, 7*time.Second, mt.direct.primary.TLSHandshakeTimeout,
"NB_PROXY_TLS_HANDSHAKE_TIMEOUT must propagate to the direct transport")
assert.Equal(t, 42, mt.insecure.MaxIdleConns,
assert.Equal(t, 42, mt.insecure.primary.MaxIdleConns,
"env tuning must also apply to the insecure-skip-verify direct transport")
}
// TestMultiTransport_UpstreamHTTPVersion pins the protocol actually
// negotiated with an HTTPS upstream that offers both h2 and http/1.1.
// The request rides the insecure clone, so this also covers the version
// surviving http.Transport.Clone.
func TestMultiTransport_UpstreamHTTPVersion(t *testing.T) {
tests := []struct {
name string
env string
wantProto string
}{
{name: "unset negotiates h2", env: "", wantProto: "HTTP/2.0"},
{name: "auto negotiates h2", env: "auto", wantProto: "HTTP/2.0"},
{name: "1.1 pins http/1.1", env: "1.1", wantProto: "HTTP/1.1"},
{name: "2 negotiates h2", env: "2", wantProto: "HTTP/2.0"},
{name: "unsupported value keeps the default", env: "http3", wantProto: "HTTP/2.0"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
// t.Setenv registers the restore for whatever the process
// inherited; unsetting afterwards lets the default row
// exercise a genuinely absent variable.
t.Setenv(EnvUpstreamHTTPVersion, tc.env)
if tc.env == "" {
require.NoError(t, os.Unsetenv(EnvUpstreamHTTPVersion))
}
// The test server's certificate isn't in any root pool, so the
// request rides the insecure branch via WithSkipTLSVerify.
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, r.Proto)
}))
srv.EnableHTTP2 = true
srv.StartTLS()
defer srv.Close()
mt := NewDirectOnly(nil)
ctx := WithSkipTLSVerify(WithDirectUpstream(context.Background()))
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
require.NoError(t, err)
resp, err := mt.RoundTrip(req)
require.NoError(t, err)
body, err := io.ReadAll(resp.Body)
_ = resp.Body.Close()
require.NoError(t, err)
assert.Equal(t, tc.wantProto, resp.Proto,
"client-side protocol must follow %s=%q", EnvUpstreamHTTPVersion, tc.env)
assert.Equal(t, tc.wantProto, string(body),
"the upstream must see the same protocol the client negotiated")
})
}
}
// TestMultiTransport_NilEmbeddedErrorsWhenWGPathRequested guards
// against the previous silent fallback: a MultiTransport constructed
// without an embedded transport must reject requests that don't
+30 -13
View File
@@ -30,6 +30,12 @@ import (
const deviceNamePrefix = "ingress-proxy-"
// envProxyRosenpass toggles Rosenpass (permissive) on the embedded proxy client. Defaults to on.
const envProxyRosenpass = "NB_PROXY_ROSENPASS" //nolint:gosec // env var name, not a credential
// envProxyClientLogLevel sets the embedded NetBird client's log level.
const envProxyClientLogLevel = "NB_PROXY_CLIENT_LOG_LEVEL"
const clientStopTimeout = 30 * time.Second
const createProxyPeerTimeout = 30 * time.Second
@@ -76,10 +82,10 @@ type serviceNotification struct {
// clientEntry holds an embedded NetBird client and tracks which services use it.
type clientEntry struct {
client *embed.Client
transport *http.Transport
transport *upstreamTransport
// insecureTransport is a clone of transport with TLS verification disabled,
// used when per-target skip_tls_verify is set.
insecureTransport *http.Transport
insecureTransport *upstreamTransport
services map[ServiceKey]serviceInfo
createdAt time.Time
started bool
@@ -353,11 +359,11 @@ func (n *NetBird) createClientEntry(ctx context.Context, accountID types.Account
// NB_PROXY_CLIENT_LOG_LEVEL (e.g. "trace") to surface the embedded NetBird
// client's relay / signal / handshake detail for local debugging.
clientLogLevel := log.WarnLevel.String()
if v := strings.TrimSpace(os.Getenv("NB_PROXY_CLIENT_LOG_LEVEL")); v != "" {
if v := strings.TrimSpace(os.Getenv(envProxyClientLogLevel)); v != "" {
if lvl, err := log.ParseLevel(v); err == nil {
clientLogLevel = lvl.String()
} else {
n.logger.Warnf("invalid NB_PROXY_CLIENT_LOG_LEVEL %q, using %q: %v", v, clientLogLevel, err)
n.logger.Warnf("invalid %s %q, using %q: %v", envProxyClientLogLevel, v, clientLogLevel, err)
}
}
@@ -367,15 +373,26 @@ func (n *NetBird) createClientEntry(ctx context.Context, accountID types.Account
}
})
// Rosenpass runs in permissive mode by default so the embedded proxy can
// establish connections with Rosenpass-enabled peers (which otherwise fail
// on a PSK mismatch) while still falling back to plain WireGuard for peers
// that do not run Rosenpass. Set NB_PROXY_ROSENPASS=false to disable it.
rosenpassEnabled := true
if v, ok := envBool(envProxyRosenpass, n.logger); ok {
rosenpassEnabled = v
}
// Create embedded NetBird client with the generated private key.
// The peer has already been created via CreateProxyPeer RPC with the public key.
wgPort := int(n.clientCfg.WGPort)
embedOpts := embed.Options{
DeviceName: deviceNamePrefix + n.proxyID,
ManagementURL: n.clientCfg.MgmtAddr,
PrivateKey: privateKey.String(),
LogLevel: clientLogLevel,
BlockInbound: n.clientCfg.BlockInbound,
DeviceName: deviceNamePrefix + n.proxyID,
ManagementURL: n.clientCfg.MgmtAddr,
PrivateKey: privateKey.String(),
LogLevel: clientLogLevel,
BlockInbound: n.clientCfg.BlockInbound,
EnableRosenpass: rosenpassEnabled,
RosenpassPermissive: rosenpassEnabled,
// The embedded proxy peer must never be a stepping stone into
// the proxy host's LAN: it only exists to reach NetBird mesh
// targets or, when direct_upstream is set, the host network
@@ -397,7 +414,6 @@ func (n *NetBird) createClientEntry(ctx context.Context, accountID types.Account
// not work with reverse proxied requests.
transport := &http.Transport{
DialContext: dialWithTimeout(client.DialContext),
ForceAttemptHTTP2: true,
MaxIdleConns: n.transportCfg.maxIdleConns,
MaxIdleConnsPerHost: n.transportCfg.maxIdleConnsPerHost,
MaxConnsPerHost: n.transportCfg.maxConnsPerHost,
@@ -409,15 +425,14 @@ func (n *NetBird) createClientEntry(ctx context.Context, accountID types.Account
ReadBufferSize: n.transportCfg.readBufferSize,
DisableCompression: n.transportCfg.disableCompression,
}
insecureTransport := transport.Clone()
insecureTransport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec
return &clientEntry{
client: client,
services: map[ServiceKey]serviceInfo{key: si},
transport: transport,
insecureTransport: insecureTransport,
transport: newUpstreamTransport(transport, n.transportCfg.upstreamHTTPVersion, n.logger),
insecureTransport: newUpstreamTransport(insecureTransport, n.transportCfg.upstreamHTTPVersion, n.logger),
createdAt: time.Now(),
started: false,
inflightMap: make(map[backendKey]chan struct{}),
@@ -899,6 +914,8 @@ func logEmbedOptions(logger *log.Logger, accountID types.AccountID, serviceID ty
"mtu": mtu,
"block_inbound": opts.BlockInbound,
"block_lan_access": opts.BlockLANAccess,
"rosenpass_enabled": opts.EnableRosenpass,
"rosenpass_permissive": opts.RosenpassPermissive,
"disable_ipv6": opts.DisableIPv6,
"disable_client_routes": opts.DisableClientRoutes,
"no_userspace": opts.NoUserspace,
+108
View File
@@ -1,8 +1,11 @@
package roundtrip
import (
"crypto/tls"
"net/http"
"os"
"strconv"
"strings"
"time"
log "github.com/sirupsen/logrus"
@@ -21,6 +24,30 @@ const (
EnvReadBufferSize = "NB_PROXY_READ_BUFFER_SIZE"
EnvDisableCompression = "NB_PROXY_DISABLE_COMPRESSION"
EnvMaxInflight = "NB_PROXY_MAX_INFLIGHT"
EnvUpstreamHTTPVersion = "NB_PROXY_UPSTREAM_HTTP_VERSION"
)
// upstreamHTTPVersion selects the HTTP version the proxy uses towards an
// upstream. The explicit values are absolute: they mean the same thing
// however the transports are dialled and whatever the default becomes,
// so operator configuration survives a change of default.
type upstreamHTTPVersion string
const (
// upstreamHTTPAuto leaves the choice to the upstream: h2 is offered
// alongside http/1.1 in the TLS handshake and the upstream picks.
// An upstream that picks h2 and then fails to serve it is moved to
// HTTP/1.1 on its own (see upstreamTransport), which is the part
// ALPN cannot express. This is the only value whose meaning tracks
// the proxy's default.
upstreamHTTPAuto upstreamHTTPVersion = "auto"
// upstreamHTTP11 never offers h2, so the upstream sees HTTP/1.1.
upstreamHTTP11 upstreamHTTPVersion = "1.1"
// upstreamHTTP2 offers h2 in the TLS handshake and keeps it there:
// an upstream that negotiates h2 and then breaks is never moved to
// HTTP/1.1. Cleartext upstreams stay on HTTP/1.1 regardless: the
// proxy speaks no h2c.
upstreamHTTP2 upstreamHTTPVersion = "2"
)
// transportConfig holds tunable parameters for the per-account HTTP transport.
@@ -37,6 +64,11 @@ type transportConfig struct {
disableCompression bool
// maxInflight limits per-backend concurrent requests. 0 means unlimited.
maxInflight int
// upstreamHTTPVersion selects the HTTP version used towards HTTPS
// upstreams. The default negotiates it with each upstream; the
// explicit values are for backends whose advertised h2 support is
// unusable and whose failure mode the negotiation cannot see.
upstreamHTTPVersion upstreamHTTPVersion
}
func defaultTransportConfig() transportConfig {
@@ -47,6 +79,7 @@ func defaultTransportConfig() transportConfig {
idleConnTimeout: 90 * time.Second,
tlsHandshakeTimeout: 10 * time.Second,
expectContinueTimeout: 1 * time.Second,
upstreamHTTPVersion: upstreamHTTPAuto,
}
}
@@ -86,6 +119,9 @@ func loadTransportConfig(logger *log.Logger) transportConfig {
if v, ok := envInt(EnvMaxInflight, logger); ok {
cfg.maxInflight = v
}
if v, ok := envUpstreamHTTPVersion(EnvUpstreamHTTPVersion, logger); ok {
cfg.upstreamHTTPVersion = v
}
logger.WithFields(log.Fields{
"max_idle_conns": cfg.maxIdleConns,
@@ -99,11 +135,83 @@ func loadTransportConfig(logger *log.Logger) transportConfig {
"read_buffer_size": cfg.readBufferSize,
"disable_compression": cfg.disableCompression,
"max_inflight": cfg.maxInflight,
"upstream_http_version": cfg.upstreamHTTPVersion,
}).Debug("backend transport configuration")
return cfg
}
// applyUpstreamHTTPVersion configures t's ALPN offer for the requested
// HTTP version. It is the single place that decides which protocols a
// transport offers, so changing the proxy's default only touches this
// function and leaves every explicit operator setting intact. What
// happens when a negotiated h2 upstream then fails belongs to
// upstreamTransport, which owns the runtime half of "auto".
//
// HTTP/1.1 is pinned by clearing ForceAttemptHTTP2 and installing an
// empty TLSNextProto, which disables h2 regardless of how the transport
// is dialled. Relying on net/http's conservative default (h2 off
// whenever a custom dialer is set) would silently start negotiating h2
// again the day a transport switches to DialTLSContext.
func applyUpstreamHTTPVersion(t *http.Transport, version upstreamHTTPVersion) {
if version == upstreamHTTP11 {
t.ForceAttemptHTTP2 = false
t.TLSNextProto = map[string]func(string, *tls.Conn) http.RoundTripper{}
t.TLSClientConfig = withoutHTTP2ALPN(t.TLSClientConfig)
return
}
t.ForceAttemptHTTP2 = true
}
// withoutHTTP2ALPN drops h2 from the ALPN offer. Configuring h2 makes
// net/http append h2 to the transport's TLSClientConfig, so a transport
// cloned from one that already served a request carries that offer with
// it. Left in place, the upstream would select a protocol this
// transport then refuses to speak, and the response would come back as
// h2 frames parsed as an HTTP/1.1 message.
func withoutHTTP2ALPN(cfg *tls.Config) *tls.Config {
// A nil config offers no ALPN at all, which is already HTTP/1.1.
if cfg == nil || len(cfg.NextProtos) == 0 {
return cfg
}
protos := make([]string, 0, len(cfg.NextProtos))
for _, proto := range cfg.NextProtos {
if proto == "h2" {
continue
}
protos = append(protos, proto)
}
if len(protos) == len(cfg.NextProtos) {
return cfg
}
// Clone rather than edit in place: the caller may share this config
// with the transport it was cloned from.
stripped := cfg.Clone()
stripped.NextProtos = protos
return stripped
}
// envUpstreamHTTPVersion reads an upstream HTTP version from the
// environment. An unrecognised value warns and leaves the default in
// place rather than guessing at the operator's intent.
func envUpstreamHTTPVersion(key string, logger *log.Logger) (upstreamHTTPVersion, bool) {
s := strings.TrimSpace(os.Getenv(key))
if s == "" {
return "", false
}
switch v := upstreamHTTPVersion(strings.ToLower(s)); v {
case upstreamHTTPAuto, upstreamHTTP11, upstreamHTTP2:
return v, true
default:
logger.Warnf("ignoring unsupported %s=%q, expected one of %q, %q, %q",
key, s, upstreamHTTPAuto, upstreamHTTP11, upstreamHTTP2)
return "", false
}
}
func envInt(key string, logger *log.Logger) (int, bool) {
s := os.Getenv(key)
if s == "" {
+455
View File
@@ -0,0 +1,455 @@
package roundtrip
import (
"errors"
"net"
"net/http"
"net/netip"
"net/url"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
// upstreamDowngradeTTL is how long an upstream stays pinned to HTTP/1.1
// after an h2 failure that only implied it cannot serve h2. Bounded so a
// fixed or replaced backend returns to h2 without restarting the proxy.
// A pin the upstream asked for itself does not expire — see downgrade.
const upstreamDowngradeTTL = 10 * time.Minute
// downgrade is an upstream's HTTP/1.1 pin.
type downgrade struct {
// expiry is when the pin lapses and the upstream is offered h2
// again. The zero time means it never does: the upstream answered
// HTTP_1_1_REQUIRED, which is a statement about how it is
// configured, not a fault that may clear on its own. Re-probing
// that every upstreamDowngradeTTL would buy nothing but a failed
// request per interval, so the pin holds until the transport goes
// away with the proxy or the account's client.
expiry time.Time
}
// permanent reports whether the upstream asked for this pin itself.
func (d downgrade) permanent() bool {
return d.expiry.IsZero()
}
// active reports whether the pin still stands at now.
func (d downgrade) active(now time.Time) bool {
return d.permanent() || now.Before(d.expiry)
}
// upstreamTransport carries requests to a single upstream family (one
// TLS configuration) and implements what upstreamHTTPAuto means.
//
// ALPN already lets the upstream pick the protocol: primary offers both
// h2 and http/1.1 and the server chooses. What ALPN cannot express is
// an upstream that selects h2 and then fails to speak it — the case
// this type handles. The first h2-level failure for a host pins that
// host to fallback, an HTTP/1.1-only clone of primary, and the request
// is retried there when it can be replayed.
//
// The downgrade is per upstream host, not per transport: one broken
// backend must not drop every other backend to HTTP/1.1.
type upstreamTransport struct {
// primary is the configured transport: h2 offered in ALPN for
// upstreamHTTPAuto and upstreamHTTP2, HTTP/1.1-only for
// upstreamHTTP11.
primary *http.Transport
// version decides whether a downgrade may happen at all. Only
// upstreamHTTPAuto downgrades; the explicit values are absolute.
version upstreamHTTPVersion
logger *log.Logger
// fallbackMu guards the lazy fallback clone: most deployments never
// hit a broken h2 upstream and should not pay for a second
// connection pool.
fallbackMu sync.Mutex
fallback *http.Transport
mu sync.RWMutex
// downgraded maps an upstream host to its HTTP/1.1 pin.
downgraded map[string]downgrade
}
// newUpstreamTransport wraps base for the requested HTTP version. base
// must not be used directly afterwards: the wrapper owns it, including
// its connection pool.
func newUpstreamTransport(base *http.Transport, version upstreamHTTPVersion, logger *log.Logger) *upstreamTransport {
if logger == nil {
logger = log.StandardLogger()
}
applyUpstreamHTTPVersion(base, version)
return &upstreamTransport{
primary: base,
version: version,
logger: logger,
downgraded: make(map[string]downgrade),
}
}
// RoundTrip implements http.RoundTripper.
func (t *upstreamTransport) RoundTrip(req *http.Request) (*http.Response, error) {
if !t.mayDowngrade(req) {
return t.primary.RoundTrip(req)
}
host := upstreamKey(req.URL)
if t.isDowngraded(host) {
return t.http1().RoundTrip(req)
}
resp, err := t.primary.RoundTrip(req)
if err == nil || !isHTTP2ProtocolError(err) {
return resp, err
}
// HTTP_1_1_REQUIRED is the upstream saying it will not serve this
// request over h2 however often it is asked — IIS answers it for
// Windows Authentication and for client-certificate sites, where
// the cause is site configuration rather than a passing fault.
t.markDowngraded(host, isHTTP11Required(err))
if !safeToRetry(req, err) {
// The upstream may have carried out the request before failing
// to answer over h2, and repeating it could duplicate whatever
// it did. The host is pinned either way, so the next request
// goes out over HTTP/1.1.
t.logger.WithFields(log.Fields{
"upstream": host,
"method": req.Method,
}).Debug("not retrying over HTTP/1.1: the upstream may already have applied this request")
return nil, err
}
retry, ok := replayable(req)
if !ok {
// The body is already consumed and cannot be regenerated, so
// this request fails, and the pin carries the next one.
return nil, err
}
return t.http1().RoundTrip(retry)
}
// CloseIdleConnections closes idle connections on both pools.
func (t *upstreamTransport) CloseIdleConnections() {
t.primary.CloseIdleConnections()
if fallback := t.existingHTTP1(); fallback != nil {
fallback.CloseIdleConnections()
}
}
// mayDowngrade reports whether a failed request is a downgrade
// candidate. Only upstreamHTTPAuto downgrades, and only for TLS
// upstreams: the proxy speaks no h2c, so a cleartext upstream is
// already on HTTP/1.1 and an error there says nothing about h2.
func (t *upstreamTransport) mayDowngrade(req *http.Request) bool {
return t.version == upstreamHTTPAuto && req.URL != nil && req.URL.Scheme == "https"
}
func (t *upstreamTransport) isDowngraded(host string) bool {
t.mu.RLock()
pin, ok := t.downgraded[host]
t.mu.RUnlock()
if !ok {
return false
}
if pin.active(time.Now()) {
return true
}
t.mu.Lock()
defer t.mu.Unlock()
// Re-read under the write lock rather than trusting the expired pin
// from above: a concurrent request may have re-pinned the host since,
// and that pin decides this request too. Reporting the stale read
// would send one request back to h2 against a live pin.
pin, ok = t.downgraded[host]
if !ok {
return false
}
if pin.active(time.Now()) {
return true
}
delete(t.downgraded, host)
return false
}
// markDowngraded pins host to HTTP/1.1. permanent marks a pin the
// upstream asked for; anything else lapses after upstreamDowngradeTTL so
// a repaired backend is offered h2 again.
func (t *upstreamTransport) markDowngraded(host string, permanent bool) {
now := time.Now()
pin := downgrade{expiry: now.Add(upstreamDowngradeTTL)}
if permanent {
pin = downgrade{}
}
t.mu.Lock()
previous, pinned := t.downgraded[host]
// A permanent pin is never weakened back into an expiring one: the
// upstream has already said h2 is not on offer.
promoted := pinned && !previous.permanent() && permanent
if !pinned || !previous.permanent() {
t.downgraded[host] = pin
}
for h, existing := range t.downgraded {
if !existing.active(now) {
delete(t.downgraded, h)
}
}
t.mu.Unlock()
// Log a new pin, and a pin the upstream has since asked to make
// permanent — otherwise an operator would only ever see the "for the
// next 10m" line and never learn the upstream settled the question.
if pinned && !promoted {
return
}
entry := t.logger.WithField("upstream", host)
if permanent {
entry.Warnf("upstream answered HTTP_1_1_REQUIRED, using HTTP/1.1 for it from now on")
return
}
entry.Warnf("upstream negotiated HTTP/2 but failed to serve it, using HTTP/1.1 for the next %s (set %s=1.1 to pin it)",
upstreamDowngradeTTL, EnvUpstreamHTTPVersion)
}
// http1 returns the HTTP/1.1-only clone, creating it on first use.
func (t *upstreamTransport) http1() *http.Transport {
t.fallbackMu.Lock()
defer t.fallbackMu.Unlock()
if t.fallback == nil {
fallback := t.primary.Clone()
applyUpstreamHTTPVersion(fallback, upstreamHTTP11)
t.fallback = fallback
}
return t.fallback
}
// existingHTTP1 returns the fallback transport only if it was already
// created, so housekeeping never allocates a second connection pool for
// an upstream that never needed one.
func (t *upstreamTransport) existingHTTP1() *http.Transport {
t.fallbackMu.Lock()
defer t.fallbackMu.Unlock()
return t.fallback
}
// upstreamKey normalizes an authority for use as a pin key, so one
// upstream cannot end up with two independent pins. DNS labels compare
// case-insensitively, and the default HTTPS port is implied — every
// downgrade path is TLS-only, so a bare host and the same host on :443
// are the same upstream.
func upstreamKey(u *url.URL) string {
host := normalizeUpstreamHost(u.Hostname())
port := u.Port()
if port == "" || port == "443" {
return host
}
// JoinHostPort rather than concatenation: an IPv6 literal needs its
// brackets back after Hostname stripped them.
return net.JoinHostPort(host, port)
}
// normalizeUpstreamHost folds the spellings of one host onto a single
// key. An IP literal goes through netip so that the several textual
// forms of one address (case, leading zeroes, a compressed run) collapse
// and a v4-mapped address keys as the v4 address it is. A zone
// identifier is left exactly as written: it names an interface, and
// interface names are case-sensitive on the systems that have them, so
// %eth0 and %ETH0 may be different links and must not share a pin.
// Anything that is not an IP literal is a DNS name, which compares
// case-insensitively.
func normalizeUpstreamHost(host string) string {
if addr, err := netip.ParseAddr(host); err == nil {
return addr.Unmap().String()
}
return strings.ToLower(host)
}
// safeToRetry reports whether req may be sent a second time over
// HTTP/1.1 after err ended its h2 attempt.
//
// A failure at the h2 layer does not say whether the upstream already
// carried out the request, so replaying one that changes state could
// duplicate it. Two cases are safe: a request whose repetition is
// harmless by definition, and an upstream that told us it processed
// nothing on the connection. The second is what makes the IIS case work
// for every method — a site requiring HTTP/1.1 refuses at stream 0,
// before the request is looked at.
func safeToRetry(req *http.Request, err error) bool {
return idempotent(req) || upstreamProcessedNothing(err)
}
// idempotent reports whether repeating req is defined to be harmless.
// It mirrors net/http's own retry rule (Request.isReplayable): a method
// with no side effects, or a caller that promised the upstream
// deduplicates by key.
func idempotent(req *http.Request) bool {
if req.Header.Get("Idempotency-Key") != "" || req.Header.Get("X-Idempotency-Key") != "" {
return true
}
switch req.Method {
// An empty method means GET, as in net/http.
case "", http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace:
return true
}
return false
}
// upstreamProcessedNothing reports whether err describes a GOAWAY that
// named this request's stream as one the upstream had not received, so
// it cannot have acted on it. A stream error says the opposite: the
// stream was open, so the request had been delivered.
func upstreamProcessedNothing(err error) bool {
if err == nil {
return false
}
msg := transportError(err).Error()
return strings.Contains(msg, goAwayStreamNotReceivedMarker) ||
strings.Contains(msg, goAwayNothingProcessedMarker)
}
// replayable returns a request that can be sent a second time, or
// ok=false when the body is gone. A RoundTripper consumes and closes
// the body it was given, so a retry needs either no body at all or
// GetBody to produce a fresh one.
func replayable(req *http.Request) (*http.Request, bool) {
if req.Body == nil || req.Body == http.NoBody {
return req, true
}
if req.GetBody == nil {
return nil, false
}
body, err := req.GetBody()
if err != nil {
return nil, false
}
retry := req.Clone(req.Context())
retry.Body = body
return retry, true
}
// http2ErrorMarkers are the substrings that identify an HTTP/2 protocol
// failure. net/http bundles its own private copy of the http2 package,
// so its errors cannot be matched by type from here: http2.StreamError
// and friends in x/net are different types from the ones a
// bundled-h2 transport returns. The strings below are the formats those
// bundled errors print, and they are specific to h2 framing — a
// downgrade must never be triggered by an ordinary network or TLS
// error, which retrying on HTTP/1.1 would not fix.
var http2ErrorMarkers = []string{
// Transport-level h2 failures, e.g.
// "http2: server sent GOAWAY and closed the connection".
"http2:",
// http2.StreamError, e.g. "stream error: stream ID 1; PROTOCOL_ERROR".
"stream error: stream ID",
// http2.ConnectionError, e.g. "connection error: PROTOCOL_ERROR".
"connection error: ",
// The code an upstream sends to say the request must be retried
// over HTTP/1.1, as a GOAWAY or on the stream.
http11RequiredMarker,
}
const (
// http11RequiredMarker is the error code an upstream sends to say the
// request belongs on HTTP/1.1. Unlike the other markers it is not a
// fault: the upstream is describing its own configuration.
http11RequiredMarker = "HTTP_1_1_REQUIRED"
// A GOAWAY carrying NO_ERROR closes a connection without complaint:
// a server draining before shutdown, recycling an application pool,
// capping requests per connection. The upstream speaks h2 perfectly
// well, so this must never pin it. The two spellings are the two
// formats the bundled transport prints the code in.
goAwayNoErrorEqualsMarker = "ErrCode=NO_ERROR"
goAwayNoErrorColonMarker = "ErrCode:NO_ERROR"
// gracefulGoAwayMarker is errClientConnGotGoAway, which the bundled
// transport raises for a stream the server never received on a
// connection it is shutting down gracefully. It normally retries
// those itself on a new connection and this never surfaces.
gracefulGoAwayMarker = "Transport received Server's graceful shutdown GOAWAY"
// goAwayStreamNotReceivedMarker is the bundled transport's abort for
// the first stream on a connection whose GOAWAY carried a real error
// code — the IIS case. It sits in the same "streamID > LastStreamID"
// branch as the graceful abort, so the server had not received the
// stream (see net/http's h2_bundle.go).
goAwayStreamNotReceivedMarker = "Transport received GOAWAY from server ErrCode:"
// goAwayNothingProcessedMarker is a GoAwayError naming stream 0 as
// the last one received, which says the same thing. The trailing
// comma keeps it from matching LastStreamID=10 and the rest.
goAwayNothingProcessedMarker = "LastStreamID=0,"
)
// isHTTP11Required reports whether the upstream itself asked for
// HTTP/1.1, rather than merely failing at h2.
func isHTTP11Required(err error) bool {
return err != nil && strings.Contains(transportError(err).Error(), http11RequiredMarker)
}
// isHTTP2ProtocolError reports whether err says the upstream cannot
// serve the h2 it negotiated.
func isHTTP2ProtocolError(err error) bool {
if err == nil {
return false
}
msg := transportError(err).Error()
// A graceful GOAWAY is routine connection management, not an
// upstream that cannot serve h2.
if isGracefulGoAway(msg) {
return false
}
for _, marker := range http2ErrorMarkers {
if strings.Contains(msg, marker) {
return true
}
}
return false
}
// isGracefulGoAway reports whether msg describes a GOAWAY sent to close
// a healthy connection rather than to report an inability to serve h2.
func isGracefulGoAway(msg string) bool {
return strings.Contains(msg, goAwayNoErrorEqualsMarker) ||
strings.Contains(msg, goAwayNoErrorColonMarker) ||
strings.Contains(msg, gracefulGoAwayMarker)
}
// transportError strips a *url.Error wrapper, which prefixes the request
// URL to the message. Markers are matched as substrings, so a URL left
// in place could classify an ordinary dial or TLS failure as an h2 one
// on the strength of the path alone.
func transportError(err error) error {
var urlErr *url.Error
if errors.As(err, &urlErr) && urlErr.Err != nil {
return urlErr.Err
}
return err
}
+587
View File
@@ -0,0 +1,587 @@
package roundtrip
import (
"bufio"
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"errors"
"fmt"
"io"
"math/big"
"net"
"net/http"
"net/url"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/net/http2"
)
// TestUpstreamTransport_AutoFallsBackOnBrokenHTTP2 covers the case ALPN
// cannot express: the upstream advertises h2, picks it, and then cannot
// serve it. The request must still succeed, over HTTP/1.1, and the
// upstream must stay on HTTP/1.1 for the requests that follow.
func TestUpstreamTransport_AutoFallsBackOnBrokenHTTP2(t *testing.T) {
t.Setenv(EnvUpstreamHTTPVersion, string(upstreamHTTPAuto))
srv := startBrokenHTTP2Server(t)
mt := NewDirectOnly(nil)
ctx := WithSkipTLSVerify(WithDirectUpstream(context.Background()))
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://"+srv.addr, nil)
require.NoError(t, err)
resp, err := mt.RoundTrip(req)
require.NoError(t, err, "a replayable request must be retried on HTTP/1.1 instead of failing")
body, err := io.ReadAll(resp.Body)
_ = resp.Body.Close()
require.NoError(t, err)
assert.Equal(t, "HTTP/1.1", resp.Proto, "the retry must ride the HTTP/1.1 transport")
assert.Equal(t, "http/1.1", string(body), "the upstream must see an http/1.1 ALPN offer on the retry")
assert.True(t, mt.insecure.isDowngraded(srv.addr),
"the upstream must stay pinned to HTTP/1.1 after proving it cannot serve h2")
mt.insecure.mu.RLock()
pin := mt.insecure.downgraded[srv.addr]
mt.insecure.mu.RUnlock()
assert.True(t, pin.permanent(),
"an upstream answering HTTP_1_1_REQUIRED must not be re-probed for h2")
// The second request must not repeat the h2 attempt: the server
// counts h2 handshakes, so a repeat would show up here.
h2Attempts := srv.http2Handshakes()
req, err = http.NewRequestWithContext(ctx, http.MethodGet, "https://"+srv.addr, nil)
require.NoError(t, err)
resp, err = mt.RoundTrip(req)
require.NoError(t, err)
_ = resp.Body.Close()
assert.Equal(t, "HTTP/1.1", resp.Proto, "a pinned upstream must go straight to HTTP/1.1")
assert.Equal(t, h2Attempts, srv.http2Handshakes(),
"a pinned upstream must not be probed for h2 again until the pin expires")
}
// TestUpstreamTransport_ExplicitHTTP2NeverDowngrades pins the promise
// that the explicit values are absolute: an operator who asked for h2
// keeps h2, broken upstream or not.
func TestUpstreamTransport_ExplicitHTTP2NeverDowngrades(t *testing.T) {
t.Setenv(EnvUpstreamHTTPVersion, string(upstreamHTTP2))
srv := startBrokenHTTP2Server(t)
mt := NewDirectOnly(nil)
ctx := WithSkipTLSVerify(WithDirectUpstream(context.Background()))
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://"+srv.addr, nil)
require.NoError(t, err)
resp, err := mt.RoundTrip(req)
if err == nil {
_ = resp.Body.Close()
}
require.Error(t, err, "NB_PROXY_UPSTREAM_HTTP_VERSION=2 must not fall back to HTTP/1.1")
assert.False(t, mt.insecure.isDowngraded(srv.addr), "an explicit version must never pin an upstream")
}
// TestUpstreamTransport_AutoDoesNotReplayUnsafeRequests covers the
// other half of the fallback: an h2 failure says nothing about whether
// the upstream already applied the request, so a state-changing one is
// not replayed. The host is still pinned, so the next request rides
// HTTP/1.1 without a second h2 attempt.
func TestUpstreamTransport_AutoDoesNotReplayUnsafeRequests(t *testing.T) {
t.Setenv(EnvUpstreamHTTPVersion, string(upstreamHTTPAuto))
srv := startBrokenHTTP2Server(t)
// A stream error means the stream was open, so the upstream had the
// request in hand — unlike the GOAWAY at stream 0 the fake server
// sends, which states it processed nothing.
streamErr := http2.StreamError{StreamID: 1, Code: http2.ErrCodeProtocol}
mt := NewDirectOnly(nil)
transport := mt.insecure
ctx := WithSkipTLSVerify(WithDirectUpstream(context.Background()))
assert.False(t, safeToRetry(newTestRequest(t, ctx, http.MethodPost, srv.addr), streamErr),
"a POST must not be replayed after a failure that may have been applied")
assert.True(t, safeToRetry(newTestRequest(t, ctx, http.MethodGet, srv.addr), streamErr),
"a GET is safe to replay whatever the failure was")
// The fake server's GOAWAY names stream 0, so even a POST is safe
// there and the request must succeed over HTTP/1.1.
resp, err := transport.RoundTrip(newTestRequest(t, ctx, http.MethodPost, srv.addr))
require.NoError(t, err, "a GOAWAY at stream 0 means the upstream applied nothing, so the POST may be replayed")
body, err := io.ReadAll(resp.Body)
_ = resp.Body.Close()
require.NoError(t, err)
assert.Equal(t, "http/1.1", string(body), "the retry must reach the upstream over http/1.1")
h2Attempts := srv.http2Handshakes()
resp, err = transport.RoundTrip(newTestRequest(t, ctx, http.MethodPost, srv.addr))
require.NoError(t, err)
_ = resp.Body.Close()
assert.Equal(t, h2Attempts, srv.http2Handshakes(),
"the pin must carry later requests without another h2 attempt")
}
func newTestRequest(t *testing.T, ctx context.Context, method, addr string) *http.Request {
t.Helper()
req, err := http.NewRequestWithContext(ctx, method, "https://"+addr, strings.NewReader("payload"))
require.NoError(t, err)
return req
}
func TestUpstreamKey(t *testing.T) {
tests := []struct {
name string
url string
want string
}{
{name: "host", url: "https://backend.invalid/path", want: "backend.invalid"},
// DNS is case-insensitive, so these are one upstream and must
// share one pin.
{name: "mixed case", url: "https://Backend.INVALID/path", want: "backend.invalid"},
// The default port is implied on every path that can downgrade.
{name: "explicit default port", url: "https://backend.invalid:443/", want: "backend.invalid"},
{name: "non-default port", url: "https://backend.invalid:8443/", want: "backend.invalid:8443"},
// An IPv6 literal needs its brackets back after Hostname strips
// them, or the key is not a dialable authority.
{name: "ipv6 default port", url: "https://[2001:db8::1]/", want: "2001:db8::1"},
{name: "ipv6 with port", url: "https://[2001:db8::1]:8443/", want: "[2001:db8::1]:8443"},
// One address in three spellings: hex case, a leading zero and an
// uncompressed zero run are all the same upstream.
{name: "ipv6 upper case", url: "https://[2001:DB8::1]/", want: "2001:db8::1"},
{name: "ipv6 leading zero", url: "https://[2001:0db8::1]/", want: "2001:db8::1"},
{name: "ipv6 uncompressed", url: "https://[2001:db8:0:0:0:0:0:1]/", want: "2001:db8::1"},
// A v4-mapped address is the v4 address, not a second upstream.
{name: "v4-mapped", url: "https://[::ffff:192.0.2.1]/", want: "192.0.2.1"},
// A zone names an interface, and interface names are
// case-sensitive, so these two are different links.
{name: "ipv6 zone", url: "https://[fe80::1%25eth0]/", want: "fe80::1%eth0"},
{name: "ipv6 zone upper case", url: "https://[fe80::1%25ETH0]/", want: "fe80::1%ETH0"},
// The address before the zone still normalizes.
{name: "ipv6 zone with upper-case address", url: "https://[FE80::1%25eth0]/", want: "fe80::1%eth0"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
parsed, err := url.Parse(tc.url)
require.NoError(t, err)
assert.Equal(t, tc.want, upstreamKey(parsed))
})
}
}
func TestSafeToRetry(t *testing.T) {
streamErr := http2.StreamError{StreamID: 1, Code: http2.ErrCodeProtocol}
goAwayAtZero := errors.New(`http2: server sent GOAWAY and closed the connection; LastStreamID=0, ErrCode=HTTP_1_1_REQUIRED, debug=""`)
goAwayLater := errors.New(`http2: server sent GOAWAY and closed the connection; LastStreamID=11, ErrCode=PROTOCOL_ERROR, debug=""`)
tests := []struct {
name string
method string
headers map[string]string
err error
want bool
}{
{name: "get", method: http.MethodGet, err: streamErr, want: true},
{name: "head", method: http.MethodHead, err: streamErr, want: true},
{name: "options", method: http.MethodOptions, err: streamErr, want: true},
{name: "trace", method: http.MethodTrace, err: streamErr, want: true},
{name: "post", method: http.MethodPost, err: streamErr, want: false},
{name: "put", method: http.MethodPut, err: streamErr, want: false},
{name: "patch", method: http.MethodPatch, err: streamErr, want: false},
{name: "delete", method: http.MethodDelete, err: streamErr, want: false},
// The upstream reported it handled nothing, so repeating the
// request cannot duplicate anything.
{name: "post with goaway at stream 0", method: http.MethodPost, err: goAwayAtZero, want: true},
// What the bundled transport actually raises for the first
// stream on a connection the upstream GOAWAYs with a real error
// code, which is the shape a real IIS site produces.
{
name: "post with first-stream goaway abort",
method: http.MethodPost,
err: errors.New("http2: Transport received GOAWAY from server ErrCode:HTTP_1_1_REQUIRED"),
want: true,
},
// It handled earlier streams, so this one may have been applied.
{name: "post with goaway after other streams", method: http.MethodPost, err: goAwayLater, want: false},
{
name: "post with idempotency key",
method: http.MethodPost,
headers: map[string]string{"Idempotency-Key": "abc"},
err: streamErr,
want: true,
},
{
name: "post with prefixed idempotency key",
method: http.MethodPost,
headers: map[string]string{"X-Idempotency-Key": "abc"},
err: streamErr,
want: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
req, err := http.NewRequest(tc.method, "https://backend.invalid", nil)
require.NoError(t, err)
for k, v := range tc.headers {
req.Header.Set(k, v)
}
assert.Equal(t, tc.want, safeToRetry(req, tc.err))
})
}
}
func TestUpstreamTransport_MayDowngrade(t *testing.T) {
tests := []struct {
name string
version upstreamHTTPVersion
url string
want bool
}{
{name: "auto over TLS", version: upstreamHTTPAuto, url: "https://backend.invalid", want: true},
// The proxy speaks no h2c, so a cleartext upstream is already on
// HTTP/1.1 and its failures say nothing about h2.
{name: "auto cleartext", version: upstreamHTTPAuto, url: "http://backend.invalid", want: false},
{name: "explicit 1.1", version: upstreamHTTP11, url: "https://backend.invalid", want: false},
{name: "explicit 2", version: upstreamHTTP2, url: "https://backend.invalid", want: false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
transport := newUpstreamTransport(&http.Transport{}, tc.version, nil)
req, err := http.NewRequest(http.MethodGet, tc.url, nil)
require.NoError(t, err)
assert.Equal(t, tc.want, transport.mayDowngrade(req))
})
}
}
func TestUpstreamTransport_DowngradeExpires(t *testing.T) {
transport := newUpstreamTransport(&http.Transport{}, upstreamHTTPAuto, nil)
transport.markDowngraded("backend.invalid:443", false)
require.True(t, transport.isDowngraded("backend.invalid:443"))
transport.mu.Lock()
transport.downgraded["backend.invalid:443"] = downgrade{expiry: time.Now().Add(-time.Second)}
transport.mu.Unlock()
assert.False(t, transport.isDowngraded("backend.invalid:443"),
"an expired pin must let the upstream be offered h2 again")
transport.mu.RLock()
_, stillTracked := transport.downgraded["backend.invalid:443"]
transport.mu.RUnlock()
assert.False(t, stillTracked, "an expired pin must not be kept around")
}
func TestUpstreamTransport_DowngradeIsPerUpstream(t *testing.T) {
transport := newUpstreamTransport(&http.Transport{}, upstreamHTTPAuto, nil)
transport.markDowngraded("broken.invalid:443", false)
assert.True(t, transport.isDowngraded("broken.invalid:443"))
assert.False(t, transport.isDowngraded("healthy.invalid:443"),
"one broken upstream must not drop the others to HTTP/1.1")
}
// TestUpstreamTransport_HTTP11RequiredPinIsPermanent covers the IIS
// case: HTTP_1_1_REQUIRED describes how the upstream is configured
// (Windows Authentication, client certificates), so re-probing it every
// upstreamDowngradeTTL would only buy a failed request per interval.
func TestUpstreamTransport_HTTP11RequiredPinIsPermanent(t *testing.T) {
transport := newUpstreamTransport(&http.Transport{}, upstreamHTTPAuto, nil)
transport.markDowngraded("iis.invalid:443", true)
transport.mu.RLock()
pin := transport.downgraded["iis.invalid:443"]
transport.mu.RUnlock()
assert.True(t, pin.permanent(), "an upstream that asked for HTTP/1.1 must not be re-probed")
assert.True(t, pin.active(time.Now().Add(100*upstreamDowngradeTTL)),
"a permanent pin must outlive any TTL")
}
func TestUpstreamTransport_PermanentPinSurvivesLaterFailures(t *testing.T) {
transport := newUpstreamTransport(&http.Transport{}, upstreamHTTPAuto, nil)
transport.markDowngraded("iis.invalid:443", true)
// A later ambiguous failure for the same upstream must not turn the
// permanent pin into an expiring one.
transport.markDowngraded("iis.invalid:443", false)
transport.mu.RLock()
pin := transport.downgraded["iis.invalid:443"]
transport.mu.RUnlock()
assert.True(t, pin.permanent(), "a permanent pin must never be weakened")
}
func TestIsHTTP11Required(t *testing.T) {
tests := []struct {
name string
err error
want bool
}{
{
name: "goaway",
err: errors.New(`http2: server sent GOAWAY and closed the connection; LastStreamID=0, ErrCode=HTTP_1_1_REQUIRED, debug=""`),
want: true,
},
{
name: "stream error",
err: http2.StreamError{StreamID: 1, Code: http2.ErrCodeHTTP11Required},
want: true,
},
{
name: "other h2 failure",
err: http2.StreamError{StreamID: 1, Code: http2.ErrCodeProtocol},
want: false,
},
{name: "nil", err: nil, want: false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, isHTTP11Required(tc.err))
})
}
}
func TestIsHTTP2ProtocolError(t *testing.T) {
tests := []struct {
name string
err error
want bool
}{
{
name: "goaway demanding http/1.1",
err: errors.New(`http2: server sent GOAWAY and closed the connection; LastStreamID=0, ErrCode=HTTP_1_1_REQUIRED, debug=""`),
want: true,
},
{
name: "stream error",
err: http2.StreamError{StreamID: 1, Code: http2.ErrCodeProtocol},
want: true,
},
{
name: "connection error",
err: http2.ConnectionError(http2.ErrCodeProtocol),
want: true,
},
{
name: "wrapped h2 error",
err: errors.New("Get \"https://backend.invalid\": http2: client connection lost"),
want: true,
},
// A GOAWAY with NO_ERROR is a server draining a connection —
// recycling an application pool, capping requests per
// connection, shutting down gracefully. It speaks h2 fine.
{
name: "graceful goaway",
err: errors.New(`http2: server sent GOAWAY and closed the connection; LastStreamID=9, ErrCode=NO_ERROR, debug=""`),
want: false,
},
{
name: "graceful shutdown abort",
err: errors.New("http2: Transport received Server's graceful shutdown GOAWAY"),
want: false,
},
// Markers are substrings, so a URL carried by a *url.Error must
// not be able to classify a plain failure as an h2 one.
{
name: "url error whose path looks like a marker",
err: &url.Error{
Op: "Get",
URL: "https://backend.invalid/http2:/connection error: x",
Err: errors.New("dial tcp 10.0.0.1:443: connect: connection refused"),
},
want: false,
},
// Retrying these on HTTP/1.1 fixes nothing, so they must never
// pin an upstream.
{name: "dial failure", err: errors.New("dial tcp 10.0.0.1:443: connect: connection refused"), want: false},
{name: "tls failure", err: errors.New("tls: failed to verify certificate: x509: certificate signed by unknown authority"), want: false},
{name: "context cancelled", err: context.Canceled, want: false},
{name: "eof", err: io.EOF, want: false},
{name: "nil", err: nil, want: false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, isHTTP2ProtocolError(tc.err))
})
}
}
func TestReplayable(t *testing.T) {
t.Run("bodyless request", func(t *testing.T) {
req, err := http.NewRequest(http.MethodGet, "https://backend.invalid", nil)
require.NoError(t, err)
retry, ok := replayable(req)
require.True(t, ok)
assert.Same(t, req, retry, "a bodyless request needs no clone")
})
t.Run("request with GetBody", func(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "https://backend.invalid", strings.NewReader("payload"))
require.NoError(t, err)
// Consume the body the way a failed RoundTrip would.
_, err = io.ReadAll(req.Body)
require.NoError(t, err)
retry, ok := replayable(req)
require.True(t, ok)
body, err := io.ReadAll(retry.Body)
require.NoError(t, err)
assert.Equal(t, "payload", string(body), "the retry must carry a fresh copy of the body")
})
t.Run("streamed request", func(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "https://backend.invalid", io.NopCloser(strings.NewReader("payload")))
require.NoError(t, err)
require.Nil(t, req.GetBody, "an opaque reader must not get a GetBody")
_, ok := replayable(req)
assert.False(t, ok, "a body that cannot be regenerated must not be replayed")
})
}
// brokenHTTP2Server advertises h2 in ALPN, accepts it, and then refuses
// to serve it — the upstream behaviour that motivated the fallback. Over
// http/1.1 it answers normally, so a downgraded request succeeds.
type brokenHTTP2Server struct {
addr string
handshakes chan struct{}
}
func (s *brokenHTTP2Server) http2Handshakes() int {
return len(s.handshakes)
}
func startBrokenHTTP2Server(t *testing.T) *brokenHTTP2Server {
t.Helper()
ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
Certificates: []tls.Certificate{selfSignedCert(t)},
NextProtos: []string{"h2", "http/1.1"},
MinVersion: tls.VersionTLS12,
})
require.NoError(t, err)
t.Cleanup(func() { _ = ln.Close() })
srv := &brokenHTTP2Server{
addr: ln.Addr().String(),
// Buffered well past what the test drives so a stuck server
// never blocks the accept loop.
handshakes: make(chan struct{}, 64),
}
go func() {
for {
conn, err := ln.Accept()
if err != nil {
return
}
go srv.handle(conn)
}
}()
return srv
}
func (s *brokenHTTP2Server) handle(conn net.Conn) {
defer func() { _ = conn.Close() }()
tlsConn, ok := conn.(*tls.Conn)
if !ok {
return
}
if err := tlsConn.Handshake(); err != nil {
return
}
proto := tlsConn.ConnectionState().NegotiatedProtocol
if proto == "h2" {
select {
case s.handshakes <- struct{}{}:
default:
}
s.refuseHTTP2(tlsConn)
return
}
s.serveHTTP1(tlsConn, proto)
}
// refuseHTTP2 completes just enough of the h2 handshake for the client
// to accept the connection, then sends the GOAWAY an upstream uses to
// say the request belongs on HTTP/1.1.
//
// The client is still writing its preface and request while the GOAWAY
// goes out, so the connection is drained before the caller closes it.
// Closing a socket with unread bytes still in its receive buffer makes
// the kernel answer with RST, which reaches the client as a write error
// rather than the GOAWAY — no h2 error, so no downgrade, and the test
// fails on the error the client saw first.
func (s *brokenHTTP2Server) refuseHTTP2(conn net.Conn) {
framer := http2.NewFramer(conn, conn)
if err := framer.WriteSettings(); err != nil {
return
}
if err := framer.WriteGoAway(0, http2.ErrCodeHTTP11Required, nil); err != nil {
return
}
// The client closes its side once it has read the GOAWAY, which ends
// the drain; the deadline is only a backstop against a client that
// never does.
_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
_, _ = io.Copy(io.Discard, conn)
}
// serveHTTP1 answers a single request with the ALPN protocol the
// upstream actually settled on, so a test asserting on the body is
// checking what the upstream saw rather than a constant.
func (s *brokenHTTP2Server) serveHTTP1(conn net.Conn, alpn string) {
reader := bufio.NewReader(conn)
if _, err := http.ReadRequest(reader); err != nil {
return
}
_, _ = fmt.Fprintf(conn,
"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s",
len(alpn), alpn)
}
func selfSignedCert(t *testing.T) tls.Certificate {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
template := x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "127.0.0.1"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
IPAddresses: []net.IP{net.ParseIP("127.0.0.1")},
IsCA: true,
}
der, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
require.NoError(t, err)
return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}
}
+1
View File
@@ -571,6 +571,7 @@ func TestIntegration_ProxyConnection_ReconnectDoesNotDuplicateState(t *testing.T
proxytypes.ServiceID(mapping.GetId()),
nil,
mapping.GetPrivate(),
mapping.GetAuth().GetAllowedGroupIds(),
)
require.NoError(t, err)
+49 -10
View File
@@ -20,6 +20,7 @@ import (
"net/url"
"path/filepath"
"reflect"
"slices"
"sync"
"time"
@@ -2063,24 +2064,56 @@ 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())
maxSessionAge := time.Duration(mapping.GetAuth().GetMaxSessionAgeSeconds()) * time.Second
if err := s.auth.AddDomain(mapping.GetDomain(), schemes, mapping.GetAuth().GetSessionKey(), maxSessionAge, accountID, svcID, ipRestrictions, mapping.GetPrivate()); err != nil {
if err := s.auth.AddDomain(mapping.GetDomain(), schemes, mapping.GetAuth().GetSessionKey(), maxSessionAge, accountID, svcID, ipRestrictions, mapping.GetPrivate(), mapping.GetAuth().GetAllowedGroupIds()); err != nil {
return fmt.Errorf("auth setup for domain %s: %w", mapping.GetDomain(), err)
}
m := s.protoToMapping(ctx, mapping)
s.proxy.AddMapping(m)
// 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.rebuildMiddlewareChains(svcID, m)
s.proxy.AddMapping(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. No entry is
// ever dropped: a name that arrives blank, or without a hash, still yields a
// scheme, because a mapping that lost its only scheme would fall through
// Protect's no-schemes pass-through and serve the domain unauthenticated.
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 !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
@@ -2115,15 +2148,21 @@ 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.
func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) {
// 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 {
if s.middlewareManager == nil {
return
return nil
}
bindings := buildMiddlewareBindings(svcID, m)
if err := s.middlewareManager.Rebuild(string(svcID), bindings); err != nil {
s.Logger.WithError(err).WithField("service_id", svcID).Error("failed to rebuild middleware chains")
return fmt.Errorf("rebuild middleware chains for service %s: %w", svcID, err)
}
return nil
}
// isLiveService reports whether svcID is currently present in the live
+60
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,62 @@ 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")
}
// TestHeaderAuthSchemes_BlankNameFailsClosed covers a mapping row whose header
// name is empty. Skipping it would leave a service whose only auth is that entry
// with no schemes at all, which Protect treats as an unprotected domain, so the
// entry is kept and the domain stays gated.
func TestHeaderAuthSchemes_BlankNameFailsClosed(t *testing.T) {
schemes := headerAuthSchemes([]*proto.HeaderAuth{{Header: "", HashedValue: "$argon2id$not-a-real-hash"}})
require.Len(t, schemes, 1, "a blank header name must still register a scheme")
assert.False(t, headerSchemeAccepts(t, schemes[0], "X-Api-Key", "anything"),
"a blank header auth must not admit any request")
}
type statusUpdateOnlyClient struct {
proto.ProxyServiceClient
}
+6 -6
View File
File diff suppressed because one or more lines are too long
+6
View File
@@ -68,6 +68,12 @@ function App() {
if (res.type === "opaqueredirect" || res.status === 0) {
setSubmitting("redirect");
globalThis.location.reload();
} else if (res.status === 429) {
const seconds = Number(res.headers.get("Retry-After"));
const wait = Number.isFinite(seconds) && seconds > 0
? ` Try again in ${Math.ceil(seconds)} seconds.`
: " Please try again later.";
handleAuthError(method, `Too many authentication attempts.${wait}`);
} else {
handleAuthError(method, "Authentication failed. Please try again.");
}