Compare commits

..

42 Commits

Author SHA1 Message Date
riccardom
4e3805f535 Renames SetRemotePort to SetRemoteAddr 2026-07-27 17:09:36 +02:00
riccardom
80c7bb195e pqkem: clock data-path PSK rotation from WireGuard handshakes
Source OnDataPathRekeyed from the WGWatcher's per-handshake callback
(onWGCheckSuccess), which fires only on a fresh handshake, and OnDataPathDown
from the handshake-timeout path. A fresh handshake clocks the next chained
KEM exchange pushed over the data-path UDP transport.
2026-07-27 17:09:36 +02:00
riccardom
cac03bd80c pqkem: register data-path endpoint from signalling
Learn the peer's data-path endpoint from the signalling offer/answer: its WG
overlay IP combined with the advertised pq UDP port (SetRemotePort -> AddPeer).
Registering here is safe before the tunnel is up because sends only ever fire
once it is (clocked by OnDataPathRekeyed). RemovePeer is wired at peer teardown
(engine.removePeer), not on transient disconnect.
2026-07-27 17:09:36 +02:00
riccardom
2681f5f8e6 pqkem: apply derived PSK at WG peer-config time (pull) + keep push for rekey 2026-07-27 17:09:36 +02:00
riccardom
905b7f7914 pqkem: carry KEM offer/answer over the signalling exchange 2026-07-27 17:09:36 +02:00
riccardom
20ae4325ff pqkem: dedicated slog logger via NB_PQ_MLKEM_LOG_LEVEL 2026-07-27 17:09:36 +02:00
riccardom
4189937ac6 Homogeneous logs prefix 2026-07-27 17:09:36 +02:00
riccardom
1bccd71954 Bit of renaming
peer -> peerAddrs
have types for remoteID and localID
t.Close log error
Manager SetTransport -> Start
2026-07-27 17:09:36 +02:00
riccardom
f91e1e34ce Typo 2026-07-27 17:09:36 +02:00
riccardom
d544bfa15e Race fix 2026-07-27 17:09:36 +02:00
riccardom
a9fb4e9f0a Makes Transport just a UDP socket.
Manager owns maps for remoteID <-> remote UDP addr
Engine talks to manager only
2026-07-27 17:09:36 +02:00
riccardom
9074f36761 Adds transport 2026-07-27 17:09:36 +02:00
riccardom
a165eec3ba Communicate the port over the signal exchange 2026-07-27 17:09:36 +02:00
riccardom
4e7cbe2ef8 Ensure iface is up and with overlay ip assigned to get a valid UDP port 2026-07-27 17:09:36 +02:00
riccardom
8157b6d78f Adds real callback setter for PSK on ready 2026-07-27 17:09:36 +02:00
riccardom
b132eff867 Initializes PQ ML-KEM manager 2026-07-27 17:09:35 +02:00
riccardom
e98019bafc Adds no-op Transports and callbacks 2026-07-27 17:09:35 +02:00
riccardom
f4fddce174 Added enabled env var 2026-07-27 17:09:35 +02:00
riccardom
94adc454c1 Adds MLKEM Payload placeholder to client internals 2026-07-27 17:09:35 +02:00
riccardom
80f415519b Invert order of keys as per draft 2026-07-27 17:09:35 +02:00
riccardom
2e7436f49b Protocol update 2026-07-27 17:09:35 +02:00
riccardom
010f5281ce Removes confirm. Uses next offer to deliver confirmation/ack of previous round
We clock the next Offer initiation to the OnDataPathRekeyed, so we have 2 minutes
ahead of us to do our attempts and stuff before to give up.
On failure, we will know because we will not receive a new answer.. but more importantly
the wg handshake will fail :D
2026-07-27 17:09:35 +02:00
riccardom
e88fc05575 Leave signal offer/answer as a pull/push operation not as an actual transport 2026-07-27 17:09:35 +02:00
riccardom
7d2c71f84f Assume two transports: initial "signal" (control plane) one (no data path established yet) + data path one
Define OnDataPathRekeyed event to transition from control plane path to data plane path over the WG tunnel.

Keep confirm ALWAYS on NEW established WG tunnel (posthandshake with rekeying). We keep an active method
irrelevant of the WG handshake (we might decide that the indirect wg handshake is sufficient in the future).

Optimistic commit on responder(when sending answer), while on initiator we set it on getting the answer
2026-07-27 17:09:35 +02:00
riccardom
3e4652e528 Epurate wg refs 2026-07-27 17:09:35 +02:00
riccardom
46d4e4585d Collapse Driver and Manager in one.
- Have just one manager => one lock
 - Session state is needed in driver to => we have it available now.
 - Isomorphically align to rosenpass components and functionality

File	Role	                                  rosenpass equivalent
kem.go	primitive pure X25519MLKEM768	          crypto.go/handshake
message.go	Offer/Answer/Confirm + Encode/Decode  messages.go
manager.go	Manager stateful, single lock	      server logic
callbacks.go	WGCallbackHandler (seam output)	  Handler
Transport (interfaccia)	seam trasporto pluggable  Conn
2026-07-27 17:09:35 +02:00
riccardom
9edf2f4dea [squash] isInitial and answered can be inferred without state variables 2026-07-27 17:09:35 +02:00
riccardom
43d43e2a97 Manages convergence 2026-07-27 17:09:35 +02:00
riccardom
bca463d4c8 Models reattempts 2026-07-27 17:09:35 +02:00
riccardom
3ad12f0e58 Reuse answer, don't calculate again 2026-07-27 17:09:35 +02:00
riccardom
85df3abf0f Adds driver to glue together manager and outside world 2026-07-27 17:09:35 +02:00
riccardom
34f5756117 Defines event callbacks 2026-07-27 17:09:35 +02:00
riccardom
b728542d40 Admits possible errors on Encode 2026-07-27 17:09:35 +02:00
riccardom
b95e1aafe3 Bench key material boilerplate time/allocs
CGO_ENABLED=1 go test ./client/internal/pqkem/ -run '^$' -bench . -benchmem 2>&1 | grep -E "Benchmark|ns/op|PASS|ok" | head -20

