Compare commits

..

3 Commits

Author SHA1 Message Date
bcmmbaga
6725b02cbb Merge branch 'main' into add-atomic-cache-ops 2026-07-22 20:20:14 +03:00
bcmmbaga
da19dcf480 prevent concurrent JWT reuse 2026-07-22 20:08:48 +03:00
bcmmbaga
6426d6f03f add atomic SetNX cache operation
Split the memory and Redis cache implementations into separate files and
provide backend-native atomic set-if-absent support.
2026-07-22 20:03:03 +03:00
106 changed files with 1029 additions and 3862 deletions

View File

@@ -5,13 +5,6 @@ on:
schedule:
- cron: "0 3 * * *"
workflow_dispatch:
inputs:
bedrock_model:
description: >-
Bedrock inference-profile id to drive the matrix with, exactly as
AWS issues it. Leave empty for the Sonnet 4.6 default.
required: false
default: ""
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
@@ -58,9 +51,6 @@ jobs:
# token (and URL, for gateways) is unset, so partial coverage is fine.
OPENAI_TOKEN: ${{ secrets.E2E_OPENAI_TOKEN }}
ANTHROPIC_TOKEN: ${{ secrets.E2E_ANTHROPIC_TOKEN }}
# Moonshot AI platform key (platform.kimi.ai); drives both Kimi wire
# shapes (OpenAI /v1 and Anthropic /anthropic) through kimi_api.
KIMI_TOKEN: ${{ secrets.E2E_KIMI_TOKEN }}
VERCEL_URL: ${{ secrets.E2E_VERCEL_URL }}
VERCEL_TOKEN: ${{ secrets.E2E_VERCEL_TOKEN }}
OPENROUTER_URL: ${{ secrets.E2E_OPENROUTER_URL }}
@@ -69,8 +59,6 @@ jobs:
CLOUDFLARE_TOKEN: ${{ secrets.E2E_CLOUDFLARE_TOKEN }}
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.E2E_AWS_BEARER_TOKEN_BEDROCK }}
AWS_REGION: ${{ secrets.E2E_AWS_REGION }}
# Bedrock model override: dispatch input wins, then the repo variable, else the test default.
AWS_BEDROCK_MODEL: ${{ inputs.bedrock_model || vars.E2E_AWS_BEDROCK_MODEL }}
# Vertex (Anthropic-on-Vertex): SA + project required; region defaults
# to "global", model to a pinned claude snapshot.
GOOGLE_VERTEX_SA_BASE64: ${{ secrets.E2E_GOOGLE_VERTEX_SA_BASE64 }}

View File

@@ -464,8 +464,6 @@ func Test_RemovePeer(t *testing.T) {
}
func Test_ConnectPeers(t *testing.T) {
t.Setenv("NB_DISABLE_EBPF_WG_PROXY", "true")
peer1ifaceName := fmt.Sprintf("utun%d", WgIntNumber+400)
peer1wgIP := netip.MustParsePrefix("10.99.99.17/30")
peer1Key, _ := wgtypes.GeneratePrivateKey()
@@ -507,8 +505,12 @@ func Test_ConnectPeers(t *testing.T) {
t.Fatal(err)
}
localIP1 := "127.0.0.1"
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP1, peer1wgPort))
localIP, err := getLocalIP()
if err != nil {
t.Fatal(err)
}
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer1wgPort))
if err != nil {
t.Fatal(err)
}
@@ -544,8 +546,7 @@ func Test_ConnectPeers(t *testing.T) {
t.Fatal(err)
}
localIP2 := "127.0.0.1"
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP2, peer2wgPort))
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer2wgPort))
if err != nil {
t.Fatal(err)
}
@@ -568,17 +569,17 @@ func Test_ConnectPeers(t *testing.T) {
if err != nil {
t.Fatal(err)
}
// The peers use userspace WireGuard (stdnet transport). A tight busy-loop
// here starves the wireguard-go goroutines that process the handshake, so
// poll on a ticker instead and yield the CPU between checks. WireGuard also
// only retries a lost handshake initiation every REKEY_TIMEOUT (5s), which
// is why the overall wait can occasionally stretch to tens of seconds.
// todo: investigate why in some tests execution we need 30s
timeout := 30 * time.Second
timeoutChannel := time.After(timeout)
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-timeoutChannel:
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
default:
}
peer, gpErr := getPeer(peer1ifaceName, peer2Key.PublicKey().String())
if gpErr != nil {
t.Fatal(gpErr)
@@ -587,12 +588,6 @@ func Test_ConnectPeers(t *testing.T) {
t.Log("peers successfully handshake")
break
}
select {
case <-timeoutChannel:
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
case <-ticker.C:
}
}
}
@@ -620,3 +615,28 @@ func getPeer(ifaceName, peerPubKey string) (wgtypes.Peer, error) {
}
return wgtypes.Peer{}, fmt.Errorf("peer not found")
}
func getLocalIP() (string, error) {
// Get all interfaces
addrs, err := net.InterfaceAddrs()
if err != nil {
return "", err
}
for _, addr := range addrs {
ipNet, ok := addr.(*net.IPNet)
if !ok {
continue
}
if ipNet.IP.IsLoopback() {
continue
}
if ipNet.IP.To4() == nil {
continue
}
return ipNet.IP.String(), nil
}
return "", fmt.Errorf("no local IP found")
}

View File

@@ -49,21 +49,11 @@ type ConnMgr struct {
// engine.syncMsgMux; all other reads stay under engine.syncMsgMux only.
lazyConnMgrMu sync.RWMutex
// reconcileRoutedIPs re-applies a peer's routed allowed IPs after its lazy wake endpoint is
// (re)armed (Mode A at arm time). Injected by the engine; nil disables the reconcile.
reconcileRoutedIPs func(peerKey string) error
wg sync.WaitGroup
lazyCtx context.Context
lazyCtxCancel context.CancelFunc
}
// SetRoutedIPsReconciler injects the callback used to re-apply a peer's routed allowed IPs when
// its lazy wake endpoint is (re)armed. Must be called before the lazy manager starts.
func (e *ConnMgr) SetRoutedIPsReconciler(fn func(peerKey string) error) {
e.reconcileRoutedIPs = fn
}
func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr {
e := &ConnMgr{
peerStore: peerStore,
@@ -301,7 +291,6 @@ func (e *ConnMgr) Close() {
func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
cfg := manager.Config{
InactivityThreshold: inactivityThresholdEnv(),
ReconcileAllowedIPs: e.reconcileRoutedIPs,
}
e.lazyConnMgrMu.Lock()

View File

@@ -3,7 +3,6 @@ package dns
import (
"context"
"fmt"
"maps"
"math"
"net"
"slices"
@@ -32,26 +31,6 @@ 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
@@ -73,19 +52,8 @@ type ResponseWriterChain struct {
origPattern string
requestID string
shouldContinue bool
// 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
response *dns.Msg
meta map[string]string
}
// RequestID returns the request ID for tracing
@@ -94,64 +62,26 @@ func (w *ResponseWriterChain) RequestID() string {
}
// SetMeta sets a metadata key-value pair for logging
func (w *ResponseWriterChain) SetMeta(key resutil.MetaKey, value string) {
func (w *ResponseWriterChain) SetMeta(key, value string) {
if w.meta == nil {
w.meta = make(responseMeta)
w.meta = make(map[string]string)
}
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(resutil.MetaKeyTruncated, "true")
w.SetMeta("truncated", "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),
@@ -293,19 +223,6 @@ 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 {
@@ -328,16 +245,11 @@ 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)
}
@@ -353,59 +265,30 @@ 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(), cw.meta.format(), time.Since(startTime))
cw.response.Len(), meta, 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.
//
// "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
// 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
// (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 {

View File

@@ -3,19 +3,15 @@ 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"
)
@@ -1242,276 +1238,6 @@ 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"}

View File

@@ -19,34 +19,6 @@ 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)
@@ -106,38 +78,10 @@ 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 MetaKey, value string)
SetMeta(key, value string)
}
// GetRequestID extracts a request ID from the ResponseWriter if available,
@@ -152,30 +96,12 @@ func GetRequestID(w dns.ResponseWriter) string {
}
// SetMeta sets metadata on the ResponseWriter if it supports it.
func SetMeta(w dns.ResponseWriter, key MetaKey, value string) {
func SetMeta(w dns.ResponseWriter, key, 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
@@ -273,27 +199,11 @@ 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 outside
// SupportedRecordQtype yield NODATA so that a routed name is never poisoned
// with NXDOMAIN for a type we cannot look up.
// 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.
func LookupRecords(ctx context.Context, r RecordResolver, name string, qtype uint16, ttl uint32) ([]dns.RR, int) {
fqdn := dns.Fqdn(name)

View File

@@ -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, resutil.MetaKeyProtocol, network)
resutil.SetMeta(w, "protocol", 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, resutil.MetaKeyEDE, res.ede)
resutil.SetMeta(w, "ede", 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, resutil.MetaKeyEDE, res.ede)
resutil.SetMeta(w, "ede", 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, resutil.MetaKeyUpstream, upstream.String())
resutil.SetMeta(w, "upstream", upstream.String())
if proto != "" {
resutil.SetMeta(w, resutil.MetaKeyUpstreamProtocol, proto)
resutil.SetMeta(w, "upstream_protocol", proto)
}
// Clear Zero bit from external responses to prevent upstream servers from

View File

@@ -26,6 +26,15 @@ 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)
@@ -207,9 +216,6 @@ 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 == "" {
@@ -223,22 +229,20 @@ func (f *DNSForwarder) handleDNSQuery(logger *log.Entry, w dns.ResponseWriter, q
reqHasEdns := query.IsEdns0() != nil
switch {
case question.Qtype == dns.TypeA || question.Qtype == dns.TypeAAAA:
switch question.Qtype {
case dns.TypeA, dns.TypeAAAA:
f.handleAddressQuery(ctx, logger, w, resp, mostSpecificResId, matchingEntries, reqHasEdns, startTime)
case resutil.SupportedRecordQtype(question.Qtype):
case dns.TypeMX, dns.TypeTXT, dns.TypeNS, dns.TypeSRV, dns.TypeCNAME, dns.TypePTR:
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; 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.
// authoritative one. The OPT pseudo-record must not appear unless the
// query advertised EDNS0.
if reqHasEdns {
resutil.AttachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
attachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
}
f.writeResponse(logger, w, resp, qname, startTime)
}
@@ -437,7 +441,7 @@ func (f *DNSForwarder) handleDNSError(
}
if reqHasEdns {
resutil.AttachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
attachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
}
f.writeResponse(logger, w, resp, domain, startTime)
@@ -484,9 +488,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 resutil.EDENetbirdUpstreamTimeout
return edeNetbirdUpstreamTimeout
}
return resutil.EDENetbirdUpstreamFailure
return edeNetbirdUpstreamFailure
}
// edeText builds the EDE extra-text describing the class of upstream failure.
@@ -499,3 +503,14 @@ 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})
}

View File

@@ -649,39 +649,6 @@ 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 {
@@ -892,7 +859,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: resutil.EDENetbirdUpstreamTimeout,
wantCode: edeNetbirdUpstreamTimeout,
wantTextHas: "netbird forwarder: upstream timeout",
},
{
@@ -900,7 +867,7 @@ func TestDNSForwarder_UpstreamFailureEDE(t *testing.T) {
lookupErr: &net.DNSError{Err: "server misbehaving", Server: "10.0.0.53:53"},
reqEdns: true,
wantEDE: true,
wantCode: resutil.EDENetbirdUpstreamFailure,
wantCode: edeNetbirdUpstreamFailure,
wantTextHas: "netbird forwarder: upstream failure",
},
{

View File

@@ -52,14 +52,11 @@ int xdp_dns_fwd(struct iphdr *ip, struct udphdr *udp) {
if (udp->dest == GENERAL_DNS_PORT && ip->daddr == dns_ip) {
udp->dest = dns_port;
// Clear the now-stale checksum; zero means "not computed" for IPv4.
udp->check = 0;
return XDP_PASS;
}
if (udp->source == dns_port && ip->saddr == dns_ip) {
udp->source = GENERAL_DNS_PORT;
udp->check = 0;
return XDP_PASS;
}

View File

@@ -50,11 +50,5 @@ int xdp_wg_proxy(struct iphdr *ip, struct udphdr *udp) {
__be16 new_dst_port = htons(proxy_port);
udp->dest = new_dst_port;
udp->source = new_src_port;
// The ports are covered by the UDP checksum. This is an IPv4 loopback hop
// and the payload is already integrity-protected, so clear the checksum (a
// zero UDP checksum means "not computed" for IPv4) rather than leave a
// stale value the kernel would drop as UDP_CSUM.
udp->check = 0;
return XDP_PASS;
}

View File

@@ -663,12 +663,6 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
iceCfg := e.createICEConfig()
e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface)
e.connMgr.SetRoutedIPsReconciler(func(peerKey string) error {
if e.routeManager == nil {
return nil
}
return e.routeManager.ReconcilePeerAllowedIPs(peerKey)
})
e.connMgr.Start(e.ctx)
// Wire DNS-time lazy-connection warm-up now that the connection manager

View File

@@ -29,11 +29,6 @@ type managedPeer struct {
type Config struct {
InactivityThreshold *time.Duration
// ReconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is
// armed. The activity listener creates the wake peer with the overlay /32 only; without the
// routed prefixes WireGuard would not steer subnet-bound traffic to the wake endpoint, so an
// idle routing peer could never be woken by that traffic. Optional; nil disables the reconcile.
ReconcileAllowedIPs func(peerKey string) error
}
// Manager manages lazy connections
@@ -61,9 +56,6 @@ type Manager struct {
peerToHAGroups map[string][]route.HAUniqueID // peer ID -> HA groups they belong to
haGroupToPeers map[route.HAUniqueID][]string // HA group -> peer IDs in the group
routesMu sync.RWMutex
// reconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is armed.
reconcileAllowedIPs func(peerKey string) error
}
// NewManager creates a new lazy connection manager
@@ -81,7 +73,6 @@ func NewManager(config Config, engineCtx context.Context, peerStore *peerstore.S
activityManager: activity.NewManager(wgIface),
peerToHAGroups: make(map[string][]route.HAUniqueID),
haGroupToPeers: make(map[route.HAUniqueID][]string),
reconcileAllowedIPs: config.ReconcileAllowedIPs,
}
if wgIface.IsUserspaceBind() {
@@ -210,7 +201,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
return false, nil
}
if err := m.armActivityListener(peerCfg); err != nil {
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
return false, err
}
@@ -297,7 +288,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) {
m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
if err := m.armActivityListener(*mp.peerCfg); err != nil {
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
return
}
@@ -474,31 +465,6 @@ func (m *Manager) close() {
}
// shouldDeferIdleForHA checks if peer should stay connected due to HA group requirements
// armRoutedAllowedIPs re-applies the peer's routed allowed IPs onto its freshly armed wake
// endpoint. The activity listener creates the wake peer with the overlay /32 only, so without
// this the routed prefixes would be missing and traffic to a routed subnet could not wake the
// idle routing peer. It is a no-op when no reconciler is configured.
// armActivityListener (re)arms the peer's wake endpoint via the activity manager and then
// re-applies its routed allowed IPs, so traffic to a routed subnet can wake an idle routing
// peer. The routed prefixes must be re-applied after the wake endpoint exists because the
// listener creates it with the overlay /32 only.
func (m *Manager) armActivityListener(peerCfg lazyconn.PeerConfig) error {
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
return err
}
m.armRoutedAllowedIPs(&peerCfg)
return nil
}
func (m *Manager) armRoutedAllowedIPs(peerCfg *lazyconn.PeerConfig) {
if m.reconcileAllowedIPs == nil {
return
}
if err := m.reconcileAllowedIPs(peerCfg.PublicKey); err != nil {
peerCfg.Log.Errorf("failed to reconcile routed allowed IPs on wake endpoint: %v", err)
}
}
func (m *Manager) shouldDeferIdleForHA(inactivePeers map[string]struct{}, peerID string) bool {
m.routesMu.RLock()
defer m.routesMu.RUnlock()
@@ -611,7 +577,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
mp.peerCfg.Log.Infof("start activity monitor")
if err := m.armActivityListener(*mp.peerCfg); err != nil {
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
continue
}

View File

@@ -22,6 +22,7 @@ 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"
@@ -39,61 +40,6 @@ 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
@@ -104,7 +50,7 @@ type DnsInterceptor struct {
currentPeerKey string
interceptedDomains domainMap
wgInterface iface.WGIface
peerStore peerAllowedIPs
peerStore *peerstore.Store
firewall firewall.Manager
fakeIPManager *fakeip.Manager
forwarderPort *atomic.Uint32
@@ -280,14 +226,11 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
return
}
qtype := r.Question[0].Qtype
dispose := dispositionFor(qtype)
if dispose == dispositionChain {
d.deferToChain(w, r, logger, "unsupported-qtype")
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.
d.mu.RLock()
peerKey := d.currentPeerKey
d.mu.RUnlock()
@@ -303,32 +246,35 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
return
}
query, hadEdns := peerQuery(r)
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)
}
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, query, upstream, upstreamIP, peerKey, logger)
reply := d.queryUpstreamDNS(ctx, w, r, 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, resutil.MetaKeyEDE, fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
resutil.SetMeta(w, "ede", fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
}
if !hadEdns {
resutil.StripOPT(reply)
}
resutil.SetMeta(w, resutil.MetaKeyPeer, peerKey)
resutil.SetMeta(w, "peer", peerKey)
reply.Id = r.Id
if err := d.writeMsg(w, reply, logger); err != nil {
@@ -336,30 +282,6 @@ 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)
@@ -699,26 +621,3 @@ 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
}

View File

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

View File

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

View File

@@ -61,7 +61,6 @@ type Manager interface {
InitialRouteRange() []string
SetFirewall(firewall.Manager) error
SetDNSForwarderPort(port uint16)
ReconcilePeerAllowedIPs(peerKey string) error
Stop(stateManager *statemanager.Manager)
}
@@ -233,30 +232,6 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
)
}
// ReconcilePeerAllowedIPs re-applies every routed allowed IP currently tracked for the peer
// onto the WireGuard device. The allowed-IP refcounter only calls its AddFunc (which pushes to
// the device) on a prefix's 0->1 transition, so a peer whose device entry was rebuilt without a
// matching refcounter change — e.g. a lazy connection cycling through idle->wake, which recreates
// the WireGuard peer with the overlay /32 only — ends up missing routed prefixes the refcounter
// still considers installed, and nothing retries. Calling this when the peer's WireGuard entry is
// (re)created restores convergence. It is add-only and idempotent: AddAllowedIP is update-only, so
// prefixes are re-added to an existing peer and an absent peer is left untouched.
func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error {
if m.allowedIPsRefCounter == nil {
return nil
}
return m.allowedIPsRefCounter.ReapplyMatching(
func(out string) bool { return out == peerKey },
func(prefix netip.Prefix) error {
if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil {
return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)
}
return nil
},
)
}
// Init sets up the routing
func (m *DefaultManager) Init() error {
m.routeSelector = m.initSelector()

View File

@@ -112,11 +112,6 @@ func (m *MockManager) SetFirewall(firewall.Manager) error {
func (m *MockManager) SetDNSForwarderPort(port uint16) {
}
// ReconcilePeerAllowedIPs mock implementation of ReconcilePeerAllowedIPs from Manager interface
func (m *MockManager) ReconcilePeerAllowedIPs(peerKey string) error {
return nil
}
// Stop mock implementation of Stop from Manager interface
func (m *MockManager) Stop(stateManager *statemanager.Manager) {
if m.StopFunc != nil {

View File

@@ -1,90 +0,0 @@
//go:build !windows
package routemanager
import (
"net"
"net/netip"
"sync"
"testing"
"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/routemanager/refcounter"
)
// reconcileWGMock is a minimal iface.WGIface that only records AddAllowedIP calls; every other
// method is an inert stub because ReconcilePeerAllowedIPs exercises none of them.
type reconcileWGMock struct {
mu sync.Mutex
adds map[string][]netip.Prefix
}
func (m *reconcileWGMock) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
m.mu.Lock()
defer m.mu.Unlock()
if m.adds == nil {
m.adds = map[string][]netip.Prefix{}
}
m.adds[peerKey] = append(m.adds[peerKey], allowedIP)
return nil
}
func (m *reconcileWGMock) added(peerKey string) []netip.Prefix {
m.mu.Lock()
defer m.mu.Unlock()
return m.adds[peerKey]
}
func (m *reconcileWGMock) RemoveAllowedIP(string, netip.Prefix) error { return nil }
func (m *reconcileWGMock) Name() string { return "utun-test" }
func (m *reconcileWGMock) Address() wgaddr.Address { return wgaddr.Address{} }
func (m *reconcileWGMock) ToInterface() *net.Interface { return nil }
func (m *reconcileWGMock) IsUserspaceBind() bool { return false }
func (m *reconcileWGMock) GetFilter() device.PacketFilter { return nil }
func (m *reconcileWGMock) GetDevice() *device.FilteredDevice { return nil }
func (m *reconcileWGMock) GetNet() *netstack.Net { return nil }
// TestReconcilePeerAllowedIPs verifies the declarative reconcile re-applies every routed prefix
// tracked for the peer (self-heal, independent of refcount level) and stays scoped to that peer.
func TestReconcilePeerAllowedIPs(t *testing.T) {
wg := &reconcileWGMock{}
m := &DefaultManager{wgInterface: wg}
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 },
)
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
_, err := m.allowedIPsRefCounter.Increment(prefix, peer)
require.NoError(t, err)
}
// Extra reference: reconcile must still re-apply the prefix even though its refcount never
// hit 0 again (the exact case the plain incremental path skips).
_, err := m.allowedIPsRefCounter.Increment(peerA1, "peerA")
require.NoError(t, err)
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, wg.added("peerA"),
"reconcile must re-apply all routed prefixes of the peer")
assert.Empty(t, wg.added("peerB"), "reconcile must not touch another peer's prefixes")
}
// TestReconcilePeerAllowedIPsNoCounter verifies reconcile is a safe no-op before the refcounter is
// set up.
func TestReconcilePeerAllowedIPsNoCounter(t *testing.T) {
wg := &reconcileWGMock{}
m := &DefaultManager{wgInterface: wg}
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
assert.Empty(t, wg.added("peerA"))
}

View File

@@ -94,26 +94,6 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
return ref, ok
}
// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the
// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect
// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its
// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied.
// pred and apply are invoked under the lock, so they must not call back into the counter.
func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error {
rm.mu.Lock()
defer rm.mu.Unlock()
var merr *multierror.Error
for key, ref := range rm.refCountMap {
if pred(ref.Out) {
if err := apply(key); err != nil {
merr = multierror.Append(merr, err)
}
}
}
return nberrors.FormatErrorOrNil(merr)
}
// Increment increments the reference count for the given key.
// If this is the first reference to the key, the AddFunc is called.
func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {

View File

@@ -1,47 +0,0 @@
package refcounter
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored
// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive
// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes.
func TestReapplyMatching(t *testing.T) {
rc := New[netip.Prefix, string, string](
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
func(netip.Prefix, string) error { return nil },
)
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
_, err := rc.Increment(prefix, peer)
require.NoError(t, err)
}
// a second reference must not make the key applied twice
_, err := rc.Increment(peerA1, "peerA")
require.NoError(t, err)
var applied []netip.Prefix
err = rc.ReapplyMatching(
func(out string) bool { return out == "peerA" },
func(key netip.Prefix) error { applied = append(applied, key); return nil },
)
require.NoError(t, err)
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied)
var none []netip.Prefix
err = rc.ReapplyMatching(
func(out string) bool { return out == "missing" },
func(key netip.Prefix) error { none = append(none, key); return nil },
)
require.NoError(t, err)
assert.Empty(t, none)
}

View File

