mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-28 03:11:28 +02:00
Compare commits
19 Commits
grpc-acl
...
worktree-d
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d918235afd | ||
|
|
f7af852851 | ||
|
|
4546908ceb | ||
|
|
4f6247b5c3 | ||
|
|
1e5b0a5c89 | ||
|
|
b65ec8b68a | ||
|
|
e13bcdbd44 | ||
|
|
46568f7af8 | ||
|
|
3358138ccc | ||
|
|
6d15d0729a | ||
|
|
0936918d24 | ||
|
|
d4a4418969 | ||
|
|
31ed241a1a | ||
|
|
178e6a8530 | ||
|
|
96963b6751 | ||
|
|
9770814f39 | ||
|
|
b0c1ed31b8 | ||
|
|
8435682ac8 | ||
|
|
ed682fad87 |
12
.github/workflows/agent-network-e2e.yml
vendored
12
.github/workflows/agent-network-e2e.yml
vendored
@@ -5,6 +5,13 @@ 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 }}
|
||||
@@ -51,6 +58,9 @@ 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 }}
|
||||
@@ -59,6 +69,8 @@ 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 }}
|
||||
|
||||
2
.github/workflows/frontend-ui.yml
vendored
2
.github/workflows/frontend-ui.yml
vendored
@@ -86,7 +86,7 @@ jobs:
|
||||
${{ runner.os }}-pnpm-
|
||||
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
run: pnpm install --frozen-lockfile --ignore-scripts
|
||||
|
||||
- name: Generate Wails bindings
|
||||
run: pnpm run bindings
|
||||
|
||||
4
.github/workflows/golangci-lint.yml
vendored
4
.github/workflows/golangci-lint.yml
vendored
@@ -45,7 +45,7 @@ jobs:
|
||||
display_name: Linux
|
||||
name: ${{ matrix.display_name }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 25
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
@@ -79,4 +79,4 @@ jobs:
|
||||
skip-cache: true
|
||||
skip-save-cache: true
|
||||
cache-invalidation-interval: 0
|
||||
args: --timeout=12m
|
||||
args: --timeout=20m
|
||||
|
||||
@@ -464,6 +464,8 @@ 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()
|
||||
@@ -505,12 +507,8 @@ func Test_ConnectPeers(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
localIP, err := getLocalIP()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer1wgPort))
|
||||
localIP1 := "127.0.0.1"
|
||||
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP1, peer1wgPort))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -546,7 +544,8 @@ func Test_ConnectPeers(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer2wgPort))
|
||||
localIP2 := "127.0.0.1"
|
||||
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP2, peer2wgPort))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -569,17 +568,17 @@ func Test_ConnectPeers(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// todo: investigate why in some tests execution we need 30s
|
||||
// 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.
|
||||
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)
|
||||
@@ -588,6 +587,12 @@ 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:
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -615,28 +620,3 @@ 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")
|
||||
}
|
||||
|
||||
@@ -351,6 +351,7 @@ func (a *Auth) setSystemInfoFlags(info *system.Info) {
|
||||
a.config.BlockLANAccess,
|
||||
a.config.BlockInbound,
|
||||
a.config.DisableIPv6,
|
||||
a.config.SyncMessageVersion,
|
||||
a.config.EnableSSHRoot,
|
||||
a.config.EnableSSHSFTP,
|
||||
a.config.EnableSSHLocalPortForwarding,
|
||||
|
||||
@@ -34,6 +34,8 @@ const (
|
||||
// - Handling connection establishment based on peer signaling
|
||||
//
|
||||
// The implementation is not thread-safe; it is protected by engine.syncMsgMux.
|
||||
// The only exception is ActivatePeer, which is safe for concurrent use so the
|
||||
// DNS warm-up path can call it without contending on the engine mutex.
|
||||
type ConnMgr struct {
|
||||
peerStore *peerstore.Store
|
||||
statusRecorder *peer.Status
|
||||
@@ -42,12 +44,26 @@ type ConnMgr struct {
|
||||
rosenpassEnabled bool
|
||||
|
||||
lazyConnMgr *manager.Manager
|
||||
// lazyConnMgrMu guards the lazyConnMgr pointer for readers outside the
|
||||
// engine loop (ActivatePeer). Writers hold it in addition to
|
||||
// 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,
|
||||
@@ -238,12 +254,20 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) {
|
||||
conn.Log.Infof("removed peer from lazy conn manager")
|
||||
}
|
||||
|
||||
// ActivatePeer wakes an idle lazy connection. Unlike the rest of ConnMgr it is
|
||||
// safe for concurrent use: the lazy manager pointer is read under lazyConnMgrMu
|
||||
// and the manager itself is internally synchronized, so callers outside the
|
||||
// engine loop (DNS warm-up) do not need engine.syncMsgMux.
|
||||
func (e *ConnMgr) ActivatePeer(ctx context.Context, conn *peer.Conn) {
|
||||
if !e.isStartedWithLazyMgr() {
|
||||
e.lazyConnMgrMu.RLock()
|
||||
lazyConnMgr := e.lazyConnMgr
|
||||
started := lazyConnMgr != nil && e.lazyCtxCancel != nil
|
||||
e.lazyConnMgrMu.RUnlock()
|
||||
if !started {
|
||||
return
|
||||
}
|
||||
|
||||
if found := e.lazyConnMgr.ActivatePeer(conn.GetKey()); found {
|
||||
if found := lazyConnMgr.ActivatePeer(conn.GetKey()); found {
|
||||
if err := conn.Open(ctx); err != nil {
|
||||
conn.Log.Errorf("failed to open connection: %v", err)
|
||||
}
|
||||
@@ -268,16 +292,22 @@ func (e *ConnMgr) Close() {
|
||||
|
||||
e.lazyCtxCancel()
|
||||
e.wg.Wait()
|
||||
|
||||
e.lazyConnMgrMu.Lock()
|
||||
e.lazyConnMgr = nil
|
||||
e.lazyConnMgrMu.Unlock()
|
||||
}
|
||||
|
||||
func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
|
||||
cfg := manager.Config{
|
||||
InactivityThreshold: inactivityThresholdEnv(),
|
||||
ReconcileAllowedIPs: e.reconcileRoutedIPs,
|
||||
}
|
||||
e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
|
||||
|
||||
e.lazyConnMgrMu.Lock()
|
||||
e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
|
||||
e.lazyCtx, e.lazyCtxCancel = context.WithCancel(engineCtx)
|
||||
e.lazyConnMgrMu.Unlock()
|
||||
|
||||
e.wg.Add(1)
|
||||
go func() {
|
||||
@@ -316,7 +346,10 @@ func (e *ConnMgr) closeManager(ctx context.Context) {
|
||||
|
||||
e.lazyCtxCancel()
|
||||
e.wg.Wait()
|
||||
|
||||
e.lazyConnMgrMu.Lock()
|
||||
e.lazyConnMgr = nil
|
||||
e.lazyConnMgrMu.Unlock()
|
||||
|
||||
for _, peerID := range e.peerStore.PeersPubKey() {
|
||||
e.peerStore.PeerConnOpen(ctx, peerID)
|
||||
|
||||
@@ -1,10 +1,21 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
|
||||
|
||||
"github.com/netbirdio/netbird/client/iface/wgaddr"
|
||||
"github.com/netbirdio/netbird/client/internal/lazyconn"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/peerstore"
|
||||
"github.com/netbirdio/netbird/monotime"
|
||||
)
|
||||
|
||||
func TestResolveLazyForce(t *testing.T) {
|
||||
@@ -38,3 +49,58 @@ func TestResolveLazyForce(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type mockLazyWGIface struct{}
|
||||
|
||||
func (mockLazyWGIface) RemovePeer(string) error { return nil }
|
||||
func (mockLazyWGIface) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error {
|
||||
return nil
|
||||
}
|
||||
func (mockLazyWGIface) IsUserspaceBind() bool { return false }
|
||||
func (mockLazyWGIface) Address() wgaddr.Address { return wgaddr.Address{} }
|
||||
func (mockLazyWGIface) LastActivities() map[string]monotime.Time { return nil }
|
||||
func (mockLazyWGIface) MTU() uint16 { return 1280 }
|
||||
|
||||
// TestConnMgr_ActivatePeerConcurrentWithLifecycle exercises ActivatePeer from
|
||||
// non-engine goroutines (the DNS warm-up path) racing the manager lifecycle,
|
||||
// which stays on the engine loop. Run with -race: it fails if ActivatePeer
|
||||
// still requires engine.syncMsgMux for safety.
|
||||
func TestConnMgr_ActivatePeerConcurrentWithLifecycle(t *testing.T) {
|
||||
t.Setenv(lazyconn.EnvLazyConn, "on")
|
||||
|
||||
status := peer.NewRecorder("https://mgm")
|
||||
store := peerstore.NewConnStore()
|
||||
connMgr := NewConnMgr(&EngineConfig{}, status, store, mockLazyWGIface{})
|
||||
|
||||
conn := newTestPeerConn(t, "peerA")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
connMgr.Start(ctx)
|
||||
|
||||
done := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
for range 4 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
default:
|
||||
connMgr.ActivatePeer(ctx, conn)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Let the activators spin against the started manager, then tear it down
|
||||
// underneath them and let them spin against the stopped manager.
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
connMgr.Close()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
close(done)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
@@ -621,6 +621,7 @@ func createEngineConfig(key wgtypes.Key, config *profilemanager.Config, peerConf
|
||||
BlockLANAccess: config.BlockLANAccess,
|
||||
BlockInbound: config.BlockInbound,
|
||||
DisableIPv6: config.DisableIPv6,
|
||||
SyncMessageVersion: config.SyncMessageVersion,
|
||||
|
||||
LazyConnection: lazyconn.ParseState(config.LazyConnection),
|
||||
|
||||
@@ -696,6 +697,7 @@ func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte,
|
||||
config.BlockLANAccess,
|
||||
config.BlockInbound,
|
||||
config.DisableIPv6,
|
||||
config.SyncMessageVersion,
|
||||
config.EnableSSHRoot,
|
||||
config.EnableSSHSFTP,
|
||||
config.EnableSSHLocalPortForwarding,
|
||||
|
||||
@@ -676,6 +676,7 @@ func (g *BundleGenerator) addCommonConfigFields(configContent *strings.Builder)
|
||||
configContent.WriteString(fmt.Sprintf("BlockLANAccess: %v\n", g.internalConfig.BlockLANAccess))
|
||||
configContent.WriteString(fmt.Sprintf("BlockInbound: %v\n", g.internalConfig.BlockInbound))
|
||||
configContent.WriteString(fmt.Sprintf("DisableIPv6: %v\n", g.internalConfig.DisableIPv6))
|
||||
configContent.WriteString(fmt.Sprintf("SyncMessageVersion: %v\n", g.internalConfig.SyncMessageVersion))
|
||||
|
||||
if g.internalConfig.DisableNotifications != nil {
|
||||
configContent.WriteString(fmt.Sprintf("DisableNotifications: %v\n", *g.internalConfig.DisableNotifications))
|
||||
|
||||
@@ -887,6 +887,8 @@ func TestAddConfig_AllFieldsCovered(t *testing.T) {
|
||||
ClientCertKeyPath: "/tmp/key",
|
||||
LazyConnection: "on",
|
||||
MTU: 1280,
|
||||
DisableIPv6: true,
|
||||
SyncMessageVersion: func(v int) *int { return &v }(1),
|
||||
}
|
||||
|
||||
for _, anonymize := range []bool{false, true} {
|
||||
|
||||
@@ -3,6 +3,7 @@ package dns
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math"
|
||||
"net"
|
||||
"slices"
|
||||
@@ -31,6 +32,26 @@ type SubdomainMatcher interface {
|
||||
MatchSubdomains() bool
|
||||
}
|
||||
|
||||
// responseMeta holds the annotations handlers attach to a request to explain the
|
||||
// response the chain ends up writing. It survives a deferral, so an answer that
|
||||
// did not come from the handler that owns the name still says who stepped aside
|
||||
// and why.
|
||||
type responseMeta map[resutil.MetaKey]string
|
||||
|
||||
// format renders the annotations for the response log line. The order is stable
|
||||
// so the same event reads the same way every time; a map's own order is not.
|
||||
func (m responseMeta) format() string {
|
||||
if len(m) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
for _, k := range slices.Sorted(maps.Keys(m)) {
|
||||
b.WriteString(" " + string(k) + "=" + m[k])
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
type HandlerEntry struct {
|
||||
Handler dns.Handler
|
||||
Priority int
|
||||
@@ -52,8 +73,19 @@ type ResponseWriterChain struct {
|
||||
origPattern string
|
||||
requestID string
|
||||
shouldContinue bool
|
||||
response *dns.Msg
|
||||
meta map[string]string
|
||||
// softNegative suppresses a poisoning negative verdict for this request. A
|
||||
// handler that owns the name but cannot answer this query type sets it
|
||||
// before deferring, and it stays set for every handler that runs after.
|
||||
softNegative bool
|
||||
// clientHasEdns records whether the original query advertised EDNS0, taken
|
||||
// before any handler ran: handlers add EDNS0 to the query they forward
|
||||
// upstream, so the message itself no longer answers the question later.
|
||||
clientHasEdns bool
|
||||
response *dns.Msg
|
||||
// meta is handed to the next handler when this one defers, so the same map
|
||||
// outlives the writer that created it. Handlers must only set metadata from
|
||||
// within ServeDNS, never from a goroutine that outlives the call.
|
||||
meta responseMeta
|
||||
}
|
||||
|
||||
// RequestID returns the request ID for tracing
|
||||
@@ -62,26 +94,64 @@ func (w *ResponseWriterChain) RequestID() string {
|
||||
}
|
||||
|
||||
// SetMeta sets a metadata key-value pair for logging
|
||||
func (w *ResponseWriterChain) SetMeta(key, value string) {
|
||||
func (w *ResponseWriterChain) SetMeta(key resutil.MetaKey, value string) {
|
||||
if w.meta == nil {
|
||||
w.meta = make(map[string]string)
|
||||
w.meta = make(responseMeta)
|
||||
}
|
||||
w.meta[key] = value
|
||||
}
|
||||
|
||||
// RequestSoftNegative marks the request so a downstream NXDOMAIN is turned into
|
||||
// NODATA before it reaches the client. Set by a handler that defers a query for
|
||||
// a name it owns but a type it cannot resolve.
|
||||
func (w *ResponseWriterChain) RequestSoftNegative() {
|
||||
w.softNegative = true
|
||||
}
|
||||
|
||||
func (w *ResponseWriterChain) WriteMsg(m *dns.Msg) error {
|
||||
// Check if this is a continue signal (NXDOMAIN with Zero bit set)
|
||||
if m.Rcode == dns.RcodeNameError && m.MsgHdr.Zero {
|
||||
w.shouldContinue = true
|
||||
return nil
|
||||
}
|
||||
if w.softNegative && m.Rcode == dns.RcodeNameError {
|
||||
m = softenNegative(m, w.clientHasEdns)
|
||||
w.SetMeta(resutil.MetaKeySoftened, "nxdomain->nodata")
|
||||
}
|
||||
w.response = m
|
||||
if m.MsgHdr.Truncated {
|
||||
w.SetMeta("truncated", "true")
|
||||
w.SetMeta(resutil.MetaKeyTruncated, "true")
|
||||
}
|
||||
return w.ResponseWriter.WriteMsg(m)
|
||||
}
|
||||
|
||||
// softenNegative downgrades an NXDOMAIN to NODATA for a request a handler that
|
||||
// owns the name deferred. NXDOMAIN is cached for the name and every type below
|
||||
// it, so a resolver that has never heard of a routed name would take out the
|
||||
// record types the route does serve; NODATA is cached for this name and type
|
||||
// only. The authority section goes with it: the negative TTL of a zone we just
|
||||
// overruled does not apply, and RFC 2308 keeps a negative answer that carries no
|
||||
// SOA out of caches altogether, so the rewrite cannot outlive the route.
|
||||
//
|
||||
// The rewrite is ours, not the answering resolver's, so an EDNS0 client is told
|
||||
// as much: the reply we hand back travels a path no capture on this host sees,
|
||||
// and an empty answer is otherwise indistinguishable from a real one. Returns a
|
||||
// copy so the answering handler keeps whatever it may hold on to.
|
||||
func softenNegative(m *dns.Msg, clientHasEdns bool) *dns.Msg {
|
||||
out := m.Copy()
|
||||
out.Rcode = dns.RcodeSuccess
|
||||
out.Ns = nil
|
||||
|
||||
if clientHasEdns {
|
||||
resutil.AttachEDE(out, resutil.EDENetbirdSoftenedNegative,
|
||||
"netbird: name is served locally, NXDOMAIN from the fallthrough resolver suppressed")
|
||||
} else {
|
||||
resutil.StripOPT(out)
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
func NewHandlerChain() *HandlerChain {
|
||||
return &HandlerChain{
|
||||
handlers: make([]HandlerEntry, 0),
|
||||
@@ -223,6 +293,19 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
handlers := slices.Clone(c.handlers)
|
||||
c.mu.RUnlock()
|
||||
|
||||
// Carried across deferrals: once a handler that owns the name defers, no
|
||||
// handler after it may answer with a verdict that poisons the name. The
|
||||
// metadata of a handler that stepped aside is carried too, so the response
|
||||
// log line explains an answer that did not come from the handler that owns
|
||||
// the name.
|
||||
var softNegative bool
|
||||
var carried responseMeta
|
||||
|
||||
// Taken before any handler runs: handlers advertise EDNS0 on the query they
|
||||
// forward upstream, so afterwards the message no longer tells us whether the
|
||||
// client did.
|
||||
clientHasEdns := r.IsEdns0() != nil
|
||||
|
||||
// Try handlers in priority order
|
||||
for _, entry := range handlers {
|
||||
if entry.Priority > maxPriority {
|
||||
@@ -245,11 +328,16 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
ResponseWriter: w,
|
||||
origPattern: entry.OrigPattern,
|
||||
requestID: requestID,
|
||||
softNegative: softNegative,
|
||||
clientHasEdns: clientHasEdns,
|
||||
meta: carried,
|
||||
}
|
||||
entry.Handler.ServeDNS(chainWriter, r)
|
||||
|
||||
// If handler wants to continue, try next handler
|
||||
if chainWriter.shouldContinue {
|
||||
softNegative = softNegative || chainWriter.softNegative
|
||||
carried = chainWriter.meta
|
||||
if entry.Priority != PriorityMgmtCache {
|
||||
logger.Tracef("handler requested continue for domain=%s", qname)
|
||||
}
|
||||
@@ -265,30 +353,59 @@ func (c *HandlerChain) dispatch(w dns.ResponseWriter, r *dns.Msg, maxPriority in
|
||||
qname, dns.TypeToString[question.Qtype], dns.ClassToString[question.Qclass])
|
||||
resp := &dns.Msg{}
|
||||
resp.SetRcode(r, dns.RcodeRefused)
|
||||
// A handler that owns the name deferred and nothing below it could answer
|
||||
// (a client with no primary nameserver group). The name exists as far as
|
||||
// this client is concerned, since the route serves its addresses, so REFUSED
|
||||
// would contradict the route: a stub that takes it as "not served here" and
|
||||
// retries another resolver can come back with an NXDOMAIN that takes the
|
||||
// whole name down. Answer "no records of this type" instead, without an SOA,
|
||||
// so RFC 2308 keeps it out of negative caches, and tell an EDNS0 client the
|
||||
// empty answer is ours rather than a resolver's.
|
||||
if softNegative {
|
||||
resp.Rcode = dns.RcodeSuccess
|
||||
if clientHasEdns {
|
||||
resutil.AttachEDE(resp, resutil.EDENetbirdSoftenedNegative,
|
||||
"netbird: name is served locally, no fallthrough resolver for this query type")
|
||||
}
|
||||
// logResponse never runs on this path, so the carried metadata is
|
||||
// appended here or the reason for the deferral is lost in exactly the
|
||||
// case that is hardest to diagnose.
|
||||
logger.Tracef("no handler below the deferring one for domain=%s type=%s, answering NODATA%s",
|
||||
qname, dns.TypeToString[question.Qtype], carried.format())
|
||||
}
|
||||
if err := w.WriteMsg(resp); err != nil {
|
||||
logger.Errorf("failed to write DNS response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *HandlerChain) logResponse(logger *log.Entry, cw *ResponseWriterChain, qname string, startTime time.Time) {
|
||||
// Runs for every query, and the arguments below are not free: Len() packs
|
||||
// the message to measure it, and formatting the answers and the metadata
|
||||
// allocates. None of it is worth doing when the line is discarded.
|
||||
if !log.IsLevelEnabled(log.TraceLevel) {
|
||||
return
|
||||
}
|
||||
|
||||
if cw.response == nil {
|
||||
return
|
||||
}
|
||||
|
||||
var meta string
|
||||
for k, v := range cw.meta {
|
||||
meta += " " + k + "=" + v
|
||||
}
|
||||
|
||||
logger.Tracef("response: domain=%s rcode=%s answers=%s size=%dB%s took=%s",
|
||||
qname, dns.RcodeToString[cw.response.Rcode], resutil.FormatAnswers(cw.response.Answer),
|
||||
cw.response.Len(), meta, time.Since(startTime))
|
||||
cw.response.Len(), cw.meta.format(), time.Since(startTime))
|
||||
}
|
||||
|
||||
// ResolveInternal runs an in-process DNS query against the chain, skipping any
|
||||
// handler with priority > maxPriority. Used by internal callers (e.g. the mgmt
|
||||
// cache refresher) that must bypass themselves to avoid loops. Honors ctx
|
||||
// cancellation; on ctx.Done the dispatch goroutine is left to drain on its own
|
||||
// cache refresher) that must bypass themselves to avoid loops.
|
||||
//
|
||||
// "Nothing answered" is read off RcodeRefused, which a request soft-negatived by
|
||||
// a deferring handler never carries: it ends in an empty NOERROR instead, and
|
||||
// would look resolved. No caller can reach that today, since every handler that
|
||||
// defers sits above the maxPriority any caller passes. Lowering one below it
|
||||
// means this check needs the soft-negative case too.
|
||||
//
|
||||
// Honors ctx cancellation; on ctx.Done the dispatch goroutine is left to drain on its own
|
||||
// (bounded by the invoked handler's internal timeout).
|
||||
func (c *HandlerChain) ResolveInternal(ctx context.Context, r *dns.Msg, maxPriority int) (*dns.Msg, error) {
|
||||
if len(r.Question) == 0 {
|
||||
|
||||
@@ -3,15 +3,19 @@ package dns_test
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
log "github.com/sirupsen/logrus"
|
||||
logtest "github.com/sirupsen/logrus/hooks/test"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/client/internal/dns"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/test"
|
||||
)
|
||||
|
||||
@@ -1238,6 +1242,276 @@ func TestHandlerChain_ResolveInternal_HonorsContextTimeout(t *testing.T) {
|
||||
assert.Less(t, elapsed, 500*time.Millisecond, "ResolveInternal must return shortly after ctx deadline")
|
||||
}
|
||||
|
||||
// requestSoftNegative asks the chain to soften a negative verdict produced by
|
||||
// the handlers that run after the caller defers, reporting whether the chain
|
||||
// supports the signal. Written as a type assertion so these tests compile
|
||||
// against a chain that does not support it yet.
|
||||
func requestSoftNegative(w dns.ResponseWriter) bool {
|
||||
sn, ok := w.(interface{ RequestSoftNegative() })
|
||||
if ok {
|
||||
sn.RequestSoftNegative()
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
// deferringHandler defers to the next handler in the chain, optionally asking
|
||||
// for the negative verdict of whatever answers instead to be softened. This is
|
||||
// what a DNS route handler does for a record type its routing peer cannot
|
||||
// resolve.
|
||||
type deferringHandler struct {
|
||||
softNegative bool
|
||||
called bool
|
||||
supported bool
|
||||
}
|
||||
|
||||
func (h *deferringHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
h.called = true
|
||||
if h.softNegative {
|
||||
h.supported = requestSoftNegative(w)
|
||||
// The real handler records why it stepped aside, for the response log
|
||||
// line of whichever handler answers instead.
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredBy, "test handler")
|
||||
}
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
resp.MsgHdr.Zero = true
|
||||
_ = w.WriteMsg(resp)
|
||||
}
|
||||
|
||||
// nxdomainHandler answers an authoritative NXDOMAIN with an SOA in the
|
||||
// authority section, the way a public resolver answers for a name that only
|
||||
// exists inside the routed network.
|
||||
type nxdomainHandler struct{}
|
||||
|
||||
func (h *nxdomainHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
resp.Ns = []dns.RR{&dns.SOA{
|
||||
Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: 3600},
|
||||
Ns: "ns1.example.com.",
|
||||
Mbox: "hostmaster.example.com.",
|
||||
Minttl: 3600,
|
||||
Expire: 604800,
|
||||
Refresh: 7200,
|
||||
Retry: 3600,
|
||||
}}
|
||||
_ = w.WriteMsg(resp)
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN is the whole point
|
||||
// of the soft-negative signal: a route handler may only defer a query to the
|
||||
// public chain if the answer cannot poison the routed name. NXDOMAIN is cached
|
||||
// for the name and every type under it (RFC 2308, RFC 8020), so it has to be
|
||||
// rewritten to NODATA, which is cached per name and type only.
|
||||
func TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.True(t, route.called, "the route handler must run first")
|
||||
require.True(t, route.supported, "the chain writer must accept a soft-negative request")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "NXDOMAIN must be softened to NODATA")
|
||||
assert.Empty(t, resp.Answer, "a softened negative carries no answer")
|
||||
assert.Empty(t, resp.Ns, "the downstream zone's SOA must not set the negative TTL for a name we overrode")
|
||||
}
|
||||
|
||||
// responseLineFor returns the chain's response log line for one query. Raising
|
||||
// the level to trace also unmutes whatever else is logging in this package, so
|
||||
// the line has to be picked by the name it was asked about rather than by being
|
||||
// the last one seen.
|
||||
func responseLineFor(hook *logtest.Hook, qname string) string {
|
||||
for _, e := range hook.AllEntries() {
|
||||
if strings.HasPrefix(e.Message, "response:") && strings.Contains(e.Message, qname) {
|
||||
return e.Message
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_IsVisibleToTheClient covers observability of the
|
||||
// rewrite. The reply we hand the application travels over loopback, which the
|
||||
// bundle capture does not see, so a softened verdict has to say so on the wire:
|
||||
// without it an empty answer is indistinguishable from a real "no such record"
|
||||
// in a dig output or a capture taken next to the application.
|
||||
func TestHandlerChain_SoftNegative_IsVisibleToTheClient(t *testing.T) {
|
||||
// One hook for the whole test: logtest installs it on the standard logger
|
||||
// and logrus has no way to take it off again, so a hook per subtest would
|
||||
// leave several behind buffering every later line in the package.
|
||||
hook := logtest.NewGlobal()
|
||||
t.Cleanup(hook.Reset)
|
||||
|
||||
newChain := func() *nbdns.HandlerChain {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
chain.AddHandler("*.example.com.", &deferringHandler{softNegative: true}, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
return chain
|
||||
}
|
||||
|
||||
t.Run("EDNS0 client gets an extended error", func(t *testing.T) {
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
r.SetEdns0(dns.DefaultMsgSize, false)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp)
|
||||
require.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
|
||||
ede, ok := resutil.ExtractEDE(resp)
|
||||
require.True(t, ok, "a softened verdict must carry an extended DNS error")
|
||||
assert.Equal(t, resutil.EDENetbirdSoftenedNegative, ede.InfoCode)
|
||||
assert.Contains(t, ede.ExtraText, "netbird", "the text must name us as the source of the rewrite")
|
||||
})
|
||||
|
||||
t.Run("plain client gets no OPT", func(t *testing.T) {
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp)
|
||||
assert.Nil(t, resp.IsEdns0(), "RFC 6891 forbids an OPT toward a client that did not advertise EDNS0")
|
||||
})
|
||||
|
||||
// The response line is what support reads out of a debug bundle, and it now
|
||||
// carries several annotations at once. Built from a map, their order would
|
||||
// differ on every query, so the same event never looks the same twice.
|
||||
t.Run("log fields keep a stable order", func(t *testing.T) {
|
||||
const qname = "_mongodb._tcp.stable.example.com."
|
||||
|
||||
prev := log.GetLevel()
|
||||
log.SetLevel(log.TraceLevel)
|
||||
t.Cleanup(func() { log.SetLevel(prev) })
|
||||
|
||||
lineFor := func() string {
|
||||
hook.Reset()
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion(qname, dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
line := responseLineFor(hook, qname)
|
||||
// The duration differs per query and is not what we compare.
|
||||
line, _, _ = strings.Cut(line, " took=")
|
||||
return line
|
||||
}
|
||||
|
||||
first := lineFor()
|
||||
require.NotEmpty(t, first, "the chain must log the response it wrote")
|
||||
require.Contains(t, first, "deferred_by=", "the line must carry more than one annotation to be worth ordering")
|
||||
require.Contains(t, first, "softened=")
|
||||
|
||||
for range 20 {
|
||||
assert.Equal(t, first, lineFor(), "the same event must produce the same line")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("logged with the reason it was deferred", func(t *testing.T) {
|
||||
const qname = "_mongodb._tcp.reason.example.com."
|
||||
|
||||
hook.Reset()
|
||||
|
||||
prev := log.GetLevel()
|
||||
log.SetLevel(log.TraceLevel)
|
||||
t.Cleanup(func() { log.SetLevel(prev) })
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion(qname, dns.TypeSRV)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
response := responseLineFor(hook, qname)
|
||||
require.NotEmpty(t, response, "the chain must log the response it wrote")
|
||||
assert.Contains(t, response, "softened=", "the log must show the verdict was rewritten")
|
||||
assert.Contains(t, response, "deferred_by=", "the log must name the handler that deferred")
|
||||
})
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_NoHandlerBelow covers a client with no primary
|
||||
// nameserver group: the deferred query reaches the end of the chain unanswered.
|
||||
// REFUSED would say the name is not served here while the route serves its
|
||||
// addresses, and a stub that acts on that by asking elsewhere can bring back an
|
||||
// NXDOMAIN for the whole name. The answer must be an empty, uncacheable NODATA.
|
||||
func TestHandlerChain_SoftNegative_NoHandlerBelow(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", dns.TypeHTTPS)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "an unanswered soft-negative query must be NODATA, not REFUSED")
|
||||
assert.Empty(t, resp.Answer)
|
||||
assert.Empty(t, resp.Ns,
|
||||
"no SOA, so RFC 2308 keeps the empty answer out of negative caches and it cannot outlive the route")
|
||||
}
|
||||
|
||||
// TestHandlerChain_SoftNegative_KeepsRealAnswers guards the other direction:
|
||||
// softening applies to negative verdicts only. A real answer from a downstream
|
||||
// handler must reach the client untouched.
|
||||
func TestHandlerChain_SoftNegative_KeepsRealAnswers(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{softNegative: true}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &answeringHandler{name: "public", ip: "203.0.113.10"}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", dns.TypeA)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
require.Len(t, resp.Answer, 1, "the downstream answer must pass through")
|
||||
}
|
||||
|
||||
// TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative makes sure the
|
||||
// softening is opt-in: an ordinary chain continuation still yields NXDOMAIN, so
|
||||
// genuine non-existence keeps being reported.
|
||||
func TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
|
||||
route := &deferringHandler{}
|
||||
chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
|
||||
chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("nope.example.com.", dns.TypeA)
|
||||
|
||||
mw := &test.MockResponseWriter{}
|
||||
chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)
|
||||
|
||||
resp := mw.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must reach the client")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "without the signal a real NXDOMAIN must survive")
|
||||
}
|
||||
|
||||
func TestHandlerChain_HasRootHandlerAtOrBelow(t *testing.T) {
|
||||
chain := nbdns.NewHandlerChain()
|
||||
h := &answeringHandler{name: "h", ip: "10.0.0.1"}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -36,7 +37,43 @@ type resolver interface {
|
||||
// record is left alone (it points at something outside our mesh, e.g.
|
||||
// a non-peer upstream).
|
||||
type PeerConnectivity interface {
|
||||
IsConnectedByIP(ip string) (known, connected bool)
|
||||
IsConnectedByIP(ip netip.Addr) (known, connected bool)
|
||||
}
|
||||
|
||||
// PeerActivator wakes lazy-connection peers on demand. The local resolver calls
|
||||
// it with the tunnel IPs an answer points at, so a peer that is idle (lazily
|
||||
// disconnected) starts connecting at DNS-resolution time rather than racing the
|
||||
// client's first request packet. nil disables warm-up.
|
||||
type PeerActivator interface {
|
||||
// ActivatePeersByIP triggers wake-up for the peer(s) owning addrs and blocks
|
||||
// until one is connected or ctx (a short per-query budget) expires. It is a
|
||||
// fast no-op for unknown or already-connected addresses.
|
||||
ActivatePeersByIP(ctx context.Context, addrs []netip.Addr)
|
||||
}
|
||||
|
||||
const (
|
||||
defaultLazyWarmupTimeout = 2 * time.Second
|
||||
envLazyWarmupTimeout = "NB_DNS_LAZY_WARMUP_TIMEOUT"
|
||||
)
|
||||
|
||||
// lazyWarmupTimeoutFromEnv returns the per-query budget for waking a
|
||||
// lazy-connection peer a DNS answer points at. Tunable via
|
||||
// NB_DNS_LAZY_WARMUP_TIMEOUT (a Go duration). Parsed once at construction time.
|
||||
func lazyWarmupTimeoutFromEnv() time.Duration {
|
||||
v := os.Getenv(envLazyWarmupTimeout)
|
||||
if v == "" {
|
||||
return defaultLazyWarmupTimeout
|
||||
}
|
||||
d, err := time.ParseDuration(v)
|
||||
if err != nil {
|
||||
log.Warnf("invalid %s value %q, using default %s: %v", envLazyWarmupTimeout, v, defaultLazyWarmupTimeout, err)
|
||||
return defaultLazyWarmupTimeout
|
||||
}
|
||||
if d <= 0 {
|
||||
log.Warnf("non-positive %s value %q, using default %s", envLazyWarmupTimeout, v, defaultLazyWarmupTimeout)
|
||||
return defaultLazyWarmupTimeout
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
type Resolver struct {
|
||||
@@ -51,6 +88,12 @@ type Resolver struct {
|
||||
// filter and preserves the legacy "return whatever is registered"
|
||||
// behaviour for callers that never wire a status source.
|
||||
peerConn PeerConnectivity
|
||||
// peerActivator, when non-nil, is called at resolution time to warm the
|
||||
// lazy connection to the peer(s) an answer points at. nil disables warm-up.
|
||||
peerActivator PeerActivator
|
||||
// warmupTimeout is the per-query budget for the lazy-connection warm-up
|
||||
// wait, resolved from the environment once at construction time.
|
||||
warmupTimeout time.Duration
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
@@ -59,11 +102,12 @@ type Resolver struct {
|
||||
func NewResolver() *Resolver {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Resolver{
|
||||
records: make(map[dns.Question][]dns.RR),
|
||||
domains: make(map[domain.Domain]struct{}),
|
||||
zones: make(map[domain.Domain]bool),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
records: make(map[dns.Question][]dns.RR),
|
||||
domains: make(map[domain.Domain]struct{}),
|
||||
zones: make(map[domain.Domain]bool),
|
||||
warmupTimeout: lazyWarmupTimeoutFromEnv(),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,6 +120,14 @@ func (d *Resolver) SetPeerConnectivity(p PeerConnectivity) {
|
||||
d.peerConn = p
|
||||
}
|
||||
|
||||
// SetPeerActivator wires the DNS-time lazy-connection warm-up. Pass nil to
|
||||
// disable. Safe to call multiple times; the latest value wins.
|
||||
func (d *Resolver) SetPeerActivator(a PeerActivator) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.peerActivator = a
|
||||
}
|
||||
|
||||
func (d *Resolver) MatchSubdomains() bool {
|
||||
return true
|
||||
}
|
||||
@@ -122,6 +174,9 @@ func (d *Resolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
replyMessage.RecursionAvailable = true
|
||||
|
||||
result := d.lookupRecords(logger, question)
|
||||
// Warm before filtering: activation flips a lazily-idle target to connected,
|
||||
// which then lets it survive the disconnected-peer filter below.
|
||||
d.warmLazyPeers(question, result.records)
|
||||
result.records = d.filterDisconnectedPeerAnswers(logger, question, result.records)
|
||||
replyMessage.Authoritative = !result.hasExternalData
|
||||
replyMessage.Answer = result.records
|
||||
@@ -495,8 +550,8 @@ func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns
|
||||
kept := make([]dns.RR, 0, len(records))
|
||||
var dropped int
|
||||
for _, rr := range records {
|
||||
ip := extractRecordIP(rr)
|
||||
if ip == "" {
|
||||
ip, ok := extractRecordAddr(rr)
|
||||
if !ok {
|
||||
kept = append(kept, rr)
|
||||
continue
|
||||
}
|
||||
@@ -518,22 +573,57 @@ func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns
|
||||
return kept
|
||||
}
|
||||
|
||||
// extractRecordIP returns the dotted-decimal / colon-hex IP carried by
|
||||
// an A or AAAA record, or "" for any other record type.
|
||||
func extractRecordIP(rr dns.RR) string {
|
||||
// warmLazyPeers triggers lazy-connection wake-up for the peers a resolved
|
||||
// answer points at and waits briefly for one to connect, so the caller's first
|
||||
// request doesn't race the connection establishment. Warm-up is scoped to
|
||||
// match-only (non-authoritative) zones — the synthesized private-service zones
|
||||
// and user-created zones whose records point at specific peers. The account's
|
||||
// peer zone is authoritative, so plain peer-name lookups never trigger warm-up;
|
||||
// otherwise resolving any peer's name would wake its idle connection, defeating
|
||||
// laziness mesh-wide. No-op when no activator is wired (lazy connections
|
||||
// disabled) or the answer carries no peer IPs.
|
||||
func (d *Resolver) warmLazyPeers(question dns.Question, records []dns.RR) {
|
||||
if len(records) < 2 {
|
||||
return
|
||||
}
|
||||
d.mu.RLock()
|
||||
activator := d.peerActivator
|
||||
var nonAuth, found bool
|
||||
if activator != nil {
|
||||
nonAuth, found = d.findZone(question.Name)
|
||||
}
|
||||
d.mu.RUnlock()
|
||||
if activator == nil || !found || !nonAuth {
|
||||
return
|
||||
}
|
||||
|
||||
var addrs []netip.Addr
|
||||
for _, rr := range records {
|
||||
if addr, ok := extractRecordAddr(rr); ok {
|
||||
addrs = append(addrs, addr)
|
||||
}
|
||||
}
|
||||
if len(addrs) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(d.ctx, d.warmupTimeout)
|
||||
defer cancel()
|
||||
activator.ActivatePeersByIP(ctx, addrs)
|
||||
}
|
||||
|
||||
// extractRecordAddr returns the IP address carried by an A or AAAA record.
|
||||
// ok is false for any other record type or a record with no address.
|
||||
func extractRecordAddr(rr dns.RR) (netip.Addr, bool) {
|
||||
switch r := rr.(type) {
|
||||
case *dns.A:
|
||||
if r.A == nil {
|
||||
return ""
|
||||
}
|
||||
return r.A.String()
|
||||
addr, ok := netip.AddrFromSlice(r.A)
|
||||
return addr.Unmap(), ok
|
||||
case *dns.AAAA:
|
||||
if r.AAAA == nil {
|
||||
return ""
|
||||
}
|
||||
return r.AAAA.String()
|
||||
addr, ok := netip.AddrFromSlice(r.AAAA)
|
||||
return addr.Unmap(), ok
|
||||
}
|
||||
return ""
|
||||
return netip.Addr{}, false
|
||||
}
|
||||
|
||||
// Update replaces all zones and their records
|
||||
|
||||
@@ -37,8 +37,8 @@ type mockPeerConnectivity struct {
|
||||
byIP map[string]struct{ known, connected bool }
|
||||
}
|
||||
|
||||
func (m mockPeerConnectivity) IsConnectedByIP(ip string) (known, connected bool) {
|
||||
v, ok := m.byIP[ip]
|
||||
func (m mockPeerConnectivity) IsConnectedByIP(ip netip.Addr) (known, connected bool) {
|
||||
v, ok := m.byIP[ip.String()]
|
||||
if !ok {
|
||||
return false, false
|
||||
}
|
||||
|
||||
204
client/internal/dns/local/warmup_test.go
Normal file
204
client/internal/dns/local/warmup_test.go
Normal file
@@ -0,0 +1,204 @@
|
||||
package local
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"net/netip"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/dns/test"
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
)
|
||||
|
||||
// recordingActivator records the addresses it was asked to warm and returns
|
||||
// immediately, so ServeDNS is not blocked by the test.
|
||||
type recordingActivator struct {
|
||||
mu sync.Mutex
|
||||
called bool
|
||||
addrs []netip.Addr
|
||||
}
|
||||
|
||||
func (r *recordingActivator) ActivatePeersByIP(_ context.Context, addrs []netip.Addr) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.called = true
|
||||
r.addrs = append(r.addrs, addrs...)
|
||||
}
|
||||
|
||||
func serveA(t *testing.T, resolver *Resolver, name string) *dns.Msg {
|
||||
t.Helper()
|
||||
var resp *dns.Msg
|
||||
w := &test.MockResponseWriter{WriteMsgFunc: func(m *dns.Msg) error { resp = m; return nil }}
|
||||
resolver.ServeDNS(w, new(dns.Msg).SetQuestion(name, dns.TypeA))
|
||||
return resp
|
||||
}
|
||||
|
||||
// serviceZone registers rec in a match-only (non-authoritative) zone, the shape
|
||||
// the synthesized private-service zones arrive in.
|
||||
func serviceZone(t *testing.T, resolver *Resolver, zone string, records ...nbdns.SimpleRecord) {
|
||||
t.Helper()
|
||||
resolver.Update([]nbdns.CustomZone{{
|
||||
Domain: zone,
|
||||
Records: records,
|
||||
NonAuthoritative: true,
|
||||
}})
|
||||
}
|
||||
|
||||
func TestLocalResolver_WarmsLazyPeerOnResolve(t *testing.T) {
|
||||
// Warm-up fires only for multi-record answers (the HA / round-robin shape of
|
||||
// the synthesized private-service zones), so register two peer targets.
|
||||
const name = "svc.proxy.netbird.cloud."
|
||||
recs := []nbdns.SimpleRecord{
|
||||
{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"},
|
||||
{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.8"},
|
||||
}
|
||||
resolver := NewResolver()
|
||||
serviceZone(t, resolver, "proxy.netbird.cloud", recs...)
|
||||
|
||||
act := &recordingActivator{}
|
||||
resolver.SetPeerActivator(act)
|
||||
|
||||
resp := serveA(t, resolver, name)
|
||||
require.NotNil(t, resp, "resolver must answer")
|
||||
require.NotEmpty(t, resp.Answer, "answer must carry the A records")
|
||||
|
||||
act.mu.Lock()
|
||||
defer act.mu.Unlock()
|
||||
assert.True(t, act.called, "activator must be invoked for a multi-record service-zone answer")
|
||||
assert.Contains(t, act.addrs, netip.MustParseAddr("100.64.0.7"), "activator must receive the first peer IP")
|
||||
assert.Contains(t, act.addrs, netip.MustParseAddr("100.64.0.8"), "activator must receive the second peer IP")
|
||||
}
|
||||
|
||||
func TestLocalResolver_NoWarmupForSingleRecord(t *testing.T) {
|
||||
// A single-record answer does not trigger warm-up; the resolver only warms
|
||||
// multi-record answers.
|
||||
rec := nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"}
|
||||
resolver := NewResolver()
|
||||
serviceZone(t, resolver, "proxy.netbird.cloud", rec)
|
||||
|
||||
act := &recordingActivator{}
|
||||
resolver.SetPeerActivator(act)
|
||||
|
||||
resp := serveA(t, resolver, rec.Name)
|
||||
require.NotNil(t, resp, "resolver must answer")
|
||||
require.NotEmpty(t, resp.Answer, "answer must carry the A record")
|
||||
|
||||
act.mu.Lock()
|
||||
defer act.mu.Unlock()
|
||||
assert.False(t, act.called, "activator must not be invoked for a single-record answer")
|
||||
}
|
||||
|
||||
func TestLocalResolver_NoActivatorNoWarmup(t *testing.T) {
|
||||
// With no activator wired the resolver behaves exactly as before.
|
||||
rec := nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"}
|
||||
resolver := NewResolver()
|
||||
serviceZone(t, resolver, "proxy.netbird.cloud", rec)
|
||||
|
||||
resp := serveA(t, resolver, rec.Name)
|
||||
require.NotNil(t, resp, "resolver must still answer without an activator")
|
||||
require.NotEmpty(t, resp.Answer, "answer must carry the A record")
|
||||
}
|
||||
|
||||
func TestLocalResolver_NoWarmupForMissingRecord(t *testing.T) {
|
||||
// A query that resolves to nothing must not invoke the activator (no IPs).
|
||||
resolver := NewResolver()
|
||||
serviceZone(t, resolver, "proxy.netbird.cloud",
|
||||
nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"})
|
||||
|
||||
act := &recordingActivator{}
|
||||
resolver.SetPeerActivator(act)
|
||||
|
||||
serveA(t, resolver, "absent.proxy.netbird.cloud.")
|
||||
|
||||
act.mu.Lock()
|
||||
defer act.mu.Unlock()
|
||||
assert.False(t, act.called, "activator must not be invoked when there is no answer")
|
||||
}
|
||||
|
||||
func TestLocalResolver_NoWarmupInAuthoritativeZone(t *testing.T) {
|
||||
// The account's peer zone is authoritative; resolving a peer's name there
|
||||
// must not wake its lazy connection — warm-up is scoped to match-only
|
||||
// (non-authoritative) zones such as the synthesized private-service zones.
|
||||
// Use a multi-record answer so the authoritative-zone scoping is the only
|
||||
// reason warm-up is skipped, not the single-record guard.
|
||||
const name = "peer.netbird.cloud."
|
||||
recs := []nbdns.SimpleRecord{
|
||||
{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.9"},
|
||||
{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.10"},
|
||||
}
|
||||
resolver := NewResolver()
|
||||
resolver.Update([]nbdns.CustomZone{{
|
||||
Domain: "netbird.cloud",
|
||||
Records: recs,
|
||||
}})
|
||||
|
||||
act := &recordingActivator{}
|
||||
resolver.SetPeerActivator(act)
|
||||
|
||||
resp := serveA(t, resolver, name)
|
||||
require.NotNil(t, resp, "resolver must answer")
|
||||
require.NotEmpty(t, resp.Answer, "answer must carry the A records")
|
||||
|
||||
act.mu.Lock()
|
||||
defer act.mu.Unlock()
|
||||
assert.False(t, act.called, "activator must not be invoked for authoritative-zone answers")
|
||||
}
|
||||
|
||||
func TestLazyWarmupTimeoutFromEnv(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
value string
|
||||
envSet bool
|
||||
want time.Duration
|
||||
}{
|
||||
{name: "unset uses default", want: defaultLazyWarmupTimeout},
|
||||
{name: "valid overrides", value: "5s", envSet: true, want: 5 * time.Second},
|
||||
{name: "invalid falls back", value: "not-a-duration", envSet: true, want: defaultLazyWarmupTimeout},
|
||||
{name: "negative falls back", value: "-1s", envSet: true, want: defaultLazyWarmupTimeout},
|
||||
{name: "zero falls back", value: "0s", envSet: true, want: defaultLazyWarmupTimeout},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if tt.envSet {
|
||||
t.Setenv(envLazyWarmupTimeout, tt.value)
|
||||
}
|
||||
assert.Equal(t, tt.want, lazyWarmupTimeoutFromEnv())
|
||||
assert.Equal(t, tt.want, NewResolver().warmupTimeout, "constructor must resolve the timeout once")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRecordAddr(t *testing.T) {
|
||||
t.Run("A record yields unmapped v4", func(t *testing.T) {
|
||||
// net.ParseIP returns the 16-byte v4-in-v6 form, the same shape
|
||||
// miekg/dns stores after parsing an A record; the extracted address
|
||||
// must compare equal to a plain v4 netip.Addr.
|
||||
addr, ok := extractRecordAddr(&dns.A{A: net.ParseIP("100.64.0.7")})
|
||||
require.True(t, ok)
|
||||
assert.True(t, addr.Is4())
|
||||
assert.Equal(t, netip.MustParseAddr("100.64.0.7"), addr)
|
||||
})
|
||||
|
||||
t.Run("AAAA record yields v6", func(t *testing.T) {
|
||||
addr, ok := extractRecordAddr(&dns.AAAA{AAAA: net.ParseIP("fd00::1")})
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, netip.MustParseAddr("fd00::1"), addr)
|
||||
})
|
||||
|
||||
t.Run("A record without address", func(t *testing.T) {
|
||||
_, ok := extractRecordAddr(&dns.A{})
|
||||
assert.False(t, ok)
|
||||
})
|
||||
|
||||
t.Run("non-address record", func(t *testing.T) {
|
||||
_, ok := extractRecordAddr(&dns.CNAME{Target: "target.netbird.cloud."})
|
||||
assert.False(t, ok)
|
||||
})
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"github.com/miekg/dns"
|
||||
|
||||
dnsconfig "github.com/netbirdio/netbird/client/internal/dns/config"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/local"
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
@@ -92,6 +93,11 @@ func (m *MockServer) SetFirewall(Firewall) {
|
||||
// Mock implementation - no-op
|
||||
}
|
||||
|
||||
// SetPeerActivator mock implementation of SetPeerActivator from Server interface
|
||||
func (m *MockServer) SetPeerActivator(local.PeerActivator) {
|
||||
// Mock implementation - no-op
|
||||
}
|
||||
|
||||
// BeginBatch mock implementation of BeginBatch from Server interface
|
||||
func (m *MockServer) BeginBatch() {
|
||||
// Mock implementation - no-op
|
||||
|
||||
@@ -19,6 +19,34 @@ import (
|
||||
// uses when a resolved host has no addresses of the requested family.
|
||||
const errNoSuitableAddress = "no suitable address found"
|
||||
|
||||
// Extended DNS Error info codes NetBird emits so a client can see why an answer
|
||||
// looks the way it does without reading this peer's logs. They live in the RFC
|
||||
// 8914 Private Use range (49152-65535) and are registered here, in one place,
|
||||
// because nothing else guarantees two NetBird components pick distinct codes.
|
||||
const (
|
||||
// EDENetbirdUpstreamTimeout: a DNS forwarder's upstream did not answer.
|
||||
EDENetbirdUpstreamTimeout uint16 = 49152
|
||||
// EDENetbirdUpstreamFailure: a DNS forwarder's upstream failed.
|
||||
EDENetbirdUpstreamFailure uint16 = 49153
|
||||
// EDENetbirdSoftenedNegative: the empty answer is ours, not the answering
|
||||
// resolver's. A handler that owns the name deferred this query type, and the
|
||||
// negative verdict that came back was downgraded so it cannot poison the
|
||||
// name.
|
||||
EDENetbirdSoftenedNegative uint16 = 49154
|
||||
)
|
||||
|
||||
// AttachEDE adds an Extended DNS Error (RFC 8914) option to a message, creating
|
||||
// the OPT pseudo-record if it has none. Callers must only use it toward a client
|
||||
// that advertised EDNS0: per RFC 6891 an OPT must not appear otherwise.
|
||||
func AttachEDE(msg *dns.Msg, code uint16, text string) {
|
||||
opt := msg.IsEdns0()
|
||||
if opt == nil {
|
||||
msg.SetEdns0(dns.DefaultMsgSize, false)
|
||||
opt = msg.IsEdns0()
|
||||
}
|
||||
opt.Option = append(opt.Option, &dns.EDNS0_EDE{InfoCode: code, ExtraText: text})
|
||||
}
|
||||
|
||||
// GenerateRequestID creates a random 8-character hex string for request tracing.
|
||||
func GenerateRequestID() string {
|
||||
bytes := make([]byte, 4)
|
||||
@@ -78,10 +106,38 @@ type resolver interface {
|
||||
LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error)
|
||||
}
|
||||
|
||||
// MetaKey names an annotation a handler attaches to a request to explain the
|
||||
// response the chain ends up writing. The set is closed and rendered onto a
|
||||
// single log line, so an unrecognized key is a typo inventing a field rather
|
||||
// than a new fact.
|
||||
type MetaKey string
|
||||
|
||||
const (
|
||||
// MetaKeyProtocol: the transport a query arrived on.
|
||||
MetaKeyProtocol MetaKey = "protocol"
|
||||
// MetaKeyUpstream: the upstream that answered.
|
||||
MetaKeyUpstream MetaKey = "upstream"
|
||||
// MetaKeyUpstreamProtocol: the transport used toward that upstream.
|
||||
MetaKeyUpstreamProtocol MetaKey = "upstream_protocol"
|
||||
// MetaKeyPeer: the routing peer whose forwarder answered.
|
||||
MetaKeyPeer MetaKey = "peer"
|
||||
// MetaKeyEDE: an Extended DNS Error carried by the answer.
|
||||
MetaKeyEDE MetaKey = "ede"
|
||||
// MetaKeyTruncated: the answer did not fit and was truncated.
|
||||
MetaKeyTruncated MetaKey = "truncated"
|
||||
// MetaKeySoftened: a negative verdict was rewritten so it cannot poison a
|
||||
// name served locally.
|
||||
MetaKeySoftened MetaKey = "softened"
|
||||
// MetaKeyDeferredBy: the handler that owned the name and stepped aside.
|
||||
MetaKeyDeferredBy MetaKey = "deferred_by"
|
||||
// MetaKeyDeferredReason: why it stepped aside.
|
||||
MetaKeyDeferredReason MetaKey = "deferred_reason"
|
||||
)
|
||||
|
||||
// chainedWriter is implemented by ResponseWriters that carry request metadata
|
||||
type chainedWriter interface {
|
||||
RequestID() string
|
||||
SetMeta(key, value string)
|
||||
SetMeta(key MetaKey, value string)
|
||||
}
|
||||
|
||||
// GetRequestID extracts a request ID from the ResponseWriter if available,
|
||||
@@ -96,12 +152,30 @@ func GetRequestID(w dns.ResponseWriter) string {
|
||||
}
|
||||
|
||||
// SetMeta sets metadata on the ResponseWriter if it supports it.
|
||||
func SetMeta(w dns.ResponseWriter, key, value string) {
|
||||
func SetMeta(w dns.ResponseWriter, key MetaKey, value string) {
|
||||
if cw, ok := w.(chainedWriter); ok {
|
||||
cw.SetMeta(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
// softNegativeRequester is implemented by chain writers that can soften the
|
||||
// negative verdict of the handlers a deferring handler falls through to.
|
||||
type softNegativeRequester interface {
|
||||
RequestSoftNegative()
|
||||
}
|
||||
|
||||
// RequestSoftNegative asks the handler chain to downgrade an NXDOMAIN from the
|
||||
// handlers that run after the caller defers. A handler that owns a name but
|
||||
// cannot answer one query type for it needs this: the resolvers it falls
|
||||
// through to cannot prove the name absent, and an NXDOMAIN from them is cached
|
||||
// for the name and every type under it (RFC 2308, RFC 8020), taking the
|
||||
// addresses the handler does serve down with it.
|
||||
func RequestSoftNegative(w dns.ResponseWriter) {
|
||||
if sn, ok := w.(softNegativeRequester); ok {
|
||||
sn.RequestSoftNegative()
|
||||
}
|
||||
}
|
||||
|
||||
// LookupResult contains the result of an external DNS lookup
|
||||
type LookupResult struct {
|
||||
IPs []netip.Addr
|
||||
@@ -199,11 +273,27 @@ type RecordResolver interface {
|
||||
LookupAddr(ctx context.Context, addr string) ([]string, error)
|
||||
}
|
||||
|
||||
// SupportedRecordQtype reports whether LookupRecords can resolve qtype. The
|
||||
// set is bounded by the net.Resolver API, which exposes no way to query an
|
||||
// arbitrary record type, and going around it with a raw exchange would mean
|
||||
// picking nameservers ourselves instead of resolving the way the host does.
|
||||
//
|
||||
// Both ends read this: a DNS forwarder answers these types for the domains it
|
||||
// routes, and a client uses it to tell which types are worth forwarding to a
|
||||
// peer at all.
|
||||
func SupportedRecordQtype(qtype uint16) bool {
|
||||
switch qtype {
|
||||
case dns.TypeMX, dns.TypeTXT, dns.TypeNS, dns.TypeSRV, dns.TypeCNAME, dns.TypePTR:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// LookupRecords resolves a non-address DNS record type through the host
|
||||
// resolver and returns the resource records and the DNS rcode. Types the host
|
||||
// resolver cannot answer (anything not covered by the net.Resolver Lookup*
|
||||
// methods) yield NODATA so that a routed name is never poisoned with NXDOMAIN
|
||||
// for an unsupported type.
|
||||
// resolver and returns the resource records and the DNS rcode. Types outside
|
||||
// SupportedRecordQtype yield NODATA so that a routed name is never poisoned
|
||||
// with NXDOMAIN for a type we cannot look up.
|
||||
func LookupRecords(ctx context.Context, r RecordResolver, name string, qtype uint16, ttl uint32) ([]dns.RR, int) {
|
||||
fqdn := dns.Fqdn(name)
|
||||
|
||||
|
||||
@@ -82,6 +82,7 @@ type Server interface {
|
||||
PopulateManagementDomain(mgmtURL *url.URL) error
|
||||
SetRouteSources(selected, active func() route.HAMap)
|
||||
SetFirewall(Firewall)
|
||||
SetPeerActivator(local.PeerActivator)
|
||||
}
|
||||
|
||||
type nsGroupsByDomain struct {
|
||||
@@ -491,6 +492,13 @@ func (s *DefaultServer) SetFirewall(fw Firewall) {
|
||||
}
|
||||
}
|
||||
|
||||
// SetPeerActivator wires the DNS-time lazy-connection warm-up on the local
|
||||
// resolver. Injected after the connection manager exists (it does not at
|
||||
// DNS-server construction time). Pass nil to disable.
|
||||
func (s *DefaultServer) SetPeerActivator(a local.PeerActivator) {
|
||||
s.localResolver.SetPeerActivator(a)
|
||||
}
|
||||
|
||||
// Stop stops the server
|
||||
func (s *DefaultServer) Stop() {
|
||||
s.ctxCancel()
|
||||
@@ -1435,11 +1443,11 @@ type localPeerConnectivity struct {
|
||||
|
||||
// IsConnectedByIP looks the IP up in the peerstore and surfaces both
|
||||
// the known and connected bits. Used by Resolver.filterDisconnectedPeerAnswers.
|
||||
func (l localPeerConnectivity) IsConnectedByIP(ip string) (known, connected bool) {
|
||||
func (l localPeerConnectivity) IsConnectedByIP(ip netip.Addr) (known, connected bool) {
|
||||
if l.status == nil {
|
||||
return false, false
|
||||
}
|
||||
state, ok := l.status.PeerStateByIP(ip)
|
||||
state, ok := l.status.PeerStateByIP(ip.String())
|
||||
if !ok {
|
||||
return false, false
|
||||
}
|
||||
|
||||
@@ -295,7 +295,7 @@ func (u *upstreamResolverBase) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
if addr := w.RemoteAddr(); addr != nil {
|
||||
network := addr.Network()
|
||||
ctx = contextWithDNSProtocol(ctx, network)
|
||||
resutil.SetMeta(w, "protocol", network)
|
||||
resutil.SetMeta(w, resutil.MetaKeyProtocol, network)
|
||||
}
|
||||
|
||||
ok, failures := u.tryUpstreamServers(ctx, w, r, logger)
|
||||
@@ -331,7 +331,7 @@ func (u *upstreamResolverBase) tryOnlyRace(ctx context.Context, w dns.ResponseWr
|
||||
return false, res.failures
|
||||
}
|
||||
if res.ede != "" {
|
||||
resutil.SetMeta(w, "ede", res.ede)
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, res.ede)
|
||||
}
|
||||
u.writeSuccessResponse(w, res.msg, res.upstream, r.Question[0].Name, res.protocol, logger)
|
||||
return true, res.failures
|
||||
@@ -361,7 +361,7 @@ func (u *upstreamResolverBase) raceAll(ctx context.Context, w dns.ResponseWriter
|
||||
failures = append(failures, res.failures...)
|
||||
if res.msg != nil {
|
||||
if res.ede != "" {
|
||||
resutil.SetMeta(w, "ede", res.ede)
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, res.ede)
|
||||
}
|
||||
u.writeSuccessResponse(w, res.msg, res.upstream, r.Question[0].Name, res.protocol, logger)
|
||||
return true, failures
|
||||
@@ -550,9 +550,9 @@ func (u *upstreamResolverBase) debugUpstreamTimeout(upstream netip.AddrPort) str
|
||||
}
|
||||
|
||||
func (u *upstreamResolverBase) writeSuccessResponse(w dns.ResponseWriter, rm *dns.Msg, upstream netip.AddrPort, domain string, proto string, logger *log.Entry) {
|
||||
resutil.SetMeta(w, "upstream", upstream.String())
|
||||
resutil.SetMeta(w, resutil.MetaKeyUpstream, upstream.String())
|
||||
if proto != "" {
|
||||
resutil.SetMeta(w, "upstream_protocol", proto)
|
||||
resutil.SetMeta(w, resutil.MetaKeyUpstreamProtocol, proto)
|
||||
}
|
||||
|
||||
// Clear Zero bit from external responses to prevent upstream servers from
|
||||
|
||||
76
client/internal/dns_peer_activator.go
Normal file
76
client/internal/dns_peer_activator.go
Normal file
@@ -0,0 +1,76 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/peerstore"
|
||||
)
|
||||
|
||||
const dnsActivationPollInterval = 50 * time.Millisecond
|
||||
|
||||
// dnsPeerActivator wakes lazy-connection peers from the DNS resolution path. It
|
||||
// implements dns/local.PeerActivator. DNS queries run on their own goroutines,
|
||||
// so it only touches state that is safe for concurrent use — ConnMgr.ActivatePeer,
|
||||
// peerstore.Store and peer.Status — and never takes the engine's syncMsgMux,
|
||||
// keeping DNS resolution from contending with network-map processing.
|
||||
type dnsPeerActivator struct {
|
||||
connMgr *ConnMgr
|
||||
peerStore *peerstore.Store
|
||||
status *peer.Status
|
||||
// ctx is the engine's long-lived context. The connection dial is tied to it
|
||||
// (not the per-query DNS wait budget) so a handshake that outlasts the wait
|
||||
// still completes in the background rather than being cancelled at the deadline.
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
// ActivatePeersByIP triggers wake-up for the peer(s) owning addrs and waits
|
||||
// until one is connected or ctx (the per-query DNS wait budget) expires.
|
||||
// Activation itself is tied to the engine's long-lived context so the dial
|
||||
// survives a wait that times out. Unknown or already-connected addresses are
|
||||
// skipped, so the steady-state (warm) path adds no latency.
|
||||
func (a *dnsPeerActivator) ActivatePeersByIP(ctx context.Context, addrs []netip.Addr) {
|
||||
if a == nil || a.connMgr == nil {
|
||||
return
|
||||
}
|
||||
|
||||
var pending []string
|
||||
for _, addr := range addrs {
|
||||
ip := addr.String()
|
||||
st, ok := a.status.PeerStateByIP(ip)
|
||||
if !ok || st.ConnStatus == peer.StatusConnected {
|
||||
continue
|
||||
}
|
||||
conn, ok := a.peerStore.PeerConn(st.PubKey)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
a.connMgr.ActivatePeer(a.ctx, conn)
|
||||
pending = append(pending, ip)
|
||||
}
|
||||
|
||||
if len(pending) == 0 {
|
||||
return
|
||||
}
|
||||
a.waitConnected(ctx, pending)
|
||||
}
|
||||
|
||||
// waitConnected blocks until any of ips reports a connected peer or ctx expires.
|
||||
func (a *dnsPeerActivator) waitConnected(ctx context.Context, ips []string) {
|
||||
ticker := time.NewTicker(dnsActivationPollInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
for _, ip := range ips {
|
||||
if st, ok := a.status.PeerStateByIP(ip); ok && st.ConnStatus == peer.StatusConnected {
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
129
client/internal/dns_peer_activator_test.go
Normal file
129
client/internal/dns_peer_activator_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/peerstore"
|
||||
)
|
||||
|
||||
func newTestPeerConn(t *testing.T, key string) *peer.Conn {
|
||||
t.Helper()
|
||||
conn, err := peer.NewConn(peer.ConnConfig{
|
||||
Key: key,
|
||||
LocalKey: "local",
|
||||
WgConfig: peer.WgConfig{
|
||||
AllowedIps: []netip.Prefix{netip.MustParsePrefix("100.64.0.1/32")},
|
||||
},
|
||||
}, peer.ServiceDependencies{})
|
||||
require.NoError(t, err)
|
||||
return conn
|
||||
}
|
||||
|
||||
func newTestDNSPeerActivator(t *testing.T) (*dnsPeerActivator, *peer.Status, *peerstore.Store) {
|
||||
t.Helper()
|
||||
status := peer.NewRecorder("https://mgm")
|
||||
store := peerstore.NewConnStore()
|
||||
// ConnMgr without Start: the lazy manager is nil, so ActivatePeer is a
|
||||
// no-op — these tests exercise the activator's skip/wait logic.
|
||||
connMgr := NewConnMgr(&EngineConfig{}, status, store, nil)
|
||||
return &dnsPeerActivator{
|
||||
connMgr: connMgr,
|
||||
peerStore: store,
|
||||
status: status,
|
||||
ctx: context.Background(),
|
||||
}, status, store
|
||||
}
|
||||
|
||||
func TestDNSPeerActivator_NilSafe(t *testing.T) {
|
||||
var a *dnsPeerActivator
|
||||
a.ActivatePeersByIP(context.Background(), []netip.Addr{netip.MustParseAddr("100.64.0.1")})
|
||||
}
|
||||
|
||||
// TestDNSPeerActivator_SkipsUnknownAndConnectedPeers verifies the steady-state
|
||||
// (warm) path adds no latency: already-connected and unknown addresses never
|
||||
// enter the wait loop.
|
||||
func TestDNSPeerActivator_SkipsUnknownAndConnectedPeers(t *testing.T) {
|
||||
a, status, store := newTestDNSPeerActivator(t)
|
||||
|
||||
require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", "fd00::1"))
|
||||
require.NoError(t, status.UpdatePeerState(peer.State{PubKey: "peerA", ConnStatus: peer.StatusConnected}))
|
||||
store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
start := time.Now()
|
||||
a.ActivatePeersByIP(ctx, []netip.Addr{
|
||||
netip.MustParseAddr("100.64.0.1"), // known, connected -> skipped
|
||||
netip.MustParseAddr("fd00::1"), // known via IPv6, connected -> skipped
|
||||
netip.MustParseAddr("100.64.0.99"), // unknown -> skipped
|
||||
})
|
||||
require.Less(t, time.Since(start), time.Second, "no pending peer must mean no wait")
|
||||
}
|
||||
|
||||
// TestDNSPeerActivator_WaitsForPendingPeerToConnect verifies the wait loop
|
||||
// returns as soon as a pending peer reports connected, well before the
|
||||
// per-query budget expires.
|
||||
func TestDNSPeerActivator_WaitsForPendingPeerToConnect(t *testing.T) {
|
||||
a, status, store := newTestDNSPeerActivator(t)
|
||||
|
||||
require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
|
||||
store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
|
||||
|
||||
go func() {
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
_ = status.UpdatePeerState(peer.State{PubKey: "peerA", ConnStatus: peer.StatusConnected})
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
start := time.Now()
|
||||
a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
require.GreaterOrEqual(t, elapsed, 100*time.Millisecond, "must wait for the pending peer")
|
||||
require.Less(t, elapsed, 5*time.Second, "must return on connect, not at the deadline")
|
||||
}
|
||||
|
||||
// TestDNSPeerActivator_ReturnsAtBudgetWhenPeerStaysIdle verifies a peer that
|
||||
// never connects releases the DNS response at the per-query budget instead of
|
||||
// blocking it indefinitely.
|
||||
func TestDNSPeerActivator_ReturnsAtBudgetWhenPeerStaysIdle(t *testing.T) {
|
||||
a, status, store := newTestDNSPeerActivator(t)
|
||||
|
||||
require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
|
||||
store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
start := time.Now()
|
||||
a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
require.GreaterOrEqual(t, elapsed, 250*time.Millisecond, "must wait out the budget for a pending peer")
|
||||
require.Less(t, elapsed, 5*time.Second, "must not block past the budget")
|
||||
}
|
||||
|
||||
// TestDNSPeerActivator_NoWaitWithoutPeerConn verifies a known-but-idle peer
|
||||
// with no connection object in the store is not waited on: there is nothing to
|
||||
// activate, so waiting could only ever time out.
|
||||
func TestDNSPeerActivator_NoWaitWithoutPeerConn(t *testing.T) {
|
||||
a, status, _ := newTestDNSPeerActivator(t)
|
||||
|
||||
require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
start := time.Now()
|
||||
a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
|
||||
require.Less(t, time.Since(start), time.Second, "peer without a conn must not be waited on")
|
||||
}
|
||||
@@ -26,15 +26,6 @@ import (
|
||||
const errResolveFailed = "failed to resolve query for domain=%s: %v"
|
||||
const upstreamTimeout = 15 * time.Second
|
||||
|
||||
// EDE info codes the forwarder emits on upstream failures so the querying
|
||||
// client can see the reason without inspecting this peer's logs. They live in
|
||||
// the RFC 8914 Private Use range (49152-65535); the Go resolver never exposes a
|
||||
// real upstream EDE here, so these cannot collide with a genuine code.
|
||||
const (
|
||||
edeNetbirdUpstreamTimeout uint16 = 49152
|
||||
edeNetbirdUpstreamFailure uint16 = 49153
|
||||
)
|
||||
|
||||
type resolver interface {
|
||||
LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error)
|
||||
LookupMX(ctx context.Context, name string) ([]*net.MX, error)
|
||||
@@ -216,6 +207,9 @@ func (f *DNSForwarder) handleDNSQuery(logger *log.Entry, w dns.ResponseWriter, q
|
||||
qname, dns.TypeToString[question.Qtype], dns.ClassToString[question.Qclass])
|
||||
|
||||
resp := query.SetReply(query)
|
||||
// Every answer here comes from a recursive lookup on this peer. SetReply
|
||||
// leaves RA unset, which reads to a client as a server that cannot recurse.
|
||||
resp.RecursionAvailable = true
|
||||
|
||||
mostSpecificResId, matchingEntries := f.getMatchingEntries(strings.TrimSuffix(qname, "."))
|
||||
if mostSpecificResId == "" {
|
||||
@@ -229,20 +223,22 @@ func (f *DNSForwarder) handleDNSQuery(logger *log.Entry, w dns.ResponseWriter, q
|
||||
|
||||
reqHasEdns := query.IsEdns0() != nil
|
||||
|
||||
switch question.Qtype {
|
||||
case dns.TypeA, dns.TypeAAAA:
|
||||
switch {
|
||||
case question.Qtype == dns.TypeA || question.Qtype == dns.TypeAAAA:
|
||||
f.handleAddressQuery(ctx, logger, w, resp, mostSpecificResId, matchingEntries, reqHasEdns, startTime)
|
||||
case dns.TypeMX, dns.TypeTXT, dns.TypeNS, dns.TypeSRV, dns.TypeCNAME, dns.TypePTR:
|
||||
case resutil.SupportedRecordQtype(question.Qtype):
|
||||
f.handleRecordQuery(ctx, logger, w, resp, startTime)
|
||||
default:
|
||||
// The domain is routed here, so any other type is answered NODATA
|
||||
// (NOERROR, empty answer) rather than falling back to a resolver that
|
||||
// would poison the name with NXDOMAIN. The Extended DNS Error lets a
|
||||
// client tell this capability-driven NODATA apart from an
|
||||
// authoritative one. The OPT pseudo-record must not appear unless the
|
||||
// query advertised EDNS0.
|
||||
// authoritative one; a current client knows the type is unsupported and
|
||||
// resolves it through its own handler chain instead of asking at all.
|
||||
// The OPT pseudo-record must not appear unless the query advertised
|
||||
// EDNS0.
|
||||
if reqHasEdns {
|
||||
attachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
|
||||
resutil.AttachEDE(resp, dns.ExtendedErrorCodeNotSupported, "netbird forwarder: unsupported query type")
|
||||
}
|
||||
f.writeResponse(logger, w, resp, qname, startTime)
|
||||
}
|
||||
@@ -441,7 +437,7 @@ func (f *DNSForwarder) handleDNSError(
|
||||
}
|
||||
|
||||
if reqHasEdns {
|
||||
attachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
|
||||
resutil.AttachEDE(resp, edeCodeFor(dnsErr), edeText(dnsErr))
|
||||
}
|
||||
|
||||
f.writeResponse(logger, w, resp, domain, startTime)
|
||||
@@ -488,9 +484,9 @@ func (f *DNSForwarder) getMatchingEntries(domain string) (route.ResID, []*Forwar
|
||||
// edeCodeFor maps an upstream lookup error to the NetBird EDE info code.
|
||||
func edeCodeFor(dnsErr *net.DNSError) uint16 {
|
||||
if dnsErr != nil && dnsErr.IsTimeout {
|
||||
return edeNetbirdUpstreamTimeout
|
||||
return resutil.EDENetbirdUpstreamTimeout
|
||||
}
|
||||
return edeNetbirdUpstreamFailure
|
||||
return resutil.EDENetbirdUpstreamFailure
|
||||
}
|
||||
|
||||
// edeText builds the EDE extra-text describing the class of upstream failure.
|
||||
@@ -503,14 +499,3 @@ func edeText(dnsErr *net.DNSError) string {
|
||||
}
|
||||
return "netbird forwarder: upstream failure"
|
||||
}
|
||||
|
||||
// attachEDE adds an Extended DNS Error (RFC 8914) option to the response,
|
||||
// creating the OPT pseudo-record if the response does not already carry one.
|
||||
func attachEDE(resp *dns.Msg, code uint16, text string) {
|
||||
opt := resp.IsEdns0()
|
||||
if opt == nil {
|
||||
resp.SetEdns0(dns.DefaultMsgSize, false)
|
||||
opt = resp.IsEdns0()
|
||||
}
|
||||
opt.Option = append(opt.Option, &dns.EDNS0_EDE{InfoCode: code, ExtraText: text})
|
||||
}
|
||||
|
||||
@@ -649,6 +649,39 @@ func TestDNSForwarder_ResponseCodes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestDNSForwarder_RecursionAvailable covers the RA bit. Every answer this
|
||||
// forwarder produces comes from a recursive lookup on the routing peer, but
|
||||
// dns.Msg.SetReply leaves RA unset, so clients report "recursion not available"
|
||||
// and some stub resolvers treat the server as unable to serve the query.
|
||||
func TestDNSForwarder_RecursionAvailable(t *testing.T) {
|
||||
t.Run("record answer", func(t *testing.T) {
|
||||
mockResolver := &MockResolver{}
|
||||
forwarder := newRecordTestForwarder(t, mockResolver, "example.com")
|
||||
|
||||
mockResolver.On("LookupMX", mock.Anything, "example.com.").
|
||||
Return([]*net.MX{{Host: "mail.example.com.", Pref: 10}}, nil).Once()
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "example.com", dns.TypeMX)
|
||||
assert.True(t, resp.RecursionAvailable, "the forwarder resolves recursively")
|
||||
})
|
||||
|
||||
t.Run("unsupported type NODATA", func(t *testing.T) {
|
||||
forwarder := newRecordTestForwarder(t, &MockResolver{}, "example.com")
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "example.com", dns.TypeCAA)
|
||||
require.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
assert.True(t, resp.RecursionAvailable, "RA describes the server, not the query type")
|
||||
})
|
||||
|
||||
t.Run("unauthorized domain", func(t *testing.T) {
|
||||
forwarder := newRecordTestForwarder(t, &MockResolver{}, "example.com")
|
||||
|
||||
resp := runRecordQuery(t, forwarder, "other.com", dns.TypeMX)
|
||||
require.Equal(t, dns.RcodeRefused, resp.Rcode)
|
||||
assert.True(t, resp.RecursionAvailable, "RA describes the server, not the verdict")
|
||||
})
|
||||
}
|
||||
|
||||
func hasEDE(m *dns.Msg, code uint16) bool {
|
||||
opt := m.IsEdns0()
|
||||
if opt == nil {
|
||||
@@ -859,7 +892,7 @@ func TestDNSForwarder_UpstreamFailureEDE(t *testing.T) {
|
||||
lookupErr: &net.DNSError{Err: "i/o timeout", Server: "10.0.0.53:53", IsTimeout: true},
|
||||
reqEdns: true,
|
||||
wantEDE: true,
|
||||
wantCode: edeNetbirdUpstreamTimeout,
|
||||
wantCode: resutil.EDENetbirdUpstreamTimeout,
|
||||
wantTextHas: "netbird forwarder: upstream timeout",
|
||||
},
|
||||
{
|
||||
@@ -867,7 +900,7 @@ func TestDNSForwarder_UpstreamFailureEDE(t *testing.T) {
|
||||
lookupErr: &net.DNSError{Err: "server misbehaving", Server: "10.0.0.53:53"},
|
||||
reqEdns: true,
|
||||
wantEDE: true,
|
||||
wantCode: edeNetbirdUpstreamFailure,
|
||||
wantCode: resutil.EDENetbirdUpstreamFailure,
|
||||
wantTextHas: "netbird forwarder: upstream failure",
|
||||
},
|
||||
{
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -52,11 +52,14 @@ 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;
|
||||
}
|
||||
|
||||
|
||||
@@ -50,5 +50,11 @@ 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;
|
||||
}
|
||||
|
||||
@@ -64,7 +64,10 @@ import (
|
||||
"github.com/netbirdio/netbird/route"
|
||||
mgm "github.com/netbirdio/netbird/shared/management/client"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
|
||||
nbnetworkmap "github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
mgmProto "github.com/netbirdio/netbird/shared/management/proto"
|
||||
types "github.com/netbirdio/netbird/shared/management/types"
|
||||
"github.com/netbirdio/netbird/shared/netiputil"
|
||||
auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
|
||||
relayClient "github.com/netbirdio/netbird/shared/relay/client"
|
||||
@@ -147,6 +150,7 @@ type EngineConfig struct {
|
||||
BlockLANAccess bool
|
||||
BlockInbound bool
|
||||
DisableIPv6 bool
|
||||
SyncMessageVersion *int
|
||||
|
||||
// LazyConnection is the MDM-sourced lazy-connection override; StateUnset defers to
|
||||
// the env var and management feature flag.
|
||||
@@ -220,6 +224,13 @@ type Engine struct {
|
||||
// networkSerial is the latest CurrentSerial (state ID) of the network sent by the Management service
|
||||
networkSerial uint64
|
||||
|
||||
// latestComponents is the most-recent NetworkMapComponents decoded from
|
||||
// a NetworkMapEnvelope (capability=3 peers only). Held alongside the
|
||||
// NetworkMap that Calculate() produced from it so future incremental
|
||||
// updates have a base to apply changes against. nil for legacy-format
|
||||
// peers. Guarded by syncMsgMux.
|
||||
latestComponents *types.NetworkMapComponents
|
||||
|
||||
networkMonitor *networkmonitor.NetworkMonitor
|
||||
|
||||
sshServer sshServer
|
||||
@@ -652,8 +663,24 @@ 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
|
||||
// exists (it does not at DNS-server construction time). A DNS answer that
|
||||
// points at an idle peer then wakes it before the client's first request.
|
||||
e.dnsServer.SetPeerActivator(&dnsPeerActivator{
|
||||
connMgr: e.connMgr,
|
||||
peerStore: e.peerStore,
|
||||
status: e.statusRecorder,
|
||||
ctx: e.ctx,
|
||||
})
|
||||
|
||||
e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg)
|
||||
e.srWatcher.Start(peer.IsForceRelayed())
|
||||
|
||||
@@ -963,8 +990,12 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
|
||||
|
||||
e.ApplySessionDeadline(update.GetSessionExpiresAt())
|
||||
|
||||
if update.NetworkMap != nil && update.NetworkMap.PeerConfig != nil {
|
||||
e.handleAutoUpdateVersion(update.NetworkMap.PeerConfig.AutoUpdate)
|
||||
// Envelope sync responses carry PeerConfig at the top level; legacy
|
||||
// NetworkMap syncs carry it under NetworkMap.PeerConfig.
|
||||
if pc := update.GetPeerConfig(); pc != nil {
|
||||
e.handleAutoUpdateVersion(pc.GetAutoUpdate())
|
||||
} else if nm := update.GetNetworkMap(); nm != nil && nm.GetPeerConfig() != nil {
|
||||
e.handleAutoUpdateVersion(nm.GetPeerConfig().GetAutoUpdate())
|
||||
}
|
||||
|
||||
done := e.phase("netbird_config")
|
||||
@@ -974,12 +1005,47 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
|
||||
return err
|
||||
}
|
||||
|
||||
// Decode the network map from either the components envelope or the
|
||||
// legacy proto.NetworkMap before the posture-check gating below, so the
|
||||
// "is there a network map" decision covers both wire shapes.
|
||||
var (
|
||||
nm *mgmProto.NetworkMap
|
||||
components *types.NetworkMapComponents
|
||||
)
|
||||
if version := update.GetVersion(); version == int32(sharedgrpc.ComponentNetworkMap) {
|
||||
// Components-format peer: decode the envelope back to typed
|
||||
// components, run Calculate() locally, and convert to the wire
|
||||
// NetworkMap shape the rest of the engine consumes. Components are
|
||||
// retained so future incremental updates can apply deltas instead
|
||||
// of doing a full reconstruction.
|
||||
envelope := update.GetNetworkMapEnvelope()
|
||||
if envelope == nil {
|
||||
return fmt.Errorf("received a SyncReponse indicating use of components network map, but components are missing")
|
||||
}
|
||||
|
||||
localKey := e.config.WgPrivateKey.PublicKey().String()
|
||||
dnsName := ""
|
||||
if pc := update.GetPeerConfig(); pc != nil {
|
||||
// PeerConfig.Fqdn = "<dns_label>.<dns_domain>" — extract the
|
||||
// shared domain by stripping the peer's own label prefix. Falls
|
||||
// back to empty if the FQDN doesn't have the expected shape.
|
||||
dnsName = extractDNSDomainFromFQDN(pc.GetFqdn())
|
||||
}
|
||||
result, err := nbnetworkmap.EnvelopeToNetworkMap(e.ctx, envelope, localKey, dnsName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode network map envelope: %w", err)
|
||||
}
|
||||
nm = result.NetworkMap
|
||||
components = result.Components
|
||||
} else {
|
||||
nm = update.GetNetworkMap()
|
||||
}
|
||||
|
||||
// Posture checks are bound to the network map presence:
|
||||
// NetworkMap != nil, checks present -> apply the received checks
|
||||
// NetworkMap != nil, checks nil -> posture checks were removed, clear them
|
||||
// NetworkMap == nil -> config-only update (e.g. relay token rotation),
|
||||
// leave the previously applied checks untouched
|
||||
nm := update.GetNetworkMap()
|
||||
if nm == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -992,6 +1058,14 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
|
||||
}
|
||||
|
||||
done = e.phase("persist")
|
||||
// Only retain the components view when the server sent the envelope
|
||||
// path. A legacy proto.NetworkMap means components == nil; writing it
|
||||
// here would clobber a previously-cached snapshot, breaking the
|
||||
// incremental-delta base on a future envelope sync.
|
||||
if components != nil {
|
||||
e.latestComponents = components
|
||||
}
|
||||
|
||||
e.persistSyncResponse(update)
|
||||
done()
|
||||
|
||||
@@ -1005,6 +1079,19 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// extractDNSDomainFromFQDN returns the trailing dotted domain part of the
|
||||
// receiving peer's FQDN — the same value the management server fills as
|
||||
// dnsName when it builds the legacy NetworkMap. "peer42.netbird.cloud" →
|
||||
// "netbird.cloud". An empty string is returned for unrecognized formats.
|
||||
func extractDNSDomainFromFQDN(fqdn string) string {
|
||||
for i := 0; i < len(fqdn); i++ {
|
||||
if fqdn[i] == '.' && i+1 < len(fqdn) {
|
||||
return fqdn[i+1:]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// updateNetbirdConfig applies the management-provided NetBird configuration:
|
||||
// STUN/TURN and relay servers, flow logging and DNS settings. A nil config is a no-op,
|
||||
// which is the case for sync updates carrying only a network map.
|
||||
@@ -1164,6 +1251,7 @@ func (e *Engine) applyInfoFlags(info *system.Info) {
|
||||
e.config.BlockLANAccess,
|
||||
e.config.BlockInbound,
|
||||
e.config.DisableIPv6,
|
||||
e.config.SyncMessageVersion,
|
||||
e.config.EnableSSHRoot,
|
||||
e.config.EnableSSHSFTP,
|
||||
e.config.EnableSSHLocalPortForwarding,
|
||||
@@ -2032,6 +2120,7 @@ func (e *Engine) readInitialSettings() ([]*route.Route, *nbdns.Config, bool, err
|
||||
e.config.BlockLANAccess,
|
||||
e.config.BlockInbound,
|
||||
e.config.DisableIPv6,
|
||||
e.config.SyncMessageVersion,
|
||||
e.config.EnableSSHRoot,
|
||||
e.config.EnableSSHSFTP,
|
||||
e.config.EnableSSHLocalPortForwarding,
|
||||
|
||||
@@ -29,6 +29,11 @@ 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
|
||||
@@ -56,6 +61,9 @@ 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
|
||||
@@ -73,6 +81,7 @@ 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() {
|
||||
@@ -201,7 +210,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
|
||||
if err := m.armActivityListener(peerCfg); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
@@ -288,7 +297,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) {
|
||||
|
||||
m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
|
||||
if err := m.armActivityListener(*mp.peerCfg); err != nil {
|
||||
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
|
||||
return
|
||||
}
|
||||
@@ -465,6 +474,31 @@ 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()
|
||||
@@ -577,7 +611,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
|
||||
|
||||
mp.peerCfg.Log.Infof("start activity monitor")
|
||||
|
||||
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
|
||||
if err := m.armActivityListener(*mp.peerCfg); err != nil {
|
||||
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -96,6 +96,7 @@ type ConfigInput struct {
|
||||
BlockLANAccess *bool
|
||||
BlockInbound *bool
|
||||
DisableIPv6 *bool
|
||||
SyncMessageVersion *int
|
||||
|
||||
DisableNotifications *bool
|
||||
|
||||
@@ -137,6 +138,7 @@ type Config struct {
|
||||
BlockLANAccess bool
|
||||
BlockInbound bool
|
||||
DisableIPv6 bool
|
||||
SyncMessageVersion *int
|
||||
|
||||
DisableNotifications *bool
|
||||
|
||||
@@ -587,6 +589,12 @@ func (config *Config) apply(input ConfigInput) (updated bool, err error) {
|
||||
updated = true
|
||||
}
|
||||
|
||||
if input.SyncMessageVersion != nil && *input.SyncMessageVersion != *config.SyncMessageVersion {
|
||||
log.Infof("setting SyncMessageVersion to %v", *input.SyncMessageVersion)
|
||||
*config.SyncMessageVersion = *input.SyncMessageVersion
|
||||
updated = true
|
||||
}
|
||||
|
||||
if input.DisableNotifications != nil && (config.DisableNotifications == nil || *input.DisableNotifications != *config.DisableNotifications) {
|
||||
if *input.DisableNotifications {
|
||||
log.Infof("disabling notifications")
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
nbdns "github.com/netbirdio/netbird/client/internal/dns"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/client/internal/peerstore"
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager/common"
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager/fakeip"
|
||||
iface "github.com/netbirdio/netbird/client/internal/routemanager/iface"
|
||||
@@ -40,6 +39,61 @@ type internalDNATer interface {
|
||||
AddInternalDNATMapping(netip.Addr, netip.Addr) error
|
||||
}
|
||||
|
||||
// peerAllowedIPs reports the tunnel address a peer is reachable on.
|
||||
type peerAllowedIPs interface {
|
||||
AllowedIP(pubKey string) (netip.Addr, bool)
|
||||
}
|
||||
|
||||
// disposition says where a query for a given record type has to be answered.
|
||||
type disposition int
|
||||
|
||||
const (
|
||||
// dispositionPeer: the routing peer owns the answer. Address records are
|
||||
// what a DNS route exists for, and their answer programs the routes,
|
||||
// allowed IPs and firewall sets, so they are never resolved anywhere else.
|
||||
dispositionPeer disposition = iota
|
||||
// dispositionPeerFirst: the peer's forwarder resolves the type, and the
|
||||
// records may only exist inside the routed network, so it has to be asked.
|
||||
// A peer running a client that predates that support cannot answer, which
|
||||
// only its reply reveals.
|
||||
dispositionPeerFirst
|
||||
// dispositionChain: no forwarder resolves the type, so asking would cost a
|
||||
// tunnel round trip for a reply that says nothing. Public records for the
|
||||
// name are still better than none, so the query goes to the rest of the
|
||||
// chain.
|
||||
dispositionChain
|
||||
)
|
||||
|
||||
func dispositionFor(qtype uint16) disposition {
|
||||
switch {
|
||||
case qtype == dns.TypeA || qtype == dns.TypeAAAA:
|
||||
return dispositionPeer
|
||||
case resutil.SupportedRecordQtype(qtype):
|
||||
return dispositionPeerFirst
|
||||
default:
|
||||
return dispositionChain
|
||||
}
|
||||
}
|
||||
|
||||
// peerCannotAnswer reports whether a forwarder reply means the peer is unable to
|
||||
// resolve this query type, as opposed to having resolved it and found nothing.
|
||||
// A forwarder older than the one that learned to resolve non-address types
|
||||
// answers NOTIMP to all of them, before it even looks at the domain; FORMERR
|
||||
// covers one that cannot parse the EDNS0 we add to the query.
|
||||
//
|
||||
// REFUSED is deliberately not here. It says the peer does not hold the name,
|
||||
// which happens while the client and the peer disagree about the route set, and
|
||||
// falling through then would hand the name of an internal-only domain to a
|
||||
// public resolver. No forwarder version reports a missing record type that way.
|
||||
func peerCannotAnswer(rcode int) bool {
|
||||
switch rcode {
|
||||
case dns.RcodeNotImplemented, dns.RcodeFormatError:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
type DnsInterceptor struct {
|
||||
mu sync.RWMutex
|
||||
route *route.Route
|
||||
@@ -50,7 +104,7 @@ type DnsInterceptor struct {
|
||||
currentPeerKey string
|
||||
interceptedDomains domainMap
|
||||
wgInterface iface.WGIface
|
||||
peerStore *peerstore.Store
|
||||
peerStore peerAllowedIPs
|
||||
firewall firewall.Manager
|
||||
fakeIPManager *fakeip.Manager
|
||||
forwarderPort *atomic.Uint32
|
||||
@@ -226,11 +280,14 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
return
|
||||
}
|
||||
|
||||
// All query types for an intercepted domain are forwarded to the peer's
|
||||
// DNS forwarder, which owns the name. Falling through to the system
|
||||
// resolver would let it answer NXDOMAIN for a name it isn't authoritative
|
||||
// for, poisoning the whole name (including the A/AAAA records the route
|
||||
// does serve). The forwarder answers NODATA for types it cannot resolve.
|
||||
qtype := r.Question[0].Qtype
|
||||
dispose := dispositionFor(qtype)
|
||||
|
||||
if dispose == dispositionChain {
|
||||
d.deferToChain(w, r, logger, "unsupported-qtype")
|
||||
return
|
||||
}
|
||||
|
||||
d.mu.RLock()
|
||||
peerKey := d.currentPeerKey
|
||||
d.mu.RUnlock()
|
||||
@@ -246,35 +303,32 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
return
|
||||
}
|
||||
|
||||
if r.Extra == nil {
|
||||
r.MsgHdr.AuthenticatedData = true
|
||||
}
|
||||
|
||||
// Advertise EDNS0 to the forwarder so it may return an Extended DNS Error
|
||||
// describing why a lookup failed. The OPT is stripped from the reply when
|
||||
// the original client did not request EDNS0.
|
||||
hadEdns := r.IsEdns0() != nil
|
||||
if !hadEdns {
|
||||
r.SetEdns0(dns.DefaultMsgSize, false)
|
||||
}
|
||||
query, hadEdns := peerQuery(r)
|
||||
|
||||
upstream := net.JoinHostPort(upstreamIP.String(), strconv.FormatUint(uint64(d.forwarderPort.Load()), 10))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), dnsTimeout)
|
||||
defer cancel()
|
||||
|
||||
reply := d.queryUpstreamDNS(ctx, w, r, upstream, upstreamIP, peerKey, logger)
|
||||
reply := d.queryUpstreamDNS(ctx, w, query, upstream, upstreamIP, peerKey, logger)
|
||||
if reply == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// The peer owns the name but its forwarder cannot resolve this type. No
|
||||
// capability is announced anywhere, so the round trip above is the probe.
|
||||
if dispose == dispositionPeerFirst && peerCannotAnswer(reply.Rcode) {
|
||||
d.deferToChain(w, r, logger, "peer-rcode-"+dns.RcodeToString[reply.Rcode])
|
||||
return
|
||||
}
|
||||
|
||||
if ede, ok := resutil.ExtractEDE(reply); ok {
|
||||
resutil.SetMeta(w, "ede", fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
|
||||
resutil.SetMeta(w, resutil.MetaKeyEDE, fmt.Sprintf("%d %s", ede.InfoCode, ede.ExtraText))
|
||||
}
|
||||
if !hadEdns {
|
||||
resutil.StripOPT(reply)
|
||||
}
|
||||
|
||||
resutil.SetMeta(w, "peer", peerKey)
|
||||
resutil.SetMeta(w, resutil.MetaKeyPeer, peerKey)
|
||||
|
||||
reply.Id = r.Id
|
||||
if err := d.writeMsg(w, reply, logger); err != nil {
|
||||
@@ -282,6 +336,30 @@ func (d *DnsInterceptor) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
|
||||
}
|
||||
}
|
||||
|
||||
// deferToChain hands the query to the next handler in the chain and asks the
|
||||
// chain to soften the negative verdict of whatever answers instead. The
|
||||
// resolvers below cannot prove a name inside the routed network absent, so their
|
||||
// NXDOMAIN must not reach the client: it would be cached for the name and every
|
||||
// type under it, taking the addresses this route does serve with it.
|
||||
func (d *DnsInterceptor) deferToChain(w dns.ResponseWriter, r *dns.Msg, logger *log.Entry, reason string) {
|
||||
logger.Tracef("continuing to next handler for domain=%s type=%s reason=%s",
|
||||
r.Question[0].Name, dns.TypeToString[r.Question[0].Qtype], reason)
|
||||
|
||||
resutil.RequestSoftNegative(w)
|
||||
// Carried to the chain's response log line: without it the answer looks like
|
||||
// it came from the fallthrough resolver on its own.
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredBy, "dns-route")
|
||||
resutil.SetMeta(w, resutil.MetaKeyDeferredReason, reason)
|
||||
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(r, dns.RcodeNameError)
|
||||
// Set Zero bit to signal handler chain to continue
|
||||
resp.MsgHdr.Zero = true
|
||||
if err := w.WriteMsg(resp); err != nil {
|
||||
logger.Errorf("failed writing DNS continue response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (d *DnsInterceptor) writeDNSError(w dns.ResponseWriter, r *dns.Msg, logger *log.Entry, reason string) {
|
||||
logger.Warnf("failed to query upstream for domain=%s: %s", r.Question[0].Name, reason)
|
||||
|
||||
@@ -621,3 +699,26 @@ func (d *DnsInterceptor) debugPeerTimeout(peerIP netip.Addr, peerKey string) str
|
||||
|
||||
return fmt.Sprintf(" (peer %s)", nbdns.FormatPeerStatus(&peerState))
|
||||
}
|
||||
|
||||
// peerQuery builds the query sent to a peer's DNS forwarder and reports whether
|
||||
// the client itself advertised EDNS0. EDNS0 is added so the forwarder can return
|
||||
// an Extended DNS Error describing an upstream failure, and the OPT is stripped
|
||||
// from the reply again when the client did not ask for it. The client's message
|
||||
// is left untouched: it may still be handed to the next handler in the chain,
|
||||
// and neither the OPT nor the AD bit is ours to put on the wire on its behalf.
|
||||
func peerQuery(r *dns.Msg) (query *dns.Msg, hadEdns bool) {
|
||||
query = r.Copy()
|
||||
hadEdns = query.IsEdns0() != nil
|
||||
|
||||
// AD tells the forwarder we understand authenticated data. Only set when the
|
||||
// client sent no additional section of its own, so we never overrule what it
|
||||
// asked for.
|
||||
if len(query.Extra) == 0 {
|
||||
query.MsgHdr.AuthenticatedData = true
|
||||
}
|
||||
if !hadEdns {
|
||||
query.SetEdns0(dns.DefaultMsgSize, false)
|
||||
}
|
||||
|
||||
return query, hadEdns
|
||||
}
|
||||
|
||||
399
client/internal/routemanager/dnsinterceptor/handler_test.go
Normal file
399
client/internal/routemanager/dnsinterceptor/handler_test.go
Normal file
@@ -0,0 +1,399 @@
|
||||
package dnsinterceptor
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/miekg/dns"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.zx2c4.com/wireguard/tun/netstack"
|
||||
|
||||
"github.com/netbirdio/netbird/client/iface/device"
|
||||
"github.com/netbirdio/netbird/client/iface/wgaddr"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/resutil"
|
||||
"github.com/netbirdio/netbird/client/internal/dns/test"
|
||||
"github.com/netbirdio/netbird/client/internal/peer"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
// softNegativeWriter records what the handler told the chain: whether to soften
|
||||
// a negative verdict from the handlers it defers to, and the metadata that ends
|
||||
// up on the chain's response log line.
|
||||
type softNegativeWriter struct {
|
||||
test.MockResponseWriter
|
||||
softNegative bool
|
||||
meta map[resutil.MetaKey]string
|
||||
}
|
||||
|
||||
func (w *softNegativeWriter) RequestSoftNegative() { w.softNegative = true }
|
||||
|
||||
func (w *softNegativeWriter) RequestID() string { return "test" }
|
||||
|
||||
func (w *softNegativeWriter) SetMeta(key resutil.MetaKey, value string) {
|
||||
if w.meta == nil {
|
||||
w.meta = make(map[resutil.MetaKey]string)
|
||||
}
|
||||
w.meta[key] = value
|
||||
}
|
||||
|
||||
// TestServeDNS_QtypeTheForwarderCannotResolve covers the record types no DNS
|
||||
// forwarder can answer, because the host resolver exposes no API for them.
|
||||
// Asking the peer only burns a tunnel round trip, so the query goes straight to
|
||||
// the rest of the chain. No peer key is configured, proving no round trip is
|
||||
// attempted.
|
||||
func TestServeDNS_QtypeTheForwarderCannotResolve(t *testing.T) {
|
||||
qtypes := []uint16{
|
||||
dns.TypeHTTPS,
|
||||
dns.TypeSVCB,
|
||||
dns.TypeCAA,
|
||||
dns.TypeNAPTR,
|
||||
dns.TypeTLSA,
|
||||
dns.TypeSOA,
|
||||
}
|
||||
|
||||
for _, qtype := range qtypes {
|
||||
t.Run(dns.TypeToString[qtype], func(t *testing.T) {
|
||||
d := &DnsInterceptor{}
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", qtype)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "chain continuation is signalled as NXDOMAIN")
|
||||
assert.True(t, resp.MsgHdr.Zero, "Zero bit must be set so the chain continues")
|
||||
assert.True(t, w.softNegative,
|
||||
"a downstream NXDOMAIN must be softened, or the fallthrough poisons the routed name")
|
||||
assert.Equal(t, "dns-route", w.meta[resutil.MetaKeyDeferredBy],
|
||||
"the chain's response log line must name us as the handler that stepped aside")
|
||||
assert.Equal(t, "unsupported-qtype", w.meta[resutil.MetaKeyDeferredReason],
|
||||
"the reason must survive to the log line, or the fallthrough is invisible")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_AddressQtypeNeverFallsThrough locks in the opposite case: A and
|
||||
// AAAA are what the route exists for, and the peer is authoritative for them.
|
||||
// A missing peer is an error the client must see, never a fallthrough to a
|
||||
// resolver that knows nothing about the routed name.
|
||||
func TestServeDNS_AddressQtypeNeverFallsThrough(t *testing.T) {
|
||||
for _, qtype := range []uint16{dns.TypeA, dns.TypeAAAA} {
|
||||
t.Run(dns.TypeToString[qtype], func(t *testing.T) {
|
||||
d := &DnsInterceptor{}
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("db.example.com.", qtype)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeServerFailure, resp.Rcode, "an unusable route must fail, not fall through")
|
||||
assert.False(t, resp.MsgHdr.Zero, "the chain must not continue for address queries")
|
||||
assert.False(t, w.softNegative, "no fallthrough means nothing to soften")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDispositionFor pins the query-type policy, and ties the types we ask a
|
||||
// peer about to the types a forwarder can actually resolve, so the two ends
|
||||
// cannot drift apart.
|
||||
func TestDispositionFor(t *testing.T) {
|
||||
tests := []struct {
|
||||
qtype uint16
|
||||
want disposition
|
||||
}{
|
||||
{dns.TypeA, dispositionPeer},
|
||||
{dns.TypeAAAA, dispositionPeer},
|
||||
{dns.TypeMX, dispositionPeerFirst},
|
||||
{dns.TypeTXT, dispositionPeerFirst},
|
||||
{dns.TypeNS, dispositionPeerFirst},
|
||||
{dns.TypeSRV, dispositionPeerFirst},
|
||||
{dns.TypeCNAME, dispositionPeerFirst},
|
||||
{dns.TypePTR, dispositionPeerFirst},
|
||||
{dns.TypeHTTPS, dispositionChain},
|
||||
{dns.TypeSVCB, dispositionChain},
|
||||
{dns.TypeCAA, dispositionChain},
|
||||
{dns.TypeNAPTR, dispositionChain},
|
||||
{dns.TypeTLSA, dispositionChain},
|
||||
{dns.TypeDS, dispositionChain},
|
||||
{dns.TypeDNSKEY, dispositionChain},
|
||||
{dns.TypeSOA, dispositionChain},
|
||||
{dns.TypeANY, dispositionChain},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(dns.TypeToString[tt.qtype], func(t *testing.T) {
|
||||
assert.Equal(t, tt.want, dispositionFor(tt.qtype), "disposition for %s", dns.TypeToString[tt.qtype])
|
||||
|
||||
// Address types are resolved by the forwarder's address path, not by
|
||||
// its record path, so they are deliberately outside
|
||||
// SupportedRecordQtype and their disposition does not depend on it.
|
||||
if tt.want == dispositionPeer {
|
||||
assert.False(t, resutil.SupportedRecordQtype(tt.qtype),
|
||||
"address types take the forwarder's address path")
|
||||
return
|
||||
}
|
||||
|
||||
assert.Equal(t, tt.want == dispositionPeerFirst, resutil.SupportedRecordQtype(tt.qtype),
|
||||
"only types a forwarder resolves may be sent to a peer")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeerCannotAnswer separates a peer that cannot resolve a type from one that
|
||||
// resolved it and found nothing. Only the former may be retried elsewhere:
|
||||
// treating NODATA or NXDOMAIN as a capability failure would send the client to a
|
||||
// public resolver behind the peer's back and prefer a public record over the
|
||||
// private one the route exists to reach.
|
||||
func TestPeerCannotAnswer(t *testing.T) {
|
||||
cannot := []int{dns.RcodeNotImplemented, dns.RcodeFormatError}
|
||||
for _, rcode := range cannot {
|
||||
assert.True(t, peerCannotAnswer(rcode), "%s means the peer cannot answer", dns.RcodeToString[rcode])
|
||||
}
|
||||
|
||||
// REFUSED belongs here, not above: it says the peer does not hold the name,
|
||||
// so retrying elsewhere would leak the name of an internal-only domain to a
|
||||
// public resolver. No forwarder version reports a missing record type this
|
||||
// way, so nothing is lost by keeping it out.
|
||||
can := []int{
|
||||
dns.RcodeSuccess,
|
||||
dns.RcodeNameError,
|
||||
dns.RcodeServerFailure,
|
||||
dns.RcodeNotAuth,
|
||||
dns.RcodeRefused,
|
||||
}
|
||||
for _, rcode := range can {
|
||||
assert.False(t, peerCannotAnswer(rcode), "%s is the peer's answer, not a capability failure", dns.RcodeToString[rcode])
|
||||
}
|
||||
}
|
||||
|
||||
// fakeWGIface is the minimum an interceptor needs to reach a peer's DNS
|
||||
// forwarder over a plain socket. A nil netstack keeps the exchange on the host
|
||||
// stack so the test can point it at a loopback listener.
|
||||
type fakeWGIface struct{}
|
||||
|
||||
func (fakeWGIface) AddAllowedIP(string, netip.Prefix) error { return nil }
|
||||
func (fakeWGIface) RemoveAllowedIP(string, netip.Prefix) error { return nil }
|
||||
func (fakeWGIface) Name() string { return "wt0" }
|
||||
func (fakeWGIface) Address() wgaddr.Address { return wgaddr.Address{} }
|
||||
func (fakeWGIface) ToInterface() *net.Interface { return nil }
|
||||
func (fakeWGIface) IsUserspaceBind() bool { return false }
|
||||
func (fakeWGIface) GetFilter() device.PacketFilter { return nil }
|
||||
func (fakeWGIface) GetDevice() *device.FilteredDevice { return nil }
|
||||
func (fakeWGIface) GetNet() *netstack.Net { return nil }
|
||||
|
||||
// fakePeerIPs resolves every peer key to the loopback address the test's
|
||||
// forwarder stub listens on.
|
||||
type fakePeerIPs struct {
|
||||
addr netip.Addr
|
||||
}
|
||||
|
||||
func (p fakePeerIPs) AllowedIP(string) (netip.Addr, bool) { return p.addr, p.addr.IsValid() }
|
||||
|
||||
// forwarderStub stands in for the DNS forwarder of a routing peer, recording
|
||||
// what it was asked and answering with a canned reply.
|
||||
type forwarderStub struct {
|
||||
mu sync.Mutex
|
||||
queries []*dns.Msg
|
||||
reply func(q *dns.Msg) *dns.Msg
|
||||
port uint16
|
||||
shutdown func()
|
||||
}
|
||||
|
||||
func (s *forwarderStub) received() []*dns.Msg {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return append([]*dns.Msg(nil), s.queries...)
|
||||
}
|
||||
|
||||
func newForwarderStub(t *testing.T, reply func(q *dns.Msg) *dns.Msg) *forwarderStub {
|
||||
t.Helper()
|
||||
|
||||
conn, err := net.ListenUDP("udp", net.UDPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:0")))
|
||||
require.NoError(t, err, "listen for the forwarder stub")
|
||||
|
||||
stub := &forwarderStub{reply: reply}
|
||||
stub.port = uint16(conn.LocalAddr().(*net.UDPAddr).Port)
|
||||
|
||||
mux := dns.NewServeMux()
|
||||
mux.HandleFunc(".", func(w dns.ResponseWriter, q *dns.Msg) {
|
||||
stub.mu.Lock()
|
||||
stub.queries = append(stub.queries, q.Copy())
|
||||
stub.mu.Unlock()
|
||||
_ = w.WriteMsg(stub.reply(q))
|
||||
})
|
||||
|
||||
srv := &dns.Server{PacketConn: conn, Handler: mux}
|
||||
started := make(chan struct{})
|
||||
srv.NotifyStartedFunc = func() { close(started) }
|
||||
go func() {
|
||||
_ = srv.ActivateAndServe()
|
||||
}()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("forwarder stub did not start")
|
||||
}
|
||||
|
||||
stub.shutdown = func() { _ = srv.Shutdown() }
|
||||
t.Cleanup(stub.shutdown)
|
||||
|
||||
return stub
|
||||
}
|
||||
|
||||
func newTestInterceptor(t *testing.T, stub *forwarderStub) *DnsInterceptor {
|
||||
t.Helper()
|
||||
|
||||
port := new(atomic.Uint32)
|
||||
port.Store(uint32(stub.port))
|
||||
|
||||
return &DnsInterceptor{
|
||||
route: &route.Route{Domains: nil},
|
||||
statusRecorder: peer.NewRecorder("https://mgm"),
|
||||
currentPeerKey: "peer-key",
|
||||
interceptedDomains: make(domainMap),
|
||||
wgInterface: fakeWGIface{},
|
||||
peerStore: fakePeerIPs{addr: netip.MustParseAddr("127.0.0.1")},
|
||||
forwarderPort: port,
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerCannotResolveQtype is the reported regression: a routing peer
|
||||
// running a client older than the one that taught the forwarder about non-address
|
||||
// record types answers NOTIMP for every SRV query, and the interceptor used to
|
||||
// hand that straight to the application, breaking SRV-based service discovery.
|
||||
// The client cannot know the peer's capability up front, so it must try the peer
|
||||
// and fall back to the chain when the peer says it cannot answer.
|
||||
func TestServeDNS_PeerCannotResolveQtype(t *testing.T) {
|
||||
for _, rcode := range []int{dns.RcodeNotImplemented, dns.RcodeFormatError} {
|
||||
t.Run(dns.RcodeToString[rcode], func(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, rcode)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
require.Len(t, stub.received(), 1, "the peer must be asked before giving up on it")
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeNameError, resp.Rcode, "chain continuation is signalled as NXDOMAIN")
|
||||
assert.True(t, resp.MsgHdr.Zero, "Zero bit must be set so the chain continues")
|
||||
assert.True(t, w.softNegative, "the fallthrough must not be allowed to poison the routed name")
|
||||
assert.Equal(t, "dns-route", w.meta[resutil.MetaKeyDeferredBy])
|
||||
assert.Equal(t, "peer-rcode-"+dns.RcodeToString[rcode], w.meta[resutil.MetaKeyDeferredReason],
|
||||
"the peer's verdict must be visible in the log, it is the only trace of the probe")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerVerdictPassesThrough covers the replies that are the peer's
|
||||
// answer rather than a statement about its capability. The peer is authoritative
|
||||
// for a name routed to it, so its verdict reaches the client as it is: retrying
|
||||
// these elsewhere would send the name of an internal-only domain to a public
|
||||
// resolver and prefer whatever that resolver says over the route.
|
||||
func TestServeDNS_PeerVerdictPassesThrough(t *testing.T) {
|
||||
// REFUSED means the peer does not consider the name routed to it, which
|
||||
// happens while the client and the peer disagree about the route set. It is
|
||||
// never how a peer reports a record type it cannot resolve: a forwarder too
|
||||
// old for non-address types answers NOTIMP before it looks at the domain.
|
||||
for _, rcode := range []int{dns.RcodeRefused, dns.RcodeNameError, dns.RcodeSuccess} {
|
||||
t.Run(dns.RcodeToString[rcode], func(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, rcode)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
require.Len(t, stub.received(), 1, "the peer must be asked exactly once")
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, rcode, resp.Rcode, "the peer's verdict must reach the client unchanged")
|
||||
assert.False(t, resp.MsgHdr.Zero, "the chain must not continue past an answered query")
|
||||
assert.False(t, w.softNegative, "nothing was deferred, so nothing may be softened")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeDNS_PeerResolvesQtype guards the common case: a peer that can answer
|
||||
// the type stays authoritative, and its records reach the client unchanged.
|
||||
func TestServeDNS_PeerResolvesQtype(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetReply(q)
|
||||
resp.Answer = []dns.RR{&dns.SRV{
|
||||
Hdr: dns.RR_Header{Name: q.Question[0].Name, Rrtype: dns.TypeSRV, Class: dns.ClassINET, Ttl: 60},
|
||||
Target: "shard-00.db.example.com.",
|
||||
Port: 27017,
|
||||
}}
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
resp := w.GetLastResponse()
|
||||
require.NotNil(t, resp, "a response must be written")
|
||||
assert.Equal(t, dns.RcodeSuccess, resp.Rcode)
|
||||
require.Len(t, resp.Answer, 1, "the peer's records must pass through")
|
||||
assert.False(t, resp.MsgHdr.Zero, "an answered query must not continue the chain")
|
||||
assert.False(t, w.softNegative)
|
||||
}
|
||||
|
||||
// TestServeDNS_DeferredQueryIsPristine covers what the fallthrough hands to the
|
||||
// next handler. The interceptor advertises EDNS0 to the peer so the forwarder can
|
||||
// return an Extended DNS Error, and sets AD, but neither belongs to the client's
|
||||
// query: on the deferred path they would travel to the public resolver and the
|
||||
// reply could come back carrying an OPT the client never advertised, which
|
||||
// RFC 6891 forbids us from passing on.
|
||||
func TestServeDNS_DeferredQueryIsPristine(t *testing.T) {
|
||||
stub := newForwarderStub(t, func(q *dns.Msg) *dns.Msg {
|
||||
resp := new(dns.Msg)
|
||||
resp.SetRcode(q, dns.RcodeNotImplemented)
|
||||
return resp
|
||||
})
|
||||
d := newTestInterceptor(t, stub)
|
||||
w := &softNegativeWriter{}
|
||||
|
||||
r := new(dns.Msg)
|
||||
r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
|
||||
|
||||
d.ServeDNS(w, r)
|
||||
|
||||
sent := stub.received()
|
||||
require.Len(t, sent, 1)
|
||||
assert.NotNil(t, sent[0].IsEdns0(), "the peer must be asked with EDNS0 so it can return an EDE")
|
||||
|
||||
assert.Nil(t, r.IsEdns0(), "the deferred query must not carry an OPT the client never sent")
|
||||
assert.False(t, r.AuthenticatedData, "the deferred query must not carry an AD bit the client never set")
|
||||
assert.Empty(t, r.Extra, "the deferred query must reach the next handler as the client sent it")
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//go:build windows
|
||||
|
||||
package dnsinterceptor
|
||||
|
||||
// GetInterfaceGUIDString completes iface.WGIface on Windows, which requires the
|
||||
// interface GUID for DNS registration.
|
||||
func (fakeWGIface) GetInterfaceGUIDString() (string, error) { return "", nil }
|
||||
@@ -61,6 +61,7 @@ type Manager interface {
|
||||
InitialRouteRange() []string
|
||||
SetFirewall(firewall.Manager) error
|
||||
SetDNSForwarderPort(port uint16)
|
||||
ReconcilePeerAllowedIPs(peerKey string) error
|
||||
Stop(stateManager *statemanager.Manager)
|
||||
}
|
||||
|
||||
@@ -232,6 +233,30 @@ 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()
|
||||
|
||||
@@ -112,6 +112,11 @@ 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 {
|
||||
|
||||
90
client/internal/routemanager/reconcile_test.go
Normal file
90
client/internal/routemanager/reconcile_test.go
Normal file
@@ -0,0 +1,90 @@
|
||||
//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"))
|
||||
}
|
||||
@@ -94,6 +94,26 @@ 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) {
|
||||
|
||||
47
client/internal/routemanager/refcounter/refcounter_test.go
Normal file
47
client/internal/routemanager/refcounter/refcounter_test.go
Normal file
@@ -0,0 +1,47 @@
|
||||
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)
|
||||
}
|
||||
@@ -79,13 +79,15 @@ type Info struct {
|
||||
EnableSSHLocalPortForwarding bool
|
||||
EnableSSHRemotePortForwarding bool
|
||||
DisableSSHAuth bool
|
||||
|
||||
SyncMessageVersion *int
|
||||
}
|
||||
|
||||
func (i *Info) SetFlags(
|
||||
rosenpassEnabled, rosenpassPermissive bool,
|
||||
serverSSHAllowed *bool,
|
||||
disableClientRoutes, disableServerRoutes,
|
||||
disableDNS, disableFirewall, blockLANAccess, blockInbound, disableIPv6 bool,
|
||||
disableDNS, disableFirewall, blockLANAccess, blockInbound, disableIPv6 bool, syncMessageVersion *int,
|
||||
enableSSHRoot, enableSSHSFTP, enableSSHLocalPortForwarding, enableSSHRemotePortForwarding *bool,
|
||||
disableSSHAuth *bool,
|
||||
) {
|
||||
@@ -103,6 +105,8 @@ func (i *Info) SetFlags(
|
||||
i.BlockInbound = blockInbound
|
||||
i.DisableIPv6 = disableIPv6
|
||||
|
||||
i.SyncMessageVersion = syncMessageVersion
|
||||
|
||||
if enableSSHRoot != nil {
|
||||
i.EnableSSHRoot = *enableSSHRoot
|
||||
}
|
||||
|
||||
151
combined/cmd/admin.go
Normal file
151
combined/cmd/admin.go
Normal file
@@ -0,0 +1,151 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/dexidp/dex/storage"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
admincmd "github.com/netbirdio/netbird/management/cmd/admin"
|
||||
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server/activity"
|
||||
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
// newAdminCommands creates the admin command tree with combined-specific resource openers.
|
||||
func newAdminCommands() *cobra.Command {
|
||||
return admincmd.NewCommands(admincmd.Openers{
|
||||
Resources: withAdminResources,
|
||||
Store: withAdminStoreOnly,
|
||||
IDP: withAdminIDPOnly,
|
||||
})
|
||||
}
|
||||
|
||||
func newLegacyTokenCommand() *cobra.Command {
|
||||
cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
|
||||
cmd.Deprecated = "use 'admin token' instead"
|
||||
return cmd
|
||||
}
|
||||
|
||||
// withAdminResources loads the combined YAML config, initializes stores, and calls fn.
|
||||
func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
|
||||
return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
|
||||
mgmtConfig, err := adminManagementConfig(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
managementStore, err := openAdminStore(ctx, cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseStore(ctx, managementStore)
|
||||
|
||||
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(mgmtConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseIDPStorage(idpStorage)
|
||||
|
||||
eventStore, esErr := openAdminEventStore(ctx, cfg, mgmtConfig)
|
||||
if esErr != nil {
|
||||
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
|
||||
}
|
||||
if eventStore != nil {
|
||||
defer func() {
|
||||
if err := eventStore.Close(ctx); err != nil {
|
||||
log.Debugf("close activity event store: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
|
||||
})
|
||||
}
|
||||
|
||||
// withAdminStoreOnly opens only the management store for admin subcommands that do not
|
||||
// need embedded IdP storage.
|
||||
func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
|
||||
return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
|
||||
managementStore, err := openAdminStore(ctx, cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseStore(ctx, managementStore)
|
||||
|
||||
return fn(ctx, managementStore)
|
||||
})
|
||||
}
|
||||
|
||||
func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
|
||||
return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
|
||||
mgmtConfig, err := adminManagementConfig(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(mgmtConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseIDPStorage(idpStorage)
|
||||
|
||||
return fn(ctx, idpStorage, idpStorageFile)
|
||||
})
|
||||
}
|
||||
|
||||
func withAdminConfig(cmd *cobra.Command, fn func(ctx context.Context, cfg *CombinedConfig) error) error {
|
||||
if err := util.InitLog("error", "console"); err != nil {
|
||||
return fmt.Errorf("init log: %w", err)
|
||||
}
|
||||
|
||||
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
|
||||
|
||||
cfg, err := LoadConfig(configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load config: %w", err)
|
||||
}
|
||||
cfg.ApplyAdminDefaults()
|
||||
applyServerStoreEnv(cfg.Server.Store)
|
||||
|
||||
return fn(ctx, cfg)
|
||||
}
|
||||
|
||||
func adminManagementConfig(cfg *CombinedConfig) (*nbconfig.Config, error) {
|
||||
mgmtConfig, err := cfg.ToManagementConfig()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create management config: %w", err)
|
||||
}
|
||||
return mgmtConfig, nil
|
||||
}
|
||||
|
||||
func openAdminStore(ctx context.Context, cfg *CombinedConfig) (store.Store, error) {
|
||||
managementStore, err := store.NewStore(ctx, types.Engine(cfg.Management.Store.Engine), cfg.Management.DataDir, nil, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create store: %w", err)
|
||||
}
|
||||
return managementStore, nil
|
||||
}
|
||||
|
||||
func openAdminEventStore(ctx context.Context, cfg *CombinedConfig, config *nbconfig.Config) (activity.Store, error) {
|
||||
if config.DataStoreEncryptionKey == "" {
|
||||
return nil, fmt.Errorf("data store encryption key is not configured")
|
||||
}
|
||||
if err := applyActivityStoreEnv(cfg.Server.ActivityStore); err != nil {
|
||||
return nil, fmt.Errorf("configure activity event store: %w", err)
|
||||
}
|
||||
eventStore, err := activitystore.NewSqlStore(ctx, config.Datadir, config.DataStoreEncryptionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open activity event store: %w", err)
|
||||
}
|
||||
if eventStore == nil {
|
||||
return nil, fmt.Errorf("open activity event store: returned nil store")
|
||||
}
|
||||
return eventStore, nil
|
||||
}
|
||||
47
combined/cmd/admin_config_test.go
Normal file
47
combined/cmd/admin_config_test.go
Normal file
@@ -0,0 +1,47 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
)
|
||||
|
||||
func TestApplyAdminDefaultsCopiesServerStoreWithoutExposedAddress(t *testing.T) {
|
||||
cfg := DefaultConfig()
|
||||
cfg.Server.ExposedAddress = ""
|
||||
cfg.Server.DataDir = "/srv/netbird"
|
||||
cfg.Server.Store = StoreConfig{
|
||||
Engine: "postgres",
|
||||
DSN: "postgres://user:pass@example.com/netbird",
|
||||
}
|
||||
|
||||
cfg.ApplyAdminDefaults()
|
||||
|
||||
require.Equal(t, "/srv/netbird", cfg.Management.DataDir)
|
||||
require.Equal(t, "postgres", cfg.Management.Store.Engine)
|
||||
require.Equal(t, cfg.Server.Store.DSN, cfg.Management.Store.DSN)
|
||||
}
|
||||
|
||||
func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
|
||||
eventStore, err := openAdminEventStore(context.Background(), &CombinedConfig{}, &nbconfig.Config{})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "encryption key")
|
||||
require.Nil(t, eventStore)
|
||||
}
|
||||
|
||||
func TestApplyServerStoreEnv(t *testing.T) {
|
||||
t.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", "")
|
||||
t.Setenv("NB_STORE_ENGINE_MYSQL_DSN", "")
|
||||
t.Setenv("NB_STORE_ENGINE_SQLITE_FILE", "")
|
||||
|
||||
applyServerStoreEnv(StoreConfig{Engine: "postgres", DSN: "postgres-dsn", File: "store.db"})
|
||||
require.Equal(t, "postgres-dsn", os.Getenv("NB_STORE_ENGINE_POSTGRES_DSN"))
|
||||
require.Equal(t, "store.db", os.Getenv("NB_STORE_ENGINE_SQLITE_FILE"))
|
||||
|
||||
applyServerStoreEnv(StoreConfig{Engine: "mysql", DSN: "mysql-dsn"})
|
||||
require.Equal(t, "mysql-dsn", os.Getenv("NB_STORE_ENGINE_MYSQL_DSN"))
|
||||
}
|
||||
@@ -6,8 +6,7 @@ import (
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
filePath "path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -74,6 +73,9 @@ type ServerConfig struct {
|
||||
ActivityStore StoreConfig `yaml:"activityStore"`
|
||||
AuthStore StoreConfig `yaml:"authStore"`
|
||||
ReverseProxy ReverseProxyConfig `yaml:"reverseProxy"`
|
||||
|
||||
SupportedSyncMessageVersions *int `yaml:"supportedSyncMessageVersions,omitempty"`
|
||||
PerAccountSupportedSyncMessageVersions map[string]int `yaml:"perAccountSupportedSyncMessageVersions,omitempty"`
|
||||
}
|
||||
|
||||
// TLSConfig contains TLS/HTTPS settings
|
||||
@@ -300,6 +302,19 @@ func (c *CombinedConfig) ApplySimplifiedDefaults() {
|
||||
c.autoConfigureClientSettings(exposedProto, exposedHost, exposedHostPort, hasExternalStuns, hasExternalRelay, hasExternalSignal)
|
||||
}
|
||||
|
||||
// ApplyAdminDefaults applies the management settings needed by admin commands even
|
||||
// when the full server config is invalid and ApplySimplifiedDefaults cannot run.
|
||||
func (c *CombinedConfig) ApplyAdminDefaults() {
|
||||
if c.Management.DataDir == "" || c.Management.DataDir == "/var/lib/netbird/" {
|
||||
c.Management.DataDir = c.Server.DataDir
|
||||
}
|
||||
if c.Management.Store.Engine == "" || c.Management.Store.Engine == "sqlite" {
|
||||
if c.Server.Store.Engine != "" || c.Server.Store.File != "" || c.Server.Store.DSN != "" {
|
||||
c.Management.Store = c.Server.Store
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyRelayDefaults configures the relay service if no external relay is configured.
|
||||
func (c *CombinedConfig) applyRelayDefaults(exposedProto, exposedHostPort string, hasExternalRelay, hasExternalStuns bool) {
|
||||
if hasExternalRelay {
|
||||
@@ -577,11 +592,11 @@ func (c *CombinedConfig) buildEmbeddedIdPConfig(mgmt ManagementConfig) (*idp.Emb
|
||||
return nil, fmt.Errorf("authStore.dsn is required when authStore.engine is postgres")
|
||||
}
|
||||
} else {
|
||||
authStorageFile = path.Join(mgmt.DataDir, "idp.db")
|
||||
authStorageFile = filePath.Join(mgmt.DataDir, "idp.db")
|
||||
if c.Server.AuthStore.File != "" {
|
||||
authStorageFile = c.Server.AuthStore.File
|
||||
if !filepath.IsAbs(authStorageFile) {
|
||||
authStorageFile = filepath.Join(mgmt.DataDir, authStorageFile)
|
||||
if !filePath.IsAbs(authStorageFile) {
|
||||
authStorageFile = filePath.Join(mgmt.DataDir, authStorageFile)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -696,16 +711,18 @@ func (c *CombinedConfig) ToManagementConfig() (*nbconfig.Config, error) {
|
||||
httpConfig.AuthCallbackURL = callbackURL + types.ProxyCallbackEndpointFull
|
||||
|
||||
return &nbconfig.Config{
|
||||
Stuns: stuns,
|
||||
Relay: relayConfig,
|
||||
Signal: signalConfig,
|
||||
Datadir: mgmt.DataDir,
|
||||
DataStoreEncryptionKey: mgmt.Store.EncryptionKey,
|
||||
HttpConfig: httpConfig,
|
||||
StoreConfig: storeConfig,
|
||||
ReverseProxy: reverseProxy,
|
||||
DisableDefaultPolicy: mgmt.DisableDefaultPolicy,
|
||||
EmbeddedIdP: embeddedIdP,
|
||||
Stuns: stuns,
|
||||
Relay: relayConfig,
|
||||
Signal: signalConfig,
|
||||
Datadir: mgmt.DataDir,
|
||||
DataStoreEncryptionKey: mgmt.Store.EncryptionKey,
|
||||
HttpConfig: httpConfig,
|
||||
StoreConfig: storeConfig,
|
||||
ReverseProxy: reverseProxy,
|
||||
DisableDefaultPolicy: mgmt.DisableDefaultPolicy,
|
||||
EmbeddedIdP: embeddedIdP,
|
||||
HighestSupportedSyncMessageVersion: c.Server.SupportedSyncMessageVersions,
|
||||
PerAccountHighestSupportedSyncMessageVersion: c.Server.PerAccountSupportedSyncMessageVersions,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -729,7 +746,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config, mgmtPort
|
||||
cfg.EmbeddedIdP.Storage.Type = "sqlite3"
|
||||
}
|
||||
if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
|
||||
cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
|
||||
cfg.EmbeddedIdP.Storage.Config.File = filePath.Join(cfg.Datadir, "idp.db")
|
||||
}
|
||||
|
||||
issuer := cfg.EmbeddedIdP.Issuer
|
||||
|
||||
@@ -31,6 +31,7 @@ import (
|
||||
relayServer "github.com/netbirdio/netbird/relay/server"
|
||||
"github.com/netbirdio/netbird/relay/server/listener"
|
||||
"github.com/netbirdio/netbird/relay/server/listener/ws"
|
||||
syncgrpc "github.com/netbirdio/netbird/shared/management/grpc"
|
||||
sharedMetrics "github.com/netbirdio/netbird/shared/metrics"
|
||||
"github.com/netbirdio/netbird/shared/relay/auth"
|
||||
"github.com/netbirdio/netbird/shared/signal/proto"
|
||||
@@ -64,7 +65,8 @@ func init() {
|
||||
rootCmd.PersistentFlags().StringVarP(&configPath, "config", "c", "", "path to YAML configuration file (required)")
|
||||
_ = rootCmd.MarkPersistentFlagRequired("config")
|
||||
|
||||
rootCmd.AddCommand(newTokenCommands())
|
||||
rootCmd.AddCommand(newAdminCommands())
|
||||
rootCmd.AddCommand(newLegacyTokenCommand())
|
||||
}
|
||||
|
||||
func RootCmd() *cobra.Command {
|
||||
@@ -122,6 +124,37 @@ func execute(cmd *cobra.Command, _ []string) error {
|
||||
}
|
||||
|
||||
// initializeConfig loads and validates the configuration, then initializes logging.
|
||||
func applyServerStoreEnv(storeConfig StoreConfig) {
|
||||
if dsn := storeConfig.DSN; dsn != "" {
|
||||
switch strings.ToLower(storeConfig.Engine) {
|
||||
case "postgres":
|
||||
os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
|
||||
case "mysql":
|
||||
os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
|
||||
}
|
||||
}
|
||||
if file := storeConfig.File; file != "" {
|
||||
os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
|
||||
}
|
||||
}
|
||||
|
||||
func applyActivityStoreEnv(storeConfig StoreConfig) error {
|
||||
if engine := storeConfig.Engine; engine != "" {
|
||||
engineLower := strings.ToLower(engine)
|
||||
if engineLower == "postgres" && storeConfig.DSN == "" {
|
||||
return fmt.Errorf("activityStore.dsn is required when activityStore.engine is postgres")
|
||||
}
|
||||
os.Setenv("NB_ACTIVITY_EVENT_STORE_ENGINE", engineLower)
|
||||
if dsn := storeConfig.DSN; dsn != "" {
|
||||
os.Setenv("NB_ACTIVITY_EVENT_POSTGRES_DSN", dsn)
|
||||
}
|
||||
}
|
||||
if file := storeConfig.File; file != "" {
|
||||
os.Setenv("NB_ACTIVITY_EVENT_SQLITE_FILE", file)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func initializeConfig() error {
|
||||
var err error
|
||||
config, err = LoadConfig(configPath)
|
||||
@@ -137,30 +170,10 @@ func initializeConfig() error {
|
||||
return fmt.Errorf("failed to initialize log: %w", err)
|
||||
}
|
||||
|
||||
if dsn := config.Server.Store.DSN; dsn != "" {
|
||||
switch strings.ToLower(config.Server.Store.Engine) {
|
||||
case "postgres":
|
||||
os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
|
||||
case "mysql":
|
||||
os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
|
||||
}
|
||||
}
|
||||
if file := config.Server.Store.File; file != "" {
|
||||
os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
|
||||
}
|
||||
applyServerStoreEnv(config.Server.Store)
|
||||
|
||||
if engine := config.Server.ActivityStore.Engine; engine != "" {
|
||||
engineLower := strings.ToLower(engine)
|
||||
if engineLower == "postgres" && config.Server.ActivityStore.DSN == "" {
|
||||
return fmt.Errorf("activityStore.dsn is required when activityStore.engine is postgres")
|
||||
}
|
||||
os.Setenv("NB_ACTIVITY_EVENT_STORE_ENGINE", engineLower)
|
||||
if dsn := config.Server.ActivityStore.DSN; dsn != "" {
|
||||
os.Setenv("NB_ACTIVITY_EVENT_POSTGRES_DSN", dsn)
|
||||
}
|
||||
}
|
||||
if file := config.Server.ActivityStore.File; file != "" {
|
||||
os.Setenv("NB_ACTIVITY_EVENT_SQLITE_FILE", file)
|
||||
if err := applyActivityStoreEnv(config.Server.ActivityStore); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.Infof("Starting combined NetBird server")
|
||||
@@ -505,6 +518,16 @@ func createManagementServer(cfg *CombinedConfig, mgmtConfig *nbconfig.Config) (m
|
||||
}
|
||||
mgmtPort, _ := strconv.Atoi(portStr)
|
||||
|
||||
if err := syncgrpc.ValidateSyncMessageVersion(mgmtConfig.HighestSupportedSyncMessageVersion); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for accountId, version := range mgmtConfig.PerAccountHighestSupportedSyncMessageVersion {
|
||||
if err := syncgrpc.ValidateSyncMessageVersion(&version); err != nil {
|
||||
return nil, fmt.Errorf("unrecognized sync message version in perAccountSupportedSyncMessageVersions for account %s %w", accountId, err)
|
||||
}
|
||||
}
|
||||
|
||||
mgmtSrv := newServer(
|
||||
&mgmtServer.Config{
|
||||
NbConfig: mgmtConfig,
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
// newTokenCommands creates the token command tree with combined-specific store opener.
|
||||
func newTokenCommands() *cobra.Command {
|
||||
return tokencmd.NewCommands(withTokenStore)
|
||||
}
|
||||
|
||||
// withTokenStore loads the combined YAML config, initializes the store, and calls fn.
|
||||
func withTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
|
||||
if err := util.InitLog("error", "console"); err != nil {
|
||||
return fmt.Errorf("init log: %w", err)
|
||||
}
|
||||
|
||||
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
|
||||
|
||||
cfg, err := LoadConfig(configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load config: %w", err)
|
||||
}
|
||||
|
||||
if dsn := cfg.Server.Store.DSN; dsn != "" {
|
||||
switch strings.ToLower(cfg.Server.Store.Engine) {
|
||||
case "postgres":
|
||||
os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
|
||||
case "mysql":
|
||||
os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
|
||||
}
|
||||
}
|
||||
if file := cfg.Server.Store.File; file != "" {
|
||||
os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
|
||||
}
|
||||
|
||||
datadir := cfg.Management.DataDir
|
||||
engine := types.Engine(cfg.Management.Store.Engine)
|
||||
|
||||
s, err := store.NewStore(ctx, engine, datadir, nil, true)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create store: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := s.Close(ctx); err != nil {
|
||||
log.Debugf("close store: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
return fn(ctx, s)
|
||||
}
|
||||
@@ -53,6 +53,7 @@ type NameServerGroup struct {
|
||||
ID string `gorm:"primaryKey"`
|
||||
// AccountID is a reference to Account that this object belongs
|
||||
AccountID string `gorm:"index"`
|
||||
PublicID string `json:"-"`
|
||||
// Name group name
|
||||
Name string
|
||||
// Description group description
|
||||
|
||||
@@ -9,25 +9,234 @@ 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
|
||||
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)
|
||||
}
|
||||
|
||||
// availableProviders builds the matrix from the provider env vars that are set.
|
||||
@@ -39,6 +248,24 @@ 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})
|
||||
}
|
||||
@@ -84,14 +311,27 @@ func availableProviders() []providerCase {
|
||||
}
|
||||
}
|
||||
|
||||
// Bedrock: path-routed, bearer auth. Model is a cross-region inference
|
||||
// profile id (distinct string from the first-party Anthropic case).
|
||||
// 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.
|
||||
if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
|
||||
region := os.Getenv("AWS_REGION")
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
region = "eu-central-1"
|
||||
}
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock})
|
||||
// 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.
|
||||
model := os.Getenv("AWS_BEDROCK_MODEL")
|
||||
if model == "" {
|
||||
model = "global.anthropic.claude-sonnet-4-6"
|
||||
}
|
||||
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: model, kind: harness.WireBedrock})
|
||||
}
|
||||
return ps
|
||||
}
|
||||
@@ -108,8 +348,16 @@ func providerRequest(pc providerCase) api.AgentNetworkProviderRequest {
|
||||
Enabled: ptr(true),
|
||||
}
|
||||
if pc.kind != harness.WireVertex {
|
||||
// The router matches the normalized catalog id. Bedrock's request model
|
||||
// travels as a region-prefixed inference-profile id in the URL path
|
||||
// (us.anthropic...), which the router strips before matching, so register
|
||||
// the normalized form here or routing fails as model_not_routable.
|
||||
modelID := pc.model
|
||||
if pc.kind == harness.WireBedrock {
|
||||
modelID = catalogModel(pc)
|
||||
}
|
||||
req.Models = &[]api.AgentNetworkProviderModel{
|
||||
{Id: pc.model, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
{Id: modelID, InputPer1k: 0.001, OutputPer1k: 0.002},
|
||||
}
|
||||
}
|
||||
return req
|
||||
@@ -201,11 +449,12 @@ 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.
|
||||
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 {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
|
||||
for _, pc := range matrix {
|
||||
pc := pc
|
||||
@@ -216,6 +465,10 @@ 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
|
||||
@@ -226,11 +479,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, "Reply with exactly: pong", sessionID)
|
||||
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, matrixPrompt, sessionID)
|
||||
case harness.WireBedrock:
|
||||
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
|
||||
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, matrixPrompt, sessionID)
|
||||
default:
|
||||
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
|
||||
c, b, cerr = cl.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, pc.kind, pc.model, matrixPrompt, sessionID)
|
||||
}
|
||||
if cerr == nil {
|
||||
code, body = c, b
|
||||
@@ -249,6 +502,7 @@ 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 {
|
||||
@@ -256,11 +510,15 @@ 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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -281,4 +539,7 @@ 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)
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -15,13 +16,29 @@ import (
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// bedrockRegionPrefixes and bedrockVersionSuffix mirror the proxy's Bedrock
|
||||
// model normalization (region/inference-profile prefix + version suffix) so the
|
||||
// provider is registered under the same catalog key the router matches against.
|
||||
var (
|
||||
bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
|
||||
bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
|
||||
)
|
||||
|
||||
// catalogModel returns the normalized catalog id the proxy stamps for a
|
||||
// path-routed provider's configured model — the form the guardrail allowlist is
|
||||
// compared against (region prefix / @version stripped).
|
||||
// path-routed provider's configured model — the form the router and guardrail
|
||||
// allowlist compare against (Bedrock region prefix + version stripped, Vertex
|
||||
// @version stripped).
|
||||
func catalogModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
return strings.TrimPrefix(pc.model, "us.")
|
||||
m := pc.model
|
||||
for _, p := range bedrockRegionPrefixes {
|
||||
if strings.HasPrefix(m, p) {
|
||||
m = m[len(p):]
|
||||
break
|
||||
}
|
||||
}
|
||||
return bedrockVersionSuffix.ReplaceAllString(m, "")
|
||||
case harness.WireVertex:
|
||||
return strings.SplitN(pc.model, "@", 2)[0]
|
||||
default:
|
||||
@@ -35,7 +52,9 @@ func catalogModel(pc providerCase) string {
|
||||
func disallowedModel(pc providerCase) string {
|
||||
switch pc.kind {
|
||||
case harness.WireBedrock:
|
||||
return "us.anthropic.claude-opus-4-8"
|
||||
// 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"
|
||||
case harness.WireVertex:
|
||||
return "claude-opus-4-8@20250101"
|
||||
default:
|
||||
@@ -55,7 +74,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.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
|
||||
code, _, err = cl.ChatPrefixed(ctx, endpoint, proxyIP, pc.pathPrefix, pc.kind, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
|
||||
return code
|
||||
@@ -147,11 +166,12 @@ 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.
|
||||
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 {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
|
||||
for _, pc := range providers {
|
||||
pc := pc
|
||||
|
||||
@@ -104,11 +104,12 @@ 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.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
|
||||
// Positive: skip=true reaches the self-signed upstream. Retry to absorb
|
||||
// tunnel/DNS jitter on the first call; success also proves the path works.
|
||||
|
||||
@@ -106,11 +106,12 @@ 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.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
|
||||
before, _ := srv.ListAccessLogs(ctx)
|
||||
sessionID := "e2e-session-vllm"
|
||||
|
||||
@@ -4,6 +4,7 @@ package harness
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
@@ -167,22 +168,54 @@ 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)
|
||||
}
|
||||
|
||||
// ResolveProxyIP resolves the agent-network endpoint to the proxy peer's
|
||||
// NetBird IP from inside the client (via magic DNS).
|
||||
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.
|
||||
func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
|
||||
code, reader, err := cl.container.Exec(ctx, []string{"getent", "hosts", endpoint}, tcexec.Multiplexed())
|
||||
if err != nil {
|
||||
return "", err
|
||||
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 + "/",
|
||||
}
|
||||
out, _ := io.ReadAll(reader)
|
||||
if code != 0 {
|
||||
return "", fmt.Errorf("getent hosts %s exited %d", endpoint, code)
|
||||
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):
|
||||
}
|
||||
}
|
||||
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.
|
||||
@@ -206,18 +239,29 @@ 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":64,"messages":[{"role":"user","content":%q}]}`, model, prompt)
|
||||
body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"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, path, body, withSessionID(headers, sessionID))
|
||||
return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
|
||||
}
|
||||
|
||||
// Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
|
||||
@@ -227,7 +271,7 @@ func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prom
|
||||
// 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":64,"messages":[{"role":"user","content":%q}]}`, prompt)
|
||||
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
|
||||
}
|
||||
|
||||
@@ -238,7 +282,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":64,"messages":[{"role":"user","content":%q}]}`, prompt)
|
||||
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
|
||||
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
|
||||
}
|
||||
|
||||
@@ -297,7 +341,7 @@ func (cl *Client) Terminate(ctx context.Context) error {
|
||||
return cl.container.Terminate(ctx)
|
||||
}
|
||||
|
||||
// containerLogs reads up to 256 KiB of a container's logs for diagnostics.
|
||||
// containerLogs reads up to 4 MiB of a container's logs for diagnostics — enough for a whole provider-matrix run.
|
||||
func containerLogs(ctx context.Context, c testcontainers.Container) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
@@ -307,6 +351,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, 256<<10))
|
||||
b, _ := io.ReadAll(io.LimitReader(r, 4<<20))
|
||||
return string(b)
|
||||
}
|
||||
|
||||
@@ -221,6 +221,29 @@ 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)
|
||||
|
||||
@@ -308,7 +308,7 @@ func (s *Storage) OpenStorage(logger *slog.Logger) (storage.Storage, error) {
|
||||
if file == "" {
|
||||
return nil, fmt.Errorf("sqlite3 storage requires 'file' config")
|
||||
}
|
||||
return (&sql.SQLite3{File: file}).Open(logger)
|
||||
return newSQLite3(file).Open(logger)
|
||||
case "postgres":
|
||||
dsn, _ := s.Config["dsn"].(string)
|
||||
if dsn == "" {
|
||||
|
||||
@@ -20,7 +20,6 @@ import (
|
||||
"github.com/dexidp/dex/server"
|
||||
"github.com/dexidp/dex/server/signer"
|
||||
"github.com/dexidp/dex/storage"
|
||||
"github.com/dexidp/dex/storage/sql"
|
||||
"github.com/go-jose/go-jose/v4"
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
@@ -41,7 +40,7 @@ type Config struct {
|
||||
GRPCAddr string
|
||||
}
|
||||
|
||||
const localConnectorID = "local"
|
||||
const LocalConnectorID = "local"
|
||||
|
||||
// Provider wraps a Dex server
|
||||
type Provider struct {
|
||||
@@ -79,7 +78,7 @@ func NewProvider(ctx context.Context, config *Config) (*Provider, error) {
|
||||
|
||||
// Initialize SQLite storage
|
||||
dbPath := filepath.Join(config.DataDir, "oidc.db")
|
||||
sqliteConfig := &sql.SQLite3{File: dbPath}
|
||||
sqliteConfig := newSQLite3(dbPath)
|
||||
stor, err := sqliteConfig.Open(logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to open storage: %w", err)
|
||||
@@ -495,18 +494,60 @@ func (p *Provider) Storage() storage.Storage {
|
||||
return p.storage
|
||||
}
|
||||
|
||||
// SetClientsMFAChain updates the MFAChain field on OAuth2 clients in Dex storage.
|
||||
// Pass a non-empty slice (e.g. []string{"default-totp"}) to enable MFA, or nil to disable it.
|
||||
func SetClientsMFAChain(ctx context.Context, st storage.Storage, clientIDs []string, mfaChain []string) error {
|
||||
previousChains := make(map[string][]string, len(clientIDs))
|
||||
for _, clientID := range clientIDs {
|
||||
client, err := st.GetClient(ctx, clientID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get client %s before MFA chain update: %w", clientID, err)
|
||||
}
|
||||
previousChains[clientID] = cloneMFAChain(client.MFAChain)
|
||||
}
|
||||
|
||||
updatedClientIDs := make([]string, 0, len(clientIDs))
|
||||
for _, clientID := range clientIDs {
|
||||
if err := st.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
|
||||
old.MFAChain = cloneMFAChain(mfaChain)
|
||||
return old, nil
|
||||
}); err != nil {
|
||||
if rollbackErr := rollbackClientsMFAChain(ctx, st, updatedClientIDs, previousChains); rollbackErr != nil {
|
||||
return fmt.Errorf("failed to update MFA chain on client %s: %w (also failed to roll back previous MFA chains: %v)", clientID, err, rollbackErr)
|
||||
}
|
||||
return fmt.Errorf("failed to update MFA chain on client %s: %w", clientID, err)
|
||||
}
|
||||
updatedClientIDs = append(updatedClientIDs, clientID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func rollbackClientsMFAChain(ctx context.Context, st storage.Storage, clientIDs []string, previousChains map[string][]string) error {
|
||||
var rollbackErrs []error
|
||||
for i := len(clientIDs) - 1; i >= 0; i-- {
|
||||
clientID := clientIDs[i]
|
||||
previousChain := cloneMFAChain(previousChains[clientID])
|
||||
if err := st.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
|
||||
old.MFAChain = previousChain
|
||||
return old, nil
|
||||
}); err != nil {
|
||||
rollbackErrs = append(rollbackErrs, fmt.Errorf("client %s: %w", clientID, err))
|
||||
}
|
||||
}
|
||||
return errors.Join(rollbackErrs...)
|
||||
}
|
||||
|
||||
func cloneMFAChain(chain []string) []string {
|
||||
if chain == nil {
|
||||
return nil
|
||||
}
|
||||
return append([]string(nil), chain...)
|
||||
}
|
||||
|
||||
// SetClientsMFAChain updates the MFAChain field on the dashboard and CLI OAuth2 clients.
|
||||
// Pass a non-empty slice (e.g. []string{"default-totp"}) to enable MFA, or nil to disable it.
|
||||
func (p *Provider) SetClientsMFAChain(ctx context.Context, clientIDs []string, mfaChain []string) error {
|
||||
for _, clientID := range clientIDs {
|
||||
if err := p.storage.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
|
||||
old.MFAChain = mfaChain
|
||||
return old, nil
|
||||
}); err != nil {
|
||||
return fmt.Errorf("failed to update MFA chain on client %s: %w", clientID, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
return SetClientsMFAChain(ctx, p.storage, clientIDs, mfaChain)
|
||||
}
|
||||
|
||||
// Handler returns the Dex server as an http.Handler for embedding in another server.
|
||||
@@ -546,7 +587,7 @@ func (p *Provider) CreateUser(ctx context.Context, email, username, password str
|
||||
|
||||
// Encode the user ID in Dex's format: base64(protobuf{user_id, connector_id})
|
||||
// This matches the format Dex uses in JWT tokens
|
||||
encodedID := EncodeDexUserID(userID, localConnectorID)
|
||||
encodedID := EncodeDexUserID(userID, LocalConnectorID)
|
||||
return encodedID, nil
|
||||
}
|
||||
|
||||
@@ -625,7 +666,7 @@ func DecodeDexUserID(encodedID string) (userID, connectorID string, err error) {
|
||||
// local password connector.
|
||||
func IsLocalUserID(encodedID string) bool {
|
||||
_, connectorID, err := DecodeDexUserID(encodedID)
|
||||
return err == nil && connectorID == localConnectorID
|
||||
return err == nil && connectorID == LocalConnectorID
|
||||
}
|
||||
|
||||
// GetUser returns a user by email
|
||||
|
||||
@@ -3,6 +3,8 @@ package dex
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -11,11 +13,44 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/dexidp/dex/storage"
|
||||
"github.com/dexidp/dex/storage/memory"
|
||||
sqllib "github.com/dexidp/dex/storage/sql"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type updateFailingStorage struct {
|
||||
storage.Storage
|
||||
failClientID string
|
||||
}
|
||||
|
||||
func (s *updateFailingStorage) UpdateClient(ctx context.Context, id string, updater func(storage.Client) (storage.Client, error)) error {
|
||||
if id == s.failClientID {
|
||||
return errors.New("forced update failure")
|
||||
}
|
||||
return s.Storage.UpdateClient(ctx, id, updater)
|
||||
}
|
||||
|
||||
func TestSetClientsMFAChainRollsBackUpdatedClients(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
|
||||
require.NoError(t, st.CreateClient(ctx, storage.Client{ID: "client-1", MFAChain: []string{"old-1"}}))
|
||||
require.NoError(t, st.CreateClient(ctx, storage.Client{ID: "client-2", MFAChain: []string{"old-2"}}))
|
||||
|
||||
err := SetClientsMFAChain(ctx, &updateFailingStorage{Storage: st, failClientID: "client-2"}, []string{"client-1", "client-2"}, []string{"new"})
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "failed to update MFA chain on client client-2")
|
||||
|
||||
client1, err := st.GetClient(ctx, "client-1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"old-1"}, client1.MFAChain)
|
||||
|
||||
client2, err := st.GetClient(ctx, "client-2")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"old-2"}, client2.MFAChain)
|
||||
}
|
||||
|
||||
func TestUserCreationFlow(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
|
||||
15
idp/dex/sqlite_cgo.go
Normal file
15
idp/dex/sqlite_cgo.go
Normal file
@@ -0,0 +1,15 @@
|
||||
//go:build cgo
|
||||
|
||||
package dex
|
||||
|
||||
import (
|
||||
sql "github.com/dexidp/dex/storage/sql"
|
||||
)
|
||||
|
||||
// newSQLite3 builds the dex SQLite3 config. CGO builds use the upstream
|
||||
// struct that takes a File path. Non-CGO builds get an empty stub whose
|
||||
// Open() returns the dex "SQLite not available" error — correct behaviour
|
||||
// for binaries that can't link sqlite3 (e.g. cross-compiled ARM targets).
|
||||
func newSQLite3(file string) *sql.SQLite3 {
|
||||
return &sql.SQLite3{File: file}
|
||||
}
|
||||
15
idp/dex/sqlite_nocgo.go
Normal file
15
idp/dex/sqlite_nocgo.go
Normal file
@@ -0,0 +1,15 @@
|
||||
//go:build !cgo
|
||||
|
||||
package dex
|
||||
|
||||
import (
|
||||
sql "github.com/dexidp/dex/storage/sql"
|
||||
)
|
||||
|
||||
// newSQLite3 for non-CGO builds. The dex SQLite3 stub has no fields and its
|
||||
// Open() returns an error documenting the missing CGO support — correct
|
||||
// behaviour for cross-compiled artefacts that never actually run the
|
||||
// embedded IdP. The `file` argument is ignored.
|
||||
func newSQLite3(_ string) *sql.SQLite3 {
|
||||
return &sql.SQLite3{}
|
||||
}
|
||||
@@ -234,9 +234,6 @@ 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)
|
||||
@@ -263,10 +260,6 @@ 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
|
||||
|
||||
@@ -556,7 +556,7 @@ start_services_and_show_instructions() {
|
||||
echo "Creating proxy access token..."
|
||||
# Use docker exec with bash to run the token command directly
|
||||
PROXY_TOKEN=$($DOCKER_COMPOSE_COMMAND exec -T netbird-server \
|
||||
/go/bin/netbird-server token create --name "default-proxy" --config /etc/netbird/config.yaml 2>/dev/null | grep "^Token:" | awk '{print $2}')
|
||||
/go/bin/netbird-server admin token create --name "default-proxy" --config /etc/netbird/config.yaml 2>/dev/null | grep "^Token:" | awk '{print $2}')
|
||||
|
||||
if [[ -z "$PROXY_TOKEN" ]]; then
|
||||
echo "ERROR: Failed to create proxy token. Check netbird-server logs." > /dev/stderr
|
||||
|
||||
@@ -490,8 +490,6 @@ 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 ""
|
||||
@@ -588,11 +586,6 @@ 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
|
||||
|
||||
177
management/cmd/admin.go
Normal file
177
management/cmd/admin.go
Normal file
@@ -0,0 +1,177 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/dexidp/dex/storage"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
admincmd "github.com/netbirdio/netbird/management/cmd/admin"
|
||||
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server/activity"
|
||||
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
var adminDatadir string
|
||||
|
||||
// newAdminCommands creates the admin command tree with management-specific resource openers.
|
||||
func newAdminCommands() *cobra.Command {
|
||||
cmd := admincmd.NewCommands(admincmd.Openers{
|
||||
Resources: withAdminResources,
|
||||
Store: withAdminStoreOnly,
|
||||
IDP: withAdminIDPOnly,
|
||||
})
|
||||
cmd.PersistentFlags().StringVar(&adminDatadir, "datadir", "", "Override the data directory from config (used for store.db and the default idp.db)")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func newLegacyTokenCommand() *cobra.Command {
|
||||
cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
|
||||
cmd.Deprecated = "use 'admin token' instead"
|
||||
cmd.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
|
||||
return cmd
|
||||
}
|
||||
|
||||
// withAdminResources initializes logging, loads config, opens the management store
|
||||
// and embedded IdP storage, and calls fn.
|
||||
func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
|
||||
return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
|
||||
managementStore, err := openAdminStore(ctx, config, datadir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseStore(ctx, managementStore)
|
||||
|
||||
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseIDPStorage(idpStorage)
|
||||
|
||||
eventStore, esErr := openAdminEventStore(ctx, config, datadir)
|
||||
if esErr != nil {
|
||||
_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
|
||||
}
|
||||
if eventStore != nil {
|
||||
defer func() {
|
||||
if err := eventStore.Close(ctx); err != nil {
|
||||
log.Debugf("close activity event store: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
|
||||
})
|
||||
}
|
||||
|
||||
// withAdminStoreOnly opens only the management store for admin subcommands that do not
|
||||
// need embedded IdP storage.
|
||||
func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
|
||||
return withAdminConfig(cmd, false, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
|
||||
managementStore, err := openAdminStore(ctx, config, datadir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseStore(ctx, managementStore)
|
||||
|
||||
return fn(ctx, managementStore)
|
||||
})
|
||||
}
|
||||
|
||||
func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
|
||||
return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, _ string) error {
|
||||
idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer admincmd.CloseIDPStorage(idpStorage)
|
||||
|
||||
return fn(ctx, idpStorage, idpStorageFile)
|
||||
})
|
||||
}
|
||||
|
||||
func withAdminConfig(cmd *cobra.Command, applyIDPDefaults bool, fn func(ctx context.Context, config *nbconfig.Config, datadir string) error) error {
|
||||
if err := util.InitLog("error", "console"); err != nil {
|
||||
return fmt.Errorf("init log: %w", err)
|
||||
}
|
||||
|
||||
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
|
||||
|
||||
config, datadir, err := loadAdminMgmtConfig(ctx, applyIDPDefaults)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load config: %w", err)
|
||||
}
|
||||
|
||||
return fn(ctx, config, datadir)
|
||||
}
|
||||
|
||||
func loadAdminMgmtConfig(ctx context.Context, applyIDPDefaults bool) (*nbconfig.Config, string, error) {
|
||||
config := &nbconfig.Config{}
|
||||
if _, err := util.ReadJsonWithEnvSub(nbconfig.MgmtConfigPath, config); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if applyIDPDefaults {
|
||||
if err := ApplyEmbeddedIdPConfig(ctx, config); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
|
||||
datadir := config.Datadir
|
||||
applyAdminDatadirOverride(config, &datadir)
|
||||
return config, datadir, nil
|
||||
}
|
||||
|
||||
func applyAdminDatadirOverride(config *nbconfig.Config, datadir *string) {
|
||||
if adminDatadir == "" {
|
||||
return
|
||||
}
|
||||
|
||||
oldDatadir := *datadir
|
||||
*datadir = adminDatadir
|
||||
if config.EmbeddedIdP != nil && config.EmbeddedIdP.Storage.Type == "sqlite3" && isDefaultIDPStorageFile(config.EmbeddedIdP.Storage.Config.File, oldDatadir) {
|
||||
config.EmbeddedIdP.Storage.Config.File = filepath.Join(*datadir, "idp.db")
|
||||
}
|
||||
}
|
||||
|
||||
func isDefaultIDPStorageFile(file, datadir string) bool {
|
||||
if file == "" {
|
||||
return true
|
||||
}
|
||||
defaultFile := filepath.Join(datadir, "idp.db")
|
||||
legacyDefaultFile := path.Join(datadir, "idp.db")
|
||||
legacySlashDefaultFile := path.Join(filepath.ToSlash(datadir), "idp.db")
|
||||
return filepath.Clean(file) == filepath.Clean(defaultFile) ||
|
||||
file == legacyDefaultFile ||
|
||||
filepath.ToSlash(file) == legacySlashDefaultFile
|
||||
}
|
||||
|
||||
func openAdminStore(ctx context.Context, config *nbconfig.Config, datadir string) (store.Store, error) {
|
||||
managementStore, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create store: %w", err)
|
||||
}
|
||||
return managementStore, nil
|
||||
}
|
||||
|
||||
func openAdminEventStore(ctx context.Context, config *nbconfig.Config, datadir string) (activity.Store, error) {
|
||||
if config.DataStoreEncryptionKey == "" {
|
||||
return nil, fmt.Errorf("data store encryption key is not configured")
|
||||
}
|
||||
eventStore, err := activitystore.NewSqlStore(ctx, datadir, config.DataStoreEncryptionKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open activity event store: %w", err)
|
||||
}
|
||||
if eventStore == nil {
|
||||
return nil, fmt.Errorf("open activity event store: returned nil store")
|
||||
}
|
||||
return eventStore, nil
|
||||
}
|
||||
577
management/cmd/admin/admin.go
Normal file
577
management/cmd/admin/admin.go
Normal file
@@ -0,0 +1,577 @@
|
||||
// Package admincmd provides reusable cobra commands for self-hosted administrator helpers.
|
||||
// Both the management and combined binaries use these commands, each providing
|
||||
// their own opener to handle config loading and storage initialization.
|
||||
package admincmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/dexidp/dex/storage"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
nbdex "github.com/netbirdio/netbird/idp/dex"
|
||||
"github.com/netbirdio/netbird/management/cmd/proxy"
|
||||
"github.com/netbirdio/netbird/management/cmd/token"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server"
|
||||
"github.com/netbirdio/netbird/management/server/activity"
|
||||
"github.com/netbirdio/netbird/management/server/idp"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
// Resources contains the storages required by the admin commands.
|
||||
type Resources struct {
|
||||
Store store.Store
|
||||
IDPStorage storage.Storage
|
||||
IDPStorageFile string
|
||||
EventStore activity.Store
|
||||
}
|
||||
|
||||
// Opener initializes command resources from the command context and calls fn.
|
||||
type Opener func(cmd *cobra.Command, fn func(ctx context.Context, resources Resources) error) error
|
||||
|
||||
// StoreOpener initializes only the management store from the command context and calls fn.
|
||||
type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
|
||||
|
||||
// IDPOpener initializes only the embedded IdP storage from the command context and calls fn.
|
||||
type IDPOpener func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error
|
||||
|
||||
// Openers contains the resource openers needed by the admin command tree.
|
||||
type Openers struct {
|
||||
Resources Opener
|
||||
Store StoreOpener
|
||||
IDP IDPOpener
|
||||
}
|
||||
|
||||
type userSelector struct {
|
||||
email string
|
||||
userID string
|
||||
}
|
||||
|
||||
func (s userSelector) normalized() userSelector {
|
||||
return userSelector{
|
||||
email: strings.TrimSpace(s.email),
|
||||
userID: strings.TrimSpace(s.userID),
|
||||
}
|
||||
}
|
||||
|
||||
func (s userSelector) validate() error {
|
||||
s = s.normalized()
|
||||
if (s.email == "") == (s.userID == "") {
|
||||
return fmt.Errorf("provide exactly one of --email or --user-id")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewCommands creates the admin command tree with the given resource openers.
|
||||
func NewCommands(openers Openers) *cobra.Command {
|
||||
adminCmd := &cobra.Command{
|
||||
Use: "admin",
|
||||
Short: "Self-hosted administrator helpers",
|
||||
Long: "Administrative helpers for self-hosted deployments using the embedded identity provider.",
|
||||
}
|
||||
|
||||
userCmd := &cobra.Command{
|
||||
Use: "user",
|
||||
Short: "Manage local embedded IdP users",
|
||||
}
|
||||
|
||||
var passwordSelector userSelector
|
||||
var password string
|
||||
var passwordFile string
|
||||
passwordCmd := &cobra.Command{
|
||||
Use: "change-password (--email email | --user-id id) (--password password | --password-file path)",
|
||||
Aliases: []string{"set-password"},
|
||||
Short: "Change a local user's password",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if err := passwordSelector.validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
newPassword, err := resolvePasswordInput(cmd, password, passwordFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
|
||||
return runChangePassword(ctx, idpStorage, cmd.OutOrStdout(), passwordSelector, newPassword, storageFile)
|
||||
})
|
||||
},
|
||||
}
|
||||
addUserSelectorFlags(passwordCmd, &passwordSelector)
|
||||
passwordCmd.Flags().StringVar(&password, "password", "", "New password for the user")
|
||||
passwordCmd.Flags().StringVar(&passwordFile, "password-file", "", "Read new password from file ('-' for stdin)")
|
||||
|
||||
var resetSelector userSelector
|
||||
resetMFACmd := &cobra.Command{
|
||||
Use: "reset-mfa (--email email | --user-id id)",
|
||||
Short: "Reset a local user's MFA enrollment",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
if err := resetSelector.validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
|
||||
return runResetMFA(ctx, idpStorage, cmd.OutOrStdout(), resetSelector, storageFile)
|
||||
})
|
||||
},
|
||||
}
|
||||
addUserSelectorFlags(resetMFACmd, &resetSelector)
|
||||
|
||||
userCmd.AddCommand(passwordCmd, resetMFACmd)
|
||||
|
||||
mfaCmd := &cobra.Command{
|
||||
Use: "mfa",
|
||||
Short: "Manage local MFA for embedded IdP users",
|
||||
}
|
||||
|
||||
enableCmd := &cobra.Command{
|
||||
Use: "enable",
|
||||
Short: "Enable MFA for local embedded IdP users",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
|
||||
return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), true)
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
disableCmd := &cobra.Command{
|
||||
Use: "disable",
|
||||
Short: "Disable MFA for local embedded IdP users",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
|
||||
return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), false)
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
statusCmd := &cobra.Command{
|
||||
Use: "status",
|
||||
Short: "Show local MFA status",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
|
||||
return runMFAStatus(ctx, resources, cmd.OutOrStdout())
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
mfaCmd.AddCommand(enableCmd, disableCmd, statusCmd)
|
||||
adminCmd.AddCommand(userCmd, mfaCmd)
|
||||
if openers.Store != nil {
|
||||
adminCmd.AddCommand(tokencmd.NewCommands(tokencmd.StoreOpener(openers.Store)))
|
||||
adminCmd.AddCommand(proxycmd.NewCommands(proxycmd.StoreOpener(openers.Store)))
|
||||
}
|
||||
return adminCmd
|
||||
}
|
||||
|
||||
// OpenEmbeddedIDPStorage opens the Dex storage configured for the embedded IdP.
|
||||
func OpenEmbeddedIDPStorage(cfg *idp.EmbeddedIdPConfig) (storage.Storage, error) {
|
||||
if cfg == nil || !cfg.Enabled {
|
||||
return nil, fmt.Errorf("admin commands require the embedded IdP to be enabled")
|
||||
}
|
||||
|
||||
yamlConfig, err := cfg.ToYAMLConfig()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build embedded IdP config: %w", err)
|
||||
}
|
||||
|
||||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
st, err := yamlConfig.Storage.OpenStorage(logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open embedded IdP storage: %w", err)
|
||||
}
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// CloseStore closes the management store and logs cleanup errors at debug level.
|
||||
func CloseStore(ctx context.Context, s store.Store) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if err := s.Close(ctx); err != nil {
|
||||
log.Debugf("close store: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// OpenIDPStorage opens embedded IdP storage and returns its sqlite file path when applicable.
|
||||
func OpenIDPStorage(config *nbconfig.Config) (storage.Storage, string, error) {
|
||||
if config == nil {
|
||||
return nil, "", fmt.Errorf("management config is required")
|
||||
}
|
||||
idpStorage, err := OpenEmbeddedIDPStorage(config.EmbeddedIdP)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return idpStorage, embeddedIDPStorageFile(config), nil
|
||||
}
|
||||
|
||||
func embeddedIDPStorageFile(config *nbconfig.Config) string {
|
||||
if config.EmbeddedIdP == nil || config.EmbeddedIdP.Storage.Type != "sqlite3" {
|
||||
return ""
|
||||
}
|
||||
return config.EmbeddedIdP.Storage.Config.File
|
||||
}
|
||||
|
||||
// CloseIDPStorage closes embedded IdP storage and logs cleanup errors at debug level.
|
||||
func CloseIDPStorage(s storage.Storage) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if err := s.Close(); err != nil {
|
||||
log.Debugf("close embedded IdP storage: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func addUserSelectorFlags(cmd *cobra.Command, selector *userSelector) {
|
||||
cmd.Flags().StringVar(&selector.email, "email", "", "User email")
|
||||
cmd.Flags().StringVar(&selector.userID, "user-id", "", "User ID")
|
||||
}
|
||||
|
||||
func resolvePasswordInput(cmd *cobra.Command, password, passwordFile string) (string, error) {
|
||||
if password != "" && passwordFile != "" {
|
||||
return "", fmt.Errorf("provide only one of --password or --password-file")
|
||||
}
|
||||
if passwordFile == "" {
|
||||
return password, nil
|
||||
}
|
||||
|
||||
var data []byte
|
||||
var err error
|
||||
if passwordFile == "-" {
|
||||
data, err = io.ReadAll(cmd.InOrStdin())
|
||||
} else {
|
||||
data, err = os.ReadFile(passwordFile)
|
||||
}
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read password: %w", err)
|
||||
}
|
||||
return strings.TrimRight(string(data), "\r\n"), nil
|
||||
}
|
||||
|
||||
func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string, idpStorageFile string) error {
|
||||
if idpStorage == nil {
|
||||
return fmt.Errorf("embedded IdP storage is required")
|
||||
}
|
||||
selector = selector.normalized()
|
||||
if err := selector.validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if password == "" {
|
||||
return fmt.Errorf("password is required")
|
||||
}
|
||||
if err := server.ValidatePassword(password); err != nil {
|
||||
return fmt.Errorf("invalid password: %w", err)
|
||||
}
|
||||
|
||||
user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
|
||||
if err != nil {
|
||||
return fmt.Errorf("hash password: %w", err)
|
||||
}
|
||||
|
||||
if err := idpStorage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
|
||||
old.Hash = hash
|
||||
return old, nil
|
||||
}); err != nil {
|
||||
return fmt.Errorf("update password for %s: %w", user.Email, err)
|
||||
}
|
||||
|
||||
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(w, "Password updated for %s.\n", user.Email)
|
||||
return nil
|
||||
}
|
||||
|
||||
func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, idpStorageFile string) error {
|
||||
if idpStorage == nil {
|
||||
return fmt.Errorf("embedded IdP storage is required")
|
||||
}
|
||||
selector = selector.normalized()
|
||||
if err := selector.validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
reset := false
|
||||
err = idpStorage.UpdateUserIdentity(ctx, user.UserID, idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
|
||||
reset = reset || len(old.MFASecrets) > 0 || len(old.WebAuthnCredentials) > 0
|
||||
old.MFASecrets = map[string]*storage.MFASecret{}
|
||||
old.WebAuthnCredentials = map[string][]storage.WebAuthnCredential{}
|
||||
return old, nil
|
||||
})
|
||||
if errors.Is(err, storage.ErrNotFound) {
|
||||
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("reset MFA for %s: %w", user.Email, err)
|
||||
}
|
||||
|
||||
if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if reset {
|
||||
_, _ = fmt.Fprintf(w, "MFA reset for %s. The user will re-enroll at next login.\n", user.Email)
|
||||
} else {
|
||||
_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func runSetMFAEnabled(ctx context.Context, resources Resources, w io.Writer, enabled bool) error {
|
||||
if resources.Store == nil {
|
||||
return fmt.Errorf("management store is required")
|
||||
}
|
||||
if resources.IDPStorage == nil {
|
||||
return fmt.Errorf("embedded IdP storage is required")
|
||||
}
|
||||
|
||||
accountID, settings, err := getSingleAccountSettings(ctx, resources.Store)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
oldEnabled := settings.LocalMfaEnabled
|
||||
newSettings := settings.Copy()
|
||||
newSettings.LocalMfaEnabled = enabled
|
||||
|
||||
if err := setIDPClientsMFA(ctx, resources.IDPStorage, enabled); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := resources.Store.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
|
||||
if rollbackErr := setIDPClientsMFA(ctx, resources.IDPStorage, oldEnabled); rollbackErr != nil {
|
||||
return fmt.Errorf("save local MFA account setting: %w (also failed to roll back embedded IdP MFA state: %v)", err, rollbackErr)
|
||||
}
|
||||
return fmt.Errorf("save local MFA account setting: %w", err)
|
||||
}
|
||||
|
||||
if err := storeMFAActivity(ctx, resources.EventStore, accountID, enabled); err != nil {
|
||||
_, _ = fmt.Fprintf(w, "Warning: failed to record audit event: %v\n", err)
|
||||
}
|
||||
|
||||
state := "disabled"
|
||||
if enabled {
|
||||
state = "enabled"
|
||||
}
|
||||
_, _ = fmt.Fprintf(w, "Local MFA %s.\n", state)
|
||||
return nil
|
||||
}
|
||||
|
||||
func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
|
||||
if resources.Store == nil {
|
||||
return fmt.Errorf("management store is required")
|
||||
}
|
||||
if resources.IDPStorage == nil {
|
||||
return fmt.Errorf("embedded IdP storage is required")
|
||||
}
|
||||
|
||||
_, settings, err := getSingleAccountSettings(ctx, resources.Store)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
accountStatus := "disabled"
|
||||
if settings.LocalMfaEnabled {
|
||||
accountStatus = "enabled"
|
||||
}
|
||||
|
||||
clientStatus, err := idpClientsMFAStatus(ctx, resources.IDPStorage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(w, "Account setting: %s\n", accountStatus)
|
||||
_, _ = fmt.Fprintf(w, "Embedded IdP clients: %s\n", clientStatus)
|
||||
return nil
|
||||
}
|
||||
|
||||
func getSingleAccountSettings(ctx context.Context, s store.Store) (string, *types.Settings, error) {
|
||||
count, err := s.GetAccountsCounter(ctx)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("count accounts: %w", err)
|
||||
}
|
||||
if count != 1 {
|
||||
return "", nil, fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", count)
|
||||
}
|
||||
|
||||
accountID, err := s.GetAnyAccountID(ctx)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("get account ID: %w", err)
|
||||
}
|
||||
|
||||
settings, err := s.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("get account settings: %w", err)
|
||||
}
|
||||
if settings == nil {
|
||||
settings = &types.Settings{}
|
||||
}
|
||||
return accountID, settings, nil
|
||||
}
|
||||
|
||||
func storeMFAActivity(ctx context.Context, eventStore activity.Store, accountID string, enabled bool) error {
|
||||
if eventStore == nil {
|
||||
return nil
|
||||
}
|
||||
event := activity.AccountLocalMfaDisabled
|
||||
if enabled {
|
||||
event = activity.AccountLocalMfaEnabled
|
||||
}
|
||||
_, err := eventStore.Save(ctx, &activity.Event{
|
||||
Timestamp: time.Now().UTC(),
|
||||
Activity: event,
|
||||
InitiatorID: string(hook.SystemSource),
|
||||
TargetID: accountID,
|
||||
AccountID: accountID,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("save local MFA audit event: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector, idpStorageFile string) (storage.Password, error) {
|
||||
selector = selector.normalized()
|
||||
if err := selector.validate(); err != nil {
|
||||
return storage.Password{}, err
|
||||
}
|
||||
|
||||
if selector.email != "" {
|
||||
user, err := idpStorage.GetPassword(ctx, selector.email)
|
||||
if errors.Is(err, storage.ErrNotFound) {
|
||||
if empty, listErr := localUsersEmpty(ctx, idpStorage); listErr != nil {
|
||||
return storage.Password{}, listErr
|
||||
} else if empty {
|
||||
return storage.Password{}, noLocalUsersError(idpStorageFile)
|
||||
}
|
||||
return storage.Password{}, fmt.Errorf("local user with email %q not found", selector.email)
|
||||
}
|
||||
if err != nil {
|
||||
return storage.Password{}, fmt.Errorf("get local user by email %q: %w", selector.email, err)
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
rawUserID := selector.userID
|
||||
if decodedUserID, _, err := nbdex.DecodeDexUserID(selector.userID); err == nil && decodedUserID != "" {
|
||||
rawUserID = decodedUserID
|
||||
}
|
||||
|
||||
users, err := idpStorage.ListPasswords(ctx)
|
||||
if err != nil {
|
||||
return storage.Password{}, fmt.Errorf("list local users: %w", err)
|
||||
}
|
||||
for _, user := range users {
|
||||
if user.UserID == rawUserID || user.UserID == selector.userID {
|
||||
return user, nil
|
||||
}
|
||||
}
|
||||
|
||||
if len(users) == 0 {
|
||||
return storage.Password{}, noLocalUsersError(idpStorageFile)
|
||||
}
|
||||
|
||||
return storage.Password{}, fmt.Errorf("local user with ID %q not found", selector.userID)
|
||||
}
|
||||
|
||||
func localUsersEmpty(ctx context.Context, idpStorage storage.Storage) (bool, error) {
|
||||
users, err := idpStorage.ListPasswords(ctx)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("list local users: %w", err)
|
||||
}
|
||||
return len(users) == 0, nil
|
||||
}
|
||||
|
||||
func noLocalUsersError(idpStorageFile string) error {
|
||||
location := ""
|
||||
if idpStorageFile != "" {
|
||||
location = fmt.Sprintf(" (%s)", idpStorageFile)
|
||||
}
|
||||
return fmt.Errorf("no local users exist in the embedded IdP storage%s; the management server may never have started with this config, or --datadir points at the wrong location", location)
|
||||
}
|
||||
|
||||
func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, userID string) error {
|
||||
err := idpStorage.DeleteAuthSession(ctx, userID, idp.LocalConnectorID)
|
||||
if err == nil || errors.Is(err, storage.ErrNotFound) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("delete local auth session for user %s: %w", userID, err)
|
||||
}
|
||||
|
||||
func setIDPClientsMFA(ctx context.Context, idpStorage storage.Storage, enabled bool) error {
|
||||
var mfaChain []string
|
||||
if enabled {
|
||||
mfaChain = []string{idp.DefaultTOTPAuthenticatorID}
|
||||
}
|
||||
|
||||
clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
|
||||
if err := nbdex.SetClientsMFAChain(ctx, idpStorage, clientIDs, mfaChain); err != nil {
|
||||
if errors.Is(err, storage.ErrNotFound) {
|
||||
return fmt.Errorf("embedded IdP client not found; start the management server once before toggling MFA: %w", err)
|
||||
}
|
||||
return fmt.Errorf("update MFA chain on embedded IdP clients: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (string, error) {
|
||||
clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
|
||||
enabledCount := 0
|
||||
for _, clientID := range clientIDs {
|
||||
client, err := idpStorage.GetClient(ctx, clientID)
|
||||
if errors.Is(err, storage.ErrNotFound) {
|
||||
return "unknown", fmt.Errorf("embedded IdP client %q not found", clientID)
|
||||
}
|
||||
if err != nil {
|
||||
return "unknown", fmt.Errorf("get embedded IdP client %q: %w", clientID, err)
|
||||
}
|
||||
if hasAuthenticator(client.MFAChain, idp.DefaultTOTPAuthenticatorID) {
|
||||
enabledCount++
|
||||
}
|
||||
}
|
||||
|
||||
switch enabledCount {
|
||||
case 0:
|
||||
return "disabled", nil
|
||||
case len(clientIDs):
|
||||
return "enabled", nil
|
||||
default:
|
||||
return "partially enabled", nil
|
||||
}
|
||||
}
|
||||
|
||||
func hasAuthenticator(chain []string, authenticatorID string) bool {
|
||||
for _, id := range chain {
|
||||
if id == authenticatorID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
250
management/cmd/admin/admin_test.go
Normal file
250
management/cmd/admin/admin_test.go
Normal file
@@ -0,0 +1,250 @@
|
||||
package admincmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/dexidp/dex/storage"
|
||||
"github.com/dexidp/dex/storage/memory"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
|
||||
nbdex "github.com/netbirdio/netbird/idp/dex"
|
||||
"github.com/netbirdio/netbird/management/server/idp"
|
||||
mgmtstore "github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
func newTestIDPStorage(t *testing.T) storage.Storage {
|
||||
t.Helper()
|
||||
|
||||
st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte("OldPass1!"), bcrypt.DefaultCost)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, st.CreatePassword(context.Background(), storage.Password{
|
||||
Email: "user@example.com",
|
||||
Username: "User",
|
||||
UserID: "user-1",
|
||||
Hash: hash,
|
||||
}))
|
||||
require.NoError(t, st.CreateUserIdentity(context.Background(), storage.UserIdentity{
|
||||
UserID: "user-1",
|
||||
ConnectorID: idp.LocalConnectorID,
|
||||
MFASecrets: map[string]*storage.MFASecret{
|
||||
idp.DefaultTOTPAuthenticatorID: {
|
||||
AuthenticatorID: idp.DefaultTOTPAuthenticatorID,
|
||||
Type: "TOTP",
|
||||
Secret: "otpauth://totp/NetBird:user@example.com?secret=ABC",
|
||||
Confirmed: true,
|
||||
CreatedAt: time.Now(),
|
||||
},
|
||||
},
|
||||
WebAuthnCredentials: map[string][]storage.WebAuthnCredential{
|
||||
"webauthn": {{CredentialID: []byte("credential")}},
|
||||
},
|
||||
}))
|
||||
require.NoError(t, st.CreateAuthSession(context.Background(), storage.AuthSession{
|
||||
UserID: "user-1",
|
||||
ConnectorID: idp.LocalConnectorID,
|
||||
Nonce: "nonce",
|
||||
}))
|
||||
require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientCLI, Name: "CLI"}))
|
||||
require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientDashboard, Name: "Dashboard"}))
|
||||
|
||||
return st
|
||||
}
|
||||
|
||||
func TestRunChangePassword(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st := newTestIDPStorage(t)
|
||||
var out bytes.Buffer
|
||||
|
||||
err := runChangePassword(ctx, st, &out, userSelector{email: "user@example.com"}, "NewPass1!", "")
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, out.String(), "Password updated")
|
||||
|
||||
user, err := st.GetPassword(ctx, "user@example.com")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, bcrypt.CompareHashAndPassword(user.Hash, []byte("NewPass1!")))
|
||||
|
||||
_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
|
||||
require.ErrorIs(t, err, storage.ErrNotFound)
|
||||
}
|
||||
|
||||
func TestRunChangePasswordValidatesPassword(t *testing.T) {
|
||||
st := newTestIDPStorage(t)
|
||||
err := runChangePassword(context.Background(), st, io.Discard, userSelector{email: "user@example.com"}, "short", "")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "invalid password")
|
||||
}
|
||||
|
||||
func TestRunResetMFA(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st := newTestIDPStorage(t)
|
||||
var out bytes.Buffer
|
||||
|
||||
encodedUserID := nbdex.EncodeDexUserID("user-1", idp.LocalConnectorID)
|
||||
err := runResetMFA(ctx, st, &out, userSelector{userID: encodedUserID}, "")
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, out.String(), "MFA reset")
|
||||
|
||||
identity, err := st.GetUserIdentity(ctx, "user-1", idp.LocalConnectorID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, identity.MFASecrets)
|
||||
require.Empty(t, identity.WebAuthnCredentials)
|
||||
|
||||
_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
|
||||
require.ErrorIs(t, err, storage.ErrNotFound)
|
||||
}
|
||||
|
||||
func TestRunResetMFAWithoutEnrollment(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st := newTestIDPStorage(t)
|
||||
require.NoError(t, st.UpdateUserIdentity(ctx, "user-1", idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
|
||||
old.MFASecrets = nil
|
||||
old.WebAuthnCredentials = nil
|
||||
return old, nil
|
||||
}))
|
||||
|
||||
var out bytes.Buffer
|
||||
err := runResetMFA(ctx, st, &out, userSelector{email: "user@example.com"}, "")
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, out.String(), "No MFA enrollment found")
|
||||
}
|
||||
|
||||
func TestSetIDPClientsMFA(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
st := newTestIDPStorage(t)
|
||||
|
||||
require.NoError(t, setIDPClientsMFA(ctx, st, true))
|
||||
status, err := idpClientsMFAStatus(ctx, st)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "enabled", status)
|
||||
|
||||
require.NoError(t, setIDPClientsMFA(ctx, st, false))
|
||||
status, err = idpClientsMFAStatus(ctx, st)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "disabled", status)
|
||||
}
|
||||
|
||||
func newTestManagementStore(t *testing.T, localMFAEnabled bool) mgmtstore.Store {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
st, err := mgmtstore.NewStore(ctx, types.SqliteStoreEngine, t.TempDir(), nil, false)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, st.Close(ctx)) })
|
||||
require.NoError(t, st.SaveAccount(ctx, &types.Account{
|
||||
Id: "account-1",
|
||||
Settings: &types.Settings{LocalMfaEnabled: localMFAEnabled},
|
||||
}))
|
||||
return st
|
||||
}
|
||||
|
||||
func TestRunSetMFAEnabledDoesNotSaveWhenIDPUpdateFails(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
managementStore := newTestManagementStore(t, false)
|
||||
idpStorage := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
|
||||
err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "embedded IdP client")
|
||||
|
||||
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
|
||||
require.NoError(t, err)
|
||||
require.False(t, settings.LocalMfaEnabled)
|
||||
}
|
||||
|
||||
func TestRunSetMFAEnabledUpdatesSettingsAfterIDP(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
managementStore := newTestManagementStore(t, false)
|
||||
idpStorage := newTestIDPStorage(t)
|
||||
|
||||
err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
|
||||
require.NoError(t, err)
|
||||
|
||||
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
|
||||
require.NoError(t, err)
|
||||
require.True(t, settings.LocalMfaEnabled)
|
||||
clientStatus, err := idpClientsMFAStatus(ctx, idpStorage)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "enabled", clientStatus)
|
||||
}
|
||||
|
||||
func TestRunSetMFAEnabledSucceedsWithNilEventStore(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
managementStore := newTestManagementStore(t, false)
|
||||
idpStorage := newTestIDPStorage(t)
|
||||
var out bytes.Buffer
|
||||
var err error
|
||||
|
||||
require.NotPanics(t, func() {
|
||||
err = runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage, EventStore: nil}, &out, true)
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, out.String(), "Local MFA enabled")
|
||||
|
||||
settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
|
||||
require.NoError(t, err)
|
||||
require.True(t, settings.LocalMfaEnabled)
|
||||
}
|
||||
|
||||
func TestUserSelectorValidate(t *testing.T) {
|
||||
require.NoError(t, userSelector{email: " user@example.com "}.validate())
|
||||
require.NoError(t, userSelector{userID: "user-1"}.validate())
|
||||
require.Error(t, userSelector{}.validate())
|
||||
require.Error(t, userSelector{email: "user@example.com", userID: "user-1"}.validate())
|
||||
}
|
||||
|
||||
func TestFindLocalUserNotFound(t *testing.T) {
|
||||
st := newTestIDPStorage(t)
|
||||
_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "")
|
||||
require.Error(t, err)
|
||||
require.True(t, strings.Contains(err.Error(), "not found"))
|
||||
}
|
||||
|
||||
func TestFindLocalUserZeroUsersIncludesStoragePath(t *testing.T) {
|
||||
st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "/var/lib/netbird/idp.db")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "no local users exist")
|
||||
require.Contains(t, err.Error(), "/var/lib/netbird/idp.db")
|
||||
}
|
||||
|
||||
func TestUserCommandValidatesSelectorBeforeOpeningStorage(t *testing.T) {
|
||||
opened := false
|
||||
cmd := NewCommands(Openers{
|
||||
IDP: func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
|
||||
opened = true
|
||||
return nil
|
||||
},
|
||||
})
|
||||
cmd.SetArgs([]string{"user", "change-password", "--password", "NewPass1!"})
|
||||
cmd.SetOut(io.Discard)
|
||||
cmd.SetErr(io.Discard)
|
||||
|
||||
err := cmd.Execute()
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "provide exactly one")
|
||||
require.False(t, opened)
|
||||
}
|
||||
|
||||
func TestResolvePasswordInputFromStdin(t *testing.T) {
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetIn(strings.NewReader("NewPass1!\n"))
|
||||
|
||||
password, err := resolvePasswordInput(cmd, "", "-")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "NewPass1!", password)
|
||||
}
|
||||
|
||||
func TestResolvePasswordInputRejectsMultipleSources(t *testing.T) {
|
||||
_, err := resolvePasswordInput(&cobra.Command{}, "NewPass1!", "-")
|
||||
require.Error(t, err)
|
||||
}
|
||||
80
management/cmd/admin_config_test.go
Normal file
80
management/cmd/admin_config_test.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server/idp"
|
||||
)
|
||||
|
||||
func TestApplyAdminDatadirOverrideRelocatesDefaultIDPStorage(t *testing.T) {
|
||||
oldDatadir := filepath.Join(t.TempDir(), "old")
|
||||
newDatadir := filepath.Join(t.TempDir(), "new")
|
||||
|
||||
for _, defaultFile := range []string{
|
||||
"",
|
||||
filepath.Join(oldDatadir, "idp.db"),
|
||||
path.Join(oldDatadir, "idp.db"),
|
||||
} {
|
||||
t.Run(defaultFile, func(t *testing.T) {
|
||||
cfg := &nbconfig.Config{
|
||||
EmbeddedIdP: &idp.EmbeddedIdPConfig{
|
||||
Enabled: true,
|
||||
Storage: idp.EmbeddedStorageConfig{
|
||||
Type: "sqlite3",
|
||||
Config: idp.EmbeddedStorageTypeConfig{
|
||||
File: defaultFile,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
datadir := oldDatadir
|
||||
oldAdminDatadir := adminDatadir
|
||||
adminDatadir = newDatadir
|
||||
t.Cleanup(func() { adminDatadir = oldAdminDatadir })
|
||||
|
||||
applyAdminDatadirOverride(cfg, &datadir)
|
||||
|
||||
require.Equal(t, newDatadir, datadir)
|
||||
require.Equal(t, filepath.Join(newDatadir, "idp.db"), cfg.EmbeddedIdP.Storage.Config.File)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
|
||||
eventStore, err := openAdminEventStore(context.Background(), &nbconfig.Config{}, t.TempDir())
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "encryption key")
|
||||
require.Nil(t, eventStore)
|
||||
}
|
||||
|
||||
func TestApplyAdminDatadirOverrideKeepsExplicitIDPStorage(t *testing.T) {
|
||||
oldDatadir := filepath.Join(t.TempDir(), "old")
|
||||
newDatadir := filepath.Join(t.TempDir(), "new")
|
||||
explicitFile := filepath.Join(t.TempDir(), "custom-idp.db")
|
||||
cfg := &nbconfig.Config{
|
||||
EmbeddedIdP: &idp.EmbeddedIdPConfig{
|
||||
Enabled: true,
|
||||
Storage: idp.EmbeddedStorageConfig{
|
||||
Type: "sqlite3",
|
||||
Config: idp.EmbeddedStorageTypeConfig{
|
||||
File: explicitFile,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
datadir := oldDatadir
|
||||
oldAdminDatadir := adminDatadir
|
||||
adminDatadir = newDatadir
|
||||
t.Cleanup(func() { adminDatadir = oldAdminDatadir })
|
||||
|
||||
applyAdminDatadirOverride(cfg, &datadir)
|
||||
|
||||
require.Equal(t, newDatadir, datadir)
|
||||
require.Equal(t, explicitFile, cfg.EmbeddedIdP.Storage.Config.File)
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"os"
|
||||
"os/signal"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
@@ -25,6 +26,7 @@ import (
|
||||
"github.com/netbirdio/netbird/management/internals/server"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
nbdomain "github.com/netbirdio/netbird/shared/management/domain"
|
||||
"github.com/netbirdio/netbird/shared/management/grpc"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
"github.com/netbirdio/netbird/util/crypt"
|
||||
)
|
||||
@@ -153,8 +155,20 @@ func LoadMgmtConfig(ctx context.Context, mgmtConfigPath string) (*nbconfig.Confi
|
||||
|
||||
ApplyCommandLineOverrides(loadedConfig)
|
||||
|
||||
err := grpc.ValidateSyncMessageVersion(loadedConfig.HighestSupportedSyncMessageVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for account, version := range loadedConfig.PerAccountHighestSupportedSyncMessageVersion {
|
||||
err := grpc.ValidateSyncMessageVersion(&version)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unrecognized sync message version for account %s, %w", account, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Apply EmbeddedIdP config to HttpConfig if embedded IdP is enabled
|
||||
err := ApplyEmbeddedIdPConfig(ctx, loadedConfig)
|
||||
err = ApplyEmbeddedIdPConfig(ctx, loadedConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -209,7 +223,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config) error {
|
||||
cfg.EmbeddedIdP.Storage.Type = "sqlite3"
|
||||
}
|
||||
if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
|
||||
cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
|
||||
cfg.EmbeddedIdP.Storage.Config.File = filepath.Join(cfg.Datadir, "idp.db")
|
||||
}
|
||||
|
||||
issuer := cfg.EmbeddedIdP.Issuer
|
||||
|
||||
@@ -4,6 +4,9 @@ import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/grpc"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -20,34 +23,49 @@ const (
|
||||
"AuthAudience": "https://stageapp/",
|
||||
"AuthIssuer": "https://something.eu.auth0.com/",
|
||||
"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
|
||||
},
|
||||
"HighestSupportedSyncMessageVersion": 1,
|
||||
"PerAccountHighestSupportedSyncMessageVersion": {
|
||||
"1": 0,
|
||||
"2": 1
|
||||
}
|
||||
}`
|
||||
)
|
||||
|
||||
func Test_loadMgmtConfig(t *testing.T) {
|
||||
tmpFile, err := createConfig()
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create config: %s", err)
|
||||
}
|
||||
func Test_LoadMgmtConfig(t *testing.T) {
|
||||
tmpFile, err := createConfig(exampleConfig)
|
||||
assert.NoError(t, err)
|
||||
|
||||
cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to load management config: %s", err)
|
||||
}
|
||||
if cfg.Relay == nil {
|
||||
t.Fatalf("config is nil")
|
||||
}
|
||||
if len(cfg.Relay.Addresses) == 0 {
|
||||
t.Fatalf("relay address is empty")
|
||||
}
|
||||
assert.NoError(t, err)
|
||||
assert.NotEmpty(t, cfg.Relay)
|
||||
assert.NotEmpty(t, cfg.Relay.Addresses)
|
||||
assert.Equal(t, int(grpc.ComponentNetworkMap), *cfg.HighestSupportedSyncMessageVersion)
|
||||
assert.Equal(t, map[string]int{"1": int(grpc.Base), "2": int(grpc.ComponentNetworkMap)}, cfg.PerAccountHighestSupportedSyncMessageVersion)
|
||||
}
|
||||
|
||||
func createConfig() (string, error) {
|
||||
func Test_LoadMgmtConfig_Empty(t *testing.T) {
|
||||
tmpFile, err := createConfig(`{
|
||||
"HttpConfig": {
|
||||
"AuthAudience": "https://stageapp/",
|
||||
"AuthIssuer": "https://something.eu.auth0.com/",
|
||||
"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
|
||||
}
|
||||
}`)
|
||||
assert.NoError(t, err)
|
||||
|
||||
cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
|
||||
assert.NoError(t, err)
|
||||
assert.Nil(t, cfg.HighestSupportedSyncMessageVersion)
|
||||
assert.Nil(t, cfg.PerAccountHighestSupportedSyncMessageVersion)
|
||||
}
|
||||
|
||||
func createConfig(config string) (string, error) {
|
||||
tmpfile, err := os.CreateTemp("", "config.json")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
_, err = tmpfile.Write([]byte(exampleConfig))
|
||||
_, err = tmpfile.Write([]byte(config))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
141
management/cmd/proxy/proxy.go
Normal file
141
management/cmd/proxy/proxy.go
Normal file
@@ -0,0 +1,141 @@
|
||||
// Package proxycmd provides reusable cobra commands for managing reverse proxy instances.
|
||||
// Both the management and combined binaries use these commands, each providing
|
||||
// their own StoreOpener to handle config loading and store initialization.
|
||||
package proxycmd
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
// StoreOpener initializes a store from the command context and calls fn.
|
||||
type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
|
||||
|
||||
const disconnectAllConfirmation = "disconnect all proxies"
|
||||
|
||||
// NewCommands creates the proxy command tree with the given store opener.
|
||||
// Returns the parent "proxy" command with the disconnect-all subcommand.
|
||||
func NewCommands(opener StoreOpener) *cobra.Command {
|
||||
var dryRun bool
|
||||
var force bool
|
||||
|
||||
proxyCmd := &cobra.Command{
|
||||
Use: "proxy",
|
||||
Short: "Manage reverse proxy instances",
|
||||
Long: "Commands for inspecting and repairing the reverse proxy instances registered with the management server.",
|
||||
}
|
||||
|
||||
disconnectAllCmd := &cobra.Command{
|
||||
Use: "disconnect-all",
|
||||
Short: "Force-mark all reverse proxy instances as disconnected",
|
||||
Long: "Lists all reverse proxy instances and force-marks them as disconnected, regardless of their session state. " +
|
||||
"Use this to repair stale connection state, e.g. after an unclean management server shutdown. " +
|
||||
"By default, it asks for manual confirmation before changing state. Use --dry-run to preview without changing state, or --force to skip confirmation. " +
|
||||
"Run during a maintenance window; affected live proxies may stay hidden until their next heartbeat or reconnect/re-register.",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
return opener(cmd, func(ctx context.Context, s store.Store) error {
|
||||
return runDisconnectAll(ctx, s, cmd.OutOrStdout(), cmd.InOrStdin(), dryRun, force)
|
||||
})
|
||||
},
|
||||
}
|
||||
disconnectAllCmd.Flags().BoolVar(&dryRun, "dry-run", false, "List reverse proxy instances that would be disconnected without changing state")
|
||||
disconnectAllCmd.Flags().BoolVar(&force, "force", false, "Skip the confirmation prompt and apply the repair")
|
||||
|
||||
proxyCmd.AddCommand(disconnectAllCmd)
|
||||
return proxyCmd
|
||||
}
|
||||
|
||||
func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.Reader, dryRun, force bool) error {
|
||||
proxies, err := s.GetAllProxies(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("list proxies: %w", err)
|
||||
}
|
||||
|
||||
if len(proxies) == 0 {
|
||||
_, _ = fmt.Fprintln(out, "No reverse proxy instances found.")
|
||||
return nil
|
||||
}
|
||||
|
||||
toDisconnect := 0
|
||||
w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
|
||||
_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tACCOUNT\tSTATUS\tLAST SEEN")
|
||||
_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t------\t---------")
|
||||
|
||||
for _, p := range proxies {
|
||||
if p.Status != rpproxy.StatusDisconnected {
|
||||
toDisconnect++
|
||||
}
|
||||
|
||||
account := "-"
|
||||
if p.AccountID != nil {
|
||||
account = *p.AccountID
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
|
||||
p.ID,
|
||||
p.ClusterAddress,
|
||||
p.IPAddress,
|
||||
account,
|
||||
p.Status,
|
||||
p.LastSeen.Format("2006-01-02 15:04:05"),
|
||||
)
|
||||
}
|
||||
if err := w.Flush(); err != nil {
|
||||
return fmt.Errorf("write proxy list: %w", err)
|
||||
}
|
||||
|
||||
if dryRun {
|
||||
_, _ = fmt.Fprintf(out, "\nDry run: would force-mark %d of %d reverse proxy instance(s) as disconnected.\n", toDisconnect, len(proxies))
|
||||
return nil
|
||||
}
|
||||
|
||||
if !force {
|
||||
confirmed, err := confirmDisconnectAll(out, in)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !confirmed {
|
||||
_, _ = fmt.Fprintln(out, "Aborted. No reverse proxy instances were changed.")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
disconnected, err := s.DisconnectAllProxies(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("disconnect proxies: %w", err)
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintf(out, "\nForce-marked %d of %d reverse proxy instance(s) as disconnected.\n", disconnected, len(proxies))
|
||||
return nil
|
||||
}
|
||||
|
||||
func confirmDisconnectAll(out io.Writer, in io.Reader) (bool, error) {
|
||||
if in == nil {
|
||||
in = strings.NewReader("")
|
||||
}
|
||||
|
||||
_, _ = fmt.Fprintln(out, "\nWARNING: This command changes stored reverse proxy state for every non-disconnected instance.")
|
||||
_, _ = fmt.Fprintln(out, "Run it during a maintenance window; affected live proxies may stay hidden until "+
|
||||
"their next heartbeat or reconnect/re-register.")
|
||||
_, _ = fmt.Fprintf(out, "Type %q to continue: ", disconnectAllConfirmation)
|
||||
|
||||
scanner := bufio.NewScanner(in)
|
||||
if !scanner.Scan() {
|
||||
if err := scanner.Err(); err != nil {
|
||||
return false, fmt.Errorf("read confirmation: %w", err)
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
return strings.EqualFold(strings.TrimSpace(scanner.Text()), disconnectAllConfirmation), nil
|
||||
}
|
||||
180
management/cmd/proxy/proxy_test.go
Normal file
180
management/cmd/proxy/proxy_test.go
Normal file
@@ -0,0 +1,180 @@
|
||||
package proxycmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
func newTestStore(t *testing.T) store.Store {
|
||||
t.Helper()
|
||||
|
||||
s, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(cleanup)
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
func seedProxies(t *testing.T, ctx context.Context, s store.Store) {
|
||||
t.Helper()
|
||||
|
||||
accountID := "account-1"
|
||||
alreadyDisconnectedAt := time.Now().Add(-time.Hour)
|
||||
seed := []*rpproxy.Proxy{
|
||||
{
|
||||
ID: "proxy-1",
|
||||
SessionID: "session-1",
|
||||
ClusterAddress: "cluster-a.example.com",
|
||||
IPAddress: "10.0.0.1",
|
||||
LastSeen: time.Now(),
|
||||
Status: rpproxy.StatusConnected,
|
||||
},
|
||||
{
|
||||
ID: "proxy-2",
|
||||
SessionID: "session-2",
|
||||
ClusterAddress: "cluster-b.example.com",
|
||||
IPAddress: "10.0.0.2",
|
||||
AccountID: &accountID,
|
||||
LastSeen: time.Now(),
|
||||
Status: rpproxy.StatusConnected,
|
||||
},
|
||||
{
|
||||
ID: "proxy-3",
|
||||
SessionID: "session-3",
|
||||
ClusterAddress: "cluster-a.example.com",
|
||||
IPAddress: "10.0.0.3",
|
||||
LastSeen: time.Now().Add(-time.Hour),
|
||||
Status: rpproxy.StatusDisconnected,
|
||||
DisconnectedAt: &alreadyDisconnectedAt,
|
||||
},
|
||||
}
|
||||
for _, p := range seed {
|
||||
require.NoError(t, s.SaveProxy(ctx, p))
|
||||
}
|
||||
}
|
||||
|
||||
func proxiesByID(t *testing.T, ctx context.Context, s store.Store) map[string]*rpproxy.Proxy {
|
||||
t.Helper()
|
||||
|
||||
proxies, err := s.GetAllProxies(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, proxies, 3)
|
||||
|
||||
byID := make(map[string]*rpproxy.Proxy, len(proxies))
|
||||
for _, p := range proxies {
|
||||
byID[p.ID] = p
|
||||
}
|
||||
return byID
|
||||
}
|
||||
|
||||
func TestRunDisconnectAllWithConfirmation(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
seedProxies(t, ctx, s)
|
||||
|
||||
var out bytes.Buffer
|
||||
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(disconnectAllConfirmation+"\n"), false, false))
|
||||
|
||||
output := out.String()
|
||||
require.Contains(t, output, "proxy-1")
|
||||
require.Contains(t, output, "proxy-2")
|
||||
require.Contains(t, output, "proxy-3")
|
||||
require.Contains(t, output, "cluster-a.example.com")
|
||||
require.Contains(t, output, "account-1")
|
||||
require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
|
||||
require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
|
||||
|
||||
for _, p := range proxiesByID(t, ctx, s) {
|
||||
require.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", p.ID)
|
||||
require.NotNil(t, p.DisconnectedAt, "proxy %s should have a disconnected timestamp", p.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunDisconnectAllForceSkipsConfirmation(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
seedProxies(t, ctx, s)
|
||||
|
||||
var out bytes.Buffer
|
||||
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, true))
|
||||
|
||||
output := out.String()
|
||||
require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
|
||||
require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
|
||||
}
|
||||
|
||||
func TestRunDisconnectAllAbortLeavesProxiesUnchanged(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
seedProxies(t, ctx, s)
|
||||
|
||||
var out bytes.Buffer
|
||||
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader("no\n"), false, false))
|
||||
|
||||
output := out.String()
|
||||
require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
|
||||
require.Contains(t, output, "Aborted. No reverse proxy instances were changed.")
|
||||
|
||||
byID := proxiesByID(t, ctx, s)
|
||||
require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
|
||||
require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
|
||||
require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
|
||||
}
|
||||
|
||||
func TestRunDisconnectAllDryRunLeavesProxiesUnchanged(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
seedProxies(t, ctx, s)
|
||||
|
||||
var out bytes.Buffer
|
||||
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), true, false))
|
||||
|
||||
output := out.String()
|
||||
require.Contains(t, output, "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
|
||||
require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
|
||||
|
||||
byID := proxiesByID(t, ctx, s)
|
||||
require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
|
||||
require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
|
||||
require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
|
||||
}
|
||||
|
||||
func TestNewCommandsDisconnectAllDryRun(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
seedProxies(t, ctx, s)
|
||||
|
||||
opened := false
|
||||
cmd := NewCommands(func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
|
||||
opened = true
|
||||
return fn(cmd.Context(), s)
|
||||
})
|
||||
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&out)
|
||||
cmd.SetIn(strings.NewReader(""))
|
||||
cmd.SetArgs([]string{"disconnect-all", "--dry-run"})
|
||||
|
||||
require.NoError(t, cmd.ExecuteContext(ctx))
|
||||
require.True(t, opened)
|
||||
require.Contains(t, out.String(), "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
|
||||
}
|
||||
|
||||
func TestRunDisconnectAllEmpty(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
var out bytes.Buffer
|
||||
require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, false))
|
||||
require.Contains(t, out.String(), "No reverse proxy instances found.")
|
||||
}
|
||||
@@ -83,7 +83,8 @@ func init() {
|
||||
|
||||
rootCmd.AddCommand(migrationCmd)
|
||||
|
||||
tc := newTokenCommands()
|
||||
tc.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
|
||||
rootCmd.AddCommand(tc)
|
||||
ac := newAdminCommands()
|
||||
ac.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
|
||||
rootCmd.AddCommand(ac)
|
||||
rootCmd.AddCommand(newLegacyTokenCommand())
|
||||
}
|
||||
|
||||
@@ -1,55 +0,0 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/netbirdio/netbird/formatter/hook"
|
||||
tokencmd "github.com/netbirdio/netbird/management/cmd/token"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
)
|
||||
|
||||
var tokenDatadir string
|
||||
|
||||
// newTokenCommands creates the token command tree with management-specific store opener.
|
||||
func newTokenCommands() *cobra.Command {
|
||||
cmd := tokencmd.NewCommands(withTokenStore)
|
||||
cmd.PersistentFlags().StringVar(&tokenDatadir, "datadir", "", "Override the data directory from config (where store.db is located)")
|
||||
return cmd
|
||||
}
|
||||
|
||||
// withTokenStore initializes logging, loads config, opens the store, and calls fn.
|
||||
func withTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
|
||||
if err := util.InitLog("error", "console"); err != nil {
|
||||
return fmt.Errorf("init log: %w", err)
|
||||
}
|
||||
|
||||
ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
|
||||
|
||||
config, err := LoadMgmtConfig(ctx, nbconfig.MgmtConfigPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load config: %w", err)
|
||||
}
|
||||
|
||||
datadir := config.Datadir
|
||||
if tokenDatadir != "" {
|
||||
datadir = tokenDatadir
|
||||
}
|
||||
|
||||
s, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create store: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := s.Close(ctx); err != nil {
|
||||
log.Debugf("close store: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
return fn(ctx, s)
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/management/server/telemetry"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
"github.com/netbirdio/netbird/util"
|
||||
@@ -56,6 +57,10 @@ type Controller struct {
|
||||
proxyController port_forwarding.Controller
|
||||
|
||||
integratedPeerValidator integrated_validator.IntegratedValidator
|
||||
|
||||
serverSupportedSyncMessageVersion sharedgrpc.SyncMessageVersion
|
||||
|
||||
perAccountServerSupportedSyncMessageVersions map[string]sharedgrpc.SyncMessageVersion
|
||||
}
|
||||
|
||||
type bufferUpdate struct {
|
||||
@@ -90,8 +95,10 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
|
||||
dnsDomain: dnsDomain,
|
||||
config: config,
|
||||
|
||||
proxyController: proxyController,
|
||||
EphemeralPeersManager: ephemeralPeersManager,
|
||||
proxyController: proxyController,
|
||||
EphemeralPeersManager: ephemeralPeersManager,
|
||||
serverSupportedSyncMessageVersion: sharedgrpc.SyncMessageVersionFromConfig(config.HighestSupportedSyncMessageVersion),
|
||||
perAccountServerSupportedSyncMessageVersions: sharedgrpc.SyncMessageVersionsFromMap(config.PerAccountHighestSupportedSyncMessageVersion),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -222,18 +229,53 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
|
||||
c.metrics.CountCalcPostureChecksDuration(time.Since(start))
|
||||
start = time.Now()
|
||||
|
||||
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
peerGroups := account.GetPeerGroups(p.ID)
|
||||
proxyNetworkMap := proxyNetworkMaps[p.ID]
|
||||
var update *proto.SyncResponse
|
||||
|
||||
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
|
||||
c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
|
||||
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
|
||||
|
||||
log.WithContext(ctx).
|
||||
WithFields(log.Fields{
|
||||
"sync_message_version": commonSyncMessageVersion,
|
||||
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
|
||||
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
|
||||
}).Debug("common highest sync message version")
|
||||
|
||||
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
|
||||
components := account.GetPeerNetworkMapComponents(
|
||||
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
|
||||
|
||||
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
|
||||
|
||||
start = time.Now()
|
||||
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
|
||||
// the client merges it into Calculate()'s output the same
|
||||
// way the legacy server did via NetworkMap.Merge.
|
||||
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
|
||||
Update: update,
|
||||
MessageType: network_map.MessageTypeNetworkMap,
|
||||
})
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
nmap := account.GetPeerNetworkMapFromComponents(
|
||||
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
|
||||
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
|
||||
|
||||
proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
|
||||
if ok {
|
||||
remotePeerNetworkMap.Merge(proxyNetworkMap)
|
||||
if proxyNetworkMap != nil {
|
||||
nmap.Merge(proxyNetworkMap)
|
||||
}
|
||||
|
||||
peerGroups := account.GetPeerGroups(p.ID)
|
||||
start = time.Now()
|
||||
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
c.metrics.CountToSyncResponseDuration(time.Since(start))
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
|
||||
@@ -251,6 +293,13 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Controller) perAccountOrGlobalSupportedSyncMessageVersions(accountId string) sharedgrpc.SyncMessageVersion {
|
||||
if perAccount, ok := c.perAccountServerSupportedSyncMessageVersions[accountId]; ok {
|
||||
return perAccount
|
||||
}
|
||||
return c.serverSupportedSyncMessageVersion
|
||||
}
|
||||
|
||||
// UpdatePeers updates all peers that belong to an account.
|
||||
// Should be called when changes have to be synced to peers.
|
||||
func (c *Controller) UpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error {
|
||||
@@ -352,18 +401,53 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
|
||||
c.metrics.CountCalcPostureChecksDuration(time.Since(start))
|
||||
start = time.Now()
|
||||
|
||||
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
peerGroups := account.GetPeerGroups(p.ID)
|
||||
proxyNetworkMap := proxyNetworkMaps[p.ID]
|
||||
var update *proto.SyncResponse
|
||||
|
||||
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
|
||||
c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
|
||||
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
|
||||
|
||||
log.WithContext(ctx).
|
||||
WithFields(log.Fields{
|
||||
"sync_message_version": commonSyncMessageVersion,
|
||||
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
|
||||
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
|
||||
}).Debug("common highest sync message version")
|
||||
|
||||
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
|
||||
components := account.GetPeerNetworkMapComponents(
|
||||
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
|
||||
|
||||
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
|
||||
|
||||
start = time.Now()
|
||||
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
|
||||
// the client merges it into Calculate()'s output the same
|
||||
// way the legacy server did via NetworkMap.Merge.
|
||||
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
|
||||
Update: update,
|
||||
MessageType: network_map.MessageTypeNetworkMap,
|
||||
})
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
nmap := account.GetPeerNetworkMapFromComponents(
|
||||
ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
|
||||
c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
|
||||
|
||||
proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
|
||||
if ok {
|
||||
remotePeerNetworkMap.Merge(proxyNetworkMap)
|
||||
if proxyNetworkMap != nil {
|
||||
nmap.Merge(proxyNetworkMap)
|
||||
}
|
||||
|
||||
peerGroups := account.GetPeerGroups(p.ID)
|
||||
start = time.Now()
|
||||
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
|
||||
c.metrics.CountToSyncResponseDuration(time.Since(start))
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
|
||||
@@ -451,13 +535,7 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
|
||||
return err
|
||||
}
|
||||
|
||||
remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, peerId, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
|
||||
proxyNetworkMap, ok := proxyNetworkMaps[peer.ID]
|
||||
if ok {
|
||||
remotePeerNetworkMap.Merge(proxyNetworkMap)
|
||||
}
|
||||
|
||||
proxyNetworkMap := proxyNetworkMaps[peer.ID]
|
||||
extraSettings, err := c.settingsManager.GetExtraSettings(ctx, peer.AccountID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get extra settings: %v", err)
|
||||
@@ -466,7 +544,45 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
|
||||
peerGroups := account.GetPeerGroups(peerId)
|
||||
dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
|
||||
|
||||
update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
|
||||
var update *proto.SyncResponse
|
||||
|
||||
commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
|
||||
c.perAccountOrGlobalSupportedSyncMessageVersions(accountId),
|
||||
sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
|
||||
|
||||
log.WithContext(ctx).
|
||||
WithFields(log.Fields{
|
||||
"sync_message_version": commonSyncMessageVersion,
|
||||
"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
|
||||
"peer_sync_message_version": sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
|
||||
}).Debug("common highest sync message version")
|
||||
|
||||
if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
|
||||
components := account.GetPeerNetworkMapComponents(
|
||||
ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
|
||||
|
||||
// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
|
||||
// the client merges it into Calculate()'s output the same
|
||||
// way the legacy server did via NetworkMap.Merge.
|
||||
update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
|
||||
Update: update,
|
||||
MessageType: network_map.MessageTypeNetworkMap,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
nmap := account.GetPeerNetworkMapFromComponents(
|
||||
ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
|
||||
|
||||
if proxyNetworkMap != nil {
|
||||
nmap.Merge(proxyNetworkMap)
|
||||
}
|
||||
|
||||
update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
|
||||
|
||||
c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
|
||||
Update: update,
|
||||
MessageType: network_map.MessageTypeNetworkMap,
|
||||
@@ -513,6 +629,65 @@ func (c *Controller) BufferUpdateAccountPeers(ctx context.Context, accountID str
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetValidatedPeerWithComponents is the components-format counterpart of
|
||||
// GetValidatedPeerWithMap. It returns raw NetworkMapComponents for capable
|
||||
// peers along with the proxy NetworkMap fragment (BYOP / port-forwarding
|
||||
// data the legacy server folds in via NetworkMap.Merge). The gRPC layer
|
||||
// encodes both into the wire envelope. Callers must gate on capability
|
||||
// themselves before dispatching here — this method does NOT branch on it.
|
||||
func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
|
||||
if isRequiresApproval {
|
||||
network, err := c.repo.GetAccountNetwork(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
return peer, &types.NetworkMapComponents{Network: network.Copy()}, nil, nil, 0, nil
|
||||
}
|
||||
|
||||
account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
|
||||
c.injectAllProxyPolicies(ctx, account)
|
||||
|
||||
approvedPeersMap, err := c.integratedPeerValidator.GetValidatedPeers(ctx, account.Id, maps.Values(account.Groups), maps.Values(account.Peers), account.Settings.Extra)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
|
||||
postureChecks, err := c.getPeerPostureChecks(account, peer.ID)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
|
||||
accountZones, err := c.repo.GetAccountZones(ctx, account.Id)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
|
||||
// Fetch the proxy network map fragment for this peer alongside the
|
||||
// components — same single-account-load path the streaming controller
|
||||
// uses, so initial-sync delivers BYOP/forwarding patches synchronously
|
||||
// instead of waiting for the next streaming push.
|
||||
proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peer.ID, account.Peers)
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
|
||||
dnsDomain := c.GetDNSDomain(account.Settings)
|
||||
peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
|
||||
|
||||
resourcePolicies := account.GetResourcePoliciesMap()
|
||||
routers := account.GetResourceRoutersMap()
|
||||
groupIDToUserIDs := account.GetActiveGroupUsers()
|
||||
components := account.GetPeerNetworkMapComponents(ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
|
||||
dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
|
||||
|
||||
return peer, components, proxyNetworkMaps[peer.ID], postureChecks, dnsFwdPort, nil
|
||||
}
|
||||
|
||||
// BufferUpdateAffectedPeers accumulates peer IDs and flushes them after the buffer interval.
|
||||
func (c *Controller) BufferUpdateAffectedPeers(ctx context.Context, accountID string, peerIDs []string, reason types.UpdateReason) error {
|
||||
if len(peerIDs) == 0 {
|
||||
|
||||
@@ -24,6 +24,7 @@ type Controller interface {
|
||||
UpdateAccountPeer(ctx context.Context, accountId string, peerId string) error
|
||||
BufferUpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error
|
||||
GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error)
|
||||
GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error)
|
||||
GetDNSDomain(settings *types.Settings) string
|
||||
StartWarmup(context.Context)
|
||||
GetNetworkMap(ctx context.Context, peerID string) (*types.NetworkMap, error)
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: management/internals/controllers/network_map/interface.go
|
||||
// Source: ./interface.go
|
||||
//
|
||||
// Generated by this command:
|
||||
//
|
||||
// mockgen -package network_map -destination=management/internals/controllers/network_map/interface_mock.go -source=management/internals/controllers/network_map/interface.go -build_flags=-mod=mod
|
||||
// mockgen -package network_map -destination=interface_mock.go -source=./interface.go -build_flags=-mod=mod
|
||||
//
|
||||
|
||||
// Package network_map is a generated GoMock package.
|
||||
@@ -126,8 +126,27 @@ func (mr *MockControllerMockRecorder) GetNetworkMap(ctx, peerID any) *gomock.Cal
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNetworkMap", reflect.TypeOf((*MockController)(nil).GetNetworkMap), ctx, peerID)
|
||||
}
|
||||
|
||||
// GetValidatedPeerWithComponents mocks base method.
|
||||
func (m *MockController) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *peer.Peer) (*peer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetValidatedPeerWithComponents", ctx, isRequiresApproval, accountID, p)
|
||||
ret0, _ := ret[0].(*peer.Peer)
|
||||
ret1, _ := ret[1].(*types.NetworkMapComponents)
|
||||
ret2, _ := ret[2].(*types.NetworkMap)
|
||||
ret3, _ := ret[3].([]*posture.Checks)
|
||||
ret4, _ := ret[4].(int64)
|
||||
ret5, _ := ret[5].(error)
|
||||
return ret0, ret1, ret2, ret3, ret4, ret5
|
||||
}
|
||||
|
||||
// GetValidatedPeerWithComponents indicates an expected call of GetValidatedPeerWithComponents.
|
||||
func (mr *MockControllerMockRecorder) GetValidatedPeerWithComponents(ctx, isRequiresApproval, accountID, p any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetValidatedPeerWithComponents", reflect.TypeOf((*MockController)(nil).GetValidatedPeerWithComponents), ctx, isRequiresApproval, accountID, p)
|
||||
}
|
||||
|
||||
// GetValidatedPeerWithMap mocks base method.
|
||||
func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
|
||||
func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetValidatedPeerWithMap", ctx, isRequiresApproval, accountID, peerID)
|
||||
ret0, _ := ret[0].(*types.NetworkMap)
|
||||
@@ -171,7 +190,7 @@ func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID
|
||||
}
|
||||
|
||||
// OnPeersAdded mocks base method.
|
||||
func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
|
||||
func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "OnPeersAdded", ctx, accountID, peerIDs, affectedPeerIDs)
|
||||
ret0, _ := ret[0].(error)
|
||||
@@ -185,7 +204,7 @@ func (mr *MockControllerMockRecorder) OnPeersAdded(ctx, accountID, peerIDs, affe
|
||||
}
|
||||
|
||||
// OnPeersDeleted mocks base method.
|
||||
func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
|
||||
func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "OnPeersDeleted", ctx, accountID, peerIDs, affectedPeerIDs)
|
||||
ret0, _ := ret[0].(error)
|
||||
@@ -199,7 +218,7 @@ func (mr *MockControllerMockRecorder) OnPeersDeleted(ctx, accountID, peerIDs, af
|
||||
}
|
||||
|
||||
// OnPeersUpdated mocks base method.
|
||||
func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs []string, affectedPeerIDs []string) error {
|
||||
func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs, affectedPeerIDs []string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "OnPeersUpdated", ctx, accountId, peerIDs, affectedPeerIDs)
|
||||
ret0, _ := ret[0].(error)
|
||||
|
||||
@@ -18,21 +18,26 @@ 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
|
||||
metaKeyCostUSDTotal = "cost.usd_total"
|
||||
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
|
||||
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"
|
||||
)
|
||||
|
||||
// IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
|
||||
@@ -108,20 +113,25 @@ 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),
|
||||
CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal),
|
||||
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),
|
||||
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],
|
||||
}
|
||||
|
||||
var groups []types.AgentNetworkAccessLogGroup
|
||||
@@ -140,18 +150,23 @@ 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,
|
||||
CostUSD: e.CostUSD,
|
||||
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,
|
||||
}
|
||||
|
||||
usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))
|
||||
|
||||
@@ -28,17 +28,22 @@ 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: "150",
|
||||
metaKeyCostUSDTotal: "0.0123",
|
||||
metaKeyStream: "true",
|
||||
metaKeyRequestPrompt: "hello",
|
||||
metaKeyResponseCompletion: "world",
|
||||
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",
|
||||
// repeated id must be de-duplicated before the group rows insert.
|
||||
metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
|
||||
},
|
||||
@@ -65,7 +70,19 @@ 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.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost 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")
|
||||
|
||||
logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
|
||||
require.NoError(t, err)
|
||||
@@ -96,6 +113,14 @@ 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")
|
||||
|
||||
@@ -38,7 +38,7 @@ func accessLogRow(id, sessionID string, ts time.Time, opts ...func(*types.AgentN
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
TotalTokens: 150,
|
||||
CostUSD: 0.01,
|
||||
InputCostUSD: 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.CostUSD = cost
|
||||
e.InputCostUSD = 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.CostUSD, 1e-9, "cost summed")
|
||||
assert.InDelta(t, 0.031, a.TotalCostUSD(), 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")
|
||||
|
||||
@@ -420,6 +420,47 @@ 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,
|
||||
|
||||
@@ -41,8 +41,24 @@ type AgentNetworkAccessLog struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
CostUSD float64
|
||||
Stream bool
|
||||
// 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
|
||||
|
||||
// Prompt capture. Only populated when prompt collection is enabled
|
||||
// (account master switch AND policy guardrail). Heavy free text.
|
||||
@@ -60,19 +76,44 @@ 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,
|
||||
CostUsd: a.CostUSD,
|
||||
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,
|
||||
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,
|
||||
}
|
||||
|
||||
out.UserId = strPtr(a.UserID)
|
||||
@@ -112,20 +153,36 @@ 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
|
||||
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
|
||||
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
|
||||
}
|
||||
|
||||
// sessionKey is the grouping key for an entry: its session id, or — when the
|
||||
@@ -205,7 +262,12 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
|
||||
sess.InputTokens += e.InputTokens
|
||||
sess.OutputTokens += e.OutputTokens
|
||||
sess.TotalTokens += e.TotalTokens
|
||||
sess.CostUSD += e.CostUSD
|
||||
sess.CachedInputTokens += e.CachedInputTokens
|
||||
sess.CacheCreationTokens += e.CacheCreationTokens
|
||||
sess.InputCostUSD += e.InputCostUSD
|
||||
sess.CachedInputCostUSD += e.CachedInputCostUSD
|
||||
sess.CacheCreationCostUSD += e.CacheCreationCostUSD
|
||||
sess.OutputCostUSD += e.OutputCostUSD
|
||||
if e.Timestamp.Before(sess.StartedAt) {
|
||||
sess.StartedAt = e.Timestamp
|
||||
}
|
||||
@@ -248,15 +310,22 @@ 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,
|
||||
CostUsd: sess.CostUSD,
|
||||
Decision: sess.Decision,
|
||||
Entries: entries,
|
||||
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,
|
||||
}
|
||||
out.SessionId = strPtr(sess.SessionID)
|
||||
out.UserId = strPtr(sess.UserID)
|
||||
|
||||
@@ -54,7 +54,7 @@ var accessLogSortFields = map[string]string{
|
||||
"provider": "provider",
|
||||
"status_code": "status_code",
|
||||
"duration": "duration",
|
||||
"cost_usd": "cost_usd",
|
||||
"cost_usd": CostUSDSQLExpr,
|
||||
"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(cost_usd)",
|
||||
"cost_usd": "SUM" + CostUSDSQLExpr,
|
||||
"total_tokens": "SUM(total_tokens)",
|
||||
"duration": "SUM(duration)",
|
||||
"request_count": "COUNT(*)",
|
||||
|
||||
124
management/internals/modules/agentnetwork/types/cost_test.go
Normal file
124
management/internals/modules/agentnetwork/types/cost_test.go
Normal file
@@ -0,0 +1,124 @@
|
||||
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")
|
||||
}
|
||||
@@ -25,8 +25,19 @@ type AgentNetworkUsage struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
CostUSD float64
|
||||
CreatedAt time.Time
|
||||
// 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
|
||||
}
|
||||
|
||||
// TableName keeps usage records in their own stripped table. Named
|
||||
@@ -34,6 +45,17 @@ 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.
|
||||
|
||||
@@ -33,21 +33,45 @@ 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
|
||||
CostUSD float64
|
||||
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
|
||||
}
|
||||
|
||||
// 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,
|
||||
CostUsd: b.CostUSD,
|
||||
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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -84,7 +108,12 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
|
||||
b.InputTokens += r.InputTokens
|
||||
b.OutputTokens += r.OutputTokens
|
||||
b.TotalTokens += r.TotalTokens
|
||||
b.CostUSD += r.CostUSD
|
||||
b.CachedInputTokens += r.CachedInputTokens
|
||||
b.CacheCreationTokens += r.CacheCreationTokens
|
||||
b.InputCostUSD += r.InputCostUSD
|
||||
b.CachedInputCostUSD += r.CachedInputCostUSD
|
||||
b.CacheCreationCostUSD += r.CacheCreationCostUSD
|
||||
b.OutputCostUSD += r.OutputCostUSD
|
||||
}
|
||||
|
||||
out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))
|
||||
|
||||
@@ -61,6 +61,10 @@ type Config struct {
|
||||
// EmbeddedIdP contains configuration for the embedded Dex OIDC provider.
|
||||
// When set, Dex will be embedded in the management server and serve requests at /oauth2/
|
||||
EmbeddedIdP *idp.EmbeddedIdPConfig
|
||||
|
||||
HighestSupportedSyncMessageVersion *int
|
||||
|
||||
PerAccountHighestSupportedSyncMessageVersion map[string]int
|
||||
}
|
||||
|
||||
// GetAuthAudiences returns the audience from the http config and device authorization flow config
|
||||
|
||||
766
management/internals/shared/grpc/components_encoder.go
Normal file
766
management/internals/shared/grpc/components_encoder.go
Normal file
@@ -0,0 +1,766 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"strconv"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// wgKeyRawLen is the raw byte length of a WireGuard public key.
|
||||
const wgKeyRawLen = 32
|
||||
|
||||
// ComponentsEnvelopeInput bundles the data the component-format encoder needs.
|
||||
// The envelope is fully self-contained — every field needed by the client's
|
||||
// local Calculate() comes from the components struct itself. The only
|
||||
// externally-supplied data is the receiving peer's PeerConfig (which is
|
||||
// computed alongside the components in the network_map controller and reused
|
||||
// from the legacy proto path) and the dns_domain string.
|
||||
type ComponentsEnvelopeInput struct {
|
||||
Components *types.NetworkMapComponents
|
||||
PeerConfig *proto.PeerConfig
|
||||
DNSDomain string
|
||||
DNSForwarderPort int64
|
||||
// UserIDClaim is the OIDC claim name the client should embed in
|
||||
// SshAuth.UserIDClaim when reconstructing the NetworkMap. Empty value
|
||||
// is OK — client treats empty as "no SshAuth to build".
|
||||
UserIDClaim string
|
||||
// ProxyPatch carries pre-expanded NetworkMap fragments injected by
|
||||
// external controllers (BYOP/port-forwarding). Nil when no proxy data
|
||||
// is present; encoder skips the field in that case.
|
||||
ProxyPatch *proto.ProxyPatch
|
||||
}
|
||||
|
||||
// EncodeNetworkMapEnvelope converts NetworkMapComponents into the component
|
||||
// wire envelope. The encoder is intentionally non-deterministic: it iterates
|
||||
// Go maps in their native (random) order. Indexes inside the envelope
|
||||
// (peer_indexes, source_group_ids, agent_version_idx, router_peer_indexes)
|
||||
// are self-consistent within a single encode, so the decoder reconstructs
|
||||
// the same typed objects regardless of emit order. Tests that need to
|
||||
// compare envelopes do so semantically via proto round-trip + canonicalize,
|
||||
// not byte-equal.
|
||||
//
|
||||
// Callers must NOT concatenate or merge envelopes from different encodes —
|
||||
// index spaces are local to a single envelope.
|
||||
func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvelope {
|
||||
c := in.Components
|
||||
|
||||
// Graceful degrade when components is nil — matches the legacy path's
|
||||
// behaviour for missing/unvalidated peers (return a NetworkMap with only
|
||||
// Network populated). The receiver gets an envelope it can decode
|
||||
// without crashing; AccountSettings stays non-nil so client-side
|
||||
// dereferences are safe.
|
||||
if c.IsEmpty() {
|
||||
// Match legacy missing-peer minimum: a NetworkMap with only Network
|
||||
// populated. The receiver gets enough to bootstrap (Network
|
||||
// identifier, dns_domain, account_settings) and the peer itself.
|
||||
return &proto.NetworkMapEnvelope{
|
||||
Payload: &proto.NetworkMapEnvelope_Full{
|
||||
Full: &proto.NetworkMapComponentsFull{
|
||||
PeerConfig: in.PeerConfig,
|
||||
// components.Peers always contains the target peer
|
||||
Peers: []*proto.PeerCompact{toPeerCompact(c.Peers[c.PeerID])},
|
||||
DnsDomain: in.DNSDomain,
|
||||
DnsForwarderPort: in.DNSForwarderPort,
|
||||
UserIdClaim: in.UserIDClaim,
|
||||
AccountSettings: &proto.AccountSettingsCompact{},
|
||||
ProxyPatch: in.ProxyPatch,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 1: build dedup tables. Every routing peer (in c.RouterPeers) and
|
||||
// every regular peer (in c.Peers) must be indexed before any encoder
|
||||
// looks up indexes via e.peerOrder — otherwise routes / routers_map for
|
||||
// peers that exist only in c.RouterPeers would silently lose their
|
||||
// peer_index reference.
|
||||
enc := newComponentEncoder(c)
|
||||
enc.indexAllPeers()
|
||||
routerIdxs := enc.indexRouterPeers(c.RouterPeers)
|
||||
|
||||
// Phase 2: gather every policy that any consumer references (peer-pair
|
||||
// policies + resource-only policies) so encodeResourcePoliciesMap can
|
||||
// translate every *Policy pointer to a wire index.
|
||||
allPolicies := unionPolicies(c.Policies, c.ResourcePoliciesMap)
|
||||
policies := enc.encodePolicies(allPolicies)
|
||||
|
||||
// Phase 3: emit. Order of struct field expressions no longer matters:
|
||||
// every encoder either reads from the dedup tables or works on
|
||||
// independent input.
|
||||
full := &proto.NetworkMapComponentsFull{
|
||||
Serial: networkSerial(c.Network),
|
||||
PeerConfig: in.PeerConfig,
|
||||
Network: toAccountNetwork(c.Network),
|
||||
AccountSettings: toAccountSettingsCompact(c.AccountSettings),
|
||||
DnsForwarderPort: in.DNSForwarderPort,
|
||||
UserIdClaim: in.UserIDClaim,
|
||||
ProxyPatch: in.ProxyPatch,
|
||||
DnsSettings: enc.encodeDNSSettings(c.DNSSettings),
|
||||
DnsDomain: in.DNSDomain,
|
||||
CustomZoneDomain: c.CustomZoneDomain,
|
||||
AgentVersions: enc.agentVersions,
|
||||
Peers: enc.peers,
|
||||
RouterPeerIndexes: routerIdxs,
|
||||
Policies: policies,
|
||||
Groups: enc.encodeGroups(),
|
||||
Routes: enc.encodeRoutes(c.Routes),
|
||||
NameserverGroups: enc.encodeNameServerGroups(c.NameServerGroups),
|
||||
AllDnsRecords: encodeSimpleRecords(c.AllDNSRecords),
|
||||
AccountZones: encodeCustomZones(c.AccountZones),
|
||||
NetworkResources: enc.encodeNetworkResources(c.NetworkResources),
|
||||
RoutersMap: enc.encodeRoutersMap(c.RoutersMap),
|
||||
ResourcePoliciesMap: enc.encodeResourcePoliciesMap(c.ResourcePoliciesMap),
|
||||
GroupIdToUserIds: enc.encodeGroupIDToUserIDs(c.GroupIDToUserIDs),
|
||||
AllowedUserIds: stringSetToSlice(c.AllowedUserIDs),
|
||||
PostureFailedPeers: enc.encodePostureFailedPeers(c.PostureFailedPeers),
|
||||
}
|
||||
|
||||
return &proto.NetworkMapEnvelope{
|
||||
Payload: &proto.NetworkMapEnvelope_Full{Full: full},
|
||||
}
|
||||
}
|
||||
|
||||
// networkSerial returns c.Network.CurrentSerial() with a nil guard. The
|
||||
// production path always populates c.Network, but the encoder is exported
|
||||
// and a hand-built components struct may omit it.
|
||||
func networkSerial(n *types.Network) uint64 {
|
||||
if n == nil {
|
||||
return 0
|
||||
}
|
||||
return n.CurrentSerial()
|
||||
}
|
||||
|
||||
type componentEncoder struct {
|
||||
components *types.NetworkMapComponents
|
||||
|
||||
peerOrder map[string]uint32
|
||||
peers []*proto.PeerCompact
|
||||
|
||||
agentVersionOrder map[string]uint32
|
||||
agentVersions []string
|
||||
}
|
||||
|
||||
func newComponentEncoder(c *types.NetworkMapComponents) *componentEncoder {
|
||||
return &componentEncoder{
|
||||
components: c,
|
||||
peerOrder: make(map[string]uint32, len(c.Peers)),
|
||||
peers: make([]*proto.PeerCompact, 0, len(c.Peers)),
|
||||
agentVersionOrder: make(map[string]uint32),
|
||||
}
|
||||
}
|
||||
|
||||
func (e *componentEncoder) indexAllPeers() {
|
||||
for _, p := range e.components.Peers {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
e.appendPeer(p)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
|
||||
if idx, ok := e.peerOrder[p.ID]; ok {
|
||||
return idx
|
||||
}
|
||||
idx := uint32(len(e.peers))
|
||||
e.peerOrder[p.ID] = idx
|
||||
e.peers = append(e.peers, toPeerCompact(p))
|
||||
return idx
|
||||
}
|
||||
|
||||
// indexRouterPeers ensures every router peer is in the peer dedup table
|
||||
// (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 {
|
||||
if len(routers) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]uint32, 0, len(routers))
|
||||
for _, p := range routers {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, e.appendPeer(p))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeGroups() []*proto.GroupCompact {
|
||||
if len(e.components.Groups) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
out := make([]*proto.GroupCompact, 0, len(e.components.Groups))
|
||||
for _, g := range e.components.Groups {
|
||||
peerIdxs := make([]uint32, 0, len(g.Peers))
|
||||
for _, peerID := range g.Peers {
|
||||
if idx, ok := e.peerOrder[peerID]; ok {
|
||||
peerIdxs = append(peerIdxs, idx)
|
||||
}
|
||||
}
|
||||
out = append(out, &proto.GroupCompact{
|
||||
Id: g.PublicID,
|
||||
PeerIndexes: peerIdxs,
|
||||
IsAll: g.IsGroupAll(),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// encodePolicies flattens Policy{Rules} → []PolicyCompact. Returns the wire
|
||||
// list and a map from policy pointer to the indexes of its emitted rules in
|
||||
// that list — used by encodeResourcePoliciesMap to translate
|
||||
// ResourcePoliciesMap[resourceID][]*Policy into wire-side indexes.
|
||||
func (e *componentEncoder) encodePolicies(policies []*types.Policy) []*proto.PolicyCompact {
|
||||
if len(policies) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
out := make([]*proto.PolicyCompact, 0, len(policies))
|
||||
|
||||
for _, pol := range policies {
|
||||
if !pol.Enabled {
|
||||
continue
|
||||
}
|
||||
for _, r := range pol.Rules {
|
||||
if r == nil || !r.Enabled {
|
||||
continue
|
||||
}
|
||||
out = append(out, e.encodePolicyRule(pol, r))
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// encodePolicyRule maps a single PolicyRule under pol to a PolicyCompact entry.
|
||||
func (e *componentEncoder) encodePolicyRule(pol *types.Policy, r *types.PolicyRule) *proto.PolicyCompact {
|
||||
return &proto.PolicyCompact{
|
||||
Id: pol.PublicID,
|
||||
Action: networkmap.GetProtoAction(string(r.Action)),
|
||||
Protocol: networkmap.GetProtoProtocol(string(r.Protocol)),
|
||||
Bidirectional: r.Bidirectional,
|
||||
Ports: portsToUint32(r.Ports),
|
||||
PortRanges: portRangesToProto(r.PortRanges),
|
||||
SourceGroupIds: e.groupPublicXids(r.Sources),
|
||||
DestinationGroupIds: e.groupPublicXids(r.Destinations),
|
||||
AuthorizedUser: r.AuthorizedUser,
|
||||
AuthorizedGroups: e.encodeAuthorizedGroups(r.AuthorizedGroups),
|
||||
SourceResource: e.resourceToProto(r.SourceResource),
|
||||
DestinationResource: e.resourceToProto(r.DestinationResource),
|
||||
SourcePostureCheckIds: e.postureCheckSeqs(pol.SourcePostureChecks),
|
||||
}
|
||||
}
|
||||
|
||||
// groupPublicXids maps the xid group IDs in src to their public xids,
|
||||
// dropping any group with invalid public xid.
|
||||
func (e *componentEncoder) groupPublicXids(src []string) []string {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(src))
|
||||
for _, gid := range src {
|
||||
if id, ok := e.groupPublicXid(gid); ok {
|
||||
out = append(out, id)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// unionPolicies merges c.Policies with every policy referenced by
|
||||
// c.ResourcePoliciesMap, deduplicating by pointer identity. Resource-only
|
||||
// policies (relevant to a NetworkResource but not to peer-pair traffic)
|
||||
// only live in ResourcePoliciesMap; without this union step they'd be lost
|
||||
// from the wire and the client's resource-policy lookup would come back
|
||||
// empty.
|
||||
func unionPolicies(policies []*types.Policy, resourcePolicies map[string][]*types.Policy) []*types.Policy {
|
||||
// Fast path: non-router peers have no resource-only policies, so the
|
||||
// "union" is identical to `policies`. Skip the dedup map allocation.
|
||||
if len(resourcePolicies) == 0 {
|
||||
return policies
|
||||
}
|
||||
seen := make(map[string]struct{}, len(policies))
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[p.ID]; ok {
|
||||
continue
|
||||
}
|
||||
seen[p.ID] = struct{}{}
|
||||
out = append(out, p)
|
||||
}
|
||||
for _, list := range resourcePolicies {
|
||||
for _, p := range list {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[p.ID]; ok {
|
||||
continue
|
||||
}
|
||||
seen[p.ID] = struct{}{}
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// encodeAuthorizedGroups translates rule.AuthorizedGroups (map keyed by
|
||||
// group xid → local-user names) to the wire form (map keyed by group
|
||||
// account_seq_id → UserNameList). Groups without a seq id are dropped —
|
||||
// matches how source/destination group references handle the same case.
|
||||
func (e *componentEncoder) encodeAuthorizedGroups(m map[string][]string) map[string]*proto.UserNameList {
|
||||
if len(m) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*proto.UserNameList, len(m))
|
||||
for groupID, names := range m {
|
||||
id, ok := e.groupPublicXid(groupID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
out[id] = &proto.UserNameList{Names: names}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) groupPublicXid(groupID string) (string, bool) {
|
||||
g, ok := e.components.Groups[groupID]
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
return g.PublicID, true
|
||||
}
|
||||
|
||||
// resourceToProto translates types.Resource for the wire. For peer-typed
|
||||
// resources the peer id is converted to a peer index into the envelope's
|
||||
// peers array. For other resource types only the type string is shipped
|
||||
// today (Calculate's resource-typed rule path consults SourceResource only
|
||||
// for "peer" — other types fall through to group-based lookup).
|
||||
func (e *componentEncoder) resourceToProto(r types.Resource) *proto.ResourceCompact {
|
||||
if r.ID == "" && r.Type == "" {
|
||||
return nil
|
||||
}
|
||||
out := &proto.ResourceCompact{Type: string(r.Type)}
|
||||
if r.Type == types.ResourceTypePeer && r.ID != "" {
|
||||
if idx, ok := e.peerOrder[r.ID]; ok {
|
||||
out.PeerIndexSet = true
|
||||
out.PeerIndex = idx
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// postureCheckSeqs translates a slice of posture-check xids to their
|
||||
// public xids. Unresolvable xids are silently dropped — matches how group/peer
|
||||
// references handle the same case.
|
||||
func (e *componentEncoder) postureCheckSeqs(xids []string) []string {
|
||||
if len(xids) == 0 || len(e.components.PostureCheckXIDToPublicID) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(xids))
|
||||
for _, xid := range xids {
|
||||
if seq, ok := e.components.PostureCheckXIDToPublicID[xid]; ok {
|
||||
out = append(out, seq)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// networkSeq translates a Network xid to its public id using
|
||||
// the NetworkMapComponents.NetworkXIDToPublicID lookup. Returns (0,false) when
|
||||
// the xid isn't known — callers decide whether to skip the parent record.
|
||||
func (e *componentEncoder) networkPublicId(xid string) (string, bool) {
|
||||
if xid == "" {
|
||||
return "", false
|
||||
}
|
||||
id, ok := e.components.NetworkXIDToPublicID[xid]
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeDNSSettings(s *types.DNSSettings) *proto.DNSSettingsCompact {
|
||||
if s == nil || len(s.DisabledManagementGroups) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := &proto.DNSSettingsCompact{
|
||||
DisabledManagementGroupIds: make([]string, 0, len(s.DisabledManagementGroups)),
|
||||
}
|
||||
for _, gid := range s.DisabledManagementGroups {
|
||||
if id, ok := e.groupPublicXid(gid); ok {
|
||||
out.DisabledManagementGroupIds = append(out.DisabledManagementGroupIds, id)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeRoutes(routes []*nbroute.Route) []*proto.RouteRaw {
|
||||
if len(routes) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.RouteRaw, 0, len(routes))
|
||||
for _, r := range routes {
|
||||
if r == nil {
|
||||
continue
|
||||
}
|
||||
rr := &proto.RouteRaw{
|
||||
Id: r.PublicID,
|
||||
NetId: string(r.NetID),
|
||||
Description: r.Description,
|
||||
KeepRoute: r.KeepRoute,
|
||||
NetworkType: int32(r.NetworkType),
|
||||
Masquerade: r.Masquerade,
|
||||
Metric: int32(r.Metric),
|
||||
Enabled: r.Enabled,
|
||||
SkipAutoApply: r.SkipAutoApply,
|
||||
Domains: r.Domains.ToPunycodeList(),
|
||||
GroupIds: e.groupPublicXids(r.Groups),
|
||||
AccessControlGroupIds: e.groupPublicXids(r.AccessControlGroups),
|
||||
PeerGroupIds: e.groupPublicXids(r.PeerGroups),
|
||||
}
|
||||
if r.Network.IsValid() {
|
||||
rr.NetworkCidr = r.Network.String()
|
||||
}
|
||||
if r.Peer != "" {
|
||||
if idx, ok := e.peerOrder[r.Peer]; ok {
|
||||
rr.PeerIndexSet = true
|
||||
rr.PeerIndex = idx
|
||||
}
|
||||
}
|
||||
out = append(out, rr)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeNameServerGroups(nsgs []*nbdns.NameServerGroup) []*proto.NameServerGroupRaw {
|
||||
if len(nsgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.NameServerGroupRaw, 0, len(nsgs))
|
||||
for _, nsg := range nsgs {
|
||||
if nsg == nil {
|
||||
continue
|
||||
}
|
||||
entry := &proto.NameServerGroupRaw{
|
||||
Id: nsg.PublicID,
|
||||
Nameservers: encodeNameServers(nsg.NameServers),
|
||||
GroupIds: e.groupPublicXids(nsg.Groups),
|
||||
Primary: nsg.Primary,
|
||||
Domains: nsg.Domains,
|
||||
Enabled: nsg.Enabled,
|
||||
SearchDomainsEnabled: nsg.SearchDomainsEnabled,
|
||||
}
|
||||
out = append(out, entry)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeNameServers(servers []nbdns.NameServer) []*proto.NameServer {
|
||||
if len(servers) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.NameServer, 0, len(servers))
|
||||
for _, s := range servers {
|
||||
out = append(out, &proto.NameServer{
|
||||
IP: s.IP.String(),
|
||||
NSType: int64(s.NSType),
|
||||
Port: int64(s.Port),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeSimpleRecords(records []nbdns.SimpleRecord) []*proto.SimpleRecord {
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.SimpleRecord, 0, len(records))
|
||||
for _, r := range records {
|
||||
out = append(out, &proto.SimpleRecord{
|
||||
Name: r.Name,
|
||||
Type: int64(r.Type),
|
||||
Class: r.Class,
|
||||
TTL: int64(r.TTL),
|
||||
RData: r.RData,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeCustomZones(zones []nbdns.CustomZone) []*proto.CustomZone {
|
||||
if len(zones) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.CustomZone, 0, len(zones))
|
||||
for _, z := range zones {
|
||||
out = append(out, &proto.CustomZone{
|
||||
Domain: z.Domain,
|
||||
Records: encodeSimpleRecords(z.Records),
|
||||
SearchDomainDisabled: z.SearchDomainDisabled,
|
||||
NonAuthoritative: z.NonAuthoritative,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentResource) []*proto.NetworkResourceRaw {
|
||||
if len(resources) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.NetworkResourceRaw, 0, len(resources))
|
||||
for _, r := range resources {
|
||||
if r == nil {
|
||||
continue
|
||||
}
|
||||
entry := &proto.NetworkResourceRaw{
|
||||
Id: r.PublicID,
|
||||
Name: r.Name,
|
||||
Description: r.Description,
|
||||
Type: string(r.Type),
|
||||
Address: r.Address,
|
||||
DomainValue: r.Domain,
|
||||
Enabled: r.Enabled,
|
||||
}
|
||||
if id, ok := e.networkPublicId(r.NetworkID); ok {
|
||||
entry.NetworkSeq = id
|
||||
}
|
||||
if r.Prefix.IsValid() {
|
||||
entry.PrefixCidr = r.Prefix.String()
|
||||
}
|
||||
out = append(out, entry)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*types.ComponentRouter) map[string]*proto.NetworkRouterList {
|
||||
if len(routersMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*proto.NetworkRouterList, len(routersMap))
|
||||
for networkXID, routers := range routersMap {
|
||||
if len(routers) == 0 {
|
||||
continue
|
||||
}
|
||||
id, ok := e.networkPublicId(networkXID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
entries := make([]*proto.NetworkRouterEntry, 0, len(routers))
|
||||
for peerID, r := range routers {
|
||||
if r == nil {
|
||||
continue
|
||||
}
|
||||
entry := &proto.NetworkRouterEntry{
|
||||
Id: r.PublicID,
|
||||
PeerGroupIds: e.groupPublicXids(r.PeerGroups),
|
||||
Masquerade: r.Masquerade,
|
||||
Metric: int32(r.Metric),
|
||||
Enabled: r.Enabled,
|
||||
}
|
||||
if idx, ok := e.peerOrder[peerID]; ok {
|
||||
entry.PeerIndexSet = true
|
||||
entry.PeerIndex = idx
|
||||
}
|
||||
entries = append(entries, entry)
|
||||
}
|
||||
out[id] = &proto.NetworkRouterList{Entries: entries}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeResourcePoliciesMap(rpm map[string][]*types.Policy) map[string]*proto.PolicyIds {
|
||||
if len(rpm) == 0 {
|
||||
return nil
|
||||
}
|
||||
// resourceXIDToPublicID is local to one encode — built from components.NetworkResources
|
||||
// (small slice). Network resources without seq id are dropped, matching how
|
||||
// other components-without-seq are silently filtered.
|
||||
resourceXIDToPublicID := make(map[string]string, len(e.components.NetworkResources))
|
||||
for _, r := range e.components.NetworkResources {
|
||||
if r != nil {
|
||||
resourceXIDToPublicID[r.ID] = r.PublicID
|
||||
}
|
||||
}
|
||||
out := make(map[string]*proto.PolicyIds, len(rpm))
|
||||
for resourceXID, policies := range rpm {
|
||||
resId, ok := resourceXIDToPublicID[resourceXID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
ids := make([]string, 0, len(policies))
|
||||
for _, pol := range policies {
|
||||
ids = append(ids, pol.PublicID)
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
continue
|
||||
}
|
||||
out[resId] = &proto.PolicyIds{Ids: ids}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeGroupIDToUserIDs(m map[string][]string) map[string]*proto.UserIDList {
|
||||
if len(m) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*proto.UserIDList, len(m))
|
||||
for groupID, userIDs := range m {
|
||||
id, ok := e.groupPublicXid(groupID)
|
||||
if !ok || len(userIDs) == 0 {
|
||||
continue
|
||||
}
|
||||
out[id] = &proto.UserIDList{UserIds: userIDs}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stringSetToSlice(s map[string]struct{}) []string {
|
||||
if len(s) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(s))
|
||||
for k := range s {
|
||||
out = append(out, k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodePostureFailedPeers(m map[string]map[string]struct{}) map[string]*proto.PeerIndexSet {
|
||||
if len(m) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*proto.PeerIndexSet, len(m))
|
||||
for checkXID, failedPeerIDs := range m {
|
||||
id, ok := e.components.PostureCheckXIDToPublicID[checkXID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
idxs := make([]uint32, 0, len(failedPeerIDs))
|
||||
for peerID := range failedPeerIDs {
|
||||
if idx, ok := e.peerOrder[peerID]; ok {
|
||||
idxs = append(idxs, idx)
|
||||
}
|
||||
}
|
||||
if len(idxs) == 0 {
|
||||
continue
|
||||
}
|
||||
out[id] = &proto.PeerIndexSet{PeerIndexes: idxs}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// toAccountSettingsCompact always returns a non-nil message — the client
|
||||
// dereferences it unconditionally during Calculate(), so a nil here would
|
||||
// crash the receiver. A missing types.AccountSettingsInfo on the server
|
||||
// (which shouldn't happen in production but the encoder is exported)
|
||||
// degrades to login_expiration_enabled = false, which makes
|
||||
// LoginExpired() return false for every peer.
|
||||
func toAccountSettingsCompact(s *types.AccountSettingsInfo) *proto.AccountSettingsCompact {
|
||||
if s == nil {
|
||||
return &proto.AccountSettingsCompact{}
|
||||
}
|
||||
return &proto.AccountSettingsCompact{
|
||||
PeerLoginExpirationEnabled: s.PeerLoginExpirationEnabled,
|
||||
PeerLoginExpirationNs: int64(s.PeerLoginExpiration),
|
||||
}
|
||||
}
|
||||
|
||||
func toAccountNetwork(n *types.Network) *proto.AccountNetwork {
|
||||
if n == nil {
|
||||
return nil
|
||||
}
|
||||
out := &proto.AccountNetwork{
|
||||
Identifier: n.Identifier,
|
||||
NetCidr: n.Net.String(),
|
||||
Dns: n.Dns,
|
||||
Serial: n.CurrentSerial(),
|
||||
}
|
||||
if len(n.NetV6.IP) > 0 {
|
||||
out.NetV6Cidr = n.NetV6.String()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func toPeerCompact(p *types.ComponentPeer) *proto.PeerCompact {
|
||||
pc := &proto.PeerCompact{
|
||||
WgPubKey: decodeWgKey(p.Key),
|
||||
SshPubKey: []byte(p.SSHKey),
|
||||
DnsLabel: p.DNSLabel,
|
||||
AgentVersion: p.AgentVersion,
|
||||
AddedWithSsoLogin: p.AddedWithSSOLogin,
|
||||
LoginExpirationEnabled: p.LoginExpirationEnabled,
|
||||
SshEnabled: p.SSHEnabled,
|
||||
SupportsIpv6: p.SupportsIPv6,
|
||||
SupportsSourcePrefixes: p.SupportsSourcePrefixes,
|
||||
ServerSshAllowed: p.ServerSSHAllowed,
|
||||
}
|
||||
if !p.LastLogin.IsZero() {
|
||||
pc.LastLoginUnixNano = p.LastLogin.UnixNano()
|
||||
}
|
||||
switch {
|
||||
case !p.IP.IsValid():
|
||||
// leave Ip nil
|
||||
case p.IP.Is4() || p.IP.Is4In6():
|
||||
ip := p.IP.Unmap().As4()
|
||||
pc.Ip = ip[:]
|
||||
default:
|
||||
ip := p.IP.As16()
|
||||
pc.Ip = ip[:]
|
||||
}
|
||||
if p.IPv6.IsValid() {
|
||||
ip := p.IPv6.As16()
|
||||
pc.Ipv6 = ip[:]
|
||||
}
|
||||
return pc
|
||||
}
|
||||
|
||||
// decodeWgKey returns the raw 32 bytes of a base64-encoded WireGuard public
|
||||
// key, or nil for an empty / malformed key.
|
||||
func decodeWgKey(s string) []byte {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
out := make([]byte, wgKeyRawLen)
|
||||
n, err := base64.StdEncoding.Decode(out, []byte(s))
|
||||
if err != nil || n != wgKeyRawLen {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func portsToUint32(ports []string) []uint32 {
|
||||
if len(ports) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]uint32, 0, len(ports))
|
||||
for _, p := range ports {
|
||||
v, err := strconv.ParseUint(p, 10, 16)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, uint32(v))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func portRangesToProto(ranges []types.RulePortRange) []*proto.PortInfo_Range {
|
||||
if len(ranges) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*proto.PortInfo_Range, 0, len(ranges))
|
||||
for _, r := range ranges {
|
||||
out = append(out, &proto.PortInfo_Range{
|
||||
Start: uint32(r.Start),
|
||||
End: uint32(r.End),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
781
management/internals/shared/grpc/components_encoder_test.go
Normal file
781
management/internals/shared/grpc/components_encoder_test.go
Normal file
@@ -0,0 +1,781 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"cmp"
|
||||
"net"
|
||||
"net/netip"
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
goproto "google.golang.org/protobuf/proto"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
const testWgKeyA = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
|
||||
const testWgKeyB = "BBCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
|
||||
const testWgKeyC = "CBCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopq="
|
||||
|
||||
// canonicalize rewrites a NetworkMapComponentsFull in place into a canonical
|
||||
// form: peers reordered by wg_pub_key, with the rest of the message rewritten
|
||||
// to reference the new peer indexes. Groups, policies, and router indexes are
|
||||
// also sorted. After canonicalize, two envelopes built from the same logical
|
||||
// input compare byte-equal via proto.Equal.
|
||||
//
|
||||
// This lives on the test side — the encoder itself emits in map-iteration
|
||||
// order. Test-side normalization is the contract for "two encodes are
|
||||
// equivalent".
|
||||
func canonicalize(full *proto.NetworkMapComponentsFull) {
|
||||
if full == nil {
|
||||
return
|
||||
}
|
||||
|
||||
type peerEntry struct {
|
||||
peer *proto.PeerCompact
|
||||
oldIdx uint32
|
||||
}
|
||||
entries := make([]peerEntry, len(full.Peers))
|
||||
for i, p := range full.Peers {
|
||||
entries[i] = peerEntry{peer: p, oldIdx: uint32(i)}
|
||||
}
|
||||
// DnsLabel is unique per peer; it tiebreaks on equal WgPubKey (e.g. both
|
||||
// nil from malformed keys, or both empty for placeholders).
|
||||
slices.SortFunc(entries, func(a, b peerEntry) int {
|
||||
if c := bytes.Compare(a.peer.WgPubKey, b.peer.WgPubKey); c != 0 {
|
||||
return c
|
||||
}
|
||||
return cmp.Compare(a.peer.DnsLabel, b.peer.DnsLabel)
|
||||
})
|
||||
|
||||
remap := make(map[uint32]uint32, len(entries))
|
||||
newPeers := make([]*proto.PeerCompact, len(entries))
|
||||
for newIdx, e := range entries {
|
||||
remap[e.oldIdx] = uint32(newIdx)
|
||||
newPeers[newIdx] = e.peer
|
||||
}
|
||||
full.Peers = newPeers
|
||||
|
||||
full.RouterPeerIndexes = remapAndSort(full.RouterPeerIndexes, remap)
|
||||
for _, g := range full.Groups {
|
||||
g.PeerIndexes = remapAndSort(g.PeerIndexes, remap)
|
||||
}
|
||||
slices.SortFunc(full.Groups, func(a, b *proto.GroupCompact) int { return cmp.Compare(a.Id, b.Id) })
|
||||
|
||||
for _, r := range full.Routes {
|
||||
if r.PeerIndexSet {
|
||||
if newIdx, ok := remap[r.PeerIndex]; ok {
|
||||
r.PeerIndex = newIdx
|
||||
}
|
||||
}
|
||||
slices.Sort(r.GroupIds)
|
||||
slices.Sort(r.AccessControlGroupIds)
|
||||
slices.Sort(r.PeerGroupIds)
|
||||
}
|
||||
slices.SortFunc(full.Routes, func(a, b *proto.RouteRaw) int { return cmp.Compare(a.Id, b.Id) })
|
||||
|
||||
for _, list := range full.RoutersMap {
|
||||
for _, entry := range list.Entries {
|
||||
if entry.PeerIndexSet {
|
||||
if newIdx, ok := remap[entry.PeerIndex]; ok {
|
||||
entry.PeerIndex = newIdx
|
||||
}
|
||||
}
|
||||
slices.Sort(entry.PeerGroupIds)
|
||||
}
|
||||
slices.SortFunc(list.Entries, func(a, b *proto.NetworkRouterEntry) int { return cmp.Compare(a.Id, b.Id) })
|
||||
}
|
||||
|
||||
for _, set := range full.PostureFailedPeers {
|
||||
set.PeerIndexes = remapAndSort(set.PeerIndexes, remap)
|
||||
}
|
||||
|
||||
for _, p := range full.Policies {
|
||||
slices.Sort(p.SourceGroupIds)
|
||||
slices.Sort(p.DestinationGroupIds)
|
||||
}
|
||||
// Sort policies by (Id, source_group_ids, destination_group_ids) so that
|
||||
// multiple PolicyCompact entries sharing the same Id (one per rule, when
|
||||
// a Policy has multiple rules) still get a deterministic order. After
|
||||
// sorting we remap indexes in ResourcePoliciesMap.
|
||||
policyOldOrder := make(map[*proto.PolicyCompact]uint32, len(full.Policies))
|
||||
for i, p := range full.Policies {
|
||||
policyOldOrder[p] = uint32(i)
|
||||
}
|
||||
slices.SortFunc(full.Policies, func(a, b *proto.PolicyCompact) int {
|
||||
if c := cmp.Compare(a.Id, b.Id); c != 0 {
|
||||
return c
|
||||
}
|
||||
if c := slices.Compare(a.SourceGroupIds, b.SourceGroupIds); c != 0 {
|
||||
return c
|
||||
}
|
||||
return slices.Compare(a.DestinationGroupIds, b.DestinationGroupIds)
|
||||
})
|
||||
policyRemap := make(map[uint32]uint32, len(full.Policies))
|
||||
for newIdx, p := range full.Policies {
|
||||
policyRemap[policyOldOrder[p]] = uint32(newIdx)
|
||||
}
|
||||
for _, idxs := range full.ResourcePoliciesMap {
|
||||
slices.Sort(idxs.Ids)
|
||||
}
|
||||
for _, list := range full.GroupIdToUserIds {
|
||||
slices.Sort(list.UserIds)
|
||||
}
|
||||
slices.Sort(full.AllowedUserIds)
|
||||
}
|
||||
|
||||
func remapAndSort(idxs []uint32, remap map[uint32]uint32) []uint32 {
|
||||
out := make([]uint32, 0, len(idxs))
|
||||
for _, i := range idxs {
|
||||
if newIdx, ok := remap[i]; ok {
|
||||
out = append(out, newIdx)
|
||||
}
|
||||
}
|
||||
slices.Sort(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// envelopesEquivalent decodes both envelopes, canonicalizes them, and reports
|
||||
// whether they're proto.Equal. Use instead of byte-comparing marshaled output:
|
||||
// the encoder is intentionally non-deterministic.
|
||||
func envelopesEquivalent(a, b *proto.NetworkMapEnvelope) bool {
|
||||
canonicalize(a.GetFull())
|
||||
canonicalize(b.GetFull())
|
||||
return goproto.Equal(a, b)
|
||||
}
|
||||
|
||||
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",
|
||||
}
|
||||
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",
|
||||
}
|
||||
peerC := &types.ComponentPeer{
|
||||
ID: "peer-c",
|
||||
Key: testWgKeyC,
|
||||
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
|
||||
DNSLabel: "peerc",
|
||||
AgentVersion: "0.40.0",
|
||||
}
|
||||
|
||||
return &types.NetworkMapComponents{
|
||||
PeerID: "peer-a",
|
||||
Network: &types.Network{
|
||||
Identifier: "net-test",
|
||||
Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)},
|
||||
Serial: 7,
|
||||
},
|
||||
AccountSettings: &types.AccountSettingsInfo{
|
||||
PeerLoginExpirationEnabled: true,
|
||||
PeerLoginExpiration: 2 * time.Hour,
|
||||
},
|
||||
Peers: map[string]*types.ComponentPeer{
|
||||
"peer-a": peerA,
|
||||
"peer-b": peerB,
|
||||
"peer-c": peerC,
|
||||
},
|
||||
Groups: map[string]*types.ComponentGroup{
|
||||
"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"}},
|
||||
},
|
||||
Policies: []*types.Policy{
|
||||
{
|
||||
ID: "pol-1",
|
||||
PublicID: "10",
|
||||
Enabled: true,
|
||||
Rules: []*types.PolicyRule{{
|
||||
ID: "rule-1", Enabled: true, Action: types.PolicyTrafficActionAccept,
|
||||
Protocol: types.PolicyRuleProtocolTCP, Bidirectional: true,
|
||||
Ports: []string{"22", "80"},
|
||||
PortRanges: []types.RulePortRange{{Start: 8000, End: 8100}},
|
||||
Sources: []string{"group-src"},
|
||||
Destinations: []string{"group-dst"},
|
||||
}},
|
||||
},
|
||||
},
|
||||
RouterPeers: map[string]*types.ComponentPeer{"peer-c": peerC},
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_Basic(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
|
||||
Components: c,
|
||||
DNSDomain: "netbird.cloud",
|
||||
})
|
||||
|
||||
require.NotNil(t, env)
|
||||
full := env.GetFull()
|
||||
require.NotNil(t, full, "envelope must contain Full payload")
|
||||
|
||||
assert.EqualValues(t, 7, full.Serial)
|
||||
assert.Equal(t, "netbird.cloud", full.DnsDomain)
|
||||
|
||||
require.NotNil(t, full.Network)
|
||||
assert.Equal(t, "net-test", full.Network.Identifier)
|
||||
assert.Equal(t, "100.64.0.0/10", full.Network.NetCidr)
|
||||
|
||||
require.NotNil(t, full.AccountSettings)
|
||||
assert.True(t, full.AccountSettings.PeerLoginExpirationEnabled)
|
||||
assert.EqualValues(t, (2 * time.Hour).Nanoseconds(), full.AccountSettings.PeerLoginExpirationNs)
|
||||
|
||||
require.Len(t, full.Peers, 3)
|
||||
byLabel := map[string]*proto.PeerCompact{}
|
||||
for _, p := range full.Peers {
|
||||
assert.Len(t, p.WgPubKey, 32, "wg key must be raw 32 bytes")
|
||||
assert.Len(t, p.Ip, 4, "ipv4 must be raw 4 bytes")
|
||||
byLabel[p.DnsLabel] = p
|
||||
}
|
||||
assert.Len(t, byLabel["peerb"].Ipv6, 16, "peer-b has ipv6 → 16 bytes")
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RepeatEncodesEquivalent(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
|
||||
// Hammer it 100 times — Go map iteration is randomized per call, so each
|
||||
// run produces different wire bytes, but the canonicalized form must
|
||||
// match.
|
||||
for i := 0; i < 100; i++ {
|
||||
got := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
require.True(t, envelopesEquivalent(expected, got),
|
||||
"encode #%d must be semantically equivalent to first encode", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_ConcurrentEncodesEquivalent(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
|
||||
const goroutines = 50
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(goroutines)
|
||||
results := make([]*proto.NetworkMapEnvelope, goroutines)
|
||||
for i := 0; i < goroutines; i++ {
|
||||
i := i
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
results[i] = EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
for i, got := range results {
|
||||
require.NotNil(t, got, "goroutine %d returned nil", i)
|
||||
require.True(t, envelopesEquivalent(expected, got),
|
||||
"goroutine %d produced inequivalent envelope", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_GroupsByAccountPublicId(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Groups, 2)
|
||||
|
||||
groupByID := map[string]*proto.GroupCompact{}
|
||||
for _, g := range full.Groups {
|
||||
groupByID[g.Id] = g
|
||||
}
|
||||
require.Contains(t, groupByID, "1")
|
||||
require.Contains(t, groupByID, "2")
|
||||
assert.Len(t, groupByID["1"].PeerIndexes, 1)
|
||||
assert.Len(t, groupByID["2"].PeerIndexes, 2)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_PolicyExpansion(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Policies, 1)
|
||||
pc := full.Policies[0]
|
||||
assert.EqualValues(t, "10", pc.Id)
|
||||
assert.Equal(t, proto.RuleAction_ACCEPT, pc.Action)
|
||||
assert.Equal(t, proto.RuleProtocol_TCP, pc.Protocol)
|
||||
assert.True(t, pc.Bidirectional)
|
||||
assert.Equal(t, []uint32{22, 80}, pc.Ports)
|
||||
require.Len(t, pc.PortRanges, 1)
|
||||
assert.EqualValues(t, 8000, pc.PortRanges[0].Start)
|
||||
assert.EqualValues(t, 8100, pc.PortRanges[0].End)
|
||||
assert.Equal(t, []string{"1"}, pc.SourceGroupIds)
|
||||
assert.Equal(t, []string{"2"}, pc.DestinationGroupIds)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RouterIndexes(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.RouterPeerIndexes, 1)
|
||||
idx := full.RouterPeerIndexes[0]
|
||||
require.Less(t, int(idx), len(full.Peers))
|
||||
assert.Equal(t, "peerc", full.Peers[idx].DnsLabel)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_DisabledPolicySkipped(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Policies[0].Enabled = false
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
assert.Empty(t, full.Policies)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_TwoPeersSameMalformedKey(t *testing.T) {
|
||||
// Both peers have nil WgPubKey after decode; canonicalize must still
|
||||
// produce a stable order using DnsLabel as a tiebreaker, so 100 encodes
|
||||
// canonicalize identically.
|
||||
c := newTestComponents()
|
||||
c.Peers["peer-a"].Key = "garbage-a-!!!"
|
||||
c.Peers["peer-b"].Key = "garbage-b-!!!"
|
||||
|
||||
expected := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
for i := 0; i < 100; i++ {
|
||||
got := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
require.True(t, envelopesEquivalent(expected, got),
|
||||
"encode #%d with two same-key peers must canonicalize equivalently", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_MalformedWgKey(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Peers["peer-a"].Key = "not-base64-!!!"
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Peers, 3)
|
||||
|
||||
var byLabel = map[string]*proto.PeerCompact{}
|
||||
for _, p := range full.Peers {
|
||||
byLabel[p.DnsLabel] = p
|
||||
}
|
||||
assert.Nil(t, byLabel["peera"].WgPubKey, "peer with malformed key encodes nil WgPubKey")
|
||||
assert.Len(t, byLabel["peerb"].WgPubKey, 32, "other peers retain their key")
|
||||
}
|
||||
|
||||
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",
|
||||
}
|
||||
c.Peers["peer-v6"] = v6Only
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
var found *proto.PeerCompact
|
||||
for _, p := range full.Peers {
|
||||
if p.DnsLabel == "peerv6" {
|
||||
found = p
|
||||
}
|
||||
}
|
||||
require.NotNil(t, found, "ipv6-only peer must be present")
|
||||
assert.Empty(t, found.Ip, "no IPv4 address → empty Ip")
|
||||
assert.Len(t, found.Ipv6, 16)
|
||||
}
|
||||
|
||||
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",
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
var found *proto.PeerCompact
|
||||
for _, p := range full.Peers {
|
||||
if p.DnsLabel == "peernoip" {
|
||||
found = p
|
||||
}
|
||||
}
|
||||
require.NotNil(t, found)
|
||||
assert.Empty(t, found.Ip)
|
||||
assert.Empty(t, found.Ipv6)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
|
||||
c := &types.NetworkMapComponents{
|
||||
Network: &types.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
|
||||
}
|
||||
|
||||
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c})
|
||||
|
||||
full := env.GetFull()
|
||||
require.NotNil(t, full)
|
||||
assert.Empty(t, full.Peers)
|
||||
assert.Empty(t, full.Groups)
|
||||
assert.Empty(t, full.Policies)
|
||||
assert.Empty(t, full.RouterPeerIndexes)
|
||||
require.NotNil(t, full.AccountSettings, "AccountSettingsCompact must always be emitted (client dereferences it unconditionally)")
|
||||
}
|
||||
|
||||
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"].LoginExpirationEnabled = true
|
||||
c.Peers["peer-a"].LastLogin = now
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
var pa *proto.PeerCompact
|
||||
for _, p := range full.Peers {
|
||||
if p.DnsLabel == "peera" {
|
||||
pa = p
|
||||
}
|
||||
}
|
||||
require.NotNil(t, pa)
|
||||
assert.True(t, pa.AddedWithSsoLogin)
|
||||
assert.True(t, pa.LoginExpirationEnabled)
|
||||
assert.Equal(t, now.UnixNano(), pa.LastLoginUnixNano)
|
||||
|
||||
// peer-b has no UserID and no LastLogin → all fields zero-value.
|
||||
var pb *proto.PeerCompact
|
||||
for _, p := range full.Peers {
|
||||
if p.DnsLabel == "peerb" {
|
||||
pb = p
|
||||
}
|
||||
}
|
||||
require.NotNil(t, pb)
|
||||
assert.False(t, pb.AddedWithSsoLogin)
|
||||
assert.False(t, pb.LoginExpirationEnabled)
|
||||
assert.Zero(t, pb.LastLoginUnixNano)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RoutesRoundTrip(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Routes = []*nbroute.Route{
|
||||
{
|
||||
ID: "route-peer",
|
||||
PublicID: "100",
|
||||
NetID: "net-A",
|
||||
Description: "via peer-c",
|
||||
Network: netip.MustParsePrefix("10.0.0.0/16"),
|
||||
Peer: "peer-c", // peer ID, not WG key
|
||||
Groups: []string{"group-src"},
|
||||
AccessControlGroups: []string{"group-dst"},
|
||||
Enabled: true,
|
||||
},
|
||||
{
|
||||
ID: "route-peergroup",
|
||||
PublicID: "101",
|
||||
NetID: "net-B",
|
||||
Network: netip.MustParsePrefix("10.1.0.0/16"),
|
||||
PeerGroups: []string{"group-src", "group-dst"},
|
||||
Enabled: true,
|
||||
},
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Routes, 2)
|
||||
byNetID := map[string]*proto.RouteRaw{}
|
||||
for _, r := range full.Routes {
|
||||
byNetID[r.NetId] = r
|
||||
}
|
||||
|
||||
r1 := byNetID["net-A"]
|
||||
require.NotNil(t, r1)
|
||||
assert.True(t, r1.PeerIndexSet, "route with peer must set peer_index_set")
|
||||
require.Less(t, int(r1.PeerIndex), len(full.Peers))
|
||||
assert.Equal(t, "peerc", full.Peers[r1.PeerIndex].DnsLabel)
|
||||
assert.Equal(t, []string{"1"}, r1.GroupIds, "group-src has AccountSeqID 1")
|
||||
assert.Equal(t, []string{"2"}, r1.AccessControlGroupIds, "group-dst has AccountSeqID 2")
|
||||
assert.Empty(t, r1.PeerGroupIds)
|
||||
|
||||
r2 := byNetID["net-B"]
|
||||
require.NotNil(t, r2)
|
||||
assert.False(t, r2.PeerIndexSet, "route with peer_groups must NOT set peer_index_set")
|
||||
assert.ElementsMatch(t, []string{"1", "2"}, r2.PeerGroupIds)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RouteWithMissingPeerLeavesIndexUnset(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Routes = []*nbroute.Route{{
|
||||
ID: "route-x",
|
||||
PublicID: "100",
|
||||
Peer: "peer-not-in-components",
|
||||
Network: netip.MustParsePrefix("10.0.0.0/16"),
|
||||
Enabled: true,
|
||||
}}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Routes, 1)
|
||||
assert.False(t, full.Routes[0].PeerIndexSet,
|
||||
"missing peer reference must not pretend to point at peer index 0")
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
// Policy that exists ONLY in ResourcePoliciesMap, not in c.Policies. This
|
||||
// is the I1 case — without unionPolicies the encoder would silently
|
||||
// drop it from the wire.
|
||||
resourceOnlyPolicy := &types.Policy{
|
||||
ID: "pol-resource", PublicID: "99", Enabled: true,
|
||||
Rules: []*types.PolicyRule{{
|
||||
ID: "rule-r", Enabled: true, Action: types.PolicyTrafficActionAccept,
|
||||
Protocol: types.PolicyRuleProtocolTCP,
|
||||
Sources: []string{"group-src"},
|
||||
Destinations: []string{"group-dst"},
|
||||
}},
|
||||
}
|
||||
c.ResourcePoliciesMap = map[string][]*types.Policy{
|
||||
"resource-x": {c.Policies[0], resourceOnlyPolicy}, // shared + resource-only
|
||||
}
|
||||
// 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{
|
||||
{ID: "resource-x", PublicID: "77", Name: "res-x", Enabled: true},
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.Policies, 2, "encoded policies must include both peer-traffic and resource-only")
|
||||
|
||||
policyByID := map[string]*proto.PolicyCompact{}
|
||||
policyIds := make([]string, 0)
|
||||
for _, p := range full.Policies {
|
||||
policyByID[p.Id] = p
|
||||
policyIds = append(policyIds, p.Id)
|
||||
}
|
||||
require.Contains(t, policyByID, "10", "original peer-traffic policy id 10")
|
||||
require.Contains(t, policyByID, "99", "resource-only policy id 99")
|
||||
|
||||
require.Contains(t, full.ResourcePoliciesMap, "77")
|
||||
ids := full.ResourcePoliciesMap["77"].Ids
|
||||
require.Len(t, ids, 2)
|
||||
assert.ElementsMatch(t, policyIds, ids,
|
||||
"resource policies map must reference both wire policy indexes")
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_NameServerGroups(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.NameServerGroups = []*nbdns.NameServerGroup{{
|
||||
ID: "nsg-1", PublicID: "50", Name: "Main", Description: "primary",
|
||||
NameServers: []nbdns.NameServer{{
|
||||
IP: netip.MustParseAddr("8.8.8.8"), NSType: nbdns.UDPNameServerType, Port: 53,
|
||||
}},
|
||||
Groups: []string{"group-src", "group-not-persisted"},
|
||||
Primary: true, Enabled: true,
|
||||
Domains: []string{"corp.example"},
|
||||
}}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.NameserverGroups, 1)
|
||||
nsg := full.NameserverGroups[0]
|
||||
assert.EqualValues(t, "50", nsg.Id)
|
||||
assert.True(t, nsg.Primary)
|
||||
require.Len(t, nsg.Nameservers, 1)
|
||||
assert.Equal(t, "8.8.8.8", nsg.Nameservers[0].IP)
|
||||
assert.Equal(t, []string{"1"}, nsg.GroupIds)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_PostureFailedPeers(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.PostureCheckXIDToPublicID = map[string]string{"check-1": "33"}
|
||||
c.PostureFailedPeers = map[string]map[string]struct{}{
|
||||
"check-1": {
|
||||
"peer-a": {},
|
||||
"peer-b": {},
|
||||
"peer-not-in-account": {},
|
||||
},
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Contains(t, full.PostureFailedPeers, "33")
|
||||
idxs := full.PostureFailedPeers["33"].PeerIndexes
|
||||
assert.Len(t, idxs, 2, "missing peer is silently dropped (filterPostureFailedPeers guarantees presence in real data)")
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RoutersMap(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
|
||||
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
|
||||
"net-1": {
|
||||
"peer-c": {
|
||||
PublicID: "200",
|
||||
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Contains(t, full.RoutersMap, "5")
|
||||
entries := full.RoutersMap["5"].Entries
|
||||
require.Len(t, entries, 1)
|
||||
e := entries[0]
|
||||
assert.EqualValues(t, "200", e.Id)
|
||||
assert.True(t, e.PeerIndexSet)
|
||||
require.Less(t, int(e.PeerIndex), len(full.Peers))
|
||||
assert.Equal(t, "peerc", full.Peers[e.PeerIndex].DnsLabel)
|
||||
assert.True(t, e.Masquerade)
|
||||
assert.EqualValues(t, 10, e.Metric)
|
||||
assert.True(t, e.Enabled)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RouterPeerNotInComponentsPeers(t *testing.T) {
|
||||
// Router peer in c.RouterPeers but NOT in c.Peers (validation may have
|
||||
// filtered it). indexRouterPeers runs before encodeRoutersMap, so the
|
||||
// peer_index reference must still resolve.
|
||||
c := newTestComponents()
|
||||
delete(c.Peers, "peer-c")
|
||||
routerPeer := &types.ComponentPeer{
|
||||
ID: "peer-c", Key: testWgKeyC, IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
|
||||
DNSLabel: "peerc", AgentVersion: "0.40.0",
|
||||
}
|
||||
c.RouterPeers = map[string]*types.ComponentPeer{"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}},
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Contains(t, full.RoutersMap, "5")
|
||||
require.Len(t, full.RoutersMap["5"].Entries, 1)
|
||||
e := full.RoutersMap["5"].Entries[0]
|
||||
assert.True(t, e.PeerIndexSet, "router peer must be indexed even when not in c.Peers")
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_GroupIDToUserIDs(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.GroupIDToUserIDs = map[string][]string{
|
||||
"group-src": {"user-1", "user-2"},
|
||||
"group-missing": {"user-4"}, // group not in components → drop
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.Len(t, full.GroupIdToUserIds, 1, "only present groups survive")
|
||||
require.Contains(t, full.GroupIdToUserIds, "1")
|
||||
assert.ElementsMatch(t, []string{"user-1", "user-2"}, full.GroupIdToUserIds["1"].UserIds)
|
||||
}
|
||||
|
||||
func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
|
||||
assert.Nil(t, toProxyPatch(nil, "netbird.cloud", false, false))
|
||||
assert.Nil(t, toProxyPatch(&types.NetworkMap{}, "netbird.cloud", false, false),
|
||||
"empty NetworkMap (no peers, rules, routes etc) → nil patch so proto3 omits the field")
|
||||
}
|
||||
|
||||
func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
|
||||
nm := &types.NetworkMap{
|
||||
Peers: []*types.ComponentPeer{{
|
||||
ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
|
||||
DNSLabel: "extpeer", AgentVersion: "0.40.0",
|
||||
}},
|
||||
FirewallRules: []*types.FirewallRule{{
|
||||
PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
|
||||
}},
|
||||
}
|
||||
|
||||
patch := toProxyPatch(nm, "netbird.cloud", false, false)
|
||||
|
||||
require.NotNil(t, patch)
|
||||
assert.Len(t, patch.Peers, 1)
|
||||
assert.Len(t, patch.FirewallRules, 1)
|
||||
}
|
||||
|
||||
// TestEncodeNetworkMapEnvelope_ProxyPatchPropagated covers the ProxyPatch
|
||||
// pass-through in both encoder branches (normal path + nil-Components
|
||||
// graceful-degrade). Guards against a regression that drops `ProxyPatch:`
|
||||
// from one of the envelope struct literals.
|
||||
func TestEncodeNetworkMapEnvelope_ProxyPatchPropagated(t *testing.T) {
|
||||
patch := &proto.ProxyPatch{
|
||||
ForwardingRules: []*proto.ForwardingRule{{
|
||||
Protocol: proto.RuleProtocol_TCP,
|
||||
DestinationPort: &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 80}},
|
||||
TranslatedAddress: net.IPv4(10, 0, 0, 1).To4(),
|
||||
TranslatedPort: &proto.PortInfo{PortSelection: &proto.PortInfo_Port{Port: 8080}},
|
||||
}},
|
||||
}
|
||||
|
||||
t.Run("normal_path", func(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
|
||||
Components: c,
|
||||
ProxyPatch: patch,
|
||||
}).GetFull()
|
||||
|
||||
require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the normal encode path")
|
||||
assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
|
||||
})
|
||||
|
||||
t.Run("empty_components_graceful_degrade", func(t *testing.T) {
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
|
||||
Components: emptyNetworkMapComponents(),
|
||||
ProxyPatch: patch,
|
||||
}).GetFull()
|
||||
|
||||
require.NotNil(t, full.ProxyPatch, "ProxyPatch must propagate through the nil-Components branch too")
|
||||
assert.Len(t, full.ProxyPatch.ForwardingRules, 1)
|
||||
})
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_NilComponentsGracefulDegrade(t *testing.T) {
|
||||
// nil Components → minimal envelope, no crash. Matches the legacy
|
||||
// behaviour for missing/unvalidated peers.
|
||||
env := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
|
||||
Components: emptyNetworkMapComponents(),
|
||||
DNSDomain: "netbird.cloud",
|
||||
})
|
||||
|
||||
require.NotNil(t, env)
|
||||
full := env.GetFull()
|
||||
require.NotNil(t, full)
|
||||
require.NotNil(t, full.AccountSettings, "AccountSettings must always be non-nil")
|
||||
assert.Equal(t, "netbird.cloud", full.DnsDomain)
|
||||
assert.Len(t, full.Peers, 1)
|
||||
assert.Empty(t, full.Policies)
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
|
||||
c := &types.NetworkMapComponents{
|
||||
Network: &types.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
|
||||
// AccountSettings deliberately nil
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
require.NotNil(t, full.AccountSettings, "client dereferences AccountSettings unconditionally during Calculate(); a nil here would crash the receiver")
|
||||
assert.False(t, full.AccountSettings.PeerLoginExpirationEnabled)
|
||||
assert.Zero(t, full.AccountSettings.PeerLoginExpirationNs)
|
||||
}
|
||||
|
||||
func emptyNetworkMapComponents() *types.NetworkMapComponents {
|
||||
return types.EmptyNetworkMapComponents(
|
||||
&types.NetworkMapComponents{
|
||||
PeerID: "peer-id", Peers: map[string]*types.ComponentPeer{"peer-id": {}}},
|
||||
)
|
||||
}
|
||||
200
management/internals/shared/grpc/components_envelope_response.go
Normal file
200
management/internals/shared/grpc/components_envelope_response.go
Normal file
@@ -0,0 +1,200 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ssh/auth"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// ToComponentSyncResponse builds a SyncResponse carrying the compact
|
||||
// NetworkMapEnvelope for capability-aware peers. The legacy proto.NetworkMap
|
||||
// field is intentionally left empty — capable peers ignore it and the
|
||||
// envelope alone is the authoritative wire shape.
|
||||
//
|
||||
// PeerConfig is computed once server-side using the receiving peer's own
|
||||
// account-level network metadata. EnableSSH inside PeerConfig is left at
|
||||
// peer.SSHEnabled (the peer's local setting); account-policy-driven SSH is
|
||||
// computed by the client from the envelope's GroupIDToUserIDs / AllowedUserIDs
|
||||
// inside Calculate(), so the SshConfig.SshEnabled bit may flip true on the
|
||||
// client even though the server-side PeerConfig reports false.
|
||||
func ToComponentSyncResponse(
|
||||
ctx context.Context,
|
||||
config *nbconfig.Config,
|
||||
httpConfig *nbconfig.HttpServerConfig,
|
||||
deviceFlowConfig *nbconfig.DeviceAuthorizationFlow,
|
||||
peer *nbpeer.Peer,
|
||||
turnCredentials *Token,
|
||||
relayCredentials *Token,
|
||||
components *types.NetworkMapComponents,
|
||||
proxyPatch *types.NetworkMap,
|
||||
dnsName string,
|
||||
checks []*posture.Checks,
|
||||
settings *types.Settings,
|
||||
extraSettings *types.ExtraSettings,
|
||||
peerGroups []string,
|
||||
dnsFwdPort int64,
|
||||
) *proto.SyncResponse {
|
||||
//
|
||||
// 'component' parameter is expected to never be nil
|
||||
// 'peer' parameter is expected to never be nil
|
||||
//
|
||||
// TODO (dmitri) consider using invariants?
|
||||
//
|
||||
enableSSH := computeSSHEnabledForPeer(components, peer)
|
||||
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH)
|
||||
|
||||
includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
|
||||
useSourcePrefixes := peer.SupportsSourcePrefixes()
|
||||
|
||||
userIDClaim := auth.DefaultUserIDClaim
|
||||
if httpConfig != nil && httpConfig.AuthUserIDClaim != "" {
|
||||
userIDClaim = httpConfig.AuthUserIDClaim
|
||||
}
|
||||
|
||||
envelope := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{
|
||||
Components: components,
|
||||
PeerConfig: peerConfig,
|
||||
DNSDomain: dnsName,
|
||||
DNSForwarderPort: dnsFwdPort,
|
||||
UserIDClaim: userIDClaim,
|
||||
ProxyPatch: toProxyPatch(proxyPatch, dnsName, includeIPv6, useSourcePrefixes),
|
||||
})
|
||||
|
||||
resp := &proto.SyncResponse{
|
||||
PeerConfig: peerConfig,
|
||||
NetworkMapEnvelope: envelope,
|
||||
Checks: toProtocolChecks(ctx, checks),
|
||||
Version: int32(sharedgrpc.ComponentNetworkMap),
|
||||
}
|
||||
|
||||
nbConfig := toNetbirdConfig(config, turnCredentials, relayCredentials, extraSettings, settings)
|
||||
resp.NetbirdConfig = integrationsConfig.ExtendNetBirdConfig(peer.ID, peerGroups, nbConfig, extraSettings)
|
||||
|
||||
// settings == nil → field stays nil → "no info in this snapshot", client
|
||||
// preserves the deadline it already had. settings non-nil → emit either a
|
||||
// valid deadline or the explicit-zero "disabled" sentinel via
|
||||
// encodeSessionExpiresAt.
|
||||
if settings != nil {
|
||||
resp.SessionExpiresAt = encodeSessionExpiresAt(
|
||||
peer.SessionExpiresAt(settings.PeerLoginExpirationEnabled, settings.PeerLoginExpiration),
|
||||
)
|
||||
}
|
||||
|
||||
return resp
|
||||
}
|
||||
|
||||
// toProxyPatch converts a proxy-injected *types.NetworkMap into the wire
|
||||
// patch the components envelope ships alongside. Returns nil when there are
|
||||
// no fragments to merge — proto3 omits a nil message field, so the receiver
|
||||
// sees no patch and skips the merge step entirely.
|
||||
//
|
||||
// We reuse the legacy proto-conversion helpers (toProtocolRoutes,
|
||||
// toProtocolFirewallRules, toProtocolRoutesFirewallRules,
|
||||
// appendRemotePeerConfig, ForwardingRule.ToProto) because the proxy
|
||||
// delivers fragments pre-expanded — there's no raw component shape to
|
||||
// derive them from. Components purity isn't violated: proxy data isn't
|
||||
// policy-graph-derived, it's externally injected post-Calculate, so the
|
||||
// client merges it on top of its locally-computed NetworkMap.
|
||||
func toProxyPatch(nm *types.NetworkMap, dnsName string, includeIPv6, useSourcePrefixes bool) *proto.ProxyPatch {
|
||||
if nm == nil {
|
||||
return nil
|
||||
}
|
||||
if len(nm.Peers) == 0 && len(nm.OfflinePeers) == 0 && len(nm.FirewallRules) == 0 &&
|
||||
len(nm.Routes) == 0 && len(nm.RoutesFirewallRules) == 0 && len(nm.ForwardingRules) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
patch := &proto.ProxyPatch{
|
||||
Peers: networkmap.AppendRemotePeerConfig(nil, nm.Peers, dnsName, includeIPv6),
|
||||
OfflinePeers: networkmap.AppendRemotePeerConfig(nil, nm.OfflinePeers, dnsName, includeIPv6),
|
||||
FirewallRules: networkmap.ToProtocolFirewallRules(nm.FirewallRules, includeIPv6, useSourcePrefixes),
|
||||
Routes: networkmap.ToProtocolRoutes(nm.Routes),
|
||||
RouteFirewallRules: networkmap.ToProtocolRoutesFirewallRules(nm.RoutesFirewallRules),
|
||||
}
|
||||
if len(nm.ForwardingRules) > 0 {
|
||||
patch.ForwardingRules = make([]*proto.ForwardingRule, 0, len(nm.ForwardingRules))
|
||||
for _, r := range nm.ForwardingRules {
|
||||
patch.ForwardingRules = append(patch.ForwardingRules, r.ToProto())
|
||||
}
|
||||
}
|
||||
return patch
|
||||
}
|
||||
|
||||
// computeSSHEnabledForPeer mirrors the SSH-server-activation bit that
|
||||
// Calculate() folds into NetworkMap.EnableSSH. Components-format peers
|
||||
// receive a freshly-computed PeerConfig.SshConfig.SshEnabled at sync time;
|
||||
// without this helper the field would be incorrectly false for any peer
|
||||
// that's the destination of an SSH-enabling policy without having
|
||||
// peer.SSHEnabled set locally.
|
||||
//
|
||||
// Mirrors the two activation paths Calculate() uses:
|
||||
// 1. Explicit: rule.Protocol == NetbirdSSH and peer is in the rule's
|
||||
// destinations.
|
||||
// 2. Legacy implicit: rule covers TCP/22 or TCP/22022 (or ALL), peer is in
|
||||
// destinations, AND the peer has SSHEnabled set locally — this is the
|
||||
// "allow-all/TCP-22 implies SSH activation for SSH-capable peers" path.
|
||||
//
|
||||
// The full SSH AuthorizedUsers map is still produced by the client when it
|
||||
// runs Calculate() over the envelope.
|
||||
func computeSSHEnabledForPeer(c *types.NetworkMapComponents, peer *nbpeer.Peer) bool {
|
||||
if c == nil || peer == nil {
|
||||
return false
|
||||
}
|
||||
// Mirror Calculate's `getAllPeersFromGroups` invariant: target peer must
|
||||
// exist in c.Peers, otherwise no rule applies to it.
|
||||
if _, ok := c.Peers[peer.ID]; !ok {
|
||||
return false
|
||||
}
|
||||
for _, policy := range c.Policies {
|
||||
if policy == nil || !policy.Enabled {
|
||||
continue
|
||||
}
|
||||
for _, rule := range policy.Rules {
|
||||
if ruleEnablesSSHForPeer(c, rule, peer) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ruleEnablesSSHForPeer returns true when rule is active, targets peer, and
|
||||
// either explicitly authorises SSH or covers the legacy TCP/22 path while the
|
||||
// peer itself has SSH enabled locally.
|
||||
func ruleEnablesSSHForPeer(c *types.NetworkMapComponents, rule *types.PolicyRule, peer *nbpeer.Peer) bool {
|
||||
if rule == nil || !rule.Enabled {
|
||||
return false
|
||||
}
|
||||
if !peerInDestinations(c, rule, peer.ID) {
|
||||
return false
|
||||
}
|
||||
if rule.Protocol == types.PolicyRuleProtocolNetbirdSSH {
|
||||
return true
|
||||
}
|
||||
return peer.SSHEnabled && types.PolicyRuleImpliesLegacySSH(rule)
|
||||
}
|
||||
|
||||
// peerInDestinations reports whether peerID is in any of rule.Destinations'
|
||||
// groups (or matches DestinationResource if it's a peer-typed resource —
|
||||
// for non-peer types Calculate falls through to group lookup, so we mirror
|
||||
// that exactly to avoid silent divergence).
|
||||
func peerInDestinations(c *types.NetworkMapComponents, rule *types.PolicyRule, peerID string) bool {
|
||||
if rule.DestinationResource.Type == types.ResourceTypePeer && rule.DestinationResource.ID != "" {
|
||||
return rule.DestinationResource.ID == peerID
|
||||
}
|
||||
for _, groupID := range rule.Destinations {
|
||||
if c.IsPeerInGroup(peerID, groupID) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
)
|
||||
|
||||
// TestComputeSSHEnabledForPeer covers both Calculate-mirroring branches:
|
||||
// explicit NetbirdSSH protocol, and the legacy implicit case where a
|
||||
// TCP/22 (or 22022 / ALL / port-range-covering-22) rule activates SSH when
|
||||
// the destination peer has SSHEnabled=true locally.
|
||||
func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
const targetPeerID = "target"
|
||||
const targetGroupID = "g_dst"
|
||||
|
||||
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}}
|
||||
return &types.NetworkMapComponents{
|
||||
Peers: map[string]*types.ComponentPeer{targetPeerID: peer.ToComponent()},
|
||||
Groups: map[string]*types.ComponentGroup{targetGroupID: group},
|
||||
Policies: []*types.Policy{{
|
||||
ID: "p",
|
||||
Enabled: true,
|
||||
Rules: []*types.PolicyRule{rule},
|
||||
}},
|
||||
}, peer
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
peerSSH bool
|
||||
rule types.PolicyRule
|
||||
wantEnabled bool
|
||||
}{
|
||||
{
|
||||
name: "explicit-netbird-ssh-activates-regardless-of-peer-ssh",
|
||||
peerSSH: false,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "implicit-tcp-22-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "implicit-tcp-22-without-peer-ssh-disabled",
|
||||
peerSSH: false,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
},
|
||||
{
|
||||
name: "implicit-tcp-22022-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22022"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "implicit-all-protocol-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolALL,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "implicit-port-range-covers-22",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: types.PolicyRuleProtocolTCP,
|
||||
PortRanges: []types.RulePortRange{{Start: 20, End: 30}},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "tcp-80-no-ssh",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"80"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
},
|
||||
{
|
||||
name: "disabled-rule-skipped",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: false, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
},
|
||||
{
|
||||
name: "peer-not-in-destinations",
|
||||
peerSSH: true,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{"g_other"}, // target not in this group
|
||||
},
|
||||
wantEnabled: false,
|
||||
},
|
||||
{
|
||||
name: "peer-typed-destination-resource-matches",
|
||||
peerSSH: false,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
DestinationResource: types.Resource{ID: targetPeerID, Type: types.ResourceTypePeer},
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
{
|
||||
name: "non-peer-destination-resource-falls-through-to-groups",
|
||||
peerSSH: false,
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
DestinationResource: types.Resource{ID: targetPeerID, Type: "host"}, // wrong type
|
||||
Destinations: []string{targetGroupID}, // saved by group fallback
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, peer := mkComponents(&tc.rule, tc.peerSSH)
|
||||
got := computeSSHEnabledForPeer(c, peer)
|
||||
assert.Equal(t, tc.wantEnabled, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestComputeSSHEnabledForPeer_TargetMissingFromComponents covers the
|
||||
// belt-and-suspenders presence guard mirroring Calculate's
|
||||
// getAllPeersFromGroups invariant.
|
||||
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{
|
||||
"g": {ID: "g", Peers: []string{"missing"}},
|
||||
},
|
||||
Policies: []*types.Policy{{
|
||||
ID: "p", Enabled: true,
|
||||
Rules: []*types.PolicyRule{{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{"g"},
|
||||
}},
|
||||
}},
|
||||
}
|
||||
assert.False(t, computeSSHEnabledForPeer(c, peer),
|
||||
"missing target peer must short-circuit to false, not consult policies")
|
||||
}
|
||||
|
||||
// TestComputeSSHEnabledForPeer_NilInputs guards the cheap nil-checks at
|
||||
// function entry — Calculate doesn't accept nil either, but the helper is
|
||||
// exported indirectly via ToComponentSyncResponse and may receive nil
|
||||
// components on graceful-degrade paths.
|
||||
func TestComputeSSHEnabledForPeer_NilInputs(t *testing.T) {
|
||||
assert.False(t, computeSSHEnabledForPeer(nil, &nbpeer.Peer{ID: "x"}))
|
||||
assert.False(t, computeSSHEnabledForPeer(&types.NetworkMapComponents{}, nil))
|
||||
}
|
||||
@@ -10,24 +10,20 @@ import (
|
||||
|
||||
"github.com/hashicorp/go-version"
|
||||
nbversion "github.com/netbirdio/netbird/version"
|
||||
log "github.com/sirupsen/logrus"
|
||||
goproto "google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
|
||||
integrationsConfig "github.com/netbirdio/management-integrations/integrations/config"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ssh/auth"
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
"github.com/netbirdio/netbird/shared/netiputil"
|
||||
"github.com/netbirdio/netbird/shared/sshauth"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -169,8 +165,8 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
|
||||
NetworkMap: &proto.NetworkMap{
|
||||
Serial: networkMap.Network.CurrentSerial(),
|
||||
Routes: toProtocolRoutes(networkMap.Routes),
|
||||
DNSConfig: toProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
|
||||
Routes: networkmap.ToProtocolRoutes(networkMap.Routes),
|
||||
DNSConfig: networkmap.ToProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
|
||||
},
|
||||
Checks: toProtocolChecks(ctx, checks),
|
||||
@@ -183,7 +179,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
|
||||
response.NetworkMap.PeerConfig = response.PeerConfig
|
||||
|
||||
remotePeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers))
|
||||
remotePeers = appendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6)
|
||||
remotePeers = networkmap.AppendRemotePeerConfig(remotePeers, networkMap.Peers, dnsName, includeIPv6)
|
||||
|
||||
if !shouldSkipSendingDeprecatedRemotePeers(peer.Meta.WtVersion) {
|
||||
response.RemotePeers = remotePeers
|
||||
@@ -193,13 +189,13 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
|
||||
response.RemotePeersIsEmpty = len(remotePeers) == 0
|
||||
response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty
|
||||
|
||||
response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6)
|
||||
response.NetworkMap.OfflinePeers = networkmap.AppendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName, includeIPv6)
|
||||
|
||||
firewallRules := toProtocolFirewallRules(networkMap.FirewallRules, includeIPv6, useSourcePrefixes)
|
||||
firewallRules := networkmap.ToProtocolFirewallRules(networkMap.FirewallRules, includeIPv6, useSourcePrefixes)
|
||||
response.NetworkMap.FirewallRules = firewallRules
|
||||
response.NetworkMap.FirewallRulesIsEmpty = len(firewallRules) == 0
|
||||
|
||||
routesFirewallRules := toProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules)
|
||||
routesFirewallRules := networkmap.ToProtocolRoutesFirewallRules(networkMap.RoutesFirewallRules)
|
||||
response.NetworkMap.RoutesFirewallRules = routesFirewallRules
|
||||
response.NetworkMap.RoutesFirewallRulesIsEmpty = len(routesFirewallRules) == 0
|
||||
|
||||
@@ -212,7 +208,7 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
|
||||
}
|
||||
|
||||
if networkMap.AuthorizedUsers != nil {
|
||||
hashedUsers, machineUsers := buildAuthorizedUsersProto(ctx, networkMap.AuthorizedUsers)
|
||||
hashedUsers, machineUsers := networkmap.BuildAuthorizedUsersProto(ctx, networkMap.AuthorizedUsers)
|
||||
userIDClaim := auth.DefaultUserIDClaim
|
||||
if httpConfig != nil && httpConfig.AuthUserIDClaim != "" {
|
||||
userIDClaim = httpConfig.AuthUserIDClaim
|
||||
@@ -252,33 +248,6 @@ func encodeSessionExpiresAt(deadline time.Time) *timestamppb.Timestamp {
|
||||
return timestamppb.New(deadline)
|
||||
}
|
||||
|
||||
func buildAuthorizedUsersProto(ctx context.Context, authorizedUsers map[string]map[string]struct{}) ([][]byte, map[string]*proto.MachineUserIndexes) {
|
||||
userIDToIndex := make(map[string]uint32)
|
||||
var hashedUsers [][]byte
|
||||
machineUsers := make(map[string]*proto.MachineUserIndexes, len(authorizedUsers))
|
||||
|
||||
for machineUser, users := range authorizedUsers {
|
||||
indexes := make([]uint32, 0, len(users))
|
||||
for userID := range users {
|
||||
idx, exists := userIDToIndex[userID]
|
||||
if !exists {
|
||||
hash, err := sshauth.HashUserID(userID)
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to hash user id %s: %v", userID, err)
|
||||
continue
|
||||
}
|
||||
idx = uint32(len(hashedUsers))
|
||||
userIDToIndex[userID] = idx
|
||||
hashedUsers = append(hashedUsers, hash[:])
|
||||
}
|
||||
indexes = append(indexes, idx)
|
||||
}
|
||||
machineUsers[machineUser] = &proto.MachineUserIndexes{Indexes: indexes}
|
||||
}
|
||||
|
||||
return hashedUsers, machineUsers
|
||||
}
|
||||
|
||||
func shouldSkipSendingDeprecatedRemotePeers(peerVersion string) bool {
|
||||
if nbversion.IsDevelopmentVersion(peerVersion) {
|
||||
return true
|
||||
@@ -292,51 +261,6 @@ func shouldSkipSendingDeprecatedRemotePeers(peerVersion string) bool {
|
||||
return precomputedDeprecatedRemotePeersConstraint.Check(peerNBVersion)
|
||||
}
|
||||
|
||||
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() {
|
||||
allowedIPs = append(allowedIPs, rPeer.IPv6.String()+"/128")
|
||||
}
|
||||
dst = append(dst, &proto.RemotePeerConfig{
|
||||
WgPubKey: rPeer.Key,
|
||||
AllowedIps: allowedIPs,
|
||||
SshConfig: &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
|
||||
Fqdn: rPeer.FQDN(dnsName),
|
||||
AgentVersion: rPeer.Meta.WtVersion,
|
||||
})
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
// toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache
|
||||
func toProtocolDNSConfig(update nbdns.Config, cache *cache.DNSConfigCache, forwardPort int64) *proto.DNSConfig {
|
||||
protoUpdate := &proto.DNSConfig{
|
||||
ServiceEnable: update.ServiceEnable,
|
||||
CustomZones: make([]*proto.CustomZone, 0, len(update.CustomZones)),
|
||||
NameServerGroups: make([]*proto.NameServerGroup, 0, len(update.NameServerGroups)),
|
||||
ForwarderPort: forwardPort,
|
||||
}
|
||||
|
||||
for _, zone := range update.CustomZones {
|
||||
protoZone := convertToProtoCustomZone(zone)
|
||||
protoUpdate.CustomZones = append(protoUpdate.CustomZones, protoZone)
|
||||
}
|
||||
|
||||
for _, nsGroup := range update.NameServerGroups {
|
||||
cacheKey := nsGroup.ID
|
||||
if cachedGroup, exists := cache.GetNameServerGroup(cacheKey); exists {
|
||||
protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, cachedGroup)
|
||||
} else {
|
||||
protoGroup := convertToProtoNameServerGroup(nsGroup)
|
||||
cache.SetNameServerGroup(cacheKey, protoGroup)
|
||||
protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, protoGroup)
|
||||
}
|
||||
}
|
||||
|
||||
return protoUpdate
|
||||
}
|
||||
|
||||
func ToResponseProto(configProto nbconfig.Protocol) proto.HostConfig_Protocol {
|
||||
switch configProto {
|
||||
case nbconfig.UDP:
|
||||
@@ -354,203 +278,6 @@ func ToResponseProto(configProto nbconfig.Protocol) proto.HostConfig_Protocol {
|
||||
}
|
||||
}
|
||||
|
||||
func toProtocolRoutes(routes []*nbroute.Route) []*proto.Route {
|
||||
protoRoutes := make([]*proto.Route, 0, len(routes))
|
||||
for _, r := range routes {
|
||||
protoRoutes = append(protoRoutes, toProtocolRoute(r))
|
||||
}
|
||||
return protoRoutes
|
||||
}
|
||||
|
||||
func toProtocolRoute(route *nbroute.Route) *proto.Route {
|
||||
return &proto.Route{
|
||||
ID: string(route.ID),
|
||||
NetID: string(route.NetID),
|
||||
Network: route.Network.String(),
|
||||
Domains: route.Domains.ToPunycodeList(),
|
||||
NetworkType: int64(route.NetworkType),
|
||||
Peer: route.Peer,
|
||||
Metric: int64(route.Metric),
|
||||
Masquerade: route.Masquerade,
|
||||
KeepRoute: route.KeepRoute,
|
||||
SkipAutoApply: route.SkipAutoApply,
|
||||
}
|
||||
}
|
||||
|
||||
// toProtocolFirewallRules converts the firewall rules to the protocol firewall rules.
|
||||
// When useSourcePrefixes is true, the compact SourcePrefixes field is populated
|
||||
// alongside the deprecated PeerIP for forward compatibility.
|
||||
// Wildcard rules ("0.0.0.0") are expanded into separate v4 and v6 SourcePrefixes
|
||||
// when includeIPv6 is true.
|
||||
func toProtocolFirewallRules(rules []*types.FirewallRule, includeIPv6, useSourcePrefixes bool) []*proto.FirewallRule {
|
||||
result := make([]*proto.FirewallRule, 0, len(rules))
|
||||
for i := range rules {
|
||||
rule := rules[i]
|
||||
|
||||
fwRule := &proto.FirewallRule{
|
||||
PolicyID: []byte(rule.PolicyID),
|
||||
PeerIP: rule.PeerIP, //nolint:staticcheck // populated for backward compatibility
|
||||
Direction: getProtoDirection(rule.Direction),
|
||||
Action: getProtoAction(rule.Action),
|
||||
Protocol: getProtoProtocol(rule.Protocol),
|
||||
Port: rule.Port,
|
||||
}
|
||||
|
||||
if useSourcePrefixes && rule.PeerIP != "" {
|
||||
result = append(result, populateSourcePrefixes(fwRule, rule, includeIPv6)...)
|
||||
}
|
||||
|
||||
if shouldUsePortRange(fwRule) {
|
||||
fwRule.PortInfo = rule.PortRange.ToProto()
|
||||
}
|
||||
|
||||
result = append(result, fwRule)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// populateSourcePrefixes sets SourcePrefixes on fwRule and returns any
|
||||
// additional rules needed (e.g. a v6 wildcard clone when the peer IP is unspecified).
|
||||
func populateSourcePrefixes(fwRule *proto.FirewallRule, rule *types.FirewallRule, includeIPv6 bool) []*proto.FirewallRule {
|
||||
addr, err := netip.ParseAddr(rule.PeerIP)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if !addr.IsUnspecified() {
|
||||
fwRule.SourcePrefixes = [][]byte{netiputil.EncodeAddr(addr.Unmap())}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IPv4Unspecified/0 is always valid, error is impossible.
|
||||
v4Wildcard, _ := netiputil.EncodePrefix(netip.PrefixFrom(netip.IPv4Unspecified(), 0))
|
||||
fwRule.SourcePrefixes = [][]byte{v4Wildcard}
|
||||
|
||||
if !includeIPv6 {
|
||||
return nil
|
||||
}
|
||||
|
||||
v6Rule := goproto.Clone(fwRule).(*proto.FirewallRule)
|
||||
v6Rule.PeerIP = "::" //nolint:staticcheck // populated for backward compatibility
|
||||
// IPv6Unspecified/0 is always valid, error is impossible.
|
||||
v6Wildcard, _ := netiputil.EncodePrefix(netip.PrefixFrom(netip.IPv6Unspecified(), 0))
|
||||
v6Rule.SourcePrefixes = [][]byte{v6Wildcard}
|
||||
if shouldUsePortRange(v6Rule) {
|
||||
v6Rule.PortInfo = rule.PortRange.ToProto()
|
||||
}
|
||||
return []*proto.FirewallRule{v6Rule}
|
||||
}
|
||||
|
||||
// getProtoDirection converts the direction to proto.RuleDirection.
|
||||
func getProtoDirection(direction int) proto.RuleDirection {
|
||||
if direction == types.FirewallRuleDirectionOUT {
|
||||
return proto.RuleDirection_OUT
|
||||
}
|
||||
return proto.RuleDirection_IN
|
||||
}
|
||||
|
||||
func toProtocolRoutesFirewallRules(rules []*types.RouteFirewallRule) []*proto.RouteFirewallRule {
|
||||
result := make([]*proto.RouteFirewallRule, len(rules))
|
||||
for i := range rules {
|
||||
rule := rules[i]
|
||||
result[i] = &proto.RouteFirewallRule{
|
||||
SourceRanges: rule.SourceRanges,
|
||||
Action: getProtoAction(rule.Action),
|
||||
Destination: rule.Destination,
|
||||
Protocol: getProtoProtocol(rule.Protocol),
|
||||
PortInfo: getProtoPortInfo(rule),
|
||||
IsDynamic: rule.IsDynamic,
|
||||
Domains: rule.Domains.ToPunycodeList(),
|
||||
PolicyID: []byte(rule.PolicyID),
|
||||
RouteID: string(rule.RouteID),
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// getProtoAction converts the action to proto.RuleAction.
|
||||
func getProtoAction(action string) proto.RuleAction {
|
||||
if action == string(types.PolicyTrafficActionDrop) {
|
||||
return proto.RuleAction_DROP
|
||||
}
|
||||
return proto.RuleAction_ACCEPT
|
||||
}
|
||||
|
||||
// getProtoProtocol converts the protocol to proto.RuleProtocol.
|
||||
func getProtoProtocol(protocol string) proto.RuleProtocol {
|
||||
switch types.PolicyRuleProtocolType(protocol) {
|
||||
case types.PolicyRuleProtocolALL:
|
||||
return proto.RuleProtocol_ALL
|
||||
case types.PolicyRuleProtocolTCP:
|
||||
return proto.RuleProtocol_TCP
|
||||
case types.PolicyRuleProtocolUDP:
|
||||
return proto.RuleProtocol_UDP
|
||||
case types.PolicyRuleProtocolICMP:
|
||||
return proto.RuleProtocol_ICMP
|
||||
default:
|
||||
return proto.RuleProtocol_UNKNOWN
|
||||
}
|
||||
}
|
||||
|
||||
// getProtoPortInfo converts the port info to proto.PortInfo.
|
||||
func getProtoPortInfo(rule *types.RouteFirewallRule) *proto.PortInfo {
|
||||
var portInfo proto.PortInfo
|
||||
if rule.Port != 0 {
|
||||
portInfo.PortSelection = &proto.PortInfo_Port{Port: uint32(rule.Port)}
|
||||
} else if portRange := rule.PortRange; portRange.Start != 0 && portRange.End != 0 {
|
||||
portInfo.PortSelection = &proto.PortInfo_Range_{
|
||||
Range: &proto.PortInfo_Range{
|
||||
Start: uint32(portRange.Start),
|
||||
End: uint32(portRange.End),
|
||||
},
|
||||
}
|
||||
}
|
||||
return &portInfo
|
||||
}
|
||||
|
||||
func shouldUsePortRange(rule *proto.FirewallRule) bool {
|
||||
return rule.Port == "" && (rule.Protocol == proto.RuleProtocol_UDP || rule.Protocol == proto.RuleProtocol_TCP)
|
||||
}
|
||||
|
||||
// Helper function to convert nbdns.CustomZone to proto.CustomZone
|
||||
func convertToProtoCustomZone(zone nbdns.CustomZone) *proto.CustomZone {
|
||||
protoZone := &proto.CustomZone{
|
||||
Domain: zone.Domain,
|
||||
Records: make([]*proto.SimpleRecord, 0, len(zone.Records)),
|
||||
SearchDomainDisabled: zone.SearchDomainDisabled,
|
||||
NonAuthoritative: zone.NonAuthoritative,
|
||||
}
|
||||
for _, record := range zone.Records {
|
||||
protoZone.Records = append(protoZone.Records, &proto.SimpleRecord{
|
||||
Name: record.Name,
|
||||
Type: int64(record.Type),
|
||||
Class: record.Class,
|
||||
TTL: int64(record.TTL),
|
||||
RData: record.RData,
|
||||
})
|
||||
}
|
||||
return protoZone
|
||||
}
|
||||
|
||||
// Helper function to convert nbdns.NameServerGroup to proto.NameServerGroup
|
||||
func convertToProtoNameServerGroup(nsGroup *nbdns.NameServerGroup) *proto.NameServerGroup {
|
||||
protoGroup := &proto.NameServerGroup{
|
||||
Primary: nsGroup.Primary,
|
||||
Domains: nsGroup.Domains,
|
||||
SearchDomainsEnabled: nsGroup.SearchDomainsEnabled,
|
||||
NameServers: make([]*proto.NameServer, 0, len(nsGroup.NameServers)),
|
||||
}
|
||||
for _, ns := range nsGroup.NameServers {
|
||||
protoGroup.NameServers = append(protoGroup.NameServers, &proto.NameServer{
|
||||
IP: ns.IP.String(),
|
||||
Port: int64(ns.Port),
|
||||
NSType: int64(ns.NSType),
|
||||
})
|
||||
}
|
||||
return protoGroup
|
||||
}
|
||||
|
||||
// buildJWTConfig constructs JWT configuration for SSH servers from management server config
|
||||
func buildJWTConfig(config *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow) *proto.JWTConfig {
|
||||
if config == nil || config.AuthAudience == "" {
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
|
||||
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
)
|
||||
|
||||
func TestToProtocolDNSConfigWithCache(t *testing.T) {
|
||||
@@ -64,13 +65,13 @@ func TestToProtocolDNSConfigWithCache(t *testing.T) {
|
||||
}
|
||||
|
||||
// First run with config1
|
||||
result1 := toProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
|
||||
result1 := networkmap.ToProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
|
||||
|
||||
// Second run with config2
|
||||
result2 := toProtocolDNSConfig(config2, &cache, int64(network_map.DnsForwarderPort))
|
||||
result2 := networkmap.ToProtocolDNSConfig(config2, &cache, int64(network_map.DnsForwarderPort))
|
||||
|
||||
// Third run with config1 again
|
||||
result3 := toProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
|
||||
result3 := networkmap.ToProtocolDNSConfig(config1, &cache, int64(network_map.DnsForwarderPort))
|
||||
|
||||
// Verify that result1 and result3 are identical
|
||||
if !reflect.DeepEqual(result1, result3) {
|
||||
@@ -102,15 +103,14 @@ func BenchmarkToProtocolDNSConfig(b *testing.B) {
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
toProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
|
||||
networkmap.ToProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
|
||||
}
|
||||
})
|
||||
|
||||
b.Run(fmt.Sprintf("WithoutCache-Size%d", size), func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
cache := &cache.DNSConfigCache{}
|
||||
toProtocolDNSConfig(testData, cache, int64(network_map.DnsForwarderPort))
|
||||
networkmap.ToProtocolDNSConfig(testData, nil, int64(network_map.DnsForwarderPort))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -608,11 +608,11 @@ func (s *ProxyServiceServer) disconnectProxy(conn *proxyConnection) {
|
||||
if err := s.proxyController.UnregisterProxyFromCluster(context.Background(), conn.address, conn.proxyID); err != nil {
|
||||
log.Warnf("Failed to unregister proxy %s from cluster: %v", conn.proxyID, err)
|
||||
}
|
||||
conn.cancel()
|
||||
if err := s.proxyManager.Disconnect(context.Background(), conn.proxyID, conn.sessionID); err != nil {
|
||||
log.Warnf("Failed to mark proxy %s as disconnected: %v", conn.proxyID, err)
|
||||
}
|
||||
|
||||
conn.cancel()
|
||||
log.Infof("Proxy %s session %s disconnected", conn.proxyID, conn.sessionID)
|
||||
}
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"google.golang.org/grpc/status"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/client/common"
|
||||
"github.com/netbirdio/netbird/shared/management/grpc"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
|
||||
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
||||
@@ -245,6 +246,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
|
||||
realIP := getRealIP(ctx)
|
||||
sRealIP := realIP.String()
|
||||
peerMeta := extractPeerMeta(ctx, syncReq.GetMeta())
|
||||
|
||||
userID, err := s.accountManager.GetUserIDByPeerKey(ctx, peerKey.String())
|
||||
if err != nil {
|
||||
s.syncSem.Add(-1)
|
||||
@@ -683,8 +685,9 @@ func extractPeerMeta(ctx context.Context, meta *proto.PeerSystemMeta) nbpeer.Pee
|
||||
LazyConnectionEnabled: meta.GetFlags().GetLazyConnectionEnabled(),
|
||||
DisableIPv6: meta.GetFlags().GetDisableIPv6(),
|
||||
},
|
||||
Files: files,
|
||||
Capabilities: capabilitiesToInt32(meta.GetCapabilities()),
|
||||
Files: files,
|
||||
Capabilities: capabilitiesToInt32(meta.GetCapabilities()),
|
||||
SyncMessageVersion: int(meta.GetSyncMessageVersion()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1016,7 +1019,43 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
|
||||
return status.Errorf(codes.Internal, "failed to get peer groups %s", err)
|
||||
}
|
||||
|
||||
plainResp := ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, peer, turnToken, relayToken, networkMap, s.networkMapController.GetDNSDomain(settings), postureChecks, nil, settings, settings.Extra, peerGroups, dnsFwdPort)
|
||||
dnsName := s.networkMapController.GetDNSDomain(settings)
|
||||
|
||||
var plainResp *proto.SyncResponse
|
||||
|
||||
commonSyncMessageVersion := grpc.HighestCommonSyncMessageVersion(
|
||||
s.perAccountOrGlobalSyncMessageVersions(peer.AccountID),
|
||||
grpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
|
||||
|
||||
log.WithContext(ctx).
|
||||
WithFields(log.Fields{
|
||||
"sync_message_version": commonSyncMessageVersion,
|
||||
"server_sync_message_version": s.perAccountOrGlobalSyncMessageVersions(peer.AccountID),
|
||||
"peer_sync_message_version": grpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
|
||||
}).Debug("common highest sync message version")
|
||||
|
||||
if commonSyncMessageVersion == grpc.ComponentNetworkMap {
|
||||
// Capable peer: discard the legacy NetworkMap that SyncAndMarkPeer
|
||||
// computed and recompute the raw components instead. This wastes one
|
||||
// Calculate() call per initial-sync — the component-based wire
|
||||
// format is what the peer actually consumes. The streaming path
|
||||
// (network_map.Controller.UpdateAccountPeers) skips this duplication
|
||||
// because it dispatches by capability before computing.
|
||||
//
|
||||
// TODO: refactor SyncPeer / SyncAndMarkPeer / their mocks + manager
|
||||
// interfaces to return PeerNetworkMapResult so the initial-sync path
|
||||
// stops doing duplicate work. Deferred until the client-side
|
||||
// decoder lands and there's a real deployment of capability=3 peers
|
||||
// worth optimizing for.
|
||||
freshPeer, components, proxyPatch, freshPostureChecks, freshDnsFwdPort, err := s.networkMapController.GetValidatedPeerWithComponents(ctx, false, peer.AccountID, peer)
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("failed to build components for peer %s on initial sync: %v", peer.ID, err)
|
||||
return status.Errorf(codes.Internal, "failed to build initial sync envelope")
|
||||
}
|
||||
plainResp = ToComponentSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, freshPeer, turnToken, relayToken, components, proxyPatch, dnsName, freshPostureChecks, settings, settings.Extra, peerGroups, freshDnsFwdPort)
|
||||
} else {
|
||||
plainResp = ToSyncResponse(ctx, s.config, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, peer, turnToken, relayToken, networkMap, dnsName, postureChecks, nil, settings, settings.Extra, peerGroups, dnsFwdPort)
|
||||
}
|
||||
|
||||
key, err := s.secretsManager.GetWGKey()
|
||||
if err != nil {
|
||||
@@ -1041,6 +1080,13 @@ func (s *Server) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, peer
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) perAccountOrGlobalSyncMessageVersions(accountId string) grpc.SyncMessageVersion {
|
||||
if version, ok := s.config.PerAccountHighestSupportedSyncMessageVersion[accountId]; ok {
|
||||
return grpc.SyncMessageVersionFromConfig(&version)
|
||||
}
|
||||
return grpc.SyncMessageVersionFromConfig(s.config.HighestSupportedSyncMessageVersion)
|
||||
}
|
||||
|
||||
// GetDeviceAuthorizationFlow returns a device authorization flow information
|
||||
// This is used for initiating an Oauth 2 device authorization grant flow
|
||||
// which will be used by our clients to Login
|
||||
|
||||
@@ -1648,6 +1648,10 @@ func (am *DefaultAccountManager) SyncUserJWTGroups(ctx context.Context, userAuth
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, g := range newGroupsToCreate {
|
||||
g.PublicID = xid.New().String()
|
||||
}
|
||||
|
||||
if err = transaction.CreateGroups(ctx, userAuth.AccountId, newGroupsToCreate); err != nil {
|
||||
return fmt.Errorf("error saving groups: %w", err)
|
||||
}
|
||||
|
||||
@@ -3170,6 +3170,16 @@ func TestAccount_SetJWTGroups(t *testing.T) {
|
||||
user, err := manager.Store.GetUserByUserID(context.Background(), store.LockingStrengthNone, "user2")
|
||||
assert.NoError(t, err, "unable to get user")
|
||||
assert.Len(t, user.AutoGroups, 1, "new group should be added")
|
||||
|
||||
var newJWTGroup *types.Group
|
||||
for _, g := range groups {
|
||||
if g.Name == "group3" {
|
||||
newJWTGroup = g
|
||||
break
|
||||
}
|
||||
}
|
||||
require.NotNil(t, newJWTGroup, "JIT-created JWT group not found")
|
||||
assert.NotEqual(t, "", newJWTGroup.PublicID, "JIT-created JWT group must have a non-empty PublicID")
|
||||
})
|
||||
|
||||
t.Run("remove all JWT groups when list is empty", func(t *testing.T) {
|
||||
|
||||
@@ -93,6 +93,8 @@ func (am *DefaultAccountManager) CreateGroup(ctx context.Context, accountID, use
|
||||
events := am.prepareGroupEvents(ctx, transaction, accountID, userID, newGroup)
|
||||
eventsToStore = append(eventsToStore, events...)
|
||||
|
||||
newGroup.PublicID = xid.New().String()
|
||||
|
||||
if err := transaction.CreateGroup(ctx, newGroup); err != nil {
|
||||
return status.Errorf(status.Internal, "failed to create group: %v", err)
|
||||
}
|
||||
@@ -158,6 +160,8 @@ func (am *DefaultAccountManager) UpdateGroup(ctx context.Context, accountID, use
|
||||
return err
|
||||
}
|
||||
|
||||
newGroup.PublicID = oldGroup.PublicID
|
||||
|
||||
if err = transaction.UpdateGroup(ctx, newGroup); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -235,6 +239,7 @@ func (am *DefaultAccountManager) CreateGroups(ctx context.Context, accountID, us
|
||||
}
|
||||
|
||||
newGroup.AccountID = accountID
|
||||
newGroup.PublicID = xid.New().String()
|
||||
|
||||
if err = transaction.CreateGroup(ctx, newGroup); err != nil {
|
||||
return err
|
||||
@@ -327,6 +332,12 @@ func (am *DefaultAccountManager) updateSingleGroup(ctx context.Context, accountI
|
||||
|
||||
newGroup.AccountID = accountID
|
||||
|
||||
oldGroup, err := transaction.GetGroupByID(ctx, store.LockingStrengthNone, accountID, newGroup.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
newGroup.PublicID = oldGroup.PublicID
|
||||
|
||||
if err := transaction.UpdateGroup(ctx, newGroup); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -341,7 +352,6 @@ func (am *DefaultAccountManager) updateSingleGroup(ctx context.Context, accountI
|
||||
|
||||
events = am.prepareGroupEvents(ctx, transaction, accountID, userID, newGroup)
|
||||
|
||||
var err error
|
||||
snap, err = affectedpeers.Load(ctx, transaction, accountID, change)
|
||||
return err
|
||||
})
|
||||
|
||||
@@ -6,6 +6,7 @@ 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"
|
||||
@@ -30,6 +31,10 @@ 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 {
|
||||
}
|
||||
|
||||
@@ -109,7 +114,7 @@ func (m *managerImpl) AddResourceToGroupInTransaction(ctx context.Context, trans
|
||||
}
|
||||
|
||||
event := func() {
|
||||
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, group.EventMetaResource(networkResource))
|
||||
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, eventMetaResource(group, networkResource))
|
||||
}
|
||||
|
||||
return event, nil
|
||||
@@ -133,7 +138,7 @@ func (m *managerImpl) RemoveResourceFromGroupInTransaction(ctx context.Context,
|
||||
}
|
||||
|
||||
event := func() {
|
||||
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, group.EventMetaResource(networkResource))
|
||||
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, eventMetaResource(group, networkResource))
|
||||
}
|
||||
|
||||
return event, nil
|
||||
|
||||
@@ -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, dnsDomain))
|
||||
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, account.Peers, dnsDomain))
|
||||
}
|
||||
|
||||
func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -534,15 +534,20 @@ func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request)
|
||||
util.WriteJSONObject(r.Context(), w, resp)
|
||||
}
|
||||
|
||||
func toAccessiblePeers(netMap *types.NetworkMap, dnsDomain string) []api.AccessiblePeer {
|
||||
// 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 {
|
||||
accessiblePeers := make([]api.AccessiblePeer, 0, len(netMap.Peers)+len(netMap.OfflinePeers))
|
||||
for _, p := range netMap.Peers {
|
||||
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
|
||||
}
|
||||
|
||||
for _, p := range netMap.OfflinePeers {
|
||||
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
|
||||
add := func(peers []*types.ComponentPeer) {
|
||||
for _, p := range peers {
|
||||
if peer := accountPeers[p.ID]; peer != nil {
|
||||
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(peer, dnsDomain))
|
||||
}
|
||||
}
|
||||
}
|
||||
add(netMap.Peers)
|
||||
add(netMap.OfflinePeers)
|
||||
|
||||
return accessiblePeers
|
||||
}
|
||||
|
||||
@@ -21,8 +21,11 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
staticClientDashboard = "netbird-dashboard"
|
||||
staticClientCLI = "netbird-cli"
|
||||
StaticClientDashboard = "netbird-dashboard"
|
||||
StaticClientCLI = "netbird-cli"
|
||||
DefaultTOTPAuthenticatorID = "default-totp"
|
||||
LocalConnectorID = dex.LocalConnectorID
|
||||
|
||||
defaultCLIRedirectURL1 = "http://localhost:53000/"
|
||||
defaultCLIRedirectURL2 = "http://localhost:54000/"
|
||||
defaultScopes = "openid profile email groups"
|
||||
@@ -189,14 +192,14 @@ func (c *EmbeddedIdPConfig) ToYAMLConfig() (*dex.YAMLConfig, error) {
|
||||
EnablePasswordDB: true,
|
||||
StaticClients: []storage.Client{
|
||||
{
|
||||
ID: staticClientDashboard,
|
||||
ID: StaticClientDashboard,
|
||||
Name: "NetBird Dashboard",
|
||||
Public: true,
|
||||
RedirectURIs: redirectURIs,
|
||||
PostLogoutRedirectURIs: sanitizePostLogoutRedirectURIs(dashboardPostLogoutRedirectURIs),
|
||||
},
|
||||
{
|
||||
ID: staticClientCLI,
|
||||
ID: StaticClientCLI,
|
||||
Name: "NetBird CLI",
|
||||
Public: true,
|
||||
RedirectURIs: redirectURIs,
|
||||
@@ -258,13 +261,13 @@ func sanitizePostLogoutRedirectURIs(uris []string) []string {
|
||||
|
||||
func configureMFA(cfg *dex.YAMLConfig, sessionMaxLifetime, sessionIdleTimeout string, rememberMe bool, sessionCookieEncryptionKey string) error {
|
||||
cfg.MFA.Authenticators = []dex.MFAAuthenticator{{
|
||||
ID: "default-totp",
|
||||
ID: DefaultTOTPAuthenticatorID,
|
||||
// Has to be caps otherwise it will fail
|
||||
Type: "TOTP",
|
||||
Config: map[string]interface{}{
|
||||
"issuer": "NetBird",
|
||||
},
|
||||
ConnectorTypes: []string{"local"},
|
||||
ConnectorTypes: []string{LocalConnectorID},
|
||||
}}
|
||||
|
||||
if sessionMaxLifetime == "" {
|
||||
@@ -740,7 +743,7 @@ func (m *EmbeddedIdPManager) GetDefaultScopes() string {
|
||||
|
||||
// GetCLIClientID returns the client ID for CLI authentication.
|
||||
func (m *EmbeddedIdPManager) GetCLIClientID() string {
|
||||
return staticClientCLI
|
||||
return StaticClientCLI
|
||||
}
|
||||
|
||||
// GetCLIRedirectURLs returns the redirect URLs configured for the CLI client.
|
||||
@@ -779,7 +782,7 @@ func (m *EmbeddedIdPManager) GetLocalKeysLocation() string {
|
||||
|
||||
// GetClientIDs returns the OAuth2 client IDs configured for this provider.
|
||||
func (m *EmbeddedIdPManager) GetClientIDs() []string {
|
||||
return []string{staticClientDashboard, staticClientCLI}
|
||||
return []string{StaticClientDashboard, StaticClientCLI}
|
||||
}
|
||||
|
||||
// GetUserIDClaim returns the JWT claim name used for user identification.
|
||||
@@ -796,11 +799,11 @@ func (m *EmbeddedIdPManager) IsLocalAuthDisabled() bool {
|
||||
func (m *EmbeddedIdPManager) SetMFAEnabled(ctx context.Context, enabled bool) error {
|
||||
var mfaChain []string
|
||||
if enabled {
|
||||
mfaChain = []string{"default-totp"}
|
||||
mfaChain = []string{DefaultTOTPAuthenticatorID}
|
||||
}
|
||||
if err := m.provider.SetClientsMFAChain(ctx, []string{
|
||||
staticClientCLI,
|
||||
staticClientDashboard,
|
||||
StaticClientCLI,
|
||||
StaticClientDashboard,
|
||||
}, mfaChain); err != nil {
|
||||
return fmt.Errorf("failed to set MFA enabled=%v: %w", enabled, err)
|
||||
}
|
||||
|
||||
@@ -331,7 +331,7 @@ func TestEmbeddedIdPConfig_ToYAMLConfig_IncludesDeviceCallbackRedirectURI(t *tes
|
||||
|
||||
var cliRedirectURIs []string
|
||||
for _, client := range yamlConfig.StaticClients {
|
||||
if client.ID == staticClientCLI {
|
||||
if client.ID == StaticClientCLI {
|
||||
cliRedirectURIs = client.RedirectURIs
|
||||
break
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/rs/xid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
@@ -635,3 +636,128 @@ func RemoveDuplicatePeerKeys(ctx context.Context, db *gorm.DB) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func BackfillPublicIDs[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 backfill needed", model)
|
||||
return nil
|
||||
}
|
||||
|
||||
stmt := &gorm.Statement{DB: db}
|
||||
err := stmt.Parse(&model)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse model: %w", err)
|
||||
}
|
||||
tableName := stmt.Schema.Table
|
||||
|
||||
if err := db.Transaction(func(tx *gorm.DB) error {
|
||||
if !tx.Migrator().HasColumn(&model, "public_id") {
|
||||
log.WithContext(ctx).Infof("Column public_id does not exist in table %s, adding it", tableName)
|
||||
if err := tx.Migrator().AddColumn(&model, "public_id"); err != nil {
|
||||
return fmt.Errorf("add column public_id: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var rows []map[string]any
|
||||
if err := tx.Table(tableName).Select("id", "public_id").Where("public_id IS NULL").Or("public_id = ''").Find(&rows).Error; err != nil {
|
||||
return fmt.Errorf("failed to find rows with empty public_id: %w", err)
|
||||
}
|
||||
|
||||
if len(rows) == 0 {
|
||||
log.WithContext(ctx).Infof("No rows with empty public_id found in table %s, no migration needed", tableName)
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, row := range rows {
|
||||
if err := tx.Table(tableName).Where("id = ?", row["id"]).Update("public_id", xid.New().String()).Error; err != nil {
|
||||
return fmt.Errorf("failed to update row with id %v: %w", row["id"], err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ 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"
|
||||
@@ -639,3 +640,99 @@ 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")
|
||||
}
|
||||
|
||||
@@ -67,6 +67,8 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
|
||||
return err
|
||||
}
|
||||
|
||||
newNSGroup.PublicID = xid.New().String()
|
||||
|
||||
if err = transaction.SaveNameServerGroup(ctx, newNSGroup); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -116,6 +118,8 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
|
||||
return err
|
||||
}
|
||||
|
||||
nsGroupToSave.PublicID = oldNSGroup.PublicID
|
||||
|
||||
if err = transaction.SaveNameServerGroup(ctx, nsGroupToSave); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user