BenchmarkX25519Keygen-14    	   33795	     34966 ns/op	     224 B/op	       5 allocs/op
BenchmarkX25519ECDH-14      	   33855	     33973 ns/op	      32 B/op	       1 allocs/op
BenchmarkMLKEMKeygen-14     	   21817	     67778 ns/op	    8200 B/op	       2 allocs/op
BenchmarkMLKEMEncaps-14     	   29918	     43235 ns/op	    1216 B/op	       2 allocs/op
BenchmarkMLKEMDecaps-14     	   26048	     56291 ns/op	      64 B/op	       2 allocs/op
PASS
ok  	github.com/netbirdio/netbird/client/internal/pqkem	9.751s
Shell cwd was reset to /home/riccardo/Desktop/Personal/netbirdio/netbird
2026-07-27 17:09:35 +02:00
riccardom
35ec435658 Pure mechanics of manager 2026-07-27 17:09:35 +02:00
riccardom
49922a8831 Messages definition 2026-07-27 17:09:35 +02:00
riccardom
3e7f52d80b ML-KEM encapsulate/decapsulate module 2026-07-27 17:09:35 +02:00
Misha Bragin
1816a020c4 [management, proxy] Add Claude Opus 5 (#6895) 2026-07-27 16:15:36 +02:00
Zoltan Papp
aa13928b76 [client] Export agent version info for iOS (#6918)
## Describe your changes

Export agent version info for iOS

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [ ] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [x] It is a refactor
- [ ] Created tests that fail without the change (if possible)
- [ ] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature
— **OR** I have discussed it with the NetBird team beforehand (link the
issue / Slack thread in the description). See
[CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first).

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6918"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787750541&installation_model_id=427504&pr_number=6918&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6918&signature=c2de74d89c36b01aac5866422b0be0b7525f7c6e26e46174efc280339eb864b6"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->
2026-07-27 15:53:09 +02:00
Viktor Liu
d681670a9d [misc] Restore the rootless-latest docker tag (#6914) 2026-07-27 12:07:58 +02:00
Maycon Santos
4f6247b5c3 [management, proxy] Add prompt-cache token and cost accounting to agent network usage (#6900)
Co-authored-by: braginini <bangvalo@gmail.com>
2026-07-26 21:42:41 +02:00
Riccardo Manfrin
1e5b0a5c89 [client] Make Test_ConnectPeers deterministic under Docker/eBPF kernel / Darwin CI (#6884)
## Describe your changes

Make `Test_ConnectPeers` deterministic. Two issues, both surfaced once
the
privileged suite moved into a `--privileged` Docker container (#6425):

1. The peers used `getLocalIP()` as their WireGuard endpoint, i.e. the
host's
routable NIC IP (the docker bridge IP `172.17.0.2` in CI). That address
might
not hairpin reliably inside the container, so the handshake
intermittently
timed out (flaky). Use loopback (`127.1.0.x`) instead — always
self-reachable.
2. On a Linux runner with the WG kernel module the iface uses the eBPF
proxy
factory. Its manager is a singleton with one shared XDP program +
settings
map, so bringing up the two ifaces makes the second factory overwrite
the
first's `wg_port`/`proxy_port` and the handshake is dropped. The test is
   incompatible with the eBPF factory, so disable it via
`NB_DISABLE_EBPF_WG_PROXY` (peers then handshake directly over
loopback).
Running the suite across all three modes (eBPF / UDP proxy / ICE bind)
would
   need a larger refactor.

Also fixes a typo in the `ErrSharedSockStopped` message (`socked` →
`socket`).

## Issue ticket number and link

No public issue — CI flakiness follow-up to #6871 on `Test_ConnectPeers`

(https://github.com/netbirdio/netbird/blob/main/client/iface/iface_test.go).

## Stack

<!-- branch-stack -->

### Checklist
- [x] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

Test-only change (plus a log-string typo). No public API, CLI, config,
or
behavior change.

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

N/A

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6884"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787485967&installation_model_id=427504&pr_number=6884&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6884&signature=09fb996715d039bd02775a140e8cc03deca5cb42540790b82a744527264a30ad"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Corrected the “shared socket stopped” error message for clearer
output.
* **Tests**
* Improved peer connection test reliability in privileged CI by
disabling the eBPF WireGuard proxy and using fixed loopback UDP
endpoints for deterministic setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-24 15:23:22 +02:00
127 changed files with 5926 additions and 4195 deletions

View File

@@ -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 }}

View File

@@ -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

View File

@@ -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")
}

View File

@@ -50,6 +50,7 @@ import (
icemaker "github.com/netbirdio/netbird/client/internal/peer/ice"
"github.com/netbirdio/netbird/client/internal/peerstore"
"github.com/netbirdio/netbird/client/internal/portforward"
"github.com/netbirdio/netbird/client/internal/pqkem"
"github.com/netbirdio/netbird/client/internal/profilemanager"
"github.com/netbirdio/netbird/client/internal/relay"
"github.com/netbirdio/netbird/client/internal/rosenpass"
@@ -197,6 +198,10 @@ type Engine struct {
// rpManager is a Rosenpass manager
rpManager *rosenpass.Manager
// pqkemManager runs the ML-KEM post-quantum PSK exchange (gated by NB_ENABLE_PQ_MLKEM).
// It owns the data-path transport and peer endpoint routing.
pqkemManager *pqkem.Manager
// syncMsgMux is used to guarantee sequential Management Service message processing
syncMsgMux *sync.Mutex
@@ -651,6 +656,19 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
e.rpManager.SetInterface(e.wgInterface)
}
// Start the ML-KEM PQ manager after the interface is up so its dedicated UDP
// transport can bind on the WG overlay IP.
if pqkem.Enabled() {
tr, pqErr := newPQTransport(e.config.WgAddr.IP)
if pqErr != nil {
log.Errorf("pqkem: transport bind failed, exchange disabled: %v", pqErr)
} else {
e.pqkemManager = pqkem.NewManager(pqkem.LocalID(publicKey.String()), pqCallbackHandler{wg: e.wgInterface}, pqkem.NewLogger())
e.pqkemManager.Start(tr)
log.Infof("pqkem: enabled (udp port %d on overlay %s)", e.pqkemManager.LocalPort(), e.config.WgAddr.IP)
}
}
// if inbound conns are blocked there is no need to create the ACL manager
if e.firewall != nil && !e.config.BlockInbound {
e.acl = acl.NewDefaultManager(e.firewall)
@@ -914,6 +932,10 @@ func (e *Engine) removePeer(peerKey string) error {
e.connMgr.RemovePeerConn(peerKey)
if e.pqkemManager != nil {
e.pqkemManager.RemovePeer(pqkem.RemoteID(peerKey))
}
err := e.statusRecorder.RemovePeer(peerKey)
if err != nil {
log.Warnf("received error when removing peer %s from status recorder: %v", peerKey, err)
@@ -1900,6 +1922,9 @@ func (e *Engine) createPeerConn(pubKey string, allowedIPs []netip.Prefix, agentV
},
ICEConfig: e.createICEConfig(),
}
if e.pqkemManager != nil {
config.PQ = pqHandshaker{mgr: e.pqkemManager}
}
serviceDependencies := peer.ServiceDependencies{
StatusRecorder: e.statusRecorder,
@@ -2083,6 +2108,10 @@ func (e *Engine) close() {
_ = e.rpManager.Close()
}
if e.pqkemManager != nil {
e.pqkemManager.Stop()
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := e.portForwardManager.GracefullyStop(ctx); err != nil {
@@ -2895,6 +2924,8 @@ func convertToOfferAnswer(msg *sProto.Message) (*peer.OfferAnswer, error) {
Version: msg.GetBody().GetNetBirdVersion(),
RosenpassPubKey: rosenpassPubKey,
RosenpassAddr: rosenpassAddr,
MlkemPayload: msg.GetBody().GetMlkemPayload(),
MlkemPort: int(msg.GetBody().GetMlkemPort()),
RelaySrvAddress: msg.GetBody().GetRelayServerAddress(),
RelaySrvIP: relayIP,
SessionID: sessionID,

View File

@@ -74,6 +74,32 @@ type RosenpassConfig struct {
PermissiveMode bool
}
// PQHandshaker attaches post-quantum ML-KEM material to signalling offers/answers and
// feeds received material back. It is implemented by the engine over the pqkem
// manager and is nil when the PQ exchange is disabled. remoteKey is the peer's
// WireGuard public key.
type PQHandshaker interface {
// OfferPayload returns the KEM offer to embed in an outgoing offer (nil if this
// peer is not the KEM initiator) and the local PQ data-path port to announce.
OfferPayload(remoteKey string) (payload []byte, port int)
// AnswerPayload processes a received KEM offer (nil if absent) and returns the KEM
// answer to embed in the outgoing answer (nil if none) and the local PQ port.
AnswerPayload(remoteKey string, recvOffer []byte) (payload []byte, port int)
// OnAnswer feeds a received KEM answer (nil if absent).
OnAnswer(remoteKey string, recvAnswer []byte)
// PSK returns the peer's latest derived post-quantum PSK to program at WG
// peer-config time (the pull path). ok is false until one has been derived.
PSK(remoteKey string) (wgtypes.Key, bool)
// SetRemoteAddr registers the peer's data-path endpoint learned from signalling:
// its WG overlay IP with the advertised pq UDP port.
SetRemoteAddr(remoteKey string, addr netip.AddrPort)
// OnDataPathRekeyed signals a fresh WireGuard handshake for the peer; it clocks the
// next chained PSK rotation pushed over the data path.
OnDataPathRekeyed(remoteKey string)
// OnDataPathDown signals the peer's tunnel went down.
OnDataPathDown(remoteKey string)
}
// ConnConfig is a peer Connection configuration
type ConnConfig struct {
// Key is a public key of a remote peer
@@ -91,6 +117,9 @@ type ConnConfig struct {
RosenpassConfig RosenpassConfig
// PQ carries post-quantum ML-KEM material on offers/answers; nil when disabled.
PQ PQHandshaker
// ICEConfig ICE protocol configuration
ICEConfig icemaker.Config
}
@@ -681,6 +710,10 @@ func (conn *Conn) onWGDisconnected(watcherCtx context.Context) {
conn.Log.Warnf("WireGuard handshake timeout detected, closing current connection")
if conn.config.PQ != nil {
conn.config.PQ.OnDataPathDown(conn.config.Key)
}
// Close the active connection based on current priority
switch conn.currentConnPriority {
case conntype.Relay:
@@ -946,6 +979,11 @@ func (conn *Conn) onWGCheckSuccess() {
conn.mu.Lock()
conn.wgTimeouts = 0
conn.mu.Unlock()
// A fresh WireGuard handshake is the clock for the post-quantum PSK rotation.
if conn.config.PQ != nil {
conn.config.PQ.OnDataPathRekeyed(conn.config.Key)
}
}
// recordConnectionMetrics records connection stage timestamps as metrics
@@ -987,6 +1025,15 @@ func (conn *Conn) AgentVersionString() string {
}
func (conn *Conn) presharedKey(remoteRosenpassKey []byte) *wgtypes.Key {
// Post-quantum: once the ML-KEM exchange has derived a PSK for this peer, program
// it here so the peer's next WireGuard handshake adopts it. Applied at peer-config
// time (bootstrap / reconnect); steady-state rotation is pushed separately.
if conn.config.PQ != nil {
if psk, ok := conn.config.PQ.PSK(conn.config.Key); ok {
return &psk
}
}
if conn.config.RosenpassConfig.PubKey == nil {
return conn.config.WgConfig.PreSharedKey
}

View File

@@ -39,6 +39,16 @@ type OfferAnswer struct {
// This value is the local Rosenpass server address when sending the message
RosenpassAddr string
// MlkemPayload carries the post-quantum X25519MLKEM768 handshake message
// (pqkem-framed offer on an OFFER, answer on an ANSWER) that seeds the
// WireGuard PSK. Opaque here — the pqkem library frames and parses it. Nil
// when the peer does not run the ML-KEM PQ exchange.
MlkemPayload []byte
// MlkemPort is the peer's ML-KEM PQ service UDP port (bound on its WG overlay
// IP) where data-path rekey messages are sent. Zero when not running the exchange.
MlkemPort int
// relay server address
RelaySrvAddress string
// RelaySrvIP is the IP the remote peer is connected to on its
@@ -120,6 +130,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
h.updateRemoteICEState(&remoteOfferAnswer)
h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
if h.relayListener != nil {
h.relayListener.Notify(&remoteOfferAnswer)
}
@@ -128,7 +140,7 @@ func (h *Handshaker) Listen(ctx context.Context) {
h.iceListener(&remoteOfferAnswer)
}
if err := h.sendAnswer(); err != nil {
if err := h.sendAnswer(&remoteOfferAnswer); err != nil {
h.log.Errorf("failed to send remote offer confirmation: %s", err)
continue
}
@@ -142,6 +154,8 @@ func (h *Handshaker) Listen(ctx context.Context) {
h.updateRemoteICEState(&remoteOfferAnswer)
h.pqRegisterEndpoint(remoteOfferAnswer.MlkemPort)
if h.relayListener != nil {
h.relayListener.Notify(&remoteOfferAnswer)
}
@@ -149,6 +163,10 @@ func (h *Handshaker) Listen(ctx context.Context) {
if h.iceListener != nil && h.RemoteICESupported() {
h.iceListener(&remoteOfferAnswer)
}
if h.config.PQ != nil {
h.config.PQ.OnAnswer(h.config.Key, remoteOfferAnswer.MlkemPayload)
}
case <-ctx.Done():
h.log.Infof("stop listening for remote offers and answers")
return
@@ -156,6 +174,16 @@ func (h *Handshaker) Listen(ctx context.Context) {
}
}
// pqRegisterEndpoint feeds the post-quantum handshaker the peer's data-path endpoint
// (its WG overlay IP plus the advertised pq UDP port) learned from a remote offer/answer.
func (h *Handshaker) pqRegisterEndpoint(remotePort int) {
if h.config.PQ == nil || remotePort <= 0 || remotePort > 65535 || len(h.config.WgConfig.AllowedIps) == 0 {
return
}
addr := netip.AddrPortFrom(h.config.WgConfig.AllowedIps[0].Addr(), uint16(remotePort))
h.config.PQ.SetRemoteAddr(h.config.Key, addr)
}
func (h *Handshaker) SendOffer() error {
h.mu.Lock()
defer h.mu.Unlock()
@@ -195,13 +223,23 @@ func (h *Handshaker) sendOffer() error {
}
offer := h.buildOfferAnswer()
if h.config.PQ != nil {
offer.MlkemPayload, offer.MlkemPort = h.config.PQ.OfferPayload(h.config.Key)
}
h.log.Debugf("sending offer with serial: %s", offer.SessionIDString())
return h.signaler.SignalOffer(offer, h.config.Key)
}
func (h *Handshaker) sendAnswer() error {
func (h *Handshaker) sendAnswer(remoteOffer *OfferAnswer) error {
answer := h.buildOfferAnswer()
if h.config.PQ != nil {
var recvOffer []byte
if remoteOffer != nil {
recvOffer = remoteOffer.MlkemPayload
}
answer.MlkemPayload, answer.MlkemPort = h.config.PQ.AnswerPayload(h.config.Key, recvOffer)
}
h.log.Debugf("sending answer with serial: %s", answer.SessionIDString())
return h.signaler.SignalAnswer(answer, h.config.Key)

View File

@@ -63,6 +63,8 @@ func (s *Signaler) signalOfferAnswer(offerAnswer OfferAnswer, remoteKey string,
},
RosenpassPubKey: offerAnswer.RosenpassPubKey,
RosenpassAddr: offerAnswer.RosenpassAddr,
MlkemPayload: offerAnswer.MlkemPayload,
MlkemPort: offerAnswer.MlkemPort,
RelaySrvAddress: offerAnswer.RelaySrvAddress,
RelaySrvIP: offerAnswer.RelaySrvIP,
SessionID: sessionIDBytes,

View File

@@ -0,0 +1,59 @@
package pqkem
import (
"crypto/ecdh"
"crypto/mlkem"
"crypto/rand"
"testing"
)
func BenchmarkX25519Keygen(b *testing.B) {
c := ecdh.X25519()
b.ResetTimer()
for i := 0; i < b.N; i++ {
if _, err := c.GenerateKey(rand.Reader); err != nil {
b.Fatal(err)
}
}
}
func BenchmarkX25519ECDH(b *testing.B) {
c := ecdh.X25519()
a, _ := c.GenerateKey(rand.Reader)
p, _ := c.GenerateKey(rand.Reader)
pub := p.PublicKey()
b.ResetTimer()
for i := 0; i < b.N; i++ {
if _, err := a.ECDH(pub); err != nil {
b.Fatal(err)
}
}
}
func BenchmarkMLKEMKeygen(b *testing.B) {
for i := 0; i < b.N; i++ {
if _, err := mlkem.GenerateKey768(); err != nil {
b.Fatal(err)
}
}
}
func BenchmarkMLKEMEncaps(b *testing.B) {
dk, _ := mlkem.GenerateKey768()
ek := dk.EncapsulationKey()
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = ek.Encapsulate()
}
}
func BenchmarkMLKEMDecaps(b *testing.B) {
dk, _ := mlkem.GenerateKey768()
_, ct := dk.EncapsulationKey().Encapsulate()
b.ResetTimer()
for i := 0; i < b.N; i++ {
if _, err := dk.Decapsulate(ct); err != nil {
b.Fatal(err)
}
}
}

View File

@@ -0,0 +1,18 @@
package pqkem
// CallbackHandler is implemented by the host and invoked by the library. The
// library only reports events; the host owns the reaction. Keeping this an
// interface — rather than touching the transport or keying directly — is what lets
// the KEM code be extracted as a standalone library.
type CallbackHandler interface {
// OnNewPSKReady fires when a fresh post-quantum PSK has been derived for a peer
// and must be programmed into the consumer's secure channel. It is invoked at
// the commit point of each side: the initiator on receiving the answer, the
// responder on receiving the confirm.
OnNewPSKReady(remoteID RemoteID, psk PSK) error
// OnRekeyFailed fires when an exchange fails to converge within the allotted
// time. The host should tear the peer connection down so it re-establishes, and
// log a WARN. The library reports the event; it does not dictate the reaction.
OnRekeyFailed(remoteID RemoteID) error
}

View File

@@ -0,0 +1,224 @@
package pqkem
import (
"context"
"time"
)
// startExchange creates a fresh initiator exchange (acknowledging ackID, zero for a
// bootstrap) and returns the framed offer for the caller to send — pushed over the
// data path for a chained rekey, or handed to the host for signalling when viaSignal
// is set. Any previous in-flight exchange for the peer is cancelled.
func (m *Manager) startExchange(remoteID RemoteID, viaSignal bool, ackID ExchangeID) ([]byte, error) {
init, err := NewInitiator()
if err != nil {
return nil, err
}
id, err := newExchangeID()
if err != nil {
return nil, err
}
raw, err := (&OfferMsg{ExchangeID: id, AckID: ackID, KEMOffer: init.Offer()}).Encode()
if err != nil {
return nil, err
}
ctx, cancel := context.WithCancel(m.rootCtx)
m.mu.Lock()
if old := m.exchanges[remoteID]; old != nil && old.cancel != nil {
old.cancel()
}
m.exchanges[remoteID] = &exchangeCtl{
id: id,
state: stateAwaitingAnswer,
startedAt: time.Now(),
cancel: cancel,
lastSent: raw,
initiator: init,
viaSignal: viaSignal,
}
m.mu.Unlock()
m.wait.Add(1)
go m.initiatorLoop(ctx, remoteID, id)
return raw, nil
}
// processOffer (responder) first acknowledges the previous exchange the offer names
// (that offer riding the data path under the freshly adopted key proves it worked),
// then derives the PSK for the new offer, commits it optimistically, and returns the
// framed answer. A duplicate offer returns the cached answer without re-deriving.
func (m *Manager) processOffer(remoteID RemoteID, o *OfferMsg) ([]byte, error) {
if o.AckID != (ExchangeID{}) {
m.ackConverged(remoteID, o.AckID)
}
m.mu.Lock()
if ex := m.exchanges[remoteID]; ex != nil && ex.id == o.ExchangeID {
state, last := ex.state, ex.lastSent
m.mu.Unlock()
if state == stateReserved {
return nil, nil
}
return last, nil
}
// Reserve the slot so a concurrent duplicate offer bails.
m.exchanges[remoteID] = &exchangeCtl{id: o.ExchangeID, state: stateReserved, startedAt: time.Now()}
m.mu.Unlock()
answerBytes, psk, err := Respond(o.KEMOffer, m.binding(remoteID))
if err != nil {
return nil, err
}
raw, err := (&AnswerMsg{ExchangeID: o.ExchangeID, KEMAnswer: answerBytes}).Encode()
if err != nil {
return nil, err
}
m.mu.Lock()
ex := m.exchanges[remoteID]
if ex == nil || ex.id != o.ExchangeID {
m.mu.Unlock()
return nil, nil
}
ex.state = stateAwaitingAck
ex.lastSent = raw
ex.pendingPSK = psk
m.psks[remoteID] = psk
m.mu.Unlock()
// Commit optimistically so our data path can rekey to the new PSK.
if err := m.cbHandler.OnNewPSKReady(remoteID, psk); err != nil {
return nil, err
}
return raw, nil
}
// processAnswer (initiator) derives and commits the PSK and parks in
// stateAwaitingRekey; the next offer (chained from OnDataPathRekeyed) will acknowledge
// this exchange. Only valid in stateAwaitingAnswer; advancing the state under the
// lock makes a concurrent/duplicate answer bail.
func (m *Manager) processAnswer(remoteID RemoteID, a *AnswerMsg) error {
m.mu.Lock()
ex := m.exchanges[remoteID]
if ex == nil || ex.id != a.ExchangeID || ex.state != stateAwaitingAnswer {
m.mu.Unlock()
return nil
}
ex.state = stateAwaitingRekey
init := ex.initiator
ex.initiator = nil
m.mu.Unlock()
psk, err := init.Finish(a.KEMAnswer, m.binding(remoteID))
if err != nil {
return err
}
// The initiator has converged: the responder must have derived the key to answer.
m.mu.Lock()
m.established[remoteID] = true
m.failures[remoteID] = 0
m.psks[remoteID] = psk
m.mu.Unlock()
return m.cbHandler.OnNewPSKReady(remoteID, psk)
}
// ackConverged (responder) records convergence of the exchange named by ackID: a
// later offer acknowledging it proves both sides operate on that exchange's key. Only
// acts on a matching stateAwaitingAck exchange; anything else is ignored.
func (m *Manager) ackConverged(remoteID RemoteID, ackID ExchangeID) {
m.mu.Lock()
ex := m.exchanges[remoteID]
if ex == nil || ex.id != ackID || ex.state != stateAwaitingAck {
m.mu.Unlock()
return
}
delete(m.exchanges, remoteID)
m.established[remoteID] = true
m.failures[remoteID] = 0
_ = time.Since(ex.startedAt) // convergence latency (metrics hook, later step)
m.mu.Unlock()
}
// initiatorLoop enforces the offer->answer convergence deadline and retransmits the
// initiator's outstanding data-path offer while awaiting the answer (a
// signalling-bootstrapped offer is retransmitted by the host, so it is not resent
// here). Exhausting the deadline before the answer arrives is a failure. Once the
// answer is in (state past awaitingAnswer) the loop exits: the next rotation is driven
// by OnDataPathRekeyed, and the idle wait for it has no deadline.
func (m *Manager) initiatorLoop(ctx context.Context, remoteID RemoteID, id ExchangeID) {
defer m.wait.Done()
t := time.NewTicker(m.retryInterval)
defer t.Stop()
attempts := 0
for {
select {
case <-ctx.Done():
return
case <-t.C:
m.mu.Lock()
ex := m.exchanges[remoteID]
if ex == nil || ex.id != id {
m.mu.Unlock()
return
}
switch ex.state {
case stateAwaitingAnswer:
if attempts >= m.maxRetries {
delete(m.exchanges, remoteID)
fail := m.registerFailureLocked(remoteID)
m.mu.Unlock()
m.raiseFailure(remoteID, fail)
return
}
viaSignal := ex.viaSignal
msg := ex.lastSent
attempts++
m.mu.Unlock()
if !viaSignal {
if err := m.pushDataPath(remoteID, msg); err != nil {
m.logger.Warn("pqkem: offer retransmit failed", "peer", remoteID, "err", err)
}
}
default:
// Past awaiting the answer (converged) or superseded: the loop's job
// is done. The next rotation is driven externally by OnDataPathRekeyed,
// so there is no deadline while idle-waiting for it (that wait can be
// as long as the transport's natural rekey interval).
m.mu.Unlock()
return
}
}
}
}
// registerFailureLocked applies policy B and reports whether OnRekeyFailed is due:
// an initial exchange (peer never established) fails immediately; a rekey tolerates
// up to maxRekeyFailures consecutive misses (we stay on the still-valid previous
// PSK) before failing. Assumes m.mu is held.
func (m *Manager) registerFailureLocked(remoteID RemoteID) bool {
if !m.established[remoteID] {
return true
}
m.failures[remoteID]++
if m.failures[remoteID] >= m.maxRekeyFailures {
m.failures[remoteID] = 0
return true
}
return false
}
func (m *Manager) raiseFailure(remoteID RemoteID, fail bool) {
if !fail {
m.logger.Warn("pqkem: rekey attempt timed out, will retry next cycle", "peer", remoteID)
return
}
if err := m.cbHandler.OnRekeyFailed(remoteID); err != nil {
m.logger.Error("pqkem: OnRekeyFailed handler error", "peer", remoteID, "err", err)
}
}

View File

@@ -0,0 +1,74 @@
package pqkem
import (
"net/netip"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// dropTransport is a pqkem.Transport that silently discards everything.
type dropTransport struct{}
func (dropTransport) Send(netip.AddrPort, []byte) error { return nil }
func (dropTransport) LocalPort() int { return 0 }
func (dropTransport) Run(func(netip.AddrPort, []byte)) {}
func (dropTransport) Close() error { return nil }
func failedCount(f *fakeWG) int {
f.mu.Lock()
defer f.mu.Unlock()
return len(f.failed)
}
func TestManager_InitialTimeoutFailsImmediately(t *testing.T) {
wg := newFakeWG()
d := NewManager("bbbb", wg, nil) // bbbb > aaaa -> initiator
d.Start(dropTransport{})
d.retryInterval = 5 * time.Millisecond
d.maxRetries = 3
defer d.Stop()
// Bootstrap offer is produced for signalling; no answer ever comes back -> the
// initial exchange fails fast.
offer, err := d.SignalOffer("aaaa")
require.NoError(t, err)
require.NotNil(t, offer)
require.Eventually(t, func() bool { return failedCount(wg) == 1 }, time.Second, 5*time.Millisecond)
}
func TestManager_RekeyToleratesKFailures(t *testing.T) {
dA, dB, _, wgB, lbB := pair(t)
defer dA.Stop()
defer dB.Stop()
// Tighten B's timings before any exchange loop spawns (the loop reads these
// fields, so writing them after a loop is running would race).
dB.retryInterval = 5 * time.Millisecond
dB.maxRetries = 2
// Establish: bootstrap + data-path-rekeyed so B becomes established and its data
// path is usable.
bootstrap(t, dA, dB)
dA.OnDataPathRekeyed("bbbb")
dB.OnDataPathRekeyed("aaaa")
require.NotEqual(t, PSK{}, wgB.psk("aaaa"))
// Drop B's outbound so rekeys can no longer converge.
lbB.drop.Store(true)
// K-1 data-path rekeys must NOT raise OnRekeyFailed.
for i := 0; i < DefaultMaxRekeyFailures-1; i++ {
_, err := dB.startExchange("aaaa", false, ExchangeID{})
require.NoError(t, err)
time.Sleep(50 * time.Millisecond)
}
require.Equal(t, 0, failedCount(wgB), "no failure before K attempts")
// The K-th failure raises it once.
_, err := dB.startExchange("aaaa", false, ExchangeID{})
require.NoError(t, err)
require.Eventually(t, func() bool { return failedCount(wgB) == 1 }, time.Second, 5*time.Millisecond)
}

View File

@@ -0,0 +1,59 @@
package pqkem
import (
"log/slog"
"os"
"strconv"
"strings"
log "github.com/sirupsen/logrus"
)
// EnvEnabled is the environment variable that turns the ML-KEM post-quantum
// exchange on for this client. Accepts on/off aliases plus anything
// strconv.ParseBool understands (true/false/1/0).
const EnvEnabled = "NB_ENABLE_PQ_MLKEM"
// Enabled reports whether the ML-KEM PQ exchange is enabled via the environment.
// An empty or unrecognized value is treated as disabled.
func Enabled() bool {
raw := strings.ToLower(strings.TrimSpace(os.Getenv(EnvEnabled)))
switch raw {
case "":
return false
case "on":
return true
case "off":
return false
}
enabled, err := strconv.ParseBool(raw)
if err != nil {
log.Warnf("failed to parse %s value %q: %v", EnvEnabled, raw, err)
return false
}
return enabled
}
// EnvLogLevel overrides the ML-KEM manager's slog level (debug/info/warn/error).
// Defaults to info.
const EnvLogLevel = "NB_PQ_MLKEM_LOG_LEVEL"
// NewLogger builds the slog logger for the ML-KEM manager: a text handler to stdout
// at the level from EnvLogLevel. Mirrors the Rosenpass manager's logger setup so PQ
// components log consistently.
func NewLogger() *slog.Logger {
return slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel()}))
}
func logLevel() slog.Level {
switch strings.ToLower(strings.TrimSpace(os.Getenv(EnvLogLevel))) {
case "debug":
return slog.LevelDebug
case "warn":
return slog.LevelWarn
case "error":
return slog.LevelError
default:
return slog.LevelInfo
}
}

View File

@@ -0,0 +1,167 @@
// Package pqkem is a spike (NET-1406) for a post-quantum pre-shared-key exchange
// that could replace Rosenpass. It performs an X25519MLKEM768 hybrid key
// encapsulation and derives a 32-byte pre-shared key (PSK).
//
// The exchange is a single round trip designed to ride the (already
// authenticated) Signal offer/answer channel:
//
// initiator --Offer(1216B)--> responder
// initiator <--Answer(1120B)-- responder
//
// Both sides then hold the same PSK, which is bound to the two peers' identities
// (their peer identity keys) so the derived key cannot be transplanted
// to a different peer pair even if the transport authentication were bypassed.
//
// Combiner note: this follows draft-ietf-tls-ecdhe-mlkem for X25519MLKEM768 — on
// the wire ML-KEM ‖ X25519 (the draft deliberately reversed the share order for
// this group), and ML-KEM_ss ‖ X25519_ss fed into the KDF. The spike uses SHA-256
// (also binding the transcript and peer identities); a production version should
// use HKDF — see TODO below.
package pqkem
import (
"crypto/ecdh"
"crypto/mlkem"
"crypto/rand"
"crypto/sha256"
"fmt"
)
const (
// OfferSize is the initiator message: ML-KEM-768 encapsulation key ‖ X25519 public key
// (share order per draft-ietf-tls-ecdhe-mlkem for X25519MLKEM768).
OfferSize = mlkem.EncapsulationKeySize768 + 32 // 1216
// AnswerSize is the responder message: ML-KEM-768 ciphertext ‖ X25519 public key.
AnswerSize = mlkem.CiphertextSize768 + 32 // 1120
pskLabel = "netbird-pq-psk-v1"
)
// PSK is the 32-byte derived pre-shared key handed to the consumer to key its channel.
type PSK [32]byte
// Binding identifies the peer pair the PSK is derived for. Callers set both
// peer identity keys; the order does not matter (it is canonicalised).
type Binding struct {
LocalID []byte
RemoteID []byte
}
// Initiator holds the ephemeral secrets between Offer and Finish.
type Initiator struct {
x25519 *ecdh.PrivateKey
mlkemDK *mlkem.DecapsulationKey768
offer []byte
}
// NewInitiator generates the ephemeral X25519 + ML-KEM-768 keypairs.
func NewInitiator() (*Initiator, error) {
x, err := ecdh.X25519().GenerateKey(rand.Reader)
if err != nil {
return nil, fmt.Errorf("x25519 keygen: %w", err)
}
dk, err := mlkem.GenerateKey768()
if err != nil {
return nil, fmt.Errorf("ml-kem keygen: %w", err)
}
offer := make([]byte, 0, OfferSize)
offer = append(offer, dk.EncapsulationKey().Bytes()...)
offer = append(offer, x.PublicKey().Bytes()...)
return &Initiator{x25519: x, mlkemDK: dk, offer: offer}, nil
}
// Offer returns the initiator message to send over Signal.
func (i *Initiator) Offer() []byte {
return i.offer
}
// Finish consumes the responder's answer and derives the PSK.
func (i *Initiator) Finish(answer []byte, b Binding) (PSK, error) {
if len(answer) != AnswerSize {
return PSK{}, fmt.Errorf("answer: got %d bytes, want %d", len(answer), AnswerSize)
}
ct := answer[:mlkem.CiphertextSize768]
peerX := answer[mlkem.CiphertextSize768:]
ssMLKEM, err := i.mlkemDK.Decapsulate(ct)
if err != nil {
return PSK{}, fmt.Errorf("ml-kem decapsulate: %w", err)
}
pub, err := ecdh.X25519().NewPublicKey(peerX)
if err != nil {
return PSK{}, fmt.Errorf("parse peer x25519: %w", err)
}
ssX, err := i.x25519.ECDH(pub)
if err != nil {
return PSK{}, fmt.Errorf("x25519 ecdh: %w", err)
}
return derivePSK(ssMLKEM, ssX, i.offer, answer, b), nil
}
// Respond consumes an initiator offer, produces the answer, and derives the PSK.
func Respond(offer []byte, b Binding) (answer []byte, psk PSK, err error) {
if len(offer) != OfferSize {
return nil, PSK{}, fmt.Errorf("offer: got %d bytes, want %d", len(offer), OfferSize)
}
peerEK := offer[:mlkem.EncapsulationKeySize768]
peerX := offer[mlkem.EncapsulationKeySize768:]
ek, err := mlkem.NewEncapsulationKey768(peerEK)
if err != nil {
return nil, PSK{}, fmt.Errorf("parse peer ml-kem key: %w", err)
}
ssMLKEM, ct := ek.Encapsulate()
x, err := ecdh.X25519().GenerateKey(rand.Reader)
if err != nil {
return nil, PSK{}, fmt.Errorf("x25519 keygen: %w", err)
}
pub, err := ecdh.X25519().NewPublicKey(peerX)
if err != nil {
return nil, PSK{}, fmt.Errorf("parse peer x25519: %w", err)
}
ssX, err := x.ECDH(pub)
if err != nil {
return nil, PSK{}, fmt.Errorf("x25519 ecdh: %w", err)
}
answer = make([]byte, 0, AnswerSize)
answer = append(answer, ct...)
answer = append(answer, x.PublicKey().Bytes()...)
// derivePSK uses the same argument order on both sides; the responder's local
// binding is the mirror of the initiator's, canonicalised inside derivePSK.
return answer, derivePSK(ssMLKEM, ssX, offer, answer, b), nil
}
// derivePSK combines the two shared secrets and binds the result to the full
// transcript (offer ‖ answer) and the canonicalised peer identities.
//
// TODO(NET-1406): replace the SHA-256 concat with the RFC HKDF combiner
// (crypto/hkdf, Go 1.24+) and proper labels before this leaves spike status.
func derivePSK(ssMLKEM, ssX, offer, answer []byte, b Binding) PSK {
lo, hi := canonicalPair(b.LocalID, b.RemoteID)
h := sha256.New()
h.Write([]byte(pskLabel))
h.Write(ssMLKEM)
h.Write(ssX)
h.Write(offer)
h.Write(answer)
h.Write(lo)
h.Write(hi)
var psk PSK
copy(psk[:], h.Sum(nil))
return psk
}
func canonicalPair(a, b []byte) (lo, hi []byte) {
if string(a) <= string(b) {
return a, b
}
return b, a
}

View File

@@ -0,0 +1,89 @@
package pqkem
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
var (
wgA = []byte("peer-A-wireguard-pubkey-32bytes!")
wgB = []byte("peer-B-wireguard-pubkey-32bytes!")
)
func TestExchange_DerivesMatchingPSK(t *testing.T) {
init, err := NewInitiator()
require.NoError(t, err)
require.Len(t, init.Offer(), OfferSize)
answer, pskB, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
require.NoError(t, err)
require.Len(t, answer, AnswerSize)
pskA, err := init.Finish(answer, Binding{LocalID: wgA, RemoteID: wgB})
require.NoError(t, err)
require.Equal(t, pskB, pskA, "both sides must derive the same PSK")
require.NotEqual(t, PSK{}, pskA, "PSK must not be zero")
}
func TestExchange_PSKBoundToPeerIdentities(t *testing.T) {
init, err := NewInitiator()
require.NoError(t, err)
// responder computes with the honest pair...
_, pskHonest, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
require.NoError(t, err)
// ...a second responder run with a different peer identity yields a different PSK,
// even though the KEM material would otherwise combine identically.
wgC := []byte("peer-C-wireguard-pubkey-32bytes!")
_, pskWrong, err := Respond(init.Offer(), Binding{LocalID: wgC, RemoteID: wgA})
require.NoError(t, err)
require.NotEqual(t, pskHonest, pskWrong, "PSK must be bound to the peer pair")
}
func TestExchange_RejectsMalformedMessages(t *testing.T) {
init, err := NewInitiator()
require.NoError(t, err)
_, _, err = Respond(init.Offer()[:10], Binding{})
require.Error(t, err)
_, err = init.Finish([]byte("too short"), Binding{})
require.Error(t, err)
}
// TestExchange_ReportSizesAndTiming is a spike measurement, not a pass/fail gate.
// Run with: go test -run TestExchange_ReportSizesAndTiming -v ./client/internal/pqkem/
func TestExchange_ReportSizesAndTiming(t *testing.T) {
const iters = 200
var tInit, tResp, tFinish time.Duration
for i := 0; i < iters; i++ {
s0 := time.Now()
init, err := NewInitiator()
require.NoError(t, err)
tInit += time.Since(s0)
s1 := time.Now()
answer, _, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
require.NoError(t, err)
tResp += time.Since(s1)
s2 := time.Now()
_, err = init.Finish(answer, Binding{LocalID: wgA, RemoteID: wgB})
require.NoError(t, err)
tFinish += time.Since(s2)
}
t.Logf("wire sizes: offer=%d B answer=%d B (Rosenpass static pubkey ~524160 B)", OfferSize, AnswerSize)
t.Logf("total on-wire per handshake: %d B (~%.0fx smaller than RP static key)", OfferSize+AnswerSize, 524160.0/float64(OfferSize+AnswerSize))
t.Logf("avg NewInitiator (keygen): %s", tInit/iters)
t.Logf("avg Respond (encaps+dh): %s", tResp/iters)
t.Logf("avg Finish (decaps+dh): %s", tFinish/iters)
t.Logf("avg full handshake CPU: %s", (tInit+tResp+tFinish)/iters)
}

View File

@@ -0,0 +1,372 @@
package pqkem
import (
"context"
"crypto/rand"
"fmt"
"log/slog"
"net/netip"
"sync"
"time"
)
const (
// DefaultRetryInterval is how often the initiator retransmits its outstanding
// data-path offer while awaiting the answer.
DefaultRetryInterval = 2 * time.Second
// DefaultMaxRetries bounds how many ticks an exchange may run before it is
// declared failed. The convergence deadline is thus MaxRetries * RetryInterval.
DefaultMaxRetries = 10
// DefaultMaxRekeyFailures is how many consecutive rekey (non-initial) failures
// are tolerated before OnRekeyFailed. The initial exchange fails immediately.
DefaultMaxRekeyFailures = 3
)
// LocalID and RemoteID are peer identity keys (e.g. WireGuard public keys). They are
// distinct types so the local and a remote identity cannot be mixed up.
type (
LocalID string
RemoteID string
)
// Transport is the data-path socket the Manager drives (the analogue of
// go-rosenpass's Conn). It is a dumb mover of bytes to/from endpoints: the Manager
// owns the remoteID<->endpoint routing and hands the transport a resolved endpoint
// to Send, and reverse-resolves the source of each inbound datagram. Its lifecycle
// belongs to the Manager (Run at Start, Close at Stop).
type Transport interface {
// Send delivers msg to the given data-path endpoint.
Send(endpoint netip.AddrPort, msg []byte) error
// LocalPort is the bound local UDP port, announced to peers so they know where
// to send data-path messages.
LocalPort() int
// Run starts delivering inbound datagrams as (source endpoint, msg) to onInbound
// and returns immediately; it runs until Close.
Run(onInbound func(src netip.AddrPort, msg []byte))
// Close stops delivery and releases the socket.
Close() error
}
// exchangeState is the single source of truth for an exchange's role and phase.
type exchangeState uint8
const (
stateReserved exchangeState = iota // responder: deriving the answer
stateAwaitingAnswer // initiator: offer sent, awaiting the answer
stateAwaitingRekey // initiator: PSK derived+set, awaiting OnDataPathRekeyed to chain the next offer
stateAwaitingAck // responder: answer sent, awaiting the next offer that acks this exchange
)
// exchangeCtl holds all state for one in-flight exchange with a peer, under the
// Manager's single lock. state drives every decision. lastSent is the current
// data-path retransmit payload (the offer, for the initiator). initiator is the
// ephemeral handle used at Finish; pendingPSK is the responder's derived key.
// viaSignal records that the offer went to the host for the signalling channel, so
// the loop does not retransmit it on the data path. Only the initiator runs a
// retransmit loop, so only it sets cancel.
type exchangeCtl struct {
id ExchangeID
state exchangeState
startedAt time.Time
cancel context.CancelFunc
lastSent []byte
initiator *Initiator
pendingPSK PSK
viaSignal bool
}
// Manager is the stateful orchestrator — the analogue of go-rosenpass's Server. It
// drives the X25519MLKEM768 exchange, owns the peer endpoint routing and the data-path
// transport, and surfaces the derived PSK and convergence to the host via
// CallbackHandler. It is event-driven: the bootstrap is triggered by the host
// (SignalOffer) and each rotation is clocked by OnDataPathRekeyed. The cryptography is
// the pure kem.go primitives; all state lives here under one lock.
type Manager struct {
localID LocalID
cbHandler CallbackHandler
logger *slog.Logger
retryInterval time.Duration
maxRetries int
maxRekeyFailures int
rootCtx context.Context
rootCancel context.CancelFunc
mu sync.Mutex
transport Transport
exchanges map[RemoteID]*exchangeCtl // in-flight exchange per peer
established map[RemoteID]bool // peer has completed at least one exchange
failures map[RemoteID]int // consecutive rekey failures per peer
psks map[RemoteID]PSK // latest derived PSK per peer (pulled at WG peer-config time)
peerAddrs map[RemoteID]netip.AddrPort // remoteID -> data-path endpoint (send routing)
peersByAddr map[netip.AddrPort]RemoteID // reverse: source endpoint -> remoteID (inbound)
wait sync.WaitGroup
}
// NewManager builds a manager for the local peer identified by its peer identity key
// (used for the deterministic initiator role and the identity binding). A nil logger
// falls back to slog.Default(). Install the data-path transport with Start.
func NewManager(localID LocalID, h CallbackHandler, logger *slog.Logger) *Manager {
if logger == nil {
logger = slog.Default()
}
ctx, cancel := context.WithCancel(context.Background())
return &Manager{
localID: localID,
cbHandler: h,
logger: logger,
retryInterval: DefaultRetryInterval,
maxRetries: DefaultMaxRetries,
maxRekeyFailures: DefaultMaxRekeyFailures,
rootCtx: ctx,
rootCancel: cancel,
exchanges: make(map[RemoteID]*exchangeCtl),
established: make(map[RemoteID]bool),
failures: make(map[RemoteID]int),
psks: make(map[RemoteID]PSK),
peerAddrs: make(map[RemoteID]netip.AddrPort),
peersByAddr: make(map[netip.AddrPort]RemoteID),
}
}
// Start installs the data-path transport and begins its inbound delivery. The Manager
// owns it from here; Stop closes it. Start/Stop are the transport lifecycle pair.
func (m *Manager) Start(t Transport) {
m.mu.Lock()
m.transport = t
m.mu.Unlock()
if t != nil {
t.Run(m.onDataPathInbound)
}
}
// LocalPort is the data-path transport's bound UDP port (0 if no transport), to be
// announced to peers.
func (m *Manager) LocalPort() int {
m.mu.Lock()
t := m.transport
m.mu.Unlock()
if t == nil {
return 0
}
return t.LocalPort()
}
// IsInitiator reports whether the local peer drives the exchange for this remote
// peer. Roles are deterministic (lexicographic identity-key compare) so exactly one
// side initiates, mirroring how Rosenpass picks its handshake initiator.
func (m *Manager) IsInitiator(remoteID RemoteID) bool {
return string(m.localID) > string(remoteID)
}
// PSK returns the latest PSK derived for the peer, for the host to program at WG
// peer-config time (the pull path). ok is false until an exchange has derived one.
func (m *Manager) PSK(remoteID RemoteID) (PSK, bool) {
m.mu.Lock()
defer m.mu.Unlock()
psk, ok := m.psks[remoteID]
return psk, ok
}
// AddPeer registers where a peer's data-path messages are sent and received: its
// overlay endpoint (IP:port). Re-adding updates the endpoint.
func (m *Manager) AddPeer(remoteID RemoteID, endpoint netip.AddrPort) {
if !endpoint.IsValid() {
return
}
m.mu.Lock()
if old, ok := m.peerAddrs[remoteID]; ok {
delete(m.peersByAddr, old)
}
m.peerAddrs[remoteID] = endpoint
m.peersByAddr[endpoint] = remoteID
m.mu.Unlock()
}
// RemovePeer stops any in-flight exchange for a peer and drops its state and routing.
func (m *Manager) RemovePeer(remoteID RemoteID) {
m.mu.Lock()
if ex, ok := m.exchanges[remoteID]; ok {
if ex.cancel != nil {
ex.cancel()
}
delete(m.exchanges, remoteID)
}
delete(m.established, remoteID)
delete(m.failures, remoteID)
delete(m.psks, remoteID)
if ep, ok := m.peerAddrs[remoteID]; ok {
delete(m.peersByAddr, ep)
delete(m.peerAddrs, remoteID)
}
m.mu.Unlock()
}
// Stop cancels all in-flight exchanges, closes the transport, and waits for the
// exchange goroutines to exit.
func (m *Manager) Stop() {
m.rootCancel()
m.wait.Wait()
m.mu.Lock()
t := m.transport
m.transport = nil
m.exchanges = make(map[RemoteID]*exchangeCtl)
m.psks = make(map[RemoteID]PSK)
m.mu.Unlock()
if t != nil {
if err := t.Close(); err != nil {
m.logger.Warn("pqkem: closing data-path transport", "err", err)
}
}
}
// ---- Signalling channel (host-driven; rides the host's negotiation) ----
// SignalOffer returns the KEM offer for the host to embed in its outgoing offer to
// remoteID (bootstrap). It returns (nil, nil) when the local peer is not the
// initiator. It is idempotent for an in-flight bootstrap: a repeat call returns the
// same offer rather than starting a new exchange.
func (m *Manager) SignalOffer(remoteID RemoteID) ([]byte, error) {
if !m.IsInitiator(remoteID) {
return nil, nil
}
m.mu.Lock()
if ex := m.exchanges[remoteID]; ex != nil && ex.viaSignal && ex.state == stateAwaitingAnswer {
last := ex.lastSent
m.mu.Unlock()
return last, nil
}
m.mu.Unlock()
// bootstrap offer acknowledges nothing (zero AckID).
return m.startExchange(remoteID, true, ExchangeID{})
}
// SignalOnOffer processes a KEM offer the host extracted from an incoming offer and
// returns the KEM answer for the host to embed in its outgoing answer.
func (m *Manager) SignalOnOffer(remoteID RemoteID, offer []byte) ([]byte, error) {
typ, msg, err := Decode(offer)
if err != nil {
return nil, fmt.Errorf("decode signal offer from %s: %w", remoteID, err)
}
if typ != MsgOffer {
return nil, fmt.Errorf("expected offer from %s, got type %d", remoteID, typ)
}
return m.processOffer(remoteID, msg.(*OfferMsg))
}
// SignalOnAnswer processes a KEM answer the host extracted from an incoming answer.
// There is no reply: the next offer (over the data path) acknowledges this exchange.
func (m *Manager) SignalOnAnswer(remoteID RemoteID, answer []byte) error {
typ, msg, err := Decode(answer)
if err != nil {
return fmt.Errorf("decode signal answer from %s: %w", remoteID, err)
}
if typ != MsgAnswer {
return fmt.Errorf("expected answer from %s, got type %d", remoteID, typ)
}
return m.processAnswer(remoteID, msg.(*AnswerMsg))
}
// ---- Data path ----
// onDataPathInbound is the transport's inbound handler: it reverse-resolves the
// source endpoint to a peer and dispatches. Unknown sources are dropped.
func (m *Manager) onDataPathInbound(src netip.AddrPort, msg []byte) {
m.mu.Lock()
remoteID, ok := m.peersByAddr[src]
m.mu.Unlock()
if !ok {
return
}
if err := m.OnDataPathMessage(remoteID, msg); err != nil {
m.logger.Debug("pqkem: inbound", "peer", remoteID, "err", err)
}
}
// OnDataPathMessage handles a KEM message received over the data path from remoteID
// and pushes any reply back over the data path.
func (m *Manager) OnDataPathMessage(remoteID RemoteID, raw []byte) error {
typ, msg, err := Decode(raw)
if err != nil {
return fmt.Errorf("decode data-path msg from %s: %w", remoteID, err)
}
switch typ {
case MsgOffer:
answer, err := m.processOffer(remoteID, msg.(*OfferMsg))
if err != nil {
return err
}
if answer == nil {
return nil
}
return m.pushDataPath(remoteID, answer)
case MsgAnswer:
return m.processAnswer(remoteID, msg.(*AnswerMsg))
default:
return fmt.Errorf("unhandled data-path message type %d from %s", typ, remoteID)
}
}
// OnDataPathRekeyed notifies that the peer's data path is up and freshly keyed with
// the latest PSK (fired on first establishment AND every rekey). If we are the
// initiator that just derived a PSK, it chains the next exchange: a fresh offer over
// the data path that acknowledges the just-completed one (its arrival under the new
// key proves to the responder that the key works).
func (m *Manager) OnDataPathRekeyed(remoteID RemoteID) {
m.mu.Lock()
ex := m.exchanges[remoteID]
chain := ex != nil && ex.state == stateAwaitingRekey
var ackID ExchangeID
if chain {
ackID = ex.id
}
m.mu.Unlock()
if !chain {
return
}
offer, err := m.startExchange(remoteID, false, ackID)
if err != nil {
m.logger.Error("pqkem: chain offer failed to start", "peer", remoteID, "err", err)
return
}
if err := m.pushDataPath(remoteID, offer); err != nil {
m.logger.Warn("pqkem: send chain offer failed", "peer", remoteID, "err", err)
}
}
// OnDataPathDown notifies that the peer's data path went down. Rotations resume once
// the host re-bootstraps over signalling on reconnect; in-flight data-path sends will
// simply fail until then. Reserved as an explicit hook.
func (m *Manager) OnDataPathDown(remoteID RemoteID) {}
// ---- internals ----
// pushDataPath resolves the peer's endpoint and sends over the data-path transport,
// erroring if the peer is unknown or no transport is set.
func (m *Manager) pushDataPath(remoteID RemoteID, msg []byte) error {
m.mu.Lock()
ep, ok := m.peerAddrs[remoteID]
t := m.transport
m.mu.Unlock()
if !ok {
return fmt.Errorf("no data-path endpoint for peer %s", remoteID)
}
if t == nil {
return fmt.Errorf("no data-path transport")
}
return t.Send(ep, msg)
}
func (m *Manager) binding(remoteID RemoteID) Binding {
return Binding{LocalID: []byte(m.localID), RemoteID: []byte(remoteID)}
}
func newExchangeID() (ExchangeID, error) {
var id ExchangeID
if _, err := rand.Read(id[:]); err != nil {
return ExchangeID{}, fmt.Errorf("generate exchange id: %w", err)
}
return id, nil
}

View File

@@ -0,0 +1,169 @@
package pqkem
import (
"fmt"
"net/netip"
"sync"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
)
// netSwitch is an in-memory UDP fabric: transports register their endpoint and get
// datagrams delivered to their inbound handler.
type netSwitch struct {
mu sync.Mutex
h map[netip.AddrPort]func(netip.AddrPort, []byte)
}
func newSwitch() *netSwitch {
return &netSwitch{h: map[netip.AddrPort]func(netip.AddrPort, []byte){}}
}
func (s *netSwitch) register(ep netip.AddrPort, fn func(netip.AddrPort, []byte)) {
s.mu.Lock()
s.h[ep] = fn
s.mu.Unlock()
}
func (s *netSwitch) deliver(dst, src netip.AddrPort, msg []byte) error {
s.mu.Lock()
fn := s.h[dst]
s.mu.Unlock()
if fn == nil {
return fmt.Errorf("no route to %s", dst)
}
fn(src, msg)
return nil
}
// loopback is an endpoint-based pqkem.Transport over a netSwitch, with a switchable
// drop flag.
type loopback struct {
ep netip.AddrPort
sw *netSwitch
drop atomic.Bool
}
func (l *loopback) Send(dst netip.AddrPort, msg []byte) error {
if l.drop.Load() {
return nil
}
return l.sw.deliver(dst, l.ep, append([]byte(nil), msg...))
}
func (l *loopback) LocalPort() int { return int(l.ep.Port()) }
func (l *loopback) Run(onInbound func(netip.AddrPort, []byte)) { l.sw.register(l.ep, onInbound) }
func (l *loopback) Close() error { return nil }
type fakeWG struct {
mu sync.Mutex
psks map[RemoteID]PSK
failed []RemoteID
}
func newFakeWG() *fakeWG { return &fakeWG{psks: map[RemoteID]PSK{}} }
func (f *fakeWG) OnNewPSKReady(remoteID RemoteID, psk PSK) error {
f.mu.Lock()
defer f.mu.Unlock()
f.psks[remoteID] = psk
return nil
}
func (f *fakeWG) OnRekeyFailed(remoteID RemoteID) error {
f.mu.Lock()
defer f.mu.Unlock()
f.failed = append(f.failed, remoteID)
return nil
}
func (f *fakeWG) psk(peer RemoteID) PSK {
f.mu.Lock()
defer f.mu.Unlock()
return f.psks[peer]
}
var (
epA = netip.MustParseAddrPort("100.64.0.1:51833")
epB = netip.MustParseAddrPort("100.64.0.2:51833")
)
// pair builds two wired managers (B is the initiator, "bbbb" > "aaaa") sharing a
// netSwitch, with each peer's data-path endpoint registered. lbB is B's loopback
// (for toggling drop).
func pair(t *testing.T) (dA, dB *Manager, wgA, wgB *fakeWG, lbB *loopback) {
t.Helper()
sw := newSwitch()
wgA = newFakeWG()
wgB = newFakeWG()
dA = NewManager("aaaa", wgA, nil)
dB = NewManager("bbbb", wgB, nil)
dA.Start(&loopback{ep: epA, sw: sw})
lbB = &loopback{ep: epB, sw: sw}
dB.Start(lbB)
dA.AddPeer("bbbb", epB)
dB.AddPeer("aaaa", epA)
return dA, dB, wgA, wgB, lbB
}
// bootstrap runs the signalling offer/answer (the test plays the host carrying bytes).
func bootstrap(t *testing.T, dA, dB *Manager) {
t.Helper()
offer, err := dB.SignalOffer("aaaa")
require.NoError(t, err)
require.NotNil(t, offer)
answer, err := dA.SignalOnOffer("bbbb", offer)
require.NoError(t, err)
require.NotNil(t, answer)
require.NoError(t, dB.SignalOnAnswer("aaaa", answer))
}
func TestManager_BootstrapDerivesSamePSK(t *testing.T) {
dA, dB, wgA, wgB, _ := pair(t)
defer dA.Stop()
defer dB.Stop()
bootstrap(t, dA, dB)
pskA := wgA.psk("bbbb")
pskB := wgB.psk("aaaa")
require.NotEqual(t, PSK{}, pskA)
require.Equal(t, pskB, pskA, "both sides derive the same PSK from the bootstrap exchange")
}
func TestManager_ChainRotatesAndAcks(t *testing.T) {
dA, dB, wgA, wgB, _ := pair(t)
defer dA.Stop()
defer dB.Stop()
bootstrap(t, dA, dB)
psk1 := wgB.psk("aaaa")
// Data path up: B (initiator) chains the next offer over the data path, which
// rotates both to a fresh PSK and acknowledges A.
dA.OnDataPathRekeyed("bbbb")
dB.OnDataPathRekeyed("aaaa")
psk2A := wgA.psk("bbbb")
psk2B := wgB.psk("aaaa")
require.Equal(t, psk2B, psk2A, "both sides converge on the rotated PSK")
require.NotEqual(t, psk1, psk2B, "the chain rotated to a new PSK")
}
func TestManager_NonInitiatorReturnsNoOffer(t *testing.T) {
dA := NewManager("aaaa", newFakeWG(), nil)
defer dA.Stop()
offer, err := dA.SignalOffer("bbbb") // not the initiator vs "bbbb"
require.NoError(t, err)
require.Nil(t, offer)
}
func TestManager_StopIsIdempotent(t *testing.T) {
dA := NewManager("aaaa", newFakeWG(), nil)
dA.Start(&loopback{ep: epA, sw: newSwitch()})
dA.Stop()
dA.Stop() // must not panic or hang
}

View File

@@ -0,0 +1,121 @@
package pqkem
import (
"crypto/mlkem"
"fmt"
)
// Wire framing for the PQ-KEM exchange. Messages are self-contained, versioned,
// transport-agnostic byte blobs: the same bytes ride the signalling channel
// (initial bootstrap) or a data-tunnel packet (rekey). The library only ever sees
// opaque []byte at the transport seam.
//
// Layout (all messages): [type:1][version:1][exchangeID:16][payload...]
//
// There is no confirm message: an exchange is acknowledged by the NEXT offer, which
// carries the acked exchange's id (see OfferMsg.AckID) and — riding the data path
// under the freshly adopted key — proves that key works.
const (
// ProtocolVersion is bumped on any wire-incompatible change; a peer rejects
// messages it does not understand rather than misparsing them.
ProtocolVersion uint8 = 1
// ExchangeIDSize identifies one exchange so answers/acks correlate and stale
// messages are dropped.
ExchangeIDSize = 16
headerSize = 1 + 1 + ExchangeIDSize
)
// MsgType tags the two message kinds of the exchange.
type MsgType uint8
const (
MsgOffer MsgType = iota + 1
MsgAnswer
)
// ExchangeID is the per-exchange correlator. The zero value means "none" (an offer
// that acknowledges nothing, i.e. the first exchange of a connection).
type ExchangeID [ExchangeIDSize]byte
// OfferMsg carries the initiator's public material (X25519 pub ‖ ML-KEM encap key)
// and AckID, the id of the previous exchange this offer acknowledges (zero if none).
type OfferMsg struct {
ExchangeID ExchangeID
AckID ExchangeID
// KEMOffer is the raw Initiator.Offer() blob (OfferSize bytes).
KEMOffer []byte
}
// AnswerMsg carries the responder's reply (ML-KEM ciphertext ‖ X25519 pub) for the
// round identified by ExchangeID.
type AnswerMsg struct {
ExchangeID ExchangeID
// KEMAnswer is the raw Respond() answer blob (AnswerSize bytes).
KEMAnswer []byte
}
// Encode serialises the offer with its framed header (payload = AckID ‖ KEMOffer).
func (m *OfferMsg) Encode() ([]byte, error) {
if len(m.KEMOffer) != OfferSize {
return nil, fmt.Errorf("offer payload: got %d, want %d", len(m.KEMOffer), OfferSize)
}
payload := make([]byte, 0, ExchangeIDSize+OfferSize)
payload = append(payload, m.AckID[:]...)
payload = append(payload, m.KEMOffer...)
return frame(MsgOffer, m.ExchangeID, payload), nil
}
// Encode serialises the answer with its framed header.
func (m *AnswerMsg) Encode() ([]byte, error) {
if len(m.KEMAnswer) != AnswerSize {
return nil, fmt.Errorf("answer payload: got %d, want %d", len(m.KEMAnswer), AnswerSize)
}
return frame(MsgAnswer, m.ExchangeID, m.KEMAnswer), nil
}
// Decode parses a framed message into one of *OfferMsg / *AnswerMsg.
func Decode(buf []byte) (MsgType, any, error) {
if len(buf) < headerSize {
return 0, nil, fmt.Errorf("message too short: %d bytes", len(buf))
}
typ := MsgType(buf[0])
if ver := buf[1]; ver != ProtocolVersion {
return typ, nil, fmt.Errorf("unsupported protocol version %d (want %d)", ver, ProtocolVersion)
}
var id ExchangeID
copy(id[:], buf[2:headerSize])
payload := buf[headerSize:]
switch typ {
case MsgOffer:
if len(payload) != ExchangeIDSize+OfferSize {
return typ, nil, fmt.Errorf("offer payload: got %d, want %d", len(payload), ExchangeIDSize+OfferSize)
}
var ack ExchangeID
copy(ack[:], payload[:ExchangeIDSize])
return typ, &OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: payload[ExchangeIDSize:]}, nil
case MsgAnswer:
if len(payload) != AnswerSize {
return typ, nil, fmt.Errorf("answer payload: got %d, want %d", len(payload), AnswerSize)
}
return typ, &AnswerMsg{ExchangeID: id, KEMAnswer: payload}, nil
default:
return typ, nil, fmt.Errorf("unknown message type %d", typ)
}
}
func frame(typ MsgType, id ExchangeID, payload []byte) []byte {
buf := make([]byte, headerSize+len(payload))
buf[0] = byte(typ)
buf[1] = ProtocolVersion
copy(buf[2:], id[:])
copy(buf[headerSize:], payload)
return buf
}
// compile-time assurance the KEM blob sizes referenced here stay in sync with kem.go.
var _ = [1]struct{}{}[OfferSize-(32+mlkem.EncapsulationKeySize768)]

View File

@@ -0,0 +1,57 @@
package pqkem
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestMessageRoundTrip(t *testing.T) {
init, err := NewInitiator()
require.NoError(t, err)
answer, _, err := Respond(init.Offer(), Binding{LocalID: wgB, RemoteID: wgA})
require.NoError(t, err)
id := ExchangeID{1, 2, 3, 4}
ack := ExchangeID{9, 9, 9}
offBytes, err := (&OfferMsg{ExchangeID: id, AckID: ack, KEMOffer: init.Offer()}).Encode()
require.NoError(t, err)
typ, decoded, err := Decode(offBytes)
require.NoError(t, err)
require.Equal(t, MsgOffer, typ)
require.Equal(t, id, decoded.(*OfferMsg).ExchangeID)
require.Equal(t, ack, decoded.(*OfferMsg).AckID)
require.Equal(t, init.Offer(), decoded.(*OfferMsg).KEMOffer)
ansBytes, err := (&AnswerMsg{ExchangeID: id, KEMAnswer: answer}).Encode()
require.NoError(t, err)
typ, decoded, err = Decode(ansBytes)
require.NoError(t, err)
require.Equal(t, MsgAnswer, typ)
require.Equal(t, answer, decoded.(*AnswerMsg).KEMAnswer)
}
func TestDecodeRejects(t *testing.T) {
// too short
_, _, err := Decode([]byte{1, 1})
require.Error(t, err)
// wrong version
bad := make([]byte, headerSize+ExchangeIDSize+OfferSize)
bad[0] = byte(MsgOffer)
bad[1] = ProtocolVersion + 1
_, _, err = Decode(bad)
require.Error(t, err)
// unknown type
bad2 := make([]byte, headerSize)
bad2[0] = 99
bad2[1] = ProtocolVersion
_, _, err = Decode(bad2)
require.Error(t, err)
// offer with wrong payload size
_, err = (&OfferMsg{KEMOffer: []byte{1, 2, 3}}).Encode()
require.Error(t, err)
}

View File

@@ -0,0 +1,102 @@
package internal
import (
"net/netip"
log "github.com/sirupsen/logrus"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"github.com/netbirdio/netbird/client/internal/pqkem"
)
// pqPresharedKeySetter is the subset of the WireGuard interface the ML-KEM callback
// needs: programming a peer's preshared key. *iface.WGIface satisfies it.
type pqPresharedKeySetter interface {
SetPresharedKey(peerKey string, psk wgtypes.Key, updateOnly bool) error
}
// pqCallbackHandler programs the derived PQ PSK onto the WireGuard peer. It is the
// engine-side implementation of pqkem.CallbackHandler.
type pqCallbackHandler struct {
wg pqPresharedKeySetter
}
// OnNewPSKReady programs the freshly derived PSK for the peer (updateOnly: a no-op
// if the peer is not present, mirroring Rosenpass). remoteID is the peer's WG pubkey.
func (h pqCallbackHandler) OnNewPSKReady(remoteID pqkem.RemoteID, psk pqkem.PSK) error {
// updateOnly: applies to an already-configured peer (rotation). At bootstrap the
// peer is not configured yet, so this is a no-op there and the PSK is instead
// pulled at peer-config time (pqHandshaker.PSK / conn.presharedKey).
log.Debugf("pqkem: programming PSK for peer %s", remoteID)
return h.wg.SetPresharedKey(string(remoteID), wgtypes.Key(psk), true)
}
// OnRekeyFailed reports a failed PQ (re)key convergence.
// TODO(NET-1406): tear the peer connection down / trigger ICE reconnect.
func (h pqCallbackHandler) OnRekeyFailed(remoteID pqkem.RemoteID) error {
log.Warnf("pqkem: post-quantum rekey failed for peer %s", remoteID)
return nil
}
// pqHandshaker adapts the pqkem manager to peer.PQHandshaker (string peer keys),
// wiring the host's signalling offers/answers to the KEM exchange.
type pqHandshaker struct {
mgr *pqkem.Manager
}
func (p pqHandshaker) OfferPayload(remoteKey string) ([]byte, int) {
payload, err := p.mgr.SignalOffer(pqkem.RemoteID(remoteKey))
if err != nil {
log.Warnf("pqkem: build offer for %s: %v", remoteKey, err)
}
return payload, p.mgr.LocalPort()
}
func (p pqHandshaker) AnswerPayload(remoteKey string, recvOffer []byte) ([]byte, int) {
if len(recvOffer) == 0 {
return nil, p.mgr.LocalPort()
}
payload, err := p.mgr.SignalOnOffer(pqkem.RemoteID(remoteKey), recvOffer)
if err != nil {
log.Warnf("pqkem: build answer for %s: %v", remoteKey, err)
}
return payload, p.mgr.LocalPort()
}
func (p pqHandshaker) OnAnswer(remoteKey string, recvAnswer []byte) {
if len(recvAnswer) == 0 {
return
}
if err := p.mgr.SignalOnAnswer(pqkem.RemoteID(remoteKey), recvAnswer); err != nil {
log.Warnf("pqkem: process answer from %s: %v", remoteKey, err)
}
}
// PSK exposes the peer's derived PSK for the conn to program at WG peer-config time.
func (p pqHandshaker) PSK(remoteKey string) (wgtypes.Key, bool) {
psk, ok := p.mgr.PSK(pqkem.RemoteID(remoteKey))
if !ok {
return wgtypes.Key{}, false
}
return wgtypes.Key(psk), true
}
// SetRemoteAddr registers the peer's data-path endpoint (overlay IP + pq UDP port)
// learned from signalling. Sends only ever fire once the tunnel is up (clocked by
// OnDataPathRekeyed), so registering here is safe even before connection-up.
func (p pqHandshaker) SetRemoteAddr(remoteKey string, addr netip.AddrPort) {
if !addr.IsValid() || addr.Port() == 0 {
return
}
p.mgr.AddPeer(pqkem.RemoteID(remoteKey), addr)
}
// OnDataPathRekeyed clocks the next chained PSK rotation on a fresh WG handshake.
func (p pqHandshaker) OnDataPathRekeyed(remoteKey string) {
p.mgr.OnDataPathRekeyed(pqkem.RemoteID(remoteKey))
}
// OnDataPathDown signals the peer's tunnel went down.
func (p pqHandshaker) OnDataPathDown(remoteKey string) {
p.mgr.OnDataPathDown(pqkem.RemoteID(remoteKey))
}

View File

@@ -0,0 +1,72 @@
package internal
import (
"fmt"
"net"
"net/netip"
log "github.com/sirupsen/logrus"
)
// DefaultPort is the preferred UDP port for the ML-KEM data-path service, bound on
// the WG overlay IP. Since each client owns a distinct overlay IP, this port is
// almost always free, so it need not be announced (peers assume it). A peer only
// announces Body.mlkemPort when a collision forced it onto a different port.
const DefaultPort = 51833
// pqTransport is the ML-KEM data-path transport: a dumb UDP socket bound on the WG
// overlay IP. It implements pqkem.Transport — the manager owns the remoteID<->endpoint
// routing and drives this socket's lifecycle (Run / Close).
type pqTransport struct {
conn *net.UDPConn
port int
}
// newPQTransport binds a UDP socket on the WG overlay IP, preferring DefaultPort and
// falling back to an OS-assigned ephemeral port if it is in use. Call it after the WG
// interface is up so the overlay IP is assigned; when the bound port is not
// DefaultPort it must be announced to peers via Body.mlkemPort.
func newPQTransport(overlayIP netip.Addr) (*pqTransport, error) {
if !overlayIP.IsValid() {
return nil, fmt.Errorf("invalid overlay IP for pqkem transport")
}
ip := net.IP(overlayIP.AsSlice())
conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: ip, Port: DefaultPort})
if err != nil {
log.Debugf("pqkem: default port %d unavailable on %s (%v), using an ephemeral port", DefaultPort, overlayIP, err)
conn, err = net.ListenUDP("udp4", &net.UDPAddr{IP: ip, Port: 0})
if err != nil {
return nil, fmt.Errorf("bind pqkem udp on overlay %s: %w", overlayIP, err)
}
}
return &pqTransport{conn: conn, port: conn.LocalAddr().(*net.UDPAddr).Port}, nil
}
// Send implements pqkem.Transport.
func (t *pqTransport) Send(endpoint netip.AddrPort, msg []byte) error {
_, err := t.conn.WriteToUDPAddrPort(msg, endpoint)
return err
}
// LocalPort implements pqkem.Transport.
func (t *pqTransport) LocalPort() int { return t.port }
// Run implements pqkem.Transport: the receive loop, delivering each datagram as
// (source endpoint, msg). Exits when the socket is closed.
func (t *pqTransport) Run(onInbound func(src netip.AddrPort, msg []byte)) {
go func() {
buf := make([]byte, 2048)
for {
n, src, err := t.conn.ReadFromUDPAddrPort(buf)
if err != nil {
return
}
msg := make([]byte, n)
copy(msg, buf[:n])
onInbound(src, msg)
}
}()
}
// Close implements pqkem.Transport.
func (t *pqTransport) Close() error { return t.conn.Close() }

View 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()
}

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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)

View File

@@ -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
}

View File

@@ -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))

View File

@@ -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")

View File

@@ -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")

View File

@@ -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)

View File

@@ -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(*)",

View 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")
}

View File

@@ -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.

View File

@@ -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))

View File

@@ -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
}

View File

@@ -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)
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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,'');

View File

@@ -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
}

View File

@@ -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": {}}},
)
}