@@ -9,234 +9,25 @@ import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// per1k is a model's published USD rates per 1k tokens. read is the prompt-cache read rate
// (OpenAI: the cached-input discount rate); write is the cache-creation rate where one exists.
type per1k struct{ in, out, read, write float64 }
// publishedPer1k hardcodes the vendors' PUBLISHED rates for the models the live matrix can drive,
// keyed by the normalized model id the proxy stamps. Deliberately independent of the proxy's
// pricing table so a wrong embedded rate or a broken normalization fails the run.
var publishedPer1k = map[string]per1k{
"gpt-4o-mini": {0.00015, 0.0006, 0.000075, 0},
"gpt-4o": {0.0025, 0.01, 0.00125, 0},
"claude-haiku-4-5": {0.001, 0.005, 0.0001, 0.00125},
"claude-sonnet-4-5": {0.003, 0.015, 0.0003, 0.00375},
"claude-sonnet-4-6": {0.003, 0.015, 0.0003, 0.00375},
"kimi-k3": {0.003, 0.015, 0.0003, 0.003}, // no published write rate: bills at the input rate
"anthropic.claude-haiku-4-5": {0.001, 0.005, 0.0001, 0.00125},
"anthropic.claude-sonnet-4-5": {0.003, 0.015, 0.0003, 0.00375},
"anthropic.claude-sonnet-4-6": {0.003, 0.015, 0.0003, 0.00375},
}
// rawCostVerificationSQL is the operator-facing double-check, run straight against the management
// sqlite store: recompute each usage row's expected total and cache cost from its own persisted
// token buckets and hardcoded published rates. OpenAI counts cached tokens as a subset of input;
// Anthropic-shape providers count cache buckets additively.
const rawCostVerificationSQL = `
WITH rates(model, in_rate, out_rate, read_rate, write_rate) AS (
VALUES
('gpt-4o-mini', 0.00015, 0.0006, 0.000075, 0.0),
('gpt-4o', 0.0025, 0.01, 0.00125, 0.0),
('claude-haiku-4-5', 0.001, 0.005, 0.0001, 0.00125),
('claude-sonnet-4-5', 0.003, 0.015, 0.0003, 0.00375),
('claude-sonnet-4-6', 0.003, 0.015, 0.0003, 0.00375),
('kimi-k3', 0.003, 0.015, 0.0003, 0.003),
('anthropic.claude-haiku-4-5', 0.001, 0.005, 0.0001, 0.00125),
('anthropic.claude-sonnet-4-5', 0.003, 0.015, 0.0003, 0.00375),
('anthropic.claude-sonnet-4-6', 0.003, 0.015, 0.0003, 0.00375)
)
SELECT
u.provider,
u.model,
u.input_tokens,
u.output_tokens,
u.cached_input_tokens,
u.cache_creation_tokens,
u.input_cost_usd,
u.cached_input_cost_usd,
u.cache_creation_cost_usd,
u.output_cost_usd,
-- No cost_usd / cache_cost_usd columns are stored: both are derived from the
-- four per-bucket columns above, exactly as the API renders them.
(u.input_cost_usd + u.cached_input_cost_usd + u.cache_creation_cost_usd + u.output_cost_usd) AS cost_usd,
(u.cached_input_cost_usd + u.cache_creation_cost_usd) AS cache_cost_usd,
CASE WHEN u.provider = 'openai' THEN
(u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
ELSE
u.input_tokens*r.in_rate/1000.0
END AS expected_input,
CASE WHEN u.provider = 'openai' THEN
MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
ELSE
u.cached_input_tokens*r.read_rate/1000.0
END AS expected_cached_input,
CASE WHEN u.provider = 'openai' THEN
0.0
ELSE
u.cache_creation_tokens*r.write_rate/1000.0
END AS expected_cache_creation,
u.output_tokens*r.out_rate/1000.0 AS expected_output,
CASE WHEN u.provider = 'openai' THEN
(u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
+ MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
+ u.output_tokens*r.out_rate/1000.0
ELSE
u.input_tokens*r.in_rate/1000.0 + u.cached_input_tokens*r.read_rate/1000.0
+ u.cache_creation_tokens*r.write_rate/1000.0 + u.output_tokens*r.out_rate/1000.0
END AS expected_total,
CASE WHEN u.provider = 'openai' THEN
MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
ELSE
u.cached_input_tokens*r.read_rate/1000.0 + u.cache_creation_tokens*r.write_rate/1000.0
END AS expected_cache
FROM agent_network_request_usage u
JOIN rates r ON r.model = u.model
ORDER BY u.timestamp`
// verifyUsageRowsSQL re-checks every persisted usage row directly in the management sqlite store,
// bypassing the API path — the same audit an operator can run on a production store.db.
func verifyUsageRowsSQL(t *testing.T, srv *harness.Combined) {
t.Helper()
dbPath, err := srv.SnapshotStoreDB(t.TempDir())
require.NoError(t, err, "snapshot management sqlite store")
db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
require.NoError(t, err, "open store snapshot")
sqlDB, err := db.DB()
require.NoError(t, err)
defer func() { _ = sqlDB.Close() }()
rows, err := db.Raw(rawCostVerificationSQL).Rows()
require.NoError(t, err, "run raw cost verification query")
defer func() { _ = rows.Close() }()
verified := 0
for rows.Next() {
var provider, model string
var inTok, outTok, readTok, writeTok int64
var inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost float64
var wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache float64
require.NoError(t, rows.Scan(&provider, &model, &inTok, &outTok, &readTok, &writeTok,
&inCost, &cachedInCost, &cacheCreateCost, &outCost, &cost, &cacheCost,
&wantInput, &wantCachedInput, &wantCacheCreation, &wantOutput, &wantTotal, &wantCache), "scan usage row")
t.Logf("[sql] %s/%s: in=%d out=%d cache_read=%d cache_write=%d stored in/cached/create/out=$%.6f/$%.6f/$%.6f/$%.6f total=$%.6f cache=$%.6f expected total=$%.6f cache=$%.6f",
provider, model, inTok, outTok, readTok, writeTok,
inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost, wantTotal, wantCache)
assert.InDeltaf(t, wantInput, inCost, 1e-6, "stored input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCachedInput, cachedInCost, 1e-6, "stored cached_input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCacheCreation, cacheCreateCost, 1e-6, "stored cache_creation_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantOutput, outCost, 1e-6, "stored output_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantTotal, cost, 1e-6, "derived cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCache, cacheCost, 1e-6, "derived cache_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, inCost+cachedInCost+cacheCreateCost+outCost, cost, 1e-9,
"stored buckets must sum to the derived cost_usd for %s/%s", provider, model)
verified++
}
require.NoError(t, rows.Err(), "iterate usage rows")
require.Positive(t, verified, "raw SQL check must cover at least one usage row")
t.Logf("[sql] verified %d usage rows in store.db against published rates", verified)
gwRows, err := db.Raw(`SELECT model,
(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd) AS cost_usd
FROM agent_network_request_usage WHERE model LIKE '%/%'`).Rows()
require.NoError(t, err, "query gateway-prefixed usage rows")
defer func() { _ = gwRows.Close() }()
for gwRows.Next() {
var model string
var cost float64
require.NoError(t, gwRows.Scan(&model, &cost), "scan gateway usage row")
t.Logf("[sql] gateway %s: stored=$%.6f (must be 0 — deliberately unpriced)", model, cost)
assert.Zerof(t, cost, "gateway-prefixed model %q must store cost 0, never a guessed rate", model)
}
require.NoError(t, gwRows.Err(), "iterate gateway usage rows")
}
// validateAccessLogCost recomputes a live access-log row's expected total and cache cost from the
// published per-1k rates and the row's persisted token buckets, and asserts both stored values.
// Gateway-prefixed model ids the proxy deliberately does not price must store cost 0.
func validateAccessLogCost(t *testing.T, pc providerCase, row api.AgentNetworkAccessLog) {
t.Helper()
model := catalogModel(pc)
provider := ""
if row.Provider != nil {
provider = *row.Provider
}
t.Logf("[cost] %s: provider=%s model=%s in=%d out=%d total=%d cache_read=%d cache_write=%d cost=$%.6f cache_cost=$%.6f",
pc.name, provider, model, row.InputTokens, row.OutputTokens, row.TotalTokens,
row.CachedInputTokens, row.CacheCreationTokens, row.CostUsd, row.CacheCostUsd)
rates, known := publishedPer1k[model]
if !known {
if strings.Contains(model, "/") {
assert.Zerof(t, row.CostUsd, "gateway-prefixed model %q is not priced so the cost meter must skip (cost 0)", model)
return
}
t.Logf("[cost] %s: no published rate on file for model %q (env-overridden?); skipping cost validation", pc.name, model)
return
}
// input_tokens may legitimately be 0: Moonshot/Kimi reports fully cached prompts under the cache
// buckets only. Output and total must always be present on a priced row.
require.Positive(t, row.OutputTokens, "priced row must carry output tokens")
require.Positive(t, row.TotalTokens, "priced row must carry total tokens")
var wantInput, wantCachedInput, wantCacheCreation float64
if provider == "openai" {
cached := min(row.CachedInputTokens, row.InputTokens) // cached is a subset of input
wantInput = float64(row.InputTokens-cached) / 1000 * rates.in
wantCachedInput = float64(cached) / 1000 * rates.read
// OpenAI has no cache-write bucket; wantCacheCreation stays 0.
} else {
// Anthropic / Bedrock shape: cache buckets are additive to input_tokens.
wantInput = float64(row.InputTokens) / 1000 * rates.in
wantCachedInput = float64(row.CachedInputTokens) / 1000 * rates.read
wantCacheCreation = float64(row.CacheCreationTokens) / 1000 * rates.write
}
wantOutput := float64(row.OutputTokens) / 1000 * rates.out
wantCache := wantCachedInput + wantCacheCreation
wantTotal := wantInput + wantCache + wantOutput
t.Logf("[cost] %s: expecting input=$%.6f cached_input=$%.6f cache_creation=$%.6f output=$%.6f total=$%.6f cache=$%.6f from published rates",
pc.name, wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache)
assert.InDeltaf(t, wantInput, row.InputCostUsd, 1e-6, "stored input_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCachedInput, row.CachedInputCostUsd, 1e-6, "stored cached_input_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCacheCreation, row.CacheCreationCostUsd, 1e-6, "stored cache_creation_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantOutput, row.OutputCostUsd, 1e-6, "stored output_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantTotal, row.CostUsd, 1e-6, "derived cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCache, row.CacheCostUsd, 1e-6, "derived cache_cost_usd for %s (%s)", pc.name, model)
// The aggregates must be exactly the sum of the stored components, not an
// independently-computed figure that could drift from the breakdown.
assert.InDeltaf(t, row.InputCostUsd+row.CachedInputCostUsd+row.CacheCreationCostUsd+row.OutputCostUsd,
row.CostUsd, 1e-9, "stored buckets must sum to the derived cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, row.CachedInputCostUsd+row.CacheCreationCostUsd,
row.CacheCostUsd, 1e-9, "stored cache buckets must sum to the derived cache_cost_usd for %s (%s)", pc.name, model)
}
// providerCase is one entry in the live provider matrix. The same scenario runs
// for every available provider; availability is keyed off env vars so the suite
// covers whatever credentials are present (source ~/.llm-keys locally / set the
// Actions secrets in CI).
type providerCase struct {
name string
catalogID string
upstream string
apiKey string
model string // body model (chat/messages) or path model@version (vertex)
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
project string // vertex only: GCP project for the rawPredict path
region string // vertex only: GCP region for the rawPredict path
pathPrefix string // base-URL path prefix the agent carries (e.g. "/anthropic" for Kimi)
name string
catalogID string
upstream string
apiKey string
model string // body model (chat/messages) or path model@version (vertex)
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
project string // vertex only: GCP project for the rawPredict path
region string // vertex only: GCP region for the rawPredict path
}
// availableProviders builds the matrix from the provider env vars that are set.
@@ -248,24 +39,6 @@ func availableProviders() []providerCase {
if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" {
ps = append(ps, providerCase{name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k, model: "claude-haiku-4-5", kind: harness.WireMessages})
}
if k := os.Getenv("KIMI_TOKEN"); k != "" {
// Kimi (Moonshot AI) serves two body shapes from the same key: OpenAI
// Chat Completions on the bare host (/v1/...) and the Anthropic
// Messages API under the /anthropic path prefix (the endpoint
// Moonshot's Claude Code guide uses). The provider keeps the bare
// default upstream and the AGENT carries the /anthropic prefix in
// its base URL — exactly the documented Claude Code / Kimi CLI
// setup (ANTHROPIC_BASE_URL=https://<endpoint>/anthropic) — so one
// provider serves both shapes and the prefix rides through to
// Moonshot. Run the Anthropic shape, the flagship Claude Code path;
// the OpenAI wire shape is covered live by the other chat-shaped
// matrix providers, and Kimi-over-chat passed with kimi-k3 before
// the single-model constraint surfaced (run #73 on the kimi feature
// branch). The platform serves this account exactly ONE model —
// kimi-k3 (kimi-k2-thinking and even kimi-latest return
// resource_not_found_error on both surfaces).
ps = append(ps, providerCase{name: "kimi", catalogID: "kimi_api", upstream: "https://api.moonshot.ai", apiKey: k, model: "kimi-k3", kind: harness.WireMessages, pathPrefix: "/anthropic"})
}
if k, u := os.Getenv("VERCEL_TOKEN"), os.Getenv("VERCEL_URL"); k != "" && u != "" {
ps = append(ps, providerCase{name: "vercel", catalogID: "vercel_ai_gateway", upstream: u, apiKey: k, model: "openai/gpt-4o-mini", kind: harness.WireChat})
}
@@ -311,25 +84,19 @@ func availableProviders() []providerCase {
}
}
// Bedrock: path-routed, bearer auth. Model is the FULL cross-region
// inference-profile id exactly as AWS issues it — region-family prefix
// plus the date/version suffix. A bare or wrong-region id makes Bedrock
// reject the request with "The provided model identifier is invalid"
// before any inference runs. The proxy normalizes this id to the catalog
// key (anthropic.claude-haiku-4-5) for routing/pricing/allowlists.
// Defaults pair eu-central-1 with the eu.* profile; AWS_REGION overrides
// the region and the prefix follows its family.
// Bedrock: path-routed, bearer auth. Model is a cross-region inference
// profile id (distinct string from the first-party Anthropic case).
if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
region := os.Getenv("AWS_REGION")
if region == "" {
region = "eu-central-1"
region = "us-east-1"
}
// A valid Bedrock inference-profile id, overridable per account (AWS_BEDROCK_MODEL, also the
// workflow's bedrock_model dispatch input). `global.` profiles work from any region. Defaults to
// Sonnet 4.6, whose id convention dropped the -YYYYMMDD-v1:0 suffix that Haiku 4.5 still carries.
// A valid Bedrock inference-profile id (region prefix + date + version),
// overridable per account. `global.` profiles can be invoked from any
// region; set AWS_BEDROCK_MODEL to match the enabled profile for the token.
model := os.Getenv("AWS_BEDROCK_MODEL")
if model == "" {
model = "global.anthropic.claude-sonnet-4-6"
model = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
}
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: model, kind: harness.WireBedrock})
}
@@ -449,7 +216,8 @@ func TestProvidersMatrix(t *testing.T) {
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 (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
@@ -465,10 +233,6 @@ func TestProvidersMatrix(t *testing.T) {
// session id and confirm the marker propagated end-to-end.
sessionID := "e2e-session-" + pc.name
// A long-form prompt so completions carry realistic token counts for cost validation;
// max_tokens in the harness bodies (2048) lets the full answer through.
const matrixPrompt = "explain GitHub workflow in 1000 words"
// Retry briefly to absorb tunnel/DNS jitter on the first call.
var code int
var body string
@@ -479,11 +243,11 @@ func TestProvidersMatrix(t *testing.T) {
var cerr error
switch pc.kind {
case harness.WireVertex:
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, matrixPrompt, sessionID)
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
case harness.WireBedrock:
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, matrixPrompt, sessionID)
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
default:
c, b, cerr = cl.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, pc.kind, pc.model, matrixPrompt, sessionID)
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
}
if cerr == nil {
code, body = c, b
@@ -502,7 +266,6 @@ func TestProvidersMatrix(t *testing.T) {
// The session id sent as x-session-id must round-trip into the
// access-log row for this provider.
var row api.AgentNetworkAccessLog
require.Eventually(t, func() bool {
logs, lerr := srv.ListAccessLogs(ctx)
if lerr != nil {
@@ -510,15 +273,11 @@ func TestProvidersMatrix(t *testing.T) {
}
for _, r := range logs.Data {
if r.SessionId != nil && *r.SessionId == sessionID {
row = r
return true
}
}
return false
}, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row for %s", sessionID, pc.name)
// Stored total and cache cost must match the published rates applied to the row's buckets.
validateAccessLogCost(t, pc, row)
})
}
@@ -539,7 +298,4 @@ func TestProvidersMatrix(t *testing.T) {
}
return false
}, 60*time.Second, 3*time.Second, "consumption must be recorded with positive token counts after live traffic")
// Final raw-SQL audit: bypass the API and re-verify every persisted usage row in the store.
verifyUsageRowsSQL(t, srv)
}

View File

