mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-28 11:21:29 +02:00
Compare commits
3 Commits
main
...
worktree-d
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d918235afd | ||
|
|
f7af852851 | ||
|
|
4546908ceb |
@@ -273,8 +273,8 @@ dockers_v2:
|
||||
- netbirdio/netbird
|
||||
- ghcr.io/netbirdio/netbird
|
||||
tags:
|
||||
- "{{ .Version }}-rootless"
|
||||
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}rootless-latest{{ end }}"
|
||||
- "v{{ .Version }}-rootless"
|
||||
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}latest{{ end }}"
|
||||
dockerfile: client/Dockerfile-rootless
|
||||
extra_files:
|
||||
- client/netbird-entrypoint.sh
|
||||
|
||||
@@ -3,6 +3,7 @@ package dns
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math"
|
||||
"net"
|
||||
"slices"
|
||||
@@ -31,6 +32,26 @@ type SubdomainMatcher interface {
|
||||
MatchSubdomains() bool
|
||||
}
|
||||
|
||||
// responseMeta holds the annotations handlers attach to a request to explain the
|
||||
// response the chain ends up writing. It survives a deferral, so an answer that
|
||||
// did not come from the handler that owns the name still says who stepped aside
|
||||
// and why.
|
||||
type responseMeta map[resutil.MetaKey]string
|
||||
|
||||
// format renders the annotations for the response log line. The order is stable
|
||||
// so the same event reads the same way every time; a map's own order is not.
|
||||
func (m responseMeta) format() string {
|
||||
if len(m) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
for _, k := range slices.Sorted(maps.Keys(m)) {
|
||||
b.WriteString(" " + string(k) + "=" + m[k])
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
type HandlerEntry struct {
|
||||
Handler dns.Handler
|
||||
Priority int
|
||||
@@ -52,8 +73,19 @@ type ResponseWriterChain struct {
|
||||
origPattern string
|
||||
requestID string
|
||||
shouldContinue bool
|
||||
response *dns.Msg
|
||||
meta map[string]string
|
||||
// softNegative suppresses a poisoning negative verdict for this request. A
|
||||
// handler that owns the name but cannot answer this query type sets it
|
||||
// before deferring, and it stays set for every handler that runs after.
|
||||
softNegative bool
|
||||
// clientHasEdns records whether the original query advertised EDNS0, taken
|
||||
// before any handler ran: handlers add EDNS0 to the query they forward
|
||||
// upstream, so the message itself no longer answers the question later.
|
||||
clientHasEdns bool
|
||||
response *dns.Msg
|
||||
// meta is handed to the next handler when this one defers, so the same map
|
||||
// outlives the writer that created it. Handlers must only set metadata from
|
||||
// within ServeDNS, never from a goroutine that outlives the call.
|
||||
meta responseMeta
|
||||
}
|
||||
|
||||
// RequestID returns the request ID for tracing
|
||||
@@ -62,26 +94,64 @@ func (w *ResponseWriterChain) RequestID() string {
|
||||
}
|
||||
|
||||
// SetMeta sets a metadata key-value pair for logging
|
||||
func (w *ResponseWriterChain) SetMeta(key, value string) {
|
||||
func (w *ResponseWriterChain) SetMeta(key resutil.MetaKey, value string) {
|
||||
if w.meta == nil {
|
||||
w.meta = make(map[string]string)
|
||||
w.meta = make(responseMeta)
|
||||
}
|
||||
w.meta[key] = value
|
||||
}
|
||||
|
||||
// RequestSoftNegative marks the request so a downstream NXDOMAIN is turned into
|
||||
// NODATA before it reaches the client. Set by a handler that defers a query for
|
||||
// a name it owns but a type it cannot resolve.
|
||||
func (w *ResponseWriterChain) RequestSoftNegative() {
|
||||
w.softNegative = true
|
||||
}
|
||||
|
||||
func (w *ResponseWriterChain) WriteMsg(m *dns.Msg) error {
|
||||
// Check if this is a continue signal (NXDOMAIN with Zero bit set)
|
||||
if m.Rcode == dns.RcodeNameError && m.MsgHdr.Zero {
|
||||
w.shouldContinue = true
|
||||
return nil
|
||||
}
|
||||
if w.softNegative && m.Rcode == dns.RcodeNameError {
|
||||
m = softenNegative(m, w.clientHasEdns)
|
||||
w.SetMeta(resutil.MetaKeySoftened, "nxdomain->nodata")
|
||||
}
|
||||
w.response = m
|
||||
if m.MsgHdr.Truncated {
|
||||
w.SetMeta("truncated", "true")
|
||||
w.SetMeta(resutil.MetaKeyTruncated, "true")
|
||||
}
|
||||
return w.ResponseWriter.WriteMsg(m)
|
||||
}
|
||||
|
||||
// softenNegative downgrades an NXDOMAIN to NODATA for a request a handler that
|
||||
// owns the name deferred. NXDOMAIN is cached for the name and every type below
|
||||
// it, so a resolver that has never heard of a routed name would take out the
|
||||
// record types the route does serve; NODATA is cached for this name and type
|
||||
// only. The authority section goes with it: the negative TTL of a zone we just
|
||||
// overruled does not apply, and RFC 2308 keeps a negative answer that carries no
|
||||
// SOA out of caches altogether, so the rewrite cannot outlive the route.
|
||||
//
|
||||
// The rewrite is ours, not the answering resolver's, so an EDNS0 client is told
|
||||
// as much: the reply we hand back travels a path no capture on this host sees,
|
||||
// and an empty answer is otherwise indistinguishable from a real one. Returns a
|
||||
// copy so the answering handler keeps whatever it may hold on to.
|
||||
func softenNegative(m *dns.Msg, clientHasEdns bool) *dns.Msg {
|
||||
out := m.Copy()
|
||||
out.Rcode = dns.RcodeSuccess
|
||||
out.Ns = nil
|
||||
|
||||
if clientHasEdns {
|
||||
resutil.AttachEDE(out, resutil.EDENetbirdSoftenedNegative,
|
||||
"netbird: name is served locally, NXDOMAIN from the fallthrough resolver suppressed")
|
||||
} else {
|
||||
resutil.StripOPT(out)
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
func NewHandlerChain() *HandlerChain {
|
||||
return &HandlerChain{
|
||||
handlers: make([]HandlerEntry, 0),
|
||||
@@ -223,6 +293,19 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
handlers := slices.Clone(c.handlers)
|
||||
c.mu.RUnlock()
|
||||
|
||||
// Carried across deferrals: once a handler that owns the name defers, no
|
||||
// handler after it may answer with a verdict that poisons the name. The
|
||||
// metadata of a handler that stepped aside is carried too, so the response
|
||||
// log line explains an answer that did not come from the handler that owns
|
||||
// the name.
|
||||
var softNegative bool
|
||||
var carried responseMeta
|
||||
|
||||
// Taken before any handler runs: handlers advertise EDNS0 on the query they
|
||||
// forward upstream, so afterwards the message no longer tells us whether the
|
||||
// client did.
|
||||
clientHasEdns := r.IsEdns0() != nil
|
||||
|
||||
// Try handlers in priority order
|
||||
for _, entry := range handlers {
|
||||
if entry.Priority > maxPriority {
|
||||
@@ -245,11 +328,16 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
ResponseWriter: w,
|
||||
origPattern: entry.OrigPattern,
|
||||
requestID: requestID,
|
||||
softNegative: softNegative,
|
||||
clientHasEdns: clientHasEdns,
|
||||
meta: carried,
|
||||
}
|
||||
entry.Handler.ServeDNS(chainWriter, r)
|
||||
|
||||
// If handler wants to continue, try next handler
|
||||
if chainWriter.shouldContinue {
|
||||
softNegative = softNegative || chainWriter.softNegative
|
||||
carried = chainWriter.meta
|
||||
if entry.Priority != PriorityMgmtCache {
|
||||
logger.Tracef("handler requested continue for domain=%s", qname)
|
||||
}
|
||||
@@ -265,30 +353,59 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
qname, dns.TypeToString[question.Qtype], dns.ClassToString[question.Qclass])
|
||||
resp := &dns.Msg{}
|
||||
resp.SetRcode(r, dns.RcodeRefused)
|
||||
// A handler that owns the name deferred and nothing below it could answer
|
||||
// (a client with no primary nameserver group). The name exists as far as
|
||||
// this client is concerned, since the route serves its addresses, so REFUSED
|
||||
// would contradict the route: a stub that takes it as "not served here" and
|
||||
// retries another resolver can come back with an NXDOMAIN that takes the
|
||||
// whole name down. Answer "no records of this type" instead, without an SOA,
|
||||
// so RFC 2308 keeps it out of negative caches, and tell an EDNS0 client the
|
||||
// empty answer is ours rather than a resolver's.
|
||||
if softNegative {
|
||||
resp.Rcode = dns.RcodeSuccess
|
||||
if clientHasEdns {
|
||||
resutil.AttachEDE(resp, resutil.EDENetbirdSoftenedNegative,
|
||||
"netbird: name is served locally, no fallthrough resolver for this query type")
|
||||
}
|
||||
// logResponse never runs on this path, so the carried metadata is
|
||||
// appended here or the reason for the deferral is lost in exactly the
|
||||
// case that is hardest to diagnose.
|
||||
logger.Tracef("no handler below the deferring one for domain=%s type=%s, answering NODATA%s",
|
||||
qname, dns.TypeToString[question.Qtype], carried.format())
|
||||
}
|
||||
if err := w.WriteMsg(resp); err != nil {
|
||||
logger.Errorf("failed to write DNS response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *HandlerChain) logResponse(logger *log.Entry, cw *ResponseWriterChain, qname string, startTime time.Time) {
|
||||
// Runs for every query, and the arguments below are not free: Len() packs
|
||||
// the message to measure it, and formatting the answers and the metadata
|
||||
// allocates. None of it is worth doing when the line is discarded.
|
||||
if !log.IsLevelEnabled(log.TraceLevel) {
|
||||
return
|
||||
}
|
||||
|
||||
if cw.response == nil {
|
||||
return
|
||||
}
|
||||
|
||||
var meta string
|
||||
for k, v := range cw.meta {
|
||||
meta += " " + k + "=" + v
|
||||
}
|
||||
|
||||
logger.Tracef("response: domain=%s rcode=%s answers=%s size=%dB%s took=%s",
|
||||
qname, dns.RcodeToString[cw.response.Rcode], resutil.FormatAnswers(cw.response.Answer),
|
||||
cw.response.Len(), meta, time.Since(startTime))
|
||||
cw.response.Len(), cw.meta.format(), time.Since(startTime))
|
||||
}
|
||||
|
||||
// ResolveInternal runs an in-process DNS query against the chain, skipping any
|
||||
// handler with priority > maxPriority. Used by internal callers (e.g. the mgmt
|
||||
// cache refresher) that must bypass themselves to avoid loops. Honors ctx
|
||||
// cancellation; on ctx.Done the dispatch goroutine is left to drain on its own
|
||||
// cache refresher) that must bypass themselves to avoid loops.
|
||||
//
|
||||
// "Nothing answered" is read off RcodeRefused, which a request soft-negatived by
|
||||
// a deferring handler never carries: it ends in an empty NOERROR instead, and
|
||||
// would look resolved. No caller can reach that today, since every handler that
|
||||
// defers sits above the maxPriority any caller passes. Lowering one below it
|
||||
// means this check needs the soft-negative case too.
|
||||
//
|
||||
// Honors ctx cancellation; on ctx.Done the dispatch goroutine is left to drain on its own
|
||||
// (bounded by the invoked handler's internal timeout).
|
||||
func (c *HandlerChain) ResolveInternal(ctx context.Context, r *dns.Msg, maxPriority int) (*dns.Msg, error) {
|
||||
if len(r.Question) == 0 {
|
||||
|
||||
@@ -3,15 +3,19 @@ package dns_test
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
log "github.com/sirupsen/logrus"
|
||||
logtest "github.com/sirupsen/logrus/hooks/test"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/client/internal/dns"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/test"
|
||||
)
|
||||
|
||||
@@ -1238,6 +1242,276 @@ func TestHandlerChain_ResolveInternal_HonorsContextTimeout(t *testing.T) {
|
||||
assert.Less(t, elapsed, 500*time.Millisecond, "ResolveInternal must return shortly after ctx deadline")
|
||||
}
|
||||
|
||||
// requestSoftNegative asks the chain to soften a negative verdict produced by
|
||||
// the handlers that run after the caller defers, reporting whether the chain
|
||||
// supports the signal. Written as a type assertion so these tests compile
|
||||
// against a chain that does not support it yet.
|
||||
func requestSoftNegative(w dns.ResponseWriter) bool {
|
||||
sn, ok := w.(interface{ RequestSoftNegative() })
|
||||
if ok {
|
||||
sn.RequestSoftNegative()
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
// deferringHandler defers to the next handler in the chain, optionally asking
|
||||
// for the negative verdict of whatever answers instead to be softened. This is
|
||||
// what a DNS route handler does for a record type its routing peer cannot
|
||||
// resolve.
|
||||
type deferringHandler struct {
|
||||
softNegative bool
|
||||
called bool
|
||||
supported bool
|
||||
}
|
||||
|
||||
func (h *deferringHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
h.called = true
|
||||
if h.softNegative {
|
||||
h.supported = requestSoftNegative(w)
|
||||
// The real handler records why it stepped aside, for the response log
|
||||
// line of whichever handler answers instead.
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredBy, "test handler")
|
||||
}
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
resp.MsgHdr.Zero = true
|
||||
_ = w.WriteMsg(resp)
|
||||
}
|
||||
|
||||
// nxdomainHandler answers an authoritative NXDOMAIN with an SOA in the
|
||||
// authority section, the way a public resolver answers for a name that only
|
||||
// exists inside the routed network.
|
||||
type nxdomainHandler struct{}
|
||||
|
||||
func (h *nxdomainHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
resp.Ns = []dns.RR{&dns.SOA{
|
||||
Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: 3600},
|
||||
Ns: "ns1.example.com.",
|
||||
Mbox: "hostmaster.example.com.",
|
||||
Minttl: 3600,
|
||||
Expire: 604800,
|
||||
Refresh: 7200,
|
||||
Retry: 3600,
|
||||
}}
|
||||
_ = w.WriteMsg(resp)
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN is the whole point
|
||||
// of the soft-negative signal: a route handler may only defer a query to the
|
||||
// public chain if the answer cannot poison the routed name. NXDOMAIN is cached
|
||||
// for the name and every type under it (RFC 2308, RFC 8020), so it has to be
|
||||
// rewritten to NODATA, which is cached per name and type only.
|
||||
func TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.True(t, route.called, "the route handler must run first")
|
||||
require.True(t, route.supported, "the chain writer must accept a soft-negative request")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "NXDOMAIN must be softened to NODATA")
|
||||
assert.Empty(t, resp.Answer, "a softened negative carries no answer")
|
||||
assert.Empty(t, resp.Ns, "the downstream zone's SOA must not set the negative TTL for a name we overrode")
|
||||
}
|
||||
|
||||
// responseLineFor returns the chain's response log line for one query. Raising
|
||||
// the level to trace also unmutes whatever else is logging in this package, so
|
||||
// the line has to be picked by the name it was asked about rather than by being
|
||||
// the last one seen.
|
||||
func responseLineFor(hook *logtest.Hook, qname string) string {
|
||||
for _, e := range hook.AllEntries() {
|
||||
if strings.HasPrefix(e.Message, "response:") && strings.Contains(e.Message, qname) {
|
||||
return e.Message
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_IsVisibleToTheClient covers observability of the
|
||||
// rewrite. The reply we hand the application travels over loopback, which the
|
||||
// bundle capture does not see, so a softened verdict has to say so on the wire:
|
||||
// without it an empty answer is indistinguishable from a real "no such record"
|
||||
// in a dig output or a capture taken next to the application.
|
||||
func TestHandlerChain_SoftNegative_IsVisibleToTheClient(t *testing.T) {
|
||||
// One hook for the whole test: logtest installs it on the standard logger
|
||||
// and logrus has no way to take it off again, so a hook per subtest would
|
||||
// leave several behind buffering every later line in the package.
|
||||
hook := logtest.NewGlobal()
|
||||
t.Cleanup(hook.Reset)
|
||||
|
||||
newChain := func() *nbdns.HandlerChain {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
chain.AddHandler("*.example.com.", &deferringHandler{softNegative: true}, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
return chain
|
||||
}
|
||||
|
||||
t.Run("EDNS0 client gets an extended error", func(t *testing.T) {
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
r.SetEdns0(dns.DefaultMsgSize, false)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp)
|
||||
require.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
|
||||
ede, ok := resutil.ExtractEDE(resp)
|
||||
require.True(t, ok, "a softened verdict must carry an extended DNS error")
|
||||
assert.Equal(t, resutil.EDENetbirdSoftenedNegative, ede.InfoCode)
|
||||
assert.Contains(t, ede.ExtraText, "netbird", "the text must name us as the source of the rewrite")
|
||||
})
|
||||
|
||||
t.Run("plain client gets no OPT", func(t *testing.T) {
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp)
|
||||
assert.Nil(t, resp.IsEdns0(), "RFC 6891 forbids an OPT toward a client that did not advertise EDNS0")
|
||||
})
|
||||
|
||||
// The response line is what support reads out of a debug bundle, and it now
|
||||
// carries several annotations at once. Built from a map, their order would
|
||||
// differ on every query, so the same event never looks the same twice.
|
||||
t.Run("log fields keep a stable order", func(t *testing.T) {
|
||||
const qname = "_mongodb._tcp.stable.example.com."
|
||||
|
||||
prev := log.GetLevel()
|
||||
log.SetLevel(log.TraceLevel)
|
||||
t.Cleanup(func() { log.SetLevel(prev) })
|
||||
|
||||
lineFor := func() string {
|
||||
hook.Reset()
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion(qname, dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
line := responseLineFor(hook, qname)
|
||||
// The duration differs per query and is not what we compare.
|
||||
line, _, _ = strings.Cut(line, " took=")
|
||||
return line
|
||||
}
|
||||
|
||||
first := lineFor()
|
||||
require.NotEmpty(t, first, "the chain must log the response it wrote")
|
||||
require.Contains(t, first, "deferred_by=", "the line must carry more than one annotation to be worth ordering")
|
||||
require.Contains(t, first, "softened=")
|
||||
|
||||
for range 20 {
|
||||
assert.Equal(t, first, lineFor(), "the same event must produce the same line")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("logged with the reason it was deferred", func(t *testing.T) {
|
||||
const qname = "_mongodb._tcp.reason.example.com."
|
||||
|
||||
hook.Reset()
|
||||
|
||||
prev := log.GetLevel()
|
||||
log.SetLevel(log.TraceLevel)
|
||||
t.Cleanup(func() { log.SetLevel(prev) })
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion(qname, dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
response := responseLineFor(hook, qname)
|
||||
require.NotEmpty(t, response, "the chain must log the response it wrote")
|
||||
assert.Contains(t, response, "softened=", "the log must show the verdict was rewritten")
|
||||
assert.Contains(t, response, "deferred_by=", "the log must name the handler that deferred")
|
||||
})
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_NoHandlerBelow covers a client with no primary
|
||||
// nameserver group: the deferred query reaches the end of the chain unanswered.
|
||||
// REFUSED would say the name is not served here while the route serves its
|
||||
// addresses, and a stub that acts on that by asking elsewhere can bring back an
|
||||
// NXDOMAIN for the whole name. The answer must be an empty, uncacheable NODATA.
|
||||
func TestHandlerChain_SoftNegative_NoHandlerBelow(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", dns.TypeHTTPS)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "an unanswered soft-negative query must be NODATA, not REFUSED")
|
||||
assert.Empty(t, resp.Answer)
|
||||
assert.Empty(t, resp.Ns,
|
||||
"no SOA, so RFC 2308 keeps the empty answer out of negative caches and it cannot outlive the route")
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_KeepsRealAnswers guards the other direction:
|
||||
// softening applies to negative verdicts only. A real answer from a downstream
|
||||
// handler must reach the client untouched.
|
||||
func TestHandlerChain_SoftNegative_KeepsRealAnswers(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &answeringHandler{name: "public", ip: "203.0.113.10"}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", dns.TypeA)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
require.Len(t, resp.Answer, 1, "the downstream answer must pass through")
|
||||
}
|
||||
|
||||
// TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative makes sure the
|
||||
// softening is opt-in: an ordinary chain continuation still yields NXDOMAIN, so
|
||||
// genuine non-existence keeps being reported.
|
||||
func TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("nope.example.com.", dns.TypeA)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "without the signal a real NXDOMAIN must survive")
|
||||
}
|
||||
|
||||
func TestHandlerChain_HasRootHandlerAtOrBelow(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
h := &answeringHandler{name: "h", ip: "10.0.0.1"}
|
||||
|
||||
@@ -19,6 +19,34 @@ import (
|
||||
// uses when a resolved host has no addresses of the requested family.
|
||||
const errNoSuitableAddress = "no suitable address found"
|
||||
|
||||
// Extended DNS Error info codes NetBird emits so a client can see why an answer
|
||||
// looks the way it does without reading this peer's logs. They live in the RFC
|
||||
// 8914 Private Use range (49152-65535) and are registered here, in one place,
|
||||
// because nothing else guarantees two NetBird components pick distinct codes.
|
||||
const (
|
||||
// EDENetbirdUpstreamTimeout: a DNS forwarder's upstream did not answer.
|
||||
EDENetbirdUpstreamTimeout uint16 = 49152
|
||||
// EDENetbirdUpstreamFailure: a DNS forwarder's upstream failed.
|
||||
EDENetbirdUpstreamFailure uint16 = 49153
|
||||
// EDENetbirdSoftenedNegative: the empty answer is ours, not the answering
|
||||
// resolver's. A handler that owns the name deferred this query type, and the
|
||||
// negative verdict that came back was downgraded so it cannot poison the
|
||||
// name.
|
||||
EDENetbirdSoftenedNegative uint16 = 49154
|
||||
)
|
||||
|
||||
// AttachEDE adds an Extended DNS Error (RFC 8914) option to a message, creating
|
||||
// the OPT pseudo-record if it has none. Callers must only use it toward a client
|
||||
// that advertised EDNS0: per RFC 6891 an OPT must not appear otherwise.
|
||||
func AttachEDE(msg *dns.Msg, code uint16, text string) {
|
||||
opt := msg.IsEdns0()
|
||||
if opt == nil {
|
||||
msg.SetEdns0(dns.DefaultMsgSize, false)
|
||||
opt = msg.IsEdns0()
|
||||
}
|
||||
opt.Option = append(opt.Option, &dns.EDNS0_EDE{InfoCode: code, ExtraText: text})
|
||||
}
|
||||
|
||||
// GenerateRequestID creates a random 8-character hex string for request tracing.
|
||||
func GenerateRequestID() string {
|
||||
bytes := make([]byte, 4)
|
||||
@@ -78,10 +106,38 @@ type resolver interface {
|
||||
LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error)
|
||||
}
|
||||
|
||||
// MetaKey names an annotation a handler attaches to a request to explain the
|
||||
// response the chain ends up writing. The set is closed and rendered onto a
|
||||
// single log line, so an unrecognized key is a typo inventing a field rather
|
||||
// than a new fact.
|
||||
type MetaKey string
|
||||
|
||||
const (
|
||||
// MetaKeyProtocol: the transport a query arrived on.
|
||||
MetaKeyProtocol MetaKey = "protocol"
|
||||
// MetaKeyUpstream: the upstream that answered.
|
||||
MetaKeyUpstream MetaKey = "upstream"
|
||||
// MetaKeyUpstreamProtocol: the transport used toward that upstream.
|
||||
MetaKeyUpstreamProtocol MetaKey = "upstream_protocol"
|
||||
// MetaKeyPeer: the routing peer whose forwarder answered.
|
||||
MetaKeyPeer MetaKey = "peer"
|
||||
// MetaKeyEDE: an Extended DNS Error carried by the answer.
|
||||
MetaKeyEDE MetaKey = "ede"
|
||||
// MetaKeyTruncated: the answer did not fit and was truncated.
|
||||
MetaKeyTruncated MetaKey = "truncated"
|
||||
// MetaKeySoftened: a negative verdict was rewritten so it cannot poison a
|
||||
// name served locally.
|
||||
MetaKeySoftened MetaKey = "softened"
|
||||
// MetaKeyDeferredBy: the handler that owned the name and stepped aside.
|
||||
MetaKeyDeferredBy MetaKey = "deferred_by"
|
||||
// MetaKeyDeferredReason: why it stepped aside.
|
||||
MetaKeyDeferredReason MetaKey = "deferred_reason"
|
||||
)
|
||||
|
||||
// chainedWriter is implemented by ResponseWriters that carry request metadata
|
||||
type chainedWriter interface {
|
||||
RequestID() string
|
||||
SetMeta(key, value string)
|
||||
SetMeta(key MetaKey, value string)
|
||||
}
|
||||
|
||||
// GetRequestID extracts a request ID from the ResponseWriter if available,
|
||||
@@ -96,12 +152,30 @@ func GetRequestID(w dns.ResponseWriter) string {
|
||||
}
|
||||
|
||||
// SetMeta sets metadata on the ResponseWriter if it supports it.
|
||||
func SetMeta(w dns.ResponseWriter, key, value string) {
|
||||
func SetMeta(w dns.ResponseWriter, key MetaKey, value string) {
|
||||
if cw, ok := w.(chainedWriter); ok {
|
||||
cw.SetMeta(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
// softNegativeRequester is implemented by chain writers that can soften the
|
||||
// negative verdict of the handlers a deferring handler falls through to.
|
||||
type softNegativeRequester interface {
|
||||
RequestSoftNegative()
|
||||
}
|
||||
|
||||
// RequestSoftNegative asks the handler chain to downgrade an NXDOMAIN from the
|
||||
// handlers that run after the caller defers. A handler that owns a name but
|
||||
// cannot answer one query type for it needs this: the resolvers it falls
|
||||
// through to cannot prove the name absent, and an NXDOMAIN from them is cached
|
||||
// for the name and every type under it (RFC 2308, RFC 8020), taking the
|
||||
// addresses the handler does serve down with it.
|
||||
func RequestSoftNegative(w dns.ResponseWriter) {
|
||||
if sn, ok := w.(softNegativeRequester); ok {
|
||||
sn.RequestSoftNegative()
|
||||
}
|
||||
}
|
||||
|
||||
// LookupResult contains the result of an external DNS lookup
|
||||
type LookupResult struct {
|
||||
IPs []netip.Addr
|
||||
@@ -199,11 +273,27 @@ type RecordResolver interface {
|
||||
LookupAddr(ctx context.Context, addr string) ([]string, error)
|
||||
}
|
||||
|
||||
// SupportedRecordQtype reports whether LookupRecords can resolve qtype. The
|
||||
// set is bounded by the net.Resolver API, which exposes no way to query an
|
||||
// arbitrary record type, and going around it with a raw exchange would mean
|
||||
// picking nameservers ourselves instead of resolving the way the host does.
|
||||
//
|
||||
// Both ends read this: a DNS forwarder answers these types for the domains it
|
||||
// routes, and a client uses it to tell which types are worth forwarding to a
|
||||
// peer at all.
|
||||
func SupportedRecordQtype(qtype uint16) bool {
|
||||
switch qtype {
|
||||
case dns.TypeMX, dns.TypeTXT, dns.TypeNS, dns.TypeSRV, dns.TypeCNAME, dns.TypePTR:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// LookupRecords resolves a non-address DNS record type through the host
|
||||
// resolver and returns the resource records and the DNS rcode. Types the host
|
||||
// resolver cannot answer (anything not covered by the net.Resolver Lookup*
|
||||
// methods) yield NODATA so that a routed name is never poisoned with NXDOMAIN
|
||||
// for an unsupported type.
|
||||
// resolver and returns the resource records and the DNS rcode. Types outside
|
||||
// SupportedRecordQtype yield NODATA so that a routed name is never poisoned
|
||||
// with NXDOMAIN for a type we cannot look up.
|
||||
func LookupRecords(ctx context.Context, r RecordResolver, name string, qtype uint16, ttl uint32) ([]dns.RR, int) {
|
||||
fqdn := dns.Fqdn(name)
|
||||
|
||||
|
||||
@@ -295,7 +295,7 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
if addr := w.RemoteAddr(); addr != nil {
|
||||
network := addr.Network()
|
||||
ctx = contextWithDNSProtocol(ctx, network)
|
||||
resutil.SetMeta(w, "protocol", network)
|
||||
resutil.SetMeta(w, resutil.MetaKeyProtocol, network)
|
||||
}
|
||||
|
||||
ok, failures := u.tryUpstreamServers(ctx, w, r, logger)
|
||||
@@ -331,7 +331,7 @@ func (u *upstreamResolverBase) tryOnlyRace(ctx context.Context, w dns.ResponseWr
|
||||
return false, res.failures
|
||||
}
|
||||
if res.ede != "" {
|
||||
resutil.SetMeta(w, "ede", res.ede)
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, res.ede)
|
||||
}
|
||||
u.writeSuccessResponse(w, res.msg, res.upstream, r.Question[0].Name, res.protocol, logger)
|
||||
return true, res.failures
|
||||
@@ -361,7 +361,7 @@ func (u *upstreamResolverBase) raceAll(ctx context.Context, w dns.ResponseWriter
|
||||
failures = append(failures, res.failures...)
|
||||
if res.msg != nil {
|
||||
if res.ede != "" {
|
||||
resutil.SetMeta(w, "ede", res.ede)
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, res.ede)
|
||||
}
|
||||
u.writeSuccessResponse(w, res.msg, res.upstream, r.Question[0].Name, res.protocol, logger)
|
||||
return true, failures
|
||||
@@ -550,9 +550,9 @@ func (u *upstreamResolverBase) debugUpstreamTimeout(upstream netip.AddrPort) str
|
||||
}
|
||||
|
||||
func (u *upstreamResolverBase) writeSuccessResponse(w dns.ResponseWriter, rm *dns.Msg, upstream netip.AddrPort, domain string, proto string, logger *log.Entry) {
|
||||
resutil.SetMeta(w, "upstream", upstream.String())
|
||||
resutil.SetMeta(w, resutil.MetaKeyUpstream, upstream.String())
|
||||
if proto != "" {
|
||||
resutil.SetMeta(w, "upstream_protocol", proto)
|
||||
resutil.SetMeta(w, resutil.MetaKeyUpstreamProtocol, proto)
|
||||
}
|
||||
|
||||
// Clear Zero bit from external responses to prevent upstream servers from
|
||||
|
||||
@@ -26,15 +26,6 @@ import (
|
||||
const errResolveFailed = "failed to resolve query for domain=%s: %v"
|
||||
const upstreamTimeout = 15 * time.Second
|
||||
|
||||
// EDE info codes the forwarder emits on upstream failures so the querying
|
||||
// client can see the reason without inspecting this peer's logs. They live in
|
||||
// the RFC 8914 Private Use range (49152-65535); the Go resolver never exposes a
|
||||
// real upstream EDE here, so these cannot collide with a genuine code.
|
||||
const (
|
||||
edeNetbirdUpstreamTimeout uint16 = 49152
|
||||
edeNetbirdUpstreamFailure uint16 = 49153
|
||||
)
|
||||
|
||||
type resolver interface {
|
||||
LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error)
|
||||
LookupMX(ctx context.Context, name string) ([]*net.MX, error)
|
||||
@@ -216,6 +207,9 @@ func (f *DNSForwarder) handleDNSQuery(logger *log.Entry, w dns.ResponseWriter, q
|
||||
qname, dns.TypeToString[question.Qtype], dns.ClassToString[question.Qclass])
|
||||
|
||||
resp := query.SetReply(query)
|
||||
// Every answer here comes from a recursive lookup on this peer. SetReply
|
||||
// leaves RA unset, which reads to a client as a server that cannot recurse.
|
||||
resp.RecursionAvailable = true
|
||||
|
||||
mostSpecificResId, matchingEntries := f.getMatchingEntries(strings.TrimSuffix(qname, "."))
|
||||
if mostSpecificResId == "" {
|
||||
@@ -229,20 +223,22 @@ func (f *DNSForwarder) handleDNSQuery(logger *log.Entry, w dns.ResponseWriter, q
|
||||
|
||||
reqHasEdns := query.IsEdns0() != nil
|
||||
|
||||
switch question.Qtype {
|
||||
case dns.TypeA, dns.TypeAAAA:
|
||||
switch {
|
||||
case question.Qtype == dns.TypeA || question.Qtype == dns.TypeAAAA:
|
||||
f.handleAddressQuery(ctx, logger, w, resp, mostSpecificResId, matchingEntries, reqHasEdns, startTime)
|
||||
case dns.TypeMX, dns.TypeTXT, dns.TypeNS, dns.TypeSRV, dns.TypeCNAME, dns.TypePTR:
|
||||
case resutil.SupportedRecordQtype(question.Qtype):
|
||||
f.handleRecordQuery(ctx, logger, w, resp, startTime)
|
||||
default:
|
||||
// The domain is routed here, so any other type is answered NODATA
|
||||
// (NOERROR, empty answer) rather than falling back to a resolver that
|
||||
// would poison the name with NXDOMAIN. The Extended DNS Error lets a
|
||||
// client tell this capability-driven NODATA apart from an
|
||||
// authoritative one. The OPT pseudo-record must not appear unless the
|
||||
// query advertised EDNS0.
|
||||
// authoritative one; a current client knows the type is unsupported and
|
||||
// resolves it through its own handler chain instead of asking at all.
|
||||
// The OPT pseudo-record must not appear unless the query advertised
|
||||
// EDNS0.
|
||||
if reqHasEdns {
|
||||
attachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
|
||||
resutil.AttachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
|
||||
}
|
||||
f.writeResponse(logger, w, resp, qname, startTime)
|
||||
}
|
||||
@@ -441,7 +437,7 @@ func (f *DNSForwarder) handleDNSError(
|
||||
}
|
||||
|
||||
if reqHasEdns {
|
||||
attachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
|
||||
resutil.AttachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
|
||||
}
|
||||
|
||||
f.writeResponse(logger, w, resp, domain, startTime)
|
||||
@@ -488,9 +484,9 @@ func (f *DNSForwarder) getMatchingEntries(domain string) (route.ResID, []*Forwar
|
||||
// edeCodeFor maps an upstream lookup error to the NetBird EDE info code.
|
||||
func edeCodeFor(dnsErr *net.DNSError) uint16 {
|
||||
if dnsErr != nil && dnsErr.IsTimeout {
|
||||
return edeNetbirdUpstreamTimeout
|
||||
return resutil.EDENetbirdUpstreamTimeout
|
||||
}
|
||||
return edeNetbirdUpstreamFailure
|
||||
return resutil.EDENetbirdUpstreamFailure
|
||||
}
|
||||
|
||||
// edeText builds the EDE extra-text describing the class of upstream failure.
|
||||
@@ -503,14 +499,3 @@ func edeText(dnsErr *net.DNSError) string {
|
||||
}
|
||||
return "netbird forwarder: upstream failure"
|
||||
}
|
||||
|
||||
// attachEDE adds an Extended DNS Error (RFC 8914) option to the response,
|
||||
// creating the OPT pseudo-record if the response does not already carry one.
|
||||
func attachEDE(resp *dns.Msg, code uint16, text string) {
|
||||
opt := resp.IsEdns0()
|
||||
if opt == nil {
|
||||
resp.SetEdns0(dns.DefaultMsgSize, false)
|
||||
opt = resp.IsEdns0()
|
||||
}
|
||||
opt.Option = append(opt.Option, &dns.EDNS0_EDE{InfoCode: code, ExtraText: text})
|
||||
}
|
||||
|
||||
@@ -649,6 +649,39 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestDNSForwarder_RecursionAvailable covers the RA bit. Every answer this
|
||||
// forwarder produces comes from a recursive lookup on the routing peer, but
|
||||
// dns.Msg.SetReply leaves RA unset, so clients report "recursion not available"
|
||||
// and some stub resolvers treat the server as unable to serve the query.
|
||||
func TestDNSForwarder_RecursionAvailable(t *testing.T) {
|
||||
t.Run("record answer", func(t *testing.T) {
|
||||
mockResolver := &MockResolver{}
|
||||
forwarder := newRecordTestForwarder(t, mockResolver, "example.com")
|
||||
|
||||
mockResolver.On("LookupMX", mock.Anything, "example.com.").
|
||||
Return([]*net.MX{{Host: "mail.example.com.", Pref: 10}}, nil).Once()
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "example.com", dns.TypeMX)
|
||||
assert.True(t, resp.RecursionAvailable, "the forwarder resolves recursively")
|
||||
})
|
||||
|
||||
t.Run("unsupported type NODATA", func(t *testing.T) {
|
||||
forwarder := newRecordTestForwarder(t, &MockResolver{}, "example.com")
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "example.com", dns.TypeCAA)
|
||||
require.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
assert.True(t, resp.RecursionAvailable, "RA describes the server, not the query type")
|
||||
})
|
||||
|
||||
t.Run("unauthorized domain", func(t *testing.T) {
|
||||
forwarder := newRecordTestForwarder(t, &MockResolver{}, "example.com")
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "other.com", dns.TypeMX)
|
||||
require.Equal(t, dns.RcodeRefused, resp.Rcode)
|
||||
assert.True(t, resp.RecursionAvailable, "RA describes the server, not the verdict")
|
||||
})
|
||||
}
|
||||
|
||||
func hasEDE(m *dns.Msg, code uint16) bool {
|
||||
opt := m.IsEdns0()
|
||||
if opt == nil {
|
||||
@@ -859,7 +892,7 @@ func TestDNSForwarder_UpstreamFailureEDE(t *testing.T) {
|
||||
lookupErr: &net.DNSError{Err: "i/o timeout", Server: "10.0.0.53:53", IsTimeout: true},
|
||||
reqEdns: true,
|
||||
wantEDE: true,
|
||||
wantCode: edeNetbirdUpstreamTimeout,
|
||||
wantCode: resutil.EDENetbirdUpstreamTimeout,
|
||||
wantTextHas: "netbird forwarder: upstream timeout",
|
||||
},
|
||||
{
|
||||
@@ -867,7 +900,7 @@ func TestDNSForwarder_UpstreamFailureEDE(t *testing.T) {
|
||||
lookupErr: &net.DNSError{Err: "server misbehaving", Server: "10.0.0.53:53"},
|
||||
reqEdns: true,
|
||||
wantEDE: true,
|
||||
wantCode: edeNetbirdUpstreamFailure,
|
||||
wantCode: resutil.EDENetbirdUpstreamFailure,
|
||||
wantTextHas: "netbird forwarder: upstream failure",
|
||||
},
|
||||
{
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
nbdns "github.com/netbirdio/netbird/client/internal/dns"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/peerstore"
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager/common"
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager/fakeip"
|
||||
iface "github.com/netbirdio/netbird/client/internal/routemanager/iface"
|
||||
@@ -40,6 +39,61 @@ type internalDNATer interface {
|
||||
AddInternalDNATMapping(netip.Addr, netip.Addr) error
|
||||
}
|
||||
|
||||
// peerAllowedIPs reports the tunnel address a peer is reachable on.
|
||||
type peerAllowedIPs interface {
|
||||
AllowedIP(pubKey string) (netip.Addr, bool)
|
||||
}
|
||||
|
||||
// disposition says where a query for a given record type has to be answered.
|
||||
type disposition int
|
||||
|
||||
const (
|
||||
// dispositionPeer: the routing peer owns the answer. Address records are
|
||||
// what a DNS route exists for, and their answer programs the routes,
|
||||
// allowed IPs and firewall sets, so they are never resolved anywhere else.
|
||||
dispositionPeer disposition = iota
|
||||
// dispositionPeerFirst: the peer's forwarder resolves the type, and the
|
||||
// records may only exist inside the routed network, so it has to be asked.
|
||||
// A peer running a client that predates that support cannot answer, which
|
||||
// only its reply reveals.
|
||||
dispositionPeerFirst
|
||||
// dispositionChain: no forwarder resolves the type, so asking would cost a
|
||||
// tunnel round trip for a reply that says nothing. Public records for the
|
||||
// name are still better than none, so the query goes to the rest of the
|
||||
// chain.
|
||||
dispositionChain
|
||||
)
|
||||
|
||||
func dispositionFor(qtype uint16) disposition {
|
||||
switch {
|
||||
case qtype == dns.TypeA || qtype == dns.TypeAAAA:
|
||||
return dispositionPeer
|
||||
case resutil.SupportedRecordQtype(qtype):
|
||||
return dispositionPeerFirst
|
||||
default:
|
||||
return dispositionChain
|
||||
}
|
||||
}
|
||||
|
||||
// peerCannotAnswer reports whether a forwarder reply means the peer is unable to
|
||||
// resolve this query type, as opposed to having resolved it and found nothing.
|
||||
// A forwarder older than the one that learned to resolve non-address types
|
||||
// answers NOTIMP to all of them, before it even looks at the domain; FORMERR
|
||||
// covers one that cannot parse the EDNS0 we add to the query.
|
||||
//
|
||||
// REFUSED is deliberately not here. It says the peer does not hold the name,
|
||||
// which happens while the client and the peer disagree about the route set, and
|
||||
// falling through then would hand the name of an internal-only domain to a
|
||||
// public resolver. No forwarder version reports a missing record type that way.
|
||||
func peerCannotAnswer(rcode int) bool {
|
||||
switch rcode {
|
||||
case dns.RcodeNotImplemented, dns.RcodeFormatError:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
type DnsInterceptor struct {
|
||||
mu sync.RWMutex
|
||||
route *route.Route
|
||||
@@ -50,7 +104,7 @@ type DnsInterceptor struct {
|
||||
currentPeerKey string
|
||||
interceptedDomains domainMap
|
||||
wgInterface iface.WGIface
|
||||
peerStore *peerstore.Store
|
||||
peerStore peerAllowedIPs
|
||||
firewall firewall.Manager
|
||||
fakeIPManager *fakeip.Manager
|
||||
forwarderPort *atomic.Uint32
|
||||
@@ -95,7 +149,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
|
||||
|
||||
// AllowedIPs should use real IPs
|
||||
if d.currentPeerKey != "" {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
|
||||
}
|
||||
}
|
||||
@@ -172,7 +226,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error {
|
||||
}
|
||||
|
||||
// AllowedIPs use real IPs
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil {
|
||||
return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err)
|
||||
}
|
||||
|
||||
@@ -205,7 +259,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error {
|
||||
for _, prefixes := range d.interceptedDomains {
|
||||
for _, prefix := range prefixes {
|
||||
// AllowedIPs use real IPs
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
|
||||
}
|
||||
}
|
||||
@@ -226,11 +280,14 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
return
|
||||
}
|
||||
|
||||
// All query types for an intercepted domain are forwarded to the peer's
|
||||
// DNS forwarder, which owns the name. Falling through to the system
|
||||
// resolver would let it answer NXDOMAIN for a name it isn't authoritative
|
||||
// for, poisoning the whole name (including the A/AAAA records the route
|
||||
// does serve). The forwarder answers NODATA for types it cannot resolve.
|
||||
qtype := r.Question[0].Qtype
|
||||
dispose := dispositionFor(qtype)
|
||||
|
||||
if dispose == dispositionChain {
|
||||
d.deferToChain(w, r, logger, "unsupported-qtype")
|
||||
return
|
||||
}
|
||||
|
||||
d.mu.RLock()
|
||||
peerKey := d.currentPeerKey
|
||||
d.mu.RUnlock()
|
||||
@@ -246,35 +303,32 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
return
|
||||
}
|
||||
|
||||
if r.Extra == nil {
|
||||
r.MsgHdr.AuthenticatedData = true
|
||||
}
|
||||
|
||||
// Advertise EDNS0 to the forwarder so it may return an Extended DNS Error
|
||||
// describing why a lookup failed. The OPT is stripped from the reply when
|
||||
// the original client did not request EDNS0.
|
||||
hadEdns := r.IsEdns0() != nil
|
||||
if !hadEdns {
|
||||
r.SetEdns0(dns.DefaultMsgSize, false)
|
||||
}
|
||||
query, hadEdns := peerQuery(r)
|
||||
|
||||
upstream := net.JoinHostPort(upstreamIP.String(), strconv.FormatUint(uint64(d.forwarderPort.Load()), 10))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), dnsTimeout)
|
||||
defer cancel()
|
||||
|
||||
reply := d.queryUpstreamDNS(ctx, w, r, upstream, upstreamIP, peerKey, logger)
|
||||
reply := d.queryUpstreamDNS(ctx, w, query, upstream, upstreamIP, peerKey, logger)
|
||||
if reply == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// The peer owns the name but its forwarder cannot resolve this type. No
|
||||
// capability is announced anywhere, so the round trip above is the probe.
|
||||
if dispose == dispositionPeerFirst && peerCannotAnswer(reply.Rcode) {
|
||||
d.deferToChain(w, r, logger, "peer-rcode-"+dns.RcodeToString[reply.Rcode])
|
||||
return
|
||||
}
|
||||
|
||||
if ede, ok := resutil.ExtractEDE(reply); ok {
|
||||
resutil.SetMeta(w, "ede", fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
|
||||
}
|
||||
if !hadEdns {
|
||||
resutil.StripOPT(reply)
|
||||
}
|
||||
|
||||
resutil.SetMeta(w, "peer", peerKey)
|
||||
resutil.SetMeta(w, resutil.MetaKeyPeer, peerKey)
|
||||
|
||||
reply.Id = r.Id
|
||||
if err := d.writeMsg(w, reply, logger); err != nil {
|
||||
@@ -282,6 +336,30 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
}
|
||||
}
|
||||
|
||||
// deferToChain hands the query to the next handler in the chain and asks the
|
||||
// chain to soften the negative verdict of whatever answers instead. The
|
||||
// resolvers below cannot prove a name inside the routed network absent, so their
|
||||
// NXDOMAIN must not reach the client: it would be cached for the name and every
|
||||
// type under it, taking the addresses this route does serve with it.
|
||||
func (d *DnsInterceptor) deferToChain(w dns.ResponseWriter, r *dns.Msg, logger *log.Entry, reason string) {
|
||||
logger.Tracef("continuing to next handler for domain=%s type=%s reason=%s",
|
||||
r.Question[0].Name, dns.TypeToString[r.Question[0].Qtype], reason)
|
||||
|
||||
resutil.RequestSoftNegative(w)
|
||||
// Carried to the chain's response log line: without it the answer looks like
|
||||
// it came from the fallthrough resolver on its own.
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredBy, "dns-route")
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredReason, reason)
|
||||
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
// Set Zero bit to signal handler chain to continue
|
||||
resp.MsgHdr.Zero = true
|
||||
if err := w.WriteMsg(resp); err != nil {
|
||||
logger.Errorf("failed writing DNS continue response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (d *DnsInterceptor) writeDNSError(w dns.ResponseWriter, r *dns.Msg, logger *log.Entry, reason string) {
|
||||
logger.Warnf("failed to query upstream for domain=%s: %s", r.Question[0].Name, reason)
|
||||
|
||||
@@ -621,3 +699,26 @@ func (d *DnsInterceptor) debugPeerTimeout(peerIP netip.Addr, peerKey string) str
|
||||
|
||||
return fmt.Sprintf(" (peer %s)", nbdns.FormatPeerStatus(&peerState))
|
||||
}
|
||||
|
||||
// peerQuery builds the query sent to a peer's DNS forwarder and reports whether
|
||||
// the client itself advertised EDNS0. EDNS0 is added so the forwarder can return
|
||||
// an Extended DNS Error describing an upstream failure, and the OPT is stripped
|
||||
// from the reply again when the client did not ask for it. The client's message
|
||||
// is left untouched: it may still be handed to the next handler in the chain,
|
||||
// and neither the OPT nor the AD bit is ours to put on the wire on its behalf.
|
||||
func peerQuery(r *dns.Msg) (query *dns.Msg, hadEdns bool) {
|
||||
query = r.Copy()
|
||||
hadEdns = query.IsEdns0() != nil
|
||||
|
||||
// AD tells the forwarder we understand authenticated data. Only set when the
|
||||
// client sent no additional section of its own, so we never overrule what it
|
||||
// asked for.
|
||||
if len(query.Extra) == 0 {
|
||||
query.MsgHdr.AuthenticatedData = true
|
||||
}
|
||||
if !hadEdns {
|
||||
query.SetEdns0(dns.DefaultMsgSize, false)
|
||||
}
|
||||
|
||||
return query, hadEdns
|
||||
}
|
||||
|
||||
399
client/internal/routemanager/dnsinterceptor/handler_test.go
Normal file
399
client/internal/routemanager/dnsinterceptor/handler_test.go
Normal file
@@ -0,0 +1,399 @@
|
||||
package dnsinterceptor
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.zx2c4.com/wireguard/tun/netstack"
|
||||
|
||||
"github.com/netbirdio/netbird/client/iface/device"
|
||||
"github.com/netbirdio/netbird/client/iface/wgaddr"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/test"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
// softNegativeWriter records what the handler told the chain: whether to soften
|
||||
// a negative verdict from the handlers it defers to, and the metadata that ends
|
||||
// up on the chain's response log line.
|
||||
type softNegativeWriter struct {
|
||||
test.MockResponseWriter
|
||||
softNegative bool
|
||||
meta map[resutil.MetaKey]string
|
||||
}
|
||||
|
||||
func (w *softNegativeWriter) RequestSoftNegative() { w.softNegative = true }
|
||||
|
||||
func (w *softNegativeWriter) RequestID() string { return "test" }
|
||||
|
||||
func (w *softNegativeWriter) SetMeta(key resutil.MetaKey, value string) {
|
||||
if w.meta == nil {
|
||||
w.meta = make(map[resutil.MetaKey]string)
|
||||
}
|
||||
w.meta[key] = value
|
||||
}
|
||||
|
||||
// TestServeDNS_QtypeTheForwarderCannotResolve covers the record types no DNS
|
||||
// forwarder can answer, because the host resolver exposes no API for them.
|
||||
// Asking the peer only burns a tunnel round trip, so the query goes straight to
|
||||
// the rest of the chain. No peer key is configured, proving no round trip is
|
||||
// attempted.
|
||||
func TestServeDNS_QtypeTheForwarderCannotResolve(t *testing.T) {
|
||||
qtypes := []uint16{
|
||||
dns.TypeHTTPS,
|
||||
dns.TypeSVCB,
|
||||
dns.TypeCAA,
|
||||
dns.TypeNAPTR,
|
||||
dns.TypeTLSA,
|
||||
dns.TypeSOA,
|
||||
}
|
||||
|
||||
for _, qtype := range qtypes {
|
||||
t.Run(dns.TypeToString[qtype], func(t *testing.T) {
|
||||
d := &DnsInterceptor{}
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", qtype)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "chain continuation is signalled as NXDOMAIN")
|
||||
assert.True(t, resp.MsgHdr.Zero, "Zero bit must be set so the chain continues")
|
||||
assert.True(t, w.softNegative,
|
||||
"a downstream NXDOMAIN must be softened, or the fallthrough poisons the routed name")
|
||||
assert.Equal(t, "dns-route", w.meta[resutil.MetaKeyDeferredBy],
|
||||
"the chain's response log line must name us as the handler that stepped aside")
|
||||
assert.Equal(t, "unsupported-qtype", w.meta[resutil.MetaKeyDeferredReason],
|
||||
"the reason must survive to the log line, or the fallthrough is invisible")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_AddressQtypeNeverFallsThrough locks in the opposite case: A and
|
||||
// AAAA are what the route exists for, and the peer is authoritative for them.
|
||||
// A missing peer is an error the client must see, never a fallthrough to a
|
||||
// resolver that knows nothing about the routed name.
|
||||
func TestServeDNS_AddressQtypeNeverFallsThrough(t *testing.T) {
|
||||
for _, qtype := range []uint16{dns.TypeA, dns.TypeAAAA} {
|
||||
t.Run(dns.TypeToString[qtype], func(t *testing.T) {
|
||||
d := &DnsInterceptor{}
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", qtype)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeServerFailure, resp.Rcode, "an unusable route must fail, not fall through")
|
||||
assert.False(t, resp.MsgHdr.Zero, "the chain must not continue for address queries")
|
||||
assert.False(t, w.softNegative, "no fallthrough means nothing to soften")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDispositionFor pins the query-type policy, and ties the types we ask a
|
||||
// peer about to the types a forwarder can actually resolve, so the two ends
|
||||
// cannot drift apart.
|
||||
func TestDispositionFor(t *testing.T) {
|
||||
tests := []struct {
|
||||
qtype uint16
|
||||
want disposition
|
||||
}{
|
||||
{dns.TypeA, dispositionPeer},
|
||||
{dns.TypeAAAA, dispositionPeer},
|
||||
{dns.TypeMX, dispositionPeerFirst},
|
||||
{dns.TypeTXT, dispositionPeerFirst},
|
||||
{dns.TypeNS, dispositionPeerFirst},
|
||||
{dns.TypeSRV, dispositionPeerFirst},
|
||||
{dns.TypeCNAME, dispositionPeerFirst},
|
||||
{dns.TypePTR, dispositionPeerFirst},
|
||||
{dns.TypeHTTPS, dispositionChain},
|
||||
{dns.TypeSVCB, dispositionChain},
|
||||
{dns.TypeCAA, dispositionChain},
|
||||
{dns.TypeNAPTR, dispositionChain},
|
||||
{dns.TypeTLSA, dispositionChain},
|
||||
{dns.TypeDS, dispositionChain},
|
||||
{dns.TypeDNSKEY, dispositionChain},
|
||||
{dns.TypeSOA, dispositionChain},
|
||||
{dns.TypeANY, dispositionChain},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(dns.TypeToString[tt.qtype], func(t *testing.T) {
|
||||
assert.Equal(t, tt.want, dispositionFor(tt.qtype), "disposition for %s", dns.TypeToString[tt.qtype])
|
||||
|
||||
// Address types are resolved by the forwarder's address path, not by
|
||||
// its record path, so they are deliberately outside
|
||||
// SupportedRecordQtype and their disposition does not depend on it.
|
||||
if tt.want == dispositionPeer {
|
||||
assert.False(t, resutil.SupportedRecordQtype(tt.qtype),
|
||||
"address types take the forwarder's address path")
|
||||
return
|
||||
}
|
||||
|
||||
assert.Equal(t, tt.want == dispositionPeerFirst, resutil.SupportedRecordQtype(tt.qtype),
|
||||
"only types a forwarder resolves may be sent to a peer")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeerCannotAnswer separates a peer that cannot resolve a type from one that
|
||||
// resolved it and found nothing. Only the former may be retried elsewhere:
|
||||
// treating NODATA or NXDOMAIN as a capability failure would send the client to a
|
||||
// public resolver behind the peer's back and prefer a public record over the
|
||||
// private one the route exists to reach.
|
||||
func TestPeerCannotAnswer(t *testing.T) {
|
||||
cannot := []int{dns.RcodeNotImplemented, dns.RcodeFormatError}
|
||||
for _, rcode := range cannot {
|
||||
assert.True(t, peerCannotAnswer(rcode), "%s means the peer cannot answer", dns.RcodeToString[rcode])
|
||||
}
|
||||
|
||||
// REFUSED belongs here, not above: it says the peer does not hold the name,
|
||||
// so retrying elsewhere would leak the name of an internal-only domain to a
|
||||
// public resolver. No forwarder version reports a missing record type this
|
||||
// way, so nothing is lost by keeping it out.
|
||||
can := []int{
|
||||
dns.RcodeSuccess,
|
||||
dns.RcodeNameError,
|
||||
dns.RcodeServerFailure,
|
||||
dns.RcodeNotAuth,
|
||||
dns.RcodeRefused,
|
||||
}
|
||||
for _, rcode := range can {
|
||||
assert.False(t, peerCannotAnswer(rcode), "%s is the peer's answer, not a capability failure", dns.RcodeToString[rcode])
|
||||
}
|
||||
}
|
||||
|
||||
// fakeWGIface is the minimum an interceptor needs to reach a peer's DNS
|
||||
// forwarder over a plain socket. A nil netstack keeps the exchange on the host
|
||||
// stack so the test can point it at a loopback listener.
|
||||
type fakeWGIface struct{}
|
||||
|
||||
func (fakeWGIface) AddAllowedIP(string, netip.Prefix) error { return nil }
|
||||
func (fakeWGIface) RemoveAllowedIP(string, netip.Prefix) error { return nil }
|
||||
func (fakeWGIface) Name() string { return "wt0" }
|
||||
func (fakeWGIface) Address() wgaddr.Address { return wgaddr.Address{} }
|
||||
func (fakeWGIface) ToInterface() *net.Interface { return nil }
|
||||
func (fakeWGIface) IsUserspaceBind() bool { return false }
|
||||
func (fakeWGIface) GetFilter() device.PacketFilter { return nil }
|
||||
func (fakeWGIface) GetDevice() *device.FilteredDevice { return nil }
|
||||
func (fakeWGIface) GetNet() *netstack.Net { return nil }
|
||||
|
||||
// fakePeerIPs resolves every peer key to the loopback address the test's
|
||||
// forwarder stub listens on.
|
||||
type fakePeerIPs struct {
|
||||
addr netip.Addr
|
||||
}
|
||||
|
||||
func (p fakePeerIPs) AllowedIP(string) (netip.Addr, bool) { return p.addr, p.addr.IsValid() }
|
||||
|
||||
// forwarderStub stands in for the DNS forwarder of a routing peer, recording
|
||||
// what it was asked and answering with a canned reply.
|
||||
type forwarderStub struct {
|
||||
mu sync.Mutex
|
||||
queries []*dns.Msg
|
||||
reply func(q *dns.Msg) *dns.Msg
|
||||
port uint16
|
||||
shutdown func()
|
||||
}
|
||||
|
||||
func (s *forwarderStub) received() []*dns.Msg {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return append([]*dns.Msg(nil), s.queries...)
|
||||
}
|
||||
|
||||
func newForwarderStub(t *testing.T, reply func(q *dns.Msg) *dns.Msg) *forwarderStub {
|
||||
t.Helper()
|
||||
|
||||
conn, err := net.ListenUDP("udp", net.UDPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:0")))
|
||||
require.NoError(t, err, "listen for the forwarder stub")
|
||||
|
||||
stub := &forwarderStub{reply: reply}
|
||||
stub.port = uint16(conn.LocalAddr().(*net.UDPAddr).Port)
|
||||
|
||||
mux := dns.NewServeMux()
|
||||
mux.HandleFunc(".", func(w dns.ResponseWriter, q *dns.Msg) {
|
||||
stub.mu.Lock()
|
||||
stub.queries = append(stub.queries, q.Copy())
|
||||
stub.mu.Unlock()
|
||||
_ = w.WriteMsg(stub.reply(q))
|
||||
})
|
||||
|
||||
srv := &dns.Server{PacketConn: conn, Handler: mux}
|
||||
started := make(chan struct{})
|
||||
srv.NotifyStartedFunc = func() { close(started) }
|
||||
go func() {
|
||||
_ = srv.ActivateAndServe()
|
||||
}()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("forwarder stub did not start")
|
||||
}
|
||||
|
||||
stub.shutdown = func() { _ = srv.Shutdown() }
|
||||
t.Cleanup(stub.shutdown)
|
||||
|
||||
return stub
|
||||
}
|
||||
|
||||
func newTestInterceptor(t *testing.T, stub *forwarderStub) *DnsInterceptor {
|
||||
t.Helper()
|
||||
|
||||
port := new(atomic.Uint32)
|
||||
port.Store(uint32(stub.port))
|
||||
|
||||
return &DnsInterceptor{
|
||||
route: &route.Route{Domains: nil},
|
||||
statusRecorder: peer.NewRecorder("https://mgm"),
|
||||
currentPeerKey: "peer-key",
|
||||
interceptedDomains: make(domainMap),
|
||||
wgInterface: fakeWGIface{},
|
||||
peerStore: fakePeerIPs{addr: netip.MustParseAddr("127.0.0.1")},
|
||||
forwarderPort: port,
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerCannotResolveQtype is the reported regression: a routing peer
|
||||
// running a client older than the one that taught the forwarder about non-address
|
||||
// record types answers NOTIMP for every SRV query, and the interceptor used to
|
||||
// hand that straight to the application, breaking SRV-based service discovery.
|
||||
// The client cannot know the peer's capability up front, so it must try the peer
|
||||
// and fall back to the chain when the peer says it cannot answer.
|
||||
func TestServeDNS_PeerCannotResolveQtype(t *testing.T) {
|
||||
for _, rcode := range []int{dns.RcodeNotImplemented, dns.RcodeFormatError} {
|
||||
t.Run(dns.RcodeToString[rcode], func(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, rcode)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
require.Len(t, stub.received(), 1, "the peer must be asked before giving up on it")
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "chain continuation is signalled as NXDOMAIN")
|
||||
assert.True(t, resp.MsgHdr.Zero, "Zero bit must be set so the chain continues")
|
||||
assert.True(t, w.softNegative, "the fallthrough must not be allowed to poison the routed name")
|
||||
assert.Equal(t, "dns-route", w.meta[resutil.MetaKeyDeferredBy])
|
||||
assert.Equal(t, "peer-rcode-"+dns.RcodeToString[rcode], w.meta[resutil.MetaKeyDeferredReason],
|
||||
"the peer's verdict must be visible in the log, it is the only trace of the probe")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerVerdictPassesThrough covers the replies that are the peer's
|
||||
// answer rather than a statement about its capability. The peer is authoritative
|
||||
// for a name routed to it, so its verdict reaches the client as it is: retrying
|
||||
// these elsewhere would send the name of an internal-only domain to a public
|
||||
// resolver and prefer whatever that resolver says over the route.
|
||||
func TestServeDNS_PeerVerdictPassesThrough(t *testing.T) {
|
||||
// REFUSED means the peer does not consider the name routed to it, which
|
||||
// happens while the client and the peer disagree about the route set. It is
|
||||
// never how a peer reports a record type it cannot resolve: a forwarder too
|
||||
// old for non-address types answers NOTIMP before it looks at the domain.
|
||||
for _, rcode := range []int{dns.RcodeRefused, dns.RcodeNameError, dns.RcodeSuccess} {
|
||||
t.Run(dns.RcodeToString[rcode], func(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, rcode)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
require.Len(t, stub.received(), 1, "the peer must be asked exactly once")
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, rcode, resp.Rcode, "the peer's verdict must reach the client unchanged")
|
||||
assert.False(t, resp.MsgHdr.Zero, "the chain must not continue past an answered query")
|
||||
assert.False(t, w.softNegative, "nothing was deferred, so nothing may be softened")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerResolvesQtype guards the common case: a peer that can answer
|
||||
// the type stays authoritative, and its records reach the client unchanged.
|
||||
func TestServeDNS_PeerResolvesQtype(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetReply(q)
|
||||
resp.Answer = []dns.RR{&dns.SRV{
|
||||
Hdr: dns.RR_Header{Name: q.Question[0].Name, Rrtype: dns.TypeSRV, Class: dns.ClassINET, Ttl: 60},
|
||||
Target: "shard-00.db.example.com.",
|
||||
Port: 27017,
|
||||
}}
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
require.Len(t, resp.Answer, 1, "the peer's records must pass through")
|
||||
assert.False(t, resp.MsgHdr.Zero, "an answered query must not continue the chain")
|
||||
assert.False(t, w.softNegative)
|
||||
}
|
||||
|
||||
// TestServeDNS_DeferredQueryIsPristine covers what the fallthrough hands to the
|
||||
// next handler. The interceptor advertises EDNS0 to the peer so the forwarder can
|
||||
// return an Extended DNS Error, and sets AD, but neither belongs to the client's
|
||||
// query: on the deferred path they would travel to the public resolver and the
|
||||
// reply could come back carrying an OPT the client never advertised, which
|
||||
// RFC 6891 forbids us from passing on.
|
||||
func TestServeDNS_DeferredQueryIsPristine(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, dns.RcodeNotImplemented)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
sent := stub.received()
|
||||
require.Len(t, sent, 1)
|
||||
assert.NotNil(t, sent[0].IsEdns0(), "the peer must be asked with EDNS0 so it can return an EDE")
|
||||
|
||||
assert.Nil(t, r.IsEdns0(), "the deferred query must not carry an OPT the client never sent")
|
||||
assert.False(t, r.AuthenticatedData, "the deferred query must not carry an AD bit the client never set")
|
||||
assert.Empty(t, r.Extra, "the deferred query must reach the next handler as the client sent it")
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//go:build windows
|
||||
|
||||
package dnsinterceptor
|
||||
|
||||
// GetInterfaceGUIDString completes iface.WGIface on Windows, which requires the
|
||||
// interface GUID for DNS registration.
|
||||
func (fakeWGIface) GetInterfaceGUIDString() (string, error) { return "", nil }
|
||||
@@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error {
|
||||
var merr *multierror.Error
|
||||
for _, domainPrefixes := range r.dynamicDomains {
|
||||
for _, prefix := range domainPrefixes {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
|
||||
}
|
||||
}
|
||||
@@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err))
|
||||
}
|
||||
if r.currentPeerKey != "" {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -216,7 +216,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
|
||||
)
|
||||
}
|
||||
|
||||
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
|
||||
m.allowedIPsRefCounter = refcounter.New(
|
||||
func(prefix netip.Prefix, peerKey string) (string, error) {
|
||||
// save peerKey to use it in the remove function
|
||||
return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix)
|
||||
|
||||
@@ -54,7 +54,7 @@ func (m *reconcileWGMock) GetNet() *netstack.Net { return n
|
||||
func TestReconcilePeerAllowedIPs(t *testing.T) {
|
||||
wg := &reconcileWGMock{}
|
||||
m := &DefaultManager{wgInterface: wg}
|
||||
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
|
||||
m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
|
||||
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
|
||||
func(netip.Prefix, string) error { return nil },
|
||||
)
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
package refcounter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/hashicorp/go-multierror"
|
||||
|
||||
nberrors "github.com/netbirdio/netbird/client/errors"
|
||||
)
|
||||
|
||||
// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is
|
||||
// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most
|
||||
// one peer is active at a time even when several peers reference the prefix.
|
||||
type allowedIPsEntry struct {
|
||||
// peers maps a peerKey to the number of references holding the prefix for that peer.
|
||||
peers map[string]int
|
||||
// active is the peerKey currently installed in WireGuard for this prefix ("" if none).
|
||||
active string
|
||||
// total is the sum of all per-peer reference counts (kept in sync with peers).
|
||||
total int
|
||||
}
|
||||
|
||||
// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs.
|
||||
//
|
||||
// The generic Counter keys only by prefix and remembers a single Out value set by the first
|
||||
// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or
|
||||
// multiple resolved domains) can reference the same prefix through different peers, and when the
|
||||
// peer currently installed in WireGuard releases its last reference the prefix must be handed over
|
||||
// to a surviving peer instead of being left pointing at the released one.
|
||||
//
|
||||
// It calls add/remove (which program WireGuard) only on the transitions that matter:
|
||||
// - add on the first reference for a prefix, or when swapping the active peer;
|
||||
// - remove on the last reference for a prefix, or on the old peer during a swap.
|
||||
type AllowedIPsRefCounter struct {
|
||||
mu sync.Mutex
|
||||
entries map[netip.Prefix]*allowedIPsEntry
|
||||
add AddFunc[netip.Prefix, string, string]
|
||||
remove RemoveFunc[netip.Prefix, string]
|
||||
}
|
||||
|
||||
// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter.
|
||||
// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer.
|
||||
// remove unprograms the prefix from the given peer.
|
||||
func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter {
|
||||
return &AllowedIPsRefCounter{
|
||||
entries: map[netip.Prefix]*allowedIPsEntry{},
|
||||
add: add,
|
||||
remove: remove,
|
||||
}
|
||||
}
|
||||
|
||||
// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first
|
||||
// reference to a prefix; while a different peer is already installed the prefix is left with it
|
||||
// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept.
|
||||
func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
e, ok := rm.entries[prefix]
|
||||
if !ok {
|
||||
e = &allowedIPsEntry{peers: map[string]int{}}
|
||||
rm.entries[prefix] = e
|
||||
}
|
||||
|
||||
logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
|
||||
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active)
|
||||
|
||||
// Program WireGuard only when nothing is installed yet for this prefix.
|
||||
if e.active == "" {
|
||||
out, err := rm.add(prefix, peerKey)
|
||||
if errors.Is(err, ErrIgnore) {
|
||||
if e.total == 0 {
|
||||
delete(rm.entries, prefix)
|
||||
}
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
if err != nil {
|
||||
if e.total == 0 {
|
||||
delete(rm.entries, prefix)
|
||||
}
|
||||
return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err)
|
||||
}
|
||||
e.active = out
|
||||
}
|
||||
|
||||
e.peers[peerKey]++
|
||||
e.total++
|
||||
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
// Decrement removes a reference to prefix for peerKey. When the peer currently installed in
|
||||
// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists,
|
||||
// otherwise it is removed from WireGuard.
|
||||
func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
e, ok := rm.entries[prefix]
|
||||
if !ok {
|
||||
logCallerF("No allowed IP reference found for prefix %v", prefix)
|
||||
return Ref[string]{}, nil
|
||||
}
|
||||
|
||||
if e.peers[peerKey] > 0 {
|
||||
logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
|
||||
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active)
|
||||
e.peers[peerKey]--
|
||||
e.total--
|
||||
if e.peers[peerKey] == 0 {
|
||||
delete(e.peers, peerKey)
|
||||
}
|
||||
} else {
|
||||
logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey)
|
||||
}
|
||||
|
||||
// If the peer currently installed in WireGuard still holds references, nothing to reprogram.
|
||||
// Keying the check on the active peer (not the one just released) makes this self-healing:
|
||||
// a prior swap whose remove/add failed leaves e.active pointing at a peer with no references,
|
||||
// and this retries the hand-off on the next Decrement instead of getting stuck.
|
||||
if e.active != "" && e.peers[e.active] > 0 {
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
// Detach the stale/gone active peer from WireGuard before reprogramming.
|
||||
if e.active != "" {
|
||||
if err := rm.remove(prefix, e.active); err != nil {
|
||||
return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)
|
||||
}
|
||||
e.active = ""
|
||||
}
|
||||
|
||||
// Hand the prefix over to a surviving peer, or drop the entry when none remain.
|
||||
if survivor, ok := pickSurvivor(e.peers); ok {
|
||||
out, err := rm.add(prefix, survivor)
|
||||
if err != nil {
|
||||
return Ref[string]{Count: e.total, Out: ""}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err)
|
||||
}
|
||||
e.active = out
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
delete(rm.entries, prefix)
|
||||
return Ref[string]{Count: 0, Out: ""}, nil
|
||||
}
|
||||
|
||||
// Flush removes all prefixes from WireGuard and clears the counter.
|
||||
func (rm *AllowedIPsRefCounter) Flush() error {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
var merr *multierror.Error
|
||||
for prefix, e := range rm.entries {
|
||||
if e.active == "" {
|
||||
continue
|
||||
}
|
||||
logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active)
|
||||
if err := rm.remove(prefix, e.active); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err))
|
||||
}
|
||||
}
|
||||
|
||||
clear(rm.entries)
|
||||
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
// ReapplyMatching calls apply for every prefix whose currently installed (active) peer satisfies
|
||||
// pred, holding the lock for the whole pass. It is used to re-push allowed IPs onto a peer whose
|
||||
// WireGuard entry was rebuilt (e.g. a lazy connection cycling idle->wake) without a matching
|
||||
// refcounter change, which would otherwise leave the prefix installed in the counter but missing
|
||||
// on the device. Only the active peer is considered — a prefix that lost its installed peer to a
|
||||
// failed swap is skipped here and reconciled by the next Increment/Decrement.
|
||||
func (rm *AllowedIPsRefCounter) ReapplyMatching(pred func(out string) bool, apply func(key netip.Prefix) error) error {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
var merr *multierror.Error
|
||||
for prefix, e := range rm.entries {
|
||||
if e.active != "" && pred(e.active) {
|
||||
if err := apply(prefix); err != nil {
|
||||
merr = multierror.Append(merr, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do
|
||||
// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable
|
||||
// (lowest peerKey) for predictable behavior and testability.
|
||||
func pickSurvivor(peers map[string]int) (string, bool) {
|
||||
if len(peers) == 0 {
|
||||
return "", false
|
||||
}
|
||||
keys := make([]string, 0, len(peers))
|
||||
for k := range peers {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys[0], true
|
||||
}
|
||||
@@ -1,241 +0,0 @@
|
||||
package refcounter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/netip"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer.
|
||||
// failAdd/failRemove make the next add/remove fail once, to exercise the self-healing error paths.
|
||||
type fakeWG struct {
|
||||
installed map[netip.Prefix]string
|
||||
adds int
|
||||
removes int
|
||||
failAdd bool
|
||||
failRemove bool
|
||||
}
|
||||
|
||||
func newFakeWG() *fakeWG {
|
||||
return &fakeWG{installed: map[netip.Prefix]string{}}
|
||||
}
|
||||
|
||||
func (f *fakeWG) counter() *AllowedIPsRefCounter {
|
||||
return NewAllowedIPs(
|
||||
func(prefix netip.Prefix, peerKey string) (string, error) {
|
||||
if f.failAdd {
|
||||
f.failAdd = false
|
||||
return "", errors.New("add failed")
|
||||
}
|
||||
f.adds++
|
||||
f.installed[prefix] = peerKey
|
||||
return peerKey, nil
|
||||
},
|
||||
func(prefix netip.Prefix, peerKey string) error {
|
||||
if f.failRemove {
|
||||
f.failRemove = false
|
||||
return errors.New("remove failed")
|
||||
}
|
||||
f.removes++
|
||||
// only clear if this peer is the one installed, mirroring wg semantics
|
||||
if f.installed[prefix] == peerKey {
|
||||
delete(f.installed, prefix)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func mustPrefix(t *testing.T, s string) netip.Prefix {
|
||||
t.Helper()
|
||||
p, err := netip.ParsePrefix(s)
|
||||
if err != nil {
|
||||
t.Fatalf("parse prefix %q: %v", s, err)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func mustIncrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
|
||||
t.Helper()
|
||||
ref, err := c.Increment(p, peer)
|
||||
if err != nil {
|
||||
t.Fatalf("Increment(%v, %s): %v", p, peer, err)
|
||||
}
|
||||
return ref
|
||||
}
|
||||
|
||||
func mustDecrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
|
||||
t.Helper()
|
||||
ref, err := c.Decrement(p, peer)
|
||||
if err != nil {
|
||||
t.Fatalf("Decrement(%v, %s): %v", p, peer, err)
|
||||
}
|
||||
return ref
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same
|
||||
// prefix routed by different peers. Removing the network whose peer is installed must hand the
|
||||
// prefix over to the surviving peer instead of leaving it on the removed one.
|
||||
func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
// First peer wins while both are present.
|
||||
if got := f.installed[p]; got != "peerA" {
|
||||
t.Fatalf("expected peerA installed, got %q", got)
|
||||
}
|
||||
|
||||
// Remove the active peer's network -> must swap to peerB.
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if got := f.installed[p]; got != "peerB" {
|
||||
t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got)
|
||||
}
|
||||
|
||||
// Remove the last one -> prefix gone.
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("expected prefix removed, still installed on %q", f.installed[p])
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard.
|
||||
func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
removesBefore := f.removes
|
||||
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
if f.installed[p] != "peerA" {
|
||||
t.Fatalf("active peer must stay peerA, got %q", f.installed[p])
|
||||
}
|
||||
if f.removes != removesBefore {
|
||||
t.Fatalf("removing a non-active peer must not call wg remove")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until
|
||||
// the last reference is released (the reason the per-peer count must be an int, not a set).
|
||||
func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
if f.adds != 1 {
|
||||
t.Fatalf("expected a single wg add for the same peer, got %d", f.adds)
|
||||
}
|
||||
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if f.installed[p] != "peerA" {
|
||||
t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p])
|
||||
}
|
||||
if f.removes != 0 {
|
||||
t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes)
|
||||
}
|
||||
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("prefix must be removed after last reference")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log).
|
||||
func TestAllowedIPs_RefCountAndActive(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
ref := mustIncrement(t, c, p, "peerA")
|
||||
if ref.Count != 1 || ref.Out != "peerA" {
|
||||
t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out)
|
||||
}
|
||||
ref = mustIncrement(t, c, p, "peerB")
|
||||
if ref.Count != 2 || ref.Out != "peerA" {
|
||||
t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_Flush removes everything installed and clears the counter.
|
||||
func TestAllowedIPs_Flush(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p1 := mustPrefix(t, "10.44.8.0/24")
|
||||
p2 := mustPrefix(t, "10.44.9.0/24")
|
||||
|
||||
mustIncrement(t, c, p1, "peerA")
|
||||
mustIncrement(t, c, p2, "peerB")
|
||||
|
||||
if err := c.Flush(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(f.installed) != 0 {
|
||||
t.Fatalf("expected all prefixes removed, got %v", f.installed)
|
||||
}
|
||||
// After flush, a fresh increment must add again.
|
||||
mustIncrement(t, c, p1, "peerC")
|
||||
if f.installed[p1] != "peerC" {
|
||||
t.Fatalf("counter not reset after flush")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SelfHealAfterSwapAddError ensures a failed add during a swap does not permanently
|
||||
// strand the prefix: the next Decrement (or Increment) must retry and install a surviving peer.
|
||||
func TestAllowedIPs_SelfHealAfterSwapAddError(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
mustIncrement(t, c, p, "peerC")
|
||||
|
||||
// Removing the active peerA triggers a swap to a survivor; make the add fail once.
|
||||
f.failAdd = true
|
||||
if _, err := c.Decrement(p, "peerA"); err == nil {
|
||||
t.Fatalf("expected error from failed swap add")
|
||||
}
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("nothing should be installed after a failed swap add, got %q", f.installed[p])
|
||||
}
|
||||
|
||||
// A later Decrement of a non-active survivor must retry the hand-off (self-heal), not stay stuck.
|
||||
ref := mustDecrement(t, c, p, "peerC")
|
||||
if got := f.installed[p]; got == "" {
|
||||
t.Fatalf("self-heal failed: prefix left unrouted after add recovered")
|
||||
}
|
||||
if ref.Out == "" {
|
||||
t.Fatalf("expected an active peer after self-heal, got empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SelfHealAfterRemoveError ensures a failed remove during a swap is retried instead
|
||||
// of leaving e.active stuck on a peer that no longer holds references.
|
||||
func TestAllowedIPs_SelfHealAfterRemoveError(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
|
||||
// Releasing active peerA must detach it (remove) then add peerB; fail the remove once.
|
||||
f.failRemove = true
|
||||
if _, err := c.Decrement(p, "peerA"); err == nil {
|
||||
t.Fatalf("expected error from failed remove")
|
||||
}
|
||||
|
||||
// Next Decrement of the non-active survivor retries: removes stale peerA, installs peerB.
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
// peerB had only one ref, so after retry the prefix is fully released.
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("expected prefix released after self-heal, still on %q", f.installed[p])
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,5 @@ import "net/netip"
|
||||
// RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement
|
||||
type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}]
|
||||
|
||||
// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware:
|
||||
// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer,
|
||||
// so the counter records the per-peer reference count and swaps the installed peer when the active one is released.
|
||||
// See allowedips.go.
|
||||
// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement
|
||||
type AllowedIPsRefCounter = Counter[netip.Prefix, string, string]
|
||||
|
||||
@@ -15,11 +15,6 @@ type Route struct {
|
||||
route *route.Route
|
||||
routeRefCounter *refcounter.RouteRefCounter
|
||||
allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
|
||||
// currentPeerKey is the routing peer this watcher currently has the prefix installed on
|
||||
// (the HA winner elected by the watcher). It can differ from route.Peer and change on
|
||||
// failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement
|
||||
// the exact peer that was incremented.
|
||||
currentPeerKey string
|
||||
}
|
||||
|
||||
func NewRoute(params common.HandlerParams) *Route {
|
||||
@@ -57,15 +52,12 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
|
||||
ref.Out,
|
||||
)
|
||||
}
|
||||
r.currentPeerKey = peerKey
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Route) RemoveAllowedIPs() error {
|
||||
var err error
|
||||
if _, decErr := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); decErr != nil {
|
||||
err = fmt.Errorf("remove allowed IP %s: %w", r.route.Network, decErr)
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
|
||||
return err
|
||||
}
|
||||
r.currentPeerKey = ""
|
||||
return err
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1,12 +0,0 @@
|
||||
//go:build ios
|
||||
|
||||
package NetBirdSDK
|
||||
|
||||
import "github.com/netbirdio/netbird/version"
|
||||
|
||||
// GoClientVersion returns the NetBird Go client version that was baked into
|
||||
// the framework at compile time via
|
||||
// -ldflags "-X github.com/netbirdio/netbird/version.version=<version>".
|
||||
func GoClientVersion() string {
|
||||
return version.NetbirdVersion()
|
||||
}
|
||||
@@ -2,7 +2,6 @@ import { type ReactNode } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { Browser } from "@wailsio/runtime";
|
||||
import { DownloadIcon, NotepadText } from "lucide-react";
|
||||
import { Update as UpdateSvc } from "@bindings/services";
|
||||
import { Button } from "@/components/buttons/Button";
|
||||
import { useClientVersion } from "@/contexts/ClientVersionContext";
|
||||
import { cn } from "@/lib/cn";
|
||||
@@ -15,12 +14,6 @@ function openUrl(url: string) {
|
||||
});
|
||||
}
|
||||
|
||||
function openInstallerDownload() {
|
||||
UpdateSvc.DownloadURL()
|
||||
.then(openUrl)
|
||||
.catch(() => openUrl(GITHUB_RELEASES));
|
||||
}
|
||||
|
||||
export function UpdateVersionCard() {
|
||||
const { t } = useTranslation();
|
||||
const { updateVersion, enforced, triggerUpdate } = useClientVersion();
|
||||
@@ -44,7 +37,11 @@ export function UpdateVersionCard() {
|
||||
{t("update.card.installNow")}
|
||||
</Button>
|
||||
) : (
|
||||
<Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
|
||||
<Button
|
||||
variant={"primary"}
|
||||
size={"xs"}
|
||||
onClick={() => openUrl(GITHUB_RELEASES)}
|
||||
>
|
||||
<DownloadIcon size={14} />
|
||||
{t("update.card.getInstaller")}
|
||||
</Button>
|
||||
|
||||
@@ -281,9 +281,6 @@ func newApplication(onSecondInstance func()) *application.App {
|
||||
Linux: application.LinuxOptions{
|
||||
ProgramName: "netbird",
|
||||
},
|
||||
Windows: application.WindowsOptions{
|
||||
WndProcInterceptor: endSessionInterceptor(),
|
||||
},
|
||||
SingleInstance: &application.SingleInstanceOptions{
|
||||
UniqueID: "io.netbird.ui",
|
||||
OnSecondInstanceLaunch: func(_ application.SecondInstanceData) {
|
||||
@@ -370,9 +367,6 @@ func newMainWindow(app *application.App, prefStore *preferences.Store) *applicat
|
||||
|
||||
// Hide instead of quit on close; "really quit" is reached via tray -> Quit.
|
||||
window.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
|
||||
if services.ShuttingDown() {
|
||||
return
|
||||
}
|
||||
e.Cancel()
|
||||
window.Hide()
|
||||
})
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
package services
|
||||
|
||||
import "sync/atomic"
|
||||
|
||||
var (
|
||||
sessionEnding atomic.Bool
|
||||
quitting atomic.Bool
|
||||
)
|
||||
|
||||
func BeginSessionEnd() {
|
||||
sessionEnding.Store(true)
|
||||
}
|
||||
|
||||
func AbortSessionEnd() {
|
||||
sessionEnding.Store(false)
|
||||
}
|
||||
|
||||
func BeginShutdown() {
|
||||
quitting.Store(true)
|
||||
}
|
||||
|
||||
func ShuttingDown() bool {
|
||||
return sessionEnding.Load() || quitting.Load()
|
||||
}
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/client/proto"
|
||||
"github.com/netbirdio/netbird/client/ui/updater"
|
||||
"github.com/netbirdio/netbird/version"
|
||||
)
|
||||
|
||||
// UpdateResult mirrors TriggerUpdateResponse.
|
||||
@@ -34,12 +33,6 @@ func (s *Update) GetState() updater.State {
|
||||
return s.holder.Get()
|
||||
}
|
||||
|
||||
// DownloadURL returns the platform-appropriate installer download link for
|
||||
// manual (non-enforced) updates.
|
||||
func (s *Update) DownloadURL() string {
|
||||
return version.DownloadUrl()
|
||||
}
|
||||
|
||||
// Quit exits the app. Scheduled off the calling goroutine so the JS caller's
|
||||
// response returns before the runtime tears down.
|
||||
func (s *Update) Quit() {
|
||||
|
||||
@@ -154,9 +154,6 @@ func NewWindowManager(app *application.App, mainWindow *application.WebviewWindo
|
||||
})
|
||||
// Hide (not destroy) on close to keep React state; reset to General for a flash-free reopen.
|
||||
s.settings.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
|
||||
if ShuttingDown() {
|
||||
return
|
||||
}
|
||||
e.Cancel()
|
||||
s.app.Event.Emit(EventSettingsOpen, "general")
|
||||
s.settings.Hide()
|
||||
@@ -396,9 +393,6 @@ func (s *WindowManager) CloseWelcome() {
|
||||
// OpenError shows the custom error dialog; title/message are pre-localised and ride in the
|
||||
// start URL. A second error replaces the open one via SetURL. Singleton, destroyed on close.
|
||||
func (s *WindowManager) OpenError(title, message string) {
|
||||
if ShuttingDown() {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
startURL := errorDialogURL(title, message)
|
||||
|
||||
@@ -1,7 +0,0 @@
|
||||
//go:build !windows && !android && !ios && !freebsd && !js
|
||||
|
||||
package main
|
||||
|
||||
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
|
||||
return nil
|
||||
}
|
||||
@@ -1,36 +0,0 @@
|
||||
//go:build windows
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ui/services"
|
||||
)
|
||||
|
||||
const (
|
||||
wmQueryEndSession = 0x0011
|
||||
wmEndSession = 0x0016
|
||||
)
|
||||
|
||||
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
|
||||
return func(_ uintptr, msg uint32, wParam, _ uintptr) (uintptr, bool) {
|
||||
switch msg {
|
||||
case wmQueryEndSession:
|
||||
services.BeginSessionEnd()
|
||||
return 1, true
|
||||
case wmEndSession:
|
||||
if wParam == 0 {
|
||||
services.AbortSessionEnd()
|
||||
return 0, true
|
||||
}
|
||||
log.Info("windows session is ending; exiting immediately")
|
||||
os.Exit(0)
|
||||
return 0, true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -32,8 +32,9 @@ const (
|
||||
|
||||
quitDownTimeout = 5 * time.Second
|
||||
|
||||
urlGitHubRepo = "https://github.com/netbirdio/netbird"
|
||||
urlDocs = "https://docs.netbird.io"
|
||||
urlGitHubRepo = "https://github.com/netbirdio/netbird"
|
||||
urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
|
||||
urlDocs = "https://docs.netbird.io"
|
||||
)
|
||||
|
||||
// TrayServices bundles the services the tray menu needs, grouped so NewTray
|
||||
@@ -452,7 +453,6 @@ func (t *Tray) buildMenu() *application.Menu {
|
||||
}
|
||||
|
||||
func (t *Tray) handleQuit() {
|
||||
services.BeginShutdown()
|
||||
t.profileMu.Lock()
|
||||
if t.switchCancel != nil {
|
||||
t.switchCancel()
|
||||
|
||||
@@ -25,9 +25,6 @@ type sendFn func(notifications.NotificationOptions) error
|
||||
// event-dispatch goroutine that panic is fatal process-wide; recover() turns
|
||||
// it into a logged no-op.
|
||||
func safeSendNotification(send sendFn, what string, opts notifications.NotificationOptions) (err error) {
|
||||
if services.ShuttingDown() {
|
||||
return nil
|
||||
}
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Errorf("notify %s: recovered from panic (notification bus unavailable): %v", what, r)
|
||||
|
||||
@@ -13,7 +13,6 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/client/ui/services"
|
||||
"github.com/netbirdio/netbird/client/ui/updater"
|
||||
"github.com/netbirdio/netbird/version"
|
||||
)
|
||||
|
||||
// trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
|
||||
@@ -77,15 +76,15 @@ func (u *trayUpdater) applyLanguage() {
|
||||
u.refreshMenuItem(state)
|
||||
}
|
||||
|
||||
// handleClick opens the installer download link when not Enforced, otherwise
|
||||
// shows the progress page and asks the daemon to start the installer.
|
||||
// handleClick opens the GitHub releases page when not Enforced, otherwise shows
|
||||
// the progress page and asks the daemon to start the installer.
|
||||
func (u *trayUpdater) handleClick() {
|
||||
u.mu.Lock()
|
||||
state := u.state
|
||||
u.mu.Unlock()
|
||||
|
||||
if !state.Enforced {
|
||||
_ = u.app.Browser.OpenURL(version.DownloadUrl())
|
||||
_ = u.app.Browser.OpenURL(urlGitHubReleases)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -109,7 +109,7 @@ sequenceDiagram
|
||||
Chk->>Inj: continue
|
||||
Inj->>Inj: inject NetBird identity headers per provider config
|
||||
Inj->>Grd: continue
|
||||
Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
|
||||
Grd->>Grd: enforce model allowlist
|
||||
Grd->>Up: forward (over WireGuard)
|
||||
Up-->>Resp: response (JSON or SSE stream)
|
||||
Resp->>Resp: parse usage tokens, completion
|
||||
@@ -135,21 +135,6 @@ sequenceDiagram
|
||||
(`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
|
||||
PII names — see `redact.go` for the full set. See
|
||||
[`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
|
||||
- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
|
||||
is authoritative: it resolves the policy that governs this
|
||||
(provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
|
||||
when no applicable policy permits the model — so an allowlist scoped to
|
||||
one group/provider never leaks to another, and an un-guardrailed policy
|
||||
is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
|
||||
backstop: it only carries an allowlist for a provider every authorising
|
||||
policy restricts, and blocks unknown/undetermined models even when
|
||||
management is unreachable. Because that backstop allowlist is the UNION
|
||||
of every restricting policy's models, per-group narrowing lives only in
|
||||
the authoritative check: during a `CheckLLMPolicyLimits` outage
|
||||
`llm_limit_check` fails open, so a caller can reach any model in the
|
||||
provider's union — a group scoped to model A could reach model B if
|
||||
another group restricts the same provider to B. This is the documented
|
||||
fail-open trade-off; a future flag may switch it to fail-closed.
|
||||
- SSE streaming requires special handling on the response side; the
|
||||
parser must handle partial chunks without buffering the whole
|
||||
stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).
|
||||
|
||||
@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
| on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
|
||||
| on_request | 2 | `llm_limit_check` | `{}` | – |
|
||||
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
|
||||
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
|
||||
| on_response | 6 | `cost_meter` | `{}` | – |
|
||||
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |
|
||||
|
||||
@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
|
||||
| `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
|
||||
| `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
|
||||
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
|
||||
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
|
||||
| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
|
||||
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
|
||||
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
|
||||
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |
|
||||
|
||||
@@ -1,209 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
|
||||
// provider, its own guardrail whose allowlist holds ONLY that provider's
|
||||
// allowed model, and its own policy. Each case runs in isolation (its own
|
||||
// proxy + client), so the guardrail the proxy enforces contains exactly this
|
||||
// provider's model — never a mixed cross-provider list.
|
||||
type pathRoutedGuardrailCase struct {
|
||||
name string
|
||||
catalogID string // agent-network catalog provider id
|
||||
wire string // harness.WireVertex | harness.WireBedrock
|
||||
allowEntry string // the single model id put on the guardrail allowlist
|
||||
allowModel string // model id sent that MUST be served (200)
|
||||
blockModel string // model id sent that MUST be denied (403 model_blocked)
|
||||
}
|
||||
|
||||
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
|
||||
// guard for the customer report that a model-allowlist guardrail attached to a
|
||||
// policy has no effect for PATH-ROUTED providers — where the model travels in
|
||||
// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
|
||||
// AWS Bedrock (/model/{id}/invoke).
|
||||
//
|
||||
// Each provider is tested in isolation with a guardrail allowlisting a single
|
||||
// model of its own: the allowed model (in the URL path) is served (200) and an
|
||||
// unselected model (in the URL path) is denied 403 by the guardrail
|
||||
// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
|
||||
// customer verbatim — allow Sonnet, and the unselected model is the exact
|
||||
// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
|
||||
// sends a region-prefixed, versioned inference-profile id so URL-path model
|
||||
// normalization is exercised too.
|
||||
//
|
||||
// The provider is catch-all (no models), so the router forwards any model and a
|
||||
// 403 can only come from the guardrail, never model_not_routable. Only the
|
||||
// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
|
||||
// synth/reconcile, the proxy middleware chain (URL-path model extraction,
|
||||
// router, guardrail) and the tunnel are all real, and the guardrail denies
|
||||
// before the upstream is dialed so the mock cannot influence the block. A
|
||||
// static bearer api key is used so the router injects a static Authorization
|
||||
// header instead of minting a GCP token — the only reason path-routed providers
|
||||
// normally need live credentials — so the test runs with none and is always on.
|
||||
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
|
||||
cases := []pathRoutedGuardrailCase{
|
||||
{
|
||||
name: "vertex",
|
||||
catalogID: "vertex_ai_api",
|
||||
wire: harness.WireVertex,
|
||||
allowEntry: "claude-sonnet-4-5",
|
||||
allowModel: "claude-sonnet-4-5",
|
||||
blockModel: "claude-opus-4-6", // the customer-reported model
|
||||
},
|
||||
{
|
||||
name: "bedrock",
|
||||
catalogID: "bedrock_api",
|
||||
wire: harness.WireBedrock,
|
||||
allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
|
||||
allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
|
||||
blockModel: "us.anthropic.claude-opus-4-8-v1:0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
runPathRoutedGuardrailCase(t, tc)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
|
||||
t.Helper()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-" + tc.name + "-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// Catch-all provider (no models) so the router forwards any model; a static
|
||||
// bearer key means the router injects a static auth header instead of minting
|
||||
// a GCP token. Bootstraps the cluster if it isn't already.
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: tc.name,
|
||||
ProviderId: tc.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create %s provider", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// Guardrail allowlisting ONLY this provider's allowed model.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-guardrail-" + tc.name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-" + tc.name,
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Probe first: the GET resolves the endpoint and its first packet wakes the
|
||||
// lazy proxy peer, so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch tc.wire {
|
||||
case harness.WireVertex:
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case harness.WireBedrock:
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
t.Fatalf("unsupported wire %q", tc.wire)
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
|
||||
// jitter on the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(tc.allowModel)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
|
||||
// Unselected model (in the URL path) must be blocked by the guardrail.
|
||||
code, body = send(tc.blockModel)
|
||||
assert.Equal(t, 403, code,
|
||||
"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
|
||||
}
|
||||
@@ -1,205 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
|
||||
// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
|
||||
// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
|
||||
// the proxy caches; the switch is invisible until that cache expires. The proxy
|
||||
// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
|
||||
// instead of the 5-minute default.
|
||||
//
|
||||
// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
|
||||
// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
|
||||
// served and modelB denied; after switching the client grpA -> grpB, modelB is
|
||||
// served and modelA denied. The cross-group deny comes from management's
|
||||
// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
|
||||
func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelA = "e2e-model-a"
|
||||
modelB = "e2e-model-b"
|
||||
// Short tunnel-cache TTL so a group switch propagates in seconds.
|
||||
// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
|
||||
cacheTTL = 3 * time.Second
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
|
||||
require.NoError(t, err, "create group A")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
|
||||
|
||||
grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
|
||||
require.NoError(t, err, "create group B")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-gswitch-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpA.Id}, // client starts in group A
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "gswitch",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
},
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
mkGuard := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gA := mkGuard("e2e-gswitch-a", modelA)
|
||||
gB := mkGuard("e2e-gswitch-b", modelB)
|
||||
|
||||
enabled := true
|
||||
polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-a",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpA.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gA.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy A")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
|
||||
|
||||
polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-b",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpB.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gB.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy B")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
|
||||
"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
|
||||
})
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == want {
|
||||
return code, body
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Phase 1 — client is in group A: modelA served, modelB denied.
|
||||
code, body := sendUntil(modelA, 200, 90*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelB)
|
||||
assert.Equal(t, 403, code,
|
||||
"group-B model must be denied while the client is in group A; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
|
||||
// Switch the client peer from group A to group B.
|
||||
peerID := clientPeerInGroup(t, ctx, grpA.Id)
|
||||
_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpB.Name,
|
||||
Peers: &[]string{peerID},
|
||||
})
|
||||
require.NoError(t, err, "add peer to group B")
|
||||
_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpA.Name,
|
||||
Peers: &[]string{},
|
||||
})
|
||||
require.NoError(t, err, "remove peer from group A")
|
||||
|
||||
// Phase 2 — after the short TTL expires the proxy re-validates the peer,
|
||||
// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
|
||||
code, body = sendUntil(modelB, 200, 60*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelA)
|
||||
assert.Equal(t, 403, code,
|
||||
"after the switch, the old group-A model must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
}
|
||||
|
||||
// clientPeerInGroup returns the id of the single peer that is a member of the
|
||||
// given group — the test client. The proxy peer is never added to test groups.
|
||||
func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
|
||||
t.Helper()
|
||||
peers, err := srv.API().Peers.List(ctx)
|
||||
require.NoError(t, err, "list peers")
|
||||
for _, p := range peers {
|
||||
for _, g := range p.Groups {
|
||||
if g.Id == groupID {
|
||||
return p.Id
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Fatalf("no peer found in group %s", groupID)
|
||||
return ""
|
||||
}
|
||||
@@ -1,201 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
|
||||
// modelOther denied for the grpMain client (403 model_blocked, no cross-group
|
||||
// leak), and openModel on the un-guardrailed policy's provider served (200).
|
||||
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelSelected = "e2e-selected"
|
||||
modelOther = "e2e-other"
|
||||
openModel = "e2e-open"
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-mp-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// pRestricted declares the two guardrailed models so routing is deterministic
|
||||
// (model -> provider). Created first, so it carries the bootstrap cluster.
|
||||
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "restricted",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(modelSelected, modelOther),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create restricted provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
|
||||
|
||||
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "open",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(openModel),
|
||||
})
|
||||
require.NoError(t, err, "create open provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
|
||||
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
|
||||
|
||||
enabled := true
|
||||
// polMain: grpMain restricted to modelSelected on pRestricted.
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// polOther: grpOther restricted to modelOther on the SAME provider. The
|
||||
// client is not in grpOther, so modelOther must never be usable by it.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
|
||||
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-open",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pOpen.Id},
|
||||
})
|
||||
require.NoError(t, err, "create open policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("selected model allowed for its group", func(t *testing.T) {
|
||||
code, body := sendUntil200(modelSelected)
|
||||
assert.Equal(t, 200, code,
|
||||
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
|
||||
t.Run("other group's model does not leak", func(t *testing.T) {
|
||||
// modelOther is allowlisted only for grpOther. The grpMain client must be
|
||||
// denied by management's per-policy/group check — not waved through by an
|
||||
// account-wide union. This is the security-critical wrong-ALLOW guard.
|
||||
code, body := send(modelOther)
|
||||
assert.Equal(t, 403, code,
|
||||
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"denial must be a model-allowlist decision; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
|
||||
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
|
||||
// old account-wide union would have blocked openModel (it is on no
|
||||
// allowlist); it must now be served — the false-DENY guard.
|
||||
code, body := sendUntil200(openModel)
|
||||
assert.Equal(t, 200, code,
|
||||
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
@@ -1,422 +0,0 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pergroupCase describes one provider surface for the per-group allowlist matrix.
|
||||
// selectedReq/otherReq are the model identifiers as they travel in the request
|
||||
// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
|
||||
// otherAllow are the (normalized) forms the guardrail allowlist holds — for
|
||||
// Bedrock these differ from the request form so path normalization is exercised.
|
||||
type pergroupCase struct {
|
||||
name string
|
||||
catalogID string
|
||||
wire string // "chat", "messages", "vertex", "bedrock"
|
||||
models *[]api.AgentNetworkProviderModel
|
||||
selectedReq string
|
||||
selectedAllow string
|
||||
otherReq string
|
||||
otherAllow string
|
||||
|
||||
providerID string // filled during setup
|
||||
}
|
||||
|
||||
// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
|
||||
// allowlist end to end across every always-on provider surface, including the
|
||||
// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
|
||||
//
|
||||
// For each provider two policies target it: grpMain (the client) is allowed only
|
||||
// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
|
||||
// The client must get selectedReq served (200) and otherReq denied (403,
|
||||
// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
|
||||
// authoritative per-policy/group decision from management (the proxy per-provider
|
||||
// backstop carries the union of both models), so this also confirms management
|
||||
// receives the correct normalized model for path-routed providers.
|
||||
func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
cases := []*pergroupCase{
|
||||
{
|
||||
name: "openai", catalogID: "openai_api", wire: harness.WireChat,
|
||||
models: priced("oai-model-a", "oai-model-b"),
|
||||
selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
|
||||
otherReq: "oai-model-b", otherAllow: "oai-model-b",
|
||||
},
|
||||
{
|
||||
name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
|
||||
models: priced("ant-model-a", "ant-model-b"),
|
||||
selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
|
||||
otherReq: "ant-model-b", otherAllow: "ant-model-b",
|
||||
},
|
||||
{
|
||||
// Vertex catalog ids travel bare in the rawPredict path.
|
||||
name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
|
||||
selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
|
||||
otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
|
||||
},
|
||||
{
|
||||
// Bedrock request ids are region-prefixed/versioned; the parser
|
||||
// normalizes them to the catalog key the allowlist holds.
|
||||
name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
|
||||
selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
|
||||
otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
|
||||
},
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-pergroup-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
for i, c := range cases {
|
||||
req := api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-pergroup-" + c.name,
|
||||
ProviderId: c.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: c.models,
|
||||
}
|
||||
if i == 0 {
|
||||
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
|
||||
}
|
||||
prov, perr := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, perr, "create provider %s", c.name)
|
||||
c.providerID = prov.Id
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
|
||||
gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
|
||||
|
||||
polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gSel.Id},
|
||||
})
|
||||
require.NoError(t, merr, "create main policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gOth.Id},
|
||||
})
|
||||
require.NoError(t, oerr, "create other policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
}
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(c *pergroupCase, model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch c.wire {
|
||||
case "vertex":
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case "bedrock":
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
|
||||
// the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(c, c.selectedReq)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
|
||||
// grpOther's model must NOT leak to the grpMain client.
|
||||
code, body = send(c, c.otherReq)
|
||||
assert.Equal(t, 403, code,
|
||||
"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
|
||||
// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
|
||||
//
|
||||
// - union across the user's groups: the client is in gUX and gUY, each with
|
||||
// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
|
||||
// client may use BOTH models (the union of its groups' allowlists) while a
|
||||
// third, un-allowlisted model is denied.
|
||||
// - an un-guardrailed group lifts the restriction: the client is in gMP and
|
||||
// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
|
||||
// NO guardrail. Because one applicable policy is unrestricted, the client may
|
||||
// use a model on no allowlist (mix-z) as well as mix-a.
|
||||
func TestGuardrailMultiGroupUser(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
unionA = "mg-union-a"
|
||||
unionB = "mg-union-b"
|
||||
unionC = "mg-union-c" // allowlisted by neither group
|
||||
mixA = "mg-mix-a"
|
||||
mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
mkGroup := func(name string) *api.Group {
|
||||
g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
|
||||
require.NoError(t, gerr, "create group %s", name)
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gUX := mkGroup("e2e-mg-union-x")
|
||||
gUY := mkGroup("e2e-mg-union-y")
|
||||
gMP := mkGroup("e2e-mg-mix-p")
|
||||
gMQ := mkGroup("e2e-mg-mix-q")
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-mg-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
// P1 — union scenario: two restricting policies, one per group.
|
||||
p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-union",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(unionA, unionB, unionC),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create union provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
|
||||
|
||||
polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-x",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUX.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy X")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
|
||||
|
||||
polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-y",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUY.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy Y")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
|
||||
|
||||
// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
|
||||
p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-mix",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(mixA, mixZ),
|
||||
})
|
||||
require.NoError(t, err, "create mix provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
|
||||
|
||||
polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-p",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMP.Id},
|
||||
DestinationProviderIds: []string{p2.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create mix policy P")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
|
||||
|
||||
polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-q",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMQ.Id},
|
||||
DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
|
||||
})
|
||||
require.NoError(t, err, "create mix policy Q")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("union across the user's groups", func(t *testing.T) {
|
||||
code, body := sendUntil200(unionA)
|
||||
assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(unionB)
|
||||
assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
|
||||
|
||||
code, body = send(unionC)
|
||||
assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
})
|
||||
|
||||
t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
|
||||
code, body := sendUntil200(mixA)
|
||||
assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(mixZ)
|
||||
assert.Equal(t, 200, code,
|
||||
"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
|
||||
// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
|
||||
// exactly the given model, registering cleanup.
|
||||
func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
|
||||
t.Helper()
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
@@ -38,11 +38,7 @@ type Proxy struct {
|
||||
// network, registered via the given account proxy token and serving the
|
||||
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
|
||||
// peer connectivity — callers poll management for the proxy peer.
|
||||
// StartProxy launches the reverse-proxy container. Optional envOverrides are
|
||||
// merged into the container environment after the defaults, so callers can set
|
||||
// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
|
||||
// need a short authorization-cache window).
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
|
||||
root, err := repoRoot()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -97,12 +93,6 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverride
|
||||
WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
|
||||
}
|
||||
|
||||
for _, ov := range envOverrides {
|
||||
for k, v := range ov {
|
||||
req.Env[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
ContainerRequest: req,
|
||||
Started: true,
|
||||
|
||||
@@ -86,17 +86,12 @@ type Manager interface {
|
||||
|
||||
// PolicySelectionInput is the per-request selection envelope. The
|
||||
// proxy populates it from CapturedData (account, user, groups) plus
|
||||
// the provider llm_router resolved and the model it extracted.
|
||||
// the provider llm_router resolved.
|
||||
type PolicySelectionInput struct {
|
||||
AccountID string
|
||||
UserID string
|
||||
GroupIDs []string
|
||||
ProviderID string
|
||||
// Model is the already-normalised upstream model id the proxy extracted
|
||||
// (parser strips Bedrock region/version, Vertex @version), so a
|
||||
// case-insensitive compare suffices. Empty = undetermined → not permitted
|
||||
// (fail closed).
|
||||
Model string
|
||||
}
|
||||
|
||||
// PolicySelectionResult names the policy that "pays" for this request
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
@@ -36,10 +35,6 @@ const (
|
||||
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
|
||||
//nolint:gosec // account deny code label, not a credential
|
||||
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
|
||||
// denyCodeModelBlocked is returned when policies govern the request's
|
||||
// (provider, caller-groups) but none permits the model. Matches the proxy
|
||||
// guardrail's code so both layers surface the same label.
|
||||
denyCodeModelBlocked = "llm_policy.model_blocked"
|
||||
)
|
||||
|
||||
// consumptionCache holds the consumption counters prefetched for one
|
||||
@@ -164,25 +159,6 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
|
||||
}
|
||||
candidates := filterApplicablePolicies(policies, in)
|
||||
|
||||
// Model-allowlist gate scoped to the matched policies: keep candidates whose
|
||||
// guardrails permit the model (none enabled = unrestricted), deny when
|
||||
// policies apply but none permits it. Skip the load when none has a guardrail.
|
||||
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
|
||||
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if gErr != nil {
|
||||
return nil, gErr
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
|
||||
if len(permitted) == 0 {
|
||||
return &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
candidates = permitted
|
||||
}
|
||||
|
||||
// Prefetch every consumption counter the ceiling + candidate policies will
|
||||
// read, in a single store round-trip, then score against the cache.
|
||||
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
|
||||
@@ -274,90 +250,6 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
|
||||
return out
|
||||
}
|
||||
|
||||
// anyPolicyHasGuardrails reports whether any policy references at least one
|
||||
// guardrail, so the selector can skip loading guardrails when none do.
|
||||
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
|
||||
for _, p := range policies {
|
||||
if p != nil && len(p.GuardrailIDs) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
|
||||
// model-allowlist gate to resolve each candidate policy's attached guardrails.
|
||||
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
|
||||
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list account guardrails: %w", err)
|
||||
}
|
||||
byID := make(map[string]*types.Guardrail, len(guardrails))
|
||||
for _, g := range guardrails {
|
||||
if g != nil {
|
||||
byID[g.ID] = g
|
||||
}
|
||||
}
|
||||
return byID, nil
|
||||
}
|
||||
|
||||
// filterModelPermittedPolicies returns the subset of policies whose guardrails
|
||||
// permit the model. Order is preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if policyPermitsModel(p, byID, model) {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyPermitsModel reports whether a policy permits the model. No
|
||||
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
|
||||
// otherwise the model must be in the union of its allowlists, so an
|
||||
// empty/undetermined model fails closed.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
wanted := normaliseModelID(model)
|
||||
restricted := false
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
if wanted == "" {
|
||||
continue
|
||||
}
|
||||
for _, allowed := range g.Checks.ModelAllowlist.Models {
|
||||
if normaliseModelID(allowed) == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return !restricted
|
||||
}
|
||||
|
||||
// normaliseModelID lowercases and trims a model identifier so the allowlist
|
||||
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
|
||||
// normaliseModel so both layers agree on what "same model" means.
|
||||
func normaliseModelID(model string) string {
|
||||
return strings.ToLower(strings.TrimSpace(model))
|
||||
}
|
||||
|
||||
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
|
||||
// rejection. The model is quoted when known; an undetermined model is reported
|
||||
// as such so the access log distinguishes "wrong model" from "no model".
|
||||
func modelBlockedReason(model string) string {
|
||||
if normaliseModelID(model) == "" {
|
||||
return "request model could not be determined for the policy allowlist"
|
||||
}
|
||||
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
|
||||
}
|
||||
|
||||
// candidate is the per-policy intermediate the selector ranks. A
|
||||
// policy that's been exhausted on any enabled cap never makes it
|
||||
// into this slice; the selector's deny envelope carries the latest
|
||||
|
||||
@@ -1,329 +0,0 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
|
||||
// to reach providerID under the given guardrails. Uncapped keeps the selector's
|
||||
// headroom scoring trivial so these tests isolate the model-allowlist gate.
|
||||
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Enabled: true,
|
||||
SourceGroups: sourceGroups,
|
||||
DestinationProviderIDs: []string{providerID},
|
||||
GuardrailIDs: guardrailIDs,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
}
|
||||
|
||||
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
|
||||
// carries the given models.
|
||||
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
|
||||
return &types.Guardrail{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
|
||||
Return(policies, nil)
|
||||
}
|
||||
|
||||
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
|
||||
Return(guardrails, nil)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
|
||||
// decision: a policy authorises the (provider, group) but restricts the model,
|
||||
// and the requested model isn't on the list, so the request is denied.
|
||||
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
|
||||
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
|
||||
// is on the applicable policy's allowlist, so selection proceeds normally.
|
||||
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
|
||||
// and surrounding whitespace, matching the proxy guardrail's normalisation.
|
||||
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
|
||||
// policies authorise the same (provider, group) and one has no guardrail, that
|
||||
// policy makes the request unrestricted — not caught by the other's allowlist.
|
||||
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
|
||||
expectPolicies(mockStore, "acc-1", restricted, open)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
|
||||
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
|
||||
// model allowlisted only for grp-b must not be usable by a grp-a caller. The
|
||||
// selector considers only policies applicable to the caller's groups.
|
||||
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
|
||||
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
|
||||
expectPolicies(mockStore, "acc-1", polA, polB)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-a"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only allowed for grp-b
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
|
||||
// mirrors the proxy: with a restricted applicable policy and an empty model
|
||||
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
|
||||
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "", // undetermined
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
|
||||
// model allowlist is disabled imposes no model restriction, even though the
|
||||
// policy references it.
|
||||
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
disabled := &types.Guardrail{
|
||||
ID: "g-1",
|
||||
AccountID: "acc-1",
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
|
||||
},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", disabled)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
|
||||
// allowlist guardrails permits the union of their models (not just the first).
|
||||
func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only in the second guardrail's list
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
|
||||
// resolving the candidate policies' guardrails surfaces as an error, not a
|
||||
// silent allow/deny.
|
||||
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
|
||||
Return(nil, errors.New("store unavailable"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
|
||||
assert.Nil(t, res)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
|
||||
// policy referencing a guardrail ID absent from the account's set (a stale
|
||||
// reference) imposes no model restriction — same as no guardrail.
|
||||
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
|
||||
// gate narrows candidates before cap scoring: the permitting policy is selected
|
||||
// even though the blocked one has a larger, more attractive cap.
|
||||
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
|
||||
polBig.Limits = types.PolicyLimits{
|
||||
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
|
||||
}
|
||||
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
|
||||
polSmall.Limits = types.PolicyLimits{
|
||||
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", polBig, polSmall)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only pol-small's guardrail permits this
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow)
|
||||
assert.Equal(t, "pol-small", res.SelectedPolicyID,
|
||||
"the model filter must exclude pol-big before cap scoring")
|
||||
}
|
||||
@@ -235,12 +235,7 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in this map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
|
||||
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -785,12 +780,10 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
|
||||
|
||||
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
|
||||
// middleware expects. Mirrors the proxy registration documented in
|
||||
// the management→proxy contract. provider_allowlists is keyed by the
|
||||
// resolved provider id llm_router stamps; a provider absent from the map is
|
||||
// unrestricted at the proxy layer.
|
||||
// the management→proxy contract.
|
||||
type guardrailConfig struct {
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
type guardrailPromptCapture struct {
|
||||
@@ -835,10 +828,13 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
|
||||
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
|
||||
}
|
||||
|
||||
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
|
||||
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
cfg := guardrailConfig{
|
||||
ProviderAllowlists: providerAllowlists,
|
||||
PromptCapture: guardrailPromptCapture(capture),
|
||||
ModelAllowlist: merged.ModelAllowlist,
|
||||
PromptCapture: guardrailPromptCapture{
|
||||
Enabled: merged.PromptCapture.Enabled,
|
||||
RedactPii: merged.PromptCapture.RedactPii,
|
||||
},
|
||||
}
|
||||
out, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
@@ -847,74 +843,6 @@ func marshalGuardrailConfig(providerAllowlists map[string][]string, capture Merg
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
|
||||
// is included only when every authorising policy restricts models (their union);
|
||||
// if any leaves it unrestricted it is omitted, so management decides per group.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
|
||||
type providerAcc struct {
|
||||
models map[string]struct{}
|
||||
anyUnrestricted bool
|
||||
}
|
||||
accs := make(map[string]*providerAcc)
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
acc, ok := accs[providerID]
|
||||
if !ok {
|
||||
acc = &providerAcc{models: make(map[string]struct{})}
|
||||
accs[providerID] = acc
|
||||
}
|
||||
if !restricted {
|
||||
acc.anyUnrestricted = true
|
||||
continue
|
||||
}
|
||||
for _, m := range models {
|
||||
acc.models[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
out := make(map[string][]string, len(accs))
|
||||
for providerID, acc := range accs {
|
||||
if acc.anyUnrestricted {
|
||||
continue
|
||||
}
|
||||
models := make([]string, 0, len(acc.models))
|
||||
for m := range acc.models {
|
||||
models = append(models, m)
|
||||
}
|
||||
sort.Strings(models)
|
||||
out[providerID] = models
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyModelAllowlist reports whether a policy restricts models (has an
|
||||
// allowlist-enabled guardrail) and the union of allowed models. Models are
|
||||
// verbatim; the proxy factory lowercases/trims them at decode time.
|
||||
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := false
|
||||
var models []string
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
models = append(models, m)
|
||||
}
|
||||
}
|
||||
}
|
||||
return restricted, models
|
||||
}
|
||||
|
||||
// buildAccountService composes the per-account gateway Service. The
|
||||
// target carries the noop placeholder URL — the router middleware
|
||||
// rewrites every request to the matched provider's upstream before the
|
||||
@@ -1058,11 +986,38 @@ func unionSourceGroups(policies []*types.Policy) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
|
||||
// merged here — the model allowlist is emitted per-provider, and
|
||||
// token/budget/retention moved onto Policy.Limits and account Settings.
|
||||
// MergedGuardrails is the JSON shape passed to the proxy via the
|
||||
// guardrail middleware's config_json. Mirrors the proxy-side
|
||||
// expectations and is intentionally distinct from
|
||||
// types.GuardrailChecks so we can evolve either side independently.
|
||||
type MergedGuardrails struct {
|
||||
PromptCapture MergedPromptCapture
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
}
|
||||
|
||||
type MergedTokenLimits struct {
|
||||
Hourly *MergedTokenWindow `json:"hourly,omitempty"`
|
||||
Daily *MergedTokenWindow `json:"daily,omitempty"`
|
||||
Monthly *MergedTokenWindow `json:"monthly,omitempty"`
|
||||
}
|
||||
|
||||
type MergedTokenWindow struct {
|
||||
MaxInputTokens int `json:"max_input_tokens,omitempty"`
|
||||
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
|
||||
}
|
||||
|
||||
type MergedBudget struct {
|
||||
Hourly *MergedBudgetWindow `json:"hourly,omitempty"`
|
||||
Daily *MergedBudgetWindow `json:"daily,omitempty"`
|
||||
Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
|
||||
}
|
||||
|
||||
type MergedBudgetWindow struct {
|
||||
SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
|
||||
HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
|
||||
}
|
||||
|
||||
type MergedPromptCapture struct {
|
||||
@@ -1070,31 +1025,64 @@ type MergedPromptCapture struct {
|
||||
RedactPii bool `json:"redact_pii"`
|
||||
}
|
||||
|
||||
// mergeGuardrails folds the referencing policies' guardrails into the
|
||||
// prompt-capture decision only. The model allowlist is enforced per-policy/group
|
||||
// in management and shipped per-provider; token/budget/retention live off
|
||||
// guardrails now.
|
||||
type MergedRetention struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Days int `json:"days"`
|
||||
}
|
||||
|
||||
// mergeGuardrails computes the effective guardrail spec applied at the
|
||||
// proxy, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only
|
||||
// enabled policies should be passed in.
|
||||
//
|
||||
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
|
||||
// sticks if any enabling policy turns it on.
|
||||
// Merge rules:
|
||||
// - Model allowlist: union of allowlists across policies that enable it.
|
||||
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
|
||||
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
|
||||
// if any enabling policy turns it on.
|
||||
// - Retention: enabled if any enables it; smallest non-zero days wins.
|
||||
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
|
||||
merged := MergedGuardrails{}
|
||||
allowlist := make(map[string]struct{})
|
||||
allowlistEnabled := false
|
||||
|
||||
for _, policy := range policies {
|
||||
for _, gID := range policy.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil {
|
||||
continue
|
||||
}
|
||||
mergeGuardrail(g, &merged)
|
||||
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
|
||||
}
|
||||
}
|
||||
|
||||
if allowlistEnabled {
|
||||
merged.ModelAllowlist = make([]string, 0, len(allowlist))
|
||||
for m := range allowlist {
|
||||
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
|
||||
}
|
||||
sort.Strings(merged.ModelAllowlist)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
|
||||
// running merge: enabled / redact-pii stick once any enabling guardrail turns
|
||||
// them on.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
// mergeGuardrail folds a single guardrail's enabled checks into the
|
||||
// running merge: model-allowlist models join the shared set (and flip
|
||||
// allowlistEnabled), and prompt-capture / redact-pii stick once any
|
||||
// enabling guardrail turns them on.
|
||||
//
|
||||
// TokenLimits, Budget, and Retention have moved off guardrails — token
|
||||
// and budget caps now live on the Policy itself (Policy.Limits) and
|
||||
// retention moves to account-level Settings — so they are not merged here.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
|
||||
if g.Checks.ModelAllowlist.Enabled {
|
||||
*allowlistEnabled = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
allowlist[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
if g.Checks.PromptCapture.Enabled {
|
||||
merged.PromptCapture.Enabled = true
|
||||
if g.Checks.PromptCapture.RedactPii {
|
||||
|
||||
@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
|
||||
require.Len(t, services, 1)
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
|
||||
"model allowlist is a pure policy guardrail and must reach the per-provider config")
|
||||
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
|
||||
"model allowlist is a pure policy guardrail and must always reach the config")
|
||||
assert.False(t, cfg.PromptCapture.Enabled,
|
||||
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
|
||||
}
|
||||
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
|
||||
require.Len(t, services, 1, "exactly one synth service expected")
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
|
||||
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
|
||||
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
|
||||
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
|
||||
}
|
||||
|
||||
@@ -1,95 +0,0 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// policyForProviders builds an enabled policy authorising the given providers
|
||||
// under the given guardrails (both optional). Groups are irrelevant to
|
||||
// buildProviderAllowlists, which keys purely on destination provider.
|
||||
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildProviderAllowlists(t *testing.T) {
|
||||
byID := map[string]*types.Guardrail{
|
||||
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
|
||||
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
|
||||
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
|
||||
}
|
||||
|
||||
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
|
||||
"a provider every policy restricts carries the sorted union of their models")
|
||||
})
|
||||
|
||||
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", nil, "prov-x"), // no guardrail
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a policy whose only guardrail has a disabled allowlist is unrestricted")
|
||||
})
|
||||
|
||||
t.Run("providers are isolated from one another", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
|
||||
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
|
||||
})
|
||||
|
||||
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
|
||||
})
|
||||
|
||||
t.Run("union across a single policy's guardrails", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
|
||||
"a policy's own multiple allowlist guardrails union together")
|
||||
})
|
||||
|
||||
t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
|
||||
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
|
||||
got := buildProviderAllowlists([]*types.Policy{
|
||||
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
|
||||
}, empty)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
|
||||
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
|
||||
})
|
||||
}
|
||||
@@ -1031,12 +1031,8 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
|
||||
|
||||
var cfg guardrailConfig
|
||||
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
|
||||
// Both policies restrict the same provider, so the per-provider backstop
|
||||
// carries the union of their models — a coarse gate that management's
|
||||
// per-policy/group check narrows; it only blocks models outside the union
|
||||
// when management is down.
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
|
||||
"per-provider allowlist union must keep both models")
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
|
||||
"model allowlist union must keep both models")
|
||||
}
|
||||
|
||||
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {
|
||||
|
||||
@@ -285,7 +285,6 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
|
||||
UserID: req.GetUserId(),
|
||||
GroupIDs: req.GetGroupIds(),
|
||||
ProviderID: req.GetProviderId(),
|
||||
Model: req.GetModel(),
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("select policy for request: %v", err)
|
||||
|
||||
@@ -1,138 +0,0 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/grpc/codes"
|
||||
grpcstatus "google.golang.org/grpc/status"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
|
||||
// and returns a pre-programmed result, so tests can assert what the handler
|
||||
// forwards to the selector.
|
||||
type fakeAgentNetworkLimits struct {
|
||||
gotInput agentnetwork.PolicySelectionInput
|
||||
result *agentnetwork.PolicySelectionResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
|
||||
f.gotInput = in
|
||||
if f.err != nil {
|
||||
return nil, f.err
|
||||
}
|
||||
return f.result, nil
|
||||
}
|
||||
|
||||
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
|
||||
// the model the proxy extracted must reach the selector's Model unchanged,
|
||||
// alongside the account/user/group/provider fields.
|
||||
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
req := &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
UserId: "user-1",
|
||||
GroupIds: []string{"grp-a", "grp-b"},
|
||||
ProviderId: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
}
|
||||
|
||||
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
|
||||
assert.Equal(t, "user-1", fake.gotInput.UserID)
|
||||
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
|
||||
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
|
||||
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
|
||||
"the request's model must be forwarded to the selector")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
|
||||
// model (empty string) is forwarded as-is; the selector decides how to treat it.
|
||||
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
req := &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
}
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
|
||||
// envelope surfaces the model-allowlist deny code + reason through the response.
|
||||
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: "llm_policy.model_blocked",
|
||||
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
|
||||
}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Equal(t, "deny", resp.Decision)
|
||||
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
|
||||
assert.NotEmpty(t, resp.DenyReason)
|
||||
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
|
||||
// failure surfaces as an Internal gRPC error rather than a silent allow.
|
||||
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.Error(t, err)
|
||||
st, ok := grpcstatus.FromError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
|
||||
// fallback: with no limits service wired the RPC returns Unimplemented.
|
||||
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
|
||||
s := &ProxyServiceServer{}
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
})
|
||||
require.Error(t, err)
|
||||
st, ok := grpcstatus.FromError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, codes.Unimplemented, st.Code())
|
||||
}
|
||||
@@ -3,30 +3,20 @@ package auth
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/sync/singleflight"
|
||||
|
||||
"github.com/netbirdio/netbird/proxy/internal/types"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer
|
||||
// result is reused before re-fetching from management. 5 minutes balances
|
||||
// freshness against management load on busy mesh networks. Override it with
|
||||
// envTunnelCacheTTL when an account needs authorization changes to take effect
|
||||
// sooner (at the cost of more ValidateTunnelPeer RPCs).
|
||||
// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
|
||||
// reused before re-fetching from management. 5 minutes balances freshness
|
||||
// against management load on busy mesh networks.
|
||||
const tunnelCacheTTL = 300 * time.Second
|
||||
|
||||
// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string
|
||||
// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the
|
||||
// default.
|
||||
const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL"
|
||||
|
||||
// tunnelCachePerAccount caps the number of cached identities per account.
|
||||
// Bounded eviction avoids memory growth in pathological cases (huge peer
|
||||
// roster, brief request bursts) while staying generous for normal use.
|
||||
@@ -70,35 +60,16 @@ type accountBucket struct {
|
||||
order []tunnelCacheKey
|
||||
}
|
||||
|
||||
// newTunnelValidationCache constructs a cache with the configured TTL
|
||||
// (envTunnelCacheTTL override or default) and default bounds.
|
||||
// newTunnelValidationCache constructs a cache with default TTL and bounds.
|
||||
func newTunnelValidationCache() *tunnelValidationCache {
|
||||
return &tunnelValidationCache{
|
||||
entries: make(map[types.AccountID]*accountBucket),
|
||||
ttl: tunnelCacheTTLFromEnv(),
|
||||
ttl: tunnelCacheTTL,
|
||||
maxSize: tunnelCachePerAccount,
|
||||
now: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the
|
||||
// envTunnelCacheTTL override. The override must be a positive Go duration
|
||||
// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive
|
||||
// falls back to tunnelCacheTTL.
|
||||
func tunnelCacheTTLFromEnv() time.Duration {
|
||||
raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL))
|
||||
if raw == "" {
|
||||
return tunnelCacheTTL
|
||||
}
|
||||
d, err := time.ParseDuration(raw)
|
||||
if err != nil || d <= 0 {
|
||||
log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
|
||||
envTunnelCacheTTL, raw, tunnelCacheTTL)
|
||||
return tunnelCacheTTL
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// get returns a cached response for the key, or nil when missing or
|
||||
// expired. Expired entries are evicted lazily on read.
|
||||
func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {
|
||||
|
||||
@@ -169,32 +169,3 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
|
||||
assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
|
||||
assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
|
||||
}
|
||||
|
||||
func TestTunnelCacheTTLFromEnv(t *testing.T) {
|
||||
t.Run("unset uses default", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, "")
|
||||
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
|
||||
})
|
||||
t.Run("valid duration overrides", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, "45s")
|
||||
assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv())
|
||||
})
|
||||
t.Run("whitespace trimmed", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, " 2m ")
|
||||
assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv())
|
||||
})
|
||||
t.Run("unparseable uses default", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, "nonsense")
|
||||
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
|
||||
})
|
||||
t.Run("non-positive uses default", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, "0s")
|
||||
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
|
||||
t.Setenv(envTunnelCacheTTL, "-30s")
|
||||
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
|
||||
})
|
||||
t.Run("constructor honors override", func(t *testing.T) {
|
||||
t.Setenv(envTunnelCacheTTL, "90s")
|
||||
assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -28,12 +28,12 @@ func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
|
||||
"ministral-8b-latest", "mistral-embed",
|
||||
},
|
||||
"anthropic": {
|
||||
"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
|
||||
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
|
||||
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
|
||||
},
|
||||
// bedrock keys are the normalized ids the request parser emits.
|
||||
"bedrock": {
|
||||
"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
|
||||
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
|
||||
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
|
||||
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
|
||||
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",
|
||||
|
||||
@@ -180,11 +180,6 @@ anthropic:
|
||||
output_per_1k: 0.050
|
||||
cache_read_per_1k: 0.001
|
||||
cache_creation_per_1k: 0.0125
|
||||
claude-opus-5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
@@ -241,11 +236,6 @@ bedrock:
|
||||
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
|
||||
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
|
||||
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
|
||||
anthropic.claude-opus-5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
|
||||
@@ -10,15 +10,11 @@ import (
|
||||
)
|
||||
|
||||
// Config is the JSON-decoded shape accepted by the factory. The
|
||||
// runtime path consumes the normalised allowlists; raw config is not
|
||||
// runtime path consumes the normalised allowlist; raw config is not
|
||||
// retained beyond construction.
|
||||
type Config struct {
|
||||
// ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to
|
||||
// its model allowlist. A provider present is restricted to those models; one
|
||||
// absent is unrestricted. Kept per-provider so one provider's list can't leak
|
||||
// onto another.
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture PromptCapture `json:"prompt_capture"`
|
||||
ModelAllowlist []string `json:"model_allowlist"`
|
||||
PromptCapture PromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
// PromptCapture toggles the optional prompt capture + redaction step
|
||||
@@ -58,28 +54,21 @@ func isEmptyJSON(raw []byte) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// normaliseConfig lowercases and trims allowlist entries for case-insensitive
|
||||
// matching; empty entries drop. A provider whose entries all drop keeps an empty
|
||||
// (non-nil) list — "deny every model" — distinct from an absent provider
|
||||
// (unrestricted).
|
||||
// normaliseConfig lowercases and trims allowlist entries so the runtime
|
||||
// match is case-insensitive. Empty entries are dropped.
|
||||
func normaliseConfig(cfg Config) Config {
|
||||
if len(cfg.ProviderAllowlists) == 0 {
|
||||
cfg.ProviderAllowlists = nil
|
||||
if len(cfg.ModelAllowlist) == 0 {
|
||||
return cfg
|
||||
}
|
||||
cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
|
||||
for provider, models := range cfg.ProviderAllowlists {
|
||||
list := make([]string, 0, len(models))
|
||||
for _, entry := range models {
|
||||
n := normaliseModel(entry)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
list = append(list, n)
|
||||
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
|
||||
for _, entry := range cfg.ModelAllowlist {
|
||||
n := normaliseModel(entry)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
cleaned[provider] = list
|
||||
cleaned = append(cleaned, n)
|
||||
}
|
||||
cfg.ProviderAllowlists = cleaned
|
||||
cfg.ModelAllowlist = cleaned
|
||||
return cfg
|
||||
}
|
||||
|
||||
|
||||
@@ -83,9 +83,8 @@ func (m *Middleware) MutationsSupported() bool { return false }
|
||||
// prompt capture only affects the metadata emitted alongside an allow.
|
||||
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)
|
||||
|
||||
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
|
||||
if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
|
||||
return denial, nil
|
||||
}
|
||||
|
||||
@@ -111,32 +110,20 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
|
||||
// is a no-op.
|
||||
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 {
|
||||
if len(m.cfg.ProviderAllowlists) == 0 {
|
||||
// evaluateAllowlist returns a deny Output when the configured allowlist
|
||||
// rejects the model. A nil return means the request should proceed.
|
||||
func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output {
|
||||
if len(m.cfg.ModelAllowlist) == 0 {
|
||||
return nil
|
||||
}
|
||||
// Restrictions exist but the resolved provider is unknown, so we can't tell
|
||||
// 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)
|
||||
}
|
||||
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
|
||||
if !restricted {
|
||||
// This provider has no allowlist (some authorising policy left it
|
||||
// unrestricted); management owns any per-policy/group decision.
|
||||
return nil
|
||||
}
|
||||
// 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.
|
||||
// Fail closed: with an allowlist configured, a request whose model the
|
||||
// upstream parser could not extract (absent or empty) must be denied rather
|
||||
// than allowed. This is what enforces the allowlist for URL/path-routed
|
||||
// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
|
||||
if !modelPresent || normaliseModel(model) == "" {
|
||||
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
|
||||
}
|
||||
if modelInAllowlist(allowlist, model) {
|
||||
if m.modelInAllowlist(model) {
|
||||
return nil
|
||||
}
|
||||
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
|
||||
@@ -164,15 +151,14 @@ func denyModel(model, code, message, reason string) *middleware.Output {
|
||||
}
|
||||
}
|
||||
|
||||
// modelInAllowlist reports whether the model matches any entry in the supplied
|
||||
// (already-normalised) allowlist under the case-insensitive, trim-tolerant
|
||||
// comparison rule.
|
||||
func modelInAllowlist(allowlist []string, model string) bool {
|
||||
// modelInAllowlist reports whether the model matches any allowlist
|
||||
// entry under the case-insensitive, trim-tolerant comparison rule.
|
||||
func (m *Middleware) modelInAllowlist(model string) bool {
|
||||
normalised := normaliseModel(model)
|
||||
if normalised == "" {
|
||||
return false
|
||||
}
|
||||
for _, allowed := range allowlist {
|
||||
for _, allowed := range m.cfg.ModelAllowlist {
|
||||
if allowed == normalised {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -26,25 +26,6 @@ func newInput(meta ...middleware.KV) *middleware.Input {
|
||||
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
|
||||
}
|
||||
|
||||
const (
|
||||
testProvider = "prov-1"
|
||||
otherProvider = "prov-2"
|
||||
)
|
||||
|
||||
// providerCfg builds a Config restricting testProvider to the given models.
|
||||
func providerCfg(models ...string) Config {
|
||||
return Config{ProviderAllowlists: map[string][]string{testProvider: models}}
|
||||
}
|
||||
|
||||
// newInputProvider builds an input that carries a resolved provider id (as
|
||||
// llm_router would stamp) plus any extra metadata.
|
||||
func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input {
|
||||
all := make([]middleware.KV, 0, len(meta)+1)
|
||||
all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider})
|
||||
all = append(all, meta...)
|
||||
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all}
|
||||
}
|
||||
|
||||
func TestMiddlewareIdentity(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
|
||||
@@ -66,12 +47,12 @@ func TestMiddlewareIdentity(t *testing.T) {
|
||||
|
||||
func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists must allow any model")
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model")
|
||||
v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
|
||||
require.True(t, ok, "decision metadata must be emitted")
|
||||
assert.Equal(t, "allow", v, "decision must be allow")
|
||||
@@ -81,8 +62,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMatchAllows(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o", "claude-opus-4"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -90,8 +71,8 @@ func TestAllowlistMatchAllows(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMissDenies(t *testing.T) {
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -110,10 +91,10 @@ func TestAllowlistMissDenies(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistCaseInsensitive(t *testing.T) {
|
||||
mw := New(providerCfg(" GPT-4o ", "Claude-OPUS-4"))
|
||||
mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}})
|
||||
cases := []string{"gpt-4o", "GPT-4O", " claude-opus-4 "}
|
||||
for _, model := range cases {
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: model},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
@@ -122,15 +103,14 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAllowlistMissingModelKeyDenies(t *testing.T) {
|
||||
// Fail closed: with an allowlist in effect for the resolved provider, a
|
||||
// request whose model the parser could not extract (URL/path-routed
|
||||
// providers such as Bedrock or Vertex whose shape wasn't recognised) must be
|
||||
// denied, not allowed.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider))
|
||||
// Fail closed: with an allowlist configured, a request whose model the
|
||||
// parser could not extract (URL/path-routed providers such as Bedrock or
|
||||
// Vertex whose shape wasn't recognised) must be denied, not allowed.
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted")
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
|
||||
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
|
||||
require.NotNil(t, out.DenyReason, "deny reason must be populated")
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
@@ -142,101 +122,26 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) {
|
||||
|
||||
func TestAllowlistEmptyModelValueDenies(t *testing.T) {
|
||||
// A present-but-empty model is as undeterminable as an absent one.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: " "},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when the provider is restricted")
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
|
||||
require.NotNil(t, out.DenyReason, "deny reason must be populated")
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
}
|
||||
|
||||
func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
|
||||
// Without any provider allowlists there is nothing to enforce, so a missing
|
||||
// model is still allowed — the fail-closed rule only applies when a
|
||||
// restriction is in effect.
|
||||
// Without an allowlist there is nothing to enforce, so a missing model is
|
||||
// still allowed — the fail-closed rule only applies when a list is set.
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInput())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
|
||||
}
|
||||
|
||||
func TestUnrestrictedProviderAllowsAnyModel(t *testing.T) {
|
||||
// The request resolved to otherProvider, which has no allowlist, so its
|
||||
// traffic must not be caught by testProvider's restriction — the
|
||||
// cross-provider-leak / false-deny guard.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist")
|
||||
}
|
||||
|
||||
func TestPerProviderAllowlistsAreIsolated(t *testing.T) {
|
||||
// gpt-4o is allowed only on testProvider; claude-opus-4 only on
|
||||
// otherProvider. A model allowlisted for one provider must not be usable on
|
||||
// the other — the fail-closed layer never unions allowlists across providers.
|
||||
mw := New(Config{ProviderAllowlists: map[string][]string{
|
||||
testProvider: {"gpt-4o"},
|
||||
otherProvider: {"claude-opus-4"},
|
||||
}})
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown")
|
||||
}
|
||||
|
||||
func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
|
||||
// Restrictions exist for the account but the resolved provider id is absent,
|
||||
// so the request cannot be scoped to a provider. Fail closed rather than
|
||||
// wave it through.
|
||||
mw := New(providerCfg("gpt-4o"))
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
|
||||
}
|
||||
|
||||
func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) {
|
||||
// An allowlist-enabled provider with zero models is distinct from an
|
||||
// unrestricted (absent) provider: it must deny every model.
|
||||
mw := New(providerCfg())
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown")
|
||||
}
|
||||
|
||||
func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) {
|
||||
// All the provider's entries are blank; they collapse to a non-nil empty
|
||||
// list (deny everything for that provider), not "no restriction".
|
||||
raw := []byte(`{"provider_allowlists":{"prov-1":[""," "]}}`)
|
||||
mw, err := Factory{}.New(raw)
|
||||
require.NoError(t, err)
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, out)
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider")
|
||||
require.NotNil(t, out.DenyReason)
|
||||
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked")
|
||||
}
|
||||
|
||||
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
|
||||
mw := New(Config{})
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
@@ -312,8 +217,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) {
|
||||
|
||||
func TestFactoryDecodesValidConfig(t *testing.T) {
|
||||
cfg := Config{
|
||||
ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}},
|
||||
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
|
||||
ModelAllowlist: []string{"gpt-4o"},
|
||||
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
|
||||
}
|
||||
raw, err := json.Marshal(cfg)
|
||||
require.NoError(t, err, "marshalling test config must succeed")
|
||||
@@ -329,15 +234,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestFactoryNormalisesAllowlist(t *testing.T) {
|
||||
raw := []byte(`{"provider_allowlists":{"prov-1":[" GPT-4o ","",""," Claude-3 "]}}`)
|
||||
raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`)
|
||||
mw, err := Factory{}.New(raw)
|
||||
require.NoError(t, err)
|
||||
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries")
|
||||
out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
|
||||
out2, err := mw.Invoke(context.Background(), newInput(
|
||||
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
|
||||
))
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -175,7 +175,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
|
||||
DenyStatus: 403,
|
||||
DenyReason: &middleware.DenyReason{
|
||||
Code: code,
|
||||
Message: denyMessageForCode(code),
|
||||
Message: "LLM policy limit exceeded",
|
||||
},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
|
||||
@@ -184,21 +184,6 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou
|
||||
}
|
||||
}
|
||||
|
||||
// denyMessageForCode maps a management deny code to a public message.
|
||||
// Model-allowlist rejections get a model-specific message matching the
|
||||
// local guardrail; everything else keeps the generic quota wording. The
|
||||
// message stays generic so it never leaks internal quota detail.
|
||||
func denyMessageForCode(code string) string {
|
||||
switch code {
|
||||
case "llm_policy.model_blocked":
|
||||
return "model is not in the policy allowlist"
|
||||
case "llm_policy.model_unknown":
|
||||
return "request model could not be determined for the policy allowlist"
|
||||
default:
|
||||
return "LLM policy limit exceeded"
|
||||
}
|
||||
}
|
||||
|
||||
// lookupKV returns the value associated with key, or the empty
|
||||
// string when absent.
|
||||
func lookupKV(kvs []middleware.KV, key string) string {
|
||||
|
||||
@@ -115,46 +115,6 @@ func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) {
|
||||
assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller")
|
||||
}
|
||||
|
||||
// TestInvoke_ModelDenyMessages proves a model-allowlist rejection gets a
|
||||
// model-specific public message rather than the generic quota wording, so a
|
||||
// blocked or undetermined model reads consistently with the local guardrail.
|
||||
func TestInvoke_ModelDenyMessages(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
code string
|
||||
message string
|
||||
}{
|
||||
{"blocked", "llm_policy.model_blocked", "model is not in the policy allowlist"},
|
||||
{"unknown", "llm_policy.model_unknown", "request model could not be determined for the policy allowlist"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
mgmt := &fakeMgmt{
|
||||
checkResp: &proto.CheckLLMPolicyLimitsResponse{
|
||||
Decision: "deny",
|
||||
DenyCode: tc.code,
|
||||
},
|
||||
}
|
||||
m := New(mgmt, nil)
|
||||
|
||||
out := runInvoke(t, m, &middleware.Input{
|
||||
AccountID: "acc-1",
|
||||
UserGroups: []string{"grp-engineers"},
|
||||
Metadata: []middleware.KV{
|
||||
{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
|
||||
{Key: middleware.KeyLLMModel, Value: "some-model"},
|
||||
},
|
||||
})
|
||||
|
||||
assert.Equal(t, middleware.DecisionDeny, out.Decision)
|
||||
require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload")
|
||||
assert.Equal(t, tc.code, out.DenyReason.Code, "canonical deny code surfaces to the caller")
|
||||
assert.Equal(t, tc.message, out.DenyReason.Message,
|
||||
"model denials must use a model-specific message, matching the local guardrail")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvoke_NoMgmtClientPassesThrough proves the partial-wiring
|
||||
// safety: a middleware constructed without a management client
|
||||
// allows every request without attribution. This makes a half-set-up
|
||||
|
||||
@@ -25,17 +25,10 @@ func runParserGuardrail(t *testing.T, url string, body []byte, allowlist []strin
|
||||
})
|
||||
require.NoError(t, err, "parser must not error")
|
||||
|
||||
const providerID = "prov-under-test"
|
||||
guard := llm_guardrail.New(llm_guardrail.Config{
|
||||
ProviderAllowlists: map[string][]string{providerID: allowlist},
|
||||
})
|
||||
// The real chain has llm_router stamp the resolved provider id before the
|
||||
// guardrail runs; the parser doesn't, so add it here so the guardrail can
|
||||
// scope the allowlist to this provider.
|
||||
meta := append([]middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: providerID}}, parsed.Metadata...)
|
||||
guard := llm_guardrail.New(llm_guardrail.Config{ModelAllowlist: allowlist})
|
||||
out, err := guard.Invoke(context.Background(), &middleware.Input{
|
||||
Slot: middleware.SlotOnRequest,
|
||||
Metadata: meta,
|
||||
Metadata: parsed.Metadata,
|
||||
})
|
||||
require.NoError(t, err, "guardrail must not error")
|
||||
require.NotNil(t, out, "guardrail must return an output")
|
||||
|
||||
Reference in New Issue
Block a user