View File

@@ -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"
)
@@ -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 {

View File

@@ -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"},
}},
}},

View File

@@ -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) *proto.PeerConfig {
netmask, _ := network.Net.Mask.Size()
fqdn := peer.FQDN(dnsName)

View File

@@ -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),
Checks: toProtocolChecks(ctx, postureChecks),
}

View File

@@ -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

View File

@@ -446,7 +446,7 @@ func (h *Handler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request) {
netMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, dns.CustomZone{}, nil, validPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil, account.GetActiveGroupUsers())
util.WriteJSONObject(ctx, w, toAccessiblePeers(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
}

View File

@@ -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
}

View File

@@ -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")
}

View File

@@ -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 {

View File

@@ -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(),

View File

@@ -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

View File

@@ -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)

View File

@@ -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"),

View File

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

View File

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

View File

@@ -1201,7 +1201,7 @@ func TestGetNetworkMap_RouteSync(t *testing.T) {
peer1Routes, err := am.GetNetworkMap(context.Background(), peer1ID)
require.NoError(t, err)
require.Len(t, peer1Routes.Routes, 1, "we should receive one route for peer1")
require.True(t, types.TwinRoute(expectedRoute).Equal(peer1Routes.Routes[0]), "received route should be equal")
require.True(t, expectedRoute.Equal(peer1Routes.Routes[0]), "received route should be equal")
peer2Routes, err := am.GetNetworkMap(context.Background(), peer2ID)
require.NoError(t, err)

View File

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

View File

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

View File

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

View File

@@ -18,6 +18,8 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
proxydomain "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/internals/modules/zones"
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -1080,6 +1082,7 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
peersExists := make(map[string]struct{})
rules := make([]*FirewallRule, 0)
peers := make([]*nbpeer.Peer, 0)
targetComponent := targetPeer.ToComponent()
return func(rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int) {
for _, peer := range groupPeers {
@@ -1115,10 +1118,10 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
if len(rule.Ports) == 0 && len(rule.PortRanges) == 0 {
rules = append(rules, &fr)
} else {
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetPeer)...)
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetComponent)...)
}
rules = AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, FirewallRuleContext{
rules = AppendIPv6FirewallRule(rules, rulesExists, peer.ToComponent(), targetComponent, rule, FirewallRuleContext{
Direction: direction,
DirStr: strconv.Itoa(direction),
ProtocolStr: string(protocol),
@@ -1278,7 +1281,7 @@ func (a *Account) getRouteFirewallRules(ctx context.Context, peerID string, poli
return fwRules
}
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*nbpeer.Peer {
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*ComponentPeer {
distPeersWithPolicy := make(map[string]struct{})
for _, id := range rule.Sources {
group := a.Groups[id]
@@ -1305,13 +1308,13 @@ func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID
}
}
distributionGroupPeers := make([]*nbpeer.Peer, 0, len(distPeersWithPolicy))
distributionGroupPeers := make([]*ComponentPeer, 0, len(distPeersWithPolicy))
for pID := range distPeersWithPolicy {
peer := a.Peers[pID]
if peer == nil {
continue
}
distributionGroupPeers = append(distributionGroupPeers, peer)
distributionGroupPeers = append(distributionGroupPeers, peer.ToComponent())
}
return distributionGroupPeers
}
@@ -1796,3 +1799,66 @@ func filterZoneRecordsForPeers(peer *nbpeer.Peer, customZone nbdns.CustomZone, p
return filteredRecords
}
// filterPeerAppliedZones filters account zones based on the peer's group membership
func filterPeerAppliedZones(ctx context.Context, accountZones []*zones.Zone, peerGroups LookupMap) []nbdns.CustomZone {
var customZones []nbdns.CustomZone
if len(peerGroups) == 0 {
return customZones
}
for _, zone := range accountZones {
if !zone.Enabled || len(zone.Records) == 0 {
continue
}
hasAccess := false
for _, distGroupID := range zone.DistributionGroups {
if _, found := peerGroups[distGroupID]; found {
hasAccess = true
break
}
}
if !hasAccess {
continue
}
simpleRecords := make([]nbdns.SimpleRecord, 0, len(zone.Records))
for _, record := range zone.Records {
var recordType int
rData := record.Content
switch record.Type {
case records.RecordTypeA:
recordType = int(dns.TypeA)
case records.RecordTypeAAAA:
recordType = int(dns.TypeAAAA)
case records.RecordTypeCNAME:
recordType = int(dns.TypeCNAME)
rData = dns.Fqdn(record.Content)
default:
log.WithContext(ctx).Warnf("unknown DNS record type %s for record %s", record.Type, record.ID)
continue
}
simpleRecords = append(simpleRecords, nbdns.SimpleRecord{
Name: dns.Fqdn(record.Name),
Type: recordType,
Class: nbdns.DefaultClass,
TTL: record.TTL,
RData: rData,
})
}
customZones = append(customZones, nbdns.CustomZone{
Domain: dns.Fqdn(zone.Domain),
Records: simpleRecords,
SearchDomainDisabled: !zone.EnableSearchDomain,
NonAuthoritative: true,
})
}
return customZones
}

View File

@@ -2,6 +2,7 @@ package types
import (
"context"
"slices"
"time"
log "github.com/sirupsen/logrus"
@@ -10,6 +11,7 @@ import (
"github.com/netbirdio/netbird/management/internals/modules/zones"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/route"
)
// GetPeerNetworkMapResult dispatches to either the legacy-NetworkMap path or
@@ -90,9 +92,6 @@ func (a *Account) GetPeerNetworkMapFromComponents(
return nm
}
// GetPeerNetworkMapComponents builds the account's slim twin store and computes
// the peer's components on it. The calculation itself lives on
// networkmap.NetworkMapData and never touches the Account.
func (a *Account) GetPeerNetworkMapComponents(
ctx context.Context,
peerID string,
@@ -103,6 +102,623 @@ func (a *Account) GetPeerNetworkMapComponents(
routers map[string]map[string]*routerTypes.NetworkRouter,
groupIDToUserIDs map[string][]string,
) *NetworkMapComponents {
nmd := a.toNetworkMapData(accountZones, validatedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
return nmd.GetPeerNetworkMapComponents(peerID, toTwinCustomZone(peersCustomZone))
peer := a.Peers[peerID]
// this can never happen, things are very wrong if it did
// TODO (dmitri) maybe consider using invariants?
if peer == nil {
log.WithField("peer id", peerID).Error("NetworkMapComponents are computed for a peer missing from the account")
return EmptyNetworkMapComponents(&NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
if _, ok := validatedPeersMap[peerID]; !ok {
// Mirror legacy graceful-degrade: GetPeerNetworkMapFromComponents
// returns &NetworkMap{Network: a.Network.Copy()} when components is
// nil. Match that floor so the receiving client always sees the
// account Network identifier, not a fully-empty envelope.
return EmptyNetworkMapComponents(&NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
components := &NetworkMapComponents{
PeerID: peerID,
Network: a.Network.Copy(),
NameServerGroups: make([]*nbdns.NameServerGroup, 0),
CustomZoneDomain: peersCustomZone.Domain,
ResourcePoliciesMap: make(map[string][]*Policy),
RoutersMap: make(map[string]map[string]*ComponentRouter),
NetworkResources: make([]*ComponentResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(a.PostureChecks)),
RouterPeers: make(map[string]*ComponentPeer),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
}
for _, n := range a.Networks {
if n != nil {
components.NetworkXIDToPublicID[n.ID] = n.PublicID
}
}
for _, pc := range a.PostureChecks {
if pc != nil {
components.PostureCheckXIDToPublicID[pc.ID] = pc.PublicID
}
}
components.AccountSettings = &AccountSettingsInfo{
PeerLoginExpirationEnabled: a.Settings.PeerLoginExpirationEnabled,
PeerLoginExpiration: a.Settings.PeerLoginExpiration,
PeerInactivityExpirationEnabled: a.Settings.PeerInactivityExpirationEnabled,
PeerInactivityExpiration: a.Settings.PeerInactivityExpiration,
}
components.DNSSettings = &a.DNSSettings
// relevantPeers always contains the target peer (peerID)
relevantPeers, relevantGroups, relevantPolicies, relevantRoutes, sshReqs := a.getPeersGroupsPoliciesRoutes(ctx, peerID, peer.SSHEnabled, validatedPeersMap, &components.PostureFailedPeers)
if len(sshReqs.neededGroupIDs) > 0 {
components.GroupIDToUserIDs = filterGroupIDToUserIDs(groupIDToUserIDs, sshReqs.neededGroupIDs)
}
if sshReqs.needAllowedUserIDs {
components.AllowedUserIDs = a.getAllowedUserIDs()
}
components.Peers = relevantPeers
components.Groups = GroupsToComponent(relevantGroups)
components.Policies = relevantPolicies
components.Routes = relevantRoutes
components.AllDNSRecords = filterDNSRecordsByPeers(peersCustomZone.Records, relevantPeers, peer.SupportsIPv6() && peer.IPv6.IsValid())
peerGroups := a.GetPeerGroups(peerID)
components.AccountZones = filterPeerAppliedZones(ctx, accountZones, peerGroups)
components.AccountZones = append(components.AccountZones, a.SynthesizePrivateServiceZones(peerID)...)
for _, nsGroup := range a.NameServerGroups {
if nsGroup.Enabled {
for _, gID := range nsGroup.Groups {
if _, found := relevantGroups[gID]; found {
components.NameServerGroups = append(components.NameServerGroups, nsGroup)
break
}
}
}
}
for _, resource := range a.NetworkResources {
if !resource.Enabled {
continue
}
policies, exists := resourcePolicies[resource.ID]
if !exists {
continue
}
addSourcePeers := false
networkRoutingPeers, routerExists := routers[resource.NetworkID]
if routerExists {
if _, ok := networkRoutingPeers[peerID]; ok {
addSourcePeers = true
}
}
for _, policy := range policies {
if addSourcePeers {
var peers []string
if policy.Rules[0].SourceResource.Type == ResourceTypePeer && policy.Rules[0].SourceResource.ID != "" {
peers = []string{policy.Rules[0].SourceResource.ID}
} else {
peers = a.getUniquePeerIDsFromGroupsIDs(ctx, policy.SourceGroups())
}
for _, pID := range a.getPostureValidPeersSaveFailed(peers, policy.SourcePostureChecks, validatedPeersMap, &components.PostureFailedPeers) {
if _, exists := components.Peers[pID]; !exists {
components.Peers[pID] = a.GetPeer(pID).ToComponent()
}
}
} else {
peerInSources := false
if policy.Rules[0].SourceResource.Type == ResourceTypePeer && policy.Rules[0].SourceResource.ID != "" {
peerInSources = policy.Rules[0].SourceResource.ID == peerID
} else {
for _, groupID := range policy.SourceGroups() {
if group := a.GetGroup(groupID); group != nil && slices.Contains(group.Peers, peerID) {
peerInSources = true
break
}
}
}
if !peerInSources {
continue
}
isValid, pname := a.validatePostureChecksOnPeerGetFailed(ctx, policy.SourcePostureChecks, peerID)
if !isValid && len(pname) > 0 {
if _, ok := components.PostureFailedPeers[pname]; !ok {
components.PostureFailedPeers[pname] = make(map[string]struct{})
}
components.PostureFailedPeers[pname][peer.ID] = struct{}{}
continue
}
addSourcePeers = true
}
for _, rule := range policy.Rules {
for _, srcGroupID := range rule.Sources {
if g := a.Groups[srcGroupID]; g != nil {
if _, exists := components.Groups[srcGroupID]; !exists {
components.Groups[srcGroupID] = g.ToComponent()
}
}
}
for _, dstGroupID := range rule.Destinations {
if g := a.Groups[dstGroupID]; g != nil {
if _, exists := components.Groups[dstGroupID]; !exists {
components.Groups[dstGroupID] = g.ToComponent()
}
}
}
}
components.ResourcePoliciesMap[resource.ID] = policies
}
// Only expose router peers and the per-network routers_map when this
// target peer actually has access to the resource (either as a router
// itself or via a policy that includes it as a source). Without this
// gate, every peer's envelope was leaking router peers of every
// network in the account — accounts with many tenants/networks
// shipped tens of unrelated peers in `peers[]` and `routers_map`.
if addSourcePeers {
components.RoutersMap[resource.NetworkID] = routerTypes.ToComponentMap(networkRoutingPeers)
for peerIDKey := range networkRoutingPeers {
if p := a.Peers[peerIDKey]; p != nil {
cp := components.RouterPeers[peerIDKey]
if cp == nil {
cp = p.ToComponent()
components.RouterPeers[peerIDKey] = cp
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := validatedPeersMap[peerIDKey]; validated {
components.Peers[peerIDKey] = cp
}
}
}
}
components.NetworkResources = append(components.NetworkResources, resource.ToComponent())
}
}
filterGroupPeers(&components.Groups, components.Peers)
filterPostureFailedPeers(&components.PostureFailedPeers, components.Policies, components.ResourcePoliciesMap, components.Peers)
return components
}
type sshRequirements struct {
neededGroupIDs map[string]struct{}
needAllowedUserIDs bool
}
func (a *Account) getPeersGroupsPoliciesRoutes(
ctx context.Context,
peerID string,
peerSSHEnabled bool,
validatedPeersMap map[string]struct{},
postureFailedPeers *map[string]map[string]struct{},
) (map[string]*ComponentPeer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*ComponentPeer, len(a.Peers)/4)
relevantGroupIDs := make(map[string]*Group, len(a.Groups)/4)
relevantPolicies := make([]*Policy, 0, len(a.Policies))
relevantRoutes := make([]*route.Route, 0, len(a.Routes))
sshReqs := sshRequirements{neededGroupIDs: make(map[string]struct{})}
relevantPeerIDs[peerID] = a.GetPeer(peerID).ToComponent()
peerGroupSet := make(map[string]struct{}, 8)
for groupID, group := range a.Groups {
if slices.Contains(group.Peers, peerID) {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
peerGroupSet[groupID] = struct{}{}
}
}
routeAccessControlGroups := make(map[string]struct{})
for _, r := range a.Routes {
if r == nil {
continue
}
relevant := r.Peer == peerID
if !relevant {
for _, groupID := range r.PeerGroups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant && r.Enabled {
for _, groupID := range r.Groups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant {
continue
}
for _, groupID := range r.PeerGroups {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
}
for _, groupID := range r.Groups {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
}
if r.Enabled {
for _, groupID := range r.AccessControlGroups {
relevantGroupIDs[groupID] = a.GetGroup(groupID)
routeAccessControlGroups[groupID] = struct{}{}
}
}
// Include route advertisers in relevantPeerIDs. The envelope
// encoder writes route.peer_index by looking up r.Peer in the
// shipped peers list; if the advertiser is policy-isolated from
// the target peer (no rule edge between them), it would otherwise
// be omitted and the decoder would fail to resolve r.Peer, leaving
// the client without a WG tunnel target for this route. Legacy
// NetworkMap.Routes shipped the WG public key inline, so the
// equivalence path doesn't surface this — but the dependency is
// real once a client actually tries to use the route.
// Gate by validatedPeersMap so non-validated advertisers stay out
// (matches the network-resource router behaviour at the bottom of
// this loop, and the legacy invariant that only validated peers
// reach a client's view).
if r.Peer != "" {
if _, ok := validatedPeersMap[r.Peer]; ok {
if p := a.GetPeer(r.Peer); p != nil {
relevantPeerIDs[r.Peer] = p.ToComponent()
}
}
}
for _, groupID := range r.PeerGroups {
g := a.GetGroup(groupID)
if g == nil {
continue
}
for _, pid := range g.Peers {
if _, exists := relevantPeerIDs[pid]; exists {
continue
}
if _, ok := validatedPeersMap[pid]; !ok {
continue
}
if p := a.GetPeer(pid); p != nil {
relevantPeerIDs[pid] = p.ToComponent()
}
}
}
relevantRoutes = append(relevantRoutes, r)
}
for _, policy := range a.Policies {
if !policy.Enabled {
continue
}
policyRelevant := false
for _, rule := range policy.Rules {
if !rule.Enabled {
continue
}
if len(routeAccessControlGroups) > 0 {
for _, destGroupID := range rule.Destinations {
if _, needed := routeAccessControlGroups[destGroupID]; needed {
policyRelevant = true
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = a.GetGroup(srcGroupID)
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = a.GetGroup(dstGroupID)
}
break
}
}
}
var sourcePeers, destinationPeers []string
var peerInSources, peerInDestinations bool
if rule.SourceResource.Type == ResourceTypePeer && rule.SourceResource.ID != "" {
sourcePeers = []string{rule.SourceResource.ID}
if rule.SourceResource.ID == peerID {
peerInSources = true
}
} else {
sourcePeers, peerInSources = a.getPeersFromGroups(ctx, rule.Sources, peerID, policy.SourcePostureChecks, validatedPeersMap, postureFailedPeers)
}
if rule.DestinationResource.Type == ResourceTypePeer && rule.DestinationResource.ID != "" {
destinationPeers = []string{rule.DestinationResource.ID}
if rule.DestinationResource.ID == peerID {
peerInDestinations = true
}
} else {
destinationPeers, peerInDestinations = a.getPeersFromGroups(ctx, rule.Destinations, peerID, nil, validatedPeersMap, postureFailedPeers)
}
if peerInSources {
policyRelevant = true
for _, pid := range destinationPeers {
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = a.GetGroup(dstGroupID)
}
}
if peerInDestinations {
policyRelevant = true
for _, pid := range sourcePeers {
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = a.GetGroup(srcGroupID)
}
if rule.Protocol == PolicyRuleProtocolNetbirdSSH {
switch {
case len(rule.AuthorizedGroups) > 0:
for groupID := range rule.AuthorizedGroups {
sshReqs.neededGroupIDs[groupID] = struct{}{}
}
case rule.AuthorizedUser != "":
default:
sshReqs.needAllowedUserIDs = true
}
} else if PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
sshReqs.needAllowedUserIDs = true
}
}
}
if policyRelevant {
relevantPolicies = append(relevantPolicies, policy)
}
}
return relevantPeerIDs, relevantGroupIDs, relevantPolicies, relevantRoutes, sshReqs
}
func (a *Account) getPeersFromGroups(ctx context.Context, groups []string, peerID string, sourcePostureChecksIDs []string,
validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) ([]string, bool) {
peerInGroups := false
filteredPeerIDs := make([]string, 0, len(groups))
seenPeerIds := make(map[string]struct{}, len(groups))
for _, gid := range groups {
group := a.GetGroup(gid)
if group == nil {
continue
}
if group.IsGroupAll() || len(groups) == 1 {
filteredPeerIDs = make([]string, 0, len(group.Peers))
peerInGroups = false
for _, pid := range group.Peers {
peer, ok := a.Peers[pid]
if !ok || peer == nil {
continue
}
if _, ok := validatedPeersMap[peer.ID]; !ok {
continue
}
isValid, pname := a.validatePostureChecksOnPeerGetFailed(ctx, sourcePostureChecksIDs, peer.ID)
if !isValid && len(pname) > 0 {
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peer.ID] = struct{}{}
continue
}
if peer.ID == peerID {
peerInGroups = true
continue
}
filteredPeerIDs = append(filteredPeerIDs, peer.ID)
}
return filteredPeerIDs, peerInGroups
}
for _, pid := range group.Peers {
if _, seen := seenPeerIds[pid]; seen {
continue
}
seenPeerIds[pid] = struct{}{}
peer, ok := a.Peers[pid]
if !ok || peer == nil {
continue
}
if _, ok := validatedPeersMap[peer.ID]; !ok {
continue
}
isValid, pname := a.validatePostureChecksOnPeerGetFailed(ctx, sourcePostureChecksIDs, peer.ID)
if !isValid && len(pname) > 0 {
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peer.ID] = struct{}{}
continue
}
if peer.ID == peerID {
peerInGroups = true
continue
}
filteredPeerIDs = append(filteredPeerIDs, peer.ID)
}
}
return filteredPeerIDs, peerInGroups
}
func (a *Account) validatePostureChecksOnPeerGetFailed(ctx context.Context, sourcePostureChecksID []string, peerID string) (bool, string) {
peer, ok := a.Peers[peerID]
if !ok || peer == nil {
return false, ""
}
for _, postureChecksID := range sourcePostureChecksID {
postureChecks := a.GetPostureChecks(postureChecksID)
if postureChecks == nil {
continue
}
for _, check := range postureChecks.GetChecks() {
isValid, _ := check.Check(ctx, *peer)
if !isValid {
return false, postureChecksID
}
}
}
return true, ""
}
func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChecksIDs []string, validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) []string {
var dest []string
for _, peerID := range inputPeers {
if _, validated := validatedPeersMap[peerID]; !validated {
continue
}
valid, pname := a.validatePostureChecksOnPeerGetFailed(context.Background(), postureChecksIDs, peerID)
if valid {
dest = append(dest, peerID)
continue
}
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peerID] = struct{}{}
}
return dest
}
// filterGroupPeers trims each group's Peers slice to only those peers that
// also appear in `peers`. Groups whose filtered list is empty are NOT
// deleted from the map — they're kept so the components wire encoder can
// still resolve seq references from routes/policies/access-control groups
// that name them. Calculate() tolerates groups with empty Peers (the inner
// loops simply iterate zero times), so retaining them is behaviourally a
// no-op for the legacy path that consumes the same NetworkMapComponents.
func filterGroupPeers(groups *map[string]*ComponentGroup, peers map[string]*ComponentPeer) {
for groupID, groupInfo := range *groups {
filteredPeers := make([]string, 0, len(groupInfo.Peers))
for _, pid := range groupInfo.Peers {
if _, exists := peers[pid]; exists {
filteredPeers = append(filteredPeers, pid)
}
}
if len(filteredPeers) != len(groupInfo.Peers) {
ng := *groupInfo
ng.Peers = filteredPeers
(*groups)[groupID] = &ng
}
}
}
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*ComponentPeer) {
if len(*postureFailedPeers) == 0 {
return
}
referencedPostureChecks := make(map[string]struct{})
for _, policy := range policies {
for _, checkID := range policy.SourcePostureChecks {
referencedPostureChecks[checkID] = struct{}{}
}
}
for _, resPolicies := range resourcePoliciesMap {
for _, policy := range resPolicies {
for _, checkID := range policy.SourcePostureChecks {
referencedPostureChecks[checkID] = struct{}{}
}
}
}
for checkID, failedPeers := range *postureFailedPeers {
if _, referenced := referencedPostureChecks[checkID]; !referenced {
delete(*postureFailedPeers, checkID)
continue
}
for peerID := range failedPeers {
if _, exists := peers[peerID]; !exists {
delete(failedPeers, peerID)
}
}
if len(failedPeers) == 0 {
delete(*postureFailedPeers, checkID)
}
}
}
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*ComponentPeer, includeIPv6 bool) []nbdns.SimpleRecord {
if len(records) == 0 || len(peers) == 0 {
return nil
}
// Include both v4 and v6 addresses so AAAA records (whose RData is an IPv6
// address) are not filtered out when peers have IPv6 assigned. When the
// requesting peer doesn't have IPv6, omit v6 IPs so AAAA records get dropped.
peerIPs := make(map[string]struct{}, len(peers)*2)
for _, peer := range peers {
if peer == nil {
continue
}
peerIPs[peer.IP.String()] = struct{}{}
if includeIPv6 && peer.IPv6.IsValid() {
peerIPs[peer.IPv6.String()] = struct{}{}
}
}
filteredRecords := make([]nbdns.SimpleRecord, 0, len(records))
for _, record := range records {
if _, exists := peerIPs[record.RData]; exists {
filteredRecords = append(filteredRecords, record)
}
}
return filteredRecords
}
func filterGroupIDToUserIDs(fullMap map[string][]string, neededGroupIDs map[string]struct{}) map[string][]string {
if len(neededGroupIDs) == 0 {
return nil
}
filtered := make(map[string][]string, len(neededGroupIDs))
for groupID := range neededGroupIDs {
if users, ok := fullMap[groupID]; ok {
filtered[groupID] = users
}
}
return filtered
}

View File

@@ -1,515 +0,0 @@
package types
import (
"github.com/miekg/dns"
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/modules/zones"
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
// toNetworkMapData builds the slim twin store from the account once per
// account. The per-peer components calculation then runs on the twin.
func (a *Account) toNetworkMapData(
accountZones []*zones.Zone,
validatedPeersMap map[string]struct{},
resourcePolicies map[string][]*Policy,
routers map[string]map[string]*routerTypes.NetworkRouter,
groupIDToUserIDs map[string][]string,
) *networkmap.NetworkMapData {
nmd := &networkmap.NetworkMapData{
Peers: make(map[string]*nmdata.Peer, len(a.Peers)),
Groups: make(map[string]*nmdata.Group, len(a.Groups)),
Policies: make([]*nmdata.Policy, 0, len(a.Policies)),
Routes: make([]*nmdata.Route, 0, len(a.Routes)),
NameServerGroups: make([]*nmdata.NameServerGroup, 0, len(a.NameServerGroups)),
NetworkResources: make([]*nmdata.NetworkResource, 0, len(a.NetworkResources)),
PostureChecks: make(map[string]*nmdata.PostureChecks, len(a.PostureChecks)),
ResourcePolicies: make(map[string][]*nmdata.Policy, len(resourcePolicies)),
Routers: make(map[string]map[string]*nmdata.NetworkRouter, len(routers)),
ValidatedPeers: validatedPeersMap,
GroupIDToUserIDs: groupIDToUserIDs,
AllowedUserIDs: a.getAllowedUserIDs(),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
}
if a.Network != nil {
nmd.Network = TwinNetwork(a.Network)
}
nmd.DNSSettings = &nmdata.DNSSettings{DisabledManagementGroups: a.DNSSettings.DisabledManagementGroups}
if a.Settings != nil {
nmd.AccountSettings = &nmdata.AccountSettingsInfo{
PeerLoginExpirationEnabled: a.Settings.PeerLoginExpirationEnabled,
PeerLoginExpiration: a.Settings.PeerLoginExpiration,
PeerInactivityExpirationEnabled: a.Settings.PeerInactivityExpirationEnabled,
PeerInactivityExpiration: a.Settings.PeerInactivityExpiration,
}
}
for id, p := range a.Peers {
nmd.Peers[id] = twinPeer(p)
}
for id, g := range a.Groups {
nmd.Groups[id] = twinGroup(g)
}
policyCache := make(map[string]*nmdata.Policy, len(a.Policies))
twinPol := func(p *Policy) *nmdata.Policy {
if p == nil {
return nil
}
if tp, ok := policyCache[p.ID]; ok {
return tp
}
tp := twinPolicy(p)
policyCache[p.ID] = tp
return tp
}
for _, p := range a.Policies {
nmd.Policies = append(nmd.Policies, twinPol(p))
}
for resID, pols := range resourcePolicies {
twinPols := make([]*nmdata.Policy, 0, len(pols))
for _, p := range pols {
twinPols = append(twinPols, twinPol(p))
}
nmd.ResourcePolicies[resID] = twinPols
}
for _, r := range a.Routes {
if r == nil {
continue
}
nmd.Routes = append(nmd.Routes, twinRoute(r))
}
for _, nsg := range a.NameServerGroups {
nmd.NameServerGroups = append(nmd.NameServerGroups, twinNSG(nsg))
}
for _, res := range a.NetworkResources {
nmd.NetworkResources = append(nmd.NetworkResources, twinNetworkResource(res))
}
for _, pc := range a.PostureChecks {
if pc != nil {
nmd.PostureChecks[pc.ID] = twinPostureChecks(pc)
nmd.PostureCheckXIDToPublicID[pc.ID] = pc.PublicID
}
}
for _, n := range a.Networks {
if n != nil {
nmd.NetworkXIDToPublicID[n.ID] = n.PublicID
}
}
for networkID, inner := range routers {
twinInner := make(map[string]*nmdata.NetworkRouter, len(inner))
for peerID, router := range inner {
twinInner[peerID] = twinRouter(router)
}
nmd.Routers[networkID] = twinInner
}
nmd.AppliedZoneCandidates = buildAppliedZoneCandidates(accountZones)
nmd.PrivateServiceCandidates = a.buildPrivateServiceCandidates()
return nmd
}
func twinPeer(p *nbpeer.Peer) *nmdata.Peer {
if p == nil {
return nil
}
networkAddresses := make([]nmdata.NetworkAddress, 0, len(p.Meta.NetworkAddresses))
for _, na := range p.Meta.NetworkAddresses {
networkAddresses = append(networkAddresses, nmdata.NetworkAddress{NetIP: na.NetIP})
}
files := make([]nmdata.File, 0, len(p.Meta.Files))
for _, f := range p.Meta.Files {
files = append(files, nmdata.File{Path: f.Path, ProcessIsRunning: f.ProcessIsRunning})
}
return &nmdata.Peer{
ID: p.ID,
Key: p.Key,
SSHKey: p.SSHKey,
DNSLabel: p.DNSLabel,
UserID: p.UserID,
SSHEnabled: p.SSHEnabled,
LoginExpirationEnabled: p.LoginExpirationEnabled,
LastLogin: p.LastLogin,
IP: p.IP,
IPv6: p.IPv6,
Meta: nmdata.PeerSystemMeta{
WtVersion: p.Meta.WtVersion,
GoOS: p.Meta.GoOS,
OSVersion: p.Meta.OSVersion,
KernelVersion: p.Meta.KernelVersion,
NetworkAddresses: networkAddresses,
Files: files,
Capabilities: p.Meta.Capabilities,
Flags: nmdata.Flags{
ServerSSHAllowed: p.Meta.Flags.ServerSSHAllowed,
DisableIPv6: p.Meta.Flags.DisableIPv6,
},
},
Location: nmdata.PeerLocation{
CountryCode: p.Location.CountryCode,
CityName: p.Location.CityName,
ConnectionIP: p.Location.ConnectionIP,
},
}
}
// TwinPeer converts a real peer to its slim nmdata twin. Exported for the
// port-forwarding integration, which builds proxy NetworkMaps holding twins.
func TwinPeer(p *nbpeer.Peer) *nmdata.Peer {
return twinPeer(p)
}
func twinPeers(peers []*nbpeer.Peer) []*nmdata.Peer {
out := make([]*nmdata.Peer, len(peers))
for i, p := range peers {
out[i] = twinPeer(p)
}
return out
}
func twinGroup(g *Group) *nmdata.Group {
if g == nil {
return nil
}
return &nmdata.Group{
Name: g.Name,
PublicID: g.PublicID,
Peers: g.Peers,
}
}
func twinPolicy(p *Policy) *nmdata.Policy {
if p == nil {
return nil
}
rules := make([]*nmdata.PolicyRule, 0, len(p.Rules))
for _, r := range p.Rules {
rules = append(rules, twinRule(r))
}
return &nmdata.Policy{
ID: p.ID,
PublicID: p.PublicID,
Enabled: p.Enabled,
SourcePostureChecks: p.SourcePostureChecks,
Rules: rules,
}
}
func twinRule(r *PolicyRule) *nmdata.PolicyRule {
if r == nil {
return nil
}
var portRanges []nmdata.RulePortRange
if r.PortRanges != nil {
portRanges = make([]nmdata.RulePortRange, len(r.PortRanges))
for i, pr := range r.PortRanges {
portRanges[i] = nmdata.RulePortRange{Start: pr.Start, End: pr.End}
}
}
return &nmdata.PolicyRule{
ID: r.ID,
PolicyID: r.PolicyID,
Enabled: r.Enabled,
Action: string(r.Action),
Protocol: string(r.Protocol),
Bidirectional: r.Bidirectional,
Sources: r.Sources,
Destinations: r.Destinations,
SourceResource: nmdata.Resource{ID: r.SourceResource.ID, Type: string(r.SourceResource.Type)},
DestinationResource: nmdata.Resource{ID: r.DestinationResource.ID, Type: string(r.DestinationResource.Type)},
Ports: r.Ports,
PortRanges: portRanges,
AuthorizedGroups: r.AuthorizedGroups,
AuthorizedUser: r.AuthorizedUser,
}
}
func twinRoute(r *nbroute.Route) *nmdata.Route {
return &nmdata.Route{
ID: string(r.ID),
AccountID: r.AccountID,
PublicID: r.PublicID,
Network: r.Network,
Domains: r.Domains,
KeepRoute: r.KeepRoute,
NetID: string(r.NetID),
Description: r.Description,
Peer: r.Peer,
PeerID: r.PeerID,
PeerGroups: r.PeerGroups,
NetworkType: int(r.NetworkType),
Masquerade: r.Masquerade,
Metric: r.Metric,
Enabled: r.Enabled,
Groups: r.Groups,
AccessControlGroups: r.AccessControlGroups,
SkipAutoApply: r.SkipAutoApply,
}
}
// TwinRoute converts a real *route.Route to its slim nmdata twin. Exported for
// tests that assert against twin routes returned in a NetworkMap.
func TwinRoute(r *nbroute.Route) *nmdata.Route {
return twinRoute(r)
}
func twinNetworkResource(r *resourceTypes.NetworkResource) *nmdata.NetworkResource {
if r == nil {
return nil
}
return &nmdata.NetworkResource{
ID: r.ID,
NetworkID: r.NetworkID,
AccountID: r.AccountID,
PublicID: r.PublicID,
Name: r.Name,
Description: r.Description,
Type: string(r.Type),
Address: r.Address,
Domain: r.Domain,
Prefix: r.Prefix,
Enabled: r.Enabled,
}
}
func twinRouter(r *routerTypes.NetworkRouter) *nmdata.NetworkRouter {
if r == nil {
return nil
}
return &nmdata.NetworkRouter{
PublicID: r.PublicID,
PeerGroups: r.PeerGroups,
Masquerade: r.Masquerade,
Metric: r.Metric,
Enabled: r.Enabled,
}
}
func twinNSG(n *nbdns.NameServerGroup) *nmdata.NameServerGroup {
if n == nil {
return nil
}
nameServers := make([]nmdata.NameServer, 0, len(n.NameServers))
for _, ns := range n.NameServers {
nameServers = append(nameServers, nmdata.NameServer{
IP: ns.IP,
NSType: int(ns.NSType),
Port: ns.Port,
})
}
return &nmdata.NameServerGroup{
ID: n.ID,
PublicID: n.PublicID,
Name: n.Name,
Description: n.Description,
NameServers: nameServers,
Groups: n.Groups,
Primary: n.Primary,
Domains: n.Domains,
Enabled: n.Enabled,
SearchDomainsEnabled: n.SearchDomainsEnabled,
}
}
// TwinNetwork converts a real *Network to its slim twin. Exported for the
// graceful-degrade path that builds a minimal NetworkMapComponents directly.
func TwinNetwork(n *Network) *nmdata.Network {
nc := n.Copy()
return &nmdata.Network{
Identifier: nc.Identifier,
Net: nc.Net,
NetV6: nc.NetV6,
Dns: nc.Dns,
Serial: nc.Serial,
}
}
func twinPostureChecks(pc *posture.Checks) *nmdata.PostureChecks {
if pc == nil {
return nil
}
out := &nmdata.PostureChecks{ID: pc.ID}
def := pc.Checks
if def.NBVersionCheck != nil {
out.Checks.NBVersionCheck = &nmdata.NBVersionCheck{MinVersion: def.NBVersionCheck.MinVersion}
}
if def.OSVersionCheck != nil {
oc := &nmdata.OSVersionCheck{}
if def.OSVersionCheck.Android != nil {
oc.Android = &nmdata.MinVersionCheck{MinVersion: def.OSVersionCheck.Android.MinVersion}
}
if def.OSVersionCheck.Darwin != nil {
oc.Darwin = &nmdata.MinVersionCheck{MinVersion: def.OSVersionCheck.Darwin.MinVersion}
}
if def.OSVersionCheck.Ios != nil {
oc.Ios = &nmdata.MinVersionCheck{MinVersion: def.OSVersionCheck.Ios.MinVersion}
}
if def.OSVersionCheck.Linux != nil {
oc.Linux = &nmdata.MinKernelVersionCheck{MinKernelVersion: def.OSVersionCheck.Linux.MinKernelVersion}
}
if def.OSVersionCheck.Windows != nil {
oc.Windows = &nmdata.MinKernelVersionCheck{MinKernelVersion: def.OSVersionCheck.Windows.MinKernelVersion}
}
out.Checks.OSVersionCheck = oc
}
if def.GeoLocationCheck != nil {
gc := &nmdata.GeoLocationCheck{Action: def.GeoLocationCheck.Action}
for _, loc := range def.GeoLocationCheck.Locations {
gc.Locations = append(gc.Locations, nmdata.GeoLocation{CountryCode: loc.CountryCode, CityName: loc.CityName})
}
out.Checks.GeoLocationCheck = gc
}
if def.PeerNetworkRangeCheck != nil {
out.Checks.PeerNetworkRangeCheck = &nmdata.PeerNetworkRangeCheck{
Action: def.PeerNetworkRangeCheck.Action,
Ranges: def.PeerNetworkRangeCheck.Ranges,
}
}
if def.ProcessCheck != nil {
procs := make([]nmdata.Process, 0, len(def.ProcessCheck.Processes))
for _, p := range def.ProcessCheck.Processes {
procs = append(procs, nmdata.Process{LinuxPath: p.LinuxPath, MacPath: p.MacPath, WindowsPath: p.WindowsPath})
}
out.Checks.ProcessCheck = &nmdata.ProcessCheck{Processes: procs}
}
return out
}
// buildAppliedZoneCandidates precomputes the account-level custom DNS zones
// (record conversion) once; the per-peer distribution-group gate runs in the
// components calc. Mirrors the account-level half of filterPeerAppliedZones.
func buildAppliedZoneCandidates(accountZones []*zones.Zone) []networkmap.AppliedZoneCandidate {
var out []networkmap.AppliedZoneCandidate
for _, zone := range accountZones {
if !zone.Enabled || len(zone.Records) == 0 {
continue
}
simpleRecords := make([]nmdata.SimpleRecord, 0, len(zone.Records))
for _, record := range zone.Records {
var recordType int
rData := record.Content
switch record.Type {
case records.RecordTypeA:
recordType = int(dns.TypeA)
case records.RecordTypeAAAA:
recordType = int(dns.TypeAAAA)
case records.RecordTypeCNAME:
recordType = int(dns.TypeCNAME)
rData = dns.Fqdn(record.Content)
default:
continue
}
simpleRecords = append(simpleRecords, nmdata.SimpleRecord{
Name: dns.Fqdn(record.Name),
Type: recordType,
Class: nbdns.DefaultClass,
TTL: record.TTL,
RData: rData,
})
}
out = append(out, networkmap.AppliedZoneCandidate{
DistributionGroups: zone.DistributionGroups,
Zone: nmdata.CustomZone{
Domain: dns.Fqdn(zone.Domain),
Records: simpleRecords,
SearchDomainDisabled: !zone.EnableSearchDomain,
NonAuthoritative: true,
},
})
}
return out
}
// buildPrivateServiceCandidates precomputes the connected-proxy A records per
// private service (account-level); the per-peer access-group gate + apex merge
// run in the components calc. Mirrors the account-level half of
// SynthesizePrivateServiceZones.
func (a *Account) buildPrivateServiceCandidates() []networkmap.PrivateServiceCandidate {
if len(a.Services) == 0 {
return nil
}
proxyPeersByCluster := a.GetProxyPeers()
if len(proxyPeersByCluster) == 0 {
return nil
}
var out []networkmap.PrivateServiceCandidate
for _, svc := range a.Services {
if svc == nil || !svc.Enabled || !svc.Private {
continue
}
if len(svc.AccessGroups) == 0 {
continue
}
proxyPeers := proxyPeersByCluster[svc.ProxyCluster]
if len(proxyPeers) == 0 {
continue
}
apex := a.privateServiceDomainZone(svc)
if apex == "" {
continue
}
var recs []nmdata.SimpleRecord
for _, p := range proxyPeers {
if p == nil || !p.IP.IsValid() {
continue
}
if p.Status == nil || !p.Status.Connected {
continue
}
recs = append(recs, nmdata.SimpleRecord{
Name: dns.Fqdn(svc.Domain),
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: privateServiceDNSRecordTTL,
RData: p.IP.String(),
})
}
if len(recs) == 0 {
continue
}
out = append(out, networkmap.PrivateServiceCandidate{
AccessGroups: svc.AccessGroups,
Zone: nmdata.CustomZone{
Domain: dns.Fqdn(apex),
Records: recs,
NonAuthoritative: true,
SearchDomainDisabled: true,
},
})
}
return out
}
func toTwinCustomZone(z nbdns.CustomZone) nmdata.CustomZone {
records := make([]nmdata.SimpleRecord, 0, len(z.Records))
for _, r := range z.Records {
records = append(records, nmdata.SimpleRecord{
Name: r.Name,
Type: r.Type,
Class: r.Class,
TTL: r.TTL,
RData: r.RData,
})
}
return nmdata.CustomZone{
Domain: z.Domain,
Records: records,
SearchDomainDisabled: z.SearchDomainDisabled,
NonAuthoritative: z.NonAuthoritative,
}
}

View File

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

View File

@@ -13,6 +13,8 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/internals/modules/zones"
"github.com/netbirdio/netbird/management/internals/modules/zones/records"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -664,7 +666,7 @@ func Test_ExpandPortsAndRanges_SSHRuleExpansion(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer)
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer.ToComponent())
var ports []string
for _, fr := range result {
@@ -1038,6 +1040,518 @@ func Test_FilterZoneRecordsForPeers(t *testing.T) {
}
}
func Test_filterPeerAppliedZones(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
accountZones []*zones.Zone
peerGroups LookupMap
expected []nbdns.CustomZone
}{
{
name: "empty peer groups returns empty custom zones",
accountZones: []*zones.Zone{},
peerGroups: LookupMap{},
expected: []nbdns.CustomZone{},
},
{
name: "peer has access to zone with A record",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "example.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.example.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "example.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.example.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "peer has access to zone with search domain enabled",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "internal.local",
Enabled: true,
EnableSearchDomain: true,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "api.internal.local",
Type: records.RecordTypeA,
Content: "10.0.0.1",
TTL: 600,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "internal.local.",
Records: []nbdns.SimpleRecord{
{
Name: "api.internal.local.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 600,
RData: "10.0.0.1",
},
},
SearchDomainDisabled: false,
},
},
},
{
name: "peer has no access to zone",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "private.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group2"},
Records: []*records.Record{
{
ID: "record1",
Name: "secret.private.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{},
},
{
name: "disabled zone is filtered out",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "disabled.com",
Enabled: false,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.disabled.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{},
},
{
name: "zone with no records is filtered out",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "empty.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{},
},
{
name: "peer has access via multiple groups",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "multi.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1", "group2", "group3"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.multi.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group2": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "multi.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.multi.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "multiple zones with mixed access",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "allowed.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.allowed.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
{
ID: "zone2",
Domain: "denied.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group2"},
Records: []*records.Record{
{
ID: "record2",
Name: "www.denied.com",
Type: records.RecordTypeA,
Content: "192.168.1.2",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "allowed.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.allowed.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "zone with multiple record types",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "mixed.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.mixed.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
{
ID: "record2",
Name: "ipv6.mixed.com",
Type: records.RecordTypeAAAA,
Content: "2001:db8::1",
TTL: 600,
},
{
ID: "record3",
Name: "alias.mixed.com",
Type: records.RecordTypeCNAME,
Content: "www.mixed.com",
TTL: 900,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "mixed.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.mixed.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
{
Name: "ipv6.mixed.com.",
Type: int(dns.TypeAAAA),
Class: nbdns.DefaultClass,
TTL: 600,
RData: "2001:db8::1",
},
{
Name: "alias.mixed.com.",
Type: int(dns.TypeCNAME),
Class: nbdns.DefaultClass,
TTL: 900,
RData: "www.mixed.com.",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "multiple zones both accessible",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "first.com",
Enabled: true,
EnableSearchDomain: true,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.first.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
{
ID: "zone2",
Domain: "second.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record2",
Name: "www.second.com",
Type: records.RecordTypeA,
Content: "192.168.1.2",
TTL: 600,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "first.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.first.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
},
SearchDomainDisabled: false,
},
{
Domain: "second.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.second.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 600,
RData: "192.168.1.2",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "zone with multiple records of same type",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "multi-a.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.multi-a.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
{
ID: "record2",
Name: "www.multi-a.com",
Type: records.RecordTypeA,
Content: "192.168.1.2",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "multi-a.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.multi-a.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
{
Name: "www.multi-a.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.2",
},
},
SearchDomainDisabled: true,
},
},
},
{
name: "peer in multiple groups accessing different zones",
accountZones: []*zones.Zone{
{
ID: "zone1",
Domain: "zone1.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group1"},
Records: []*records.Record{
{
ID: "record1",
Name: "www.zone1.com",
Type: records.RecordTypeA,
Content: "192.168.1.1",
TTL: 300,
},
},
},
{
ID: "zone2",
Domain: "zone2.com",
Enabled: true,
EnableSearchDomain: false,
DistributionGroups: []string{"group2"},
Records: []*records.Record{
{
ID: "record2",
Name: "www.zone2.com",
Type: records.RecordTypeA,
Content: "192.168.1.2",
TTL: 300,
},
},
},
},
peerGroups: LookupMap{"group1": struct{}{}, "group2": struct{}{}},
expected: []nbdns.CustomZone{
{
Domain: "zone1.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.zone1.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.1",
},
},
SearchDomainDisabled: true,
},
{
Domain: "zone2.com.",
Records: []nbdns.SimpleRecord{
{
Name: "www.zone2.com.",
Type: int(dns.TypeA),
Class: nbdns.DefaultClass,
TTL: 300,
RData: "192.168.1.2",
},
},
SearchDomainDisabled: true,
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := filterPeerAppliedZones(ctx, tt.accountZones, tt.peerGroups)
require.Equal(t, len(tt.expected), len(result), "number of custom zones should match")
for i, expectedZone := range tt.expected {
assert.Equal(t, expectedZone.Domain, result[i].Domain, "domain should match")
assert.Equal(t, expectedZone.SearchDomainDisabled, result[i].SearchDomainDisabled, "search domain disabled flag should match")
assert.Equal(t, len(expectedZone.Records), len(result[i].Records), "number of records should match")
for j, expectedRecord := range expectedZone.Records {
assert.Equal(t, expectedRecord.Name, result[i].Records[j].Name, "record name should match")
assert.Equal(t, expectedRecord.Type, result[i].Records[j].Type, "record type should match")
assert.Equal(t, expectedRecord.Class, result[i].Records[j].Class, "record class should match")
assert.Equal(t, expectedRecord.TTL, result[i].Records[j].TTL, "record TTL should match")
assert.Equal(t, expectedRecord.RData, result[i].Records[j].RData, "record RData should match")
}
}
})
}
}
func TestInjectPrivateServicePolicies_ProxyPeerGetsInboundRule(t *testing.T) {
ctx := context.Background()

View File

@@ -6,7 +6,6 @@ import (
"net"
"net/netip"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
sharedtypes "github.com/netbirdio/netbird/shared/management/types"
)
@@ -18,9 +17,6 @@ type DNSSettings = sharedtypes.DNSSettings
type FirewallRule = sharedtypes.FirewallRule
type Group = sharedtypes.Group
type GroupPeer = sharedtypes.GroupPeer
type Network = sharedtypes.Network
type NetworkMap = sharedtypes.NetworkMap
type ForwardingRule = sharedtypes.ForwardingRule
@@ -42,6 +38,18 @@ type RouteFirewallRule = sharedtypes.RouteFirewallRule
type NetworkMapComponents = sharedtypes.NetworkMapComponents
type ComponentPeer = sharedtypes.ComponentPeer
type ComponentGroup = sharedtypes.ComponentGroup
type ComponentRouter = sharedtypes.ComponentRouter
type ComponentResource = sharedtypes.ComponentResource
type ComponentResourceType = sharedtypes.ComponentResourceType
const (
ComponentResourceHost = sharedtypes.ComponentResourceHost
ComponentResourceSubnet = sharedtypes.ComponentResourceSubnet
ComponentResourceDomain = sharedtypes.ComponentResourceDomain
)
var EmptyNetworkMapComponents = sharedtypes.EmptyNetworkMapComponents
type AccountSettingsInfo = sharedtypes.AccountSettingsInfo
@@ -52,12 +60,7 @@ type NetworkMapComponentsCompact = sharedtypes.NetworkMapComponentsCompact
type LookupMap = sharedtypes.LookupMap
type FirewallRuleContext = sharedtypes.FirewallRuleContext
const (
GroupIssuedAPI = sharedtypes.GroupIssuedAPI
GroupIssuedJWT = sharedtypes.GroupIssuedJWT
GroupIssuedIntegration = sharedtypes.GroupIssuedIntegration
GroupAllName = sharedtypes.GroupAllName
)
const GroupAllName = sharedtypes.GroupAllName
// Function forwarders preserve types.X(...) call sites that previously
// resolved to package-local funcs. Plain forwarders (not var aliases) keep
@@ -67,24 +70,20 @@ func PolicyRuleImpliesLegacySSH(rule *PolicyRule) bool {
return sharedtypes.PolicyRuleImpliesLegacySSH(rule)
}
// ExpandPortsAndRanges / AppendIPv6FirewallRule / GenerateRouteFirewallRules
// forward to the shared twin-typed helpers, converting the real types the
// legacy Account calc still uses to nmdata twins at this boundary.
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *nbpeer.Peer) []*FirewallRule {
return sharedtypes.ExpandPortsAndRanges(base, twinRule(rule), twinPeer(peer))
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *ComponentPeer) []*FirewallRule {
return sharedtypes.ExpandPortsAndRanges(base, rule, peer)
}
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *nbpeer.Peer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
return sharedtypes.AppendIPv6FirewallRule(rules, rulesExists, twinPeer(peer), twinPeer(targetPeer), twinRule(rule), rc)
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *ComponentPeer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
return sharedtypes.AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, rc)
}
func CalculateNetworkMapFromComponents(ctx context.Context, components *NetworkMapComponents) *NetworkMap {
return sharedtypes.CalculateNetworkMapFromComponents(ctx, components)
}
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int, includeIPv6 bool) []*RouteFirewallRule {
return sharedtypes.GenerateRouteFirewallRules(ctx, twinRoute(route), twinRule(rule), twinPeers(groupPeers), direction, includeIPv6)
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*ComponentPeer, direction int, includeIPv6 bool) []*RouteFirewallRule {
return sharedtypes.GenerateRouteFirewallRules(ctx, route, rule, groupPeers, direction, includeIPv6)
}
func AllocateIPv6Subnet(r *rand.Rand) net.IPNet {

View File

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

View File

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

View File

@@ -18,7 +18,6 @@ import (
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
func networkMapFromComponents(t *testing.T, account *types.Account, peerID string, validatedPeers map[string]struct{}) *types.NetworkMap {
@@ -50,7 +49,7 @@ func allPeersValidated(account *types.Account, excludePeerIDs ...string) map[str
return validated
}
func peerIDs(peers []*nmdata.Peer) []string {
func peerIDs(peers []*types.ComponentPeer) []string {
ids := make([]string, len(peers))
for i, p := range peers {
ids[i] = p.ID
@@ -626,7 +625,7 @@ func TestNetworkMapComponents_DomainNetworkResource(t *testing.T) {
var hasDomainRoute bool
for _, r := range nm.Routes {
if r.NetworkType == int(route.DomainNetwork) && len(r.Domains) > 0 && r.Domains[0].SafeString() == "api.example.com" {
if r.NetworkType == route.DomainNetwork && len(r.Domains) > 0 && r.Domains[0].SafeString() == "api.example.com" {
hasDomainRoute = true
}
}

View File

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

View File

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

View File

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

View File

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

View File

@@ -28,12 +28,12 @@ func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
"ministral-8b-latest", "mistral-embed",
},
"anthropic": {
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
},
// bedrock keys are the normalized ids the request parser emits.
"bedrock": {
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",

View File

@@ -180,6 +180,11 @@ anthropic:
output_per_1k: 0.050
cache_read_per_1k: 0.001
cache_creation_per_1k: 0.0125
claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025
@@ -236,6 +241,11 @@ bedrock:
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
anthropic.claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -5,15 +5,14 @@ import (
"fmt"
"net"
"net/netip"
"slices"
"strconv"
"time"
log "github.com/sirupsen/logrus"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbdns "github.com/netbirdio/netbird/dns"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/domain"
nmdata "github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/types"
)
@@ -36,28 +35,28 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
Network: decodeAccountNetwork(full.Network),
AccountSettings: decodeAccountSettings(full.AccountSettings),
CustomZoneDomain: full.CustomZoneDomain,
Peers: make(map[string]*nmdata.Peer, len(full.Peers)),
Groups: make(map[string]*nmdata.Group, len(full.Groups)),
Policies: make([]*nmdata.Policy, 0, len(full.Policies)),
Routes: make([]*nmdata.Route, 0, len(full.Routes)),
NameServerGroups: make([]*nmdata.NameServerGroup, 0, len(full.NameserverGroups)),
Peers: make(map[string]*types.ComponentPeer, len(full.Peers)),
Groups: make(map[string]*types.ComponentGroup, len(full.Groups)),
Policies: make([]*types.Policy, 0, len(full.Policies)),
Routes: make([]*nbroute.Route, 0, len(full.Routes)),
NameServerGroups: make([]*nbdns.NameServerGroup, 0, len(full.NameserverGroups)),
AllDNSRecords: decodeSimpleRecords(full.AllDnsRecords),
AccountZones: decodeCustomZones(full.AccountZones),
ResourcePoliciesMap: make(map[string][]*nmdata.Policy),
RoutersMap: make(map[string]map[string]*nmdata.NetworkRouter),
NetworkResources: make([]*nmdata.NetworkResource, 0, len(full.NetworkResources)),
RouterPeers: make(map[string]*nmdata.Peer),
ResourcePoliciesMap: make(map[string][]*types.Policy),
RoutersMap: make(map[string]map[string]*types.ComponentRouter),
NetworkResources: make([]*types.ComponentResource, 0, len(full.NetworkResources)),
RouterPeers: make(map[string]*types.ComponentPeer),
AllowedUserIDs: stringSliceToSet(full.AllowedUserIds),
PostureFailedPeers: make(map[string]map[string]struct{}, len(full.PostureFailedPeers)),
GroupIDToUserIDs: make(map[string][]string, len(full.GroupIdToUserIds)),
}
if full.DnsSettings != nil {
c.DNSSettings = &nmdata.DNSSettings{
c.DNSSettings = &types.DNSSettings{
DisabledManagementGroups: full.DnsSettings.DisabledManagementGroupIds,
}
} else {
c.DNSSettings = &nmdata.DNSSettings{}
c.DNSSettings = &types.DNSSettings{}
}
// Phase 1: peers. The envelope's peers slice is index-addressed on the
@@ -99,21 +98,10 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
log.WithField("peer idx", idx).Error("unrecognized peer idx during decoding")
}
}
fromCompactResources := func() []nmdata.Resource {
var toret []nmdata.Resource
for _, r := range gc.Resources {
toret = append(toret, resourceFromProto(r, peerIDByIndex))
}
return toret
}
group := &nmdata.Group{
PublicID: gc.Id,
Peers: peerIDs,
Resources: fromCompactResources(),
group := &types.ComponentGroup{
ID: groupID,
PublicID: gc.Id,
Peers: peerIDs,
}
if gc.IsAll {
group.Name = types.GroupAllName
@@ -123,7 +111,7 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
// Phase 3: policies (PolicyCompact = one rule per entry; current data
// model is 1 rule per policy).
policyByID := make(map[string]*nmdata.Policy, len(full.Policies))
policyByID := make(map[string]*types.Policy, len(full.Policies))
for i, pc := range full.Policies {
if pc == nil {
return nil, fmt.Errorf("invalid envelope: policies[%d] is nil", i)
@@ -160,7 +148,7 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
// Phase 7: routers_map (outer key = network seq id, inner key = peer-id
// reconstructed from peer_index). Synthesized network id is "net_<seq>".
for networkID, list := range full.RoutersMap {
inner := make(map[string]*nmdata.NetworkRouter, len(list.Entries))
inner := make(map[string]*types.ComponentRouter, len(list.Entries))
for _, entry := range list.Entries {
if !entry.PeerIndexSet {
continue
@@ -170,8 +158,10 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
continue
}
peerID := peerIDByIndex[entry.PeerIndex]
inner[peerID] = &nmdata.NetworkRouter{
inner[peerID] = &types.ComponentRouter{
NetworkID: networkID,
PublicID: entry.Id,
Peer: peerID,
PeerGroups: entry.PeerGroupIds,
Masquerade: entry.Masquerade,
Metric: int(entry.Metric),
@@ -190,7 +180,7 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
if len(ids.Ids) == 0 {
continue
}
policies := make([]*nmdata.Policy, 0, len(ids.Ids))
policies := make([]*types.Policy, 0, len(ids.Ids))
for _, id := range ids.Ids {
if p, ok := policyByID[id]; ok {
policies = append(policies, p)
@@ -203,15 +193,6 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
}
}
// Phase 8: rebuild resource_policies_map
for _, r := range c.NetworkResources {
policies := policiesForNetworkResource(r.ID, c.Policies, c.Groups)
if len(policies) == 0 {
continue
}
c.ResourcePoliciesMap[r.ID] = policies
}
// Phase 9: group_id_to_user_ids — wire keys are seq ids, synth to strings.
for groupId, list := range full.GroupIdToUserIds {
c.GroupIDToUserIDs[groupId] = append([]string(nil), list.UserIds...)
@@ -247,48 +228,11 @@ func DecodeEnvelope(env *proto.NetworkMapEnvelope) (*types.NetworkMapComponents,
return c, nil
}
func networkResourceGroups(resourceId string, groups map[string]*nmdata.Group) []string {
var toret []string
for _, group := range groups {
for _, resource := range group.Resources {
if resource.ID == resourceId {
toret = append(toret, group.PublicID)
}
}
}
return toret
}
func policiesForNetworkResource(resourceId string, allPolicies []*nmdata.Policy, groups map[string]*nmdata.Group) []*nmdata.Policy {
var toret []*nmdata.Policy
networkResourceGroups := networkResourceGroups(resourceId, groups)
for _, p := range allPolicies {
if p == nil || !p.Enabled {
continue
}
// there's always only one rule in each policy
if p.Rules[0].DestinationResource.ID == resourceId {
toret = append(toret, p)
continue
}
for _, groupId := range networkResourceGroups {
if slices.Contains(p.Rules[0].Destinations, groupId) {
toret = append(toret, p)
break
}
}
}
return toret
}
func decodeAccountNetwork(an *proto.AccountNetwork) *nmdata.Network {
func decodeAccountNetwork(an *proto.AccountNetwork) *types.Network {
if an == nil {
return nil
}
n := &nmdata.Network{
n := &types.Network{
Identifier: an.Identifier,
Dns: an.Dns,
Serial: an.Serial,
@@ -306,50 +250,32 @@ func decodeAccountNetwork(an *proto.AccountNetwork) *nmdata.Network {
return n
}
func decodeAccountSettings(as *proto.AccountSettingsCompact) *nmdata.AccountSettingsInfo {
func decodeAccountSettings(as *proto.AccountSettingsCompact) *types.AccountSettingsInfo {
if as == nil {
return &nmdata.AccountSettingsInfo{}
return &types.AccountSettingsInfo{}
}
return &nmdata.AccountSettingsInfo{
return &types.AccountSettingsInfo{
PeerLoginExpirationEnabled: as.PeerLoginExpirationEnabled,
PeerLoginExpiration: time.Duration(as.PeerLoginExpirationNs),
}
}
func decodePeerCompact(pc *proto.PeerCompact, peerID string) *nmdata.Peer {
var caps []int32
if pc.SupportsSourcePrefixes {
caps = append(caps, nbpeer.PeerCapabilitySourcePrefixes)
}
if pc.SupportsIpv6 {
caps = append(caps, nbpeer.PeerCapabilityIPv6Overlay)
}
peer := &nmdata.Peer{
func decodePeerCompact(pc *proto.PeerCompact, peerID string) *types.ComponentPeer {
peer := &types.ComponentPeer{
ID: peerID,
Key: peerID,
SSHKey: string(pc.SshPubKey),
SSHEnabled: pc.SshEnabled,
DNSLabel: pc.DnsLabel,
LoginExpirationEnabled: pc.LoginExpirationEnabled,
Meta: nmdata.PeerSystemMeta{
WtVersion: pc.AgentVersion,
Capabilities: caps,
Flags: nmdata.Flags{
ServerSSHAllowed: pc.ServerSshAllowed,
},
},
}
if pc.AddedWithSsoLogin {
// Set a non-empty UserID so (*Peer).AddedWithSSOLogin() returns true.
// The original UserID isn't on the wire; the value is intentionally
// visibly synthetic so any future consumer that mistakes UserID for a
// real account user xid won't silently match (or worse, write the
// sentinel into a downstream record).
peer.UserID = "<env-sso>"
AgentVersion: pc.AgentVersion,
SupportsSourcePrefixes: pc.SupportsSourcePrefixes,
SupportsIPv6: pc.SupportsIpv6,
ServerSSHAllowed: pc.ServerSshAllowed,
AddedWithSSOLogin: pc.AddedWithSsoLogin,
}
if pc.LastLoginUnixNano != 0 {
t := time.Unix(0, pc.LastLoginUnixNano)
peer.LastLogin = &t
peer.LastLogin = time.Unix(0, pc.LastLoginUnixNano)
}
switch len(pc.Ip) {
case 4:
@@ -367,13 +293,13 @@ func decodePeerCompact(pc *proto.PeerCompact, peerID string) *nmdata.Peer {
return peer
}
func decodePolicyCompact(pc *proto.PolicyCompact, policyID string, peerIDByIndex []string) *nmdata.Policy {
rule := &nmdata.PolicyRule{
func decodePolicyCompact(pc *proto.PolicyCompact, policyID string, peerIDByIndex []string) *types.Policy {
rule := &types.PolicyRule{
ID: policyID, // 1 rule per policy → reuse synthesized id
PolicyID: policyID,
Enabled: true,
Action: string(actionFromProto(pc.Action)),
Protocol: string(protocolFromProto(pc.Protocol)),
Action: actionFromProto(pc.Action),
Protocol: protocolFromProto(pc.Protocol),
Bidirectional: pc.Bidirectional,
Ports: uint32SliceToStrings(pc.Ports),
PortRanges: portRangesFromProto(pc.PortRanges),
@@ -384,11 +310,11 @@ func decodePolicyCompact(pc *proto.PolicyCompact, policyID string, peerIDByIndex
SourceResource: resourceFromProto(pc.SourceResource, peerIDByIndex),
DestinationResource: resourceFromProto(pc.DestinationResource, peerIDByIndex),
}
return &nmdata.Policy{
return &types.Policy{
ID: policyID,
PublicID: pc.Id,
Enabled: true,
Rules: []*nmdata.PolicyRule{rule},
Rules: []*types.PolicyRule{rule},
SourcePostureChecks: pc.SourcePostureCheckIds,
}
}
@@ -396,31 +322,15 @@ func decodePolicyCompact(pc *proto.PolicyCompact, policyID string, peerIDByIndex
// resourceFromProto rebuilds types.Resource. For peer-typed resources the
// peer reference is reconstructed from the envelope's peer index — wire
// format ships no xid for peers, so we use the synthesized peer id.
func resourceFromProto(r *proto.ResourceCompact, peerIDByIndex []string) nmdata.Resource {
func resourceFromProto(r *proto.ResourceCompact, peerIDByIndex []string) types.Resource {
if r == nil {
return nmdata.Resource{}
return types.Resource{}
}
t, ok := proto.ResourceCompactType_name[int32(r.Type)]
if !ok || r.Type == proto.ResourceCompactType_unknown_type {
return nmdata.Resource{}
}
if r.Type == proto.ResourceCompactType_peer && int(r.GetPeerIndex()) >= len(peerIDByIndex) {
return nmdata.Resource{}
}
if r.Type == proto.ResourceCompactType_peer && int(r.GetPeerIndex()) < len(peerIDByIndex) {
return nmdata.Resource{
Type: "peer",
ID: peerIDByIndex[int(r.GetPeerIndex())],
}
}
return nmdata.Resource{
Type: t,
ID: r.GetId(),
out := types.Resource{Type: types.ResourceType(r.Type)}
if r.PeerIndexSet && int(r.PeerIndex) < len(peerIDByIndex) {
out.ID = peerIDByIndex[r.PeerIndex]
}
return out
}
// authorizedGroupsFromProto inverts encodeAuthorizedGroups: the wire form
@@ -441,15 +351,15 @@ func authorizedGroupsFromProto(m map[string]*proto.UserNameList) map[string][]st
return out
}
func decodeRouteRaw(rr *proto.RouteRaw, peerIDByIndex []string) *nmdata.Route {
r := &nmdata.Route{
ID: rr.Id,
func decodeRouteRaw(rr *proto.RouteRaw, peerIDByIndex []string) *nbroute.Route {
r := &nbroute.Route{
ID: nbroute.ID(rr.Id),
PublicID: rr.Id,
NetID: rr.NetId,
NetID: nbroute.NetID(rr.NetId),
Description: rr.Description,
Domains: domainsFromPunycode(rr.Domains),
KeepRoute: rr.KeepRoute,
NetworkType: int(rr.NetworkType),
NetworkType: nbroute.NetworkType(rr.NetworkType),
Masquerade: rr.Masquerade,
Metric: int(rr.Metric),
Enabled: rr.Enabled,
@@ -469,8 +379,8 @@ func decodeRouteRaw(rr *proto.RouteRaw, peerIDByIndex []string) *nmdata.Route {
return r
}
func decodeNameServerGroupRaw(nsg *proto.NameServerGroupRaw) *nmdata.NameServerGroup {
out := &nmdata.NameServerGroup{
func decodeNameServerGroupRaw(nsg *proto.NameServerGroupRaw) *nbdns.NameServerGroup {
out := &nbdns.NameServerGroup{
ID: nsg.Id,
PublicID: nsg.Id,
Groups: nsg.GroupIds,
@@ -478,13 +388,13 @@ func decodeNameServerGroupRaw(nsg *proto.NameServerGroupRaw) *nmdata.NameServerG
Domains: nsg.Domains,
Enabled: nsg.Enabled,
SearchDomainsEnabled: nsg.SearchDomainsEnabled,
NameServers: make([]nmdata.NameServer, 0, len(nsg.Nameservers)),
NameServers: make([]nbdns.NameServer, 0, len(nsg.Nameservers)),
}
for _, ns := range nsg.Nameservers {
if addr, err := netip.ParseAddr(ns.IP); err == nil {
out.NameServers = append(out.NameServers, nmdata.NameServer{
out.NameServers = append(out.NameServers, nbdns.NameServer{
IP: addr,
NSType: int(ns.NSType),
NSType: nbdns.NameServerType(ns.NSType),
Port: int(ns.Port),
})
}
@@ -492,14 +402,14 @@ func decodeNameServerGroupRaw(nsg *proto.NameServerGroupRaw) *nmdata.NameServerG
return out
}
func decodeNetworkResource(nr *proto.NetworkResourceRaw) *nmdata.NetworkResource {
out := &nmdata.NetworkResource{
func decodeNetworkResource(nr *proto.NetworkResourceRaw) *types.ComponentResource {
out := &types.ComponentResource{
ID: nr.Id,
PublicID: nr.Id,
NetworkID: nr.NetworkSeq,
Name: nr.Name,
Description: nr.Description,
Type: nr.Type,
Type: types.ComponentResourceType(nr.Type),
Address: nr.Address,
Domain: nr.DomainValue,
Enabled: nr.Enabled,
@@ -512,10 +422,10 @@ func decodeNetworkResource(nr *proto.NetworkResourceRaw) *nmdata.NetworkResource
return out
}
func decodeSimpleRecords(records []*proto.SimpleRecord) []nmdata.SimpleRecord {
out := make([]nmdata.SimpleRecord, 0, len(records))
func decodeSimpleRecords(records []*proto.SimpleRecord) []nbdns.SimpleRecord {
out := make([]nbdns.SimpleRecord, 0, len(records))
for _, r := range records {
out = append(out, nmdata.SimpleRecord{
out = append(out, nbdns.SimpleRecord{
Name: r.Name,
Type: int(r.Type),
Class: r.Class,
@@ -526,10 +436,10 @@ func decodeSimpleRecords(records []*proto.SimpleRecord) []nmdata.SimpleRecord {
return out
}
func decodeCustomZones(zones []*proto.CustomZone) []nmdata.CustomZone {
out := make([]nmdata.CustomZone, 0, len(zones))
func decodeCustomZones(zones []*proto.CustomZone) []nbdns.CustomZone {
out := make([]nbdns.CustomZone, 0, len(zones))
for _, z := range zones {
out = append(out, nmdata.CustomZone{
out = append(out, nbdns.CustomZone{
Domain: z.Domain,
Records: decodeSimpleRecords(z.Records),
SearchDomainDisabled: z.SearchDomainDisabled,
@@ -550,16 +460,16 @@ func uint32SliceToStrings(ports []uint32) []string {
return out
}
func portRangesFromProto(ranges []*proto.PortInfo_Range) []nmdata.RulePortRange {
func portRangesFromProto(ranges []*proto.PortInfo_Range) []types.RulePortRange {
if len(ranges) == 0 {
return nil
}
out := make([]nmdata.RulePortRange, 0, len(ranges))
out := make([]types.RulePortRange, 0, len(ranges))
for _, r := range ranges {
if r == nil || r.Start > 65535 || r.End > 65535 {
continue
}
out = append(out, nmdata.RulePortRange{
out = append(out, types.RulePortRange{
Start: uint16(r.Start),
End: uint16(r.End),
})

View File

@@ -1,40 +0,0 @@
package networkmap
import (
"testing"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/stretchr/testify/assert"
)
func TestDecodePolicy(t *testing.T) {
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_peer, ResourceId: &proto.ResourceCompact_PeerIndex{PeerIndex: uint32(1)}},
[]string{"invalid-id-0", "valid-id", "invalid-id-2"}),
nmdata.Resource{Type: "peer", ID: "valid-id"})
// check invalid peer index returns an empty resource
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_peer, ResourceId: &proto.ResourceCompact_PeerIndex{PeerIndex: uint32(100)}},
[]string{"invalid-id-0", "valid-id", "invalid-id-2"}),
nmdata.Resource{})
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_domain, ResourceId: &proto.ResourceCompact_Id{Id: "domain"}}, []string{}),
nmdata.Resource{Type: "domain", ID: "domain"})
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_host, ResourceId: &proto.ResourceCompact_Id{Id: "host"}}, []string{}),
nmdata.Resource{Type: "host", ID: "host"})
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_subnet, ResourceId: &proto.ResourceCompact_Id{Id: "subnet"}}, []string{}),
nmdata.Resource{Type: "subnet", ID: "subnet"})
// an unknown resource type return an empty resource
assert.Equal(t,
resourceFromProto(
&proto.ResourceCompact{Type: proto.ResourceCompactType_unknown_type, ResourceId: &proto.ResourceCompact_Id{Id: "boom"}}, []string{}),
nmdata.Resource{})
}

View File

@@ -20,7 +20,7 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"net/netip"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/types"
"github.com/netbirdio/netbird/shared/netiputil"
@@ -28,7 +28,7 @@ import (
)
// ToProtocolRoutes converts a slice of typed routes to their proto form.
func ToProtocolRoutes(routes []*nmdata.Route) []*proto.Route {
func ToProtocolRoutes(routes []*nbroute.Route) []*proto.Route {
protoRoutes := make([]*proto.Route, 0, len(routes))
for _, r := range routes {
protoRoutes = append(protoRoutes, ToProtocolRoute(r))
@@ -37,7 +37,7 @@ func ToProtocolRoutes(routes []*nmdata.Route) []*proto.Route {
}
// ToProtocolRoute converts one typed route to its proto form.
func ToProtocolRoute(route *nmdata.Route) *proto.Route {
func ToProtocolRoute(route *nbroute.Route) *proto.Route {
return &proto.Route{
ID: string(route.ID),
NetID: string(route.NetID),
@@ -273,7 +273,7 @@ func ToProtocolDNSConfig(update nbdns.Config, cache DNSConfigCache, forwardPort
// AppendRemotePeerConfig appends typed peers as proto.RemotePeerConfig
// entries to dst and returns the result.
func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nmdata.Peer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*types.ComponentPeer, dnsName string, includeIPv6 bool) []*proto.RemotePeerConfig {
for _, rPeer := range peers {
allowedIPs := []string{rPeer.IP.String() + "/32"}
if includeIPv6 && rPeer.IPv6.IsValid() {
@@ -284,7 +284,7 @@ func AppendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nmdata.Peer,
AllowedIps: allowedIPs,
SshConfig: &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
Fqdn: rPeer.FQDN(dnsName),
AgentVersion: rPeer.Meta.WtVersion,
AgentVersion: rPeer.AgentVersion,
})
}
return dst

View File

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

View File

@@ -15,7 +15,6 @@ import (
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/types"
nbnetworkmap "github.com/netbirdio/netbird/shared/management/networkmap"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
"github.com/netbirdio/netbird/shared/management/proto"
)
@@ -56,13 +55,13 @@ func TestEnvelopeToNetworkMap_RoundTrip(t *testing.T) {
func TestCalculate_FirewallRuleProtocol_NeverNetbirdSSH(t *testing.T) {
c, localPeerKey := buildSmokeComponents(t)
// Replace the smoke policy with a NetbirdSSH-protocol allow.
c.Policies = []*nmdata.Policy{{
c.Policies = []*types.Policy{{
ID: "pol-ssh", PublicID: "2", Enabled: true,
Rules: []*nmdata.PolicyRule{{
Rules: []*types.PolicyRule{{
ID: "rule-ssh",
Enabled: true,
Action: string(types.PolicyTrafficActionAccept),
Protocol: string(types.PolicyRuleProtocolNetbirdSSH),
Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolNetbirdSSH,
Bidirectional: true,
Sources: []string{"group-all"},
Destinations: []string{"group-all"},
@@ -144,39 +143,39 @@ func TestDecodeEnvelope_MalformedWgKeyPeerSkipped(t *testing.T) {
func TestEnvelopeRoundTrip_AllGroupShortCircuitParity(t *testing.T) {
ctx := context.Background()
peers := map[string]*nmdata.Peer{}
peers := map[string]*types.ComponentPeer{}
for i, id := range []string{"peer-T", "peer-S", "peer-ALL", "peer-O"} {
peers[id] = &nmdata.Peer{
ID: id,
Key: randomWgKey(t),
IP: netip.AddrFrom4([4]byte{100, 64, 0, byte(i + 1)}),
DNSLabel: id,
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
peers[id] = &types.ComponentPeer{
ID: id,
Key: randomWgKey(t),
IP: netip.AddrFrom4([4]byte{100, 64, 0, byte(i + 1)}),
DNSLabel: id,
AgentVersion: "0.40.0",
}
}
c := &types.NetworkMapComponents{
PeerID: "peer-T",
Network: &nmdata.Network{
Network: &types.Network{
Identifier: "net-all-groups",
Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)},
Serial: 1,
},
AccountSettings: &nmdata.AccountSettingsInfo{},
DNSSettings: &nmdata.DNSSettings{},
AccountSettings: &types.AccountSettingsInfo{},
DNSSettings: &types.DNSSettings{},
Peers: peers,
Groups: map[string]*nmdata.Group{
"g-src": {PublicID: "1", Name: "staff", Peers: []string{"peer-T", "peer-S"}},
"g-all": {PublicID: "2", Name: "All", Peers: []string{"peer-ALL"}},
"g-two": {PublicID: "3", Name: "second", Peers: []string{"peer-T", "peer-O"}},
Groups: map[string]*types.ComponentGroup{
"g-src": {ID: "g-src", PublicID: "1", Name: "staff", Peers: []string{"peer-T", "peer-S"}},
"g-all": {ID: "g-all", PublicID: "2", Name: "All", Peers: []string{"peer-ALL"}},
"g-two": {ID: "g-two", PublicID: "3", Name: "second", Peers: []string{"peer-T", "peer-O"}},
},
Policies: []*nmdata.Policy{{
Policies: []*types.Policy{{
ID: "pol-multi-dest", PublicID: "10", Enabled: true,
Rules: []*nmdata.PolicyRule{{
Rules: []*types.PolicyRule{{
ID: "rule-multi-dest",
Enabled: true,
Action: string(types.PolicyTrafficActionAccept),
Protocol: string(types.PolicyRuleProtocolALL),
Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolALL,
Sources: []string{"g-src"},
Destinations: []string{"g-all", "g-two"},
}},
@@ -232,33 +231,33 @@ func buildSmokeComponents(t *testing.T) (*types.NetworkMapComponents, string) {
peerAKey := randomWgKey(t)
peerBKey := randomWgKey(t)
peerA := &nmdata.Peer{
ID: "peer-A",
Key: peerAKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peerA",
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
peerA := &types.ComponentPeer{
ID: "peer-A",
Key: peerAKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peerA",
AgentVersion: "0.40.0",
}
peerB := &nmdata.Peer{
ID: "peer-B",
Key: peerBKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
DNSLabel: "peerB",
Meta: nmdata.PeerSystemMeta{WtVersion: "0.40.0"},
peerB := &types.ComponentPeer{
ID: "peer-B",
Key: peerBKey,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
DNSLabel: "peerB",
AgentVersion: "0.40.0",
}
group := &nmdata.Group{
PublicID: "1", Name: "All",
group := &types.ComponentGroup{
ID: "group-all", PublicID: "1", Name: "All",
Peers: []string{"peer-A", "peer-B"},
}
policy := &nmdata.Policy{
policy := &types.Policy{
ID: "pol-allow", PublicID: "1", Enabled: true,
Rules: []*nmdata.PolicyRule{{
Rules: []*types.PolicyRule{{
ID: "rule-allow",
Enabled: true,
Action: string(types.PolicyTrafficActionAccept),
Protocol: string(types.PolicyRuleProtocolALL),
Action: types.PolicyTrafficActionAccept,
Protocol: types.PolicyRuleProtocolALL,
Bidirectional: true,
Sources: []string{"group-all"},
Destinations: []string{"group-all"},
@@ -267,21 +266,21 @@ func buildSmokeComponents(t *testing.T) (*types.NetworkMapComponents, string) {
c := &types.NetworkMapComponents{
PeerID: "peer-A",
Network: &nmdata.Network{
Network: &types.Network{
Identifier: "net-smoke",
Net: net.IPNet{IP: net.IP{100, 64, 0, 0}, Mask: net.CIDRMask(10, 32)},
Serial: 1,
},
AccountSettings: &nmdata.AccountSettingsInfo{},
DNSSettings: &nmdata.DNSSettings{},
Peers: map[string]*nmdata.Peer{
AccountSettings: &types.AccountSettingsInfo{},
DNSSettings: &types.DNSSettings{},
Peers: map[string]*types.ComponentPeer{
"peer-A": peerA,
"peer-B": peerB,
},
Groups: map[string]*nmdata.Group{
Groups: map[string]*types.ComponentGroup{
"group-all": group,
},
Policies: []*nmdata.Policy{policy},
Policies: []*types.Policy{policy},
}
return c, peerAKey
}

View File

@@ -1,659 +0,0 @@
package networkmap
import (
"slices"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
"github.com/netbirdio/netbird/shared/management/types"
)
type sshRequirements struct {
neededGroupIDs map[string]struct{}
needAllowedUserIDs bool
}
// GetPeerNetworkMapComponents computes the peer's NetworkMapComponents from the
// slim twin store. It mirrors the former Account.GetPeerNetworkMapComponents
// exactly, operating on nmdata twins throughout — no Account reference and no
// twin↔real conversion, since the produced components hold twins.
func (nmd *NetworkMapData) GetPeerNetworkMapComponents(peerID string, peersCustomZone nmdata.CustomZone) *types.NetworkMapComponents {
peer := nmd.Peers[peerID]
if peer == nil {
return types.EmptyNetworkMapComponents(&types.NetworkMapComponents{
PeerID: peerID,
Network: nmd.Network,
Peers: map[string]*nmdata.Peer{peerID: peer},
})
}
if _, ok := nmd.ValidatedPeers[peerID]; !ok {
return types.EmptyNetworkMapComponents(&types.NetworkMapComponents{
PeerID: peerID,
Network: nmd.Network,
Peers: map[string]*nmdata.Peer{peerID: peer},
})
}
components := &types.NetworkMapComponents{
PeerID: peerID,
Network: nmd.Network,
AccountSettings: nmd.AccountSettings,
DNSSettings: nmd.DNSSettings,
CustomZoneDomain: peersCustomZone.Domain,
NameServerGroups: make([]*nmdata.NameServerGroup, 0),
ResourcePoliciesMap: make(map[string][]*nmdata.Policy),
RoutersMap: make(map[string]map[string]*nmdata.NetworkRouter),
NetworkResources: make([]*nmdata.NetworkResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(nmd.PostureChecks)),
RouterPeers: make(map[string]*nmdata.Peer),
NetworkXIDToPublicID: nmd.NetworkXIDToPublicID,
PostureCheckXIDToPublicID: nmd.PostureCheckXIDToPublicID,
}
relevantPeers, relevantGroups, relevantPolicies, relevantRoutes, sshReqs := nmd.getPeersGroupsPoliciesRoutes(peerID, peer.SSHEnabled, &components.PostureFailedPeers)
if len(sshReqs.neededGroupIDs) > 0 {
components.GroupIDToUserIDs = filterGroupIDToUserIDs(nmd.GroupIDToUserIDs, sshReqs.neededGroupIDs)
}
if sshReqs.needAllowedUserIDs {
components.AllowedUserIDs = nmd.getAllowedUserIDs()
}
components.Peers = relevantPeers
components.Groups = relevantGroups
components.Policies = relevantPolicies
components.Routes = relevantRoutes
components.AllDNSRecords = filterDNSRecordsByPeers(peersCustomZone.Records, relevantPeers, peer.SupportsIPv6() && peer.IPv6.IsValid())
peerGroups := nmd.getPeerGroups(peerID)
components.AccountZones = nmd.appliedZones(peerGroups)
components.AccountZones = append(components.AccountZones, nmd.privateServiceZones(peerGroups)...)
for _, nsGroup := range nmd.NameServerGroups {
if nsGroup.Enabled {
for _, gID := range nsGroup.Groups {
if _, found := relevantGroups[gID]; found {
components.NameServerGroups = append(components.NameServerGroups, nsGroup)
break
}
}
}
}
for _, resource := range nmd.NetworkResources {
if !resource.Enabled {
continue
}
policies, exists := nmd.ResourcePolicies[resource.ID]
if !exists {
continue
}
addSourcePeers := false
networkRoutingPeers, routerExists := nmd.Routers[resource.NetworkID]
if routerExists {
if _, ok := networkRoutingPeers[peerID]; ok {
addSourcePeers = true
}
}
for _, policy := range policies {
if addSourcePeers {
var peers []string
if policy.Rules[0].SourceResource.Type == string(types.ResourceTypePeer) && policy.Rules[0].SourceResource.ID != "" {
peers = []string{policy.Rules[0].SourceResource.ID}
} else {
peers = nmd.getUniquePeerIDsFromGroupsIDs(policy.SourceGroups())
}
for _, pID := range nmd.getPostureValidPeersSaveFailed(peers, policy.SourcePostureChecks, &components.PostureFailedPeers) {
if _, exists := components.Peers[pID]; !exists {
components.Peers[pID] = nmd.Peers[pID]
}
}
} else {
peerInSources := false
if policy.Rules[0].SourceResource.Type == string(types.ResourceTypePeer) && policy.Rules[0].SourceResource.ID != "" {
peerInSources = policy.Rules[0].SourceResource.ID == peerID
} else {
for _, groupID := range policy.SourceGroups() {
if group := nmd.Groups[groupID]; group != nil && slices.Contains(group.Peers, peerID) {
peerInSources = true
break
}
}
}
if !peerInSources {
continue
}
isValid, pname := nmd.validatePostureChecksOnPeerGetFailed(policy.SourcePostureChecks, peerID)
if !isValid && len(pname) > 0 {
if _, ok := components.PostureFailedPeers[pname]; !ok {
components.PostureFailedPeers[pname] = make(map[string]struct{})
}
components.PostureFailedPeers[pname][peer.ID] = struct{}{}
continue
}
addSourcePeers = true
}
for _, rule := range policy.Rules {
for _, srcGroupID := range rule.Sources {
if g := nmd.Groups[srcGroupID]; g != nil {
if _, exists := components.Groups[srcGroupID]; !exists {
components.Groups[srcGroupID] = g
}
}
}
for _, dstGroupID := range rule.Destinations {
if g := nmd.Groups[dstGroupID]; g != nil {
if _, exists := components.Groups[dstGroupID]; !exists {
components.Groups[dstGroupID] = g
}
}
}
}
components.ResourcePoliciesMap[resource.ID] = policies
}
if addSourcePeers {
components.RoutersMap[resource.NetworkID] = networkRoutingPeers
for peerIDKey := range networkRoutingPeers {
if p := nmd.Peers[peerIDKey]; p != nil {
if _, exists := components.RouterPeers[peerIDKey]; !exists {
components.RouterPeers[peerIDKey] = p
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := nmd.ValidatedPeers[peerIDKey]; validated {
components.Peers[peerIDKey] = p
}
}
}
}
components.NetworkResources = append(components.NetworkResources, resource)
}
}
filterGroupPeers(&components.Groups, components.Peers)
filterPostureFailedPeers(&components.PostureFailedPeers, components.Policies, components.ResourcePoliciesMap, components.Peers)
return components
}
func (nmd *NetworkMapData) getPeersGroupsPoliciesRoutes(
peerID string,
peerSSHEnabled bool,
postureFailedPeers *map[string]map[string]struct{},
) (map[string]*nmdata.Peer, map[string]*nmdata.Group, []*nmdata.Policy, []*nmdata.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*nmdata.Peer, len(nmd.Peers)/4)
relevantGroupIDs := make(map[string]*nmdata.Group, len(nmd.Groups)/4)
relevantPolicies := make([]*nmdata.Policy, 0, len(nmd.Policies))
relevantRoutes := make([]*nmdata.Route, 0, len(nmd.Routes))
sshReqs := sshRequirements{neededGroupIDs: make(map[string]struct{})}
relevantPeerIDs[peerID] = nmd.Peers[peerID]
peerGroupSet := make(map[string]struct{}, 8)
for groupID, group := range nmd.Groups {
if slices.Contains(group.Peers, peerID) {
relevantGroupIDs[groupID] = group
peerGroupSet[groupID] = struct{}{}
}
}
routeAccessControlGroups := make(map[string]struct{})
for _, r := range nmd.Routes {
if r == nil {
continue
}
relevant := r.Peer == peerID
if !relevant {
for _, groupID := range r.PeerGroups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant && r.Enabled {
for _, groupID := range r.Groups {
if _, ok := peerGroupSet[groupID]; ok {
relevant = true
break
}
}
}
if !relevant {
continue
}
for _, groupID := range r.PeerGroups {
relevantGroupIDs[groupID] = nmd.Groups[groupID]
}
for _, groupID := range r.Groups {
relevantGroupIDs[groupID] = nmd.Groups[groupID]
}
if r.Enabled {
for _, groupID := range r.AccessControlGroups {
relevantGroupIDs[groupID] = nmd.Groups[groupID]
routeAccessControlGroups[groupID] = struct{}{}
}
}
if r.Peer != "" {
if _, ok := nmd.ValidatedPeers[r.Peer]; ok {
if p := nmd.Peers[r.Peer]; p != nil {
relevantPeerIDs[r.Peer] = p
}
}
}
for _, groupID := range r.PeerGroups {
g := nmd.Groups[groupID]
if g == nil {
continue
}
for _, pid := range g.Peers {
if _, exists := relevantPeerIDs[pid]; exists {
continue
}
if _, ok := nmd.ValidatedPeers[pid]; !ok {
continue
}
if p := nmd.Peers[pid]; p != nil {
relevantPeerIDs[pid] = p
}
}
}
relevantRoutes = append(relevantRoutes, r)
}
for _, policy := range nmd.Policies {
if !policy.Enabled {
continue
}
policyRelevant := false
for _, rule := range policy.Rules {
if !rule.Enabled {
continue
}
if len(routeAccessControlGroups) > 0 {
for _, destGroupID := range rule.Destinations {
if _, needed := routeAccessControlGroups[destGroupID]; needed {
policyRelevant = true
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = nmd.Groups[srcGroupID]
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = nmd.Groups[dstGroupID]
}
break
}
}
}
var sourcePeers, destinationPeers []string
var peerInSources, peerInDestinations bool
if rule.SourceResource.Type == string(types.ResourceTypePeer) && rule.SourceResource.ID != "" {
sourcePeers = []string{rule.SourceResource.ID}
if rule.SourceResource.ID == peerID {
peerInSources = true
}
} else {
sourcePeers, peerInSources = nmd.getPeersFromGroups(rule.Sources, peerID, policy.SourcePostureChecks, postureFailedPeers)
}
if rule.DestinationResource.Type == string(types.ResourceTypePeer) && rule.DestinationResource.ID != "" {
destinationPeers = []string{rule.DestinationResource.ID}
if rule.DestinationResource.ID == peerID {
peerInDestinations = true
}
} else {
destinationPeers, peerInDestinations = nmd.getPeersFromGroups(rule.Destinations, peerID, nil, postureFailedPeers)
}
if peerInSources {
policyRelevant = true
for _, pid := range destinationPeers {
relevantPeerIDs[pid] = nmd.Peers[pid]
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = nmd.Groups[dstGroupID]
}
}
if peerInDestinations {
policyRelevant = true
for _, pid := range sourcePeers {
relevantPeerIDs[pid] = nmd.Peers[pid]
}
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = nmd.Groups[srcGroupID]
}
if rule.Protocol == string(types.PolicyRuleProtocolNetbirdSSH) {
switch {
case len(rule.AuthorizedGroups) > 0:
for groupID := range rule.AuthorizedGroups {
sshReqs.neededGroupIDs[groupID] = struct{}{}
}
case rule.AuthorizedUser != "":
default:
sshReqs.needAllowedUserIDs = true
}
} else if nmdata.PolicyRuleImpliesLegacySSH(rule) && peerSSHEnabled {
sshReqs.needAllowedUserIDs = true
}
}
}
if policyRelevant {
relevantPolicies = append(relevantPolicies, policy)
}
}
return relevantPeerIDs, relevantGroupIDs, relevantPolicies, relevantRoutes, sshReqs
}
func (nmd *NetworkMapData) getPeersFromGroups(groups []string, peerID string, sourcePostureChecksIDs []string,
postureFailedPeers *map[string]map[string]struct{}) ([]string, bool) {
peerInGroups := false
filteredPeerIDs := make([]string, 0, len(groups))
seenPeerIds := make(map[string]struct{}, len(groups))
for _, gid := range groups {
group := nmd.Groups[gid]
if group == nil {
continue
}
if group.IsGroupAll() || len(groups) == 1 {
filteredPeerIDs = make([]string, 0, len(group.Peers))
peerInGroups = false
for _, pid := range group.Peers {
peer, ok := nmd.Peers[pid]
if !ok || peer == nil {
continue
}
if _, ok := nmd.ValidatedPeers[peer.ID]; !ok {
continue
}
isValid, pname := nmd.validatePostureChecksOnPeerGetFailed(sourcePostureChecksIDs, peer.ID)
if !isValid && len(pname) > 0 {
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peer.ID] = struct{}{}
continue
}
if peer.ID == peerID {
peerInGroups = true
continue
}
filteredPeerIDs = append(filteredPeerIDs, peer.ID)
}
return filteredPeerIDs, peerInGroups
}
for _, pid := range group.Peers {
if _, seen := seenPeerIds[pid]; seen {
continue
}
seenPeerIds[pid] = struct{}{}
peer, ok := nmd.Peers[pid]
if !ok || peer == nil {
continue
}
if _, ok := nmd.ValidatedPeers[peer.ID]; !ok {
continue
}
isValid, pname := nmd.validatePostureChecksOnPeerGetFailed(sourcePostureChecksIDs, peer.ID)
if !isValid && len(pname) > 0 {
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peer.ID] = struct{}{}
continue
}
if peer.ID == peerID {
peerInGroups = true
continue
}
filteredPeerIDs = append(filteredPeerIDs, peer.ID)
}
}
return filteredPeerIDs, peerInGroups
}
func (nmd *NetworkMapData) validatePostureChecksOnPeerGetFailed(sourcePostureChecksID []string, peerID string) (bool, string) {
peer, ok := nmd.Peers[peerID]
if !ok || peer == nil {
return false, ""
}
for _, postureChecksID := range sourcePostureChecksID {
postureChecks := nmd.PostureChecks[postureChecksID]
if postureChecks == nil {
continue
}
if !postureChecks.Passes(peer) {
return false, postureChecksID
}
}
return true, ""
}
func (nmd *NetworkMapData) getPostureValidPeersSaveFailed(inputPeers []string, postureChecksIDs []string, postureFailedPeers *map[string]map[string]struct{}) []string {
var dest []string
for _, peerID := range inputPeers {
if _, validated := nmd.ValidatedPeers[peerID]; !validated {
continue
}
valid, pname := nmd.validatePostureChecksOnPeerGetFailed(postureChecksIDs, peerID)
if valid {
dest = append(dest, peerID)
continue
}
if _, ok := (*postureFailedPeers)[pname]; !ok {
(*postureFailedPeers)[pname] = make(map[string]struct{})
}
(*postureFailedPeers)[pname][peerID] = struct{}{}
}
return dest
}
func (nmd *NetworkMapData) getPeerGroups(peerID string) map[string]struct{} {
groups := make(map[string]struct{})
for groupID, group := range nmd.Groups {
if slices.Contains(group.Peers, peerID) {
groups[groupID] = struct{}{}
}
}
return groups
}
func (nmd *NetworkMapData) getUniquePeerIDsFromGroupsIDs(groups []string) []string {
peerIDs := make(map[string]struct{}, len(groups))
for _, groupID := range groups {
group := nmd.Groups[groupID]
if group == nil {
continue
}
if group.IsGroupAll() || len(groups) == 1 {
return group.Peers
}
for _, peerID := range group.Peers {
peerIDs[peerID] = struct{}{}
}
}
ids := make([]string, 0, len(peerIDs))
for peerID := range peerIDs {
ids = append(ids, peerID)
}
return ids
}
func (nmd *NetworkMapData) getAllowedUserIDs() map[string]struct{} {
return nmd.AllowedUserIDs
}
func (nmd *NetworkMapData) appliedZones(peerGroups map[string]struct{}) []nmdata.CustomZone {
if len(peerGroups) == 0 {
return nil
}
var out []nmdata.CustomZone
for _, cand := range nmd.AppliedZoneCandidates {
if peerInDistributionGroups(peerGroups, cand.DistributionGroups) {
out = append(out, cand.Zone)
}
}
return out
}
func (nmd *NetworkMapData) privateServiceZones(peerGroups map[string]struct{}) []nmdata.CustomZone {
byApex := make(map[string]*nmdata.CustomZone)
var order []string
for _, cand := range nmd.PrivateServiceCandidates {
if !peerInDistributionGroups(peerGroups, cand.AccessGroups) {
continue
}
zone, exists := byApex[cand.Zone.Domain]
if !exists {
nz := nmdata.CustomZone{
Domain: cand.Zone.Domain,
SearchDomainDisabled: cand.Zone.SearchDomainDisabled,
NonAuthoritative: cand.Zone.NonAuthoritative,
}
byApex[cand.Zone.Domain] = &nz
zone = &nz
order = append(order, cand.Zone.Domain)
}
zone.Records = append(zone.Records, cand.Zone.Records...)
}
var out []nmdata.CustomZone
for _, apex := range order {
zone := byApex[apex]
if len(zone.Records) == 0 {
continue
}
out = append(out, *zone)
}
return out
}
func peerInDistributionGroups(peerGroups map[string]struct{}, groups []string) bool {
for _, g := range groups {
if _, ok := peerGroups[g]; ok {
return true
}
}
return false
}
func filterGroupPeers(groups *map[string]*nmdata.Group, peers map[string]*nmdata.Peer) {
for groupID, groupInfo := range *groups {
filteredPeers := make([]string, 0, len(groupInfo.Peers))
for _, pid := range groupInfo.Peers {
if _, exists := peers[pid]; exists {
filteredPeers = append(filteredPeers, pid)
}
}
if len(filteredPeers) != len(groupInfo.Peers) {
ng := groupInfo.Copy()
ng.Peers = filteredPeers
(*groups)[groupID] = ng
}
}
}
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*nmdata.Policy, resourcePoliciesMap map[string][]*nmdata.Policy, peers map[string]*nmdata.Peer) {
if len(*postureFailedPeers) == 0 {
return
}
referencedPostureChecks := make(map[string]struct{})
for _, policy := range policies {
for _, checkID := range policy.SourcePostureChecks {
referencedPostureChecks[checkID] = struct{}{}
}
}
for _, resPolicies := range resourcePoliciesMap {
for _, policy := range resPolicies {
for _, checkID := range policy.SourcePostureChecks {
referencedPostureChecks[checkID] = struct{}{}
}
}
}
for checkID, failedPeers := range *postureFailedPeers {
if _, referenced := referencedPostureChecks[checkID]; !referenced {
delete(*postureFailedPeers, checkID)
continue
}
for peerID := range failedPeers {
if _, exists := peers[peerID]; !exists {
delete(failedPeers, peerID)
}
}
if len(failedPeers) == 0 {
delete(*postureFailedPeers, checkID)
}
}
}
func filterDNSRecordsByPeers(records []nmdata.SimpleRecord, peers map[string]*nmdata.Peer, includeIPv6 bool) []nmdata.SimpleRecord {
if len(records) == 0 || len(peers) == 0 {
return nil
}
peerIPs := make(map[string]struct{}, len(peers)*2)
for _, peer := range peers {
if peer == nil {
continue
}
peerIPs[peer.IP.String()] = struct{}{}
if includeIPv6 && peer.IPv6.IsValid() {
peerIPs[peer.IPv6.String()] = struct{}{}
}
}
filteredRecords := make([]nmdata.SimpleRecord, 0, len(records))
for _, record := range records {
if _, exists := peerIPs[record.RData]; exists {
filteredRecords = append(filteredRecords, record)
}
}
return filteredRecords
}
func filterGroupIDToUserIDs(fullMap map[string][]string, neededGroupIDs map[string]struct{}) map[string][]string {
if len(neededGroupIDs) == 0 {
return nil
}
filtered := make(map[string][]string, len(neededGroupIDs))
for groupID := range neededGroupIDs {
if users, ok := fullMap[groupID]; ok {
filtered[groupID] = users
}
}
return filtered
}

View File

@@ -1,55 +0,0 @@
package networkmap
import (
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
)
// NetworkMapData is a dependency-light, slim twin of the server Account. It
// carries only the state GetPeerNetworkMapComponents needs, expressed in the
// fresh nmdata twin types. A builder converts an Account into a NetworkMapData
// once per account; the per-peer components calculation then runs on this twin
// with no reference back to the Account.
type NetworkMapData struct {
Peers map[string]*nmdata.Peer
Groups map[string]*nmdata.Group
Policies []*nmdata.Policy
Routes []*nmdata.Route
NameServerGroups []*nmdata.NameServerGroup
NetworkResources []*nmdata.NetworkResource
Network *nmdata.Network
DNSSettings *nmdata.DNSSettings
AccountSettings *nmdata.AccountSettingsInfo
PostureChecks map[string]*nmdata.PostureChecks
AllowedUserIDs map[string]struct{}
NetworkXIDToPublicID map[string]string
PostureCheckXIDToPublicID map[string]string
ValidatedPeers map[string]struct{}
ResourcePolicies map[string][]*nmdata.Policy
Routers map[string]map[string]*nmdata.NetworkRouter
GroupIDToUserIDs map[string][]string
DNSDomain string
AppliedZoneCandidates []AppliedZoneCandidate
PrivateServiceCandidates []PrivateServiceCandidate
}
// AppliedZoneCandidate is an account-level custom DNS zone reduced to the
// per-peer decision the components calc still makes: include the zone only when
// the peer belongs to one of its distribution groups. Record conversion is done
// once at build time.
type AppliedZoneCandidate struct {
DistributionGroups []string
Zone nmdata.CustomZone
}
// PrivateServiceCandidate is a single private service's synthesized records,
// carried per apex zone. The builder resolves proxy-cluster connectivity and
// domain-suffix matching once; the calc merges the candidates whose AccessGroups
// the peer belongs to, grouped by Zone.Domain.
type PrivateServiceCandidate struct {
AccessGroups []string
Zone nmdata.CustomZone
}

View File

@@ -1,11 +0,0 @@
package nmdata
import "time"
// AccountSettingsInfo is the slim twin of types.AccountSettingsInfo.
type AccountSettingsInfo struct {
PeerLoginExpirationEnabled bool
PeerLoginExpiration time.Duration
PeerInactivityExpirationEnabled bool
PeerInactivityExpiration time.Duration
}

View File

@@ -1,18 +0,0 @@
package nmdata
// SimpleRecord is the slim twin of dns.SimpleRecord.
type SimpleRecord struct {
Name string
Type int
Class string
TTL int
RData string
}
// CustomZone is the slim twin of dns.CustomZone.
type CustomZone struct {
Domain string
Records []SimpleRecord
SearchDomainDisabled bool
NonAuthoritative bool
}

View File

@@ -1,6 +0,0 @@
package nmdata
// DNSSettings is the slim twin of types.DNSSettings.
type DNSSettings struct {
DisabledManagementGroups []string
}

View File

@@ -1,25 +0,0 @@
package nmdata
import "slices"
const groupAllName = "All"
// Group is the slim twin of types.Group.
type Group struct {
Name string
PublicID string
Peers []string
Resources []Resource
}
func (g *Group) IsGroupAll() bool {
return g.Name == groupAllName
}
func (g *Group) Copy() *Group {
return &Group{
Name: g.Name,
PublicID: g.PublicID,
Peers: slices.Clone(g.Peers),
}
}

View File

@@ -1,24 +0,0 @@
package nmdata
import "net/netip"
// NameServerGroup is the slim twin of dns.NameServerGroup.
type NameServerGroup struct {
ID string
PublicID string
Name string
Description string
NameServers []NameServer
Groups []string
Primary bool
Domains []string
Enabled bool
SearchDomainsEnabled bool
}
// NameServer is the slim twin of dns.NameServer.
type NameServer struct {
IP netip.Addr
NSType int
Port int
}

View File

@@ -1,16 +0,0 @@
package nmdata
import "net"
// Network is the slim twin of types.Network.
type Network struct {
Identifier string
Net net.IPNet
NetV6 net.IPNet
Dns string
Serial uint64
}
func (n *Network) CurrentSerial() uint64 {
return n.Serial
}

View File

@@ -1,18 +0,0 @@
package nmdata
import "net/netip"
// NetworkResource is the slim twin of resources/types.NetworkResource.
type NetworkResource struct {
ID string
NetworkID string
AccountID string
PublicID string
Name string
Description string
Type string
Address string
Domain string
Prefix netip.Prefix
Enabled bool
}

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