@@ -52,9 +52,7 @@ func catalogModel(pc providerCase) string {
func disallowedModel(pc providerCase) string {
switch pc.kind {
case harness.WireBedrock:
// Same profile prefix as the allowed model so only the model name
// differs; the guardrail must deny it before it reaches AWS.
return strings.SplitN(pc.model, ".", 2)[0] + ".anthropic.claude-opus-4-8"
return "us.anthropic.claude-opus-4-8"
case harness.WireVertex:
return "claude-opus-4-8@20250101"
default:
@@ -74,7 +72,7 @@ func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint,
case harness.WireVertex:
code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "")
default:
code, _, err = cl.ChatPrefixed(ctx, endpoint, proxyIP, pc.pathPrefix, pc.kind, model, "Reply with exactly: pong", "")
code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
}
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
return code
@@ -166,7 +164,8 @@ func TestModelAllowlistEnforced(t *testing.T) {
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 (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {

View File

@@ -104,7 +104,8 @@ func TestProviderSkipTLSVerification(t *testing.T) {
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 (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the 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 {

View File

@@ -106,7 +106,8 @@ func TestVLLMProvider(t *testing.T) {
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 (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the 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 {

View File

@@ -4,7 +4,6 @@ package harness
import (
"context"
"errors"
"fmt"
"io"
"os/exec"
@@ -168,54 +167,22 @@ func (cl *Client) pollStatus(ctx context.Context, timeout time.Duration, want st
return fmt.Errorf("timed out waiting for %q; last status:\n%s", want, last)
}
const (
// curlExitCouldNotResolve is curl's exit code for a DNS resolution failure, distinct from connection-level failures.
curlExitCouldNotResolve = 6
// dnsProbeRetryWindow bounds DNS-failure retries: the synthesized zone lands a beat after management connects, so early NXDOMAIN is propagation; a zone still absent after this window is a real failure.
dnsProbeRetryWindow = 30 * time.Second
dnsProbeRetryInterval = 2 * time.Second
)
// ResolveProxyIP GETs https://<endpoint>/ from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; only DNS failures retry, within dnsProbeRetryWindow. Returns the connected IP for --resolve pinning.
// ResolveProxyIP resolves the agent-network endpoint to the proxy peer's
// NetBird IP from inside the client (via magic DNS).
func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
args := []string{
"run", "--rm",
"--network", "container:" + cl.container.GetContainerID(),
curlImage,
"-ksS", "-o", "/dev/null",
"--connect-timeout", "30", "--max-time", "60",
"-w", "%{remote_ip}",
"https://" + endpoint + "/",
code, reader, err := cl.container.Exec(ctx, []string{"getent", "hosts", endpoint}, tcexec.Multiplexed())
if err != nil {
return "", err
}
deadline := time.Now().Add(dnsProbeRetryWindow)
for {
cmd := exec.CommandContext(ctx, "docker", args...)
var stdout, stderr strings.Builder
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err := cmd.Run()
if err == nil {
ip := strings.TrimSpace(stdout.String())
if ip == "" {
return "", fmt.Errorf("got an HTTP response from %s but no remote IP", endpoint)
}
return ip, nil
}
var exitErr *exec.ExitError
if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve {
return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String()))
}
dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String()))
if time.Until(deadline) < dnsProbeRetryInterval {
return "", dnsErr
}
select {
case <-ctx.Done():
return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err())
case <-time.After(dnsProbeRetryInterval):
}
out, _ := io.ReadAll(reader)
if code != 0 {
return "", fmt.Errorf("getent hosts %s exited %d", endpoint, code)
}
fields := strings.Fields(string(out))
if len(fields) == 0 {
return "", fmt.Errorf("no address for %s", endpoint)
}
return fields[0], nil
}
// Wire shapes for Chat.
@@ -239,29 +206,18 @@ const (
// the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). A non-empty
// sessionID is sent as the universal x-session-id header the proxy records.
func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) {
return cl.ChatPrefixed(ctx, endpoint, proxyIP, "", kind, model, prompt, sessionID)
}
// ChatPrefixed is Chat with a base-URL path prefix prepended to the wire
// path, mirroring agents whose base URL carries a shape-selecting prefix that
// rides through to the upstream — e.g. Claude Code against a Kimi provider
// sets ANTHROPIC_BASE_URL=https://<endpoint>/anthropic so the proxy forwards
// /anthropic/v1/messages to Moonshot's Anthropic surface while the provider's
// upstream URL stays the bare https://api.moonshot.ai. Empty prefix is plain
// Chat.
func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefix, kind, model, prompt, sessionID string) (int, string, error) {
var path, body string
var headers []string
switch kind {
case WireMessages:
path = "/v1/messages"
headers = []string{"anthropic-version: 2023-06-01"}
body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, model, prompt)
body = fmt.Sprintf(`{"model":%q,"max_tokens":64,"messages":[{"role":"user","content":%q}]}`, model, prompt)
default:
path = "/v1/chat/completions"
body = fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":%q}]}`, model, prompt)
}
return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID))
}
// Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
@@ -271,7 +227,7 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi
// is sent as the universal x-session-id header the proxy records.
func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt, sessionID string) (int, string, error) {
path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", project, region, model)
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
@@ -282,7 +238,7 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region
// header the proxy records.
func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) {
path := "/model/" + model + "/invoke"
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
@@ -341,7 +297,7 @@ func (cl *Client) Terminate(ctx context.Context) error {
return cl.container.Terminate(ctx)
}
// containerLogs reads up to 4 MiB of a container's logs for diagnostics — enough for a whole provider-matrix run.
// containerLogs reads up to 256 KiB of a container's logs for diagnostics.
func containerLogs(ctx context.Context, c testcontainers.Container) string {
if c == nil {
return ""
@@ -351,6 +307,6 @@ func containerLogs(ctx context.Context, c testcontainers.Container) string {
return fmt.Sprintf("<logs error: %v>", err)
}
defer r.Close()
b, _ := io.ReadAll(io.LimitReader(r, 4<<20))
b, _ := io.ReadAll(io.LimitReader(r, 256<<10))
return string(b)
}

View File

@@ -221,29 +221,6 @@ func (c *Combined) CreateProxyTokenCLI(ctx context.Context, name string) (string
return "", fmt.Errorf("token not found in CLI output: %s", string(out))
}
// SnapshotStoreDB copies the management sqlite store (with WAL/SHM sidecars) out of the bind-mounted
// data dir into dstDir and returns the copy's path; reading a copy avoids locking against live writes.
func (c *Combined) SnapshotStoreDB(dstDir string) (string, error) {
src := filepath.Join(c.workDir, "data", "store.db")
if _, err := os.Stat(src); err != nil {
return "", fmt.Errorf("management store not found at %s: %w", src, err)
}
dst := filepath.Join(dstDir, "store.db")
for _, suffix := range []string{"", "-wal", "-shm"} {
data, err := os.ReadFile(src + suffix)
if err != nil {
if os.IsNotExist(err) && suffix != "" {
continue // sidecar only exists in WAL mode
}
return "", fmt.Errorf("read %s: %w", src+suffix, err)
}
if err := os.WriteFile(dst+suffix, data, 0o600); err != nil {
return "", fmt.Errorf("write %s: %w", dst+suffix, err)
}
}
return dst, nil
}
// Logs returns the combined server container logs, for diagnostics.
func (c *Combined) Logs(ctx context.Context) string {
return containerLogs(ctx, c.container)

View File

@@ -234,6 +234,9 @@ init_environment() {
NETBIRD_LICENSE_KEY=$(read_secret "Enter license key (input hidden)")
GHCR_USERNAME="netbirdExtAccess1"
GHCR_TOKEN=$(read_secret "Enter GHCR token (input hidden)")
POSTGRES_USER="netbird"
POSTGRES_DB="netbird"
POSTGRES_PASSWORD=$(rand_secret)
@@ -260,6 +263,10 @@ init_environment() {
install -m 600 /dev/null config.yaml
render_config_yaml >> config.yaml
echo "Logging in to ghcr.io ..."
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
unset GHCR_TOKEN
echo ""
echo "Pulling images ..."
$DOCKER_COMPOSE_COMMAND pull

View File

@@ -490,6 +490,8 @@ init_migration() {
echo ""
echo "Step 1: Image swap (community → Enterprise). License key required."
NB_LICENSE_KEY=$(read_secret " License key")
GHCR_USERNAME="netbirdExtAccess1"
GHCR_TOKEN=$(read_secret " GHCR token (input hidden)")
# Step 2 — optional
echo ""
@@ -586,6 +588,11 @@ apply_changes() {
fi
} >> "$ENV_FILE"
echo ""
echo "Logging in to ghcr.io ..."
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
unset GHCR_TOKEN
echo ""
echo "Pulling enterprise images ..."
$DOCKER_COMPOSE_COMMAND pull

View File

@@ -18,26 +18,21 @@ import (
// contract between the proxy and management; management flattens them into
// queryable columns. Keep in sync with the proxy side.
const (
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCachedInputTokens = "llm.cached_input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDInput = "cost.usd_input"
metaKeyCostUSDCachedInput = "cost.usd_cached_input"
metaKeyCostUSDCacheCreate = "cost.usd_cache_creation"
metaKeyCostUSDOutput = "cost.usd_output"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDTotal = "cost.usd_total"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
)
// IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -113,25 +108,20 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
BytesUpload: e.BytesUpload,
BytesDownload: e.BytesDownload,
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CachedInputTokens: parseMetaInt(meta, metaKeyCachedInputTokens),
CacheCreationTokens: parseMetaInt(meta, metaKeyCacheCreationTokens),
InputCostUSD: parseMetaFloat(meta, metaKeyCostUSDInput),
CachedInputCostUSD: parseMetaFloat(meta, metaKeyCostUSDCachedInput),
CacheCreationCostUSD: parseMetaFloat(meta, metaKeyCostUSDCacheCreate),
OutputCostUSD: parseMetaFloat(meta, metaKeyCostUSDOutput),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
}
var groups []types.AgentNetworkAccessLogGroup
@@ -150,23 +140,18 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
// log's ID so the two correlate.
func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
usage := &types.AgentNetworkUsage{
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CachedInputTokens: e.CachedInputTokens,
CacheCreationTokens: e.CacheCreationTokens,
InputCostUSD: e.InputCostUSD,
CachedInputCostUSD: e.CachedInputCostUSD,
CacheCreationCostUSD: e.CacheCreationCostUSD,
OutputCostUSD: e.OutputCostUSD,
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CostUSD: e.CostUSD,
}
usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))

View File

@@ -28,22 +28,17 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
UserId: "user-1",
AgentNetwork: true,
Metadata: map[string]string{
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "1174",
metaKeyCachedInputTokens: "256",
metaKeyCacheCreationTokens: "768",
metaKeyCostUSDInput: "0.0071",
metaKeyCostUSDCachedInput: "0.0009",
metaKeyCostUSDCacheCreate: "0.0020",
metaKeyCostUSDOutput: "0.0023",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "150",
metaKeyCostUSDTotal: "0.0123",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
// repeated id must be de-duplicated before the group rows insert.
metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
},
@@ -70,19 +65,7 @@ func TestIngestAccessLog_RealStore_LogCollectionOff(t *testing.T) {
require.Len(t, usage, 1, "usage row must be written even with log collection off")
assert.Equal(t, int64(100), usage[0].InputTokens, "input tokens must round-trip from metadata")
assert.Equal(t, int64(50), usage[0].OutputTokens, "output tokens must round-trip from metadata")
assert.Equal(t, int64(256), usage[0].CachedInputTokens, "cache-read tokens must round-trip from metadata")
assert.Equal(t, int64(768), usage[0].CacheCreationTokens, "cache-write tokens must round-trip from metadata")
// The per-bucket breakdown is the only cost state stored, and must survive
// the write/read cycle as real columns — usage rows are the only cost
// record for accounts with log collection off, so a dropped column here
// loses the split permanently.
assert.InDelta(t, 0.0071, usage[0].InputCostUSD, 1e-9, "input cost must round-trip from metadata")
assert.InDelta(t, 0.0009, usage[0].CachedInputCostUSD, 1e-9, "cache-read cost must round-trip from metadata")
assert.InDelta(t, 0.0020, usage[0].CacheCreationCostUSD, 1e-9, "cache-write cost must round-trip from metadata")
assert.InDelta(t, 0.0023, usage[0].OutputCostUSD, 1e-9, "output cost must round-trip from metadata")
// Aggregates are derived from the stored columns, never stored themselves.
assert.InDelta(t, 0.0123, usage[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0029, usage[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
assert.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost must round-trip from metadata")
logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
require.NoError(t, err)
@@ -113,14 +96,6 @@ func TestIngestAccessLog_RealStore_LogCollectionOn(t *testing.T) {
require.Equal(t, int64(1), total, "exactly one access-log row expected")
require.Len(t, logs, 1, "full access-log row must be written when log collection is on")
assert.Equal(t, "gpt-5.4", logs[0].Model, "model must flatten from metadata")
assert.Equal(t, int64(256), logs[0].CachedInputTokens, "cache-read tokens must flatten from metadata")
assert.Equal(t, int64(768), logs[0].CacheCreationTokens, "cache-write tokens must flatten from metadata")
assert.InDelta(t, 0.0029, logs[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
assert.InDelta(t, 0.0123, logs[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0071, logs[0].InputCostUSD, 1e-9, "input cost must flatten from metadata")
assert.InDelta(t, 0.0009, logs[0].CachedInputCostUSD, 1e-9, "cache-read cost must flatten from metadata")
assert.InDelta(t, 0.0020, logs[0].CacheCreationCostUSD, 1e-9, "cache-write cost must flatten from metadata")
assert.InDelta(t, 0.0023, logs[0].OutputCostUSD, 1e-9, "output cost must flatten from metadata")
assert.Equal(t, "hello", logs[0].RequestPrompt, "prompt must be retained when log collection is on")
assert.Equal(t, "world", logs[0].ResponseCompletion, "completion must be retained when log collection is on")
assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")

View File

@@ -38,7 +38,7 @@ func accessLogRow(id, sessionID string, ts time.Time, opts ...func(*types.AgentN
InputTokens: 100,
OutputTokens: 50,
TotalTokens: 150,
InputCostUSD: 0.01,
CostUSD: 0.01,
}
for _, o := range opts {
o(e)
@@ -74,7 +74,7 @@ func withTokens(in, out, total int64, cost float64) func(*types.AgentNetworkAcce
e.InputTokens = in
e.OutputTokens = out
e.TotalTokens = total
e.InputCostUSD = cost
e.CostUSD = cost
}
}
@@ -155,7 +155,7 @@ func TestAccessLogSessions_FoldAndAggregate(t *testing.T) {
assert.Equal(t, int64(310), a.InputTokens, "input tokens summed")
assert.Equal(t, int64(135), a.OutputTokens, "output tokens summed")
assert.Equal(t, int64(445), a.TotalTokens, "total tokens summed")
assert.InDelta(t, 0.031, a.TotalCostUSD(), 1e-9, "cost summed")
assert.InDelta(t, 0.031, a.CostUSD, 1e-9, "cost summed")
assert.Equal(t, "deny", a.Decision, "any deny makes the session a deny")
assert.ElementsMatch(t, []string{"openai", "anthropic"}, a.Providers, "distinct providers")
assert.ElementsMatch(t, []string{"gpt-5.4", "claude-haiku-4-5"}, a.Models, "distinct models")

View File

@@ -420,47 +420,6 @@ var providers = []Provider{
{ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192},
},
},
{
ID: "kimi_api",
Kind: KindProvider,
Name: "Kimi (Moonshot AI) API",
Description: "Kimi K3 / K2 models via the Moonshot AI platform",
DefaultHost: "api.moonshot.ai",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#1A1A2E",
// ParserID empty on purpose: Moonshot serves two body shapes on
// the same host and key, and the proxy's URL sniffer dispatches
// both (same pattern as Bifrost). /v1/chat/completions matches
// OpenAIParser; the Anthropic-compatible endpoint the official
// Claude Code guide uses (/anthropic/v1/messages) contains
// "/v1/messages" and matches AnthropicParser. Pinning "openai"
// here would misparse the Claude Code path — the primary way
// teams consume Kimi for coding today. Both endpoints accept the
// same Moonshot key via Authorization: Bearer (Claude Code's
// ANTHROPIC_AUTH_TOKEN rides that header too).
//
// api.moonshot.ai is the international platform; mainland-China
// accounts live on api.moonshot.cn with separate billing —
// operators there override the host on the provider record. The
// kimi.com subscription coding endpoint (api.kimi.com/coding,
// model id "k3") is account-bound seat licensing rather than a
// meterable platform key, so it's deliberately not the default.
ParserID: "",
// Pricing per Moonshot's platform rates at K3 launch (July 2026):
// $3/$15 per MTok with $0.30 cached input, flat across the 1M-token
// window. kimi-k3 is the ONLY model the platform serves newer
// accounts — K2-era ids (kimi-k2-thinking) and even the kimi-latest
// alias return resource_not_found_error, verified live 2026-07-21 —
// so it's the only catalog entry. Grandfathered accounts with K2
// access can still type those ids on the provider's model rows.
// The consumer app's "K3 Swarm Max" mode is not an API SKU, so it
// doesn't appear here.
Models: []Model{
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
},
},
{
ID: "litellm_proxy",
Kind: KindGateway,

View File

@@ -41,24 +41,8 @@ type AgentNetworkAccessLog struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown — one column per token bucket the provider
// bills separately. These four are the only cost state stored: the total
// and the cache portion are derived on read (TotalCostUSD / CacheCostUSD)
// rather than stored alongside, so a stored aggregate can never drift out
// of step with the components it summarises.
//
// default:0 matters on upgrade: these columns are ALTER TABLE ADD COLUMN
// on an existing table, and without it every historical row holds NULL —
// which a raw SUM()/scan into float64 can't read. The default backfills
// them as 0, so pre-upgrade rows report an unknown split, not an error.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
Stream bool
CostUSD float64
Stream bool
// Prompt capture. Only populated when prompt collection is enabled
// (account master switch AND policy guardrail). Heavy free text.
@@ -76,44 +60,19 @@ type AgentNetworkAccessLog struct {
// the reverse-proxy AccessLogEntry table.
func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_log" }
// CostUSDSQLExpr is the SQL sum of the per-bucket cost columns — the total cost
// of a row. Used wherever a query has to sort or aggregate on total cost now
// that no cost_usd column is stored. Plain arithmetic over NOT NULL columns, so
// it stays portable across SQLite and Postgres.
const CostUSDSQLExpr = "(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd)"
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (a *AgentNetworkAccessLog) TotalCostUSD() float64 {
return a.InputCostUSD + a.CachedInputCostUSD + a.CacheCreationCostUSD + a.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets:
// cache reads plus cache writes.
func (a *AgentNetworkAccessLog) CacheCostUSD() float64 {
return a.CachedInputCostUSD + a.CacheCreationCostUSD
}
// ToAPIResponse renders the flattened entry as the API representation.
func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
out := api.AgentNetworkAccessLog{
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CachedInputTokens: a.CachedInputTokens,
CacheCreationTokens: a.CacheCreationTokens,
InputCostUsd: a.InputCostUSD,
CachedInputCostUsd: a.CachedInputCostUSD,
CacheCreationCostUsd: a.CacheCreationCostUSD,
OutputCostUsd: a.OutputCostUSD,
CostUsd: a.TotalCostUSD(),
CacheCostUsd: a.CacheCostUSD(),
Stream: &a.Stream,
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CostUsd: a.CostUSD,
Stream: &a.Stream,
}
out.UserId = strPtr(a.UserID)
@@ -153,36 +112,20 @@ func strPtr(s string) *string {
// summary plus its ordered entries. Assembled in Go from a page of entries — it
// is not a stored table.
type AgentNetworkAccessLogSession struct {
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
}
// TotalCostUSD is the session's total cost: the sum of the four per-bucket
// costs accumulated across its entries.
func (sess *AgentNetworkAccessLogSession) TotalCostUSD() float64 {
return sess.InputCostUSD + sess.CachedInputCostUSD + sess.CacheCreationCostUSD + sess.OutputCostUSD
}
// CacheCostUSD is the session's prompt-cache spend: cache reads plus writes.
func (sess *AgentNetworkAccessLogSession) CacheCostUSD() float64 {
return sess.CachedInputCostUSD + sess.CacheCreationCostUSD
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
}
// sessionKey is the grouping key for an entry: its session id, or — when the
@@ -262,12 +205,7 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
sess.InputTokens += e.InputTokens
sess.OutputTokens += e.OutputTokens
sess.TotalTokens += e.TotalTokens
sess.CachedInputTokens += e.CachedInputTokens
sess.CacheCreationTokens += e.CacheCreationTokens
sess.InputCostUSD += e.InputCostUSD
sess.CachedInputCostUSD += e.CachedInputCostUSD
sess.CacheCreationCostUSD += e.CacheCreationCostUSD
sess.OutputCostUSD += e.OutputCostUSD
sess.CostUSD += e.CostUSD
if e.Timestamp.Before(sess.StartedAt) {
sess.StartedAt = e.Timestamp
}
@@ -310,22 +248,15 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
}
out := api.AgentNetworkAccessLogSession{
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CachedInputTokens: sess.CachedInputTokens,
CacheCreationTokens: sess.CacheCreationTokens,
InputCostUsd: sess.InputCostUSD,
CachedInputCostUsd: sess.CachedInputCostUSD,
CacheCreationCostUsd: sess.CacheCreationCostUSD,
OutputCostUsd: sess.OutputCostUSD,
CostUsd: sess.TotalCostUSD(),
CacheCostUsd: sess.CacheCostUSD(),
Decision: sess.Decision,
Entries: entries,
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CostUsd: sess.CostUSD,
Decision: sess.Decision,
Entries: entries,
}
out.SessionId = strPtr(sess.SessionID)
out.UserId = strPtr(sess.UserID)

View File

@@ -54,7 +54,7 @@ var accessLogSortFields = map[string]string{
"provider": "provider",
"status_code": "status_code",
"duration": "duration",
"cost_usd": CostUSDSQLExpr,
"cost_usd": "cost_usd",
"total_tokens": "total_tokens",
"user_id": "user_id",
"decision": "decision",
@@ -70,7 +70,7 @@ var accessLogSortFields = map[string]string{
var sessionSortExprs = map[string]string{ //nolint:gosec // G101 false positive: "total_tokens" sort key, not a credential
"timestamp": "MAX(timestamp)",
"started_at": "MIN(timestamp)",
"cost_usd": "SUM" + CostUSDSQLExpr,
"cost_usd": "SUM(cost_usd)",
"total_tokens": "SUM(total_tokens)",
"duration": "SUM(duration)",
"request_count": "COUNT(*)",

View File

@@ -1,124 +0,0 @@
package types
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// costRow builds an access-log entry carrying only a cost breakdown — the rest
// of the row is irrelevant to the summation identities under test.
func costRow(id, session string, ts time.Time, in, cachedIn, cacheCreate, out float64) *AgentNetworkAccessLog {
return &AgentNetworkAccessLog{
ID: id,
SessionID: session,
Timestamp: ts,
InputCostUSD: in,
CachedInputCostUSD: cachedIn,
CacheCreationCostUSD: cacheCreate,
OutputCostUSD: out,
}
}
// TestAPIResponse_CostComponentsSumToAggregates is the contract a client adding
// up an API response depends on: within a single rendered object, the four
// per-bucket fields sum to cost_usd, and the two cache fields sum to
// cache_cost_usd. Uses rates that are not exactly representable in binary
// floating point, so the identity is checked against real arithmetic rather
// than round numbers.
func TestAPIResponse_CostComponentsSumToAggregates(t *testing.T) {
row := costRow("r1", "s1", time.Now(), 0.000768, 0.0002304, 0.00192, 0.003)
api := row.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "rendered buckets must sum to the rendered cost_usd")
assert.InDelta(t, api.CachedInputCostUsd+api.CacheCreationCostUsd, api.CacheCostUsd, 1e-12,
"rendered cache buckets must sum to the rendered cache_cost_usd")
assert.InDelta(t, 0.0059184, api.CostUsd, 1e-12, "total is the exact sum, not a separately rounded figure")
assert.InDelta(t, 0.0021504, api.CacheCostUsd, 1e-12, "cache cost is the exact sum of the two cache buckets")
}
// TestSessionSummary_SumsMatchSummedEntries proves a session summary equals the
// sum of the entries it renders: a client that adds up the entries itself must
// land on the same number the summary reports, per bucket and in total.
func TestSessionSummary_SumsMatchSummedEntries(t *testing.T) {
base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
entries := []*AgentNetworkAccessLog{
costRow("r1", "s1", base, 0.000768, 0.0002304, 0.00192, 0.003),
costRow("r2", "s1", base.Add(time.Minute), 0.000625, 0.0009375, 0, 0.005),
costRow("r3", "s1", base.Add(2*time.Minute), 0.0000016, 0, 0, 0.0000032),
}
sessions := FoldAccessLogSessions([]string{"s1"}, entries)
require.Len(t, sessions, 1)
sess := sessions[0].ToAPIResponse()
var wantInput, wantCachedInput, wantCacheCreation, wantOutput float64
for _, e := range entries {
wantInput += e.InputCostUSD
wantCachedInput += e.CachedInputCostUSD
wantCacheCreation += e.CacheCreationCostUSD
wantOutput += e.OutputCostUSD
}
assert.InDelta(t, wantInput, sess.InputCostUsd, 1e-12, "session input cost is the sum of its entries")
assert.InDelta(t, wantCachedInput, sess.CachedInputCostUsd, 1e-12, "session cache-read cost is the sum of its entries")
assert.InDelta(t, wantCacheCreation, sess.CacheCreationCostUsd, 1e-12, "session cache-write cost is the sum of its entries")
assert.InDelta(t, wantOutput, sess.OutputCostUsd, 1e-12, "session output cost is the sum of its entries")
assert.InDelta(t, wantInput+wantCachedInput+wantCacheCreation+wantOutput, sess.CostUsd, 1e-12,
"session total equals the summed entry buckets")
// Summing the rendered entries must give the same answer as reading the
// summary — the property a UI relies on when it totals a table itself.
var fromEntries float64
for _, e := range sess.Entries {
fromEntries += e.CostUsd
}
assert.InDelta(t, sess.CostUsd, fromEntries, 1e-12, "summary total must match the summed rendered entries")
// The sub-microdollar row must still contribute; it would vanish under
// 6-decimal quantisation.
assert.Greater(t, sess.InputCostUsd, 0.001393, "small-cost rows must not be quantised away")
}
// TestUsageBuckets_SumsMatchSummedRows proves the same identity one level up:
// a usage bucket equals the sum of the ledger rows folded into it, and the
// buckets together equal the whole range.
func TestUsageBuckets_SumsMatchSummedRows(t *testing.T) {
day1 := time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC)
day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
rows := []*AgentNetworkUsage{
{ID: "u1", Timestamp: day1, InputCostUSD: 0.000768, CachedInputCostUSD: 0.0002304, CacheCreationCostUSD: 0.00192, OutputCostUSD: 0.003},
{ID: "u2", Timestamp: day1.Add(time.Hour), InputCostUSD: 0.000625, CachedInputCostUSD: 0.0009375, OutputCostUSD: 0.005},
{ID: "u3", Timestamp: day2, InputCostUSD: 0.0000016, OutputCostUSD: 0.0000032},
}
buckets := AggregateUsageByGranularity(rows, UsageGranularityDay)
require.Len(t, buckets, 2, "two distinct days expected")
var total, cache float64
for _, b := range buckets {
api := b.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "each bucket's components must sum to its cost_usd")
total += api.CostUsd
cache += api.CacheCostUsd
}
var wantTotal, wantCache float64
for _, r := range rows {
wantTotal += r.TotalCostUSD()
wantCache += r.CacheCostUSD()
}
assert.InDelta(t, wantTotal, total, 1e-12, "buckets must sum to the total across all ledger rows")
assert.InDelta(t, wantCache, cache, 1e-12, "buckets must sum to the cache spend across all ledger rows")
// A month bucket over the same rows must total identically — regrouping
// changes the partition, never the sum.
monthly := AggregateUsageByGranularity(rows, UsageGranularityMonth)
require.Len(t, monthly, 1)
assert.InDelta(t, wantTotal, monthly[0].ToAPIResponse().CostUsd, 1e-12,
"re-bucketing at a different granularity must preserve the total")
}

View File

@@ -25,19 +25,8 @@ type AgentNetworkUsage struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown, mirroring AgentNetworkAccessLog — the only
// cost state stored; total and cache portion are derived on read. Kept on
// the usage ledger too so spend can be attributed per bucket even for
// accounts with log collection turned off. See AgentNetworkAccessLog for
// why the columns carry a zero default.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
CreatedAt time.Time
CostUSD float64
CreatedAt time.Time
}
// TableName keeps usage records in their own stripped table. Named
@@ -45,17 +34,6 @@ type AgentNetworkUsage struct {
// agent_network_usage table in a shared database.
func (AgentNetworkUsage) TableName() string { return "agent_network_request_usage" }
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (u *AgentNetworkUsage) TotalCostUSD() float64 {
return u.InputCostUSD + u.CachedInputCostUSD + u.CacheCreationCostUSD + u.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets.
func (u *AgentNetworkUsage) CacheCostUSD() float64 {
return u.CachedInputCostUSD + u.CacheCreationCostUSD
}
// AgentNetworkUsageGroup is the normalised many-to-many row linking a usage
// record to one authorising group, mirroring AgentNetworkAccessLogGroup so the
// usage overview can filter by group with a `group_id IN (...)` join.

View File

@@ -33,45 +33,21 @@ func ParseUsageGranularity(s string) UsageGranularity {
// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
// the UTC start of the bucket as YYYY-MM-DD.
type AgentNetworkUsageBucket struct {
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
}
// TotalCostUSD is the bucket's total spend: the sum of the four per-bucket
// costs. Derived rather than accumulated separately so it cannot disagree with
// the components.
func (b *AgentNetworkUsageBucket) TotalCostUSD() float64 {
return b.InputCostUSD + b.CachedInputCostUSD + b.CacheCreationCostUSD + b.OutputCostUSD
}
// CacheCostUSD is the bucket's prompt-cache spend: cache reads plus writes.
func (b *AgentNetworkUsageBucket) CacheCostUSD() float64 {
return b.CachedInputCostUSD + b.CacheCreationCostUSD
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
}
// ToAPIResponse renders the bucket as the API representation.
func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
return api.AgentNetworkUsageBucket{
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CachedInputTokens: b.CachedInputTokens,
CacheCreationTokens: b.CacheCreationTokens,
InputCostUsd: b.InputCostUSD,
CachedInputCostUsd: b.CachedInputCostUSD,
CacheCreationCostUsd: b.CacheCreationCostUSD,
OutputCostUsd: b.OutputCostUSD,
CostUsd: b.TotalCostUSD(),
CacheCostUsd: b.CacheCostUSD(),
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CostUsd: b.CostUSD,
}
}
@@ -108,12 +84,7 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
b.InputTokens += r.InputTokens
b.OutputTokens += r.OutputTokens
b.TotalTokens += r.TotalTokens
b.CachedInputTokens += r.CachedInputTokens
b.CacheCreationTokens += r.CacheCreationTokens
b.InputCostUSD += r.InputCostUSD
b.CachedInputCostUSD += r.CachedInputCostUSD
b.CacheCreationCostUSD += r.CacheCreationCostUSD
b.OutputCostUSD += r.OutputCostUSD
b.CostUSD += r.CostUSD
}
out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))

View File

@@ -20,17 +20,15 @@ import (
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/keepalive"
cachestore "github.com/eko/gocache/lib/v4/store"
"github.com/netbirdio/netbird/encryption"
"github.com/netbirdio/netbird/formatter/hook"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/activity"
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
nbcache "github.com/netbirdio/netbird/management/server/cache"
nbContext "github.com/netbirdio/netbird/management/server/context"
nbhttp "github.com/netbirdio/netbird/management/server/http"
@@ -70,8 +68,8 @@ func (s *BaseServer) Metrics() telemetry.AppMetrics {
// CacheStore returns a shared cache store backed by Redis or in-memory depending on the environment.
// All consumers should reuse this store to avoid creating multiple Redis connections.
func (s *BaseServer) CacheStore() cachestore.StoreInterface {
return Create(s, func() cachestore.StoreInterface {
func (s *BaseServer) CacheStore() nbcache.Store {
return Create(s, func() nbcache.Store {
cs, err := nbcache.NewStore(context.Background(), nbcache.DefaultStoreMaxTimeout, nbcache.DefaultStoreCleanupInterval, nbcache.DefaultStoreMaxConn)
if err != nil {
log.Fatalf("failed to create shared cache store: %v", err)

View File

@@ -5,6 +5,9 @@ import (
"strconv"
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap"
@@ -163,7 +166,7 @@ func (e *componentEncoder) indexAllPeers() {
}
}
func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
func (e *componentEncoder) appendPeer(p *nbpeer.Peer) uint32 {
if idx, ok := e.peerOrder[p.ID]; ok {
return idx
}
@@ -177,7 +180,7 @@ func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
// (c.RouterPeers may contain peers not in c.Peers when validation rules drop
// them) and returns their wire indexes for the RouterPeerIndexes field. Must
// run before any encoder that resolves peer ids via e.peerOrder.
func (e *componentEncoder) indexRouterPeers(routers map[string]*types.ComponentPeer) []uint32 {
func (e *componentEncoder) indexRouterPeers(routers map[string]*nbpeer.Peer) []uint32 {
if len(routers) == 0 {
return nil
}
@@ -511,7 +514,7 @@ func encodeCustomZones(zones []nbdns.CustomZone) []*proto.CustomZone {
return out
}
func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentResource) []*proto.NetworkResourceRaw {
func (e *componentEncoder) encodeNetworkResources(resources []*resourceTypes.NetworkResource) []*proto.NetworkResourceRaw {
if len(resources) == 0 {
return nil
}
@@ -540,7 +543,7 @@ func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentRe
return out
}
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*types.ComponentRouter) map[string]*proto.NetworkRouterList {
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*routerTypes.NetworkRouter) map[string]*proto.NetworkRouterList {
if len(routersMap) == 0 {
return nil
}
@@ -689,20 +692,20 @@ func toAccountNetwork(n *types.Network) *proto.AccountNetwork {
return out
}
func toPeerCompact(p *types.ComponentPeer) *proto.PeerCompact {
func toPeerCompact(p *nbpeer.Peer) *proto.PeerCompact {
pc := &proto.PeerCompact{
WgPubKey: decodeWgKey(p.Key),
SshPubKey: []byte(p.SSHKey),
DnsLabel: p.DNSLabel,
AgentVersion: p.AgentVersion,
AddedWithSsoLogin: p.AddedWithSSOLogin,
AgentVersion: p.Meta.WtVersion,
AddedWithSsoLogin: p.UserID != "",
LoginExpirationEnabled: p.LoginExpirationEnabled,
SshEnabled: p.SSHEnabled,
SupportsIpv6: p.SupportsIPv6,
SupportsSourcePrefixes: p.SupportsSourcePrefixes,
ServerSshAllowed: p.ServerSSHAllowed,
SupportsIpv6: p.SupportsIPv6(),
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
ServerSshAllowed: p.Meta.Flags.ServerSSHAllowed,
}
if !p.LastLogin.IsZero() {
if p.LastLogin != nil {
pc.LastLoginUnixNano = p.LastLogin.UnixNano()
}
switch {

View File

@@ -15,6 +15,9 @@ import (
goproto "google.golang.org/protobuf/proto"
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
@@ -152,28 +155,29 @@ func envelopesEquivalent(a, b *proto.NetworkMapEnvelope) bool {
}
func newTestComponents() *types.NetworkMapComponents {
peerA := &types.ComponentPeer{
ID: "peer-a",
Key: testWgKeyA,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peera",
SSHKey: "ssh-a",
AgentVersion: "0.40.0",
peerA := &nbpeer.Peer{
ID: "peer-a",
Key: testWgKeyA,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peera",
SSHKey: "ssh-a",
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now()},
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
peerB := &types.ComponentPeer{
ID: "peer-b",
Key: testWgKeyB,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
DNSLabel: "peerb",
AgentVersion: "0.25.0",
peerB := &nbpeer.Peer{
ID: "peer-b",
Key: testWgKeyB,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
DNSLabel: "peerb",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.25.0"},
}
peerC := &types.ComponentPeer{
ID: "peer-c",
Key: testWgKeyC,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc",
AgentVersion: "0.40.0",
peerC := &nbpeer.Peer{
ID: "peer-c",
Key: testWgKeyC,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
return &types.NetworkMapComponents{
@@ -187,12 +191,12 @@ func newTestComponents() *types.NetworkMapComponents {
PeerLoginExpirationEnabled: true,
PeerLoginExpiration: 2 * time.Hour,
},
Peers: map[string]*types.ComponentPeer{
Peers: map[string]*nbpeer.Peer{
"peer-a": peerA,
"peer-b": peerB,
"peer-c": peerC,
},
Groups: map[string]*types.ComponentGroup{
Groups: map[string]*types.Group{
"group-src": {ID: "group-src", PublicID: "1", Name: "Src", Peers: []string{"peer-a"}},
"group-dst": {ID: "group-dst", PublicID: "2", Name: "Dst", Peers: []string{"peer-b", "peer-c"}},
},
@@ -211,7 +215,7 @@ func newTestComponents() *types.NetworkMapComponents {
}},
},
},
RouterPeers: map[string]*types.ComponentPeer{"peer-c": peerC},
RouterPeers: map[string]*nbpeer.Peer{"peer-c": peerC},
}
}
@@ -377,12 +381,12 @@ func TestEncodeNetworkMapEnvelope_MalformedWgKey(t *testing.T) {
func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
c := newTestComponents()
v6Only := &types.ComponentPeer{
ID: "peer-v6",
Key: testWgKeyA,
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
DNSLabel: "peerv6",
AgentVersion: "0.40.0",
v6Only := &nbpeer.Peer{
ID: "peer-v6",
Key: testWgKeyA,
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
DNSLabel: "peerv6",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
c.Peers["peer-v6"] = v6Only
@@ -401,11 +405,11 @@ func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
func TestEncodeNetworkMapEnvelope_PeerWithoutIP(t *testing.T) {
c := newTestComponents()
c.Peers["peer-noip"] = &types.ComponentPeer{
ID: "peer-noip",
Key: testWgKeyA,
DNSLabel: "peernoip",
AgentVersion: "0.40.0",
c.Peers["peer-noip"] = &nbpeer.Peer{
ID: "peer-noip",
Key: testWgKeyA,
DNSLabel: "peernoip",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -440,9 +444,9 @@ func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
func TestEncodeNetworkMapEnvelope_PeerLoginExpirationFields(t *testing.T) {
c := newTestComponents()
now := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
c.Peers["peer-a"].AddedWithSSOLogin = true
c.Peers["peer-a"].UserID = "user-1"
c.Peers["peer-a"].LoginExpirationEnabled = true
c.Peers["peer-a"].LastLogin = now
c.Peers["peer-a"].LastLogin = &now
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -553,7 +557,7 @@ func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing
}
// Resource must appear in components.NetworkResources with a seq id —
// encoder uses that to translate the xid map key to uint32.
c.NetworkResources = []*types.ComponentResource{
c.NetworkResources = []*resourceTypes.NetworkResource{
{ID: "resource-x", PublicID: "77", Name: "res-x", Enabled: true},
}
@@ -621,11 +625,11 @@ func TestEncodeNetworkMapEnvelope_PostureFailedPeers(t *testing.T) {
func TestEncodeNetworkMapEnvelope_RoutersMap(t *testing.T) {
c := newTestComponents()
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
c.RoutersMap = map[string]map[string]*routerTypes.NetworkRouter{
"net-1": {
"peer-c": {
PublicID: "200",
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
ID: "router-1", PublicID: "200",
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
},
},
}
@@ -651,14 +655,14 @@ func TestEncodeNetworkMapEnvelope_RouterPeerNotInComponentsPeers(t *testing.T) {
// peer_index reference must still resolve.
c := newTestComponents()
delete(c.Peers, "peer-c")
routerPeer := &types.ComponentPeer{
routerPeer := &nbpeer.Peer{
ID: "peer-c", Key: testWgKeyC, IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc", AgentVersion: "0.40.0",
DNSLabel: "peerc", Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
c.RouterPeers = map[string]*types.ComponentPeer{"peer-c": routerPeer}
c.RouterPeers = map[string]*nbpeer.Peer{"peer-c": routerPeer}
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
"net-1": {"peer-c": {PublicID: "1", Peer: "peer-c", Enabled: true}},
c.RoutersMap = map[string]map[string]*routerTypes.NetworkRouter{
"net-1": {"peer-c": {ID: "r-1", PublicID: "1", Peer: "peer-c", Enabled: true}},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -691,9 +695,9 @@ func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
nm := &types.NetworkMap{
Peers: []*types.ComponentPeer{{
Peers: []*nbpeer.Peer{{
ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
DNSLabel: "extpeer", AgentVersion: "0.40.0",
DNSLabel: "extpeer", Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}},
FirewallRules: []*types.FirewallRule{{
PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
@@ -776,6 +780,6 @@ func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
func emptyNetworkMapComponents() *types.NetworkMapComponents {
return types.EmptyNetworkMapComponents(
&types.NetworkMapComponents{
PeerID: "peer-id", Peers: map[string]*types.ComponentPeer{"peer-id": {}}},
PeerID: "peer-id", Peers: map[string]*nbpeer.Peer{"peer-id": {}}},
)
}

View File

@@ -19,10 +19,10 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
mkComponents := func(rule *types.PolicyRule, sshEnabled bool) (*types.NetworkMapComponents, *nbpeer.Peer) {
peer := &nbpeer.Peer{ID: targetPeerID, SSHEnabled: sshEnabled}
group := &types.ComponentGroup{ID: targetGroupID, Name: "dst", Peers: []string{targetPeerID}}
group := &types.Group{ID: targetGroupID, Name: "dst", Peers: []string{targetPeerID}}
return &types.NetworkMapComponents{
Peers: map[string]*types.ComponentPeer{targetPeerID: peer.ToComponent()},
Groups: map[string]*types.ComponentGroup{targetGroupID: group},
Peers: map[string]*nbpeer.Peer{targetPeerID: peer},
Groups: map[string]*types.Group{targetGroupID: group},
Policies: []*types.Policy{{
ID: "p",
Enabled: true,
@@ -158,8 +158,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
func TestComputeSSHEnabledForPeer_TargetMissingFromComponents(t *testing.T) {
peer := &nbpeer.Peer{ID: "missing", SSHEnabled: true}
c := &types.NetworkMapComponents{
Peers: map[string]*types.ComponentPeer{}, // target peer NOT present
Groups: map[string]*types.ComponentGroup{
Peers: map[string]*nbpeer.Peer{}, // target peer NOT present
Groups: map[string]*types.Group{
"g": {ID: "g", Peers: []string{"missing"}},
},
Policies: []*types.Policy{{

View File

@@ -7,9 +7,6 @@ import (
"errors"
"fmt"
"time"
"github.com/eko/gocache/lib/v4/cache"
"github.com/eko/gocache/lib/v4/store"
)
const (
@@ -22,12 +19,17 @@ var (
ErrTokenExpired = errors.New("JWT expired")
)
type SessionStore struct {
cache *cache.Cache[string]
// TokenCache atomically records used JWTs until their expiration.
type TokenCache interface {
SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error)
}
func NewSessionStore(cacheStore store.StoreInterface) *SessionStore {
return &SessionStore{cache: cache.New[string](cacheStore)}
type SessionStore struct {
cache TokenCache
}
func NewSessionStore(cacheStore TokenCache) *SessionStore {
return &SessionStore{cache: cacheStore}
}
// RegisterToken records a JWT until its exp time and rejects reuse.
@@ -38,19 +40,13 @@ func (s *SessionStore) RegisterToken(ctx context.Context, token string, expiresA
}
key := usedTokenKeyPrefix + hashToken(token)
_, err := s.cache.Get(ctx, key)
if err == nil {
return ErrTokenAlreadyUsed
}
var notFound *store.NotFound
if !errors.As(err, &notFound) {
return fmt.Errorf("failed to lookup used token entry: %w", err)
}
if err := s.cache.Set(ctx, key, usedTokenMarker, store.WithExpiration(ttl)); err != nil {
created, err := s.cache.SetNX(ctx, key, usedTokenMarker, ttl)
if err != nil {
return fmt.Errorf("failed to store used token entry: %w", err)
}
if !created {
return ErrTokenAlreadyUsed
}
return nil
}

View File

@@ -2,6 +2,7 @@ package auth
import (
"context"
"errors"
"testing"
"time"
@@ -38,6 +39,39 @@ func TestSessionStore_RegisterSameTokenTwiceIsRejected(t *testing.T) {
assert.ErrorIs(t, err, ErrTokenAlreadyUsed)
}
func TestSessionStore_ConcurrentRegistrationAllowsOneCaller(t *testing.T) {
s := newTestSessionStore(t)
ctx := context.Background()
const attempts = 100
start := make(chan struct{})
results := make(chan error, attempts)
for range attempts {
go func() {
<-start
results <- s.RegisterToken(ctx, "token", time.Now().Add(time.Hour))
}()
}
close(start)
succeeded := 0
alreadyUsed := 0
for range attempts {
err := <-results
switch {
case err == nil:
succeeded++
case errors.Is(err, ErrTokenAlreadyUsed):
alreadyUsed++
default:
require.NoError(t, err)
}
}
assert.Equal(t, 1, succeeded)
assert.Equal(t, attempts-1, alreadyUsed)
}
func TestSessionStore_RegisterDifferentTokensAreIndependent(t *testing.T) {
s := newTestSessionStore(t)
ctx := context.Background()
@@ -72,6 +106,23 @@ func TestSessionStore_EntryEvictsAtTTLAndAllowsReRegistration(t *testing.T) {
require.NoError(t, s.RegisterToken(ctx, token, time.Now().Add(time.Hour)))
}
type failingTokenCache struct {
err error
}
func (f failingTokenCache) SetNX(context.Context, string, string, time.Duration) (bool, error) {
return false, f.err
}
func TestSessionStore_CacheErrorIsReturned(t *testing.T) {
cacheErr := errors.New("cache unavailable")
s := NewSessionStore(failingTokenCache{err: cacheErr})
err := s.RegisterToken(context.Background(), "token", time.Now().Add(time.Hour))
require.Error(t, err)
assert.ErrorIs(t, err, cacheErr)
}
func TestHashToken_StableAndDoesNotLeak(t *testing.T) {
a := hashToken("tokenA")
b := hashToken("tokenB")

31
management/server/cache/memory.go vendored Normal file
View File

@@ -0,0 +1,31 @@
package cache
import (
"context"
"time"
"github.com/eko/gocache/lib/v4/store"
gocachestore "github.com/eko/gocache/store/go_cache/v4"
gocache "github.com/patrickmn/go-cache"
)
type goCacheStore struct {
store.StoreInterface
client *gocache.Cache
}
func newMemoryStore(maxTimeout, cleanupInterval time.Duration) Store {
client := gocache.New(maxTimeout, cleanupInterval)
return &goCacheStore{
StoreInterface: gocachestore.NewGoCache(client),
client: client,
}
}
func (s *goCacheStore) SetNX(_ context.Context, key, value string, ttl time.Duration) (bool, error) {
// Add only returns an error when a non-expired entry already exists.
if err := s.client.Add(key, value, ttl); err != nil {
return false, nil
}
return true, nil
}

49
management/server/cache/memory_test.go vendored Normal file
View File

@@ -0,0 +1,49 @@
package cache_test
import (
"context"
"testing"
"time"
"github.com/netbirdio/netbird/management/server/cache"
)
func TestMemoryStore(t *testing.T) {
memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
if err != nil {
t.Fatalf("couldn't create memory store: %s", err)
}
ctx := context.Background()
key, value := "testing", "tested"
err = memStore.Set(ctx, key, value)
if err != nil {
t.Errorf("couldn't set testing data: %s", err)
}
result, err := memStore.Get(ctx, key)
if err != nil {
t.Errorf("couldn't get testing data: %s", err)
}
if value != result.(string) {
t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
}
created, err := memStore.SetNX(ctx, "atomic", value, 100*time.Millisecond)
if err != nil {
t.Fatalf("couldn't atomically set testing data: %s", err)
}
if !created {
t.Fatal("first atomic set should create the entry")
}
created, err = memStore.SetNX(ctx, "atomic", value, 100*time.Millisecond)
if err != nil {
t.Fatalf("couldn't atomically check testing data: %s", err)
}
if created {
t.Fatal("second atomic set should not replace the entry")
}
// test expiration
time.Sleep(300 * time.Millisecond)
_, err = memStore.Get(ctx, key)
if err == nil {
t.Error("value should not be found")
}
}

51
management/server/cache/redis.go vendored Normal file
View File

@@ -0,0 +1,51 @@
package cache
import (
"context"
"fmt"
"math"
"time"
"github.com/eko/gocache/lib/v4/store"
redisstore "github.com/eko/gocache/store/redis/v4"
"github.com/redis/go-redis/v9"
log "github.com/sirupsen/logrus"
)
type redisStore struct {
store.StoreInterface
client *redis.Client
}
func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (Store, error) {
options, err := redis.ParseURL(redisEnvAddr)
if err != nil {
return nil, fmt.Errorf("parsing redis cache url: %s", err)
}
options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns
options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns
options.MaxActiveConns = maxConn
options.ConnMaxIdleTime = 30 * time.Minute
options.ConnMaxLifetime = 0
options.PoolTimeout = 10 * time.Second
redisClient := redis.NewClient(options)
subCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
_, err = redisClient.Ping(subCtx).Result()
if err != nil {
return nil, err
}
log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr)
return &redisStore{
StoreInterface: redisstore.NewRedis(redisClient),
client: redisClient,
}, nil
}
func (s *redisStore) SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error) {
return s.client.SetNX(ctx, key, value, ttl).Result()
}

View File

@@ -12,32 +12,6 @@ import (
"github.com/netbirdio/netbird/management/server/cache"
)
func TestMemoryStore(t *testing.T) {
memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
if err != nil {
t.Fatalf("couldn't create memory store: %s", err)
}
ctx := context.Background()
key, value := "testing", "tested"
err = memStore.Set(ctx, key, value)
if err != nil {
t.Errorf("couldn't set testing data: %s", err)
}
result, err := memStore.Get(ctx, key)
if err != nil {
t.Errorf("couldn't get testing data: %s", err)
}
if value != result.(string) {
t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
}
// test expiration
time.Sleep(300 * time.Millisecond)
_, err = memStore.Get(ctx, key)
if err == nil {
t.Error("value should not be found")
}
}
func TestRedisStoreConnectionFailure(t *testing.T) {
t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379")
_, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond, 100)
@@ -94,6 +68,47 @@ func TestRedisStoreConnectionSuccess(t *testing.T) {
if value != r {
t.Errorf("value returned from redis doesn't match testing data, got %s, expected %s", r, value)
}
secondRedisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
if err != nil {
t.Fatalf("couldn't create second redis store: %s", err)
}
start := make(chan struct{})
type setResult struct {
created bool
err error
}
results := make(chan setResult, 2)
for _, cacheStore := range []cache.Store{redisStore, secondRedisStore} {
go func() {
<-start
created, err := cacheStore.SetNX(ctx, "atomic", value, time.Second)
results <- setResult{created: created, err: err}
}()
}
close(start)
created := 0
for range 2 {
result := <-results
if result.err != nil {
t.Fatalf("atomic redis set failed: %s", result.err)
}
if result.created {
created++
}
}
if created != 1 {
t.Fatalf("expected exactly one redis client to create the entry, got %d", created)
}
ttl, err := redisClient.PTTL(ctx, "atomic").Result()
if err != nil {
t.Fatalf("couldn't read atomic entry TTL: %s", err)
}
if ttl <= 0 {
t.Fatalf("atomic entry should have a positive TTL, got %s", ttl)
}
// test expiration
time.Sleep(300 * time.Millisecond)
_, err = redisStore.Get(ctx, key)

View File

@@ -2,17 +2,10 @@ package cache
import (
"context"
"fmt"
"math"
"os"
"time"
"github.com/eko/gocache/lib/v4/store"
gocache_store "github.com/eko/gocache/store/go_cache/v4"
redis_store "github.com/eko/gocache/store/redis/v4"
gocache "github.com/patrickmn/go-cache"
"github.com/redis/go-redis/v9"
log "github.com/sirupsen/logrus"
)
// RedisStoreEnvVar is the environment variable that determines if a redis store should be used.
@@ -31,15 +24,21 @@ const (
DefaultStoreMaxConn = 1000
)
// Store extends the shared cache interface with atomic insertion support.
type Store interface {
store.StoreInterface
// SetNX atomically stores a value with a TTL only when the key does not exist.
SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error)
}
// NewStore creates a new cache store with the given max timeout and cleanup interval. It checks for the environment Variable RedisStoreEnvVar
// to determine if a redis store should be used. If the environment variable is set, it will attempt to connect to the redis store.
func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (store.StoreInterface, error) {
func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (Store, error) {
redisAddr := GetAddrFromEnv()
if redisAddr != "" {
return getRedisStore(ctx, redisAddr, maxConn)
}
goc := gocache.New(maxTimeout, cleanupInterval)
return gocache_store.NewGoCache(goc), nil
return newMemoryStore(maxTimeout, cleanupInterval), nil
}
// GetAddrFromEnv returns the redis address from the environment variable RedisStoreEnvVar or its legacy counterpart.
@@ -50,29 +49,3 @@ func GetAddrFromEnv() string {
}
return addr
}
func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (store.StoreInterface, error) {
options, err := redis.ParseURL(redisEnvAddr)
if err != nil {
return nil, fmt.Errorf("parsing redis cache url: %s", err)
}
options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns
options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns
options.MaxActiveConns = maxConn
options.ConnMaxIdleTime = 30 * time.Minute
options.ConnMaxLifetime = 0
options.PoolTimeout = 10 * time.Second
redisClient := redis.NewClient(options)
subCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
_, err = redisClient.Ping(subCtx).Result()
if err != nil {
return nil, err
}
log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr)
return redis_store.NewRedis(redisClient), nil
}

View File

@@ -6,7 +6,6 @@ import (
"github.com/netbirdio/netbird/management/server/account"
"github.com/netbirdio/netbird/management/server/activity"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/permissions/modules"
"github.com/netbirdio/netbird/management/server/permissions/operations"
@@ -31,10 +30,6 @@ type managerImpl struct {
accountManager account.Manager
}
func eventMetaResource(group *types.Group, resource *resourceTypes.NetworkResource) map[string]any {
return map[string]any{"name": group.Name, "id": group.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
type mockManager struct {
}
@@ -114,7 +109,7 @@ func (m *managerImpl) AddResourceToGroupInTransaction(ctx context.Context, trans
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, eventMetaResource(group, networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, group.EventMetaResource(networkResource))
}
return event, nil
@@ -138,7 +133,7 @@ func (m *managerImpl) RemoveResourceFromGroupInTransaction(ctx context.Context,
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, eventMetaResource(group, networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, group.EventMetaResource(networkResource))
}
return event, nil

View File

@@ -446,7 +446,7 @@ func (h *Handler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request) {
netMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, dns.CustomZone{}, nil, validPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil, account.GetActiveGroupUsers())
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, account.Peers, dnsDomain))
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, dnsDomain))
}
func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request) {
@@ -534,20 +534,15 @@ func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request)
util.WriteJSONObject(r.Context(), w, resp)
}
// toAccessiblePeers rehydrates the calculated map's component peers into the
// account's full peer objects, which carry the location/status/meta fields
// the API response needs.
func toAccessiblePeers(netMap *types.NetworkMap, accountPeers map[string]*nbpeer.Peer, dnsDomain string) []api.AccessiblePeer {
func toAccessiblePeers(netMap *types.NetworkMap, dnsDomain string) []api.AccessiblePeer {
accessiblePeers := make([]api.AccessiblePeer, 0, len(netMap.Peers)+len(netMap.OfflinePeers))
add := func(peers []*types.ComponentPeer) {
for _, p := range peers {
if peer := accountPeers[p.ID]; peer != nil {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(peer, dnsDomain))
}
}
for _, p := range netMap.Peers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
}
for _, p := range netMap.OfflinePeers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
}
add(netMap.Peers)
add(netMap.OfflinePeers)
return accessiblePeers
}

View File

@@ -683,81 +683,3 @@ func BackfillPublicIDs[T any](ctx context.Context, db *gorm.DB) error {
log.WithContext(ctx).Infof("Backfill of empty public_id in table %s completed", tableName)
return nil
}
// FoldCostAggregatesIntoBuckets migrates a per-request cost table from the old
// "stored aggregate" shape (cost_usd + cache_cost_usd columns) to the per-bucket
// breakdown, where the total and cache portion are derived on read instead.
//
// The fold preserves both aggregates exactly for historical rows: the cache
// total moves into cached_input_cost_usd and the remainder into
// input_cost_usd, so a row's derived total and cache cost still match what it
// reported before the upgrade. The finer split is genuinely unknown for those
// rows — the old schema never recorded a read/write or input/output division —
// so it is lumped rather than guessed; only rows written after the upgrade
// carry a true four-way split.
//
// Dropping the columns before folding would zero every historical row's cost,
// so the update runs first and the drop only happens once it succeeds. A table
// with no cost_usd column has already been migrated (or was created fresh) and
// is skipped.
func FoldCostAggregatesIntoBuckets[T any](ctx context.Context, db *gorm.DB) error {
var model T
if !db.Migrator().HasTable(&model) {
log.WithContext(ctx).Debugf("table for %T does not exist, no cost-bucket migration needed", model)
return nil
}
if !db.Migrator().HasColumn(&model, "cost_usd") {
log.WithContext(ctx).Debugf("table for %T has no cost_usd column, cost buckets already migrated", model)
return nil
}
stmt := &gorm.Statement{DB: db}
if err := stmt.Parse(&model); err != nil {
return fmt.Errorf("parse model schema: %w", err)
}
tableName := stmt.Schema.Table
// COALESCE guards rows whose new columns were added as NULL by an earlier
// AutoMigrate run that predates the NOT NULL default.
hasCacheColumn := db.Migrator().HasColumn(&model, "cache_cost_usd")
cacheExpr := "0"
if hasCacheColumn {
cacheExpr = "COALESCE(cache_cost_usd, 0)"
}
if err := db.Transaction(func(tx *gorm.DB) error {
// Only touch rows that carry a legacy total and no breakdown yet, so
// the migration is idempotent and never overwrites a true split.
update := fmt.Sprintf(`UPDATE %s
SET input_cost_usd = COALESCE(cost_usd, 0) - %s,
cached_input_cost_usd = %s,
cache_creation_cost_usd = 0,
output_cost_usd = 0
WHERE COALESCE(cost_usd, 0) <> 0
AND COALESCE(input_cost_usd, 0) = 0
AND COALESCE(cached_input_cost_usd, 0) = 0
AND COALESCE(cache_creation_cost_usd, 0) = 0
AND COALESCE(output_cost_usd, 0) = 0`, tableName, cacheExpr, cacheExpr)
res := tx.Exec(update)
if res.Error != nil {
return fmt.Errorf("fold legacy cost aggregates in %s: %w", tableName, res.Error)
}
log.WithContext(ctx).Infof("folded legacy cost aggregates into per-bucket columns for %d rows in table %s", res.RowsAffected, tableName)
if err := tx.Migrator().DropColumn(&model, "cost_usd"); err != nil {
return fmt.Errorf("drop cost_usd from %s: %w", tableName, err)
}
if hasCacheColumn {
if err := tx.Migrator().DropColumn(&model, "cache_cost_usd"); err != nil {
return fmt.Errorf("drop cache_cost_usd from %s: %w", tableName, err)
}
}
return nil
}); err != nil {
return err
}
log.WithContext(ctx).Infof("migration of stored cost aggregates to per-bucket columns in table %s completed", tableName)
return nil
}

View File

@@ -16,7 +16,6 @@ import (
"gorm.io/driver/sqlite"
"gorm.io/gorm"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/migration"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/testutil"
@@ -640,99 +639,3 @@ func TestCleanupOrphanedResources_SkipsWhenForeignKeyExists(t *testing.T) {
db.Model(&testChildWithFK{}).Count(&count)
assert.Equal(t, int64(2), count, "Both rows should survive — migration must skip when FK constraint exists")
}
// legacyCostRow is the pre-breakdown shape of the usage table: cost was stored
// as a total plus a cache portion, with no per-bucket columns. Used to build a
// realistic pre-upgrade table for the fold migration to run against.
type legacyCostRow struct {
ID string `gorm:"primaryKey"`
AccountID string
Model string
CostUSD float64
CacheCostUSD float64
}
func (legacyCostRow) TableName() string { return "agent_network_request_usage" }
// TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost covers the upgrade
// path: a table written under the old schema must come out with its per-row
// total and cache cost unchanged, because dropping cost_usd without folding it
// forward would silently zero every historical row's spend.
func TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
// setupDatabase hands back a process-shared database, so start from a clean
// table rather than inheriting rows from another test.
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&legacyCostRow{}), "legacy table must be created")
require.NoError(t, db.Create(&legacyCostRow{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6", CostUSD: 0.0123, CacheCostUSD: 0.0029,
}).Error)
require.NoError(t, db.Create(&legacyCostRow{
ID: "u2", AccountID: "acct-1", Model: "gpt-4o", CostUSD: 0.5, CacheCostUSD: 0,
}).Error)
// A zero-cost row (denied / unpriced request) must stay zero, not be touched.
require.NoError(t, db.Create(&legacyCostRow{ID: "u3", AccountID: "acct-1", Model: "gw/unpriced"}).Error)
// AutoMigrate adds the per-bucket columns alongside the legacy ones, exactly
// as a real upgrade does before the post-auto migrations run.
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}), "new columns must be added")
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db))
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cost_usd"),
"legacy cost_usd column must be dropped once folded")
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cache_cost_usd"),
"legacy cache_cost_usd column must be dropped once folded")
var rows []*agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.Order("id").Find(&rows).Error)
require.Len(t, rows, 3)
// u1: total and cache portion both preserved; the read/write and
// input/output splits are unknowable for a legacy row, so the cache total
// lands on cached_input and the remainder on input.
assert.InDelta(t, 0.0123, rows[0].TotalCostUSD(), 1e-9, "historical total must survive the fold")
assert.InDelta(t, 0.0029, rows[0].CacheCostUSD(), 1e-9, "historical cache cost must survive the fold")
assert.InDelta(t, 0.0094, rows[0].InputCostUSD, 1e-9, "non-cache remainder lands on input")
assert.InDelta(t, 0.0029, rows[0].CachedInputCostUSD, 1e-9, "legacy cache total lands on cached input")
assert.Zero(t, rows[0].CacheCreationCostUSD, "legacy rows carry no read/write split to recover")
assert.Zero(t, rows[0].OutputCostUSD, "legacy rows carry no input/output split to recover")
// u2: no cache spend — the whole total is the non-cache remainder.
assert.InDelta(t, 0.5, rows[1].TotalCostUSD(), 1e-9, "cache-free historical total must survive")
assert.Zero(t, rows[1].CacheCostUSD(), "a cache-free row must stay cache-free")
// u3: zero stays zero rather than being rewritten.
assert.Zero(t, rows[2].TotalCostUSD(), "an unpriced row must remain unpriced")
}
// TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated proves the migration is
// safe to re-run: with no legacy column present it is a no-op that leaves a
// true four-way split untouched.
func TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}))
// Timestamp must be set explicitly: a zero time.Time serialises as
// '0000-00-00 00:00:00', which MySQL rejects under strict mode.
require.NoError(t, db.Create(&agentNetworkTypes.AgentNetworkUsage{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6",
Timestamp: time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC),
InputCostUSD: 0.001, CachedInputCostUSD: 0.002, CacheCreationCostUSD: 0.003, OutputCostUSD: 0.004,
}).Error)
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db),
"running against an already-migrated table must be a no-op, not an error")
var row agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.First(&row, "id = ?", "u1").Error)
assert.InDelta(t, 0.001, row.InputCostUSD, 1e-9, "a true split must not be rewritten")
assert.InDelta(t, 0.002, row.CachedInputCostUSD, 1e-9)
assert.InDelta(t, 0.003, row.CacheCreationCostUSD, 1e-9)
assert.InDelta(t, 0.004, row.OutputCostUSD, 1e-9)
assert.InDelta(t, 0.01, row.TotalCostUSD(), 1e-9, "derived total sums the four buckets")
}

View File

@@ -14,7 +14,6 @@ import (
nbDomain "github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkResourceType string
@@ -65,27 +64,6 @@ func NewNetworkResource(accountID, networkID, name, description, address string,
}, nil
}
// ToComponent converts the resource to its self-contained components
// representation. Returns nil for a nil resource.
func (n *NetworkResource) ToComponent() *sharedTypes.ComponentResource {
if n == nil {
return nil
}
return &sharedTypes.ComponentResource{
ID: n.ID,
PublicID: n.PublicID,
NetworkID: n.NetworkID,
AccountID: n.AccountID,
Name: n.Name,
Description: n.Description,
Type: sharedTypes.ComponentResourceType(n.Type),
Address: n.Address,
Domain: n.Domain,
Prefix: n.Prefix,
Enabled: n.Enabled,
}
}
func (n *NetworkResource) ToAPIResponse(groups []api.GroupMinimum) *api.NetworkResource {
addr := n.Prefix.String()
if n.Type == Domain {

View File

@@ -7,7 +7,6 @@ import (
"github.com/netbirdio/netbird/management/server/networks/types"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkRouter struct {
@@ -22,36 +21,6 @@ type NetworkRouter struct {
Enabled bool
}
// ToComponent converts the router to its self-contained components
// representation. Returns nil for a nil router.
func (n *NetworkRouter) ToComponent() *sharedTypes.ComponentRouter {
if n == nil {
return nil
}
return &sharedTypes.ComponentRouter{
NetworkID: n.NetworkID,
PublicID: n.PublicID,
Peer: n.Peer,
PeerGroups: n.PeerGroups,
Masquerade: n.Masquerade,
Metric: n.Metric,
Enabled: n.Enabled,
}
}
// ToComponentMap converts a peer-keyed router map to its components
// representation.
func ToComponentMap(routers map[string]*NetworkRouter) map[string]*sharedTypes.ComponentRouter {
if routers == nil {
return nil
}
out := make(map[string]*sharedTypes.ComponentRouter, len(routers))
for id, r := range routers {
out[id] = r.ToComponent()
}
return out
}
func NewNetworkRouter(accountID string, networkID string, peer string, peerGroups []string, masquerade bool, metric int, enabled bool) (*NetworkRouter, error) {
r := &NetworkRouter{
ID: xid.New().String(),

View File

@@ -405,7 +405,7 @@ func (am *DefaultAccountManager) CreatePeerJob(ctx context.Context, accountID, p
return status.NewPeerNotPartOfAccountError()
}
meetMinVer, err := version.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
meetMinVer, err := posture.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
if !version.IsDevelopmentVersion(p.Meta.WtVersion) && (!meetMinVer || err != nil) {
return status.Errorf(status.PreconditionFailed, "peer version %s does not meet the minimum required version %s for remote jobs", p.Meta.WtVersion, remoteJobsMinVer)
}
@@ -1588,7 +1588,7 @@ func affectedPeerIDsFromNetworkMap(nmap *types.NetworkMap, selfPeerID string) []
}
seen := make(map[string]struct{}, len(nmap.Peers)+len(nmap.OfflinePeers))
ids := make([]string, 0, len(nmap.Peers)+len(nmap.OfflinePeers))
add := func(peers []*types.ComponentPeer) {
add := func(peers []*nbpeer.Peer) {
for _, p := range peers {
if p == nil || p.ID == "" || p.ID == selfPeerID {
continue

View File

@@ -13,7 +13,6 @@ import (
"github.com/netbirdio/netbird/management/server/util"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
// Peer capability constants mirror the proto enum values.
@@ -206,35 +205,6 @@ func (p *Peer) AddedWithSSOLogin() bool {
return p.UserID != ""
}
// ToComponent converts the peer to its self-contained components
// representation, carrying exactly the subset of peer data that crosses the
// components wire format. Returns nil for a nil peer so callers can convert
// possibly-missing peers without guarding.
func (p *Peer) ToComponent() *sharedTypes.ComponentPeer {
if p == nil {
return nil
}
cp := &sharedTypes.ComponentPeer{
ID: p.ID,
Key: p.Key,
IP: p.IP,
IPv6: p.IPv6,
DNSLabel: p.DNSLabel,
SSHKey: p.SSHKey,
SSHEnabled: p.SSHEnabled,
ServerSSHAllowed: p.Meta.Flags.ServerSSHAllowed,
AgentVersion: p.Meta.WtVersion,
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
SupportsIPv6: p.SupportsIPv6(),
LoginExpirationEnabled: p.LoginExpirationEnabled,
AddedWithSSOLogin: p.AddedWithSSOLogin(),
}
if p.LastLogin != nil {
cp.LastLogin = *p.LastLogin
}
return cp
}
// HasCapability reports whether the peer has the given capability.
func (p *Peer) HasCapability(capability int32) bool {
return slices.Contains(p.Meta.Capabilities, capability)

View File

@@ -1092,14 +1092,14 @@ func TestToSyncResponse(t *testing.T) {
}
networkMap := &types.NetworkMap{
Network: &types.Network{Net: *ipnet, Serial: 1000},
Peers: []*types.ComponentPeer{{
Peers: []*nbpeer.Peer{{
IP: netip.MustParseAddr("192.168.1.2"),
IPv6: netip.MustParseAddr("fd00::2"),
Key: "peer2-key",
DNSLabel: "peer2",
SSHEnabled: true,
SSHKey: "peer2-ssh-key"}},
OfflinePeers: []*types.ComponentPeer{{
OfflinePeers: []*nbpeer.Peer{{
IP: netip.MustParseAddr("192.168.1.3"),
IPv6: netip.MustParseAddr("fd00::3"),
Key: "peer3-key",

View File

@@ -3,9 +3,11 @@ package posture
import (
"context"
"fmt"
"strings"
"github.com/hashicorp/go-version"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbversion "github.com/netbirdio/netbird/version"
)
type NBVersionCheck struct {
@@ -14,8 +16,14 @@ type NBVersionCheck struct {
var _ Check = (*NBVersionCheck)(nil)
// sanitizeVersion removes anything after the pre-release tag (e.g., "-dev", "-alpha", etc.)
func sanitizeVersion(version string) string {
parts := strings.Split(version, "-")
return parts[0]
}
func (n *NBVersionCheck) Check(ctx context.Context, peer nbpeer.Peer) (bool, error) {
meetsMin, err := nbversion.MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
meetsMin, err := MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
if err != nil {
return false, err
}
@@ -40,3 +48,21 @@ func (n *NBVersionCheck) Validate() error {
}
return nil
}
// MeetsMinVersion checks if the peer's version meets or exceeds the minimum required version
func MeetsMinVersion(minVer, peerVer string) (bool, error) {
peerVer = sanitizeVersion(peerVer)
minVer = sanitizeVersion(minVer)
peerNBVer, err := version.NewVersion(peerVer)
if err != nil {
return false, err
}
constraints, err := version.NewConstraint(">= " + minVer)
if err != nil {
return false, err
}
return constraints.Check(peerNBVer), nil
}

View File

@@ -139,3 +139,68 @@ func TestNBVersionCheck_Validate(t *testing.T) {
})
}
}
func TestMeetsMinVersion(t *testing.T) {
tests := []struct {
name string
minVer string
peerVer string
want bool
wantErr bool
}{
{
name: "Peer version greater than min version",
minVer: "0.26.0",
peerVer: "0.60.1",
want: true,
wantErr: false,
},
{
name: "Peer version equals min version",
minVer: "1.0.0",
peerVer: "1.0.0",
want: true,
wantErr: false,
},
{
name: "Peer version less than min version",
minVer: "1.0.0",
peerVer: "0.9.9",
want: false,
wantErr: false,
},
{
name: "Peer version with pre-release tag greater than min version",
minVer: "1.0.0",
peerVer: "1.0.1-alpha",
want: true,
wantErr: false,
},
{
name: "Invalid peer version format",
minVer: "1.0.0",
peerVer: "dev",
want: false,
wantErr: true,
},
{
name: "Invalid min version format",
minVer: "invalid.version",
peerVer: "1.0.0",
want: false,
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := MeetsMinVersion(tt.minVer, tt.peerVer)
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
assert.Equal(t, tt.want, got)
})
}
}

View File

@@ -71,7 +71,7 @@ func (s *SqlStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetr
usageRow := db.Model(&agentNetworkTypes.AgentNetworkUsage{}).
Select("COALESCE(SUM(input_tokens), 0) AS input_tokens, " +
"COALESCE(SUM(output_tokens), 0) AS output_tokens, " +
"COALESCE(SUM" + agentNetworkTypes.CostUSDSQLExpr + ", 0) AS cost_usd").Row()
"COALESCE(SUM(cost_usd), 0) AS cost_usd").Row()
if err := usageRow.Scan(&m.InputTokens, &m.OutputTokens, &m.CostUSD); err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("scan agent network usage metrics: %w", err)
}

View File

@@ -37,7 +37,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
InputCostUSD: 0.0231,
CostUSD: 0.0231,
}
usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -71,7 +71,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
InputCostUSD: 0.0231,
CostUSD: 0.0231,
}
entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -127,7 +127,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
mk := func(id string, ts time.Time, model string, in, out int64, cost float64) *agentNetworkTypes.AgentNetworkUsage {
return &agentNetworkTypes.AgentNetworkUsage{
ID: id, AccountID: accountID, Timestamp: ts, Model: model,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, InputCostUSD: cost,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, CostUSD: cost,
}
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
@@ -143,7 +143,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
assert.Equal(t, "2026-05-05", buckets[0].PeriodStart, "oldest-first ordering")
assert.Equal(t, int64(300), buckets[0].InputTokens, "same-day input tokens summed")
assert.Equal(t, int64(130), buckets[0].OutputTokens)
assert.InDelta(t, 0.30, buckets[0].TotalCostUSD(), 1e-9, "same-day cost summed")
assert.InDelta(t, 0.30, buckets[0].CostUSD, 1e-9, "same-day cost summed")
assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
assert.Equal(t, int64(15), buckets[1].TotalTokens)
@@ -174,7 +174,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts,
UserID: user, StatusCode: 200, Provider: provider, Model: model,
SessionID: session, Decision: decision,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, InputCostUSD: cost,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, CostUSD: cost,
}
}
@@ -207,7 +207,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
s1 := sessions[2]
assert.Equal(t, 2, s1.RequestCount, "s1 has two requests")
assert.Equal(t, int64(300), s1.TotalTokens, "tokens summed across the session")
assert.InDelta(t, 0.30, s1.TotalCostUSD(), 1e-9, "cost summed across the session")
assert.InDelta(t, 0.30, s1.CostUSD, 1e-9, "cost summed across the session")
assert.Equal(t, "alice", s1.UserID)
assert.Equal(t, "allow", s1.Decision)
// SQLite hands times back in time.Local; normalise to UTC so the instant is

View File

@@ -650,14 +650,6 @@ func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
func(db *gorm.DB) error {
return migration.DropIndex[proxy.Proxy](ctx, db, "idx_proxy_account_id_unique")
},
// Post-auto so the per-bucket cost columns already exist when the legacy
// aggregates are folded into them and dropped.
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkAccessLog](ctx, db)
},
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db)
},
}
}

View File

@@ -283,8 +283,8 @@ func (a *Account) SynthesizePrivateServiceZones(peerID string) []nbdns.CustomZon
// it, adding a single private service would black-hole every
// other name under the zone apex.
zone = &nbdns.CustomZone{
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
NonAuthoritative: true,
SearchDomainDisabled: true,
}
@@ -1082,7 +1082,6 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
peersExists := make(map[string]struct{})
rules := make([]*FirewallRule, 0)
peers := make([]*nbpeer.Peer, 0)
targetComponent := targetPeer.ToComponent()
return func(rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int) {
for _, peer := range groupPeers {
@@ -1118,10 +1117,10 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
if len(rule.Ports) == 0 && len(rule.PortRanges) == 0 {
rules = append(rules, &fr)
} else {
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetComponent)...)
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetPeer)...)
}
rules = AppendIPv6FirewallRule(rules, rulesExists, peer.ToComponent(), targetComponent, rule, FirewallRuleContext{
rules = AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, FirewallRuleContext{
Direction: direction,
DirStr: strconv.Itoa(direction),
ProtocolStr: string(protocol),
@@ -1281,7 +1280,7 @@ func (a *Account) getRouteFirewallRules(ctx context.Context, peerID string, poli
return fwRules
}
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*ComponentPeer {
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*nbpeer.Peer {
distPeersWithPolicy := make(map[string]struct{})
for _, id := range rule.Sources {
group := a.Groups[id]
@@ -1308,13 +1307,13 @@ func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID
}
}
distributionGroupPeers := make([]*ComponentPeer, 0, len(distPeersWithPolicy))
distributionGroupPeers := make([]*nbpeer.Peer, 0, len(distPeersWithPolicy))
for pID := range distPeersWithPolicy {
peer := a.Peers[pID]
if peer == nil {
continue
}
distributionGroupPeers = append(distributionGroupPeers, peer.ToComponent())
distributionGroupPeers = append(distributionGroupPeers, peer)
}
return distributionGroupPeers
}

View File

@@ -9,7 +9,9 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/modules/zones"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/route"
)
@@ -111,7 +113,7 @@ func (a *Account) GetPeerNetworkMapComponents(
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
Peers: map[string]*nbpeer.Peer{peerID: peer},
})
}
@@ -124,7 +126,7 @@ func (a *Account) GetPeerNetworkMapComponents(
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
Peers: map[string]*nbpeer.Peer{peerID: peer},
})
}
@@ -134,10 +136,10 @@ func (a *Account) GetPeerNetworkMapComponents(
NameServerGroups: make([]*nbdns.NameServerGroup, 0),
CustomZoneDomain: peersCustomZone.Domain,
ResourcePoliciesMap: make(map[string][]*Policy),
RoutersMap: make(map[string]map[string]*ComponentRouter),
NetworkResources: make([]*ComponentResource, 0),
RoutersMap: make(map[string]map[string]*routerTypes.NetworkRouter),
NetworkResources: make([]*resourceTypes.NetworkResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(a.PostureChecks)),
RouterPeers: make(map[string]*ComponentPeer),
RouterPeers: make(map[string]*nbpeer.Peer),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
}
@@ -172,7 +174,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
components.Peers = relevantPeers
components.Groups = GroupsToComponent(relevantGroups)
components.Groups = relevantGroups
components.Policies = relevantPolicies
components.Routes = relevantRoutes
components.AllDNSRecords = filterDNSRecordsByPeers(peersCustomZone.Records, relevantPeers, peer.SupportsIPv6() && peer.IPv6.IsValid())
@@ -221,7 +223,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
for _, pID := range a.getPostureValidPeersSaveFailed(peers, policy.SourcePostureChecks, validatedPeersMap, &components.PostureFailedPeers) {
if _, exists := components.Peers[pID]; !exists {
components.Peers[pID] = a.GetPeer(pID).ToComponent()
components.Peers[pID] = a.GetPeer(pID)
}
}
} else {
@@ -254,14 +256,14 @@ func (a *Account) GetPeerNetworkMapComponents(
for _, srcGroupID := range rule.Sources {
if g := a.Groups[srcGroupID]; g != nil {
if _, exists := components.Groups[srcGroupID]; !exists {
components.Groups[srcGroupID] = g.ToComponent()
components.Groups[srcGroupID] = g
}
}
}
for _, dstGroupID := range rule.Destinations {
if g := a.Groups[dstGroupID]; g != nil {
if _, exists := components.Groups[dstGroupID]; !exists {
components.Groups[dstGroupID] = g.ToComponent()
components.Groups[dstGroupID] = g
}
}
}
@@ -276,22 +278,20 @@ func (a *Account) GetPeerNetworkMapComponents(
// network in the account — accounts with many tenants/networks
// shipped tens of unrelated peers in `peers[]` and `routers_map`.
if addSourcePeers {
components.RoutersMap[resource.NetworkID] = routerTypes.ToComponentMap(networkRoutingPeers)
components.RoutersMap[resource.NetworkID] = networkRoutingPeers
for peerIDKey := range networkRoutingPeers {
if p := a.Peers[peerIDKey]; p != nil {
cp := components.RouterPeers[peerIDKey]
if cp == nil {
cp = p.ToComponent()
components.RouterPeers[peerIDKey] = cp
if _, exists := components.RouterPeers[peerIDKey]; !exists {
components.RouterPeers[peerIDKey] = p
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := validatedPeersMap[peerIDKey]; validated {
components.Peers[peerIDKey] = cp
components.Peers[peerIDKey] = p
}
}
}
}
components.NetworkResources = append(components.NetworkResources, resource.ToComponent())
components.NetworkResources = append(components.NetworkResources, resource)
}
}
@@ -312,14 +312,14 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
peerSSHEnabled bool,
validatedPeersMap map[string]struct{},
postureFailedPeers *map[string]map[string]struct{},
) (map[string]*ComponentPeer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*ComponentPeer, len(a.Peers)/4)
) (map[string]*nbpeer.Peer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*nbpeer.Peer, len(a.Peers)/4)
relevantGroupIDs := make(map[string]*Group, len(a.Groups)/4)
relevantPolicies := make([]*Policy, 0, len(a.Policies))
relevantRoutes := make([]*route.Route, 0, len(a.Routes))
sshReqs := sshRequirements{neededGroupIDs: make(map[string]struct{})}
relevantPeerIDs[peerID] = a.GetPeer(peerID).ToComponent()
relevantPeerIDs[peerID] = a.GetPeer(peerID)
peerGroupSet := make(map[string]struct{}, 8)
for groupID, group := range a.Groups {
@@ -384,7 +384,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if r.Peer != "" {
if _, ok := validatedPeersMap[r.Peer]; ok {
if p := a.GetPeer(r.Peer); p != nil {
relevantPeerIDs[r.Peer] = p.ToComponent()
relevantPeerIDs[r.Peer] = p
}
}
}
@@ -401,7 +401,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
continue
}
if p := a.GetPeer(pid); p != nil {
relevantPeerIDs[pid] = p.ToComponent()
relevantPeerIDs[pid] = p
}
}
}
@@ -458,9 +458,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInSources {
policyRelevant = true
for _, pid := range destinationPeers {
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
relevantPeerIDs[pid] = a.GetPeer(pid)
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = a.GetGroup(dstGroupID)
@@ -470,9 +468,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInDestinations {
policyRelevant = true
for _, pid := range sourcePeers {
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
relevantPeerIDs[pid] = a.GetPeer(pid)
}
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = a.GetGroup(srcGroupID)
@@ -628,7 +624,7 @@ func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChe
// that name them. Calculate() tolerates groups with empty Peers (the inner
// loops simply iterate zero times), so retaining them is behaviourally a
// no-op for the legacy path that consumes the same NetworkMapComponents.
func filterGroupPeers(groups *map[string]*ComponentGroup, peers map[string]*ComponentPeer) {
func filterGroupPeers(groups *map[string]*Group, peers map[string]*nbpeer.Peer) {
for groupID, groupInfo := range *groups {
filteredPeers := make([]string, 0, len(groupInfo.Peers))
for _, pid := range groupInfo.Peers {
@@ -638,14 +634,14 @@ func filterGroupPeers(groups *map[string]*ComponentGroup, peers map[string]*Comp
}
if len(filteredPeers) != len(groupInfo.Peers) {
ng := *groupInfo
ng := groupInfo.Copy()
ng.Peers = filteredPeers
(*groups)[groupID] = &ng
(*groups)[groupID] = ng
}
}
}
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*ComponentPeer) {
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*nbpeer.Peer) {
if len(*postureFailedPeers) == 0 {
return
}
@@ -680,7 +676,7 @@ func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}
}
}
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*ComponentPeer, includeIPv6 bool) []nbdns.SimpleRecord {
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*nbpeer.Peer, includeIPv6 bool) []nbdns.SimpleRecord {
if len(records) == 0 || len(peers) == 0 {
return nil
}

View File

@@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/require"
nbdns "github.com/netbirdio/netbird/dns"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
)
func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
@@ -48,7 +49,7 @@ func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
})
}
func netmapPeerIDs(peers []*ComponentPeer) []string {
func netmapPeerIDs(peers []*nbpeer.Peer) []string {
ids := make([]string, 0, len(peers))
for _, p := range peers {
if p == nil {

View File

@@ -666,7 +666,7 @@ func Test_ExpandPortsAndRanges_SSHRuleExpansion(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer.ToComponent())
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer)
var ports []string
for _, fr := range result {

View File

@@ -6,6 +6,7 @@ import (
"net"
"net/netip"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
sharedtypes "github.com/netbirdio/netbird/shared/management/types"
)
@@ -17,6 +18,9 @@ type DNSSettings = sharedtypes.DNSSettings
type FirewallRule = sharedtypes.FirewallRule
type Group = sharedtypes.Group
type GroupPeer = sharedtypes.GroupPeer
type Network = sharedtypes.Network
type NetworkMap = sharedtypes.NetworkMap
type ForwardingRule = sharedtypes.ForwardingRule
@@ -38,18 +42,6 @@ type RouteFirewallRule = sharedtypes.RouteFirewallRule
type NetworkMapComponents = sharedtypes.NetworkMapComponents
type ComponentPeer = sharedtypes.ComponentPeer
type ComponentGroup = sharedtypes.ComponentGroup
type ComponentRouter = sharedtypes.ComponentRouter
type ComponentResource = sharedtypes.ComponentResource
type ComponentResourceType = sharedtypes.ComponentResourceType
const (
ComponentResourceHost = sharedtypes.ComponentResourceHost
ComponentResourceSubnet = sharedtypes.ComponentResourceSubnet
ComponentResourceDomain = sharedtypes.ComponentResourceDomain
)
var EmptyNetworkMapComponents = sharedtypes.EmptyNetworkMapComponents
type AccountSettingsInfo = sharedtypes.AccountSettingsInfo
@@ -60,7 +52,12 @@ type NetworkMapComponentsCompact = sharedtypes.NetworkMapComponentsCompact
type LookupMap = sharedtypes.LookupMap
type FirewallRuleContext = sharedtypes.FirewallRuleContext
const GroupAllName = sharedtypes.GroupAllName
const (
GroupIssuedAPI = sharedtypes.GroupIssuedAPI
GroupIssuedJWT = sharedtypes.GroupIssuedJWT
GroupIssuedIntegration = sharedtypes.GroupIssuedIntegration
GroupAllName = sharedtypes.GroupAllName
)
// Function forwarders preserve types.X(...) call sites that previously
// resolved to package-local funcs. Plain forwarders (not var aliases) keep
@@ -70,11 +67,11 @@ func PolicyRuleImpliesLegacySSH(rule *PolicyRule) bool {
return sharedtypes.PolicyRuleImpliesLegacySSH(rule)
}
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *ComponentPeer) []*FirewallRule {
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *nbpeer.Peer) []*FirewallRule {
return sharedtypes.ExpandPortsAndRanges(base, rule, peer)
}
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *ComponentPeer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *nbpeer.Peer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
return sharedtypes.AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, rc)
}
@@ -82,7 +79,7 @@ func CalculateNetworkMapFromComponents(ctx context.Context, components *NetworkM
return sharedtypes.CalculateNetworkMapFromComponents(ctx, components)
}
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*ComponentPeer, direction int, includeIPv6 bool) []*RouteFirewallRule {
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int, includeIPv6 bool) []*RouteFirewallRule {
return sharedtypes.GenerateRouteFirewallRules(ctx, route, rule, groupPeers, direction, includeIPv6)
}

View File

@@ -9,7 +9,6 @@ import (
"github.com/stretchr/testify/require"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
func TestNetworkMapComponents_IPv6EndToEnd(t *testing.T) {
@@ -105,7 +104,7 @@ func TestNetworkMapComponents_RemotePeerWithoutCapability(t *testing.T) {
require.NotNil(t, nm)
t.Run("AllowedIPs include remote v6", func(t *testing.T) {
var dst *types.ComponentPeer
var dst *nbpeer.Peer
for _, p := range nm.Peers {
if p.ID == "peer-dst-1" {
dst = p

View File

@@ -49,7 +49,7 @@ func allPeersValidated(account *types.Account, excludePeerIDs ...string) map[str
return validated
}
func peerIDs(peers []*types.ComponentPeer) []string {
func peerIDs(peers []*nbpeer.Peer) []string {
ids := make([]string, len(peers))
for i, p := range peers {
ids[i] = p.ID

View File

@@ -19,3 +19,34 @@ func Difference(a, b []string) []string {
func ToPtr[T any](value T) *T {
return &value
}
type comparableObject[T any] interface {
Equal(other T) bool
}
func MergeUnique[T comparableObject[T]](arr1, arr2 []T) []T {
var result []T
for _, item := range arr1 {
if !contains(result, item) {
result = append(result, item)
}
}
for _, item := range arr2 {
if !contains(result, item) {
result = append(result, item)
}
}
return result
}
func contains[T comparableObject[T]](slice []T, element T) bool {
for _, item := range slice {
if item.Equal(element) {
return true
}
}
return false
}

View File

@@ -1,4 +1,4 @@
package types
package util
import (
"testing"
@@ -17,7 +17,7 @@ func (t testObject) Equal(other testObject) bool {
func Test_MergeUniqueArraysWithoutDuplicates(t *testing.T) {
arr1 := []testObject{{value: 1}, {value: 2}}
arr2 := []testObject{{value: 2}, {value: 3}}
result := mergeUnique(arr1, arr2)
result := MergeUnique(arr1, arr2)
assert.Len(t, result, 3)
assert.Contains(t, result, testObject{value: 1})
assert.Contains(t, result, testObject{value: 2})
@@ -27,14 +27,14 @@ func Test_MergeUniqueArraysWithoutDuplicates(t *testing.T) {
func Test_MergeUniqueHandlesEmptyArrays(t *testing.T) {
arr1 := []testObject{}
arr2 := []testObject{}
result := mergeUnique(arr1, arr2)
result := MergeUnique(arr1, arr2)
assert.Empty(t, result)
}
func Test_MergeUniqueHandlesOneEmptyArray(t *testing.T) {
arr1 := []testObject{{value: 1}, {value: 2}}
arr2 := []testObject{}
result := mergeUnique(arr1, arr2)
result := MergeUnique(arr1, arr2)
assert.Len(t, result, 2)
assert.Contains(t, result, testObject{value: 1})
assert.Contains(t, result, testObject{value: 2})

View File

@@ -10,7 +10,6 @@ COPY encryption ./encryption
COPY flow ./flow
COPY formatter ./formatter
COPY monotime ./monotime
COPY management ./management
COPY proxy ./proxy
COPY route ./route
COPY shared ./shared

View File

@@ -221,21 +221,14 @@ func (l *Logger) allowDenyLog(serviceID types.ServiceID, reason string) bool {
// proxy/internal/middleware/keys.go — only the dimensions management needs to
// record a usage row (provider / model / tokens / cost / groups).
var usageMetadataKeys = map[string]struct{}{
"llm.provider": {},
"llm.model": {},
"llm.resolved_provider_id": {},
"llm.input_tokens": {},
"llm.output_tokens": {},
"llm.total_tokens": {},
"llm.cached_input_tokens": {},
"llm.cache_creation_tokens": {},
"cost.usd_input": {},
"cost.usd_cached_input": {},
"cost.usd_cache_creation": {},
"cost.usd_output": {},
"cost.usd_total": {},
"cost.usd_cache": {},
"llm.authorising_groups": {},
"llm.provider": {},
"llm.model": {},
"llm.resolved_provider_id": {},
"llm.input_tokens": {},
"llm.output_tokens": {},
"llm.total_tokens": {},
"cost.usd_total": {},
"llm.authorising_groups": {},
}
// stripAgentNetworkEntryForUsage returns the entry reduced to what's needed to

View File

@@ -56,12 +56,10 @@ type bedrockResponse struct {
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
// Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens).
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
CacheReadTokensCamel int64 `json:"cacheReadInputTokens"`
CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"`
// Converse — camelCase.
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
} `json:"usage"`
}
@@ -85,18 +83,16 @@ func (BedrockParser) ParseResponse(status int, contentType string, body []byte)
}
inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel)
cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel)
total := resp.Usage.TotalTokensCamel
if total == 0 {
total = inTok + outTok + cacheRead + cacheWrite
total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
}
return Usage{
InputTokens: inTok,
OutputTokens: outTok,
TotalTokens: total,
CachedInputTokens: cacheRead,
CacheCreationTokens: cacheWrite,
CachedInputTokens: resp.Usage.CacheReadInputTokens,
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
}, nil
}

View File

@@ -26,18 +26,6 @@ func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
}
// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields.
func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) {
body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(11), u.InputTokens, "converse input tokens")
require.Equal(t, int64(3), u.OutputTokens, "converse output tokens")
require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens")
require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens")
require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens")
}
func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")

View File

@@ -161,16 +161,6 @@ openai:
input_per_1k: 0.0001
output_per_1k: 0
# Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
# reports cache hits OpenAI-style when present; cached input is 10% of
# input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
# serves newer accounts (K2-era ids and kimi-latest 404), matching the
# management catalog.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cached_input_per_1k: 0.0003
anthropic:
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
# Pricing source: Anthropic's current published rates per million tokens,
@@ -216,20 +206,6 @@ anthropic:
cache_read_per_1k: 0.0001
cache_creation_per_1k: 0.00125
# Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
# (/anthropic/v1/messages — the official Claude Code setup). Same rates
# as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
# Claude Code guides set for the 1M-context alias; priced identically so
# cost metering doesn't silently skip those requests.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
"kimi-k3[1m]":
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
bedrock:
# AWS Bedrock model ids, normalised by the request parser (cross-region
# inference-profile prefix + version/throughput suffix stripped), e.g.

View File

@@ -128,46 +128,6 @@ type Table struct {
// - Other providers: cached and cacheCreation are ignored; cost is
// inTokens*InputPer1K + outTokens*OutputPer1K.
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
return c.TotalUSD, ok
}
// Costs is a per-request cost split. The four per-bucket fields are the base
// of the breakdown — one per token bucket the provider bills separately — and
// the two aggregates are derived from them:
//
// TotalUSD = InputUSD + CachedInputUSD + CacheCreationUSD + OutputUSD
// CacheUSD = CachedInputUSD + CacheCreationUSD
//
// InputUSD is always the cost of the *non-cached* input bucket, for both
// provider shapes: on OpenAI the cached subset is carved out of inTokens and
// billed as CachedInputUSD, so the two never double-count. Buckets a provider
// doesn't bill are zero, which keeps the identities above true everywhere.
type Costs struct {
InputUSD float64
CachedInputUSD float64
CacheCreationUSD float64
OutputUSD float64
TotalUSD float64
CacheUSD float64
}
// newCosts assembles a split from its per-bucket parts, deriving the two
// aggregates so TotalUSD and CacheUSD can never drift from the breakdown.
func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
return Costs{
InputUSD: input,
CachedInputUSD: cachedInput,
CacheCreationUSD: cacheCreation,
OutputUSD: output,
TotalUSD: input + cachedInput + cacheCreation + output,
CacheUSD: cachedInput + cacheCreation,
}
}
// Costs returns the estimated USD cost split for the given token counts, with
// the same semantics as Cost.
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
// Clamp negatives to zero before any pricing math so a malformed
// upstream count can never produce a negative cost.
if inTokens < 0 {
@@ -183,15 +143,15 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
cacheCreation = 0
}
if t == nil {
return Costs{}, false
return 0, false
}
byModel, ok := t.entries[provider]
if !ok {
return Costs{}, false
return 0, false
}
entry, ok := byModel[model]
if !ok {
return Costs{}, false
return 0, false
}
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
switch provider {
@@ -208,7 +168,7 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
}
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
cached := float64(clamped) / 1000.0 * cachedRate
return newCosts(nonCached, cached, 0, output), true
return nonCached + cached + output, true
case "anthropic", "bedrock":
// Bedrock-Anthropic returns the same additive cache buckets as
// first-party Anthropic; non-Anthropic Bedrock models simply report
@@ -224,10 +184,10 @@ func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput,
input := float64(inTokens) / 1000.0 * entry.InputPer1K
read := float64(cachedInput) / 1000.0 * readRate
create := float64(cacheCreation) / 1000.0 * createRate
return newCosts(input, read, create, output), true
return input + read + create + output, true
default:
input := float64(inTokens) / 1000.0 * entry.InputPer1K
return newCosts(input, 0, 0, output), true
return input + output, true
}
}

View File

@@ -1,329 +0,0 @@
package builtin_test
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"strconv"
"testing"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
)
// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the embedded default pricing
// table and asserts exact USD amounts hardcoded from the vendors' published prices, including the cache split.
func TestCostCalculation_ProviderMatrix(t *testing.T) {
// Empty data dir → embedded defaults, like a proxy with no pricing override.
builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil)
reqMW, err := llm_request_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_request_parser")
respMW, err := llm_response_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_response_parser")
costMW, err := cost_meter.Factory{}.New(nil)
require.NoError(t, err, "build cost_meter")
t.Cleanup(func() { _ = costMW.Close() })
const jsonCT = "application/json"
const sseCT = "text/event-stream"
const awsCT = "application/vnd.amazon.eventstream"
cases := []struct {
name string
url string
reqBody []byte
respCT string
respBody []byte
wantProvider string
wantModel string
wantCost float64 // exact expected USD; ignored when wantSkip is set
wantCacheCost float64 // expected cost.usd_cache portion of wantCost
wantSkip string // expected cost.skipped reason, "" when priced
}{
{
// gpt-4o-mini $0.15/$0.60 per MTok: 1000×0.15/1M + 500×0.60/1M.
name: "openai chat completions",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"choices":[{"message":{"content":"pong"}}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// OpenAI cached tokens are a SUBSET of prompt_tokens at a discount; gpt-4o $2.50/$10 per MTok, cached $1.25/M:
// 250×2.5/1M + 750×1.25/1M + 500×10/1M.
name: "openai cached subset discount",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500,"prompt_tokens_details":{"cached_tokens":750}}}`),
wantProvider: "openai",
wantModel: "gpt-4o",
wantCost: 0.0065625,
wantCacheCost: 0.0009375,
},
{
// OpenAI streaming: usage rides the final SSE frame.
name: "openai chat SSE stream",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"choices":[{"delta":{"content":"po"}}]}`, `{"choices":[{"delta":{"content":"ng"}}]}`, `{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500}}`, "[DONE]"),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// Mistral speaks the OpenAI shape: mistral-large-latest $0.50/$1.50 per MTok.
name: "mistral via openai shape",
url: "https://api.mistral.ai/v1/chat/completions",
reqBody: []byte(`{"model":"mistral-large-latest","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":1000}}`),
wantProvider: "openai",
wantModel: "mistral-large-latest",
wantCost: 0.002,
},
{
// The field report, minus caching: Bedrock Sonnet 4.6 $3/$15 per MTok, 3×3/1M + 1514×15/1M = $0.022719.
// Also covers inference-profile normalization of the region-prefixed versioned id in the URL.
name: "bedrock invoke — reported scenario, no cache",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":3,"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.022719,
},
{
// The field report as observed: the FIRST call of a session also wrote a 30,528-token prompt cache at
// 1.25× input ($3.75/M): 0.022719 + 30528×3.75/1M = $0.137199 — the reported $0.1372.
name: "bedrock invoke — reported scenario with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":3,"output_tokens":1514,"cache_creation_input_tokens":30528,"cache_read_input_tokens":0}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Same numbers over the InvokeModel event-stream: message_start carries input + cache, message_delta the output.
name: "bedrock invoke stream with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke-with-response-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: awsCT,
respBody: bedrockInvokeStream(t, `{"type":"message_start","message":{"usage":{"input_tokens":3,"output_tokens":1,"cache_creation_input_tokens":30528}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Converse camelCase usage incl. cache buckets. Haiku 4.5 $1/$5 per MTok, read $0.10/M, write $1.25/M:
// 50×1/1M + 100×5/1M + 2000×0.1/1M + 1000×1.25/1M = $0.002.
name: "bedrock converse with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: jsonCT,
respBody: []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// Same numbers over converse-stream: usage rides the trailing metadata frame.
name: "bedrock converse stream with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: awsCT,
respBody: bedrockConverseStream(t,
`{"delta":{"text":"pong"}}`,
`{"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`,
),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// First-party Anthropic, additive cache buckets. Sonnet 4.6:
// 256×3/1M + 200×15/1M + 768×0.3/1M + 512×3.75/1M.
name: "anthropic messages with cache buckets",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0059184,
wantCacheCost: 0.0021504,
},
{
// Anthropic SSE: input from message_start, output from message_delta. Haiku 4.5: 1000×1/1M + 2000×5/1M.
name: "anthropic SSE stream",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-haiku-4-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"type":"message_start","message":{"usage":{"input_tokens":1000,"output_tokens":2}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":2000}}`, `{"type":"message_stop"}`),
wantProvider: "anthropic",
wantModel: "claude-haiku-4-5",
wantCost: 0.011,
},
{
// Kimi's Anthropic-compatible endpoint: kimi-k3 $3/$15 per MTok under the anthropic table.
name: "kimi anthropic shape",
url: "https://api.moonshot.ai/anthropic/v1/messages",
reqBody: []byte(`{"model":"kimi-k3","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":1000,"output_tokens":1000}}`),
wantProvider: "anthropic",
wantModel: "kimi-k3",
wantCost: 0.018,
},
{
// Vertex path-routed model with "@version" stripped; Anthropic-on-Vertex priced under the anthropic table.
name: "vertex anthropic path-routed",
url: "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-6@20260115:rawPredict",
reqBody: []byte(`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":200,"output_tokens":100}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0021,
},
{
// Gateway-prefixed model ids are not in the pricing table: the meter must SKIP, never guess a rate.
name: "gateway-prefixed model skips pricing",
url: "https://gateway.example.com/v1/chat/completions",
reqBody: []byte(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500}}`),
wantProvider: "openai",
wantModel: "openai/gpt-4o-mini",
wantSkip: "unknown_model",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
in := &middleware.Input{
Method: "POST",
URL: tc.url,
Headers: []middleware.KV{{Key: "Content-Type", Value: "application/json"}},
Body: tc.reqBody,
}
reqOut, err := reqMW.Invoke(context.Background(), in)
require.NoError(t, err, "request parser")
in.Metadata = append(in.Metadata, reqOut.Metadata...)
require.Equal(t, tc.wantProvider, metaKV(in.Metadata, middleware.KeyLLMProvider), "detected provider")
require.Equal(t, tc.wantModel, metaKV(in.Metadata, middleware.KeyLLMModel), "detected (normalized) model")
in.Status = 200
in.RespHeaders = []middleware.KV{{Key: "Content-Type", Value: tc.respCT}}
in.RespBody = tc.respBody
respOut, err := respMW.Invoke(context.Background(), in)
require.NoError(t, err, "response parser")
in.Metadata = append(in.Metadata, respOut.Metadata...)
costOut, err := costMW.Invoke(context.Background(), in)
require.NoError(t, err, "cost meter")
if tc.wantSkip != "" {
assert.Equal(t, tc.wantSkip, metaKV(costOut.Metadata, middleware.KeyCostSkipped), "expected cost skip reason")
assert.Empty(t, metaKV(costOut.Metadata, middleware.KeyCostUSDTotal), "no cost may be emitted on skip")
return
}
raw := metaKV(costOut.Metadata, middleware.KeyCostUSDTotal)
require.NotEmpty(t, raw, "cost.usd_total must be emitted; skip=%q", metaKV(costOut.Metadata, middleware.KeyCostSkipped))
got, err := strconv.ParseFloat(raw, 64)
require.NoError(t, err, "cost must be a float")
// cost.usd_total is rendered with %.6f: allow half of the last printed digit on top of float error.
assert.InDelta(t, tc.wantCost, got, 5.1e-7, "USD cost for %s", tc.name)
rawCache := metaKV(costOut.Metadata, middleware.KeyCostUSDCache)
require.NotEmpty(t, rawCache, "cost.usd_cache must be emitted next to cost.usd_total")
gotCache, err := strconv.ParseFloat(rawCache, 64)
require.NoError(t, err, "cache cost must be a float")
assert.InDelta(t, tc.wantCacheCost, gotCache, 5.1e-7, "cache USD cost for %s", tc.name)
})
}
}
// metaKV returns the value for key in kvs, or "" when absent.
func metaKV(kvs []middleware.KV, key string) string {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value
}
}
return ""
}
// sseBody renders data frames as a text/event-stream body.
func sseBody(frames ...string) []byte {
var b bytes.Buffer
for _, f := range frames {
b.WriteString("data: ")
b.WriteString(f)
b.WriteString("\n\n")
}
return b.Bytes()
}
// awsFrame encodes one AWS event-stream frame with the given :event-type.
func awsFrame(t *testing.T, eventType string, payload []byte) []byte {
t.Helper()
var buf bytes.Buffer
enc := eventstream.NewEncoder()
require.NoError(t, enc.Encode(&buf, eventstream.Message{
Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}},
Payload: payload,
}), "encode event-stream frame")
return buf.Bytes()
}
// bedrockInvokeStream builds an invoke-with-response-stream body: each "chunk" frame wraps a base64 Anthropic event.
func bedrockInvokeStream(t *testing.T, events ...string) []byte {
t.Helper()
var body bytes.Buffer
for _, ev := range events {
wrap, err := json.Marshal(map[string]string{"bytes": base64.StdEncoding.EncodeToString([]byte(ev))})
require.NoError(t, err)
body.Write(awsFrame(t, "chunk", wrap))
}
return body.Bytes()
}
// bedrockConverseStream builds a converse-stream body: contentBlockDelta frames plus a trailing metadata usage frame.
func bedrockConverseStream(t *testing.T, deltas ...string) []byte {
t.Helper()
var body bytes.Buffer
for i, ev := range deltas {
eventType := "contentBlockDelta"
if i == len(deltas)-1 {
eventType = "metadata"
}
body.Write(awsFrame(t, eventType, []byte(ev)))
}
return body.Bytes()
}

View File

@@ -32,12 +32,7 @@ const (
)
var metadataKeys = []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
@@ -145,38 +140,18 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
}
table := m.loader.Get()
costs, ok := table.Costs(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
cost, ok := table.Cost(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
if !ok {
out.Metadata = skip(skipUnknownModel)
return out, nil
}
// Per-bucket costs first: they're the base of the breakdown, and the two
// aggregates that follow are derived from exactly these four values.
out.Metadata = []middleware.KV{
{Key: middleware.KeyCostUSDInput, Value: usd(costs.InputUSD)},
{Key: middleware.KeyCostUSDCachedInput, Value: usd(costs.CachedInputUSD)},
{Key: middleware.KeyCostUSDCacheCreation, Value: usd(costs.CacheCreationUSD)},
{Key: middleware.KeyCostUSDOutput, Value: usd(costs.OutputUSD)},
{Key: middleware.KeyCostUSDTotal, Value: usd(costs.TotalUSD)},
{Key: middleware.KeyCostUSDCache, Value: usd(costs.CacheUSD)},
{Key: middleware.KeyCostUSDTotal, Value: fmt.Sprintf("%.6f", cost)},
}
return out, nil
}
// usd renders a cost as the fixed-precision string every cost.usd_* key
// carries, so the per-bucket values and the aggregates round identically.
//
// 9 decimals, not 6: these values are summed downstream — per request, per
// session, and per usage bucket — so the rounding step is applied once per
// bucket per row and then accumulated. At 6 decimals a single row loses up to
// 2e-6 across its four buckets (enough to break a 1e-6 reconciliation against
// published rates), and a bucket smaller than half a microdollar quantises to
// zero outright: 16 cache-read tokens on a cheap model is 1.6e-9, so summing
// 10k such rows reports 0.02 instead of 0.016. Nano-dollar precision keeps the
// per-row error ~1000x below the smallest realistic bucket.
func usd(v float64) string { return fmt.Sprintf("%.9f", v) }
// skip returns a single-entry metadata slice carrying the given skip
// reason under KeyCostSkipped.
func skip(reason string) []middleware.KV {

View File

@@ -67,15 +67,7 @@ func TestMiddleware_StaticSurface(t *testing.T) {
assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op")
keys := mw.MetadataKeys()
expected := []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
expected := []string{middleware.KeyCostUSDTotal, middleware.KeyCostSkipped}
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
}
@@ -113,7 +105,7 @@ func TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted for known model")
assert.Equal(t, "0.000750000", value, "0.00015 + 0.0006 per 1k tokens, 9-decimal format")
assert.Equal(t, "0.000750", value, "0.00015 + 0.0006 per 1k tokens, 6-decimal format")
}
func TestFactory_PricingPathOverride(t *testing.T) {
@@ -137,7 +129,7 @@ func TestFactory_PricingPathOverride(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted with custom pricing path")
assert.Equal(t, "0.015000000", value, "2*0.0025 + 1*0.01 = 0.015 with 9-decimal format")
assert.Equal(t, "0.015000", value, "2*0.0025 + 1*0.01 = 0.015 with 6-decimal format")
}
func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
@@ -156,7 +148,7 @@ func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted")
assert.Equal(t, "0.018000000", value, "0.003 + 0.015 = 0.018 with 9-decimal format")
assert.Equal(t, "0.018000", value, "0.003 + 0.015 = 0.018 with 6-decimal format")
_, skipped := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
assert.False(t, skipped, "cost.skipped must not be set when cost is computed")
}
@@ -365,25 +357,8 @@ func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cached subset path must produce a cost — never a skip")
// 250 non-cached at 0.0025/1k + 750 cached at 0.00125/1k + 500 output at 0.01/1k.
assert.Equal(t, "0.006562500", value,
assert.Equal(t, "0.006563", value,
"cached subset must be billed at the discount rate, non-cached at the full rate; never double-billed")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 750 cached at 0.00125/1k = 0.0009375.
assert.Equal(t, "0.000937500", cache, "cache cost is the discounted cost of the cached subset")
// Per-bucket breakdown. On OpenAI the cached subset is carved out of the
// input bucket, so input covers only the 250 non-cached tokens — the two
// must never double-count the same 750 tokens.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000625000",
"input bucket bills only the non-cached remainder")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000937500",
"cached-input bucket bills the discounted subset")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.000000000",
"OpenAI has no cache-write bucket")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.005000000",
"output bucket bills 500 tokens at 0.01/1k")
}
// TestInvoke_AnthropicCacheBucketsAdditive proves the Anthropic
@@ -409,33 +384,9 @@ func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 256 input * 0.003 + 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 + 200 output * 0.015
// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184.
assert.Equal(t, "0.005918400", value,
// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184 → "0.005918" with 6-decimal format.
assert.Equal(t, "0.005918", value,
"each Anthropic input bucket must bill at its own rate — cache_read cheap, cache_creation expensive, regular input mid")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 = 0.0021504.
assert.Equal(t, "0.002150400", cache, "cache cost sums the read and creation buckets")
// Per-bucket breakdown: four separately-billed buckets, each at its own rate.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000768000",
"input bucket bills 256 tokens at 0.003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000230400",
"cache-read bucket bills 768 tokens at the cheap 0.0003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.001920000",
"cache-write bucket bills 512 tokens at the expensive 0.00375/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.003000000",
"output bucket bills 200 tokens at 0.015/1k")
}
// assertBucket asserts one per-bucket cost key carries the expected
// 6-decimal value.
func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
t.Helper()
got, ok := metaValue(t, md, key)
require.Truef(t, ok, "%s must be emitted", key)
assert.Equal(t, want, got, msg)
}
// TestInvoke_CachedTokensAbsentFallsBackToBaseFormula covers the
@@ -460,7 +411,7 @@ func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 1000 input * 0.0025 + 500 output * 0.01 = 0.0025 + 0.005 = 0.0075
assert.Equal(t, "0.007500000", value, "no cached metadata = same cost as before the feature landed")
assert.Equal(t, "0.007500", value, "no cached metadata = same cost as before the feature landed")
}
// TestInvoke_UnparseableCachedTokensSkippedSilently proves the
@@ -484,7 +435,7 @@ func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) {
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "garbage cache metadata must NOT switch the response from a cost to a skip — fall back to 0 cached")
assert.Equal(t, "0.007500000", value, "same as the no-cached-metadata path")
assert.Equal(t, "0.007500", value, "same as the no-cached-metadata path")
}
// TestMiddleware_CloseCancelsReloader proves Close stops the per-instance

View File

@@ -69,18 +69,15 @@ func applyBedrockInvokeChunk(payload []byte, usage *llm.Usage, completion *strin
}
// converseStreamEvent captures the Converse stream frames carrying completion
// text (contentBlockDelta) and the final token usage (metadata). Cache buckets
// are additive to inputTokens (AWS write bucket: cacheWriteInputTokens).
// text (contentBlockDelta) and the final token usage (metadata).
type converseStreamEvent struct {
Delta *struct {
Text string `json:"text"`
} `json:"delta"`
Usage *struct {
InputTokens int64 `json:"inputTokens"`
OutputTokens int64 `json:"outputTokens"`
TotalTokens int64 `json:"totalTokens"`
CacheReadTokens int64 `json:"cacheReadInputTokens"`
CacheWriteTokens int64 `json:"cacheWriteInputTokens"`
InputTokens int64 `json:"inputTokens"`
OutputTokens int64 `json:"outputTokens"`
TotalTokens int64 `json:"totalTokens"`
} `json:"usage"`
}
@@ -108,12 +105,6 @@ func applyConverseStreamEvent(eventType string, payload []byte, usage *llm.Usage
if ev.Usage.TotalTokens > 0 {
usage.TotalTokens = ev.Usage.TotalTokens
}
if ev.Usage.CacheReadTokens > 0 {
usage.CachedInputTokens = ev.Usage.CacheReadTokens
}
if ev.Usage.CacheWriteTokens > 0 {
usage.CacheCreationTokens = ev.Usage.CacheWriteTokens
}
}
}
}

View File

@@ -66,24 +66,6 @@ func TestAccumulateBedrockStream_Converse(t *testing.T) {
require.Equal(t, "pong", completion, "converse text deltas concatenated")
}
// The converse-stream metadata frame's camelCase cache fields must reach the billed cache buckets.
func TestAccumulateBedrockStream_ConverseCacheBuckets(t *testing.T) {
var body bytes.Buffer
body.Write(bedrockFrame(t, "contentBlockDelta", mustJSON(t, map[string]any{"delta": map[string]any{"text": "pong"}})))
body.Write(bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{
"inputTokens": 11, "outputTokens": 3, "totalTokens": 30,
"cacheReadInputTokens": 7, "cacheWriteInputTokens": 9,
}})))
usage, completion := accumulateBedrockStream(body.Bytes())
require.Equal(t, int64(11), usage.InputTokens, "input tokens from metadata frame")
require.Equal(t, int64(3), usage.OutputTokens, "output tokens from metadata frame")
require.Equal(t, int64(7), usage.CachedInputTokens, "cache-read tokens from metadata frame")
require.Equal(t, int64(9), usage.CacheCreationTokens, "cache-write tokens from metadata frame")
require.Equal(t, int64(30), usage.TotalTokens, "provider-reported total wins")
require.Equal(t, "pong", completion)
}
func TestAccumulateBedrockStream_Truncated(t *testing.T) {
// A body cut mid-frame must not panic; partial usage is returned.
full := bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{"inputTokens": 11, "outputTokens": 3}}))

View File

@@ -75,19 +75,8 @@ const (
KeyLLMAttributionGroupID = "llm.attribution_group_id"
KeyLLMAttributionWindowS = "llm.attribution_window_seconds"
// Cost metering (emitted by cost_meter). The four per-bucket keys are the
// base of the breakdown — one per token bucket the provider bills
// separately — and the two aggregates below are derived from them:
// usd_total is their sum, usd_cache is cached_input + cache_creation.
KeyCostUSDInput = "cost.usd_input"
// KeyCostUSDCachedInput is the cost of the cache-read bucket (Anthropic cache_read; OpenAI's discounted cached subset of input).
KeyCostUSDCachedInput = "cost.usd_cached_input"
// KeyCostUSDCacheCreation is the cost of the cache-write bucket. Zero for providers without one.
KeyCostUSDCacheCreation = "cost.usd_cache_creation"
KeyCostUSDOutput = "cost.usd_output"
KeyCostUSDTotal = "cost.usd_total"
// KeyCostUSDCache is the portion of cost.usd_total billed for prompt-cache buckets (cache read/creation, or OpenAI's cached input subset).
KeyCostUSDCache = "cost.usd_cache"
// Cost metering (emitted by cost_meter).
KeyCostUSDTotal = "cost.usd_total"
KeyCostSkipped = "cost.skipped"
// Framework-emitted error markers. Use the mw.<id>.* prefix to

View File

@@ -3,8 +3,8 @@
# FreeBSD Port Diff Generator for NetBird
#
# This script generates the diff file required for submitting a FreeBSD port update.
# It works on macOS, Linux, and FreeBSD by fetching files from the FreeBSD ports
# GitHub mirror and computing checksums from the Go module proxy.
# It works on macOS, Linux, and FreeBSD by fetching files from FreeBSD cgit and
# computing checksums from the Go module proxy.
#
# Usage: ./freebsd-port-diff.sh [new_version]
# Example: ./freebsd-port-diff.sh 0.60.7
@@ -14,7 +14,7 @@
set -e
GITHUB_REPO="netbirdio/netbird"
PORTS_MIRROR_BASE="https://raw.githubusercontent.com/freebsd/freebsd-ports/main/security/netbird"
PORTS_CGIT_BASE="https://cgit.freebsd.org/ports/plain/security/netbird"
GO_PROXY="https://proxy.golang.org/github.com/netbirdio/netbird/@v"
OUTPUT_DIR="${OUTPUT_DIR:-.}"
AWK_FIRST_FIELD='{print $1}'
@@ -30,17 +30,10 @@ fetch_all_tags() {
fetch_current_ports_version() {
echo "Fetching current version from FreeBSD ports..." >&2
local makefile version
makefile=$(fetch_ports_file "Makefile") || return 1
version=$(echo "$makefile" | \
curl -sL "${PORTS_CGIT_BASE}/Makefile" 2>/dev/null | \
grep -E "^DISTVERSION=" | \
sed 's/DISTVERSION=[[:space:]]*//' | \
tr -d '\t ')
if [[ -z "$version" ]]; then
echo "Error: Could not extract DISTVERSION from ports Makefile" >&2
return 1
fi
echo "$version"
tr -d '\t '
return 0
}
@@ -52,16 +45,7 @@ fetch_latest_github_release() {
fetch_ports_file() {
local filename="$1"
local content
if ! content=$(curl -fsL --proto '=https' --proto-redir '=https' --retry 3 "${PORTS_MIRROR_BASE}/${filename}" 2>/dev/null); then
echo "Error: Could not fetch ${filename} from ${PORTS_MIRROR_BASE}" >&2
return 1
fi
if [[ "$content" == \<* ]]; then
echo "Error: Received HTML instead of ${filename} from ${PORTS_MIRROR_BASE}" >&2
return 1
fi
printf '%s' "$content"
curl -sL "${PORTS_CGIT_BASE}/${filename}" 2>/dev/null
return 0
}

View File

@@ -9,22 +9,18 @@
# Example: ./freebsd-port-issue-body.sh 0.56.0 0.59.1
#
# If no versions are provided, the script will:
# - Fetch OLD version from the FreeBSD ports GitHub mirror (current version in ports tree)
# - Fetch OLD version from FreeBSD ports cgit (current version in ports tree)
# - Fetch NEW version from latest NetBird GitHub release tag
set -e
GITHUB_REPO="netbirdio/netbird"
PORTS_MAKEFILE_URL="https://raw.githubusercontent.com/freebsd/freebsd-ports/main/security/netbird/Makefile"
PORTS_CGIT_URL="https://cgit.freebsd.org/ports/plain/security/netbird/Makefile"
fetch_current_ports_version() {
echo "Fetching current version from FreeBSD ports..." >&2
local makefile_content
makefile_content=$(curl -fsL --proto '=https' --proto-redir '=https' --retry 3 "$PORTS_MAKEFILE_URL" 2>/dev/null) || makefile_content=""
if [[ "$makefile_content" == \<* ]]; then
echo "Error: Received HTML instead of Makefile from ${PORTS_MAKEFILE_URL}" >&2
return 1
fi
makefile_content=$(curl -sL "$PORTS_CGIT_URL" 2>/dev/null)
if [[ -z "$makefile_content" ]]; then
echo "Error: Could not fetch Makefile from FreeBSD ports" >&2
return 1

View File

@@ -187,16 +187,16 @@ func (c *GrpcClient) ready() bool {
// Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages
// Blocking request. The result will be sent via msgHandler callback function
func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler, backOff)
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key) error {
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler)
})
}
// Job wraps the real client's Job endpoint call and takes care of retries and encryption/decryption of messages
// Blocking request. The result will be sent via msgHandler callback function
func (c *GrpcClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error {
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
return c.handleJobStream(ctx, serverPubKey, msgHandler, backOff)
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key) error {
return c.handleJobStream(ctx, serverPubKey, msgHandler)
})
}
@@ -204,7 +204,7 @@ func (c *GrpcClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequ
// It takes care of retries, connection readiness, and fetching server public key.
func (c *GrpcClient) withMgmtStream(
ctx context.Context,
handler func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error,
handler func(ctx context.Context, serverPubKey wgtypes.Key) error,
) error {
backOff := defaultBackoff(ctx)
operation := func() error {
@@ -224,7 +224,7 @@ func (c *GrpcClient) withMgmtStream(
return err
}
return handler(ctx, *serverPubKey, backOff)
return handler(ctx, *serverPubKey)
}
err := backoff.Retry(operation, backOff)
@@ -239,7 +239,6 @@ func (c *GrpcClient) handleJobStream(
ctx context.Context,
serverPubKey wgtypes.Key,
msgHandler func(msg *proto.JobRequest) *proto.JobResponse,
backOff backoff.BackOff,
) error {
ctx, cancelStream := context.WithCancel(ctx)
defer cancelStream()
@@ -257,19 +256,6 @@ func (c *GrpcClient) handleJobStream(
log.Debug("job stream handshake sent successfully")
// The stream is up, so reset the backoff. This matters for two reasons,
// both caused by the backoff lib not resetting its state on a successful
// connection:
// 1. Without a reset, after a connect followed by an error the next retry
// starts from the accumulated (large) interval instead of retrying
// promptly, delaying reconnection.
// 2. Worse, once the accumulated elapsed time exceeds MaxElapsedTime, the
// next stream error makes NextBackOff() return Stop, so the retry loop
// exits immediately. That error is then mislabeled unrecoverable and
// bubbles up to trigger a full engine restart / data-plane teardown
// instead of a silent reconnection.
backOff.Reset()
// Main loop: receive, process, respond
for {
jobReq, err := c.receiveJobRequest(ctx, stream, serverPubKey)
@@ -385,7 +371,7 @@ func (c *GrpcClient) sendJobResponse(
return nil
}
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
ctx, cancelStream := context.WithCancel(ctx)
defer cancelStream()
@@ -404,19 +390,6 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
c.notifyConnected()
c.setSyncStreamConnected()
// The stream is up, so reset the backoff. This matters for two reasons,
// both caused by the backoff lib not resetting its state on a successful
// connection:
// 1. Without a reset, after a connect followed by an error the next retry
// starts from the accumulated (large) interval instead of retrying
// promptly, delaying reconnection.
// 2. Worse, once the accumulated elapsed time exceeds MaxElapsedTime, the
// next stream error makes NextBackOff() return Stop, so the retry loop
// exits immediately. That error is then mislabeled unrecoverable and
// bubbles up to trigger a full engine restart / data-plane teardown
// instead of a silent reconnection.
backOff.Reset()
// blocking until error
err = c.receiveUpdatesEvents(stream, serverPubKey, msgHandler)
if err != nil {

View File

@@ -5822,48 +5822,13 @@ components:
total_tokens:
type: integer
format: int64
description: Total tokens consumed, including prompt-cache tokens.
description: Total tokens consumed.
example: 1840
cached_input_tokens:
type: integer
format: int64
description: Input tokens read from the provider's prompt cache. Additive to input_tokens for Anthropic-shape providers; a subset of input_tokens for OpenAI.
example: 0
cache_creation_tokens:
type: integer
format: int64
description: Input tokens written to the provider's prompt cache. Zero for providers without a cache-write bucket.
example: 30528
cost_usd:
type: number
format: double
description: Estimated USD cost of the request.
example: 0.0231
input_cost_usd:
type: number
format: double
description: Cost of the non-cached input tokens. Base component of cost_usd.
example: 0.0048
cached_input_cost_usd:
type: number
format: double
description: Cost of the prompt-cache read tokens. Base component of cost_usd, and part of cache_cost_usd.
example: 0.0015
cache_creation_cost_usd:
type: number
format: double
description: Cost of the prompt-cache write tokens. Base component of cost_usd, and part of cache_cost_usd.
example: 0.1130
output_cost_usd:
type: number
format: double
description: Cost of the output tokens. Base component of cost_usd.
example: 0.0038
cache_cost_usd:
type: number
format: double
description: Portion of cost_usd billed for prompt-cache usage.
example: 0.1145
stream:
type: boolean
description: Whether the request was a streaming completion.
@@ -5887,14 +5852,7 @@ components:
- input_tokens
- output_tokens
- total_tokens
- cached_input_tokens
- cache_creation_tokens
- input_cost_usd
- cached_input_cost_usd
- cache_creation_cost_usd
- output_cost_usd
- cost_usd
- cache_cost_usd
AgentNetworkAccessLogsResponse:
type: object
properties:
@@ -5968,48 +5926,13 @@ components:
total_tokens:
type: integer
format: int64
description: Total tokens across the session, including prompt-cache tokens.
description: Total tokens across the session.
example: 12880
cached_input_tokens:
type: integer
format: int64
description: Total prompt-cache read tokens across the session.
example: 0
cache_creation_tokens:
type: integer
format: int64
description: Total prompt-cache write tokens across the session.
example: 30528
cost_usd:
type: number
format: double
description: Total estimated USD cost across the session.
example: 0.1617
input_cost_usd:
type: number
format: double
description: Total cost of non-cached input tokens across the session.
example: 0.0210
cached_input_cost_usd:
type: number
format: double
description: Total cost of prompt-cache read tokens across the session.
example: 0.0015
cache_creation_cost_usd:
type: number
format: double
description: Total cost of prompt-cache write tokens across the session.
example: 0.1130
output_cost_usd:
type: number
format: double
description: Total cost of output tokens across the session.
example: 0.0262
cache_cost_usd:
type: number
format: double
description: Portion of cost_usd billed for prompt-cache usage across the session.
example: 0.1145
providers:
type: array
items:
@@ -6036,14 +5959,7 @@ components:
- input_tokens
- output_tokens
- total_tokens
- cached_input_tokens
- cache_creation_tokens
- input_cost_usd
- cached_input_cost_usd
- cache_creation_cost_usd
- output_cost_usd
- cost_usd
- cache_cost_usd
- decision
- entries
AgentNetworkAccessLogSessionsResponse:
@@ -6097,61 +6013,19 @@ components:
total_tokens:
type: integer
format: int64
description: Total tokens in the bucket, including prompt-cache tokens.
description: Total tokens in the bucket.
example: 184000
cached_input_tokens:
type: integer
format: int64
description: Total prompt-cache read tokens in the bucket.
example: 20000
cache_creation_tokens:
type: integer
format: int64
description: Total prompt-cache write tokens in the bucket.
example: 45000
input_cost_usd:
type: number
format: double
description: Total cost of non-cached input tokens in the bucket.
example: 1.12
cached_input_cost_usd:
type: number
format: double
description: Total cost of prompt-cache read tokens in the bucket.
example: 0.06
cache_creation_cost_usd:
type: number
format: double
description: Total cost of prompt-cache write tokens in the bucket.
example: 0.36
output_cost_usd:
type: number
format: double
description: Total cost of output tokens in the bucket.
example: 0.77
cost_usd:
type: number
format: double
description: Total estimated USD spend in the bucket.
example: 2.31
cache_cost_usd:
type: number
format: double
description: Portion of cost_usd billed for prompt-cache usage in the bucket.
example: 0.42
required:
- period_start
- input_tokens
- output_tokens
- total_tokens
- cached_input_tokens
- cache_creation_tokens
- input_cost_usd
- cached_input_cost_usd
- cache_creation_cost_usd
- output_cost_usd
- cost_usd
- cache_cost_usd
AgentNetworkSettings:
type: object
description: Per-account Agent Network gateway settings. One row per account; cluster and subdomain are auto-assigned on first provider create and immutable thereafter.

View File

@@ -1732,21 +1732,6 @@ type AccountSettings struct {
// AgentNetworkAccessLog One per-request agent-network (LLM) access log entry with flattened, queryable LLM dimensions.
type AgentNetworkAccessLog struct {
// CacheCostUsd Portion of cost_usd billed for prompt-cache usage.
CacheCostUsd float64 `json:"cache_cost_usd"`
// CacheCreationCostUsd Cost of the prompt-cache write tokens. Base component of cost_usd, and part of cache_cost_usd.
CacheCreationCostUsd float64 `json:"cache_creation_cost_usd"`
// CacheCreationTokens Input tokens written to the provider's prompt cache. Zero for providers without a cache-write bucket.
CacheCreationTokens int64 `json:"cache_creation_tokens"`
// CachedInputCostUsd Cost of the prompt-cache read tokens. Base component of cost_usd, and part of cache_cost_usd.
CachedInputCostUsd float64 `json:"cached_input_cost_usd"`
// CachedInputTokens Input tokens read from the provider's prompt cache. Additive to input_tokens for Anthropic-shape providers; a subset of input_tokens for OpenAI.
CachedInputTokens int64 `json:"cached_input_tokens"`
// CostUsd Estimated USD cost of the request.
CostUsd float64 `json:"cost_usd"`
@@ -1768,9 +1753,6 @@ type AgentNetworkAccessLog struct {
// Id Unique identifier for the access log entry.
Id string `json:"id"`
// InputCostUsd Cost of the non-cached input tokens. Base component of cost_usd.
InputCostUsd float64 `json:"input_cost_usd"`
// InputTokens Input (prompt) tokens consumed.
InputTokens int64 `json:"input_tokens"`
@@ -1780,9 +1762,6 @@ type AgentNetworkAccessLog struct {
// Model Requested LLM model.
Model *string `json:"model,omitempty"`
// OutputCostUsd Cost of the output tokens. Base component of cost_usd.
OutputCostUsd float64 `json:"output_cost_usd"`
// OutputTokens Output (completion) tokens produced.
OutputTokens int64 `json:"output_tokens"`
@@ -1822,7 +1801,7 @@ type AgentNetworkAccessLog struct {
// Timestamp Timestamp when the request was made.
Timestamp time.Time `json:"timestamp"`
// TotalTokens Total tokens consumed, including prompt-cache tokens.
// TotalTokens Total tokens consumed.
TotalTokens int64 `json:"total_tokens"`
// UserId NetBird user id of the authenticated caller, if applicable.
@@ -1831,21 +1810,6 @@ type AgentNetworkAccessLog struct {
// AgentNetworkAccessLogSession A session-grouped view of agent-network access logs — all requests sharing a session id (or a single session-less request) folded into one summary plus its ordered entries.
type AgentNetworkAccessLogSession struct {
// CacheCostUsd Portion of cost_usd billed for prompt-cache usage across the session.
CacheCostUsd float64 `json:"cache_cost_usd"`
// CacheCreationCostUsd Total cost of prompt-cache write tokens across the session.
CacheCreationCostUsd float64 `json:"cache_creation_cost_usd"`
// CacheCreationTokens Total prompt-cache write tokens across the session.
CacheCreationTokens int64 `json:"cache_creation_tokens"`
// CachedInputCostUsd Total cost of prompt-cache read tokens across the session.
CachedInputCostUsd float64 `json:"cached_input_cost_usd"`
// CachedInputTokens Total prompt-cache read tokens across the session.
CachedInputTokens int64 `json:"cached_input_tokens"`
// CostUsd Total estimated USD cost across the session.
CostUsd float64 `json:"cost_usd"`
@@ -1861,18 +1825,12 @@ type AgentNetworkAccessLogSession struct {
// GroupIds Union of the authorising group ids across the session's entries.
GroupIds *[]string `json:"group_ids,omitempty"`
// InputCostUsd Total cost of non-cached input tokens across the session.
InputCostUsd float64 `json:"input_cost_usd"`
// InputTokens Total input (prompt) tokens across the session.
InputTokens int64 `json:"input_tokens"`
// Models Distinct models seen in the session.
Models *[]string `json:"models,omitempty"`
// OutputCostUsd Total cost of output tokens across the session.
OutputCostUsd float64 `json:"output_cost_usd"`
// OutputTokens Total output (completion) tokens across the session.
OutputTokens int64 `json:"output_tokens"`
@@ -1888,7 +1846,7 @@ type AgentNetworkAccessLogSession struct {
// StartedAt Timestamp of the session's earliest request.
StartedAt time.Time `json:"started_at"`
// TotalTokens Total tokens across the session, including prompt-cache tokens.
// TotalTokens Total tokens across the session.
TotalTokens int64 `json:"total_tokens"`
// UserId NetBird user id of the session's caller.
@@ -2389,40 +2347,19 @@ type AgentNetworkSettingsRequest struct {
// AgentNetworkUsageBucket One aggregated agent-network usage time bucket (UTC). The bucket width is set by the request's granularity.
type AgentNetworkUsageBucket struct {
// CacheCostUsd Portion of cost_usd billed for prompt-cache usage in the bucket.
CacheCostUsd float64 `json:"cache_cost_usd"`
// CacheCreationCostUsd Total cost of prompt-cache write tokens in the bucket.
CacheCreationCostUsd float64 `json:"cache_creation_cost_usd"`
// CacheCreationTokens Total prompt-cache write tokens in the bucket.
CacheCreationTokens int64 `json:"cache_creation_tokens"`
// CachedInputCostUsd Total cost of prompt-cache read tokens in the bucket.
CachedInputCostUsd float64 `json:"cached_input_cost_usd"`
// CachedInputTokens Total prompt-cache read tokens in the bucket.
CachedInputTokens int64 `json:"cached_input_tokens"`
// CostUsd Total estimated USD spend in the bucket.
CostUsd float64 `json:"cost_usd"`
// InputCostUsd Total cost of non-cached input tokens in the bucket.
InputCostUsd float64 `json:"input_cost_usd"`
// InputTokens Total input (prompt) tokens in the bucket.
InputTokens int64 `json:"input_tokens"`
// OutputCostUsd Total cost of output tokens in the bucket.
OutputCostUsd float64 `json:"output_cost_usd"`
// OutputTokens Total output (completion) tokens in the bucket.
OutputTokens int64 `json:"output_tokens"`
// PeriodStart Start of the bucket in YYYY-MM-DD (UTC) — the day, the week start (Monday), or the month start, depending on granularity.
PeriodStart string `json:"period_start"`
// TotalTokens Total tokens in the bucket, including prompt-cache tokens.
// TotalTokens Total tokens in the bucket.
TotalTokens int64 `json:"total_tokens"`
}

View File

@@ -11,6 +11,9 @@ import (
log "github.com/sirupsen/logrus"
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/proto"
@@ -35,17 +38,17 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
Network: decodeAccountNetwork(full.Network),
AccountSettings: decodeAccountSettings(full.AccountSettings),
CustomZoneDomain: full.CustomZoneDomain,
Peers: make(map[string]*types.ComponentPeer, len(full.Peers)),
Groups: make(map[string]*types.ComponentGroup, len(full.Groups)),
Peers: make(map[string]*nbpeer.Peer, len(full.Peers)),
Groups: make(map[string]*types.Group, len(full.Groups)),
Policies: make([]*types.Policy, 0, len(full.Policies)),
Routes: make([]*nbroute.Route, 0, len(full.Routes)),
NameServerGroups: make([]*nbdns.NameServerGroup, 0, len(full.NameserverGroups)),
AllDNSRecords: decodeSimpleRecords(full.AllDnsRecords),
AccountZones: decodeCustomZones(full.AccountZones),
ResourcePoliciesMap: make(map[string][]*types.Policy),
RoutersMap: make(map[string]map[string]*types.ComponentRouter),
NetworkResources: make([]*types.ComponentResource, 0, len(full.NetworkResources)),
RouterPeers: make(map[string]*types.ComponentPeer),
RoutersMap: make(map[string]map[string]*routerTypes.NetworkRouter),
NetworkResources: make([]*resourceTypes.NetworkResource, 0, len(full.NetworkResources)),
RouterPeers: make(map[string]*nbpeer.Peer),
AllowedUserIDs: stringSliceToSet(full.AllowedUserIds),
PostureFailedPeers: make(map[string]map[string]struct{}, len(full.PostureFailedPeers)),
GroupIDToUserIDs: make(map[string][]string, len(full.GroupIdToUserIds)),
@@ -98,7 +101,7 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
log.WithField("peer idx", idx).Error("unrecognized peer idx during decoding")
}
}
group := &types.ComponentGroup{
group := &types.Group{
ID: groupID,
PublicID: gc.Id,
Peers: peerIDs,
@@ -148,7 +151,7 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
// Phase 7: routers_map (outer key = network seq id, inner key = peer-id
// reconstructed from peer_index). Synthesized network id is "net_<seq>".
for networkID, list := range full.RoutersMap {
inner := make(map[string]*types.ComponentRouter, len(list.Entries))
inner := make(map[string]*routerTypes.NetworkRouter, len(list.Entries))
for _, entry := range list.Entries {
if !entry.PeerIndexSet {
continue
@@ -158,7 +161,8 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
continue
}
peerID := peerIDByIndex[entry.PeerIndex]
inner[peerID] = &types.ComponentRouter{
inner[peerID] = &routerTypes.NetworkRouter{
ID: "",
NetworkID: networkID,
PublicID: entry.Id,
Peer: peerID,
@@ -260,22 +264,40 @@ func decodeAccountSettings(as *proto.AccountSettingsCompact) *types.AccountSetti
}
}
func decodePeerCompact(pc *proto.PeerCompact, peerID string) *types.ComponentPeer {
peer := &types.ComponentPeer{
func decodePeerCompact(pc *proto.PeerCompact, peerID string) *nbpeer.Peer {
var caps []int32
if pc.SupportsSourcePrefixes {
caps = append(caps, nbpeer.PeerCapabilitySourcePrefixes)
}
if pc.SupportsIpv6 {
caps = append(caps, nbpeer.PeerCapabilityIPv6Overlay)
}
peer := &nbpeer.Peer{
ID: peerID,
Key: peerID,
SSHKey: string(pc.SshPubKey),
SSHEnabled: pc.SshEnabled,
DNSLabel: pc.DnsLabel,
LoginExpirationEnabled: pc.LoginExpirationEnabled,
AgentVersion: pc.AgentVersion,
SupportsSourcePrefixes: pc.SupportsSourcePrefixes,
SupportsIPv6: pc.SupportsIpv6,
ServerSSHAllowed: pc.ServerSshAllowed,
AddedWithSSOLogin: pc.AddedWithSsoLogin,
Meta: nbpeer.PeerSystemMeta{
WtVersion: pc.AgentVersion,
Capabilities: caps,
Flags: nbpeer.Flags{
ServerSSHAllowed: pc.ServerSshAllowed,
},
},
}
if pc.AddedWithSsoLogin {
// Set a non-empty UserID so (*Peer).AddedWithSSOLogin() returns true.
// The original UserID isn't on the wire; the value is intentionally
// visibly synthetic so any future consumer that mistakes UserID for a
// real account user xid won't silently match (or worse, write the
// sentinel into a downstream record).
peer.UserID = "<env-sso>"
}
if pc.LastLoginUnixNano != 0 {
peer.LastLogin = time.Unix(0, pc.LastLoginUnixNano)
t := time.Unix(0, pc.LastLoginUnixNano)
peer.LastLogin = &t
}
switch len(pc.Ip) {
case 4:
@@ -402,14 +424,14 @@ func decodeNameServerGroupRaw(nsg *proto.NameServerGroupRaw) *nbdns.NameServerGr
return out
}
func decodeNetworkResource(nr *proto.NetworkResourceRaw) *types.ComponentResource {
out := &types.ComponentResource{
func decodeNetworkResource(nr *proto.NetworkResourceRaw) *resourceTypes.NetworkResource {
out := &resourceTypes.NetworkResource{
ID: nr.Id,
PublicID: nr.Id,
NetworkID: nr.NetworkSeq,
Name: nr.Name,
Description: nr.Description,
Type: types.ComponentResourceType(nr.Type),
Type: resourceTypes.NetworkResourceType(nr.Type),
Address: nr.Address,
Domain: nr.DomainValue,
Enabled: nr.Enabled,

View File

@@ -20,9 +20,10 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"net/netip"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/shared/management/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/types"
"github.com/netbirdio/netbird/shared/netiputil"
"github.com/netbirdio/netbird/shared/sshauth"
)
@@ -273,7 +274,7 @@ func ToProtocolDNSConfig(update nbdns.Config, cache DNSConfigCache, forwardPort
// AppendRemotePeerConfig appends typed peers as proto.RemotePeerConfig
// entries to dst and returns the result.
func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*types.ComponentPeer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
for _, rPeer := range peers {
allowedIPs := []string{rPeer.IP.String() + "/32"}
if includeIPv6 && rPeer.IPv6.IsValid() {
@@ -284,7 +285,7 @@ func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*types.Compon
AllowedIps: allowedIPs,
SshConfig: &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
Fqdn: rPeer.FQDN(dnsName),
AgentVersion: rPeer.AgentVersion,
AgentVersion: rPeer.Meta.WtVersion,
})
}
return dst

View File

@@ -54,8 +54,8 @@ func EnvelopeToNetworkMap(ctx context.Context, env *proto.NetworkMapEnvelope, lo
}
components.PeerID = canonicalKey
includeIPv6 := localPeer.SupportsIPv6 && localPeer.IPv6.IsValid()
useSourcePrefixes := localPeer.SupportsSourcePrefixes
includeIPv6 := localPeer.SupportsIPv6() && localPeer.IPv6.IsValid()
useSourcePrefixes := localPeer.SupportsSourcePrefixes()
typedNM := components.Calculate(ctx)

View File

@@ -13,6 +13,7 @@ import (
goproto "google.golang.org/protobuf/proto"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
nbnetworkmap "github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/proto"
@@ -143,14 +144,14 @@ func TestDecodeEnvelope_MalformedWgKeyPeerSkipped(t *testing.T) {
func TestEnvelopeRoundTrip_AllGroupShortCircuitParity(t *testing.T) {
ctx := context.Background()
peers := map[string]*types.ComponentPeer{}
peers := map[string]*nbpeer.Peer{}
for i, id := range []string{"peer-T", "peer-S", "peer-ALL", "peer-O"} {
peers[id] = &types.ComponentPeer{
ID: id,
Key: randomWgKey(t),
IP: netip.AddrFrom4([4]byte{100, 64, 0, byte(i + 1)}),
DNSLabel: id,
AgentVersion: "0.40.0",
peers[id] = &nbpeer.Peer{
ID: id,
Key: randomWgKey(t),
IP: netip.AddrFrom4([4]byte{100, 64, 0, byte(i + 1)}),
DNSLabel: id,
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
}
@@ -164,7 +165,7 @@ func TestEnvelopeRoundTrip_AllGroupShortCircuitParity(t *testing.T) {
AccountSettings: &types.AccountSettingsInfo{},
DNSSettings: &types.DNSSettings{},
Peers: peers,
Groups: map[string]*types.ComponentGroup{
Groups: map[string]*types.Group{
"g-src": {ID: "g-src", PublicID: "1", Name: "staff", Peers: []string{"peer-T", "peer-S"}},
"g-all": {ID: "g-all", PublicID: "2", Name: "All", Peers: []string{"peer-ALL"}},
"g-two": {ID: "g-two", PublicID: "3", Name: "second", Peers: []string{"peer-T", "peer-O"}},
@@ -231,22 +232,22 @@ func buildSmokeComponents(t *testing.T) (*types.NetworkMapComponents, string) {
peerAKey := randomWgKey(t)
peerBKey := randomWgKey(t)
peerA := &types.ComponentPeer{
ID: "peer-A",
Key: peerAKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peerA",
AgentVersion: "0.40.0",
peerA := &nbpeer.Peer{
ID: "peer-A",
Key: peerAKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peerA",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
peerB := &types.ComponentPeer{
ID: "peer-B",
Key: peerBKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
DNSLabel: "peerB",
AgentVersion: "0.40.0",
peerB := &nbpeer.Peer{
ID: "peer-B",
Key: peerBKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
DNSLabel: "peerB",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
}
group := &types.ComponentGroup{
group := &types.Group{
ID: "group-all", PublicID: "1", Name: "All",
Peers: []string{"peer-A", "peer-B"},
}
@@ -273,11 +274,11 @@ func buildSmokeComponents(t *testing.T) (*types.NetworkMapComponents, string) {
},
AccountSettings: &types.AccountSettingsInfo{},
DNSSettings: &types.DNSSettings{},
Peers: map[string]*types.ComponentPeer{
Peers: map[string]*nbpeer.Peer{
"peer-A": peerA,
"peer-B": peerB,
},
Groups: map[string]*types.ComponentGroup{
Groups: map[string]*types.Group{
"group-all": group,
},
Policies: []*types.Policy{policy},

View File

@@ -1,103 +0,0 @@
package types
import (
"net/netip"
"time"
)
// ComponentPeer is the self-contained peer representation used by
// NetworkMapComponents and the calculated NetworkMap. It carries exactly the
// subset of peer data that crosses the components wire format, so the shared
// calculation layer stays independent of the management server's domain
// types.
type ComponentPeer struct {
ID string
Key string
IP netip.Addr
IPv6 netip.Addr
DNSLabel string
SSHKey string
SSHEnabled bool
ServerSSHAllowed bool
AgentVersion string
SupportsSourcePrefixes bool
SupportsIPv6 bool
LoginExpirationEnabled bool
AddedWithSSOLogin bool
LastLogin time.Time
}
// FQDN returns the peer's FQDN combined of the peer's DNS label and the system's DNS domain.
func (p *ComponentPeer) FQDN(dnsDomain string) string {
if dnsDomain == "" {
return ""
}
return p.DNSLabel + "." + dnsDomain
}
// LoginExpired indicates whether the peer's login has expired, mirroring the
// server-side peer semantics: only SSO-added peers with login expiration
// enabled can expire.
func (p *ComponentPeer) LoginExpired(expiresIn time.Duration) (bool, time.Duration) {
if !p.AddedWithSSOLogin || !p.LoginExpirationEnabled {
return false, 0
}
timeLeft := time.Until(p.LastLogin.Add(expiresIn))
return timeLeft <= 0, timeLeft
}
// GroupAllName is the reserved name of the default group that contains every peer in an account.
const GroupAllName = "All"
// ComponentGroup is the self-contained group representation used by
// NetworkMapComponents: just the membership view the network-map calculation
// needs, without the server's storage fields.
type ComponentGroup struct {
ID string
PublicID string
Name string
Peers []string
}
// IsGroupAll checks if the group is a default "All" group.
func (g *ComponentGroup) IsGroupAll() bool {
return g.Name == GroupAllName
}
// ComponentRouter is the self-contained network-router representation used by
// NetworkMapComponents.
type ComponentRouter struct {
NetworkID string
PublicID string
Peer string
PeerGroups []string
Masquerade bool
Metric int
Enabled bool
}
// ComponentResourceType mirrors the network-resource type enum on the
// components wire format.
type ComponentResourceType string
const (
ComponentResourceHost ComponentResourceType = "host"
ComponentResourceSubnet ComponentResourceType = "subnet"
ComponentResourceDomain ComponentResourceType = "domain"
)
// ComponentResource is the self-contained network-resource representation
// used by NetworkMapComponents.
type ComponentResource struct {
ID string
PublicID string
NetworkID string
AccountID string
Name string
Description string
Type ComponentResourceType
Address string
Domain string
Prefix netip.Prefix
Enabled bool
}

View File

@@ -3,6 +3,8 @@ package types
import (
"strconv"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/version"
)
@@ -46,8 +48,8 @@ func portsIncludesSSH(ports []string) bool {
}
// ExpandPortsAndRanges expands Ports and PortRanges of a rule into individual firewall rules.
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *ComponentPeer) []*FirewallRule {
features := peerSupportedFirewallFeatures(peer.AgentVersion)
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *nbpeer.Peer) []*FirewallRule {
features := peerSupportedFirewallFeatures(peer.Meta.WtVersion)
var expanded []*FirewallRule
@@ -104,8 +106,8 @@ func isPortInRule(portString string, portInt uint16, rule *FirewallRule) bool {
return rule.Port == portString || (rule.PortRange.Start <= portInt && portInt <= rule.PortRange.End)
}
func shouldCheckRulesForNativeSSH(supportsNative bool, rule *PolicyRule, peer *ComponentPeer) bool {
return supportsNative && peer.SSHEnabled && peer.ServerSSHAllowed && rule.Protocol == PolicyRuleProtocolTCP
func shouldCheckRulesForNativeSSH(supportsNative bool, rule *PolicyRule, peer *nbpeer.Peer) bool {
return supportsNative && peer.SSHEnabled && peer.Meta.Flags.ServerSSHAllowed && rule.Protocol == PolicyRuleProtocolTCP
}
func peerSupportedFirewallFeatures(peerVer string) supportedFeatures {
@@ -115,13 +117,13 @@ func peerSupportedFirewallFeatures(peerVer string) supportedFeatures {
var features supportedFeatures
meetMinVer, err := version.MeetsMinVersion(firewallRuleMinNativeSSHVer, peerVer)
meetMinVer, err := posture.MeetsMinVersion(firewallRuleMinNativeSSHVer, peerVer)
features.nativeSSH = err == nil && meetMinVer
if features.nativeSSH {
features.portRanges = true
} else {
meetMinVer, err = version.MeetsMinVersion(firewallRuleMinPortRangesVer, peerVer)
meetMinVer, err = posture.MeetsMinVersion(firewallRuleMinPortRangesVer, peerVer)
features.portRanges = err == nil && meetMinVer
}

View File

@@ -9,6 +9,7 @@ import (
log "github.com/sirupsen/logrus"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
)
@@ -50,7 +51,7 @@ func (r *FirewallRule) Equal(other *FirewallRule) bool {
// For static routes, source ranges match the destination family (v4 or v6).
// For dynamic routes (domain-based), separate v4 and v6 rules are generated
// so the routing peer's forwarding chain allows both address families.
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*ComponentPeer, direction int, includeIPv6 bool) []*RouteFirewallRule {
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int, includeIPv6 bool) []*RouteFirewallRule {
rulesExists := make(map[string]struct{})
rules := make([]*RouteFirewallRule, 0)
@@ -106,7 +107,7 @@ func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule
}
// splitPeerSourcesByFamily separates peer IPs into v4 (/32) and v6 (/128) source ranges.
func splitPeerSourcesByFamily(groupPeers []*ComponentPeer) (v4, v6 []string) {
func splitPeerSourcesByFamily(groupPeers []*nbpeer.Peer) (v4, v6 []string) {
v4 = make([]string, 0, len(groupPeers))
v6 = make([]string, 0, len(groupPeers))
for _, peer := range groupPeers {

View File

@@ -8,12 +8,13 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/domain"
)
func TestSplitPeerSourcesByFamily(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
@@ -35,7 +36,7 @@ func TestSplitPeerSourcesByFamily(t *testing.T) {
}
func TestGenerateRouteFirewallRules_V4Route(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
@@ -64,7 +65,7 @@ func TestGenerateRouteFirewallRules_V4Route(t *testing.T) {
}
func TestGenerateRouteFirewallRules_V6Route(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
@@ -92,7 +93,7 @@ func TestGenerateRouteFirewallRules_V6Route(t *testing.T) {
}
func TestGenerateRouteFirewallRules_DynamicRoute_DualStack(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),
@@ -125,7 +126,7 @@ func TestGenerateRouteFirewallRules_DynamicRoute_DualStack(t *testing.T) {
}
func TestGenerateRouteFirewallRules_DynamicRoute_NoV6Peers(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{IP: netip.MustParseAddr("100.64.0.1")},
{IP: netip.MustParseAddr("100.64.0.2")},
}
@@ -149,7 +150,7 @@ func TestGenerateRouteFirewallRules_DynamicRoute_NoV6Peers(t *testing.T) {
}
func TestGenerateRouteFirewallRules_IncludeIPv6False(t *testing.T) {
peers := []*ComponentPeer{
peers := []*nbpeer.Peer{
{
IP: netip.MustParseAddr("100.64.0.1"),
IPv6: netip.MustParseAddr("fd00::1"),

View File

@@ -2,6 +2,7 @@ package types
import (
"github.com/netbirdio/netbird/management/server/integration_reference"
"github.com/netbirdio/netbird/management/server/networks/resources/types"
)
const (
@@ -67,6 +68,10 @@ func (g *Group) EventMeta() map[string]any {
return map[string]any{"name": g.Name}
}
func (g *Group) EventMetaResource(resource *types.NetworkResource) map[string]any {
return map[string]any{"name": g.Name, "id": g.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
func (g *Group) Copy() *Group {
group := &Group{
ID: g.ID,
@@ -90,39 +95,14 @@ func (g *Group) HasPeers() bool {
return len(g.Peers) > 0
}
// GroupAllName is the reserved name of the default group that contains every peer in an account.
const GroupAllName = "All"
// IsGroupAll checks if the group is a default "All" group.
func (g *Group) IsGroupAll() bool {
return g.Name == GroupAllName
}
// ToComponent converts the group to its self-contained components
// representation. The Peers slice is shared, not copied — components are
// treated as immutable snapshots. Returns nil for a nil group.
func (g *Group) ToComponent() *ComponentGroup {
if g == nil {
return nil
}
return &ComponentGroup{
ID: g.ID,
PublicID: g.PublicID,
Name: g.Name,
Peers: g.Peers,
}
}
// GroupsToComponent converts an id-keyed group map to its components
// representation, preserving nil entries.
func GroupsToComponent(groups map[string]*Group) map[string]*ComponentGroup {
if groups == nil {
return nil
}
out := make(map[string]*ComponentGroup, len(groups))
for id, g := range groups {
out[id] = g.ToComponent()
}
return out
}
// AddPeer adds peerID to Peers if not present, returning true if added.
func (g *Group) AddPeer(peerID string) bool {
if peerID == "" {

Some files were not shown because too many files have changed in this diff Show More