mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-29 11:51:27 +02:00
Compare commits
32 Commits
revert/com
...
test/gui-m
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c81eab8472 | ||
|
|
c49333e48a | ||
|
|
8b4061cfaf | ||
|
|
175ca4f0e0 | ||
|
|
859ac60c90 | ||
|
|
dd2bdc0de3 | ||
|
|
44fef45c2f | ||
|
|
3d1f209ea3 | ||
|
|
2ef457be95 | ||
|
|
28ff2b7b3b | ||
|
|
22e0b66412 | ||
|
|
63c320b6a9 | ||
|
|
b6deed39fb | ||
|
|
4c95b20251 | ||
|
|
4b915e4a5a | ||
|
|
4acbe2670a | ||
|
|
0fb4c8c423 | ||
|
|
42e45ff9f9 | ||
|
|
e1d82fc326 | ||
|
|
29b97340d6 | ||
|
|
9269b56386 | ||
|
|
b3f9b82442 | ||
|
|
8a43f4f943 | ||
|
|
2f268c8141 | ||
|
|
e3c4128164 | ||
|
|
bab5572a74 | ||
|
|
9b4a5df925 | ||
|
|
1816a020c4 | ||
|
|
aa13928b76 | ||
|
|
d681670a9d | ||
|
|
4f6247b5c3 | ||
|
|
1e5b0a5c89 |
9
.github/workflows/agent-network-e2e.yml
vendored
9
.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 }}
|
||||
@@ -62,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 }}
|
||||
|
||||
@@ -273,8 +273,8 @@ dockers_v2:
|
||||
- netbirdio/netbird
|
||||
- ghcr.io/netbirdio/netbird
|
||||
tags:
|
||||
- "v{{ .Version }}-rootless"
|
||||
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}latest{{ end }}"
|
||||
- "{{ .Version }}-rootless"
|
||||
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}rootless-latest{{ end }}"
|
||||
dockerfile: client/Dockerfile-rootless
|
||||
extra_files:
|
||||
- client/netbird-entrypoint.sh
|
||||
|
||||
@@ -24,6 +24,8 @@ builds:
|
||||
ldflags:
|
||||
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
|
||||
mod_timestamp: "{{ .CommitTimestamp }}"
|
||||
tags:
|
||||
- production
|
||||
|
||||
- id: netbird-ui-windows-amd64
|
||||
dir: client/ui
|
||||
@@ -39,6 +41,8 @@ builds:
|
||||
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
|
||||
- -H windowsgui
|
||||
mod_timestamp: "{{ .CommitTimestamp }}"
|
||||
tags:
|
||||
- production
|
||||
|
||||
- id: netbird-ui-windows-arm64
|
||||
dir: client/ui
|
||||
@@ -55,6 +59,8 @@ builds:
|
||||
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
|
||||
- -H windowsgui
|
||||
mod_timestamp: "{{ .CommitTimestamp }}"
|
||||
tags:
|
||||
- production
|
||||
|
||||
archives:
|
||||
- id: linux-arch
|
||||
|
||||
@@ -29,6 +29,8 @@ builds:
|
||||
ldflags:
|
||||
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
|
||||
mod_timestamp: "{{ .CommitTimestamp }}"
|
||||
tags:
|
||||
- production
|
||||
|
||||
universal_binaries:
|
||||
- id: netbird-ui-darwin
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -299,6 +300,13 @@ func (c *Client) SetInfoLogLevel() {
|
||||
// PeersList return with the list of the PeerInfos
|
||||
func (c *Client) PeersList() *PeerInfoArray {
|
||||
|
||||
// The recorder only caches transfer counters and handshake times; nothing
|
||||
// refreshes them on its own, so without this they read as zero. The desktop
|
||||
// daemon does the same before serving a full peer status.
|
||||
if err := c.recorder.RefreshWireGuardStats(); err != nil {
|
||||
log.Debugf("failed to refresh WireGuard stats: %v", err)
|
||||
}
|
||||
|
||||
fullStatus := c.recorder.GetFullStatus()
|
||||
|
||||
peerInfos := make([]PeerInfo, len(fullStatus.Peers))
|
||||
@@ -309,6 +317,20 @@ func (c *Client) PeersList() *PeerInfoArray {
|
||||
FQDN: p.FQDN,
|
||||
ConnStatus: int(p.ConnStatus),
|
||||
Routes: PeerRoutes{routes: maps.Keys(p.GetRoutes())},
|
||||
|
||||
PubKey: p.PubKey,
|
||||
Latency: formatDuration(p.Latency),
|
||||
LatencyMs: p.Latency.Milliseconds(),
|
||||
BytesRx: p.BytesRx,
|
||||
BytesTx: p.BytesTx,
|
||||
ConnStatusUpdate: formatTime(p.ConnStatusUpdate),
|
||||
Relayed: p.Relayed,
|
||||
RosenpassEnabled: p.RosenpassEnabled,
|
||||
LastWireguardHandshake: formatTime(p.LastWireguardHandshake),
|
||||
LocalIceCandidateType: p.LocalIceCandidateType,
|
||||
RemoteIceCandidateType: p.RemoteIceCandidateType,
|
||||
LocalIceCandidateEndpoint: p.LocalIceCandidateEndpoint,
|
||||
RemoteIceCandidateEndpoint: p.RemoteIceCandidateEndpoint,
|
||||
}
|
||||
peerInfos[n] = pi
|
||||
}
|
||||
@@ -439,10 +461,6 @@ func (c *Client) RemoveConnectionListener() {
|
||||
c.recorder.RemoveConnectionListener()
|
||||
}
|
||||
|
||||
func (c *Client) toggleRoute(command routeCommand) error {
|
||||
return command.toggleRoute()
|
||||
}
|
||||
|
||||
func (c *Client) getRouteManager() (routemanager.Manager, error) {
|
||||
client := c.getConnectClient()
|
||||
if client == nil {
|
||||
@@ -462,22 +480,22 @@ func (c *Client) getRouteManager() (routemanager.Manager, error) {
|
||||
return manager, nil
|
||||
}
|
||||
|
||||
func (c *Client) SelectRoute(route string) error {
|
||||
func (c *Client) SelectRoute(id string) error {
|
||||
manager, err := c.getRouteManager()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return c.toggleRoute(selectRouteCommand{route: route, manager: manager})
|
||||
return manager.SelectRoutes([]route.NetID{route.NetID(id)}, true)
|
||||
}
|
||||
|
||||
func (c *Client) DeselectRoute(route string) error {
|
||||
func (c *Client) DeselectRoute(id string) error {
|
||||
manager, err := c.getRouteManager()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return c.toggleRoute(deselectRouteCommand{route: route, manager: manager})
|
||||
return manager.DeselectRoutes([]route.NetID{route.NetID(id)})
|
||||
}
|
||||
|
||||
// getNetworkDomainsFromRoute extracts domains from a route and enriches each domain
|
||||
@@ -512,3 +530,28 @@ func exportEnvList(list *EnvList) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// formatDuration renders a duration for display, trimming the fractional part
|
||||
// to two digits so latencies read as "12.34ms" rather than "12.345678ms".
|
||||
func formatDuration(d time.Duration) string {
|
||||
ds := d.String()
|
||||
dotIndex := strings.Index(ds, ".")
|
||||
if dotIndex == -1 {
|
||||
return ds
|
||||
}
|
||||
|
||||
endIndex := min(dotIndex+3, len(ds))
|
||||
|
||||
// Skip the remaining digits so only the unit suffix is appended back.
|
||||
unitStart := endIndex
|
||||
for unitStart < len(ds) && ds[unitStart] >= '0' && ds[unitStart] <= '9' {
|
||||
unitStart++
|
||||
}
|
||||
return ds[:endIndex] + ds[unitStart:]
|
||||
}
|
||||
|
||||
// formatTime renders a timestamp in UTC using a fixed layout. The zero time is
|
||||
// passed through as-is so the UI can recognise it and show "never" instead.
|
||||
func formatTime(t time.Time) string {
|
||||
return t.UTC().Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
||||
@@ -12,12 +12,30 @@ const (
|
||||
)
|
||||
|
||||
// PeerInfo describe information about the peers. It designed for the UI usage
|
||||
//
|
||||
// The fields below ConnStatus back the peer detail screen. Durations and times
|
||||
// are pre-formatted into strings so the UI does not have to know Go's layouts;
|
||||
// Latency is additionally exposed as LatencyMs for colour coding.
|
||||
type PeerInfo struct {
|
||||
IP string
|
||||
IPv6 string
|
||||
FQDN string
|
||||
ConnStatus int
|
||||
Routes PeerRoutes
|
||||
|
||||
PubKey string
|
||||
Latency string
|
||||
LatencyMs int64
|
||||
BytesRx int64
|
||||
BytesTx int64
|
||||
ConnStatusUpdate string
|
||||
Relayed bool
|
||||
RosenpassEnabled bool
|
||||
LastWireguardHandshake string
|
||||
LocalIceCandidateType string
|
||||
RemoteIceCandidateType string
|
||||
LocalIceCandidateEndpoint string
|
||||
RemoteIceCandidateEndpoint string
|
||||
}
|
||||
|
||||
func (p *PeerInfo) GetPeerRoutes() *PeerRoutes {
|
||||
|
||||
@@ -189,6 +189,19 @@ func (pm *ProfileManager) LogoutProfile(id string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// RenameProfile changes a profile's display name. The profile ID, and therefore
|
||||
// its on-disk filename, is left untouched: only the "name" field of the config
|
||||
// is rewritten. This works for the default profile too, whose config lives in
|
||||
// netbird.cfg rather than under profiles/.
|
||||
func (pm *ProfileManager) RenameProfile(id string, newName string) error {
|
||||
if err := pm.serviceMgr.RenameProfile(profilemanager.ID(id), androidUsername, newName); err != nil {
|
||||
return fmt.Errorf("failed to rename profile: %w", err)
|
||||
}
|
||||
|
||||
log.Infof("renamed profile %s to: %s", id, newName)
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemoveProfile deletes a profile
|
||||
func (pm *ProfileManager) RemoveProfile(id string) error {
|
||||
// Use ServiceManager (removes profile from profiles/ directory)
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
//go:build android
|
||||
|
||||
package android
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/exp/maps"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/routemanager"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
func executeRouteToggle(id string, manager routemanager.Manager,
|
||||
operationName string,
|
||||
routeOperation func(routes []route.NetID, allRoutes []route.NetID) error) error {
|
||||
netID := route.NetID(id)
|
||||
routes := []route.NetID{netID}
|
||||
|
||||
routesMap := manager.GetClientRoutesWithNetID()
|
||||
routes = route.ExpandV6ExitPairs(routes, routesMap)
|
||||
|
||||
log.Debugf("%s with ids: %v", operationName, routes)
|
||||
|
||||
if err := routeOperation(routes, maps.Keys(routesMap)); err != nil {
|
||||
log.Debugf("error when %s: %s", operationName, err)
|
||||
return fmt.Errorf("error %s: %w", operationName, err)
|
||||
}
|
||||
|
||||
manager.TriggerSelection(manager.GetClientRoutes())
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type routeCommand interface {
|
||||
toggleRoute() error
|
||||
}
|
||||
|
||||
type selectRouteCommand struct {
|
||||
route string
|
||||
manager routemanager.Manager
|
||||
}
|
||||
|
||||
func (s selectRouteCommand) toggleRoute() error {
|
||||
routeSelector := s.manager.GetRouteSelector()
|
||||
if routeSelector == nil {
|
||||
return fmt.Errorf("no route selector available")
|
||||
}
|
||||
|
||||
routeOperation := func(routes []route.NetID, allRoutes []route.NetID) error {
|
||||
return routeSelector.SelectRoutes(routes, true, allRoutes)
|
||||
}
|
||||
|
||||
return executeRouteToggle(s.route, s.manager, "selecting route", routeOperation)
|
||||
}
|
||||
|
||||
type deselectRouteCommand struct {
|
||||
route string
|
||||
manager routemanager.Manager
|
||||
}
|
||||
|
||||
func (d deselectRouteCommand) toggleRoute() error {
|
||||
routeSelector := d.manager.GetRouteSelector()
|
||||
if routeSelector == nil {
|
||||
return fmt.Errorf("no route selector available")
|
||||
}
|
||||
|
||||
return executeRouteToggle(d.route, d.manager, "deselecting route", routeSelector.DeselectRoutes)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -621,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")
|
||||
}
|
||||
|
||||
@@ -385,11 +385,20 @@ func inactivityThresholdEnv() *time.Duration {
|
||||
return nil
|
||||
}
|
||||
|
||||
parsedMinutes, err := strconv.Atoi(envValue)
|
||||
if err != nil || parsedMinutes <= 0 {
|
||||
return nil
|
||||
// Documented format: a Go duration such as "30m" or "1h".
|
||||
if d, err := time.ParseDuration(envValue); err == nil {
|
||||
if d <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &d
|
||||
}
|
||||
|
||||
d := time.Duration(parsedMinutes) * time.Minute
|
||||
return &d
|
||||
// Backwards compatibility: a bare integer used to be interpreted as minutes.
|
||||
if parsedMinutes, err := strconv.Atoi(envValue); err == nil && parsedMinutes > 0 {
|
||||
d := time.Duration(parsedMinutes) * time.Minute
|
||||
return &d
|
||||
}
|
||||
|
||||
log.Warnf("invalid %s value %q: expected a Go duration such as 30m or 1h", lazyconn.EnvInactivityThreshold, envValue)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -104,3 +104,38 @@ func TestConnMgr_ActivatePeerConcurrentWithLifecycle(t *testing.T) {
|
||||
close(done)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestInactivityThresholdEnv(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
val string
|
||||
want *time.Duration
|
||||
}{
|
||||
{name: "unset", val: "", want: nil},
|
||||
{name: "go duration minutes", val: "30m", want: durPtr(30 * time.Minute)},
|
||||
{name: "go duration hours", val: "1h", want: durPtr(time.Hour)},
|
||||
{name: "go duration seconds", val: "90s", want: durPtr(90 * time.Second)},
|
||||
{name: "bare integer is minutes (backwards compat)", val: "5", want: durPtr(5 * time.Minute)},
|
||||
{name: "zero duration", val: "0s", want: nil},
|
||||
{name: "zero integer", val: "0", want: nil},
|
||||
{name: "negative duration", val: "-5m", want: nil},
|
||||
{name: "garbage", val: "abc", want: nil},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Setenv(lazyconn.EnvInactivityThreshold, tc.val)
|
||||
got := inactivityThresholdEnv()
|
||||
switch {
|
||||
case tc.want == nil && got != nil:
|
||||
t.Fatalf("want nil, got %v", *got)
|
||||
case tc.want != nil && got == nil:
|
||||
t.Fatalf("want %v, got nil", *tc.want)
|
||||
case tc.want != nil && *got != *tc.want:
|
||||
t.Fatalf("want %v, got %v", *tc.want, *got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func durPtr(d time.Duration) *time.Duration { return &d }
|
||||
|
||||
@@ -34,6 +34,7 @@ import (
|
||||
"github.com/netbirdio/netbird/client/internal/profilemanager"
|
||||
"github.com/netbirdio/netbird/client/internal/statemanager"
|
||||
"github.com/netbirdio/netbird/client/internal/stdnet"
|
||||
"github.com/netbirdio/netbird/client/internal/tunnelnotifier"
|
||||
"github.com/netbirdio/netbird/client/internal/updater"
|
||||
"github.com/netbirdio/netbird/client/internal/updater/installer"
|
||||
nbnet "github.com/netbirdio/netbird/client/net"
|
||||
@@ -136,10 +137,13 @@ func (c *ConnectClient) RunOniOS(
|
||||
// Set GC percent to 5% to reduce memory usage as iOS only allows 50MB of memory for the extension.
|
||||
debug.SetGCPercent(5)
|
||||
|
||||
notifier := tunnelnotifier.New(networkChangeListener, dnsManager)
|
||||
defer notifier.Close()
|
||||
|
||||
mobileDependency := MobileDependency{
|
||||
FileDescriptor: fileDescriptor,
|
||||
NetworkChangeListener: networkChangeListener,
|
||||
DnsManager: dnsManager,
|
||||
NetworkChangeListener: notifier,
|
||||
DnsManager: notifier,
|
||||
StateFilePath: stateFilePath,
|
||||
TempDir: cacheDir,
|
||||
}
|
||||
|
||||
@@ -11,12 +11,14 @@ import (
|
||||
|
||||
// MobileDependency collect all dependencies for mobile platform
|
||||
type MobileDependency struct {
|
||||
// Android only
|
||||
TunAdapter device.TunAdapter
|
||||
IFaceDiscover stdnet.ExternalIFaceDiscover
|
||||
// Android and iOS
|
||||
NetworkChangeListener listener.NetworkChangeListener
|
||||
HostDNSAddresses []netip.AddrPort
|
||||
DnsReadyListener dns.ReadyListener
|
||||
|
||||
// Android only
|
||||
TunAdapter device.TunAdapter
|
||||
IFaceDiscover stdnet.ExternalIFaceDiscover
|
||||
HostDNSAddresses []netip.AddrPort
|
||||
DnsReadyListener dns.ReadyListener
|
||||
|
||||
// iOS only
|
||||
DnsManager dns.IosDnsManager
|
||||
|
||||
@@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
|
||||
|
||||
// AllowedIPs should use real IPs
|
||||
if d.currentPeerKey != "" {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
|
||||
}
|
||||
}
|
||||
@@ -172,7 +172,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error {
|
||||
}
|
||||
|
||||
// AllowedIPs use real IPs
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil {
|
||||
return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err)
|
||||
}
|
||||
|
||||
@@ -205,7 +205,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error {
|
||||
for _, prefixes := range d.interceptedDomains {
|
||||
for _, prefix := range prefixes {
|
||||
// AllowedIPs use real IPs
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error {
|
||||
var merr *multierror.Error
|
||||
for _, domainPrefixes := range r.dynamicDomains {
|
||||
for _, prefix := range domainPrefixes {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
|
||||
}
|
||||
}
|
||||
@@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err))
|
||||
}
|
||||
if r.currentPeerKey != "" {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -52,6 +52,10 @@ type Manager interface {
|
||||
UpdateRoutes(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
|
||||
ClassifyRoutes(newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
|
||||
TriggerSelection(route.HAMap)
|
||||
SelectRoutes(ids []route.NetID, appendRoute bool) error
|
||||
DeselectRoutes(ids []route.NetID) error
|
||||
SelectAllRoutes()
|
||||
DeselectAllRoutes()
|
||||
GetRouteSelector() *routeselector.RouteSelector
|
||||
GetClientRoutes() route.HAMap
|
||||
GetSelectedClientRoutes() route.HAMap
|
||||
@@ -216,7 +220,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
|
||||
)
|
||||
}
|
||||
|
||||
m.allowedIPsRefCounter = refcounter.New(
|
||||
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
|
||||
func(prefix netip.Prefix, peerKey string) (string, error) {
|
||||
// save peerKey to use it in the remove function
|
||||
return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix)
|
||||
@@ -800,7 +804,7 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
|
||||
var info exitNodeInfo
|
||||
|
||||
for haID, routes := range clientRoutes {
|
||||
if !m.isExitNodeRoute(routes) {
|
||||
if !isExitNodeRoutes(routes) {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -820,13 +824,6 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
|
||||
return info
|
||||
}
|
||||
|
||||
func (m *DefaultManager) isExitNodeRoute(routes []*route.Route) bool {
|
||||
if len(routes) == 0 {
|
||||
return false
|
||||
}
|
||||
return route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network)
|
||||
}
|
||||
|
||||
func (m *DefaultManager) categorizeUserSelection(netID route.NetID, info *exitNodeInfo) {
|
||||
if m.routeSelector.IsSelected(netID) {
|
||||
info.userSelected = append(info.userSelected, netID)
|
||||
|
||||
@@ -16,6 +16,8 @@ type MockManager struct {
|
||||
ClassifyRoutesFunc func(routes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
|
||||
UpdateRoutesFunc func(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
|
||||
TriggerSelectionFunc func(haMap route.HAMap)
|
||||
SelectRoutesFunc func(ids []route.NetID, appendRoute bool) error
|
||||
DeselectRoutesFunc func(ids []route.NetID) error
|
||||
GetRouteSelectorFunc func() *routeselector.RouteSelector
|
||||
GetClientRoutesFunc func() route.HAMap
|
||||
GetSelectedClientRoutesFunc func() route.HAMap
|
||||
@@ -55,6 +57,30 @@ func (m *MockManager) TriggerSelection(networks route.HAMap) {
|
||||
}
|
||||
}
|
||||
|
||||
// SelectRoutes mock implementation of SelectRoutes from Manager interface
|
||||
func (m *MockManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
|
||||
if m.SelectRoutesFunc != nil {
|
||||
return m.SelectRoutesFunc(ids, appendRoute)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeselectRoutes mock implementation of DeselectRoutes from Manager interface
|
||||
func (m *MockManager) DeselectRoutes(ids []route.NetID) error {
|
||||
if m.DeselectRoutesFunc != nil {
|
||||
return m.DeselectRoutesFunc(ids)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SelectAllRoutes mock implementation of SelectAllRoutes from Manager interface
|
||||
func (m *MockManager) SelectAllRoutes() {
|
||||
}
|
||||
|
||||
// DeselectAllRoutes mock implementation of DeselectAllRoutes from Manager interface
|
||||
func (m *MockManager) DeselectAllRoutes() {
|
||||
}
|
||||
|
||||
// GetRouteSelector mock implementation of GetRouteSelector from Manager interface
|
||||
func (m *MockManager) GetRouteSelector() *routeselector.RouteSelector {
|
||||
if m.GetRouteSelectorFunc != nil {
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
package notifier
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"net/netip"
|
||||
"slices"
|
||||
"sort"
|
||||
@@ -16,20 +15,12 @@ import (
|
||||
|
||||
type Notifier struct {
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
currentPrefixes []string
|
||||
listener listener.NetworkChangeListener
|
||||
queue *list.List
|
||||
closed bool
|
||||
}
|
||||
|
||||
func NewNotifier() *Notifier {
|
||||
n := &Notifier{
|
||||
queue: list.New(),
|
||||
}
|
||||
n.cond = sync.NewCond(&n.mu)
|
||||
go n.deliverLoop()
|
||||
return n
|
||||
return &Notifier{}
|
||||
}
|
||||
|
||||
func (n *Notifier) SetListener(listener listener.NetworkChangeListener) {
|
||||
@@ -59,44 +50,19 @@ func (n *Notifier) OnNewPrefixes(prefixes []netip.Prefix) {
|
||||
sort.Strings(newNets)
|
||||
|
||||
n.mu.Lock()
|
||||
defer n.mu.Unlock()
|
||||
if slices.Equal(n.currentPrefixes, newNets) {
|
||||
n.mu.Unlock()
|
||||
return
|
||||
}
|
||||
n.currentPrefixes = newNets
|
||||
routes := strings.Join(n.currentPrefixes, ",")
|
||||
n.queue.PushBack(routes)
|
||||
n.cond.Signal()
|
||||
n.mu.Unlock()
|
||||
if n.listener != nil {
|
||||
n.listener.OnNetworkChanged(strings.Join(n.currentPrefixes, ","))
|
||||
}
|
||||
}
|
||||
|
||||
func (n *Notifier) Close() {
|
||||
n.mu.Lock()
|
||||
n.closed = true
|
||||
n.cond.Signal()
|
||||
n.mu.Unlock()
|
||||
}
|
||||
|
||||
func (n *Notifier) GetInitialRouteRanges() []string {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (n *Notifier) deliverLoop() {
|
||||
for {
|
||||
n.mu.Lock()
|
||||
for n.queue.Len() == 0 && !n.closed {
|
||||
n.cond.Wait()
|
||||
}
|
||||
if n.closed && n.queue.Len() == 0 {
|
||||
n.mu.Unlock()
|
||||
return
|
||||
}
|
||||
routes := n.queue.Remove(n.queue.Front()).(string)
|
||||
l := n.listener
|
||||
n.mu.Unlock()
|
||||
|
||||
if l != nil {
|
||||
l.OnNetworkChanged(routes)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,7 +54,7 @@ func (m *reconcileWGMock) GetNet() *netstack.Net { return n
|
||||
func TestReconcilePeerAllowedIPs(t *testing.T) {
|
||||
wg := &reconcileWGMock{}
|
||||
m := &DefaultManager{wgInterface: wg}
|
||||
m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
|
||||
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
|
||||
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
|
||||
func(netip.Prefix, string) error { return nil },
|
||||
)
|
||||
|
||||
206
client/internal/routemanager/refcounter/allowedips.go
Normal file
206
client/internal/routemanager/refcounter/allowedips.go
Normal file
@@ -0,0 +1,206 @@
|
||||
package refcounter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/netip"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/hashicorp/go-multierror"
|
||||
|
||||
nberrors "github.com/netbirdio/netbird/client/errors"
|
||||
)
|
||||
|
||||
// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is
|
||||
// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most
|
||||
// one peer is active at a time even when several peers reference the prefix.
|
||||
type allowedIPsEntry struct {
|
||||
// peers maps a peerKey to the number of references holding the prefix for that peer.
|
||||
peers map[string]int
|
||||
// active is the peerKey currently installed in WireGuard for this prefix ("" if none).
|
||||
active string
|
||||
// total is the sum of all per-peer reference counts (kept in sync with peers).
|
||||
total int
|
||||
}
|
||||
|
||||
// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs.
|
||||
//
|
||||
// The generic Counter keys only by prefix and remembers a single Out value set by the first
|
||||
// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or
|
||||
// multiple resolved domains) can reference the same prefix through different peers, and when the
|
||||
// peer currently installed in WireGuard releases its last reference the prefix must be handed over
|
||||
// to a surviving peer instead of being left pointing at the released one.
|
||||
//
|
||||
// It calls add/remove (which program WireGuard) only on the transitions that matter:
|
||||
// - add on the first reference for a prefix, or when swapping the active peer;
|
||||
// - remove on the last reference for a prefix, or on the old peer during a swap.
|
||||
type AllowedIPsRefCounter struct {
|
||||
mu sync.Mutex
|
||||
entries map[netip.Prefix]*allowedIPsEntry
|
||||
add AddFunc[netip.Prefix, string, string]
|
||||
remove RemoveFunc[netip.Prefix, string]
|
||||
}
|
||||
|
||||
// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter.
|
||||
// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer.
|
||||
// remove unprograms the prefix from the given peer.
|
||||
func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter {
|
||||
return &AllowedIPsRefCounter{
|
||||
entries: map[netip.Prefix]*allowedIPsEntry{},
|
||||
add: add,
|
||||
remove: remove,
|
||||
}
|
||||
}
|
||||
|
||||
// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first
|
||||
// reference to a prefix; while a different peer is already installed the prefix is left with it
|
||||
// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept.
|
||||
func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
e, ok := rm.entries[prefix]
|
||||
if !ok {
|
||||
e = &allowedIPsEntry{peers: map[string]int{}}
|
||||
rm.entries[prefix] = e
|
||||
}
|
||||
|
||||
logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
|
||||
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active)
|
||||
|
||||
// Program WireGuard only when nothing is installed yet for this prefix.
|
||||
if e.active == "" {
|
||||
out, err := rm.add(prefix, peerKey)
|
||||
if errors.Is(err, ErrIgnore) {
|
||||
if e.total == 0 {
|
||||
delete(rm.entries, prefix)
|
||||
}
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
if err != nil {
|
||||
if e.total == 0 {
|
||||
delete(rm.entries, prefix)
|
||||
}
|
||||
return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err)
|
||||
}
|
||||
e.active = out
|
||||
}
|
||||
|
||||
e.peers[peerKey]++
|
||||
e.total++
|
||||
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
// Decrement removes a reference to prefix for peerKey. When the peer currently installed in
|
||||
// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists,
|
||||
// otherwise it is removed from WireGuard.
|
||||
func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
e, ok := rm.entries[prefix]
|
||||
if !ok {
|
||||
logCallerF("No allowed IP reference found for prefix %v", prefix)
|
||||
return Ref[string]{}, nil
|
||||
}
|
||||
|
||||
if e.peers[peerKey] > 0 {
|
||||
logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
|
||||
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active)
|
||||
e.peers[peerKey]--
|
||||
e.total--
|
||||
if e.peers[peerKey] == 0 {
|
||||
delete(e.peers, peerKey)
|
||||
}
|
||||
} else {
|
||||
logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey)
|
||||
}
|
||||
|
||||
// If the peer currently installed in WireGuard still holds references, nothing to reprogram.
|
||||
// Keying the check on the active peer (not the one just released) makes this self-healing:
|
||||
// a prior swap whose remove/add failed leaves e.active pointing at a peer with no references,
|
||||
// and this retries the hand-off on the next Decrement instead of getting stuck.
|
||||
if e.active != "" && e.peers[e.active] > 0 {
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
// Detach the stale/gone active peer from WireGuard before reprogramming.
|
||||
if e.active != "" {
|
||||
if err := rm.remove(prefix, e.active); err != nil {
|
||||
return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)
|
||||
}
|
||||
e.active = ""
|
||||
}
|
||||
|
||||
// Hand the prefix over to a surviving peer, or drop the entry when none remain.
|
||||
if survivor, ok := pickSurvivor(e.peers); ok {
|
||||
out, err := rm.add(prefix, survivor)
|
||||
if err != nil {
|
||||
return Ref[string]{Count: e.total, Out: ""}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err)
|
||||
}
|
||||
e.active = out
|
||||
return Ref[string]{Count: e.total, Out: e.active}, nil
|
||||
}
|
||||
|
||||
delete(rm.entries, prefix)
|
||||
return Ref[string]{Count: 0, Out: ""}, nil
|
||||
}
|
||||
|
||||
// Flush removes all prefixes from WireGuard and clears the counter.
|
||||
func (rm *AllowedIPsRefCounter) Flush() error {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
var merr *multierror.Error
|
||||
for prefix, e := range rm.entries {
|
||||
if e.active == "" {
|
||||
continue
|
||||
}
|
||||
logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active)
|
||||
if err := rm.remove(prefix, e.active); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err))
|
||||
}
|
||||
}
|
||||
|
||||
clear(rm.entries)
|
||||
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
// ReapplyMatching calls apply for every prefix whose currently installed (active) peer satisfies
|
||||
// pred, holding the lock for the whole pass. It is used to re-push allowed IPs onto a peer whose
|
||||
// WireGuard entry was rebuilt (e.g. a lazy connection cycling idle->wake) without a matching
|
||||
// refcounter change, which would otherwise leave the prefix installed in the counter but missing
|
||||
// on the device. Only the active peer is considered — a prefix that lost its installed peer to a
|
||||
// failed swap is skipped here and reconciled by the next Increment/Decrement.
|
||||
func (rm *AllowedIPsRefCounter) ReapplyMatching(pred func(out string) bool, apply func(key netip.Prefix) error) error {
|
||||
rm.mu.Lock()
|
||||
defer rm.mu.Unlock()
|
||||
|
||||
var merr *multierror.Error
|
||||
for prefix, e := range rm.entries {
|
||||
if e.active != "" && pred(e.active) {
|
||||
if err := apply(prefix); err != nil {
|
||||
merr = multierror.Append(merr, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do
|
||||
// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable
|
||||
// (lowest peerKey) for predictable behavior and testability.
|
||||
func pickSurvivor(peers map[string]int) (string, bool) {
|
||||
if len(peers) == 0 {
|
||||
return "", false
|
||||
}
|
||||
keys := make([]string, 0, len(peers))
|
||||
for k := range peers {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys[0], true
|
||||
}
|
||||
241
client/internal/routemanager/refcounter/allowedips_test.go
Normal file
241
client/internal/routemanager/refcounter/allowedips_test.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package refcounter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/netip"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer.
|
||||
// failAdd/failRemove make the next add/remove fail once, to exercise the self-healing error paths.
|
||||
type fakeWG struct {
|
||||
installed map[netip.Prefix]string
|
||||
adds int
|
||||
removes int
|
||||
failAdd bool
|
||||
failRemove bool
|
||||
}
|
||||
|
||||
func newFakeWG() *fakeWG {
|
||||
return &fakeWG{installed: map[netip.Prefix]string{}}
|
||||
}
|
||||
|
||||
func (f *fakeWG) counter() *AllowedIPsRefCounter {
|
||||
return NewAllowedIPs(
|
||||
func(prefix netip.Prefix, peerKey string) (string, error) {
|
||||
if f.failAdd {
|
||||
f.failAdd = false
|
||||
return "", errors.New("add failed")
|
||||
}
|
||||
f.adds++
|
||||
f.installed[prefix] = peerKey
|
||||
return peerKey, nil
|
||||
},
|
||||
func(prefix netip.Prefix, peerKey string) error {
|
||||
if f.failRemove {
|
||||
f.failRemove = false
|
||||
return errors.New("remove failed")
|
||||
}
|
||||
f.removes++
|
||||
// only clear if this peer is the one installed, mirroring wg semantics
|
||||
if f.installed[prefix] == peerKey {
|
||||
delete(f.installed, prefix)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func mustPrefix(t *testing.T, s string) netip.Prefix {
|
||||
t.Helper()
|
||||
p, err := netip.ParsePrefix(s)
|
||||
if err != nil {
|
||||
t.Fatalf("parse prefix %q: %v", s, err)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func mustIncrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
|
||||
t.Helper()
|
||||
ref, err := c.Increment(p, peer)
|
||||
if err != nil {
|
||||
t.Fatalf("Increment(%v, %s): %v", p, peer, err)
|
||||
}
|
||||
return ref
|
||||
}
|
||||
|
||||
func mustDecrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
|
||||
t.Helper()
|
||||
ref, err := c.Decrement(p, peer)
|
||||
if err != nil {
|
||||
t.Fatalf("Decrement(%v, %s): %v", p, peer, err)
|
||||
}
|
||||
return ref
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same
|
||||
// prefix routed by different peers. Removing the network whose peer is installed must hand the
|
||||
// prefix over to the surviving peer instead of leaving it on the removed one.
|
||||
func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
// First peer wins while both are present.
|
||||
if got := f.installed[p]; got != "peerA" {
|
||||
t.Fatalf("expected peerA installed, got %q", got)
|
||||
}
|
||||
|
||||
// Remove the active peer's network -> must swap to peerB.
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if got := f.installed[p]; got != "peerB" {
|
||||
t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got)
|
||||
}
|
||||
|
||||
// Remove the last one -> prefix gone.
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("expected prefix removed, still installed on %q", f.installed[p])
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard.
|
||||
func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
removesBefore := f.removes
|
||||
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
if f.installed[p] != "peerA" {
|
||||
t.Fatalf("active peer must stay peerA, got %q", f.installed[p])
|
||||
}
|
||||
if f.removes != removesBefore {
|
||||
t.Fatalf("removing a non-active peer must not call wg remove")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until
|
||||
// the last reference is released (the reason the per-peer count must be an int, not a set).
|
||||
func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
if f.adds != 1 {
|
||||
t.Fatalf("expected a single wg add for the same peer, got %d", f.adds)
|
||||
}
|
||||
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if f.installed[p] != "peerA" {
|
||||
t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p])
|
||||
}
|
||||
if f.removes != 0 {
|
||||
t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes)
|
||||
}
|
||||
|
||||
mustDecrement(t, c, p, "peerA")
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("prefix must be removed after last reference")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log).
|
||||
func TestAllowedIPs_RefCountAndActive(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
ref := mustIncrement(t, c, p, "peerA")
|
||||
if ref.Count != 1 || ref.Out != "peerA" {
|
||||
t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out)
|
||||
}
|
||||
ref = mustIncrement(t, c, p, "peerB")
|
||||
if ref.Count != 2 || ref.Out != "peerA" {
|
||||
t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_Flush removes everything installed and clears the counter.
|
||||
func TestAllowedIPs_Flush(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p1 := mustPrefix(t, "10.44.8.0/24")
|
||||
p2 := mustPrefix(t, "10.44.9.0/24")
|
||||
|
||||
mustIncrement(t, c, p1, "peerA")
|
||||
mustIncrement(t, c, p2, "peerB")
|
||||
|
||||
if err := c.Flush(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(f.installed) != 0 {
|
||||
t.Fatalf("expected all prefixes removed, got %v", f.installed)
|
||||
}
|
||||
// After flush, a fresh increment must add again.
|
||||
mustIncrement(t, c, p1, "peerC")
|
||||
if f.installed[p1] != "peerC" {
|
||||
t.Fatalf("counter not reset after flush")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SelfHealAfterSwapAddError ensures a failed add during a swap does not permanently
|
||||
// strand the prefix: the next Decrement (or Increment) must retry and install a surviving peer.
|
||||
func TestAllowedIPs_SelfHealAfterSwapAddError(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
mustIncrement(t, c, p, "peerC")
|
||||
|
||||
// Removing the active peerA triggers a swap to a survivor; make the add fail once.
|
||||
f.failAdd = true
|
||||
if _, err := c.Decrement(p, "peerA"); err == nil {
|
||||
t.Fatalf("expected error from failed swap add")
|
||||
}
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("nothing should be installed after a failed swap add, got %q", f.installed[p])
|
||||
}
|
||||
|
||||
// A later Decrement of a non-active survivor must retry the hand-off (self-heal), not stay stuck.
|
||||
ref := mustDecrement(t, c, p, "peerC")
|
||||
if got := f.installed[p]; got == "" {
|
||||
t.Fatalf("self-heal failed: prefix left unrouted after add recovered")
|
||||
}
|
||||
if ref.Out == "" {
|
||||
t.Fatalf("expected an active peer after self-heal, got empty")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllowedIPs_SelfHealAfterRemoveError ensures a failed remove during a swap is retried instead
|
||||
// of leaving e.active stuck on a peer that no longer holds references.
|
||||
func TestAllowedIPs_SelfHealAfterRemoveError(t *testing.T) {
|
||||
f := newFakeWG()
|
||||
c := f.counter()
|
||||
p := mustPrefix(t, "10.44.8.0/24")
|
||||
|
||||
mustIncrement(t, c, p, "peerA")
|
||||
mustIncrement(t, c, p, "peerB")
|
||||
|
||||
// Releasing active peerA must detach it (remove) then add peerB; fail the remove once.
|
||||
f.failRemove = true
|
||||
if _, err := c.Decrement(p, "peerA"); err == nil {
|
||||
t.Fatalf("expected error from failed remove")
|
||||
}
|
||||
|
||||
// Next Decrement of the non-active survivor retries: removes stale peerA, installs peerB.
|
||||
mustDecrement(t, c, p, "peerB")
|
||||
// peerB had only one ref, so after retry the prefix is fully released.
|
||||
if _, ok := f.installed[p]; ok {
|
||||
t.Fatalf("expected prefix released after self-heal, still on %q", f.installed[p])
|
||||
}
|
||||
}
|
||||
@@ -5,5 +5,7 @@ import "net/netip"
|
||||
// RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement
|
||||
type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}]
|
||||
|
||||
// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement
|
||||
type AllowedIPsRefCounter = Counter[netip.Prefix, string, string]
|
||||
// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware:
|
||||
// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer,
|
||||
// so the counter records the per-peer reference count and swaps the installed peer when the active one is released.
|
||||
// See allowedips.go.
|
||||
|
||||
138
client/internal/routemanager/selection.go
Normal file
138
client/internal/routemanager/selection.go
Normal file
@@ -0,0 +1,138 @@
|
||||
package routemanager
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
|
||||
"github.com/hashicorp/go-multierror"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/exp/maps"
|
||||
|
||||
nberrors "github.com/netbirdio/netbird/client/errors"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
// SelectRoutes selects the routes with the given network IDs and applies the
|
||||
// new selection. V4/v6 exit-node pairs are expanded automatically. Exit nodes
|
||||
// are mutually exclusive: if the selection activates an exit node, every other
|
||||
// available exit node is deselected so two can't be active at once. With
|
||||
// appendRoute=false the previous selection is replaced instead of extended.
|
||||
func (m *DefaultManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
|
||||
if err := m.selectRoutes(ids, appendRoute); err != nil {
|
||||
return err
|
||||
}
|
||||
m.TriggerSelection(m.GetClientRoutes())
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeselectRoutes removes the routes with the given network IDs from the
|
||||
// selection and applies the change. V4/v6 exit-node pairs are expanded
|
||||
// automatically.
|
||||
func (m *DefaultManager) DeselectRoutes(ids []route.NetID) error {
|
||||
if err := m.deselectRoutes(ids); err != nil {
|
||||
return err
|
||||
}
|
||||
m.TriggerSelection(m.GetClientRoutes())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *DefaultManager) deselectRoutes(ids []route.NetID) error {
|
||||
routesMap := m.GetClientRoutesWithNetID()
|
||||
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
|
||||
|
||||
log.Debugf("deselecting routes with ids: %v", routes)
|
||||
|
||||
if err := m.routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
|
||||
return fmt.Errorf("deselect routes: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SelectAllRoutes selects every available route and applies the selection.
|
||||
// Exit nodes stay mutually exclusive: at most one remains active.
|
||||
func (m *DefaultManager) SelectAllRoutes() {
|
||||
m.selectAllRoutes()
|
||||
m.TriggerSelection(m.GetClientRoutes())
|
||||
}
|
||||
|
||||
func (m *DefaultManager) selectAllRoutes() {
|
||||
m.routeSelector.SelectAllRoutes()
|
||||
|
||||
// Select-all wipes every explicit selection, so exit nodes fall back to
|
||||
// management's auto-apply flags — which may mark several at once.
|
||||
// Reconcile immediately so at most one exit node stays active instead of
|
||||
// waiting for the next network map to enforce it.
|
||||
m.mux.Lock()
|
||||
defer m.mux.Unlock()
|
||||
m.updateRouteSelectorFromManagement(m.clientRoutes)
|
||||
}
|
||||
|
||||
// DeselectAllRoutes deselects every route and applies the change.
|
||||
func (m *DefaultManager) DeselectAllRoutes() {
|
||||
m.routeSelector.DeselectAllRoutes()
|
||||
m.TriggerSelection(m.GetClientRoutes())
|
||||
}
|
||||
|
||||
func (m *DefaultManager) selectRoutes(ids []route.NetID, appendRoute bool) error {
|
||||
routesMap := m.GetClientRoutesWithNetID()
|
||||
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
|
||||
allIDs := maps.Keys(routesMap)
|
||||
|
||||
log.Debugf("selecting routes with ids: %v", routes)
|
||||
|
||||
// A partial failure (e.g. an unknown ID in the request) still selects the
|
||||
// valid routes, so exclusivity below must run regardless of the error.
|
||||
var merr *multierror.Error
|
||||
if err := m.routeSelector.SelectRoutes(routes, appendRoute, allIDs); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("select routes: %w", err))
|
||||
}
|
||||
|
||||
// Exit nodes are mutually exclusive: if this selection activates an
|
||||
// exit node, deselect every other available exit node so two can't be
|
||||
// selected at once. Non-exit route selections are left untouched.
|
||||
if requestActivatesExitNode(routes, routesMap) {
|
||||
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
|
||||
if err := m.routeSelector.DeselectRoutes(others, allIDs); err != nil {
|
||||
merr = multierror.Append(merr, fmt.Errorf("deselect sibling exit nodes: %w", err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nberrors.FormatErrorOrNil(merr)
|
||||
}
|
||||
|
||||
func isExitNodeRoutes(routes []*route.Route) bool {
|
||||
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
|
||||
}
|
||||
|
||||
// requestActivatesExitNode reports whether any requested NetID maps to an exit
|
||||
// node (default route) in the current route table.
|
||||
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
|
||||
for _, id := range requested {
|
||||
if isExitNodeRoutes(routesMap[id]) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// otherExitNodeIDs returns every available exit-node NetID that is not in the
|
||||
// requested set — the siblings to deselect so a single exit node stays active.
|
||||
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
|
||||
keep := make(map[route.NetID]struct{}, len(requested))
|
||||
for _, id := range requested {
|
||||
keep[id] = struct{}{}
|
||||
}
|
||||
var others []route.NetID
|
||||
for id, routes := range routesMap {
|
||||
if !isExitNodeRoutes(routes) {
|
||||
continue
|
||||
}
|
||||
if _, ok := keep[id]; ok {
|
||||
continue
|
||||
}
|
||||
others = append(others, id)
|
||||
}
|
||||
return others
|
||||
}
|
||||
129
client/internal/routemanager/selection_test.go
Normal file
129
client/internal/routemanager/selection_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package routemanager
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/routeselector"
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
func v6ExitRoute(netID, peer string) *route.Route {
|
||||
return &route.Route{
|
||||
NetID: route.NetID(netID),
|
||||
Network: netip.MustParsePrefix("::/0"),
|
||||
Peer: peer,
|
||||
}
|
||||
}
|
||||
|
||||
func newSelectionTestManager() *DefaultManager {
|
||||
return &DefaultManager{
|
||||
routeSelector: routeselector.NewRouteSelector(),
|
||||
clientRoutes: route.HAMap{
|
||||
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", true)},
|
||||
"exitA-v6|::/0": {v6ExitRoute("exitA-v6", "p1")},
|
||||
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", true)},
|
||||
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectRoutes_ExitNodeExclusivity(t *testing.T) {
|
||||
m := newSelectionTestManager()
|
||||
|
||||
// Selecting an exit node selects its v6 pair and deselects the sibling.
|
||||
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
|
||||
assert.True(t, m.routeSelector.IsSelected("exitA"), "exitA should be selected")
|
||||
assert.True(t, m.routeSelector.IsSelected("exitA-v6"), "the v6 pair follows its v4 base")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitB"), "the sibling exit node must be deselected")
|
||||
|
||||
// Switching to the sibling deselects the previous exit node and its v6 pair.
|
||||
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
|
||||
assert.True(t, m.routeSelector.IsSelected("exitB"), "exitB should now be selected")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitA"), "the previous exit node must be deselected")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "the previous exit node's v6 pair must be deselected")
|
||||
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
|
||||
|
||||
// Selecting a non-exit route leaves the active exit node alone.
|
||||
require.NoError(t, m.selectRoutes([]route.NetID{"lan"}, true))
|
||||
assert.True(t, m.routeSelector.IsSelected("exitB"), "selecting a non-exit route keeps the exit node")
|
||||
|
||||
// Deselecting the active exit node turns every exit node off.
|
||||
require.NoError(t, m.deselectRoutes([]route.NetID{"exitB"}))
|
||||
assert.False(t, m.routeSelector.IsSelected("exitB"), "exitB should be deselected")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitA"), "exitA stays deselected")
|
||||
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
|
||||
}
|
||||
|
||||
func TestSelectRoutes_PartialErrorStillEnforcesExclusivity(t *testing.T) {
|
||||
// The unknown ID must be reported, but the valid exit node in the same
|
||||
// request is still selected — so its sibling must still be deselected.
|
||||
// Both orderings are covered: processing must continue past the invalid
|
||||
// ID wherever it sits in the request.
|
||||
requests := map[string][]route.NetID{
|
||||
"invalid id first": {"missing", "exitB"},
|
||||
"invalid id last": {"exitB", "missing"},
|
||||
}
|
||||
|
||||
for name, ids := range requests {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
m := newSelectionTestManager()
|
||||
|
||||
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
|
||||
|
||||
err := m.selectRoutes(ids, true)
|
||||
assert.Error(t, err, "unknown id must be reported")
|
||||
assert.True(t, m.routeSelector.IsSelected("exitB"), "valid exit node from the request is selected")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitA"), "sibling exit node must be deselected despite the error")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "sibling's v6 pair must be deselected too")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectAllRoutes_KeepsSingleExitNode(t *testing.T) {
|
||||
// Both exit nodes are marked for auto-apply by management
|
||||
// (SkipAutoApply=false), the state where select-all could turn on two at
|
||||
// once without the immediate reconciliation.
|
||||
m := &DefaultManager{
|
||||
routeSelector: routeselector.NewRouteSelector(),
|
||||
clientRoutes: route.HAMap{
|
||||
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", false)},
|
||||
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", false)},
|
||||
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
|
||||
},
|
||||
}
|
||||
|
||||
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
|
||||
|
||||
m.selectAllRoutes()
|
||||
|
||||
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit routes are all selected")
|
||||
assert.True(t, m.routeSelector.IsSelected("exitA"), "the deterministic management pick stays active")
|
||||
assert.False(t, m.routeSelector.IsSelected("exitB"), "select-all must not leave a second exit node active")
|
||||
}
|
||||
|
||||
func TestSelectRoutes_UnknownRoute(t *testing.T) {
|
||||
m := newSelectionTestManager()
|
||||
|
||||
assert.Error(t, m.selectRoutes([]route.NetID{"missing"}, true), "selecting an unavailable route must fail")
|
||||
assert.Error(t, m.deselectRoutes([]route.NetID{"missing"}), "deselecting an unavailable route must fail")
|
||||
}
|
||||
|
||||
func TestExitNodeSelectionHelpers(t *testing.T) {
|
||||
routesMap := map[route.NetID][]*route.Route{
|
||||
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
|
||||
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
|
||||
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
|
||||
}
|
||||
|
||||
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
|
||||
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
|
||||
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
|
||||
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
|
||||
|
||||
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
|
||||
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
|
||||
}
|
||||
@@ -15,6 +15,11 @@ type Route struct {
|
||||
route *route.Route
|
||||
routeRefCounter *refcounter.RouteRefCounter
|
||||
allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
|
||||
// currentPeerKey is the routing peer this watcher currently has the prefix installed on
|
||||
// (the HA winner elected by the watcher). It can differ from route.Peer and change on
|
||||
// failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement
|
||||
// the exact peer that was incremented.
|
||||
currentPeerKey string
|
||||
}
|
||||
|
||||
func NewRoute(params common.HandlerParams) *Route {
|
||||
@@ -52,12 +57,15 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
|
||||
ref.Out,
|
||||
)
|
||||
}
|
||||
r.currentPeerKey = peerKey
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Route) RemoveAllowedIPs() error {
|
||||
if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
|
||||
return err
|
||||
var err error
|
||||
if _, decErr := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); decErr != nil {
|
||||
err = fmt.Errorf("remove allowed IP %s: %w", r.route.Network, decErr)
|
||||
}
|
||||
return nil
|
||||
r.currentPeerKey = ""
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -20,6 +20,8 @@ const (
|
||||
rpFilterPath = "net.ipv4.conf.all.rp_filter"
|
||||
rpFilterInterfacePath = "net.ipv4.conf.%s.rp_filter"
|
||||
srcValidMarkPath = "net.ipv4.conf.all.src_valid_mark"
|
||||
percentEscape = "%25"
|
||||
dotEscape = "%2E"
|
||||
)
|
||||
|
||||
type iface interface {
|
||||
@@ -56,7 +58,11 @@ func Setup(wgIface iface) (map[string]int, error) {
|
||||
continue
|
||||
}
|
||||
|
||||
i := fmt.Sprintf(rpFilterInterfacePath, intf.Name)
|
||||
// Escape '%' and '.' so they survive the dot-to-slash conversion in Set()
|
||||
safeName := strings.ReplaceAll(intf.Name, "%", percentEscape)
|
||||
safeName = strings.ReplaceAll(safeName, ".", dotEscape)
|
||||
|
||||
i := fmt.Sprintf(rpFilterInterfacePath, safeName)
|
||||
oldVal, err := Set(i, 2, true)
|
||||
if err != nil {
|
||||
result = multierror.Append(result, err)
|
||||
@@ -70,7 +76,11 @@ func Setup(wgIface iface) (map[string]int, error) {
|
||||
|
||||
// Set sets a sysctl configuration, if onlyIfOne is true it will only set the new value if it's set to 1
|
||||
func Set(key string, desiredValue int, onlyIfOne bool) (int, error) {
|
||||
path := fmt.Sprintf("/proc/sys/%s", strings.ReplaceAll(key, ".", "/"))
|
||||
path := strings.ReplaceAll(key, ".", "/")
|
||||
// Unescape interface dots and percent signs
|
||||
path = strings.ReplaceAll(path, dotEscape, ".")
|
||||
path = strings.ReplaceAll(path, percentEscape, "%")
|
||||
path = fmt.Sprintf("/proc/sys/%s", path)
|
||||
currentValue, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return -1, fmt.Errorf("read sysctl %s: %w", key, err)
|
||||
|
||||
124
client/internal/tunnelnotifier/notifier.go
Normal file
124
client/internal/tunnelnotifier/notifier.go
Normal file
@@ -0,0 +1,124 @@
|
||||
package tunnelnotifier
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"sync"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal/dns"
|
||||
"github.com/netbirdio/netbird/client/internal/listener"
|
||||
)
|
||||
|
||||
type eventKind int
|
||||
|
||||
const (
|
||||
eventRoutes eventKind = iota
|
||||
eventIfaceIP
|
||||
eventIfaceIPv6
|
||||
eventDNS
|
||||
)
|
||||
|
||||
var (
|
||||
_ listener.NetworkChangeListener = (*Notifier)(nil)
|
||||
_ dns.IosDnsManager = (*Notifier)(nil)
|
||||
)
|
||||
|
||||
type event struct {
|
||||
kind eventKind
|
||||
payload string
|
||||
}
|
||||
|
||||
type Notifier struct {
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
queue *list.List
|
||||
closed bool
|
||||
done chan struct{}
|
||||
|
||||
listener listener.NetworkChangeListener
|
||||
dnsManager dns.IosDnsManager
|
||||
}
|
||||
|
||||
func New(l listener.NetworkChangeListener, dm dns.IosDnsManager) *Notifier {
|
||||
n := &Notifier{
|
||||
queue: list.New(),
|
||||
done: make(chan struct{}),
|
||||
listener: l,
|
||||
dnsManager: dm,
|
||||
}
|
||||
n.cond = sync.NewCond(&n.mu)
|
||||
go n.deliverLoop()
|
||||
return n
|
||||
}
|
||||
|
||||
func (n *Notifier) OnNetworkChanged(routes string) {
|
||||
n.enqueue(event{kind: eventRoutes, payload: routes})
|
||||
}
|
||||
|
||||
func (n *Notifier) SetInterfaceIP(ip string) {
|
||||
n.enqueue(event{kind: eventIfaceIP, payload: ip})
|
||||
}
|
||||
|
||||
func (n *Notifier) SetInterfaceIPv6(ip string) {
|
||||
n.enqueue(event{kind: eventIfaceIPv6, payload: ip})
|
||||
}
|
||||
|
||||
func (n *Notifier) ApplyDns(config string) {
|
||||
n.enqueue(event{kind: eventDNS, payload: config})
|
||||
}
|
||||
|
||||
// Close stops accepting new events and blocks until the delivery loop has
|
||||
// drained all queued events and exited.
|
||||
func (n *Notifier) Close() {
|
||||
n.mu.Lock()
|
||||
n.closed = true
|
||||
n.cond.Signal()
|
||||
n.mu.Unlock()
|
||||
<-n.done
|
||||
}
|
||||
|
||||
func (n *Notifier) enqueue(ev event) {
|
||||
n.mu.Lock()
|
||||
defer n.mu.Unlock()
|
||||
if n.closed {
|
||||
return
|
||||
}
|
||||
n.queue.PushBack(ev)
|
||||
n.cond.Signal()
|
||||
}
|
||||
|
||||
func (n *Notifier) deliverLoop() {
|
||||
defer close(n.done)
|
||||
for {
|
||||
n.mu.Lock()
|
||||
for n.queue.Len() == 0 && !n.closed {
|
||||
n.cond.Wait()
|
||||
}
|
||||
if n.closed && n.queue.Len() == 0 {
|
||||
n.mu.Unlock()
|
||||
return
|
||||
}
|
||||
ev := n.queue.Remove(n.queue.Front()).(event)
|
||||
l := n.listener
|
||||
dm := n.dnsManager
|
||||
n.mu.Unlock()
|
||||
|
||||
switch ev.kind {
|
||||
case eventRoutes:
|
||||
if l != nil {
|
||||
l.OnNetworkChanged(ev.payload)
|
||||
}
|
||||
case eventIfaceIP:
|
||||
if l != nil {
|
||||
l.SetInterfaceIP(ev.payload)
|
||||
}
|
||||
case eventIfaceIPv6:
|
||||
if l != nil {
|
||||
l.SetInterfaceIPv6(ev.payload)
|
||||
}
|
||||
case eventDNS:
|
||||
if dm != nil {
|
||||
dm.ApplyDns(ev.payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
192
client/internal/tunnelnotifier/notifier_test.go
Normal file
192
client/internal/tunnelnotifier/notifier_test.go
Normal file
@@ -0,0 +1,192 @@
|
||||
package tunnelnotifier
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type call struct {
|
||||
kind string
|
||||
payload string
|
||||
}
|
||||
|
||||
type recorder struct {
|
||||
mu sync.Mutex
|
||||
calls []call
|
||||
inFlight atomic.Int32
|
||||
overlap atomic.Bool
|
||||
delay time.Duration
|
||||
}
|
||||
|
||||
func (r *recorder) record(kind, payload string) {
|
||||
if r.inFlight.Add(1) != 1 {
|
||||
r.overlap.Store(true)
|
||||
}
|
||||
if r.delay > 0 {
|
||||
time.Sleep(r.delay)
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.calls = append(r.calls, call{kind: kind, payload: payload})
|
||||
r.mu.Unlock()
|
||||
r.inFlight.Add(-1)
|
||||
}
|
||||
|
||||
func (r *recorder) count() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return len(r.calls)
|
||||
}
|
||||
|
||||
func (r *recorder) snapshot() []call {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]call, len(r.calls))
|
||||
copy(out, r.calls)
|
||||
return out
|
||||
}
|
||||
|
||||
type fakeListener struct {
|
||||
rec *recorder
|
||||
}
|
||||
|
||||
func (f *fakeListener) OnNetworkChanged(routes string) {
|
||||
f.rec.record("routes", routes)
|
||||
}
|
||||
|
||||
func (f *fakeListener) SetInterfaceIP(ip string) {
|
||||
f.rec.record("ip", ip)
|
||||
}
|
||||
|
||||
func (f *fakeListener) SetInterfaceIPv6(ip string) {
|
||||
f.rec.record("ipv6", ip)
|
||||
}
|
||||
|
||||
type fakeDNSManager struct {
|
||||
rec *recorder
|
||||
}
|
||||
|
||||
func (f *fakeDNSManager) ApplyDns(config string) {
|
||||
f.rec.record("dns", config)
|
||||
}
|
||||
|
||||
func TestFIFOOrder(t *testing.T) {
|
||||
rec := &recorder{}
|
||||
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
|
||||
defer n.Close()
|
||||
|
||||
n.SetInterfaceIP("10.0.0.1")
|
||||
n.SetInterfaceIPv6("fd00::1")
|
||||
n.ApplyDns(`{"domains":[]}`)
|
||||
n.OnNetworkChanged("10.0.0.0/8,192.168.0.0/16")
|
||||
n.ApplyDns(`{"domains":["example.com"]}`)
|
||||
|
||||
require.Eventually(t, func() bool { return rec.count() == 5 }, time.Second, time.Millisecond)
|
||||
|
||||
expected := []call{
|
||||
{kind: "ip", payload: "10.0.0.1"},
|
||||
{kind: "ipv6", payload: "fd00::1"},
|
||||
{kind: "dns", payload: `{"domains":[]}`},
|
||||
{kind: "routes", payload: "10.0.0.0/8,192.168.0.0/16"},
|
||||
{kind: "dns", payload: `{"domains":["example.com"]}`},
|
||||
}
|
||||
assert.Equal(t, expected, rec.snapshot())
|
||||
}
|
||||
|
||||
func TestNoOverlappingCalls(t *testing.T) {
|
||||
rec := &recorder{delay: 100 * time.Microsecond}
|
||||
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
|
||||
defer n.Close()
|
||||
|
||||
const producers = 8
|
||||
const perProducer = 25
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < producers; i++ {
|
||||
wg.Add(1)
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < perProducer; j++ {
|
||||
payload := fmt.Sprintf("%d-%d", id, j)
|
||||
switch j % 4 {
|
||||
case 0:
|
||||
n.OnNetworkChanged(payload)
|
||||
case 1:
|
||||
n.SetInterfaceIP(payload)
|
||||
case 2:
|
||||
n.SetInterfaceIPv6(payload)
|
||||
case 3:
|
||||
n.ApplyDns(payload)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
require.Eventually(t, func() bool { return rec.count() == producers*perProducer }, 5*time.Second, time.Millisecond)
|
||||
assert.False(t, rec.overlap.Load())
|
||||
}
|
||||
|
||||
func TestDNSAndRoutesInterleaved(t *testing.T) {
|
||||
rec := &recorder{delay: 100 * time.Microsecond}
|
||||
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
|
||||
defer n.Close()
|
||||
|
||||
const events = 50
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < events; i++ {
|
||||
n.ApplyDns(fmt.Sprintf("dns-%d", i))
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < events; i++ {
|
||||
n.OnNetworkChanged(fmt.Sprintf("routes-%d", i))
|
||||
}
|
||||
}()
|
||||
wg.Wait()
|
||||
|
||||
require.Eventually(t, func() bool { return rec.count() == 2*events }, 5*time.Second, time.Millisecond)
|
||||
assert.False(t, rec.overlap.Load())
|
||||
|
||||
var dnsSeen, routesSeen int
|
||||
for _, c := range rec.snapshot() {
|
||||
switch c.kind {
|
||||
case "dns":
|
||||
assert.Equal(t, fmt.Sprintf("dns-%d", dnsSeen), c.payload)
|
||||
dnsSeen++
|
||||
case "routes":
|
||||
assert.Equal(t, fmt.Sprintf("routes-%d", routesSeen), c.payload)
|
||||
routesSeen++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, events, dnsSeen)
|
||||
assert.Equal(t, events, routesSeen)
|
||||
}
|
||||
|
||||
func TestCloseDrainsQueue(t *testing.T) {
|
||||
rec := &recorder{delay: time.Millisecond}
|
||||
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
|
||||
|
||||
const events = 20
|
||||
for i := 0; i < events; i++ {
|
||||
n.OnNetworkChanged(fmt.Sprintf("routes-%d", i))
|
||||
}
|
||||
n.Close()
|
||||
|
||||
require.Equal(t, events, rec.count(), "Close must not return before all queued events are delivered")
|
||||
|
||||
n.OnNetworkChanged("after-close")
|
||||
n.ApplyDns("after-close")
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
assert.Equal(t, events, rec.count())
|
||||
}
|
||||
@@ -13,7 +13,6 @@ import (
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"golang.org/x/exp/maps"
|
||||
|
||||
"github.com/netbirdio/netbird/client/internal"
|
||||
"github.com/netbirdio/netbird/client/internal/auth"
|
||||
@@ -637,23 +636,18 @@ func (c *Client) SelectRoute(id string) error {
|
||||
}
|
||||
|
||||
routeManager := engine.GetRouteManager()
|
||||
routeSelector := routeManager.GetRouteSelector()
|
||||
if id == "All" {
|
||||
log.Debugf("select all routes")
|
||||
routeSelector.SelectAllRoutes()
|
||||
} else {
|
||||
log.Debugf("select route with id: %s", id)
|
||||
routes := toNetIDs([]string{id})
|
||||
routesMap := routeManager.GetClientRoutesWithNetID()
|
||||
routes = route.ExpandV6ExitPairs(routes, routesMap)
|
||||
if err := routeSelector.SelectRoutes(routes, true, maps.Keys(routesMap)); err != nil {
|
||||
log.Debugf("error when selecting routes: %s", err)
|
||||
return fmt.Errorf("select routes: %w", err)
|
||||
}
|
||||
routeManager.SelectAllRoutes()
|
||||
return nil
|
||||
}
|
||||
routeManager.TriggerSelection(routeManager.GetClientRoutes())
|
||||
return nil
|
||||
|
||||
log.Debugf("select route with id: %s", id)
|
||||
if err := routeManager.SelectRoutes(toNetIDs([]string{id}), true); err != nil {
|
||||
log.Debugf("error when selecting routes: %s", err)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) DeselectRoute(id string) error {
|
||||
@@ -667,21 +661,17 @@ func (c *Client) DeselectRoute(id string) error {
|
||||
}
|
||||
|
||||
routeManager := engine.GetRouteManager()
|
||||
routeSelector := routeManager.GetRouteSelector()
|
||||
if id == "All" {
|
||||
log.Debugf("deselect all routes")
|
||||
routeSelector.DeselectAllRoutes()
|
||||
} else {
|
||||
log.Debugf("deselect route with id: %s", id)
|
||||
routes := toNetIDs([]string{id})
|
||||
routesMap := routeManager.GetClientRoutesWithNetID()
|
||||
routes = route.ExpandV6ExitPairs(routes, routesMap)
|
||||
if err := routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
|
||||
log.Debugf("error when deselecting routes: %s", err)
|
||||
return fmt.Errorf("deselect routes: %w", err)
|
||||
}
|
||||
routeManager.DeselectAllRoutes()
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Debugf("deselect route with id: %s", id)
|
||||
if err := routeManager.DeselectRoutes(toNetIDs([]string{id})); err != nil {
|
||||
log.Debugf("error when deselecting routes: %s", err)
|
||||
return err
|
||||
}
|
||||
routeManager.TriggerSelection(routeManager.GetClientRoutes())
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
12
client/ios/NetBirdSDK/version.go
Normal file
12
client/ios/NetBirdSDK/version.go
Normal file
@@ -0,0 +1,12 @@
|
||||
//go:build ios
|
||||
|
||||
package NetBirdSDK
|
||||
|
||||
import "github.com/netbirdio/netbird/version"
|
||||
|
||||
// GoClientVersion returns the NetBird Go client version that was baked into
|
||||
// the framework at compile time via
|
||||
// -ldflags "-X github.com/netbirdio/netbird/version.version=<version>".
|
||||
func GoClientVersion() string {
|
||||
return version.NetbirdVersion()
|
||||
}
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/exp/maps"
|
||||
"google.golang.org/grpc/codes"
|
||||
gstatus "google.golang.org/grpc/status"
|
||||
|
||||
@@ -161,30 +160,11 @@ func (s *Server) SelectNetworks(_ context.Context, req *proto.SelectNetworksRequ
|
||||
return nil, fmt.Errorf("no route manager")
|
||||
}
|
||||
|
||||
routeSelector := routeManager.GetRouteSelector()
|
||||
if req.GetAll() {
|
||||
routeSelector.SelectAllRoutes()
|
||||
} else {
|
||||
routes := toNetIDs(req.GetNetworkIDs())
|
||||
routesMap := routeManager.GetClientRoutesWithNetID()
|
||||
routes = route.ExpandV6ExitPairs(routes, routesMap)
|
||||
netIdRoutes := maps.Keys(routesMap)
|
||||
if err := routeSelector.SelectRoutes(routes, req.GetAppend(), netIdRoutes); err != nil {
|
||||
return nil, fmt.Errorf("select routes: %w", err)
|
||||
}
|
||||
|
||||
// Exit nodes are mutually exclusive: if this selection activates an
|
||||
// exit node, deselect every other available exit node so two can't be
|
||||
// selected at once. Non-exit route selections are left untouched.
|
||||
if requestActivatesExitNode(routes, routesMap) {
|
||||
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
|
||||
if err := routeSelector.DeselectRoutes(others, netIdRoutes); err != nil {
|
||||
return nil, fmt.Errorf("deselect sibling exit nodes: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
routeManager.SelectAllRoutes()
|
||||
} else if err := routeManager.SelectRoutes(toNetIDs(req.GetNetworkIDs()), req.GetAppend()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
routeManager.TriggerSelection(routeManager.GetClientRoutes())
|
||||
|
||||
s.statusRecorder.PublishEvent(
|
||||
proto.SystemEvent_INFO,
|
||||
@@ -224,19 +204,11 @@ func (s *Server) DeselectNetworks(_ context.Context, req *proto.SelectNetworksRe
|
||||
return nil, fmt.Errorf("no route manager")
|
||||
}
|
||||
|
||||
routeSelector := routeManager.GetRouteSelector()
|
||||
if req.GetAll() {
|
||||
routeSelector.DeselectAllRoutes()
|
||||
} else {
|
||||
routes := toNetIDs(req.GetNetworkIDs())
|
||||
routesMap := routeManager.GetClientRoutesWithNetID()
|
||||
routes = route.ExpandV6ExitPairs(routes, routesMap)
|
||||
netIdRoutes := maps.Keys(routesMap)
|
||||
if err := routeSelector.DeselectRoutes(routes, netIdRoutes); err != nil {
|
||||
return nil, fmt.Errorf("deselect routes: %w", err)
|
||||
}
|
||||
routeManager.DeselectAllRoutes()
|
||||
} else if err := routeManager.DeselectRoutes(toNetIDs(req.GetNetworkIDs())); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
routeManager.TriggerSelection(routeManager.GetClientRoutes())
|
||||
|
||||
s.statusRecorder.PublishEvent(
|
||||
proto.SystemEvent_INFO,
|
||||
@@ -261,37 +233,3 @@ func toNetIDs(routes []string) []route.NetID {
|
||||
return netIDs
|
||||
}
|
||||
|
||||
func isExitNodeRoutes(routes []*route.Route) bool {
|
||||
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
|
||||
}
|
||||
|
||||
// requestActivatesExitNode reports whether any requested NetID maps to an exit
|
||||
// node (default route) in the current route table.
|
||||
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
|
||||
for _, id := range requested {
|
||||
if isExitNodeRoutes(routesMap[id]) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// otherExitNodeIDs returns every available exit-node NetID that is not in the
|
||||
// requested set — the siblings to deselect so a single exit node stays active.
|
||||
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
|
||||
keep := make(map[route.NetID]struct{}, len(requested))
|
||||
for _, id := range requested {
|
||||
keep[id] = struct{}{}
|
||||
}
|
||||
var others []route.NetID
|
||||
for id, routes := range routesMap {
|
||||
if !isExitNodeRoutes(routes) {
|
||||
continue
|
||||
}
|
||||
if _, ok := keep[id]; ok {
|
||||
continue
|
||||
}
|
||||
others = append(others, id)
|
||||
}
|
||||
return others
|
||||
}
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/netip"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/route"
|
||||
)
|
||||
|
||||
func TestExitNodeSelectionHelpers(t *testing.T) {
|
||||
routesMap := map[route.NetID][]*route.Route{
|
||||
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
|
||||
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
|
||||
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
|
||||
}
|
||||
|
||||
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
|
||||
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
|
||||
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
|
||||
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
|
||||
|
||||
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
|
||||
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
|
||||
}
|
||||
@@ -1,20 +1,27 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/proto"
|
||||
)
|
||||
|
||||
const statusCoalesceWindow = 200 * time.Millisecond
|
||||
|
||||
// SubscribeStatus pushes a fresh StatusResponse on every connection state
|
||||
// change. The first message is the current snapshot, so a re-subscribing
|
||||
// client doesn't need to also call Status. Subsequent messages fire when
|
||||
// the peer recorder reports any of: connected/disconnected/connecting,
|
||||
// management or signal flip, address change, or peers list change.
|
||||
//
|
||||
// The change channel coalesces bursts to a single tick. If the consumer
|
||||
// is slow the daemon drops extras (not blocks), and the next snapshot
|
||||
// the consumer pulls already reflects everything.
|
||||
// Bursts are coalesced deterministically: the first tick after a quiet
|
||||
// period is sent immediately, then a short window swallows the rest of the
|
||||
// burst and a single trailing snapshot covers whatever arrived meanwhile.
|
||||
// Every send is a full snapshot of the recorder's current state, so
|
||||
// swallowed ticks lose no information.
|
||||
func (s *Server) SubscribeStatus(req *proto.StatusRequest, stream proto.DaemonService_SubscribeStatusServer) error {
|
||||
subID, ch := s.statusRecorder.SubscribeToStateChanges()
|
||||
defer func() {
|
||||
@@ -37,12 +44,50 @@ func (s *Server) SubscribeStatus(req *proto.StatusRequest, stream proto.DaemonSe
|
||||
if err := s.sendStatusSnapshot(req, stream); err != nil {
|
||||
return err
|
||||
}
|
||||
pending, open := collectStatusBurst(stream.Context(), ch)
|
||||
if pending {
|
||||
if err := s.sendStatusSnapshot(req, stream); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if !open {
|
||||
return nil
|
||||
}
|
||||
case <-stream.Context().Done():
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// collectStatusBurst waits out the coalesce window, absorbing further ticks.
|
||||
// pending reports whether any tick arrived; open is false when the channel
|
||||
// closed or the stream context ended.
|
||||
func collectStatusBurst(ctx context.Context, ch <-chan struct{}) (pending, open bool) {
|
||||
timer := time.NewTimer(statusCoalesceWindow)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case _, ok := <-ch:
|
||||
if !ok {
|
||||
return pending, false
|
||||
}
|
||||
pending = true
|
||||
case <-timer.C:
|
||||
select {
|
||||
case _, ok := <-ch:
|
||||
if !ok {
|
||||
return pending, false
|
||||
}
|
||||
pending = true
|
||||
default:
|
||||
}
|
||||
return pending, true
|
||||
case <-ctx.Done():
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) sendStatusSnapshot(req *proto.StatusRequest, stream proto.DaemonService_SubscribeStatusServer) error {
|
||||
resp, err := s.buildStatusResponse(stream.Context(), req)
|
||||
if err != nil {
|
||||
|
||||
@@ -2,6 +2,7 @@ import { type ReactNode } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { Browser } from "@wailsio/runtime";
|
||||
import { DownloadIcon, NotepadText } from "lucide-react";
|
||||
import { Update as UpdateSvc } from "@bindings/services";
|
||||
import { Button } from "@/components/buttons/Button";
|
||||
import { useClientVersion } from "@/contexts/ClientVersionContext";
|
||||
import { cn } from "@/lib/cn";
|
||||
@@ -14,6 +15,12 @@ function openUrl(url: string) {
|
||||
});
|
||||
}
|
||||
|
||||
function openInstallerDownload() {
|
||||
UpdateSvc.DownloadURL()
|
||||
.then(openUrl)
|
||||
.catch(() => openUrl(GITHUB_RELEASES));
|
||||
}
|
||||
|
||||
export function UpdateVersionCard() {
|
||||
const { t } = useTranslation();
|
||||
const { updateVersion, enforced, triggerUpdate } = useClientVersion();
|
||||
@@ -37,11 +44,7 @@ export function UpdateVersionCard() {
|
||||
{t("update.card.installNow")}
|
||||
</Button>
|
||||
) : (
|
||||
<Button
|
||||
variant={"primary"}
|
||||
size={"xs"}
|
||||
onClick={() => openUrl(GITHUB_RELEASES)}
|
||||
>
|
||||
<Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
|
||||
<DownloadIcon size={14} />
|
||||
{t("update.card.getInstaller")}
|
||||
</Button>
|
||||
|
||||
@@ -103,6 +103,18 @@ func main() {
|
||||
updaterHolder := updater.NewHolder(app.Event)
|
||||
update := services.NewUpdate(conn, updaterHolder)
|
||||
daemonFeed := services.NewDaemonFeed(conn, app.Event, updaterHolder, debugLog)
|
||||
// Status snapshots go only to visible windows — a snapshot is full state,
|
||||
// so a hidden webview loses nothing by being skipped; it gets the cached
|
||||
// one on show (SetShowReplay below). Every other event stays on the bus.
|
||||
daemonFeed.SetWindowDispatcher(func(st services.Status) {
|
||||
ev := &application.CustomEvent{Name: services.EventStatusSnapshot, Data: st}
|
||||
for _, w := range app.Window.GetAll() {
|
||||
if w == nil || !w.IsVisible() {
|
||||
continue
|
||||
}
|
||||
w.DispatchWailsEvent(ev)
|
||||
}
|
||||
})
|
||||
notifier := notifications.New()
|
||||
compat := services.NewCompat(conn)
|
||||
// macOS shows no toast until permission is requested. Run it after
|
||||
@@ -152,8 +164,31 @@ func main() {
|
||||
// re-centering on that environment; nil leaves placement to the WM on full
|
||||
// desktops, macOS, and Windows.
|
||||
windowManager.SetRecenterOnShow(recenterOnShowPredicate())
|
||||
// Replay the latest snapshot into a window on (re)show, so a webview that
|
||||
// was hidden while pushes flowed never paints stale state. ReplayLast keeps
|
||||
// the feed's status lock across the dispatch, so a racing live push can't
|
||||
// slip in between and then be overwritten by this older cached snapshot.
|
||||
windowManager.SetShowReplay(func(w application.Window) {
|
||||
daemonFeed.ReplayLast(func(st services.Status) {
|
||||
w.DispatchWailsEvent(&application.CustomEvent{Name: services.EventStatusSnapshot, Data: st})
|
||||
})
|
||||
})
|
||||
app.RegisterService(application.NewService(windowManager))
|
||||
|
||||
// On macOS, Wails' default applicationShouldHandleReopen handler Show()s
|
||||
// every hidden window on dock-icon click, resurrecting hide-on-close
|
||||
// surfaces like Settings. Cancel it in a hook (hooks run before listeners)
|
||||
// and show only the main window. No-op elsewhere — the event never fires.
|
||||
if runtime.GOOS == "darwin" {
|
||||
app.Event.RegisterApplicationEventHook(events.Mac.ApplicationShouldHandleReopen, func(e *application.ApplicationEvent) {
|
||||
e.Cancel()
|
||||
if e.Context().HasVisibleWindows() {
|
||||
return
|
||||
}
|
||||
windowManager.ShowMain()
|
||||
})
|
||||
}
|
||||
|
||||
// Welcome window, first launch only — Continue flips OnboardingCompleted
|
||||
// so later launches skip it. ApplicationStarted hook so the Wails window
|
||||
// machinery is fully up before the window is created.
|
||||
@@ -281,6 +316,9 @@ func newApplication(onSecondInstance func()) *application.App {
|
||||
Linux: application.LinuxOptions{
|
||||
ProgramName: "netbird",
|
||||
},
|
||||
Windows: application.WindowsOptions{
|
||||
WndProcInterceptor: endSessionInterceptor(),
|
||||
},
|
||||
SingleInstance: &application.SingleInstanceOptions{
|
||||
UniqueID: "io.netbird.ui",
|
||||
OnSecondInstanceLaunch: func(_ application.SecondInstanceData) {
|
||||
@@ -367,24 +405,12 @@ func newMainWindow(app *application.App, prefStore *preferences.Store) *applicat
|
||||
|
||||
// Hide instead of quit on close; "really quit" is reached via tray -> Quit.
|
||||
window.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
|
||||
if services.ShuttingDown() {
|
||||
return
|
||||
}
|
||||
e.Cancel()
|
||||
window.Hide()
|
||||
})
|
||||
|
||||
// On macOS, Wails' default applicationShouldHandleReopen handler Show()s
|
||||
// every hidden window on dock-icon click, resurrecting hide-on-close
|
||||
// surfaces like Settings. Cancel it in a hook (hooks run before listeners)
|
||||
// and show only the main window. No-op elsewhere — the event never fires.
|
||||
if runtime.GOOS == "darwin" {
|
||||
app.Event.RegisterApplicationEventHook(events.Mac.ApplicationShouldHandleReopen, func(e *application.ApplicationEvent) {
|
||||
e.Cancel()
|
||||
if e.Context().HasVisibleWindows() {
|
||||
return
|
||||
}
|
||||
window.Show()
|
||||
window.Focus()
|
||||
})
|
||||
}
|
||||
|
||||
return window
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ package services
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -167,6 +168,14 @@ type DaemonFeed struct {
|
||||
cancel context.CancelFunc
|
||||
streamWg sync.WaitGroup
|
||||
|
||||
// statusSubs are Go-side snapshot consumers (the tray), fed directly so
|
||||
// they don't ride the window event bus. Callbacks run synchronously on
|
||||
// the stream goroutine, so pushes arrive in order.
|
||||
statusSubsMu sync.Mutex
|
||||
statusSubs []func(Status)
|
||||
lastStatus *Status
|
||||
windowDispatcher func(Status)
|
||||
|
||||
switchMu sync.Mutex
|
||||
switchInProgress bool
|
||||
switchInProgressUntil time.Time
|
||||
@@ -188,6 +197,55 @@ func NewDaemonFeed(conn DaemonConn, emitter Emitter, updaterHolder *updater.Hold
|
||||
return &DaemonFeed{conn: conn, emitter: emitter, updater: updaterHolder, logCtl: logCtl}
|
||||
}
|
||||
|
||||
// OnStatus registers a Go-side status subscriber. Not for the frontend —
|
||||
// React consumers subscribe to EventStatusSnapshot on the event bus.
|
||||
func (s *DaemonFeed) OnStatus(cb func(Status)) {
|
||||
s.statusSubsMu.Lock()
|
||||
s.statusSubs = append(s.statusSubs, cb)
|
||||
s.statusSubsMu.Unlock()
|
||||
}
|
||||
|
||||
// SetWindowDispatcher installs the frontend push path: it receives every
|
||||
// snapshot and decides which webview windows get it (visible ones). While
|
||||
// unset, pushStatus falls back to the event-bus broadcast.
|
||||
func (s *DaemonFeed) SetWindowDispatcher(fn func(Status)) {
|
||||
s.statusSubsMu.Lock()
|
||||
s.windowDispatcher = fn
|
||||
s.statusSubsMu.Unlock()
|
||||
}
|
||||
|
||||
// pushStatus delivers a snapshot to the Go-side subscribers and the frontend,
|
||||
// and caches it for ReplayLast. Hidden windows are skipped by the
|
||||
// window dispatcher; they catch up via the show replay (WindowManager).
|
||||
func (s *DaemonFeed) pushStatus(st Status) {
|
||||
s.statusSubsMu.Lock()
|
||||
s.lastStatus = &st
|
||||
subs := slices.Clone(s.statusSubs)
|
||||
dispatch := s.windowDispatcher
|
||||
s.statusSubsMu.Unlock()
|
||||
for _, cb := range subs {
|
||||
cb(st)
|
||||
}
|
||||
if dispatch != nil {
|
||||
dispatch(st)
|
||||
return
|
||||
}
|
||||
s.emitter.Emit(EventStatusSnapshot, st)
|
||||
}
|
||||
|
||||
// ReplayLast feeds the most recently pushed snapshot to dispatch; no-op before
|
||||
// the first push. The status lock is held across the dispatch so a concurrent
|
||||
// pushStatus cannot deliver a newer snapshot in between the read and the
|
||||
// dispatch — the replayed value is never older than anything already delivered.
|
||||
func (s *DaemonFeed) ReplayLast(dispatch func(Status)) {
|
||||
s.statusSubsMu.Lock()
|
||||
defer s.statusSubsMu.Unlock()
|
||||
if s.lastStatus == nil {
|
||||
return
|
||||
}
|
||||
dispatch(*s.lastStatus)
|
||||
}
|
||||
|
||||
// BeginProfileSwitch arms suppression for a switch from Connected/Connecting,
|
||||
// where the daemon emits stale Connected updates during Down's teardown then an
|
||||
// Idle before the new Up; statusStreamLoop drops those, and a synthetic
|
||||
@@ -201,7 +259,7 @@ func (s *DaemonFeed) BeginProfileSwitch() {
|
||||
s.switchLoginWatch = true
|
||||
s.switchLoginWatchUntil = now.Add(30 * time.Second)
|
||||
s.switchMu.Unlock()
|
||||
s.emitter.Emit(EventStatusSnapshot, Status{Status: StatusConnecting})
|
||||
s.pushStatus(Status{Status: StatusConnecting})
|
||||
}
|
||||
|
||||
// CancelProfileSwitch aborts a switch midway (tray Disconnect while Connecting):
|
||||
@@ -343,7 +401,7 @@ func (s *DaemonFeed) statusStreamLoop(ctx context.Context) {
|
||||
return
|
||||
}
|
||||
unavailable = true
|
||||
s.emitter.Emit(EventStatusSnapshot, Status{Status: StatusDaemonUnavailable})
|
||||
s.pushStatus(Status{Status: StatusDaemonUnavailable})
|
||||
}
|
||||
|
||||
op := func() error {
|
||||
@@ -403,7 +461,7 @@ func (s *DaemonFeed) emitStatus(st Status) {
|
||||
log.Debugf("suppressing status=%q during profile switch", st.Status)
|
||||
return
|
||||
}
|
||||
s.emitter.Emit(EventStatusSnapshot, st)
|
||||
s.pushStatus(st)
|
||||
if triggerLogin {
|
||||
s.emitter.Emit(EventTriggerLogin)
|
||||
}
|
||||
|
||||
24
client/ui/services/shutdown.go
Normal file
24
client/ui/services/shutdown.go
Normal file
@@ -0,0 +1,24 @@
|
||||
package services
|
||||
|
||||
import "sync/atomic"
|
||||
|
||||
var (
|
||||
sessionEnding atomic.Bool
|
||||
quitting atomic.Bool
|
||||
)
|
||||
|
||||
func BeginSessionEnd() {
|
||||
sessionEnding.Store(true)
|
||||
}
|
||||
|
||||
func AbortSessionEnd() {
|
||||
sessionEnding.Store(false)
|
||||
}
|
||||
|
||||
func BeginShutdown() {
|
||||
quitting.Store(true)
|
||||
}
|
||||
|
||||
func ShuttingDown() bool {
|
||||
return sessionEnding.Load() || quitting.Load()
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/client/proto"
|
||||
"github.com/netbirdio/netbird/client/ui/updater"
|
||||
"github.com/netbirdio/netbird/version"
|
||||
)
|
||||
|
||||
// UpdateResult mirrors TriggerUpdateResponse.
|
||||
@@ -33,6 +34,12 @@ func (s *Update) GetState() updater.State {
|
||||
return s.holder.Get()
|
||||
}
|
||||
|
||||
// DownloadURL returns the platform-appropriate installer download link for
|
||||
// manual (non-enforced) updates.
|
||||
func (s *Update) DownloadURL() string {
|
||||
return version.DownloadUrl()
|
||||
}
|
||||
|
||||
// Quit exits the app. Scheduled off the calling goroutine so the JS caller's
|
||||
// response returns before the runtime tears down.
|
||||
func (s *Update) Quit() {
|
||||
|
||||
@@ -115,6 +115,9 @@ type WindowManager struct {
|
||||
// recenterOnShow is set only on the minimal-WM/XEmbed path, where the WM neither centers nor
|
||||
// restores position; nil on full desktops so re-centering can't fight a user-moved window.
|
||||
recenterOnShow func() bool
|
||||
// showReplay fires whenever a live-but-hidden window is (re)shown, so the
|
||||
// caller can replay the latest status snapshot into its webview.
|
||||
showReplay func(application.Window)
|
||||
}
|
||||
|
||||
// NewWindowManager wires the manager to the main app; translator/prefs may be nil (tests). The
|
||||
@@ -154,6 +157,9 @@ func NewWindowManager(app *application.App, mainWindow *application.WebviewWindo
|
||||
})
|
||||
// Hide (not destroy) on close to keep React state; reset to General for a flash-free reopen.
|
||||
s.settings.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
|
||||
if ShuttingDown() {
|
||||
return
|
||||
}
|
||||
e.Cancel()
|
||||
s.app.Event.Emit(EventSettingsOpen, "general")
|
||||
s.settings.Hide()
|
||||
@@ -171,6 +177,7 @@ func (s *WindowManager) OpenSettings(tab string) {
|
||||
s.app.Event.Emit(EventSettingsOpen, target)
|
||||
s.settings.Show()
|
||||
s.settings.Focus()
|
||||
s.notifyShown(s.settings)
|
||||
// Re-center (minimal-WM only; see centerWhenReady).
|
||||
s.centerWhenReady(s.settings)
|
||||
}
|
||||
@@ -217,6 +224,7 @@ func (s *WindowManager) OpenBrowserLogin(uri string) {
|
||||
s.centerOnCursorScreen(s.browserLogin)
|
||||
s.browserLogin.Show()
|
||||
s.browserLogin.Focus()
|
||||
s.notifyShown(s.browserLogin)
|
||||
}
|
||||
|
||||
// BrowserLoginWindow returns the live SSO popup, or nil. While non-nil it is the
|
||||
@@ -277,6 +285,7 @@ func (s *WindowManager) OpenSessionExpiration(seconds int) {
|
||||
s.centerOnCursorScreen(s.sessionExpiration)
|
||||
s.sessionExpiration.Show()
|
||||
s.sessionExpiration.Focus()
|
||||
s.notifyShown(s.sessionExpiration)
|
||||
}
|
||||
|
||||
func (s *WindowManager) CloseSessionExpiration() {
|
||||
@@ -344,6 +353,7 @@ func (s *WindowManager) OpenInstallProgress(version string) {
|
||||
s.installProgress.SetURL(startURL)
|
||||
s.installProgress.Show()
|
||||
s.installProgress.Focus()
|
||||
s.notifyShown(s.installProgress)
|
||||
s.centerWhenReady(s.installProgress)
|
||||
}
|
||||
|
||||
@@ -377,6 +387,7 @@ func (s *WindowManager) OpenWelcome() {
|
||||
}
|
||||
s.welcome.Show()
|
||||
s.welcome.Focus()
|
||||
s.notifyShown(s.welcome)
|
||||
s.centerWhenReady(s.welcome)
|
||||
}
|
||||
|
||||
@@ -393,6 +404,9 @@ func (s *WindowManager) CloseWelcome() {
|
||||
// OpenError shows the custom error dialog; title/message are pre-localised and ride in the
|
||||
// start URL. A second error replaces the open one via SetURL. Singleton, destroyed on close.
|
||||
func (s *WindowManager) OpenError(title, message string) {
|
||||
if ShuttingDown() {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
startURL := errorDialogURL(title, message)
|
||||
@@ -411,6 +425,7 @@ func (s *WindowManager) OpenError(title, message string) {
|
||||
s.errorDialog.SetURL(startURL)
|
||||
s.errorDialog.Show()
|
||||
s.errorDialog.Focus()
|
||||
s.notifyShown(s.errorDialog)
|
||||
s.centerWhenReady(s.errorDialog)
|
||||
}
|
||||
|
||||
@@ -437,6 +452,7 @@ func (s *WindowManager) ShowMain() {
|
||||
}
|
||||
s.mainWindow.Show()
|
||||
s.mainWindow.Focus()
|
||||
s.notifyShown(s.mainWindow)
|
||||
// Re-center (minimal-WM only; see centerWhenReady).
|
||||
s.centerWhenReady(s.mainWindow)
|
||||
}
|
||||
@@ -446,6 +462,17 @@ func (s *WindowManager) SetRecenterOnShow(pred func() bool) {
|
||||
s.recenterOnShow = pred
|
||||
}
|
||||
|
||||
// SetShowReplay installs the shown-window hook (see the showReplay field).
|
||||
func (s *WindowManager) SetShowReplay(fn func(application.Window)) {
|
||||
s.showReplay = fn
|
||||
}
|
||||
|
||||
func (s *WindowManager) notifyShown(w application.Window) {
|
||||
if s.showReplay != nil && w != nil {
|
||||
s.showReplay(w)
|
||||
}
|
||||
}
|
||||
|
||||
// centerWhenReady centers w only on minimal WMs (recenterOnShow); elsewhere it
|
||||
// returns so it never fights a user-moved window. On GTK4 an inline Center()
|
||||
// no-ops until the GdkSurface is realized (async, after Show) and InvokeAsync
|
||||
@@ -566,6 +593,7 @@ func (s *WindowManager) restoreHiddenWindowsLocked() {
|
||||
continue
|
||||
}
|
||||
w.Show()
|
||||
s.notifyShown(w)
|
||||
if w == s.mainWindow {
|
||||
mainRestored = true
|
||||
}
|
||||
|
||||
7
client/ui/shutdown_other.go
Normal file
7
client/ui/shutdown_other.go
Normal file
@@ -0,0 +1,7 @@
|
||||
//go:build !windows && !android && !ios && !freebsd && !js
|
||||
|
||||
package main
|
||||
|
||||
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
|
||||
return nil
|
||||
}
|
||||
36
client/ui/shutdown_windows.go
Normal file
36
client/ui/shutdown_windows.go
Normal file
@@ -0,0 +1,36 @@
|
||||
//go:build windows
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ui/services"
|
||||
)
|
||||
|
||||
const (
|
||||
wmQueryEndSession = 0x0011
|
||||
wmEndSession = 0x0016
|
||||
)
|
||||
|
||||
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
|
||||
return func(_ uintptr, msg uint32, wParam, _ uintptr) (uintptr, bool) {
|
||||
switch msg {
|
||||
case wmQueryEndSession:
|
||||
services.BeginSessionEnd()
|
||||
return 1, true
|
||||
case wmEndSession:
|
||||
if wParam == 0 {
|
||||
services.AbortSessionEnd()
|
||||
return 0, true
|
||||
}
|
||||
log.Info("windows session is ending; exiting immediately")
|
||||
os.Exit(0)
|
||||
return 0, true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -32,9 +32,8 @@ const (
|
||||
|
||||
quitDownTimeout = 5 * time.Second
|
||||
|
||||
urlGitHubRepo = "https://github.com/netbirdio/netbird"
|
||||
urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
|
||||
urlDocs = "https://docs.netbird.io"
|
||||
urlGitHubRepo = "https://github.com/netbirdio/netbird"
|
||||
urlDocs = "https://docs.netbird.io"
|
||||
)
|
||||
|
||||
// TrayServices bundles the services the tray menu needs, grouped so NewTray
|
||||
@@ -68,9 +67,8 @@ type Tray struct {
|
||||
loc *Localizer
|
||||
|
||||
// menu and the *Item/*Submenu fields below are reassigned by buildMenu
|
||||
// on every relayout — touch them only with menuMu held. Exceptions:
|
||||
// the Connect/Disconnect OnClick closures capture their own item, and
|
||||
// refreshSessionExpiresLabel snapshots its item under menuMu.
|
||||
// on every relayout, which destroys the replaced tree — locate items and
|
||||
// call them with menuMu held, so a relayout can't destroy one mid-call.
|
||||
menu *application.Menu
|
||||
statusItem *application.MenuItem
|
||||
// sessionExpiresItem shows the SSO deadline as a remaining-time label,
|
||||
@@ -173,7 +171,7 @@ func NewTray(app *application.App, window *application.WebviewWindow, svc TraySe
|
||||
// in the right locale — no English flash then re-paint.
|
||||
loc: svc.Localizer,
|
||||
}
|
||||
t.updater = newTrayUpdater(app, window, svc.Update, svc.Notifier, t.loc, func() { t.applyIcon() }, func() { t.relayoutMenu() })
|
||||
t.updater = newTrayUpdater(app, window, svc.Update, svc.Notifier, t.loc, func() { t.applyIcon() }, func() { t.relayoutMenu() }, func() { t.showMainWindow() })
|
||||
t.tray = app.SystemTray.New()
|
||||
// Seed panel-theme detection before the first paint so the initial icon
|
||||
// matches the panel's light/dark scheme (Linux only).
|
||||
@@ -197,7 +195,7 @@ func NewTray(app *application.App, window *application.WebviewWindow, svc TraySe
|
||||
// menu (e.g. GNOME Shell AppIndicator).
|
||||
bindTrayClick(t)
|
||||
|
||||
app.Event.On(services.EventStatusSnapshot, t.onStatusEvent)
|
||||
svc.DaemonFeed.OnStatus(t.applyStatus)
|
||||
app.Event.On(services.EventDaemonNotification, t.onSystemEvent)
|
||||
// Refresh the Profiles submenu on ProfileSwitcher's change event. A
|
||||
// switch on an idle daemon drives no status transition, so without this
|
||||
@@ -254,6 +252,19 @@ func (t *Tray) ShowWindow() {
|
||||
t.window.Focus()
|
||||
}
|
||||
|
||||
// showMainWindow brings the main window forward through the WindowManager so
|
||||
// the status replay and re-centering apply; falls back to a bare Show in tests.
|
||||
func (t *Tray) showMainWindow() {
|
||||
if t.svc.WindowManager != nil {
|
||||
t.svc.WindowManager.ShowMain()
|
||||
return
|
||||
}
|
||||
if t.window != nil {
|
||||
t.window.Show()
|
||||
t.window.Focus()
|
||||
}
|
||||
}
|
||||
|
||||
// applyLanguage re-renders every translated surface in the Localizer's current
|
||||
// language. Wails dispatches menu/tray APIs onto the UI thread internally, so
|
||||
// calling them from the Localizer's background goroutine is safe; profileLoadMu
|
||||
@@ -287,6 +298,7 @@ func (t *Tray) relayoutMenu() {
|
||||
t.menuMu.Lock()
|
||||
defer t.menuMu.Unlock()
|
||||
|
||||
old := t.menu
|
||||
t.menu = t.buildMenu()
|
||||
|
||||
t.statusMu.Lock()
|
||||
@@ -357,6 +369,10 @@ func (t *Tray) relayoutMenu() {
|
||||
// Single push of the whole tree: on Linux one LayoutUpdated with fresh
|
||||
// container ids; on darwin an NSMenu rebuild against the cached pointer.
|
||||
t.tray.SetMenu(t.menu)
|
||||
|
||||
if old != nil {
|
||||
old.Destroy()
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Tray) buildMenu() *application.Menu {
|
||||
@@ -372,13 +388,11 @@ func (t *Tray) buildMenu() *application.Menu {
|
||||
|
||||
menu.AddSeparator()
|
||||
|
||||
// The OnClick closures capture the local item because t.upItem/t.downItem
|
||||
// are menuMu-guarded and must not be read from the click goroutine.
|
||||
upItem := menu.Add(t.loc.T("tray.menu.connect"))
|
||||
upItem.OnClick(func(*application.Context) { t.handleConnect(upItem) })
|
||||
upItem.OnClick(func(*application.Context) { t.handleConnect() })
|
||||
t.upItem = upItem
|
||||
downItem := menu.Add(t.loc.T("tray.menu.disconnect"))
|
||||
downItem.OnClick(func(*application.Context) { t.handleDisconnect(downItem) })
|
||||
downItem.OnClick(func(*application.Context) { t.handleDisconnect() })
|
||||
downItem.SetHidden(true)
|
||||
t.downItem = downItem
|
||||
|
||||
@@ -453,6 +467,7 @@ func (t *Tray) buildMenu() *application.Menu {
|
||||
}
|
||||
|
||||
func (t *Tray) handleQuit() {
|
||||
services.BeginShutdown()
|
||||
t.profileMu.Lock()
|
||||
if t.switchCancel != nil {
|
||||
t.switchCancel()
|
||||
@@ -469,9 +484,7 @@ func (t *Tray) handleQuit() {
|
||||
t.app.Quit()
|
||||
}
|
||||
|
||||
// handleConnect receives the clicked item from the buildMenu closure —
|
||||
// t.upItem is menuMu-guarded and must not be read here.
|
||||
func (t *Tray) handleConnect(upItem *application.MenuItem) {
|
||||
func (t *Tray) handleConnect() {
|
||||
// NeedsLogin/SessionExpired/LoginFailed won't honor a plain Up RPC — they
|
||||
// need the Login → WaitSSOLogin → Up sequence. Emit EventTriggerLogin so
|
||||
// the React startLogin() (which owns the BrowserLogin popup) drives it;
|
||||
@@ -485,7 +498,7 @@ func (t *Tray) handleConnect(upItem *application.MenuItem) {
|
||||
t.app.Event.Emit(services.EventTriggerLogin)
|
||||
return
|
||||
}
|
||||
upItem.SetEnabled(false)
|
||||
t.setItemEnabled(func() *application.MenuItem { return t.upItem }, false)
|
||||
// Arm the SSO auto-handoff: Up() is async and the daemon may flip to
|
||||
// NeedsLogin on an SSO peer with no cached token. applyStatus consumes the
|
||||
// flag on that transition to trigger browser-login without a second Connect
|
||||
@@ -500,18 +513,25 @@ func (t *Tray) handleConnect(upItem *application.MenuItem) {
|
||||
t.statusMu.Lock()
|
||||
t.pendingConnectLogin = false
|
||||
t.statusMu.Unlock()
|
||||
upItem.SetEnabled(true)
|
||||
t.setItemEnabled(func() *application.MenuItem { return t.upItem }, true)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (t *Tray) setItemEnabled(get func() *application.MenuItem, enabled bool) {
|
||||
t.menuMu.Lock()
|
||||
defer t.menuMu.Unlock()
|
||||
if item := get(); item != nil {
|
||||
item.SetEnabled(enabled)
|
||||
}
|
||||
}
|
||||
|
||||
// handleDisconnect aborts any in-flight profile switch before sending Down —
|
||||
// otherwise the switcher's queued Up would reconnect right after, making the
|
||||
// click a no-op. Also clears Peers' optimistic-Connecting guard so the daemon's
|
||||
// Idle push paints through instead of being swallowed by the suppression filter.
|
||||
// Receives the clicked item from the buildMenu closure (see handleConnect).
|
||||
func (t *Tray) handleDisconnect(downItem *application.MenuItem) {
|
||||
downItem.SetEnabled(false)
|
||||
func (t *Tray) handleDisconnect() {
|
||||
t.setItemEnabled(func() *application.MenuItem { return t.downItem }, false)
|
||||
t.profileMu.Lock()
|
||||
if t.switchCancel != nil {
|
||||
t.switchCancel()
|
||||
@@ -523,7 +543,7 @@ func (t *Tray) handleDisconnect(downItem *application.MenuItem) {
|
||||
if err := t.svc.Connection.Down(context.Background()); err != nil {
|
||||
log.Errorf("disconnect: %v", err)
|
||||
t.notifyError(t.loc.T("notify.error.disconnect"))
|
||||
downItem.SetEnabled(true)
|
||||
t.setItemEnabled(func() *application.MenuItem { return t.downItem }, true)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -25,6 +25,9 @@ type sendFn func(notifications.NotificationOptions) error
|
||||
// event-dispatch goroutine that panic is fatal process-wide; recover() turns
|
||||
// it into a logged no-op.
|
||||
func safeSendNotification(send sendFn, what string, opts notifications.NotificationOptions) (err error) {
|
||||
if services.ShuttingDown() {
|
||||
return nil
|
||||
}
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
log.Errorf("notify %s: recovered from panic (notification bus unavailable): %v", what, r)
|
||||
|
||||
@@ -5,6 +5,7 @@ package main
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
@@ -59,10 +60,11 @@ func (t *Tray) loadConfig() {
|
||||
t.profileMu.Unlock()
|
||||
}
|
||||
|
||||
// loadProfiles fetches the profile list and relayouts the menu. Also called
|
||||
// from applyStatus to catch flips from another channel (CLI, autoconnect),
|
||||
// since the daemon emits no active-profile event. Full relayout (not
|
||||
// Clear()+Add()) is required for KDE/Plasma — see relayoutMenu's doc comment.
|
||||
// loadProfiles fetches the profile list and relayouts the menu when the rows
|
||||
// changed. Also called from applyStatus to catch flips from another channel
|
||||
// (CLI, autoconnect), since the daemon emits no active-profile event. Full
|
||||
// relayout (not Clear()+Add()) is required for KDE/Plasma — see relayoutMenu's
|
||||
// doc comment.
|
||||
func (t *Tray) loadProfiles() {
|
||||
t.profileLoadMu.Lock()
|
||||
defer t.profileLoadMu.Unlock()
|
||||
@@ -80,11 +82,14 @@ func (t *Tray) loadProfiles() {
|
||||
}
|
||||
|
||||
t.profilesMu.Lock()
|
||||
changed := username != t.profilesUser || !slices.Equal(profiles, t.profiles)
|
||||
t.profiles = profiles
|
||||
t.profilesUser = username
|
||||
t.profilesMu.Unlock()
|
||||
|
||||
t.relayoutMenu()
|
||||
if changed {
|
||||
t.relayoutMenu()
|
||||
}
|
||||
}
|
||||
|
||||
// fillProfileSubmenu paints cached profile rows into the freshly built submenu.
|
||||
|
||||
@@ -30,11 +30,11 @@ const (
|
||||
// handleSessionExpired notifies and brings the window forward so the frontend's /login route drives renewal.
|
||||
func (t *Tray) handleSessionExpired() {
|
||||
t.notify(t.loc.T("notify.sessionExpired.title"), t.loc.T("notify.sessionExpired.body"), notifyIDSessionExpired)
|
||||
if t.window != nil {
|
||||
t.window.SetURL("/#/login")
|
||||
t.window.Show()
|
||||
t.window.Focus()
|
||||
if t.window == nil {
|
||||
return
|
||||
}
|
||||
t.window.SetURL("/#/login")
|
||||
t.showMainWindow()
|
||||
}
|
||||
|
||||
// applySessionExpiry refreshes the cached SSO deadline and reports whether it changed.
|
||||
@@ -104,21 +104,19 @@ func sessionRefreshInterval(remaining time.Duration) time.Duration {
|
||||
}
|
||||
|
||||
// refreshSessionExpiresLabel updates only the countdown label, no relayout, to avoid disturbing an open menu.
|
||||
// The item is snapshotted under menuMu since buildMenu reassigns it on every relayout.
|
||||
// menuMu is held across the SetLabel so a relayout can't destroy the item mid-call.
|
||||
func (t *Tray) refreshSessionExpiresLabel() {
|
||||
t.menuMu.Lock()
|
||||
item := t.sessionExpiresItem
|
||||
t.menuMu.Unlock()
|
||||
if item == nil {
|
||||
return
|
||||
}
|
||||
t.sessionMu.Lock()
|
||||
deadline := t.sessionExpiresAt
|
||||
t.sessionMu.Unlock()
|
||||
if deadline.IsZero() {
|
||||
return
|
||||
}
|
||||
item.SetLabel(t.sessionRowLabel(deadline))
|
||||
t.menuMu.Lock()
|
||||
defer t.menuMu.Unlock()
|
||||
if t.sessionExpiresItem != nil {
|
||||
t.sessionExpiresItem.SetLabel(t.sessionRowLabel(deadline))
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Tray) sessionRowLabel(deadline time.Time) string {
|
||||
@@ -310,8 +308,7 @@ func (t *Tray) openSessionExtendFlow() {
|
||||
if seconds <= 0 {
|
||||
if t.window != nil {
|
||||
t.window.SetURL("/#/login")
|
||||
t.window.Show()
|
||||
t.window.Focus()
|
||||
t.showMainWindow()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
@@ -5,19 +5,9 @@ package main
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/wailsapp/wails/v3/pkg/application"
|
||||
|
||||
"github.com/netbirdio/netbird/client/ui/services"
|
||||
)
|
||||
|
||||
func (t *Tray) onStatusEvent(ev *application.CustomEvent) {
|
||||
st, ok := ev.Data.(services.Status)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
t.applyStatus(st)
|
||||
}
|
||||
|
||||
// applyStatus repaints the tray from a daemon snapshot. Icon refresh is skipped
|
||||
// when no icon-relevant input changed: the daemon emits rapid SubscribeStatus
|
||||
// bursts during health probes that would otherwise spam Shell_NotifyIcon.
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/client/ui/services"
|
||||
"github.com/netbirdio/netbird/client/ui/updater"
|
||||
"github.com/netbirdio/netbird/version"
|
||||
)
|
||||
|
||||
// trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
|
||||
@@ -27,6 +28,9 @@ type trayUpdater struct {
|
||||
// About submenu, which KDE/Plasma caches on first open and never re-fetches
|
||||
// on a plain SetLabel/SetHidden — only a relayout (fresh submenu ids) repaints.
|
||||
onMenuChange func()
|
||||
// showMain brings the main window forward via the WindowManager (status
|
||||
// replay + centering); nil falls back to a bare window.Show.
|
||||
showMain func()
|
||||
|
||||
mu sync.Mutex
|
||||
item *application.MenuItem
|
||||
@@ -35,7 +39,7 @@ type trayUpdater struct {
|
||||
progressWindowOpen bool
|
||||
}
|
||||
|
||||
func newTrayUpdater(app *application.App, window *application.WebviewWindow, update *services.Update, notifier *notifications.NotificationService, loc *Localizer, onIconChange func(), onMenuChange func()) *trayUpdater {
|
||||
func newTrayUpdater(app *application.App, window *application.WebviewWindow, update *services.Update, notifier *notifications.NotificationService, loc *Localizer, onIconChange, onMenuChange, showMain func()) *trayUpdater {
|
||||
u := &trayUpdater{
|
||||
app: app,
|
||||
window: window,
|
||||
@@ -44,6 +48,7 @@ func newTrayUpdater(app *application.App, window *application.WebviewWindow, upd
|
||||
loc: loc,
|
||||
onIconChange: onIconChange,
|
||||
onMenuChange: onMenuChange,
|
||||
showMain: showMain,
|
||||
}
|
||||
app.Event.On(updater.EventStateChanged, u.onStateEvent)
|
||||
// Seed from cached state to cover an event that fired before wiring completed.
|
||||
@@ -76,15 +81,15 @@ func (u *trayUpdater) applyLanguage() {
|
||||
u.refreshMenuItem(state)
|
||||
}
|
||||
|
||||
// handleClick opens the GitHub releases page when not Enforced, otherwise shows
|
||||
// the progress page and asks the daemon to start the installer.
|
||||
// handleClick opens the installer download link when not Enforced, otherwise
|
||||
// shows the progress page and asks the daemon to start the installer.
|
||||
func (u *trayUpdater) handleClick() {
|
||||
u.mu.Lock()
|
||||
state := u.state
|
||||
u.mu.Unlock()
|
||||
|
||||
if !state.Enforced {
|
||||
_ = u.app.Browser.OpenURL(urlGitHubReleases)
|
||||
_ = u.app.Browser.OpenURL(version.DownloadUrl())
|
||||
return
|
||||
}
|
||||
|
||||
@@ -192,6 +197,10 @@ func (u *trayUpdater) openProgressWindow(version string) {
|
||||
url += "?version=" + version
|
||||
}
|
||||
u.window.SetURL(url)
|
||||
if u.showMain != nil {
|
||||
u.showMain()
|
||||
return
|
||||
}
|
||||
u.window.Show()
|
||||
u.window.Focus()
|
||||
}
|
||||
|
||||
@@ -109,7 +109,7 @@ sequenceDiagram
|
||||
Chk->>Inj: continue
|
||||
Inj->>Inj: inject NetBird identity headers per provider config
|
||||
Inj->>Grd: continue
|
||||
Grd->>Grd: enforce model allowlist
|
||||
Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
|
||||
Grd->>Up: forward (over WireGuard)
|
||||
Up-->>Resp: response (JSON or SSE stream)
|
||||
Resp->>Resp: parse usage tokens, completion
|
||||
@@ -135,6 +135,21 @@ sequenceDiagram
|
||||
(`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
|
||||
PII names — see `redact.go` for the full set. See
|
||||
[`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
|
||||
- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
|
||||
is authoritative: it resolves the policy that governs this
|
||||
(provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
|
||||
when no applicable policy permits the model — so an allowlist scoped to
|
||||
one group/provider never leaks to another, and an un-guardrailed policy
|
||||
is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
|
||||
backstop: it only carries an allowlist for a provider every authorising
|
||||
policy restricts, and blocks unknown/undetermined models even when
|
||||
management is unreachable. Because that backstop allowlist is the UNION
|
||||
of every restricting policy's models, per-group narrowing lives only in
|
||||
the authoritative check: during a `CheckLLMPolicyLimits` outage
|
||||
`llm_limit_check` fails open, so a caller can reach any model in the
|
||||
provider's union — a group scoped to model A could reach model B if
|
||||
another group restricts the same provider to B. This is the documented
|
||||
fail-open trade-off; a future flag may switch it to fail-closed.
|
||||
- SSE streaming requires special handling on the response side; the
|
||||
parser must handle partial chunks without buffering the whole
|
||||
stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).
|
||||
|
||||
@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
|
||||
| on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
|
||||
| on_request | 2 | `llm_limit_check` | `{}` | – |
|
||||
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
|
||||
| on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
|
||||
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
|
||||
| on_response | 6 | `cost_meter` | `{}` | – |
|
||||
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |
|
||||
|
||||
@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
|
||||
| `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
|
||||
| `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
|
||||
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
|
||||
| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
|
||||
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
|
||||
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
|
||||
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
|
||||
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |
|
||||
|
||||
@@ -9,12 +9,220 @@ 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
|
||||
@@ -116,12 +324,12 @@ func availableProviders() []providerCase {
|
||||
if region == "" {
|
||||
region = "eu-central-1"
|
||||
}
|
||||
// A valid Bedrock inference-profile id (region prefix + date + version),
|
||||
// overridable per account. `global.` profiles can be invoked from any
|
||||
// region; set AWS_BEDROCK_MODEL to match the enabled profile for the token.
|
||||
// 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-haiku-4-5-20251001-v1:0"
|
||||
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})
|
||||
}
|
||||
@@ -257,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
|
||||
@@ -267,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.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, 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
|
||||
@@ -290,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 {
|
||||
@@ -297,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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -322,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)
|
||||
}
|
||||
|
||||
209
e2e/agentnetwork/guardrail_block_test.go
Normal file
209
e2e/agentnetwork/guardrail_block_test.go
Normal file
@@ -0,0 +1,209 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
|
||||
// provider, its own guardrail whose allowlist holds ONLY that provider's
|
||||
// allowed model, and its own policy. Each case runs in isolation (its own
|
||||
// proxy + client), so the guardrail the proxy enforces contains exactly this
|
||||
// provider's model — never a mixed cross-provider list.
|
||||
type pathRoutedGuardrailCase struct {
|
||||
name string
|
||||
catalogID string // agent-network catalog provider id
|
||||
wire string // harness.WireVertex | harness.WireBedrock
|
||||
allowEntry string // the single model id put on the guardrail allowlist
|
||||
allowModel string // model id sent that MUST be served (200)
|
||||
blockModel string // model id sent that MUST be denied (403 model_blocked)
|
||||
}
|
||||
|
||||
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
|
||||
// guard for the customer report that a model-allowlist guardrail attached to a
|
||||
// policy has no effect for PATH-ROUTED providers — where the model travels in
|
||||
// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
|
||||
// AWS Bedrock (/model/{id}/invoke).
|
||||
//
|
||||
// Each provider is tested in isolation with a guardrail allowlisting a single
|
||||
// model of its own: the allowed model (in the URL path) is served (200) and an
|
||||
// unselected model (in the URL path) is denied 403 by the guardrail
|
||||
// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
|
||||
// customer verbatim — allow Sonnet, and the unselected model is the exact
|
||||
// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
|
||||
// sends a region-prefixed, versioned inference-profile id so URL-path model
|
||||
// normalization is exercised too.
|
||||
//
|
||||
// The provider is catch-all (no models), so the router forwards any model and a
|
||||
// 403 can only come from the guardrail, never model_not_routable. Only the
|
||||
// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
|
||||
// synth/reconcile, the proxy middleware chain (URL-path model extraction,
|
||||
// router, guardrail) and the tunnel are all real, and the guardrail denies
|
||||
// before the upstream is dialed so the mock cannot influence the block. A
|
||||
// static bearer api key is used so the router injects a static Authorization
|
||||
// header instead of minting a GCP token — the only reason path-routed providers
|
||||
// normally need live credentials — so the test runs with none and is always on.
|
||||
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
|
||||
cases := []pathRoutedGuardrailCase{
|
||||
{
|
||||
name: "vertex",
|
||||
catalogID: "vertex_ai_api",
|
||||
wire: harness.WireVertex,
|
||||
allowEntry: "claude-sonnet-4-5",
|
||||
allowModel: "claude-sonnet-4-5",
|
||||
blockModel: "claude-opus-4-6", // the customer-reported model
|
||||
},
|
||||
{
|
||||
name: "bedrock",
|
||||
catalogID: "bedrock_api",
|
||||
wire: harness.WireBedrock,
|
||||
allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
|
||||
allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
|
||||
blockModel: "us.anthropic.claude-opus-4-8-v1:0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
runPathRoutedGuardrailCase(t, tc)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
|
||||
t.Helper()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-" + tc.name + "-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
// Catch-all provider (no models) so the router forwards any model; a static
|
||||
// bearer key means the router injects a static auth header instead of minting
|
||||
// a GCP token. Bootstraps the cluster if it isn't already.
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: tc.name,
|
||||
ProviderId: tc.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create %s provider", tc.name)
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// Guardrail allowlisting ONLY this provider's allowed model.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-guardrail-" + tc.name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-" + tc.name,
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Probe first: the GET resolves the endpoint and its first packet wakes the
|
||||
// lazy proxy peer, so WaitProxyPeer then observes it connected.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch tc.wire {
|
||||
case harness.WireVertex:
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case harness.WireBedrock:
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
t.Fatalf("unsupported wire %q", tc.wire)
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
|
||||
// jitter on the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(tc.allowModel)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
|
||||
// Unselected model (in the URL path) must be blocked by the guardrail.
|
||||
code, body = send(tc.blockModel)
|
||||
assert.Equal(t, 403, code,
|
||||
"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
|
||||
}
|
||||
205
e2e/agentnetwork/guardrail_groupswitch_test.go
Normal file
205
e2e/agentnetwork/guardrail_groupswitch_test.go
Normal file
@@ -0,0 +1,205 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
|
||||
// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
|
||||
// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
|
||||
// the proxy caches; the switch is invisible until that cache expires. The proxy
|
||||
// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
|
||||
// instead of the 5-minute default.
|
||||
//
|
||||
// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
|
||||
// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
|
||||
// served and modelB denied; after switching the client grpA -> grpB, modelB is
|
||||
// served and modelA denied. The cross-group deny comes from management's
|
||||
// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
|
||||
func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelA = "e2e-model-a"
|
||||
modelB = "e2e-model-b"
|
||||
// Short tunnel-cache TTL so a group switch propagates in seconds.
|
||||
// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
|
||||
cacheTTL = 3 * time.Second
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
|
||||
require.NoError(t, err, "create group A")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
|
||||
|
||||
grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
|
||||
require.NoError(t, err, "create group B")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-gswitch-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpA.Id}, // client starts in group A
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "gswitch",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{
|
||||
{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
|
||||
},
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
mkGuard := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gA := mkGuard("e2e-gswitch-a", modelA)
|
||||
gB := mkGuard("e2e-gswitch-b", modelB)
|
||||
|
||||
enabled := true
|
||||
polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-a",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpA.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gA.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy A")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
|
||||
|
||||
polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-gswitch-b",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpB.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gB.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy B")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
|
||||
"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
|
||||
})
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == want {
|
||||
return code, body
|
||||
}
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
// Phase 1 — client is in group A: modelA served, modelB denied.
|
||||
code, body := sendUntil(modelA, 200, 90*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelB)
|
||||
assert.Equal(t, 403, code,
|
||||
"group-B model must be denied while the client is in group A; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
|
||||
// Switch the client peer from group A to group B.
|
||||
peerID := clientPeerInGroup(t, ctx, grpA.Id)
|
||||
_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpB.Name,
|
||||
Peers: &[]string{peerID},
|
||||
})
|
||||
require.NoError(t, err, "add peer to group B")
|
||||
_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
|
||||
Name: grpA.Name,
|
||||
Peers: &[]string{},
|
||||
})
|
||||
require.NoError(t, err, "remove peer from group A")
|
||||
|
||||
// Phase 2 — after the short TTL expires the proxy re-validates the peer,
|
||||
// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
|
||||
code, body = sendUntil(modelB, 200, 60*time.Second)
|
||||
assert.Equal(t, 200, code,
|
||||
"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
code, body = send(modelA)
|
||||
assert.Equal(t, 403, code,
|
||||
"after the switch, the old group-A model must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
}
|
||||
|
||||
// clientPeerInGroup returns the id of the single peer that is a member of the
|
||||
// given group — the test client. The proxy peer is never added to test groups.
|
||||
func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
|
||||
t.Helper()
|
||||
peers, err := srv.API().Peers.List(ctx)
|
||||
require.NoError(t, err, "list peers")
|
||||
for _, p := range peers {
|
||||
for _, g := range p.Groups {
|
||||
if g.Id == groupID {
|
||||
return p.Id
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Fatalf("no peer found in group %s", groupID)
|
||||
return ""
|
||||
}
|
||||
201
e2e/agentnetwork/guardrail_multipolicy_test.go
Normal file
201
e2e/agentnetwork/guardrail_multipolicy_test.go
Normal file
@@ -0,0 +1,201 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
|
||||
// modelOther denied for the grpMain client (403 model_blocked, no cross-group
|
||||
// leak), and openModel on the un-guardrailed policy's provider served (200).
|
||||
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
modelSelected = "e2e-selected"
|
||||
modelOther = "e2e-other"
|
||||
openModel = "e2e-open"
|
||||
)
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-guardrail-mp-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// pRestricted declares the two guardrailed models so routing is deterministic
|
||||
// (model -> provider). Created first, so it carries the bootstrap cluster.
|
||||
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "restricted",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(modelSelected, modelOther),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create restricted provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
|
||||
|
||||
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "open",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: models(openModel),
|
||||
})
|
||||
require.NoError(t, err, "create open provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
|
||||
|
||||
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, gerr := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, gerr, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
|
||||
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
|
||||
|
||||
enabled := true
|
||||
// polMain: grpMain restricted to modelSelected on pRestricted.
|
||||
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gMain.Id},
|
||||
})
|
||||
require.NoError(t, err, "create main policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
// polOther: grpOther restricted to modelOther on the SAME provider. The
|
||||
// client is not in grpOther, so modelOther must never be usable by it.
|
||||
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{pRestricted.Id},
|
||||
GuardrailIds: &[]string{gOther.Id},
|
||||
})
|
||||
require.NoError(t, err, "create other policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
|
||||
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
|
||||
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-guardrail-mp-open",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{pOpen.Id},
|
||||
})
|
||||
require.NoError(t, err, "create open policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("selected model allowed for its group", func(t *testing.T) {
|
||||
code, body := sendUntil200(modelSelected)
|
||||
assert.Equal(t, 200, code,
|
||||
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
|
||||
t.Run("other group's model does not leak", func(t *testing.T) {
|
||||
// modelOther is allowlisted only for grpOther. The grpMain client must be
|
||||
// denied by management's per-policy/group check — not waved through by an
|
||||
// account-wide union. This is the security-critical wrong-ALLOW guard.
|
||||
code, body := send(modelOther)
|
||||
assert.Equal(t, 403, code,
|
||||
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"denial must be a model-allowlist decision; body: %s", body)
|
||||
})
|
||||
|
||||
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
|
||||
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
|
||||
// old account-wide union would have blocked openModel (it is on no
|
||||
// allowlist); it must now be served — the false-DENY guard.
|
||||
code, body := sendUntil200(openModel)
|
||||
assert.Equal(t, 200, code,
|
||||
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
422
e2e/agentnetwork/guardrail_pergroup_providers_test.go
Normal file
422
e2e/agentnetwork/guardrail_pergroup_providers_test.go
Normal file
@@ -0,0 +1,422 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// pergroupCase describes one provider surface for the per-group allowlist matrix.
|
||||
// selectedReq/otherReq are the model identifiers as they travel in the request
|
||||
// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
|
||||
// otherAllow are the (normalized) forms the guardrail allowlist holds — for
|
||||
// Bedrock these differ from the request form so path normalization is exercised.
|
||||
type pergroupCase struct {
|
||||
name string
|
||||
catalogID string
|
||||
wire string // "chat", "messages", "vertex", "bedrock"
|
||||
models *[]api.AgentNetworkProviderModel
|
||||
selectedReq string
|
||||
selectedAllow string
|
||||
otherReq string
|
||||
otherAllow string
|
||||
|
||||
providerID string // filled during setup
|
||||
}
|
||||
|
||||
// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
|
||||
// allowlist end to end across every always-on provider surface, including the
|
||||
// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
|
||||
//
|
||||
// For each provider two policies target it: grpMain (the client) is allowed only
|
||||
// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
|
||||
// The client must get selectedReq served (200) and otherReq denied (403,
|
||||
// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
|
||||
// authoritative per-policy/group decision from management (the proxy per-provider
|
||||
// backstop carries the union of both models), so this also confirms management
|
||||
// receives the correct normalized model for path-routed providers.
|
||||
func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
vertexProject = "e2e-project"
|
||||
vertexRegion = "global"
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
cases := []*pergroupCase{
|
||||
{
|
||||
name: "openai", catalogID: "openai_api", wire: harness.WireChat,
|
||||
models: priced("oai-model-a", "oai-model-b"),
|
||||
selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
|
||||
otherReq: "oai-model-b", otherAllow: "oai-model-b",
|
||||
},
|
||||
{
|
||||
name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
|
||||
models: priced("ant-model-a", "ant-model-b"),
|
||||
selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
|
||||
otherReq: "ant-model-b", otherAllow: "ant-model-b",
|
||||
},
|
||||
{
|
||||
// Vertex catalog ids travel bare in the rawPredict path.
|
||||
name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
|
||||
selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
|
||||
otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
|
||||
},
|
||||
{
|
||||
// Bedrock request ids are region-prefixed/versioned; the parser
|
||||
// normalizes them to the catalog key the allowlist holds.
|
||||
name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
|
||||
selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
|
||||
otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
|
||||
},
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
|
||||
require.NoError(t, err, "create main group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
|
||||
|
||||
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
|
||||
require.NoError(t, err, "create other group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-pergroup-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grpMain.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
for i, c := range cases {
|
||||
req := api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-pergroup-" + c.name,
|
||||
ProviderId: c.catalogID,
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: c.models,
|
||||
}
|
||||
if i == 0 {
|
||||
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
|
||||
}
|
||||
prov, perr := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, perr, "create provider %s", c.name)
|
||||
c.providerID = prov.Id
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
|
||||
gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
|
||||
|
||||
polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-main",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpMain.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gSel.Id},
|
||||
})
|
||||
require.NoError(t, merr, "create main policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
|
||||
|
||||
polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-pergroup-" + c.name + "-other",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grpOther.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{gOth.Id},
|
||||
})
|
||||
require.NoError(t, oerr, "create other policy %s", c.name)
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
|
||||
}
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(c *pergroupCase, model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
var cerr error
|
||||
switch c.wire {
|
||||
case "vertex":
|
||||
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
|
||||
case "bedrock":
|
||||
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
|
||||
default:
|
||||
code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
|
||||
}
|
||||
require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
|
||||
return code, body
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
|
||||
// the first call over the freshly warmed tunnel.
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(c, c.selectedReq)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
assert.Equal(t, 200, code,
|
||||
"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
|
||||
// grpOther's model must NOT leak to the grpMain client.
|
||||
code, body = send(c, c.otherReq)
|
||||
assert.Equal(t, 403, code,
|
||||
"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
|
||||
assert.Contains(t, body, "llm_policy.model_blocked",
|
||||
"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
|
||||
// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
|
||||
//
|
||||
// - union across the user's groups: the client is in gUX and gUY, each with
|
||||
// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
|
||||
// client may use BOTH models (the union of its groups' allowlists) while a
|
||||
// third, un-allowlisted model is denied.
|
||||
// - an un-guardrailed group lifts the restriction: the client is in gMP and
|
||||
// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
|
||||
// NO guardrail. Because one applicable policy is unrestricted, the client may
|
||||
// use a model on no allowlist (mix-z) as well as mix-a.
|
||||
func TestGuardrailMultiGroupUser(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
const (
|
||||
unionA = "mg-union-a"
|
||||
unionB = "mg-union-b"
|
||||
unionC = "mg-union-c" // allowlisted by neither group
|
||||
mixA = "mg-mix-a"
|
||||
mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
|
||||
)
|
||||
|
||||
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
|
||||
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
vllm, err := harness.StartVLLM(ctx, srv)
|
||||
require.NoError(t, err, "start mock upstream")
|
||||
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
||||
|
||||
mkGroup := func(name string) *api.Group {
|
||||
g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
|
||||
require.NoError(t, gerr, "create group %s", name)
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
gUX := mkGroup("e2e-mg-union-x")
|
||||
gUY := mkGroup("e2e-mg-union-y")
|
||||
gMP := mkGroup("e2e-mg-mix-p")
|
||||
gMQ := mkGroup("e2e-mg-mix-q")
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-mg-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
||||
|
||||
staticKey := "static-e2e-token"
|
||||
enabled := true
|
||||
|
||||
// P1 — union scenario: two restricting policies, one per group.
|
||||
p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-union",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(unionA, unionB, unionC),
|
||||
BootstrapCluster: ptr(harness.AgentNetworkCluster),
|
||||
})
|
||||
require.NoError(t, err, "create union provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
|
||||
|
||||
polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-x",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUX.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy X")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
|
||||
|
||||
polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-union-y",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gUY.Id},
|
||||
DestinationProviderIds: []string{p1.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
|
||||
})
|
||||
require.NoError(t, err, "create union policy Y")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
|
||||
|
||||
// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
|
||||
p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: "e2e-mg-mix",
|
||||
ProviderId: "openai_api",
|
||||
UpstreamUrl: vllm.URL,
|
||||
ApiKey: &staticKey,
|
||||
Enabled: ptr(true),
|
||||
Models: priced(mixA, mixZ),
|
||||
})
|
||||
require.NoError(t, err, "create mix provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
|
||||
|
||||
polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-p",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMP.Id},
|
||||
DestinationProviderIds: []string{p2.Id},
|
||||
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
|
||||
})
|
||||
require.NoError(t, err, "create mix policy P")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
|
||||
|
||||
polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-mg-mix-q",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{gMQ.Id},
|
||||
DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
|
||||
})
|
||||
require.NoError(t, err, "create mix policy Q")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings")
|
||||
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
|
||||
require.NoError(t, err, "mint proxy token")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
send := func(model string) (int, string) {
|
||||
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
|
||||
require.NoError(t, cerr, "request must reach the proxy")
|
||||
return code, body
|
||||
}
|
||||
sendUntil200 := func(model string) (int, string) {
|
||||
var code int
|
||||
var body string
|
||||
deadline := time.Now().Add(90 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
code, body = send(model)
|
||||
if code == 200 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
return code, body
|
||||
}
|
||||
|
||||
t.Run("union across the user's groups", func(t *testing.T) {
|
||||
code, body := sendUntil200(unionA)
|
||||
assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(unionB)
|
||||
assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
|
||||
|
||||
code, body = send(unionC)
|
||||
assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
|
||||
assert.Contains(t, body, "llm_policy.model_blocked")
|
||||
})
|
||||
|
||||
t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
|
||||
code, body := sendUntil200(mixA)
|
||||
assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
|
||||
code, body = sendUntil200(mixZ)
|
||||
assert.Equal(t, 200, code,
|
||||
"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
|
||||
})
|
||||
}
|
||||
|
||||
// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
|
||||
// exactly the given model, registering cleanup.
|
||||
func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
|
||||
t.Helper()
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{model}
|
||||
g, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail %s", name)
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
|
||||
return g
|
||||
}
|
||||
@@ -256,7 +256,7 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi
|
||||
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)
|
||||
@@ -271,7 +271,7 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi
|
||||
// is sent as the universal x-session-id header the proxy records.
|
||||
func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt, sessionID string) (int, string, error) {
|
||||
path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", project, region, model)
|
||||
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":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))
|
||||
}
|
||||
|
||||
@@ -282,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))
|
||||
}
|
||||
|
||||
@@ -341,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 ""
|
||||
@@ -351,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)
|
||||
|
||||
@@ -38,7 +38,11 @@ type Proxy struct {
|
||||
// network, registered via the given account proxy token and serving the
|
||||
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
|
||||
// peer connectivity — callers poll management for the proxy peer.
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
|
||||
// StartProxy launches the reverse-proxy container. Optional envOverrides are
|
||||
// merged into the container environment after the defaults, so callers can set
|
||||
// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
|
||||
// need a short authorization-cache window).
|
||||
func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
|
||||
root, err := repoRoot()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
|
||||
WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
|
||||
}
|
||||
|
||||
for _, ov := range envOverrides {
|
||||
for k, v := range ov {
|
||||
req.Env[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
ContainerRequest: req,
|
||||
Started: true,
|
||||
|
||||
@@ -641,7 +641,7 @@ func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequi
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, 0, err
|
||||
}
|
||||
return peer, &types.NetworkMapComponents{Network: types.TwinNetwork(network)}, nil, nil, 0, nil
|
||||
return peer, &types.NetworkMapComponents{Network: network.Copy()}, nil, nil, 0, nil
|
||||
}
|
||||
|
||||
account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
|
||||
@@ -794,7 +794,7 @@ func (c *Controller) GetValidatedPeerWithMap(ctx context.Context, isRequiresAppr
|
||||
}
|
||||
|
||||
emptyMap := &types.NetworkMap{
|
||||
Network: types.TwinNetwork(network),
|
||||
Network: network.Copy(),
|
||||
}
|
||||
return emptyMap, nil, 0, nil
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -86,12 +86,17 @@ type Manager interface {
|
||||
|
||||
// PolicySelectionInput is the per-request selection envelope. The
|
||||
// proxy populates it from CapturedData (account, user, groups) plus
|
||||
// the provider llm_router resolved.
|
||||
// the provider llm_router resolved and the model it extracted.
|
||||
type PolicySelectionInput struct {
|
||||
AccountID string
|
||||
UserID string
|
||||
GroupIDs []string
|
||||
ProviderID string
|
||||
// Model is the already-normalised upstream model id the proxy extracted
|
||||
// (parser strips Bedrock region/version, Vertex @version), so a
|
||||
// case-insensitive compare suffices. Empty = undetermined → not permitted
|
||||
// (fail closed).
|
||||
Model string
|
||||
}
|
||||
|
||||
// PolicySelectionResult names the policy that "pays" for this request
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
@@ -35,6 +36,10 @@ const (
|
||||
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
|
||||
//nolint:gosec // account deny code label, not a credential
|
||||
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
|
||||
// denyCodeModelBlocked is returned when policies govern the request's
|
||||
// (provider, caller-groups) but none permits the model. Matches the proxy
|
||||
// guardrail's code so both layers surface the same label.
|
||||
denyCodeModelBlocked = "llm_policy.model_blocked"
|
||||
)
|
||||
|
||||
// consumptionCache holds the consumption counters prefetched for one
|
||||
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
|
||||
}
|
||||
candidates := filterApplicablePolicies(policies, in)
|
||||
|
||||
// Model-allowlist gate scoped to the matched policies: keep candidates whose
|
||||
// guardrails permit the model (none enabled = unrestricted), deny when
|
||||
// policies apply but none permits it. Skip the load when none has a guardrail.
|
||||
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
|
||||
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if gErr != nil {
|
||||
return nil, gErr
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
|
||||
if len(permitted) == 0 {
|
||||
return &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
candidates = permitted
|
||||
}
|
||||
|
||||
// Prefetch every consumption counter the ceiling + candidate policies will
|
||||
// read, in a single store round-trip, then score against the cache.
|
||||
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
|
||||
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
|
||||
return out
|
||||
}
|
||||
|
||||
// anyPolicyHasGuardrails reports whether any policy references at least one
|
||||
// guardrail, so the selector can skip loading guardrails when none do.
|
||||
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
|
||||
for _, p := range policies {
|
||||
if p != nil && len(p.GuardrailIDs) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
|
||||
// model-allowlist gate to resolve each candidate policy's attached guardrails.
|
||||
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
|
||||
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list account guardrails: %w", err)
|
||||
}
|
||||
byID := make(map[string]*types.Guardrail, len(guardrails))
|
||||
for _, g := range guardrails {
|
||||
if g != nil {
|
||||
byID[g.ID] = g
|
||||
}
|
||||
}
|
||||
return byID, nil
|
||||
}
|
||||
|
||||
// filterModelPermittedPolicies returns the subset of policies whose guardrails
|
||||
// permit the model. Order is preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if policyPermitsModel(p, byID, model) {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyPermitsModel reports whether a policy permits the model. No
|
||||
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
|
||||
// otherwise the model must be in the union of its allowlists, so an
|
||||
// empty/undetermined model fails closed.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
wanted := normaliseModelID(model)
|
||||
restricted := false
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
if wanted == "" {
|
||||
continue
|
||||
}
|
||||
for _, allowed := range g.Checks.ModelAllowlist.Models {
|
||||
if normaliseModelID(allowed) == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return !restricted
|
||||
}
|
||||
|
||||
// normaliseModelID lowercases and trims a model identifier so the allowlist
|
||||
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
|
||||
// normaliseModel so both layers agree on what "same model" means.
|
||||
func normaliseModelID(model string) string {
|
||||
return strings.ToLower(strings.TrimSpace(model))
|
||||
}
|
||||
|
||||
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
|
||||
// rejection. The model is quoted when known; an undetermined model is reported
|
||||
// as such so the access log distinguishes "wrong model" from "no model".
|
||||
func modelBlockedReason(model string) string {
|
||||
if normaliseModelID(model) == "" {
|
||||
return "request model could not be determined for the policy allowlist"
|
||||
}
|
||||
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
|
||||
}
|
||||
|
||||
// candidate is the per-policy intermediate the selector ranks. A
|
||||
// policy that's been exhausted on any enabled cap never makes it
|
||||
// into this slice; the selector's deny envelope carries the latest
|
||||
|
||||
@@ -0,0 +1,329 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
)
|
||||
|
||||
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
|
||||
// to reach providerID under the given guardrails. Uncapped keeps the selector's
|
||||
// headroom scoring trivial so these tests isolate the model-allowlist gate.
|
||||
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Enabled: true,
|
||||
SourceGroups: sourceGroups,
|
||||
DestinationProviderIDs: []string{providerID},
|
||||
GuardrailIDs: guardrailIDs,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
}
|
||||
|
||||
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
|
||||
// carries the given models.
|
||||
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
|
||||
return &types.Guardrail{
|
||||
ID: id,
|
||||
AccountID: account,
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
|
||||
Return(policies, nil)
|
||||
}
|
||||
|
||||
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
|
||||
Return(guardrails, nil)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
|
||||
// decision: a policy authorises the (provider, group) but restricts the model,
|
||||
// and the requested model isn't on the list, so the request is denied.
|
||||
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
|
||||
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
|
||||
// is on the applicable policy's allowlist, so selection proceeds normally.
|
||||
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
|
||||
// and surrounding whitespace, matching the proxy guardrail's normalisation.
|
||||
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
|
||||
// policies authorise the same (provider, group) and one has no guardrail, that
|
||||
// policy makes the request unrestricted — not caught by the other's allowlist.
|
||||
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
|
||||
expectPolicies(mockStore, "acc-1", restricted, open)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
|
||||
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
|
||||
// model allowlisted only for grp-b must not be usable by a grp-a caller. The
|
||||
// selector considers only policies applicable to the caller's groups.
|
||||
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
|
||||
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
|
||||
expectPolicies(mockStore, "acc-1", polA, polB)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-a"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only allowed for grp-b
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
|
||||
// mirrors the proxy: with a restricted applicable policy and an empty model
|
||||
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
|
||||
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "", // undetermined
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
|
||||
// model allowlist is disabled imposes no model restriction, even though the
|
||||
// policy references it.
|
||||
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
disabled := &types.Guardrail{
|
||||
ID: "g-1",
|
||||
AccountID: "acc-1",
|
||||
Checks: types.GuardrailChecks{
|
||||
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
|
||||
},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", disabled)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
|
||||
// allowlist guardrails permits the union of their models (not just the first).
|
||||
func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only in the second guardrail's list
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
|
||||
// resolving the candidate policies' guardrails surfaces as an error, not a
|
||||
// silent allow/deny.
|
||||
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
mockStore.EXPECT().
|
||||
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
|
||||
Return(nil, errors.New("store unavailable"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
|
||||
assert.Nil(t, res)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
|
||||
// policy referencing a guardrail ID absent from the account's set (a stale
|
||||
// reference) imposes no model restriction — same as no guardrail.
|
||||
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anything-goes",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
|
||||
// gate narrows candidates before cap scoring: the permitting policy is selected
|
||||
// even though the blocked one has a larger, more attractive cap.
|
||||
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
|
||||
polBig.Limits = types.PolicyLimits{
|
||||
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
|
||||
}
|
||||
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
|
||||
polSmall.Limits = types.PolicyLimits{
|
||||
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", polBig, polSmall)
|
||||
expectGuardrails(mockStore, "acc-1",
|
||||
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "claude-opus-4", // only pol-small's guardrail permits this
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow)
|
||||
assert.Equal(t, "pol-small", res.SelectedPolicyID,
|
||||
"the model filter must exclude pol-big before cap scoring")
|
||||
}
|
||||
@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
|
||||
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
|
||||
applyAccountCollectionControls(&mergedGuardrails, settings)
|
||||
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in this map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
|
||||
|
||||
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
|
||||
// middleware expects. Mirrors the proxy registration documented in
|
||||
// the management→proxy contract.
|
||||
// the management→proxy contract. provider_allowlists is keyed by the
|
||||
// resolved provider id llm_router stamps; a provider absent from the map is
|
||||
// unrestricted at the proxy layer.
|
||||
type guardrailConfig struct {
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
|
||||
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
|
||||
}
|
||||
|
||||
type guardrailPromptCapture struct {
|
||||
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
|
||||
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
|
||||
}
|
||||
|
||||
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
|
||||
cfg := guardrailConfig{
|
||||
ModelAllowlist: merged.ModelAllowlist,
|
||||
PromptCapture: guardrailPromptCapture{
|
||||
Enabled: merged.PromptCapture.Enabled,
|
||||
RedactPii: merged.PromptCapture.RedactPii,
|
||||
},
|
||||
ProviderAllowlists: providerAllowlists,
|
||||
PromptCapture: guardrailPromptCapture(capture),
|
||||
}
|
||||
out, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
|
||||
// is included only when every authorising policy restricts models (their union);
|
||||
// if any leaves it unrestricted it is omitted, so management decides per group.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
|
||||
type providerAcc struct {
|
||||
models map[string]struct{}
|
||||
anyUnrestricted bool
|
||||
}
|
||||
accs := make(map[string]*providerAcc)
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
acc, ok := accs[providerID]
|
||||
if !ok {
|
||||
acc = &providerAcc{models: make(map[string]struct{})}
|
||||
accs[providerID] = acc
|
||||
}
|
||||
if !restricted {
|
||||
acc.anyUnrestricted = true
|
||||
continue
|
||||
}
|
||||
for _, m := range models {
|
||||
acc.models[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
out := make(map[string][]string, len(accs))
|
||||
for providerID, acc := range accs {
|
||||
if acc.anyUnrestricted {
|
||||
continue
|
||||
}
|
||||
models := make([]string, 0, len(acc.models))
|
||||
for m := range acc.models {
|
||||
models = append(models, m)
|
||||
}
|
||||
sort.Strings(models)
|
||||
out[providerID] = models
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// policyModelAllowlist reports whether a policy restricts models (has an
|
||||
// allowlist-enabled guardrail) and the union of allowed models. Models are
|
||||
// verbatim; the proxy factory lowercases/trims them at decode time.
|
||||
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := false
|
||||
var models []string
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
|
||||
continue
|
||||
}
|
||||
restricted = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
models = append(models, m)
|
||||
}
|
||||
}
|
||||
}
|
||||
return restricted, models
|
||||
}
|
||||
|
||||
// buildAccountService composes the per-account gateway Service. The
|
||||
// target carries the noop placeholder URL — the router middleware
|
||||
// rewrites every request to the matched provider's upstream before the
|
||||
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// MergedGuardrails is the JSON shape passed to the proxy via the
|
||||
// guardrail middleware's config_json. Mirrors the proxy-side
|
||||
// expectations and is intentionally distinct from
|
||||
// types.GuardrailChecks so we can evolve either side independently.
|
||||
// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
|
||||
// merged here — the model allowlist is emitted per-provider, and
|
||||
// token/budget/retention moved onto Policy.Limits and account Settings.
|
||||
type MergedGuardrails struct {
|
||||
ModelAllowlist []string `json:"model_allowlist,omitempty"`
|
||||
TokenLimits MergedTokenLimits `json:"token_limits"`
|
||||
Budget MergedBudget `json:"budget"`
|
||||
PromptCapture MergedPromptCapture `json:"prompt_capture"`
|
||||
Retention MergedRetention `json:"retention"`
|
||||
}
|
||||
|
||||
type MergedTokenLimits struct {
|
||||
Hourly *MergedTokenWindow `json:"hourly,omitempty"`
|
||||
Daily *MergedTokenWindow `json:"daily,omitempty"`
|
||||
Monthly *MergedTokenWindow `json:"monthly,omitempty"`
|
||||
}
|
||||
|
||||
type MergedTokenWindow struct {
|
||||
MaxInputTokens int `json:"max_input_tokens,omitempty"`
|
||||
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
|
||||
}
|
||||
|
||||
type MergedBudget struct {
|
||||
Hourly *MergedBudgetWindow `json:"hourly,omitempty"`
|
||||
Daily *MergedBudgetWindow `json:"daily,omitempty"`
|
||||
Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
|
||||
}
|
||||
|
||||
type MergedBudgetWindow struct {
|
||||
SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
|
||||
HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
|
||||
PromptCapture MergedPromptCapture
|
||||
}
|
||||
|
||||
type MergedPromptCapture struct {
|
||||
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
|
||||
RedactPii bool `json:"redact_pii"`
|
||||
}
|
||||
|
||||
type MergedRetention struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Days int `json:"days"`
|
||||
}
|
||||
|
||||
// mergeGuardrails computes the effective guardrail spec applied at the
|
||||
// proxy, given the referencing policies and the account's guardrail
|
||||
// catalogue. Policy enabled-ness is the caller's responsibility — only
|
||||
// enabled policies should be passed in.
|
||||
// mergeGuardrails folds the referencing policies' guardrails into the
|
||||
// prompt-capture decision only. The model allowlist is enforced per-policy/group
|
||||
// in management and shipped per-provider; token/budget/retention live off
|
||||
// guardrails now.
|
||||
//
|
||||
// Merge rules:
|
||||
// - Model allowlist: union of allowlists across policies that enable it.
|
||||
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
|
||||
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
|
||||
// if any enabling policy turns it on.
|
||||
// - Retention: enabled if any enables it; smallest non-zero days wins.
|
||||
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
|
||||
// sticks if any enabling policy turns it on.
|
||||
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
|
||||
merged := MergedGuardrails{}
|
||||
allowlist := make(map[string]struct{})
|
||||
allowlistEnabled := false
|
||||
|
||||
for _, policy := range policies {
|
||||
for _, gID := range policy.GuardrailIDs {
|
||||
g, ok := byID[gID]
|
||||
if !ok || g == nil {
|
||||
continue
|
||||
}
|
||||
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
|
||||
mergeGuardrail(g, &merged)
|
||||
}
|
||||
}
|
||||
|
||||
if allowlistEnabled {
|
||||
merged.ModelAllowlist = make([]string, 0, len(allowlist))
|
||||
for m := range allowlist {
|
||||
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
|
||||
}
|
||||
sort.Strings(merged.ModelAllowlist)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// mergeGuardrail folds a single guardrail's enabled checks into the
|
||||
// running merge: model-allowlist models join the shared set (and flip
|
||||
// allowlistEnabled), and prompt-capture / redact-pii stick once any
|
||||
// enabling guardrail turns them on.
|
||||
//
|
||||
// TokenLimits, Budget, and Retention have moved off guardrails — token
|
||||
// and budget caps now live on the Policy itself (Policy.Limits) and
|
||||
// retention moves to account-level Settings — so they are not merged here.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
|
||||
if g.Checks.ModelAllowlist.Enabled {
|
||||
*allowlistEnabled = true
|
||||
for _, m := range g.Checks.ModelAllowlist.Models {
|
||||
if m != "" {
|
||||
allowlist[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
|
||||
// running merge: enabled / redact-pii stick once any enabling guardrail turns
|
||||
// them on.
|
||||
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
|
||||
if g.Checks.PromptCapture.Enabled {
|
||||
merged.PromptCapture.Enabled = true
|
||||
if g.Checks.PromptCapture.RedactPii {
|
||||
|
||||
@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
|
||||
require.Len(t, services, 1)
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
|
||||
"model allowlist is a pure policy guardrail and must always reach the config")
|
||||
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
|
||||
"model allowlist is a pure policy guardrail and must reach the per-provider config")
|
||||
assert.False(t, cfg.PromptCapture.Enabled,
|
||||
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
|
||||
}
|
||||
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
|
||||
require.Len(t, services, 1, "exactly one synth service expected")
|
||||
|
||||
cfg := decodeServiceGuardrailConfig(t, services[0])
|
||||
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
|
||||
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
|
||||
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
|
||||
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
)
|
||||
|
||||
// policyForProviders builds an enabled policy authorising the given providers
|
||||
// under the given guardrails (both optional). Groups are irrelevant to
|
||||
// buildProviderAllowlists, which keys purely on destination provider.
|
||||
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
|
||||
return &types.Policy{
|
||||
ID: id,
|
||||
Enabled: true,
|
||||
DestinationProviderIDs: providerIDs,
|
||||
GuardrailIDs: guardrailIDs,
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildProviderAllowlists(t *testing.T) {
|
||||
byID := map[string]*types.Guardrail{
|
||||
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
|
||||
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
|
||||
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
|
||||
}
|
||||
|
||||
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
|
||||
"a provider every policy restricts carries the sorted union of their models")
|
||||
})
|
||||
|
||||
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", nil, "prov-x"), // no guardrail
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a policy whose only guardrail has a disabled allowlist is unrestricted")
|
||||
})
|
||||
|
||||
t.Run("providers are isolated from one another", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
|
||||
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
|
||||
})
|
||||
|
||||
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
|
||||
})
|
||||
|
||||
t.Run("union across a single policy's guardrails", func(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
|
||||
"a policy's own multiple allowlist guardrails union together")
|
||||
})
|
||||
|
||||
t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
|
||||
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
|
||||
got := buildProviderAllowlists([]*types.Policy{
|
||||
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
|
||||
}, empty)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
|
||||
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
|
||||
})
|
||||
}
|
||||
@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
|
||||
|
||||
var cfg guardrailConfig
|
||||
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
|
||||
"model allowlist union must keep both models")
|
||||
// Both policies restrict the same provider, so the per-provider backstop
|
||||
// carries the union of their models — a coarse gate that management's
|
||||
// per-policy/group check narrows; it only blocks models outside the union
|
||||
// when management is down.
|
||||
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
|
||||
"per-provider allowlist union must keep both models")
|
||||
}
|
||||
|
||||
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
package networkmapdb
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
type NetworkMapDBStore interface {
|
||||
GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, error) // TODO: join/populate peers
|
||||
GetPeers(ctx context.Context, accountId string) ([]nmdata.Peer, error)
|
||||
GetPolicies(ctx context.Context, accountId string) ([]nmdata.Policy, error)
|
||||
}
|
||||
|
||||
type NetworkMapDBStoreImpl struct {
|
||||
store NetworkMapDBStore
|
||||
}
|
||||
@@ -1,11 +0,0 @@
|
||||
package networkmapdb
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
func (db *NetworkMapDBStoreImpl) GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, error) {
|
||||
return db.store.GetGroups(ctx, accountId)
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
package networkmap_pgsql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
const (
|
||||
GetGroupsQuery = `
|
||||
select name, public_id, resources from groups where account_id=$1
|
||||
`
|
||||
)
|
||||
|
||||
func (pg *PgStore) GetGroups(ctx context.Context, accountId string) ([]nmdata.Group, error) {
|
||||
rows, err := pg.pool.Query(ctx, GetGroupsQuery, accountId)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
group nmdata.Group
|
||||
name, publicId, resources sql.NullString
|
||||
)
|
||||
|
||||
groups, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nmdata.Group, error) {
|
||||
err := row.Scan(&name, &publicId, &resources)
|
||||
if err != nil {
|
||||
return group, err
|
||||
}
|
||||
|
||||
if name.Valid {
|
||||
group.Name = name.String
|
||||
}
|
||||
if publicId.Valid {
|
||||
group.PublicID = publicId.String
|
||||
}
|
||||
if resources.Valid {
|
||||
groupResources := make([]nmdata.Resource, 0)
|
||||
err := json.Unmarshal([]byte(resources.String), &groupResources)
|
||||
if err != nil {
|
||||
return group, err
|
||||
}
|
||||
group.Resources = groupResources
|
||||
}
|
||||
return group, nil
|
||||
})
|
||||
|
||||
return groups, err
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
package networkmap_pgsql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
_ "embed"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
//go:embed test_db.sql
|
||||
var initDb string
|
||||
|
||||
func TestXxx(t *testing.T) {
|
||||
ctx := context.TODO()
|
||||
|
||||
s, err := NewPostgresqlStore(ctx, "postgresql://root:netbird@localhost:5432/netbird")
|
||||
assert.NoError(t, err)
|
||||
// err = loadSQL(ctx, s.pool, initDb)
|
||||
//assert.NoError(t, err)
|
||||
|
||||
_, err = s.pool.Query(ctx, "insert into groups (id, account_id, name, resources, public_id) VALUES('test-group-id-1','ck7bnf2t2r9s739pkug0','test-group-1', '[{\"ID\":\"cui7q2jl0ubs73d8qpi0\",\"Type\":\"host\"}]','public-id-1')")
|
||||
assert.NoError(t, err)
|
||||
|
||||
groups, err := s.GetGroups(ctx, "ck7bnf2t2r9s739pkug0")
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t,
|
||||
groups,
|
||||
nmdata.Group{Name: "test-group-1", PublicID: "public-id-1", Resources: []nmdata.Resource{{ID: "cui7q2jl0ubs73d8qpi0", Type: "host"}}},
|
||||
)
|
||||
assert.Contains(t,
|
||||
groups,
|
||||
nmdata.Group{Name: "All", PublicID: "d9aejspvcsu517nkh4a0", Resources: []nmdata.Resource{{ID: "cui7olrl0ubs73d8qpe0", Type: "subnet"}, {ID: "cui7q2jl0ubs73d8qpi0", Type: "host"}}},
|
||||
)
|
||||
}
|
||||
|
||||
func loadSQL(ctx context.Context, pool *pgxpool.Pool, initdb string) error {
|
||||
queries := strings.Split(string(initdb), ";")
|
||||
|
||||
for _, query := range queries {
|
||||
query = strings.TrimSpace(query)
|
||||
if query != "" {
|
||||
_, err := pool.Query(ctx, query)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,132 +0,0 @@
|
||||
package networkmap_pgsql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
const (
|
||||
GetPeersQuery = `
|
||||
select id, key, ssh_key, dns_label, user_id, ssh_enabled, login_expiration_enabled, last_login, ip, ipv6,
|
||||
meta_wt_version, meta_go_os, meta_os_version, meta_kernel_version, meta_network_addresses, meta_files, meta_capabilities, meta_flags,
|
||||
location_country_code, location_city_name, location_connection_ip
|
||||
from peers
|
||||
where account_id = $1
|
||||
`
|
||||
)
|
||||
|
||||
func (pg *PgStore) GetPeers(ctx context.Context, accountId string) ([]nmdata.Peer, error) {
|
||||
rows, err := pg.pool.Query(ctx, GetPeersQuery, accountId)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
key, sshKey, dnsLabel, userId sql.NullString
|
||||
lastLogin sql.NullTime
|
||||
sshEnabled, loginExpirationEnabled sql.NullBool
|
||||
ip, ipv6, locationConnectionIp []byte
|
||||
metaFiles, metaCapabilities, metaFlags, metaNetworkAddresses []byte
|
||||
metaWtVersion, metaGoOS, metaOSVersion, metaKernelVersion sql.NullString
|
||||
locationCountryCode, locationCityName sql.NullString
|
||||
)
|
||||
|
||||
peers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nmdata.Peer, error) {
|
||||
var peer nmdata.Peer
|
||||
err := row.Scan(&peer.ID, &key, &sshKey, &dnsLabel, &userId, &sshEnabled, &loginExpirationEnabled, &lastLogin, &ip, &ipv6,
|
||||
&metaWtVersion, &metaGoOS, &metaOSVersion, &metaKernelVersion, &metaNetworkAddresses, &metaFiles, &metaCapabilities, &metaFlags,
|
||||
&locationCountryCode, &locationCityName, &locationConnectionIp)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
|
||||
if key.Valid {
|
||||
peer.Key = key.String
|
||||
}
|
||||
if sshKey.Valid {
|
||||
peer.SSHKey = sshKey.String
|
||||
}
|
||||
if dnsLabel.Valid {
|
||||
peer.DNSLabel = dnsLabel.String
|
||||
}
|
||||
if userId.Valid {
|
||||
peer.UserID = userId.String
|
||||
}
|
||||
if lastLogin.Valid {
|
||||
peer.LastLogin = &lastLogin.Time
|
||||
}
|
||||
if sshEnabled.Valid {
|
||||
peer.SSHEnabled = sshEnabled.Bool
|
||||
}
|
||||
if loginExpirationEnabled.Valid {
|
||||
peer.LoginExpirationEnabled = loginExpirationEnabled.Bool
|
||||
}
|
||||
if metaWtVersion.Valid {
|
||||
peer.Meta.WtVersion = metaWtVersion.String
|
||||
}
|
||||
if metaGoOS.Valid {
|
||||
peer.Meta.GoOS = metaGoOS.String
|
||||
}
|
||||
if metaOSVersion.Valid {
|
||||
peer.Meta.OSVersion = metaOSVersion.String
|
||||
}
|
||||
if metaKernelVersion.Valid {
|
||||
peer.Meta.KernelVersion = metaKernelVersion.String
|
||||
}
|
||||
if locationCountryCode.Valid {
|
||||
peer.Location.CountryCode = locationCountryCode.String
|
||||
}
|
||||
if locationCityName.Valid {
|
||||
peer.Location.CityName = locationCityName.String
|
||||
}
|
||||
if ip != nil {
|
||||
err := json.Unmarshal(ip, &peer.IP)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if ipv6 != nil {
|
||||
err := json.Unmarshal(ipv6, &peer.IPv6)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if locationConnectionIp != nil {
|
||||
err := json.Unmarshal(locationConnectionIp, &peer.Location.ConnectionIP)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if metaFiles != nil {
|
||||
err := json.Unmarshal(metaFiles, &peer.Meta.Files)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if metaCapabilities != nil {
|
||||
err := json.Unmarshal(metaCapabilities, &peer.Meta.Capabilities)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if metaFlags != nil {
|
||||
err := json.Unmarshal(metaFlags, &peer.Meta.Flags)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
if metaNetworkAddresses != nil {
|
||||
err := json.Unmarshal(metaNetworkAddresses, &peer.Meta.NetworkAddresses)
|
||||
if err != nil {
|
||||
return peer, err
|
||||
}
|
||||
}
|
||||
return peer, nil
|
||||
})
|
||||
|
||||
return peers, err
|
||||
}
|
||||
@@ -1,56 +0,0 @@
|
||||
package networkmap_pgsql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
|
||||
)
|
||||
|
||||
const (
|
||||
pgMaxConnections = 30
|
||||
pgMinConnections = 1
|
||||
pgMaxConnLifetime = 60 * time.Minute
|
||||
pgHealthCheckPeriod = 1 * time.Minute
|
||||
)
|
||||
|
||||
var _ networkmapdb.NetworkMapDBStore = &PgStore{}
|
||||
|
||||
type PgStore struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func NewPostgresqlStore(ctx context.Context, dsn string) (*PgStore, error) {
|
||||
pool, err := connectToPgDb(context.Background(), dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &PgStore{pool: pool}, nil
|
||||
}
|
||||
|
||||
func connectToPgDb(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
|
||||
config, err := pgxpool.ParseConfig(dsn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse database config: %w", err)
|
||||
}
|
||||
|
||||
config.MaxConns = pgMaxConnections
|
||||
config.MinConns = pgMinConnections
|
||||
config.MaxConnLifetime = pgMaxConnLifetime
|
||||
config.HealthCheckPeriod = pgHealthCheckPeriod
|
||||
|
||||
pool, err := pgxpool.NewWithConfig(ctx, config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to create connection pool: %w", err)
|
||||
}
|
||||
|
||||
if err := pool.Ping(ctx); err != nil {
|
||||
pool.Close()
|
||||
return nil, fmt.Errorf("unable to ping database: %w", err)
|
||||
}
|
||||
|
||||
return pool, nil
|
||||
}
|
||||
@@ -1,137 +0,0 @@
|
||||
package networkmap_pgsql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
const (
|
||||
GetPoliciesQuery = `
|
||||
select p.id, p.public_id, p.enabled, p.source_posture_checks, pr.enabled, pr.action, pr.protocol, pr.bidirectional,
|
||||
pr.sources, pr.destinations, pr.source_resource, pr.destination_resource, pr.ports, pr.port_ranges,
|
||||
pr.authorized_groups, pr.authorized_user
|
||||
from policies as p
|
||||
left join policy_rules as pr on p.id = pr.policy_id
|
||||
where account_id=$1
|
||||
`
|
||||
)
|
||||
|
||||
func (pg *PgStore) GetPolicies(ctx context.Context, accountId string) ([]nmdata.Policy, error) {
|
||||
rows, err := pg.pool.Query(ctx, GetPoliciesQuery, accountId)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
publicId, sourcePostureChecks sql.NullString
|
||||
enabled, ruleEnabled, bidirectional sql.NullBool
|
||||
action, protocol, sources, destinations sql.NullString
|
||||
sourceResource, destinationResource, ports, portRanges sql.NullString
|
||||
authorizedGroups, authorizedUser sql.NullString
|
||||
)
|
||||
|
||||
policies, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (nmdata.Policy, error) {
|
||||
var policy nmdata.Policy
|
||||
var policyRule *nmdata.PolicyRule
|
||||
|
||||
pr := func() *nmdata.PolicyRule {
|
||||
if policyRule != nil {
|
||||
return policyRule
|
||||
}
|
||||
|
||||
policyRule = &nmdata.PolicyRule{}
|
||||
return policyRule
|
||||
}
|
||||
|
||||
err := row.Scan(&policy.ID, &publicId, &enabled, &sourcePostureChecks, &ruleEnabled, &action, &protocol, &bidirectional,
|
||||
&sources, &destinations, &sourceResource, &destinationResource, &ports, &portRanges,
|
||||
&authorizedGroups, &authorizedUser)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
|
||||
if publicId.Valid {
|
||||
policy.PublicID = publicId.String
|
||||
}
|
||||
if sourcePostureChecks.Valid {
|
||||
err := json.Unmarshal([]byte(sourcePostureChecks.String), &policy.SourcePostureChecks)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if enabled.Valid {
|
||||
policy.Enabled = enabled.Bool
|
||||
}
|
||||
if ruleEnabled.Valid {
|
||||
pr().Enabled = ruleEnabled.Bool
|
||||
}
|
||||
if action.Valid {
|
||||
pr().Action = action.String
|
||||
}
|
||||
if protocol.Valid {
|
||||
pr().Protocol = protocol.String
|
||||
}
|
||||
if bidirectional.Valid {
|
||||
pr().Bidirectional = bidirectional.Bool
|
||||
}
|
||||
if sources.Valid {
|
||||
err := json.Unmarshal([]byte(sources.String), &pr().Sources)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if destinations.Valid {
|
||||
err := json.Unmarshal([]byte(destinations.String), &pr().Destinations)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if sourceResource.Valid {
|
||||
err := json.Unmarshal([]byte(sourceResource.String), &pr().SourceResource)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if destinationResource.Valid {
|
||||
err := json.Unmarshal([]byte(destinationResource.String), &pr().DestinationResource)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if ports.Valid {
|
||||
err := json.Unmarshal([]byte(ports.String), &pr().Ports)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if portRanges.Valid {
|
||||
err := json.Unmarshal([]byte(portRanges.String), &pr().PortRanges)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if authorizedGroups.Valid {
|
||||
err := json.Unmarshal([]byte(authorizedGroups.String), &pr().AuthorizedGroups)
|
||||
if err != nil {
|
||||
return policy, err
|
||||
}
|
||||
}
|
||||
if authorizedUser.Valid {
|
||||
pr().AuthorizedUser = authorizedUser.String
|
||||
}
|
||||
|
||||
if policyRule != nil {
|
||||
policyRule.ID = policy.ID
|
||||
policyRule.PolicyID = policy.ID
|
||||
policy.Rules = []*nmdata.PolicyRule{policyRule}
|
||||
}
|
||||
|
||||
return policy, nil
|
||||
})
|
||||
|
||||
return policies, err
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
CREATE TABLE `accounts` (`id` text,`created_by` text,`created_at` datetime,`domain` text,`domain_category` text,`is_domain_primary_account` numeric,`network_identifier` text,`network_net` text,`network_dns` text,`network_serial` integer,`dns_settings_disabled_management_groups` text,`settings_peer_login_expiration_enabled` numeric,`settings_peer_login_expiration` integer,`settings_regular_users_view_blocked` numeric,`settings_groups_propagation_enabled` numeric,`settings_jwt_groups_enabled` numeric,`settings_jwt_groups_claim_name` text,`settings_jwt_allow_groups` text,`settings_extra_peer_approval_enabled` numeric,`settings_extra_integrated_validator_groups` text,PRIMARY KEY (`id`));
|
||||
CREATE TABLE `peers` (`id` text,`account_id` text,`key` text,`setup_key` text,`ip` text,`meta_hostname` text,`meta_go_os` text,`meta_kernel` text,`meta_core` text,`meta_platform` text,`meta_os` text,`meta_os_version` text,`meta_wt_version` text,`meta_ui_version` text,`meta_kernel_version` text,`meta_network_addresses` text,`meta_system_serial_number` text,`meta_system_product_name` text,`meta_system_manufacturer` text,`meta_environment` text,`meta_files` text,`name` text,`dns_label` text,`peer_status_last_seen` datetime,`peer_status_connected` numeric,`peer_status_login_expired` numeric,`peer_status_requires_approval` numeric,`user_id` text,`ssh_key` text,`ssh_enabled` numeric,`login_expiration_enabled` numeric,`last_login` datetime,`created_at` datetime,`ephemeral` numeric,`location_connection_ip` text,`location_country_code` text,`location_city_name` text,`location_geo_name_id` integer,PRIMARY KEY (`id`),CONSTRAINT `fk_accounts_peers_g` FOREIGN KEY (`account_id`) REFERENCES `accounts`(`id`));
|
||||
CREATE TABLE `groups` (`id` text,`account_id` text,`name` text,`issued` text,`peers` text,`integration_ref_id` integer,`integration_ref_integration_type` text,PRIMARY KEY (`id`),CONSTRAINT `fk_accounts_groups_g` FOREIGN KEY (`account_id`) REFERENCES `accounts`(`id`));
|
||||
|
||||
INSERT INTO accounts VALUES('bf1c8084-ba50-4ce7-9439-34653001fc3b','edafee4e-63fb-11ec-90d6-0242ac120003','2024-10-02 16:01:38.210000+02:00','test.com','private',1,'af1c8024-ha40-4ce2-9418-34653101fc3c','{"IP":"100.64.0.0","Mask":"//8AAA=="}','',0,'[]',0,86400000000000,0,0,0,'',NULL,NULL,NULL);
|
||||
INSERT INTO "groups" VALUES('cfefqs706sqkneg59g4g','bf1c8084-ba50-4ce7-9439-34653001fc3b','All','api','[]',0,'');
|
||||
INSERT INTO "groups" VALUES('cfefqs706sqkneg59g3g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup1','api','[]',0,'');
|
||||
INSERT INTO "groups" VALUES('cfefqs706sqkneg59g2g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup2','api','[]',0,'');
|
||||
@@ -4,9 +4,10 @@ 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"
|
||||
nmdata "github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
@@ -81,7 +82,6 @@ func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvel
|
||||
enc := newComponentEncoder(c)
|
||||
enc.indexAllPeers()
|
||||
routerIdxs := enc.indexRouterPeers(c.RouterPeers)
|
||||
enc.indexAllNetworkResources()
|
||||
|
||||
// Phase 2: gather every policy that any consumer references (peer-pair
|
||||
// policies + resource-only policies) so encodeResourcePoliciesMap can
|
||||
@@ -103,6 +103,7 @@ func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvel
|
||||
DnsSettings: enc.encodeDNSSettings(c.DNSSettings),
|
||||
DnsDomain: in.DNSDomain,
|
||||
CustomZoneDomain: c.CustomZoneDomain,
|
||||
AgentVersions: enc.agentVersions,
|
||||
Peers: enc.peers,
|
||||
RouterPeerIndexes: routerIdxs,
|
||||
Policies: policies,
|
||||
@@ -127,7 +128,7 @@ func EncodeNetworkMapEnvelope(in ComponentsEnvelopeInput) *proto.NetworkMapEnvel
|
||||
// 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 *nmdata.Network) uint64 {
|
||||
func networkSerial(n *types.Network) uint64 {
|
||||
if n == nil {
|
||||
return 0
|
||||
}
|
||||
@@ -140,15 +141,16 @@ type componentEncoder struct {
|
||||
peerOrder map[string]uint32
|
||||
peers []*proto.PeerCompact
|
||||
|
||||
networkIdToPublicId map[string]string
|
||||
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)),
|
||||
networkIdToPublicId: make(map[string]string),
|
||||
components: c,
|
||||
peerOrder: make(map[string]uint32, len(c.Peers)),
|
||||
peers: make([]*proto.PeerCompact, 0, len(c.Peers)),
|
||||
agentVersionOrder: make(map[string]uint32),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -161,7 +163,7 @@ func (e *componentEncoder) indexAllPeers() {
|
||||
}
|
||||
}
|
||||
|
||||
func (e *componentEncoder) appendPeer(p *nmdata.Peer) uint32 {
|
||||
func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
|
||||
if idx, ok := e.peerOrder[p.ID]; ok {
|
||||
return idx
|
||||
}
|
||||
@@ -175,7 +177,7 @@ func (e *componentEncoder) appendPeer(p *nmdata.Peer) uint32 {
|
||||
// (c.RouterPeers may contain peers not in c.Peers when validation rules drop
|
||||
// them) and returns their wire indexes for the RouterPeerIndexes field. Must
|
||||
// run before any encoder that resolves peer ids via e.peerOrder.
|
||||
func (e *componentEncoder) indexRouterPeers(routers map[string]*nmdata.Peer) []uint32 {
|
||||
func (e *componentEncoder) indexRouterPeers(routers map[string]*types.ComponentPeer) []uint32 {
|
||||
if len(routers) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -189,15 +191,6 @@ func (e *componentEncoder) indexRouterPeers(routers map[string]*nmdata.Peer) []u
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) indexAllNetworkResources() {
|
||||
for _, r := range e.components.NetworkResources {
|
||||
if !r.Enabled {
|
||||
continue
|
||||
}
|
||||
e.networkIdToPublicId[r.ID] = r.PublicID
|
||||
}
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeGroups() []*proto.GroupCompact {
|
||||
if len(e.components.Groups) == 0 {
|
||||
return nil
|
||||
@@ -211,20 +204,10 @@ func (e *componentEncoder) encodeGroups() []*proto.GroupCompact {
|
||||
peerIdxs = append(peerIdxs, idx)
|
||||
}
|
||||
}
|
||||
|
||||
groupCompactResources := func() []*proto.ResourceCompact {
|
||||
var toret []*proto.ResourceCompact
|
||||
for _, r := range g.Resources {
|
||||
toret = append(toret, e.resourceToProto(r))
|
||||
}
|
||||
return toret
|
||||
}
|
||||
|
||||
out = append(out, &proto.GroupCompact{
|
||||
Id: g.PublicID,
|
||||
PeerIndexes: peerIdxs,
|
||||
IsAll: g.IsGroupAll(),
|
||||
Resources: groupCompactResources(),
|
||||
})
|
||||
}
|
||||
return out
|
||||
@@ -234,7 +217,7 @@ func (e *componentEncoder) encodeGroups() []*proto.GroupCompact {
|
||||
// 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 []*nmdata.Policy) []*proto.PolicyCompact {
|
||||
func (e *componentEncoder) encodePolicies(policies []*types.Policy) []*proto.PolicyCompact {
|
||||
if len(policies) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -256,7 +239,7 @@ func (e *componentEncoder) encodePolicies(policies []*nmdata.Policy) []*proto.Po
|
||||
}
|
||||
|
||||
// encodePolicyRule maps a single PolicyRule under pol to a PolicyCompact entry.
|
||||
func (e *componentEncoder) encodePolicyRule(pol *nmdata.Policy, r *nmdata.PolicyRule) *proto.PolicyCompact {
|
||||
func (e *componentEncoder) encodePolicyRule(pol *types.Policy, r *types.PolicyRule) *proto.PolicyCompact {
|
||||
return &proto.PolicyCompact{
|
||||
Id: pol.PublicID,
|
||||
Action: networkmap.GetProtoAction(string(r.Action)),
|
||||
@@ -295,14 +278,14 @@ func (e *componentEncoder) groupPublicXids(src []string) []string {
|
||||
// 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 []*nmdata.Policy, resourcePolicies map[string][]*nmdata.Policy) []*nmdata.Policy {
|
||||
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([]*nmdata.Policy, 0, len(policies))
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
@@ -360,31 +343,18 @@ func (e *componentEncoder) groupPublicXid(groupID string) (string, bool) {
|
||||
// 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 nmdata.Resource) *proto.ResourceCompact {
|
||||
t, ok := proto.ResourceCompactType_value[string(r.Type)]
|
||||
if !ok || t == 0 || r.ID == "" {
|
||||
func (e *componentEncoder) resourceToProto(r types.Resource) *proto.ResourceCompact {
|
||||
if r.ID == "" && r.Type == "" {
|
||||
return nil
|
||||
}
|
||||
if t == int32(proto.ResourceCompactType_peer) {
|
||||
idx, ok := e.peerOrder[r.ID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return &proto.ResourceCompact{
|
||||
Type: proto.ResourceCompactType_peer,
|
||||
ResourceId: &proto.ResourceCompact_PeerIndex{PeerIndex: idx},
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
publicID, ok := e.networkIdToPublicId[r.ID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
return &proto.ResourceCompact{
|
||||
Type: proto.ResourceCompactType(t),
|
||||
ResourceId: &proto.ResourceCompact_Id{Id: publicID},
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// postureCheckSeqs translates a slice of posture-check xids to their
|
||||
@@ -417,7 +387,7 @@ func (e *componentEncoder) networkPublicId(xid string) (string, bool) {
|
||||
return id, true
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeDNSSettings(s *nmdata.DNSSettings) *proto.DNSSettingsCompact {
|
||||
func (e *componentEncoder) encodeDNSSettings(s *types.DNSSettings) *proto.DNSSettingsCompact {
|
||||
if s == nil || len(s.DisabledManagementGroups) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -432,7 +402,7 @@ func (e *componentEncoder) encodeDNSSettings(s *nmdata.DNSSettings) *proto.DNSSe
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeRoutes(routes []*nmdata.Route) []*proto.RouteRaw {
|
||||
func (e *componentEncoder) encodeRoutes(routes []*nbroute.Route) []*proto.RouteRaw {
|
||||
if len(routes) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -470,7 +440,7 @@ func (e *componentEncoder) encodeRoutes(routes []*nmdata.Route) []*proto.RouteRa
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeNameServerGroups(nsgs []*nmdata.NameServerGroup) []*proto.NameServerGroupRaw {
|
||||
func (e *componentEncoder) encodeNameServerGroups(nsgs []*nbdns.NameServerGroup) []*proto.NameServerGroupRaw {
|
||||
if len(nsgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -493,7 +463,7 @@ func (e *componentEncoder) encodeNameServerGroups(nsgs []*nmdata.NameServerGroup
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeNameServers(servers []nmdata.NameServer) []*proto.NameServer {
|
||||
func encodeNameServers(servers []nbdns.NameServer) []*proto.NameServer {
|
||||
if len(servers) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -508,7 +478,7 @@ func encodeNameServers(servers []nmdata.NameServer) []*proto.NameServer {
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeSimpleRecords(records []nmdata.SimpleRecord) []*proto.SimpleRecord {
|
||||
func encodeSimpleRecords(records []nbdns.SimpleRecord) []*proto.SimpleRecord {
|
||||
if len(records) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -525,7 +495,7 @@ func encodeSimpleRecords(records []nmdata.SimpleRecord) []*proto.SimpleRecord {
|
||||
return out
|
||||
}
|
||||
|
||||
func encodeCustomZones(zones []nmdata.CustomZone) []*proto.CustomZone {
|
||||
func encodeCustomZones(zones []nbdns.CustomZone) []*proto.CustomZone {
|
||||
if len(zones) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -541,7 +511,7 @@ func encodeCustomZones(zones []nmdata.CustomZone) []*proto.CustomZone {
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeNetworkResources(resources []*nmdata.NetworkResource) []*proto.NetworkResourceRaw {
|
||||
func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentResource) []*proto.NetworkResourceRaw {
|
||||
if len(resources) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -570,7 +540,7 @@ func (e *componentEncoder) encodeNetworkResources(resources []*nmdata.NetworkRes
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*nmdata.NetworkRouter) map[string]*proto.NetworkRouterList {
|
||||
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*types.ComponentRouter) map[string]*proto.NetworkRouterList {
|
||||
if len(routersMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -606,7 +576,7 @@ func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*nm
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *componentEncoder) encodeResourcePoliciesMap(rpm map[string][]*nmdata.Policy) map[string]*proto.PolicyIds {
|
||||
func (e *componentEncoder) encodeResourcePoliciesMap(rpm map[string][]*types.Policy) map[string]*proto.PolicyIds {
|
||||
if len(rpm) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -693,7 +663,7 @@ func (e *componentEncoder) encodePostureFailedPeers(m map[string]map[string]stru
|
||||
// (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 *nmdata.AccountSettingsInfo) *proto.AccountSettingsCompact {
|
||||
func toAccountSettingsCompact(s *types.AccountSettingsInfo) *proto.AccountSettingsCompact {
|
||||
if s == nil {
|
||||
return &proto.AccountSettingsCompact{}
|
||||
}
|
||||
@@ -703,7 +673,7 @@ func toAccountSettingsCompact(s *nmdata.AccountSettingsInfo) *proto.AccountSetti
|
||||
}
|
||||
}
|
||||
|
||||
func toAccountNetwork(n *nmdata.Network) *proto.AccountNetwork {
|
||||
func toAccountNetwork(n *types.Network) *proto.AccountNetwork {
|
||||
if n == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -719,20 +689,20 @@ func toAccountNetwork(n *nmdata.Network) *proto.AccountNetwork {
|
||||
return out
|
||||
}
|
||||
|
||||
func toPeerCompact(p *nmdata.Peer) *proto.PeerCompact {
|
||||
func toPeerCompact(p *types.ComponentPeer) *proto.PeerCompact {
|
||||
pc := &proto.PeerCompact{
|
||||
WgPubKey: decodeWgKey(p.Key),
|
||||
SshPubKey: []byte(p.SSHKey),
|
||||
DnsLabel: p.DNSLabel,
|
||||
AgentVersion: p.Meta.WtVersion,
|
||||
AddedWithSsoLogin: p.UserID != "",
|
||||
AgentVersion: p.AgentVersion,
|
||||
AddedWithSsoLogin: p.AddedWithSSOLogin,
|
||||
LoginExpirationEnabled: p.LoginExpirationEnabled,
|
||||
SshEnabled: p.SSHEnabled,
|
||||
SupportsIpv6: p.SupportsIPv6(),
|
||||
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
|
||||
ServerSshAllowed: p.Meta.Flags.ServerSSHAllowed,
|
||||
SupportsIpv6: p.SupportsIPv6,
|
||||
SupportsSourcePrefixes: p.SupportsSourcePrefixes,
|
||||
ServerSshAllowed: p.ServerSSHAllowed,
|
||||
}
|
||||
if p.LastLogin != nil {
|
||||
if !p.LastLogin.IsZero() {
|
||||
pc.LastLoginUnixNano = p.LastLogin.UnixNano()
|
||||
}
|
||||
switch {
|
||||
@@ -781,7 +751,7 @@ func portsToUint32(ports []string) []uint32 {
|
||||
return out
|
||||
}
|
||||
|
||||
func portRangesToProto(ranges []nmdata.RulePortRange) []*proto.PortInfo_Range {
|
||||
func portRangesToProto(ranges []types.RulePortRange) []*proto.PortInfo_Range {
|
||||
if len(ranges) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
|
||||
nbdns "github.com/netbirdio/netbird/dns"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
@@ -152,66 +152,66 @@ func envelopesEquivalent(a, b *proto.NetworkMapEnvelope) bool {
|
||||
}
|
||||
|
||||
func newTestComponents() *types.NetworkMapComponents {
|
||||
peerA := &nmdata.Peer{
|
||||
ID: "peer-a",
|
||||
Key: testWgKeyA,
|
||||
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
|
||||
DNSLabel: "peera",
|
||||
SSHKey: "ssh-a",
|
||||
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
|
||||
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 := &nmdata.Peer{
|
||||
ID: "peer-b",
|
||||
Key: testWgKeyB,
|
||||
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
|
||||
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
|
||||
DNSLabel: "peerb",
|
||||
Meta: nmdata.PeerSystemMeta{WtVersion: "0.25.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 := &nmdata.Peer{
|
||||
ID: "peer-c",
|
||||
Key: testWgKeyC,
|
||||
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
|
||||
DNSLabel: "peerc",
|
||||
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.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: &nmdata.Network{
|
||||
Network: &types.Network{
|
||||
Identifier: "net-test",
|
||||
Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)},
|
||||
Serial: 7,
|
||||
},
|
||||
AccountSettings: &nmdata.AccountSettingsInfo{
|
||||
AccountSettings: &types.AccountSettingsInfo{
|
||||
PeerLoginExpirationEnabled: true,
|
||||
PeerLoginExpiration: 2 * time.Hour,
|
||||
},
|
||||
Peers: map[string]*nmdata.Peer{
|
||||
Peers: map[string]*types.ComponentPeer{
|
||||
"peer-a": peerA,
|
||||
"peer-b": peerB,
|
||||
"peer-c": peerC,
|
||||
},
|
||||
Groups: map[string]*nmdata.Group{
|
||||
"group-src": {PublicID: "1", Name: "Src", Peers: []string{"peer-a"}},
|
||||
"group-dst": {PublicID: "2", Name: "Dst", Peers: []string{"peer-b", "peer-c"}},
|
||||
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: []*nmdata.Policy{
|
||||
Policies: []*types.Policy{
|
||||
{
|
||||
ID: "pol-1",
|
||||
PublicID: "10",
|
||||
Enabled: true,
|
||||
Rules: []*nmdata.PolicyRule{{
|
||||
ID: "rule-1", Enabled: true, Action: string(types.PolicyTrafficActionAccept),
|
||||
Protocol: string(types.PolicyRuleProtocolTCP), Bidirectional: true,
|
||||
Rules: []*types.PolicyRule{{
|
||||
ID: "rule-1", Enabled: true, Action: types.PolicyTrafficActionAccept,
|
||||
Protocol: types.PolicyRuleProtocolTCP, Bidirectional: true,
|
||||
Ports: []string{"22", "80"},
|
||||
PortRanges: []nmdata.RulePortRange{{Start: 8000, End: 8100}},
|
||||
PortRanges: []types.RulePortRange{{Start: 8000, End: 8100}},
|
||||
Sources: []string{"group-src"},
|
||||
Destinations: []string{"group-dst"},
|
||||
}},
|
||||
},
|
||||
},
|
||||
RouterPeers: map[string]*nmdata.Peer{"peer-c": peerC},
|
||||
RouterPeers: map[string]*types.ComponentPeer{"peer-c": peerC},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -304,31 +304,6 @@ func TestEncodeNetworkMapEnvelope_GroupsByAccountPublicId(t *testing.T) {
|
||||
assert.Len(t, groupByID["2"].PeerIndexes, 2)
|
||||
}
|
||||
|
||||
func TestEncodePolicy(t *testing.T) {
|
||||
encoder := componentEncoder{peerOrder: map[string]uint32{"peerId": uint32(1234)}, networkIdToPublicId: map[string]string{"domain": "publicDomain", "host": "publicHost", "subnet": "publicSubnet"}}
|
||||
assert.Equal(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "peer", ID: "peerId"}),
|
||||
&proto.ResourceCompact{Type: proto.ResourceCompactType_peer, ResourceId: &proto.ResourceCompact_PeerIndex{PeerIndex: uint32(1234)}})
|
||||
// verify invalid peer id results in nil
|
||||
assert.Nil(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "peer", ID: "boom"}))
|
||||
assert.Equal(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "domain", ID: "domain"}),
|
||||
&proto.ResourceCompact{Type: proto.ResourceCompactType_domain, ResourceId: &proto.ResourceCompact_Id{Id: "publicDomain"}})
|
||||
assert.Equal(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "host", ID: "host"}),
|
||||
&proto.ResourceCompact{Type: proto.ResourceCompactType_host, ResourceId: &proto.ResourceCompact_Id{Id: "publicHost"}})
|
||||
assert.Equal(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "subnet", ID: "subnet"}),
|
||||
&proto.ResourceCompact{Type: proto.ResourceCompactType_subnet, ResourceId: &proto.ResourceCompact_Id{Id: "publicSubnet"}})
|
||||
// verify invalid resource type results in nil
|
||||
assert.Nil(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "boom", ID: "boom"}))
|
||||
// verify invalid networkresource id results in nil
|
||||
assert.Nil(t,
|
||||
encoder.resourceToProto(nmdata.Resource{Type: "host", ID: "boom"}))
|
||||
}
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_PolicyExpansion(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
|
||||
@@ -402,12 +377,12 @@ func TestEncodeNetworkMapEnvelope_MalformedWgKey(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
v6Only := &nmdata.Peer{
|
||||
ID: "peer-v6",
|
||||
Key: testWgKeyA,
|
||||
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
|
||||
DNSLabel: "peerv6",
|
||||
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
|
||||
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
|
||||
|
||||
@@ -426,11 +401,11 @@ func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_PeerWithoutIP(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Peers["peer-noip"] = &nmdata.Peer{
|
||||
ID: "peer-noip",
|
||||
Key: testWgKeyA,
|
||||
DNSLabel: "peernoip",
|
||||
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
|
||||
c.Peers["peer-noip"] = &types.ComponentPeer{
|
||||
ID: "peer-noip",
|
||||
Key: testWgKeyA,
|
||||
DNSLabel: "peernoip",
|
||||
AgentVersion: "0.40.0",
|
||||
}
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
@@ -448,7 +423,7 @@ func TestEncodeNetworkMapEnvelope_PeerWithoutIP(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
|
||||
c := &types.NetworkMapComponents{
|
||||
Network: &nmdata.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
|
||||
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})
|
||||
@@ -465,9 +440,9 @@ func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
|
||||
func TestEncodeNetworkMapEnvelope_PeerLoginExpirationFields(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
now := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
|
||||
c.Peers["peer-a"].UserID = "user-1"
|
||||
c.Peers["peer-a"].AddedWithSSOLogin = true
|
||||
c.Peers["peer-a"].LoginExpirationEnabled = true
|
||||
c.Peers["peer-a"].LastLogin = &now
|
||||
c.Peers["peer-a"].LastLogin = now
|
||||
|
||||
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
|
||||
|
||||
@@ -497,7 +472,7 @@ func TestEncodeNetworkMapEnvelope_PeerLoginExpirationFields(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RoutesRoundTrip(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Routes = []*nmdata.Route{
|
||||
c.Routes = []*nbroute.Route{
|
||||
{
|
||||
ID: "route-peer",
|
||||
PublicID: "100",
|
||||
@@ -544,7 +519,7 @@ func TestEncodeNetworkMapEnvelope_RoutesRoundTrip(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_RouteWithMissingPeerLeavesIndexUnset(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.Routes = []*nmdata.Route{{
|
||||
c.Routes = []*nbroute.Route{{
|
||||
ID: "route-x",
|
||||
PublicID: "100",
|
||||
Peer: "peer-not-in-components",
|
||||
@@ -564,21 +539,21 @@ func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing
|
||||
// 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 := &nmdata.Policy{
|
||||
resourceOnlyPolicy := &types.Policy{
|
||||
ID: "pol-resource", PublicID: "99", Enabled: true,
|
||||
Rules: []*nmdata.PolicyRule{{
|
||||
ID: "rule-r", Enabled: true, Action: string(types.PolicyTrafficActionAccept),
|
||||
Protocol: string(types.PolicyRuleProtocolTCP),
|
||||
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][]*nmdata.Policy{
|
||||
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 = []*nmdata.NetworkResource{
|
||||
c.NetworkResources = []*types.ComponentResource{
|
||||
{ID: "resource-x", PublicID: "77", Name: "res-x", Enabled: true},
|
||||
}
|
||||
|
||||
@@ -587,16 +562,27 @@ func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing
|
||||
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 = []*nmdata.NameServerGroup{{
|
||||
c.NameServerGroups = []*nbdns.NameServerGroup{{
|
||||
ID: "nsg-1", PublicID: "50", Name: "Main", Description: "primary",
|
||||
NameServers: []nmdata.NameServer{{
|
||||
IP: netip.MustParseAddr("8.8.8.8"), NSType: int(nbdns.UDPNameServerType), Port: 53,
|
||||
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,
|
||||
@@ -635,11 +621,11 @@ func TestEncodeNetworkMapEnvelope_PostureFailedPeers(t *testing.T) {
|
||||
func TestEncodeNetworkMapEnvelope_RoutersMap(t *testing.T) {
|
||||
c := newTestComponents()
|
||||
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
|
||||
c.RoutersMap = map[string]map[string]*nmdata.NetworkRouter{
|
||||
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
|
||||
"net-1": {
|
||||
"peer-c": {
|
||||
PublicID: "200",
|
||||
Masquerade: true, Metric: 10, Enabled: true,
|
||||
PublicID: "200",
|
||||
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -665,14 +651,14 @@ func TestEncodeNetworkMapEnvelope_RouterPeerNotInComponentsPeers(t *testing.T) {
|
||||
// peer_index reference must still resolve.
|
||||
c := newTestComponents()
|
||||
delete(c.Peers, "peer-c")
|
||||
routerPeer := &nmdata.Peer{
|
||||
routerPeer := &types.ComponentPeer{
|
||||
ID: "peer-c", Key: testWgKeyC, IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
|
||||
DNSLabel: "peerc", Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
|
||||
DNSLabel: "peerc", AgentVersion: "0.40.0",
|
||||
}
|
||||
c.RouterPeers = map[string]*nmdata.Peer{"peer-c": routerPeer}
|
||||
c.RouterPeers = map[string]*types.ComponentPeer{"peer-c": routerPeer}
|
||||
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
|
||||
c.RoutersMap = map[string]map[string]*nmdata.NetworkRouter{
|
||||
"net-1": {"peer-c": {PublicID: "1", Enabled: true}},
|
||||
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()
|
||||
@@ -705,9 +691,9 @@ func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
|
||||
|
||||
func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
|
||||
nm := &types.NetworkMap{
|
||||
Peers: []*nmdata.Peer{{
|
||||
Peers: []*types.ComponentPeer{{
|
||||
ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
|
||||
DNSLabel: "extpeer", Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
|
||||
DNSLabel: "extpeer", AgentVersion: "0.40.0",
|
||||
}},
|
||||
FirewallRules: []*types.FirewallRule{{
|
||||
PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
|
||||
@@ -776,7 +762,7 @@ func TestEncodeNetworkMapEnvelope_NilComponentsGracefulDegrade(t *testing.T) {
|
||||
|
||||
func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
|
||||
c := &types.NetworkMapComponents{
|
||||
Network: &nmdata.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
|
||||
Network: &types.Network{Identifier: "x", Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)}},
|
||||
// AccountSettings deliberately nil
|
||||
}
|
||||
|
||||
@@ -790,6 +776,6 @@ func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
|
||||
func emptyNetworkMapComponents() *types.NetworkMapComponents {
|
||||
return types.EmptyNetworkMapComponents(
|
||||
&types.NetworkMapComponents{
|
||||
PeerID: "peer-id", Peers: map[string]*nmdata.Peer{"peer-id": {}}},
|
||||
PeerID: "peer-id", Peers: map[string]*types.ComponentPeer{"peer-id": {}}},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -12,7 +12,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
nmdata "github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
@@ -51,7 +50,7 @@ func ToComponentSyncResponse(
|
||||
// TODO (dmitri) consider using invariants?
|
||||
//
|
||||
enableSSH := computeSSHEnabledForPeer(components, peer)
|
||||
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH)
|
||||
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH, components.ForceRoutingPeerDNSResolution)
|
||||
|
||||
includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
|
||||
useSourcePrefixes := peer.SupportsSourcePrefixes()
|
||||
@@ -171,25 +170,25 @@ func computeSSHEnabledForPeer(c *types.NetworkMapComponents, peer *nbpeer.Peer)
|
||||
// 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 *nmdata.PolicyRule, peer *nbpeer.Peer) bool {
|
||||
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 == string(types.PolicyRuleProtocolNetbirdSSH) {
|
||||
if rule.Protocol == types.PolicyRuleProtocolNetbirdSSH {
|
||||
return true
|
||||
}
|
||||
return peer.SSHEnabled && nmdata.PolicyRuleImpliesLegacySSH(rule)
|
||||
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 *nmdata.PolicyRule, peerID string) bool {
|
||||
if rule.DestinationResource.Type == string(types.ResourceTypePeer) && rule.DestinationResource.ID != "" {
|
||||
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 {
|
||||
|
||||
@@ -7,7 +7,6 @@ import (
|
||||
|
||||
nbpeer "github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
)
|
||||
|
||||
// TestComputeSSHEnabledForPeer covers both Calculate-mirroring branches:
|
||||
@@ -18,15 +17,16 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
const targetPeerID = "target"
|
||||
const targetGroupID = "g_dst"
|
||||
|
||||
mkComponents := func(rule *nmdata.PolicyRule, sshEnabled bool) (*types.NetworkMapComponents, *nbpeer.Peer) {
|
||||
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]*nmdata.Peer{targetPeerID: types.TwinPeer(peer)},
|
||||
Groups: map[string]*nmdata.Group{targetGroupID: {Name: "dst", Peers: []string{targetPeerID}}},
|
||||
Policies: []*nmdata.Policy{{
|
||||
Peers: map[string]*types.ComponentPeer{targetPeerID: peer.ToComponent()},
|
||||
Groups: map[string]*types.ComponentGroup{targetGroupID: group},
|
||||
Policies: []*types.Policy{{
|
||||
ID: "p",
|
||||
Enabled: true,
|
||||
Rules: []*nmdata.PolicyRule{rule},
|
||||
Rules: []*types.PolicyRule{rule},
|
||||
}},
|
||||
}, peer
|
||||
}
|
||||
@@ -34,14 +34,14 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
peerSSH bool
|
||||
rule nmdata.PolicyRule
|
||||
rule types.PolicyRule
|
||||
wantEnabled bool
|
||||
}{
|
||||
{
|
||||
name: "explicit-netbird-ssh-activates-regardless-of-peer-ssh",
|
||||
peerSSH: false,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
@@ -49,8 +49,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "implicit-tcp-22-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolTCP), Ports: []string{"22"},
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
@@ -58,8 +58,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "implicit-tcp-22-without-peer-ssh-disabled",
|
||||
peerSSH: false,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolTCP), Ports: []string{"22"},
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
@@ -67,8 +67,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "implicit-tcp-22022-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolTCP), Ports: []string{"22022"},
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"22022"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
@@ -76,8 +76,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "implicit-all-protocol-with-peer-ssh",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolALL),
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolALL,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
@@ -85,10 +85,10 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "implicit-port-range-covers-22",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: string(types.PolicyRuleProtocolTCP),
|
||||
PortRanges: []nmdata.RulePortRange{{Start: 20, End: 30}},
|
||||
Protocol: types.PolicyRuleProtocolTCP,
|
||||
PortRanges: []types.RulePortRange{{Start: 20, End: 30}},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: true,
|
||||
@@ -96,8 +96,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "tcp-80-no-ssh",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolTCP), Ports: []string{"80"},
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolTCP, Ports: []string{"80"},
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
@@ -105,8 +105,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "disabled-rule-skipped",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: false, Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
rule: types.PolicyRule{
|
||||
Enabled: false, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{targetGroupID},
|
||||
},
|
||||
wantEnabled: false,
|
||||
@@ -114,8 +114,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "peer-not-in-destinations",
|
||||
peerSSH: true,
|
||||
rule: nmdata.PolicyRule{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{"g_other"}, // target not in this group
|
||||
},
|
||||
wantEnabled: false,
|
||||
@@ -123,21 +123,21 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
{
|
||||
name: "peer-typed-destination-resource-matches",
|
||||
peerSSH: false,
|
||||
rule: nmdata.PolicyRule{
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
DestinationResource: nmdata.Resource{ID: targetPeerID, Type: string(types.ResourceTypePeer)},
|
||||
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: nmdata.PolicyRule{
|
||||
rule: types.PolicyRule{
|
||||
Enabled: true,
|
||||
Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
DestinationResource: nmdata.Resource{ID: targetPeerID, Type: "host"}, // wrong type
|
||||
Destinations: []string{targetGroupID}, // saved by group fallback
|
||||
Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
DestinationResource: types.Resource{ID: targetPeerID, Type: "host"}, // wrong type
|
||||
Destinations: []string{targetGroupID}, // saved by group fallback
|
||||
},
|
||||
wantEnabled: true,
|
||||
},
|
||||
@@ -158,14 +158,14 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
|
||||
func TestComputeSSHEnabledForPeer_TargetMissingFromComponents(t *testing.T) {
|
||||
peer := &nbpeer.Peer{ID: "missing", SSHEnabled: true}
|
||||
c := &types.NetworkMapComponents{
|
||||
Peers: map[string]*nmdata.Peer{}, // target peer NOT present
|
||||
Groups: map[string]*nmdata.Group{
|
||||
"g": {Peers: []string{"missing"}},
|
||||
Peers: map[string]*types.ComponentPeer{}, // target peer NOT present
|
||||
Groups: map[string]*types.ComponentGroup{
|
||||
"g": {ID: "g", Peers: []string{"missing"}},
|
||||
},
|
||||
Policies: []*nmdata.Policy{{
|
||||
Policies: []*types.Policy{{
|
||||
ID: "p", Enabled: true,
|
||||
Rules: []*nmdata.PolicyRule{{
|
||||
Enabled: true, Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
|
||||
Rules: []*types.PolicyRule{{
|
||||
Enabled: true, Protocol: types.PolicyRuleProtocolNetbirdSSH,
|
||||
Destinations: []string{"g"},
|
||||
}},
|
||||
}},
|
||||
|
||||
@@ -22,7 +22,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
"github.com/netbirdio/netbird/shared/netiputil"
|
||||
)
|
||||
@@ -120,7 +119,7 @@ func toNetbirdConfig(config *nbconfig.Config, turnCredentials *Token, relayToken
|
||||
return nbConfig
|
||||
}
|
||||
|
||||
func toPeerConfig(peer *nbpeer.Peer, network *nmdata.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool) *proto.PeerConfig {
|
||||
func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool, forceRoutingPeerDNS bool) *proto.PeerConfig {
|
||||
netmask, _ := network.Net.Mask.Size()
|
||||
fqdn := peer.FQDN(dnsName)
|
||||
|
||||
@@ -136,7 +135,7 @@ func toPeerConfig(peer *nbpeer.Peer, network *nmdata.Network, dnsName string, se
|
||||
Address: fmt.Sprintf("%s/%d", peer.IP.String(), netmask),
|
||||
SshConfig: sshConfig,
|
||||
Fqdn: fqdn,
|
||||
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled,
|
||||
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled || peer.ProxyMeta.Embedded || forceRoutingPeerDNS,
|
||||
LazyConnectionEnabled: settings.LazyConnectionEnabled,
|
||||
AutoUpdate: &proto.AutoUpdateSettings{
|
||||
Version: settings.AutoUpdateVersion,
|
||||
@@ -163,12 +162,12 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
|
||||
useSourcePrefixes := peer.SupportsSourcePrefixes()
|
||||
|
||||
response := &proto.SyncResponse{
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
|
||||
NetworkMap: &proto.NetworkMap{
|
||||
Serial: networkMap.Network.CurrentSerial(),
|
||||
Routes: networkmap.ToProtocolRoutes(networkMap.Routes),
|
||||
DNSConfig: networkmap.ToProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
|
||||
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
|
||||
},
|
||||
Checks: toProtocolChecks(ctx, checks),
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package grpc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"reflect"
|
||||
"testing"
|
||||
@@ -14,6 +15,7 @@ import (
|
||||
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
|
||||
"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/types"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap"
|
||||
)
|
||||
@@ -301,3 +303,35 @@ func TestToNetbirdConfig_RelayInvariant(t *testing.T) {
|
||||
assert.True(t, nbCfg.Metrics.Enabled, "metrics flag should carry the settings value")
|
||||
})
|
||||
}
|
||||
|
||||
func TestToPeerConfig_RoutingPeerDNSResolution(t *testing.T) {
|
||||
network := &types.Network{Net: net.IPNet{IP: net.IPv4(100, 0, 0, 0), Mask: net.CIDRMask(8, 32)}}
|
||||
|
||||
newPeer := func(embedded bool) *nbpeer.Peer {
|
||||
p := &nbpeer.Peer{IP: netip.MustParseAddr("100.0.0.1")}
|
||||
p.ProxyMeta.Embedded = embedded
|
||||
return p
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
globalFlag bool
|
||||
embedded bool
|
||||
forceParam bool
|
||||
wantEnabled bool
|
||||
}{
|
||||
{name: "global off, regular peer, no force", wantEnabled: false},
|
||||
{name: "global on wins", globalFlag: true, wantEnabled: true},
|
||||
{name: "embedded proxy peer forced", embedded: true, wantEnabled: true},
|
||||
{name: "routing peer forced via param", forceParam: true, wantEnabled: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
settings := &types.Settings{RoutingPeerDNSResolutionEnabled: tt.globalFlag}
|
||||
cfg := toPeerConfig(newPeer(tt.embedded), network, "netbird.selfhosted", settings, nil, nil, false, tt.forceParam)
|
||||
assert.Equal(t, tt.wantEnabled, cfg.RoutingPeerDnsResolutionEnabled,
|
||||
"RoutingPeerDnsResolutionEnabled should reflect global || embedded || forced")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
|
||||
UserID: req.GetUserId(),
|
||||
GroupIDs: req.GetGroupIds(),
|
||||
ProviderID: req.GetProviderId(),
|
||||
Model: req.GetModel(),
|
||||
})
|
||||
if err != nil {
|
||||
log.WithContext(ctx).Errorf("select policy for request: %v", err)
|
||||
|
||||
138
management/internals/shared/grpc/proxy_llm_policy_limits_test.go
Normal file
138
management/internals/shared/grpc/proxy_llm_policy_limits_test.go
Normal file
@@ -0,0 +1,138 @@
|
||||
package grpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/grpc/codes"
|
||||
grpcstatus "google.golang.org/grpc/status"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
|
||||
// and returns a pre-programmed result, so tests can assert what the handler
|
||||
// forwards to the selector.
|
||||
type fakeAgentNetworkLimits struct {
|
||||
gotInput agentnetwork.PolicySelectionInput
|
||||
result *agentnetwork.PolicySelectionResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
|
||||
f.gotInput = in
|
||||
if f.err != nil {
|
||||
return nil, f.err
|
||||
}
|
||||
return f.result, nil
|
||||
}
|
||||
|
||||
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
|
||||
// the model the proxy extracted must reach the selector's Model unchanged,
|
||||
// alongside the account/user/group/provider fields.
|
||||
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
req := &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
UserId: "user-1",
|
||||
GroupIds: []string{"grp-a", "grp-b"},
|
||||
ProviderId: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
}
|
||||
|
||||
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
|
||||
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
|
||||
assert.Equal(t, "user-1", fake.gotInput.UserID)
|
||||
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
|
||||
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
|
||||
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
|
||||
"the request's model must be forwarded to the selector")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
|
||||
// model (empty string) is forwarded as-is; the selector decides how to treat it.
|
||||
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
req := &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
}
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
|
||||
// envelope surfaces the model-allowlist deny code + reason through the response.
|
||||
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: "llm_policy.model_blocked",
|
||||
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
|
||||
}}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
Model: "claude-opus-4",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, resp)
|
||||
assert.Equal(t, "deny", resp.Decision)
|
||||
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
|
||||
assert.NotEmpty(t, resp.DenyReason)
|
||||
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
|
||||
// failure surfaces as an Internal gRPC error rather than a silent allow.
|
||||
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
|
||||
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
|
||||
s := &ProxyServiceServer{}
|
||||
s.SetAgentNetworkLimitsService(fake)
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.Error(t, err)
|
||||
st, ok := grpcstatus.FromError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
|
||||
}
|
||||
|
||||
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
|
||||
// fallback: with no limits service wired the RPC returns Unimplemented.
|
||||
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
|
||||
s := &ProxyServiceServer{}
|
||||
|
||||
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
|
||||
AccountId: "acc-1",
|
||||
ProviderId: "prov-1",
|
||||
})
|
||||
require.Error(t, err)
|
||||
st, ok := grpcstatus.FromError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, codes.Unimplemented, st.Code())
|
||||
}
|
||||
@@ -921,7 +921,7 @@ func (s *Server) prepareLoginResponse(ctx context.Context, peer *nbpeer.Peer, ne
|
||||
// if peer has reached this point then it has logged in
|
||||
loginResp := &proto.LoginResponse{
|
||||
NetbirdConfig: toNetbirdConfig(s.config, nil, relayToken, nil, settings),
|
||||
PeerConfig: toPeerConfig(peer, types.TwinNetwork(network), s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH),
|
||||
PeerConfig: toPeerConfig(peer, network, s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH, false),
|
||||
Checks: toProtocolChecks(ctx, postureChecks),
|
||||
}
|
||||
|
||||
|
||||
@@ -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(account.Peers, netMap, dnsDomain))
|
||||
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, account.Peers, dnsDomain))
|
||||
}
|
||||
|
||||
func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -534,22 +534,20 @@ func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request)
|
||||
util.WriteJSONObject(r.Context(), w, resp)
|
||||
}
|
||||
|
||||
// toAccessiblePeers resolves the twin peers in netMap back to the full account
|
||||
// peers (by ID) so the API response keeps Status/Name/OS/GeoNameID, which the
|
||||
// slim netmap twins intentionally don't carry.
|
||||
func toAccessiblePeers(accountPeers map[string]*nbpeer.Peer, 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))
|
||||
appendByID := func(id string) {
|
||||
if p, ok := accountPeers[id]; ok && p != nil {
|
||||
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))
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, p := range netMap.Peers {
|
||||
appendByID(p.ID)
|
||||
}
|
||||
for _, p := range netMap.OfflinePeers {
|
||||
appendByID(p.ID)
|
||||
}
|
||||
add(netMap.Peers)
|
||||
add(netMap.OfflinePeers)
|
||||
|
||||
return accessiblePeers
|
||||
}
|
||||
|
||||
@@ -683,3 +683,81 @@ func BackfillPublicIDs[T any](ctx context.Context, db *gorm.DB) error {
|
||||
log.WithContext(ctx).Infof("Backfill of empty public_id in table %s completed", tableName)
|
||||
return nil
|
||||
}
|
||||
|
||||
// FoldCostAggregatesIntoBuckets migrates a per-request cost table from the old
|
||||
// "stored aggregate" shape (cost_usd + cache_cost_usd columns) to the per-bucket
|
||||
// breakdown, where the total and cache portion are derived on read instead.
|
||||
//
|
||||
// The fold preserves both aggregates exactly for historical rows: the cache
|
||||
// total moves into cached_input_cost_usd and the remainder into
|
||||
// input_cost_usd, so a row's derived total and cache cost still match what it
|
||||
// reported before the upgrade. The finer split is genuinely unknown for those
|
||||
// rows — the old schema never recorded a read/write or input/output division —
|
||||
// so it is lumped rather than guessed; only rows written after the upgrade
|
||||
// carry a true four-way split.
|
||||
//
|
||||
// Dropping the columns before folding would zero every historical row's cost,
|
||||
// so the update runs first and the drop only happens once it succeeds. A table
|
||||
// with no cost_usd column has already been migrated (or was created fresh) and
|
||||
// is skipped.
|
||||
func FoldCostAggregatesIntoBuckets[T any](ctx context.Context, db *gorm.DB) error {
|
||||
var model T
|
||||
|
||||
if !db.Migrator().HasTable(&model) {
|
||||
log.WithContext(ctx).Debugf("table for %T does not exist, no cost-bucket migration needed", model)
|
||||
return nil
|
||||
}
|
||||
if !db.Migrator().HasColumn(&model, "cost_usd") {
|
||||
log.WithContext(ctx).Debugf("table for %T has no cost_usd column, cost buckets already migrated", model)
|
||||
return nil
|
||||
}
|
||||
|
||||
stmt := &gorm.Statement{DB: db}
|
||||
if err := stmt.Parse(&model); err != nil {
|
||||
return fmt.Errorf("parse model schema: %w", err)
|
||||
}
|
||||
tableName := stmt.Schema.Table
|
||||
|
||||
// COALESCE guards rows whose new columns were added as NULL by an earlier
|
||||
// AutoMigrate run that predates the NOT NULL default.
|
||||
hasCacheColumn := db.Migrator().HasColumn(&model, "cache_cost_usd")
|
||||
cacheExpr := "0"
|
||||
if hasCacheColumn {
|
||||
cacheExpr = "COALESCE(cache_cost_usd, 0)"
|
||||
}
|
||||
|
||||
if err := db.Transaction(func(tx *gorm.DB) error {
|
||||
// Only touch rows that carry a legacy total and no breakdown yet, so
|
||||
// the migration is idempotent and never overwrites a true split.
|
||||
update := fmt.Sprintf(`UPDATE %s
|
||||
SET input_cost_usd = COALESCE(cost_usd, 0) - %s,
|
||||
cached_input_cost_usd = %s,
|
||||
cache_creation_cost_usd = 0,
|
||||
output_cost_usd = 0
|
||||
WHERE COALESCE(cost_usd, 0) <> 0
|
||||
AND COALESCE(input_cost_usd, 0) = 0
|
||||
AND COALESCE(cached_input_cost_usd, 0) = 0
|
||||
AND COALESCE(cache_creation_cost_usd, 0) = 0
|
||||
AND COALESCE(output_cost_usd, 0) = 0`, tableName, cacheExpr, cacheExpr)
|
||||
res := tx.Exec(update)
|
||||
if res.Error != nil {
|
||||
return fmt.Errorf("fold legacy cost aggregates in %s: %w", tableName, res.Error)
|
||||
}
|
||||
log.WithContext(ctx).Infof("folded legacy cost aggregates into per-bucket columns for %d rows in table %s", res.RowsAffected, tableName)
|
||||
|
||||
if err := tx.Migrator().DropColumn(&model, "cost_usd"); err != nil {
|
||||
return fmt.Errorf("drop cost_usd from %s: %w", tableName, err)
|
||||
}
|
||||
if hasCacheColumn {
|
||||
if err := tx.Migrator().DropColumn(&model, "cache_cost_usd"); err != nil {
|
||||
return fmt.Errorf("drop cache_cost_usd from %s: %w", tableName, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.WithContext(ctx).Infof("migration of stored cost aggregates to per-bucket columns in table %s completed", tableName)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
nbDomain "github.com/netbirdio/netbird/shared/management/domain"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
|
||||
)
|
||||
|
||||
type NetworkResourceType string
|
||||
@@ -64,6 +65,27 @@ func NewNetworkResource(accountID, networkID, name, description, address string,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ToComponent converts the resource to its self-contained components
|
||||
// representation. Returns nil for a nil resource.
|
||||
func (n *NetworkResource) ToComponent() *sharedTypes.ComponentResource {
|
||||
if n == nil {
|
||||
return nil
|
||||
}
|
||||
return &sharedTypes.ComponentResource{
|
||||
ID: n.ID,
|
||||
PublicID: n.PublicID,
|
||||
NetworkID: n.NetworkID,
|
||||
AccountID: n.AccountID,
|
||||
Name: n.Name,
|
||||
Description: n.Description,
|
||||
Type: sharedTypes.ComponentResourceType(n.Type),
|
||||
Address: n.Address,
|
||||
Domain: n.Domain,
|
||||
Prefix: n.Prefix,
|
||||
Enabled: n.Enabled,
|
||||
}
|
||||
}
|
||||
|
||||
func (n *NetworkResource) ToAPIResponse(groups []api.GroupMinimum) *api.NetworkResource {
|
||||
addr := n.Prefix.String()
|
||||
if n.Type == Domain {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/networks/types"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
|
||||
)
|
||||
|
||||
type NetworkRouter struct {
|
||||
@@ -21,6 +22,36 @@ type NetworkRouter struct {
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
// ToComponent converts the router to its self-contained components
|
||||
// representation. Returns nil for a nil router.
|
||||
func (n *NetworkRouter) ToComponent() *sharedTypes.ComponentRouter {
|
||||
if n == nil {
|
||||
return nil
|
||||
}
|
||||
return &sharedTypes.ComponentRouter{
|
||||
NetworkID: n.NetworkID,
|
||||
PublicID: n.PublicID,
|
||||
Peer: n.Peer,
|
||||
PeerGroups: n.PeerGroups,
|
||||
Masquerade: n.Masquerade,
|
||||
Metric: n.Metric,
|
||||
Enabled: n.Enabled,
|
||||
}
|
||||
}
|
||||
|
||||
// ToComponentMap converts a peer-keyed router map to its components
|
||||
// representation.
|
||||
func ToComponentMap(routers map[string]*NetworkRouter) map[string]*sharedTypes.ComponentRouter {
|
||||
if routers == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*sharedTypes.ComponentRouter, len(routers))
|
||||
for id, r := range routers {
|
||||
out[id] = r.ToComponent()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func NewNetworkRouter(accountID string, networkID string, peer string, peerGroups []string, masquerade bool, metric int, enabled bool) (*NetworkRouter, error) {
|
||||
r := &NetworkRouter{
|
||||
ID: xid.New().String(),
|
||||
|
||||
@@ -21,7 +21,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/permissions/modules"
|
||||
"github.com/netbirdio/netbird/management/server/permissions/operations"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/posture"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
@@ -406,7 +405,7 @@ func (am *DefaultAccountManager) CreatePeerJob(ctx context.Context, accountID, p
|
||||
return status.NewPeerNotPartOfAccountError()
|
||||
}
|
||||
|
||||
meetMinVer, err := posture.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
|
||||
meetMinVer, err := version.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
|
||||
if !version.IsDevelopmentVersion(p.Meta.WtVersion) && (!meetMinVer || err != nil) {
|
||||
return status.Errorf(status.PreconditionFailed, "peer version %s does not meet the minimum required version %s for remote jobs", p.Meta.WtVersion, remoteJobsMinVer)
|
||||
}
|
||||
@@ -1589,7 +1588,7 @@ func affectedPeerIDsFromNetworkMap(nmap *types.NetworkMap, selfPeerID string) []
|
||||
}
|
||||
seen := make(map[string]struct{}, len(nmap.Peers)+len(nmap.OfflinePeers))
|
||||
ids := make([]string, 0, len(nmap.Peers)+len(nmap.OfflinePeers))
|
||||
add := func(peers []*nmdata.Peer) {
|
||||
add := func(peers []*types.ComponentPeer) {
|
||||
for _, p := range peers {
|
||||
if p == nil || p.ID == "" || p.ID == selfPeerID {
|
||||
continue
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/util"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
|
||||
)
|
||||
|
||||
// Peer capability constants mirror the proto enum values.
|
||||
@@ -205,6 +206,35 @@ func (p *Peer) AddedWithSSOLogin() bool {
|
||||
return p.UserID != ""
|
||||
}
|
||||
|
||||
// ToComponent converts the peer to its self-contained components
|
||||
// representation, carrying exactly the subset of peer data that crosses the
|
||||
// components wire format. Returns nil for a nil peer so callers can convert
|
||||
// possibly-missing peers without guarding.
|
||||
func (p *Peer) ToComponent() *sharedTypes.ComponentPeer {
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
cp := &sharedTypes.ComponentPeer{
|
||||
ID: p.ID,
|
||||
Key: p.Key,
|
||||
IP: p.IP,
|
||||
IPv6: p.IPv6,
|
||||
DNSLabel: p.DNSLabel,
|
||||
SSHKey: p.SSHKey,
|
||||
SSHEnabled: p.SSHEnabled,
|
||||
ServerSSHAllowed: p.Meta.Flags.ServerSSHAllowed,
|
||||
AgentVersion: p.Meta.WtVersion,
|
||||
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
|
||||
SupportsIPv6: p.SupportsIPv6(),
|
||||
LoginExpirationEnabled: p.LoginExpirationEnabled,
|
||||
AddedWithSSOLogin: p.AddedWithSSOLogin(),
|
||||
}
|
||||
if p.LastLogin != nil {
|
||||
cp.LastLogin = *p.LastLogin
|
||||
}
|
||||
return cp
|
||||
}
|
||||
|
||||
// HasCapability reports whether the peer has the given capability.
|
||||
func (p *Peer) HasCapability(capability int32) bool {
|
||||
return slices.Contains(p.Meta.Capabilities, capability)
|
||||
|
||||
@@ -57,7 +57,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
nbroute "github.com/netbirdio/netbird/route"
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
@@ -1092,22 +1091,22 @@ func TestToSyncResponse(t *testing.T) {
|
||||
Signature: "turn-pass",
|
||||
}
|
||||
networkMap := &types.NetworkMap{
|
||||
Network: &nmdata.Network{Net: *ipnet, Serial: 1000},
|
||||
Peers: []*nmdata.Peer{{
|
||||
Network: &types.Network{Net: *ipnet, Serial: 1000},
|
||||
Peers: []*types.ComponentPeer{{
|
||||
IP: netip.MustParseAddr("192.168.1.2"),
|
||||
IPv6: netip.MustParseAddr("fd00::2"),
|
||||
Key: "peer2-key",
|
||||
DNSLabel: "peer2",
|
||||
SSHEnabled: true,
|
||||
SSHKey: "peer2-ssh-key"}},
|
||||
OfflinePeers: []*nmdata.Peer{{
|
||||
OfflinePeers: []*types.ComponentPeer{{
|
||||
IP: netip.MustParseAddr("192.168.1.3"),
|
||||
IPv6: netip.MustParseAddr("fd00::3"),
|
||||
Key: "peer3-key",
|
||||
DNSLabel: "peer3",
|
||||
SSHEnabled: true,
|
||||
SSHKey: "peer3-ssh-key"}},
|
||||
Routes: []*nmdata.Route{
|
||||
Routes: []*nbroute.Route{
|
||||
{
|
||||
ID: "route1",
|
||||
Network: netip.MustParsePrefix("10.0.0.0/24"),
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user