mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-23 15:19:08 +02:00
Merge remote-tracking branch 'origin/main' into HEAD
# Conflicts: # shared/management/proto/proxy_service.pb.go
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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{}{}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,14 @@ func NormalizeBedrockModel(modelID string) string {
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeAnthropicModel strips the trailing "-YYYYMMDD" release-date suffix
|
||||
// from an Anthropic model id so a dated id a client pins matches the undated
|
||||
// one the operator registered. Thin delegate to shared/llm for the same
|
||||
// contract reason as the two below.
|
||||
func NormalizeAnthropicModel(modelID string) string {
|
||||
return sharedllm.NormalizeAnthropicModel(modelID)
|
||||
}
|
||||
|
||||
// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id
|
||||
// so it matches the catalog/pricing key. Thin delegate to shared/llm, kept
|
||||
// beside NormalizeBedrockModel for the same contract reason.
|
||||
|
||||
@@ -10,6 +10,8 @@ package pricing
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// Entry is a single model's input and output pricing, expressed in USD per
|
||||
@@ -92,7 +94,10 @@ func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) {
|
||||
return &Table{entries: entries}, nil
|
||||
}
|
||||
|
||||
// Lookup returns the entry for the given provider surface and model.
|
||||
// Lookup returns the entry for the given provider surface and model. A
|
||||
// dated Anthropic id falls back to its undated form, so a client pinning
|
||||
// "claude-sonnet-4-5-20250929" bills at the registered "claude-sonnet-4-5"
|
||||
// rate instead of recording no cost at all.
|
||||
func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if t == nil {
|
||||
return Entry{}, false
|
||||
@@ -101,7 +106,14 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) {
|
||||
if !ok {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[model]
|
||||
if e, found := byModel[model]; found {
|
||||
return e, true
|
||||
}
|
||||
undated := sharedllm.NormalizeAnthropicModel(model)
|
||||
if undated == model {
|
||||
return Entry{}, false
|
||||
}
|
||||
e, ok := byModel[undated]
|
||||
return e, ok
|
||||
}
|
||||
|
||||
|
||||
@@ -175,3 +175,22 @@ func TestNewTable_NilAndEmpty(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map")
|
||||
}
|
||||
|
||||
// TestLookup_DatedAnthropicIDFallsBackToUndated covers a client pinning a
|
||||
// release date on a model priced under its undated id. Without the
|
||||
// fallback the request records no cost at all.
|
||||
func TestLookup_DatedAnthropicIDFallsBackToUndated(t *testing.T) {
|
||||
table, err := NewTable(map[string]map[string]EntryJSON{
|
||||
"anthropic": {
|
||||
"claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err, "table must build from a valid defaults map")
|
||||
|
||||
entry, ok := table.Lookup("anthropic", "claude-sonnet-4-5-20250929")
|
||||
require.True(t, ok, "a dated id must resolve to the undated entry")
|
||||
assert.InDelta(t, 0.003, entry.InputPer1K, 1e-9, "dated id must bill at the registered rate")
|
||||
|
||||
_, ok = table.Lookup("anthropic", "claude-sonnet-9-9-20250929")
|
||||
assert.False(t, ok, "an unknown family must stay unpriced")
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"])
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,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,
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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}
|
||||
}
|
||||
Reference in New Issue
Block a user