Compare commits

...

19 Commits

Author SHA1 Message Date
Pascal Fischer
dd2bdc0de3 [management] explicit accountID check when deleting a user (#6944) 2026-07-28 18:34:55 +02:00
Zoltan Papp
44fef45c2f [client] Export peer details for Android (#6925)
## Describe your changes

Eexport peer details for Android

## Issue ticket number and link

## 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)
- [ ] 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/6925"><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=1787768377&installation_model_id=427504&pr_number=6925&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6925&signature=b708ba5d6f9b45ae48edf9597aad8aecac9a93c75b2a5a033e14c1c945c47f1f"><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

* **New Features**
* Enhanced Android peer details with latency, data transfer totals,
connection timestamps, handshake information, relay and security status,
and ICE connection details.
* Added formatted latency values to support clearer connection-quality
indicators.
* **Bug Fixes**
* Peer lists now refresh WireGuard statistics before displaying transfer
and handshake information, ensuring current values are shown.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 18:13:22 +02:00
Riccardo Manfrin
3d1f209ea3 [client] parse NB_LAZY_CONN_INACTIVITY_THRESHOLD as a Go duration (#6947)
## Describe your changes

`NB_LAZY_CONN_INACTIVITY_THRESHOLD` was parsed with `strconv.Atoi`, i.e.
as a bare
integer number of minutes. The documentation, however, states it takes a
Go duration
(e.g. `30m`, `1h`). As a result any documented value such as `30m` or
`5m` failed to
parse and **silently fell back to the 15m default**, so the setting
appeared to have
no effect.

`inactivityThresholdEnv()` now parses the value with
`time.ParseDuration`, matching
the docs. A bare integer is still accepted as a number of minutes for
backwards
compatibility, and an unparseable value logs a warning and falls back to
the default.

Added `TestInactivityThresholdEnv` covering Go-duration values
(`30m`/`1h`/`90s`), the
bare-integer minutes fallback, and zero/negative/garbage inputs.

## Issue ticket number and link

N/A
2026-07-28 18:12:49 +02:00
Zoltan Papp
2ef457be95 [client] Unify route selection in the route manager (#6928)
## Describe your changes

Move route select/deselect handling from the daemon server into exported
routemanager methods (SelectRoutes, DeselectRoutes, SelectAllRoutes,
DeselectAllRoutes) so every consumer shares one implementation: v4/v6
exit-pair expansion, exit-node mutual exclusion, and selection
triggering.

Previously the exit-node exclusivity lived only in the daemon's
SelectNetworks RPC, so the Android and iOS bindings could leave two exit
nodes selected until the next network map reconciliation. Both bindings
now call the shared manager methods and enforce exclusivity at toggle
time, matching the desktop behavior.

## Issue ticket number and link

## 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)
- [ ] 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/6928"><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=1787772098&installation_model_id=427504&pr_number=6928&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6928&signature=31ad59e1483e1582cd447a8db2fe21e5309230e631cbd0cad0f977cd15fb7b9b"><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

* **New Features**
* Route selection/deselection is now handled through shared
route-manager APIs for both individual routes and “all routes”.
* Exit-node selections automatically enforce mutual exclusivity while
keeping non-exit routes unaffected.
* **Bug Fixes**
* Unknown or unavailable route IDs now return errors, and exclusivity is
preserved even when some route IDs fail.
* **Tests**
* Added route-selection tests covering exclusivity (including
IPv4/IPv6), select-all behavior, partial errors, and invalid IDs.
* **Refactor / Chores**
* Simplified Android, iOS, and server routing flows to delegate to the
shared manager; updated mocks and removed redundant routing command
logic/dependencies.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 17:39:37 +02:00
Zoltan Papp
63c320b6a9 [client] Serialize iOS tunnel reconfiguration callbacks (#6870)
On iOS the tunnel reconfiguration (setTunnelNetworkSettings) was driven
from three independent Go paths without serialization: route prefix
updates via the route notifier's own delivery loop, interface IP/IPv6
set synchronously from the engine start goroutine, and DNS config
applied from the DNS apply chain through a separate Swift object. The
three sources mutated the shared Swift settings-manager state
concurrently and triggered overlapping updateTunnel() calls, losing or
half-applying route and DNS settings.

Introduce client/internal/tunnelnotifier: a single notifier that
implements both listener.NetworkChangeListener and dns.IosDnsManager,
queues all four callbacks (OnNetworkChanged, SetInterfaceIP,
SetInterfaceIPv6, ApplyDns) in one FIFO and delivers them one-by-one
from a single goroutine, so calls into Swift never overlap and arrive in
order. RunOniOS wraps the two Swift objects into the notifier and closes
it after the run loop exits.

The route notifier keeps its prefix dedup but delegates delivery to the
shared notifier instead of maintaining its own queue and loop; the dedup
check and the enqueue stay under one mutex so queue order matches
state-update order. Setting the interface IP becomes asynchronous, which
is safe: on iOS wgInterface.Create() only uses the TunFd, and the FIFO
preserves the IP -> routes/DNS relative order.

The package is build-tag free so the unit tests run on any platform
under -race. Android is unaffected: it receives all settings atomically
in one configureInterface call and serializes TUN rebuilds on a single
handler thread.

## Describe your changes

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [x] 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/6870"><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=1787404610&installation_model_id=427504&pr_number=6870&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6870&signature=7db66c4ec4262cc21137b3d3e1fd5e905950eab82ef46236a13ac3d512ff2a16"><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

* **New Features**
* Improved iOS tunnel networking by coordinating network, interface,
route, and DNS updates through a unified notifier.
* Updated iOS behavior so route/prefix changes are applied immediately
instead of via queued/background delivery.

* **Bug Fixes**
* Improved reliability by ensuring network and DNS-related callbacks are
invoked in order without overlap.
* Ensured pending updates are drained and handled correctly during
shutdown.

* **Tests**
* Added coverage for FIFO ordering, non-overlapping callback execution,
interleaved DNS/route updates, and graceful shutdown behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 16:51:01 +02:00
Stefan Fast
4acbe2670a [client] Escape dots in interface names for sysctl configuration (#6930) 2026-07-28 16:24:30 +02:00
Zoltan Papp
0fb4c8c423 [client] Build UI release binaries with the production tag (#6898)
## Describe your changes

The goreleaser UI configs passed no build tags, so released netbird-ui
binaries were built on the Wails !production path: DevTools enabled,
browser context menu forced on, and FRONTEND_DEVSERVER_URL still
honored.

## Issue ticket number and link

## 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)
- [ ] 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/6898"><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=1787616895&installation_model_id=427504&pr_number=6898&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6898&signature=c9cda21bd668dc58f307e4ef762a194fb28fdb5698786930d911a55811a2bcaf"><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

* **Chores**
* Updated release build configurations to mark Linux, Windows, and macOS
UI builds as production releases.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 16:23:09 +02:00
Zoltan Papp
42e45ff9f9 [client] Expose RenameProfile in the Android profile manager binding (#6926)
## Describe your changes
Expose RenameProfile in the Android profile manager binding

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [ ] Is it a bug fix
- [ ] Is a typo/documentation fix
- [x] Is a feature enhancement
- [ ] 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/6926"><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=1787768395&installation_model_id=427504&pr_number=6926&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6926&signature=f4684572fa5b03325c89a2ace1e1eed03df3881365acebdc31066b74fa656bdb"><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

* **New Features**
  * Added the ability to rename Android profiles.
  * Rename operations now provide clear success or failure feedback.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 16:22:48 +02:00
Viktor Liu
9269b56386 [management] Read reverse-proxy service and target columns in Postgres path (#6886) 2026-07-28 14:46:19 +02:00
Viktor Liu
b3f9b82442 [management] Force routing-peer DNS resolution for reverse-proxy domain targets (#6872) 2026-07-28 14:45:57 +02:00
Riccardo Manfrin
8a43f4f943 [client] fix build: add ReapplyMatching to dedicated AllowedIPsRefCounter (#6935)
## Describe your changes

#6799 turned `AllowedIPsRefCounter` from an alias of the generic
`Counter` into a dedicated peer-aware type. A change merged in parallel
— `DefaultManager.ReconcilePeerAllowedIPs` (lazy-connection idle→wake
reconvergence) — calls `allowedIPsRefCounter.ReapplyMatching`, which
only existed on the generic `Counter`. Each PR built alone; the merged
`main` did not:

```
client/internal/routemanager/manager.go: m.allowedIPsRefCounter.ReapplyMatching undefined
```

Add `ReapplyMatching(pred, apply)` to the dedicated type, matching the
generic contract (keyed on the active/`Out` peer): it re-applies every
prefix whose currently installed peer satisfies `pred`, skipping
prefixes with no active peer (reconciled by the next
Increment/Decrement). Also update `reconcile_test.go` to construct the
counter via `refcounter.NewAllowedIPs` — the generic `refcounter.New`
return value is no longer assignable to the dedicated type.

Verified with `cd client && CGO_ENABLED=1 go build .` (the failing CI
step) and the `TestReconcilePeerAllowedIPs` / refcounter / routemanager
tests.

## Issue ticket number and link

No public issue — fixes a `main` build break from a semantic merge
conflict between #6799 and the `ReconcilePeerAllowedIPs` change. Failing
run:
https://github.com/netbirdio/netbird/actions/runs/30330882389/job/90185591366

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

Internal build fix restoring a method on the AllowedIPs refcounter. 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/6935"><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=1787814246&installation_model_id=427504&pr_number=6935&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6935&signature=6a0f08ffe5ed39096a77d94e71a5947f68bad6f2b7abee5f1a4823fd44f35106"><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**
* Improved route reconciliation so all applicable allowed IP prefixes
are reliably re-applied for the correct peer.
* Prevented reconciliation from affecting routes assigned to other
peers.
* Improved handling of errors encountered while restoring multiple
routes, providing more consistent results.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 16:42:57 +09:00
Riccardo Manfrin
2f268c8141 [client] fix stale routing peer on overlapping-prefix network removal (#6799)
## Describe your changes

The AllowedIPs reference counter
([refcounter/types.go#L9](https://github.com/netbirdio/netbird/blob/e1a24376a/client/internal/routemanager/refcounter/types.go#L9))
was keyed only by prefix and stored a single active peer set by
the first registrar, never swapped. When two networks advertised the
same prefix via different routing peers, removing the one whose peer was
installed in WireGuard left the prefix pointing at the removed peer
instead of the
surviving one — traffic kept flowing to the old peer until a manual
`netbird down/up`.

Made the AllowedIPs counter peer-aware: it tracks a per-peer reference
count per prefix plus the installed peer, and swaps WireGuard to a
surviving peer when the active one releases its last reference (removes
the prefix when
none remain). `Decrement` now takes the peer key so the exact
incremented peer is released; the static handler records its selected
routing peer like the dynamic and DNS handlers already did. The generic
`Counter` (routes,
  exclusion, ipset) is unchanged.

  ## Issue ticket number and link

No public issue — reported internally (routes not updating without
`netbird down/up` when two networks share a subnet). Root cause is the
prefix-only key at

[refcounter/types.go#L9](https://github.com/netbirdio/netbird/blob/e1a24376a/client/internal/routemanager/refcounter/types.go#L9).

  ## Stack

  <!-- branch-stack -->

  ### Checklist
  - [x] Is it a bug fix
  - [ ] Is a typo/documentation fix
  - [ ] Is a feature enhancement
  - [ ] It is a refactor
  - [x] 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)

Internal client-side routing fix. No public API, CLI, or config change —
only the WireGuard AllowedIPs hand-off when overlapping-prefix networks
are removed.

  ### 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/6799"><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 Codesmith"
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=1786796229&installation_id=146802194&pr_number=6799&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6799&signature=322d6950b4f664b1cb3421f2efa039fbbb7946de9b79f2a1b1b131dc182ada2c"><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 Codesmith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>/codesmith</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**
* Improved routing behavior when multiple peers share the same Allowed
IP by making Allowed IP reference tracking peer-aware.
* Allowed IPs now correctly decrement using the active peer key and
transfer to another surviving active peer when the current peer is
removed.
* Prevented stale routing and incorrect reference cleanup during route
and DNS-driven teardown.

* **Tests**
* Added/extended coverage for peer handoffs, repeated references,
non-active peer removal, flushing behavior, and self-healing after swap
add/remove failures.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 14:10:58 +09:00
Zoltan Papp
e3c4128164 [client] Exit GUI immediately on Windows session end (#6878)
Wails v3 runs the full app teardown synchronously inside WM_ENDSESSION,
overrunning the 5s end-session budget and triggering the "app is
preventing shutdown" screen with a forced kill. Intercept
WM_QUERYENDSESSION/WM_ENDSESSION and exit at once instead, and suppress
error dialogs, toasts, and the hide-on-close hooks once shutdown or a
tray quit has begun.

## Describe your changes

## Issue ticket number and link

## 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)
- [ ] 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/6878"><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=1787474408&installation_model_id=427504&pr_number=6878&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6878&signature=f09980d83ad32e37b92aa30a7d4552c402a140dbd4e4662a3ca12d1868e9f04b"><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**
* Improved shutdown handling on Windows by properly responding to system
end-session requests.
* Prevented the main window, settings window, and error dialogs from
reopening or interfering while the app is closing.
* Updated tray “Quit” flow to begin shutdown immediately, ensuring
active profile/connection operations complete cleanly.
* Suppressed UI notifications during shutdown to avoid stray messages
after exit starts.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-28 14:09:16 +09:00
Maycon Santos
bab5572a74 [management, proxy] scope agent-network model allowlist per policy/group and provider (#6905)
## Describe your changes

The model-allowlist guardrail was merged into one account-wide union and
enforced flat on every request, ignoring which policy/group/provider
authorised it. With multiple policies — especially a mix of guardrailed
and un-guardrailed ones — this caused:

- **false-allow**: a model allowlisted for one group/provider leaked to
any caller; and
- **false-deny**: an un-guardrailed policy (intended unrestricted) was
blocked by another policy's allowlist.

Enforcement is now policy/group-aware, mirroring `llm_limit_check`:

- **Management (`SelectPolicyForRequest`) is authoritative.** It uses
the request model (already carried in
`CheckLLMPolicyLimitsRequest.model`, previously ignored) to keep only
applicable policies whose guardrails permit the model; no
allowlist-enabled guardrail = unrestricted. Denies
`llm_policy.model_blocked` when policies govern the (provider, groups)
but none permits the model.
- **Proxy `llm_guardrail` becomes a per-provider fail-closed backstop.**
The synthesiser emits an allowlist only for providers every authorising
policy restricts; the middleware keys off the resolved provider id and
keeps unknown-model fail-closed.
2026-07-27 20:43:59 +02:00
Zoltan Papp
9b4a5df925 [client] Use platform installer URL for manual update downloads (#6922)
The Wails UI regressed the non-enforced update download to the plain
GitHub releases page. Restore the old Fyne behavior: the tray update
item and the About card's Get installer button now open
version.DownloadUrl(), which points to the direct installer download
(pkgs.netbird.io) per OS/arch and falls back to the generic install page
where no installer exists.

Also fix the dead brew detection on darwin: exec.Command passed the
whole pipeline as a single argument to brew, so the check always failed.
Query the netbird formula and netbird-ui cask explicitly via exit codes
instead.

## Describe your changes

## Issue ticket number and link

## 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)
- [ ] 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/6922"><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=1787764680&installation_model_id=427504&pr_number=6922&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6922&signature=3d9d640b25a17589114cc704e1d631895ffcd9f554947ff5790fe90d54410d2f"><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

* **New Features**
* Added a direct installer download option for manual application
updates.
* Non-enforced updates now open the appropriate platform-specific
installer instead of the general releases page.

* **Bug Fixes**
  * Improved macOS download behavior for Homebrew installations.
* Preserved architecture-specific downloads for Intel and Apple silicon
Macs.
* Updated the tray “About” links to use the GitHub repository and
documentation instead of the releases page.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-27 20:08: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
110 changed files with 5734 additions and 751 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

@@ -24,6 +24,8 @@ builds:
ldflags:
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
mod_timestamp: "{{ .CommitTimestamp }}"
tags:
- production
- id: netbird-ui-windows-amd64
dir: client/ui
@@ -39,6 +41,8 @@ builds:
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
- -H windowsgui
mod_timestamp: "{{ .CommitTimestamp }}"
tags:
- production
- id: netbird-ui-windows-arm64
dir: client/ui
@@ -55,6 +59,8 @@ builds:
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
- -H windowsgui
mod_timestamp: "{{ .CommitTimestamp }}"
tags:
- production
archives:
- id: linux-arch

View File

@@ -29,6 +29,8 @@ builds:
ldflags:
- -s -w -X github.com/netbirdio/netbird/version.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.CommitDate}} -X main.builtBy=goreleaser
mod_timestamp: "{{ .CommitTimestamp }}"
tags:
- production
universal_binaries:
- id: netbird-ui-darwin

View File

@@ -7,6 +7,7 @@ import (
"fmt"
"os"
"slices"
"strings"
"sync"
"time"
@@ -299,6 +300,13 @@ func (c *Client) SetInfoLogLevel() {
// PeersList return with the list of the PeerInfos
func (c *Client) PeersList() *PeerInfoArray {
// The recorder only caches transfer counters and handshake times; nothing
// refreshes them on its own, so without this they read as zero. The desktop
// daemon does the same before serving a full peer status.
if err := c.recorder.RefreshWireGuardStats(); err != nil {
log.Debugf("failed to refresh WireGuard stats: %v", err)
}
fullStatus := c.recorder.GetFullStatus()
peerInfos := make([]PeerInfo, len(fullStatus.Peers))
@@ -309,6 +317,20 @@ func (c *Client) PeersList() *PeerInfoArray {
FQDN: p.FQDN,
ConnStatus: int(p.ConnStatus),
Routes: PeerRoutes{routes: maps.Keys(p.GetRoutes())},
PubKey: p.PubKey,
Latency: formatDuration(p.Latency),
LatencyMs: p.Latency.Milliseconds(),
BytesRx: p.BytesRx,
BytesTx: p.BytesTx,
ConnStatusUpdate: formatTime(p.ConnStatusUpdate),
Relayed: p.Relayed,
RosenpassEnabled: p.RosenpassEnabled,
LastWireguardHandshake: formatTime(p.LastWireguardHandshake),
LocalIceCandidateType: p.LocalIceCandidateType,
RemoteIceCandidateType: p.RemoteIceCandidateType,
LocalIceCandidateEndpoint: p.LocalIceCandidateEndpoint,
RemoteIceCandidateEndpoint: p.RemoteIceCandidateEndpoint,
}
peerInfos[n] = pi
}
@@ -439,10 +461,6 @@ func (c *Client) RemoveConnectionListener() {
c.recorder.RemoveConnectionListener()
}
func (c *Client) toggleRoute(command routeCommand) error {
return command.toggleRoute()
}
func (c *Client) getRouteManager() (routemanager.Manager, error) {
client := c.getConnectClient()
if client == nil {
@@ -462,22 +480,22 @@ func (c *Client) getRouteManager() (routemanager.Manager, error) {
return manager, nil
}
func (c *Client) SelectRoute(route string) error {
func (c *Client) SelectRoute(id string) error {
manager, err := c.getRouteManager()
if err != nil {
return err
}
return c.toggleRoute(selectRouteCommand{route: route, manager: manager})
return manager.SelectRoutes([]route.NetID{route.NetID(id)}, true)
}
func (c *Client) DeselectRoute(route string) error {
func (c *Client) DeselectRoute(id string) error {
manager, err := c.getRouteManager()
if err != nil {
return err
}
return c.toggleRoute(deselectRouteCommand{route: route, manager: manager})
return manager.DeselectRoutes([]route.NetID{route.NetID(id)})
}
// getNetworkDomainsFromRoute extracts domains from a route and enriches each domain
@@ -512,3 +530,28 @@ func exportEnvList(list *EnvList) {
}
}
}
// formatDuration renders a duration for display, trimming the fractional part
// to two digits so latencies read as "12.34ms" rather than "12.345678ms".
func formatDuration(d time.Duration) string {
ds := d.String()
dotIndex := strings.Index(ds, ".")
if dotIndex == -1 {
return ds
}
endIndex := min(dotIndex+3, len(ds))
// Skip the remaining digits so only the unit suffix is appended back.
unitStart := endIndex
for unitStart < len(ds) && ds[unitStart] >= '0' && ds[unitStart] <= '9' {
unitStart++
}
return ds[:endIndex] + ds[unitStart:]
}
// formatTime renders a timestamp in UTC using a fixed layout. The zero time is
// passed through as-is so the UI can recognise it and show "never" instead.
func formatTime(t time.Time) string {
return t.UTC().Format("2006-01-02 15:04:05")
}

View File

@@ -12,12 +12,30 @@ const (
)
// PeerInfo describe information about the peers. It designed for the UI usage
//
// The fields below ConnStatus back the peer detail screen. Durations and times
// are pre-formatted into strings so the UI does not have to know Go's layouts;
// Latency is additionally exposed as LatencyMs for colour coding.
type PeerInfo struct {
IP string
IPv6 string
FQDN string
ConnStatus int
Routes PeerRoutes
PubKey string
Latency string
LatencyMs int64
BytesRx int64
BytesTx int64
ConnStatusUpdate string
Relayed bool
RosenpassEnabled bool
LastWireguardHandshake string
LocalIceCandidateType string
RemoteIceCandidateType string
LocalIceCandidateEndpoint string
RemoteIceCandidateEndpoint string
}
func (p *PeerInfo) GetPeerRoutes() *PeerRoutes {

View File

@@ -189,6 +189,19 @@ func (pm *ProfileManager) LogoutProfile(id string) error {
return nil
}
// RenameProfile changes a profile's display name. The profile ID, and therefore
// its on-disk filename, is left untouched: only the "name" field of the config
// is rewritten. This works for the default profile too, whose config lives in
// netbird.cfg rather than under profiles/.
func (pm *ProfileManager) RenameProfile(id string, newName string) error {
if err := pm.serviceMgr.RenameProfile(profilemanager.ID(id), androidUsername, newName); err != nil {
return fmt.Errorf("failed to rename profile: %w", err)
}
log.Infof("renamed profile %s to: %s", id, newName)
return nil
}
// RemoveProfile deletes a profile
func (pm *ProfileManager) RemoveProfile(id string) error {
// Use ServiceManager (removes profile from profiles/ directory)

View File

@@ -1,70 +0,0 @@
//go:build android
package android
import (
"fmt"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
"github.com/netbirdio/netbird/client/internal/routemanager"
"github.com/netbirdio/netbird/route"
)
func executeRouteToggle(id string, manager routemanager.Manager,
operationName string,
routeOperation func(routes []route.NetID, allRoutes []route.NetID) error) error {
netID := route.NetID(id)
routes := []route.NetID{netID}
routesMap := manager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
log.Debugf("%s with ids: %v", operationName, routes)
if err := routeOperation(routes, maps.Keys(routesMap)); err != nil {
log.Debugf("error when %s: %s", operationName, err)
return fmt.Errorf("error %s: %w", operationName, err)
}
manager.TriggerSelection(manager.GetClientRoutes())
return nil
}
type routeCommand interface {
toggleRoute() error
}
type selectRouteCommand struct {
route string
manager routemanager.Manager
}
func (s selectRouteCommand) toggleRoute() error {
routeSelector := s.manager.GetRouteSelector()
if routeSelector == nil {
return fmt.Errorf("no route selector available")
}
routeOperation := func(routes []route.NetID, allRoutes []route.NetID) error {
return routeSelector.SelectRoutes(routes, true, allRoutes)
}
return executeRouteToggle(s.route, s.manager, "selecting route", routeOperation)
}
type deselectRouteCommand struct {
route string
manager routemanager.Manager
}
func (d deselectRouteCommand) toggleRoute() error {
routeSelector := d.manager.GetRouteSelector()
if routeSelector == nil {
return fmt.Errorf("no route selector available")
}
return executeRouteToggle(d.route, d.manager, "deselecting route", routeSelector.DeselectRoutes)
}

View File

@@ -385,11 +385,20 @@ func inactivityThresholdEnv() *time.Duration {
return nil
}
parsedMinutes, err := strconv.Atoi(envValue)
if err != nil || parsedMinutes <= 0 {
return nil
// Documented format: a Go duration such as "30m" or "1h".
if d, err := time.ParseDuration(envValue); err == nil {
if d <= 0 {
return nil
}
return &d
}
d := time.Duration(parsedMinutes) * time.Minute
return &d
// Backwards compatibility: a bare integer used to be interpreted as minutes.
if parsedMinutes, err := strconv.Atoi(envValue); err == nil && parsedMinutes > 0 {
d := time.Duration(parsedMinutes) * time.Minute
return &d
}
log.Warnf("invalid %s value %q: expected a Go duration such as 30m or 1h", lazyconn.EnvInactivityThreshold, envValue)
return nil
}

View File

@@ -104,3 +104,38 @@ func TestConnMgr_ActivatePeerConcurrentWithLifecycle(t *testing.T) {
close(done)
wg.Wait()
}
func TestInactivityThresholdEnv(t *testing.T) {
tests := []struct {
name string
val string
want *time.Duration
}{
{name: "unset", val: "", want: nil},
{name: "go duration minutes", val: "30m", want: durPtr(30 * time.Minute)},
{name: "go duration hours", val: "1h", want: durPtr(time.Hour)},
{name: "go duration seconds", val: "90s", want: durPtr(90 * time.Second)},
{name: "bare integer is minutes (backwards compat)", val: "5", want: durPtr(5 * time.Minute)},
{name: "zero duration", val: "0s", want: nil},
{name: "zero integer", val: "0", want: nil},
{name: "negative duration", val: "-5m", want: nil},
{name: "garbage", val: "abc", want: nil},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Setenv(lazyconn.EnvInactivityThreshold, tc.val)
got := inactivityThresholdEnv()
switch {
case tc.want == nil && got != nil:
t.Fatalf("want nil, got %v", *got)
case tc.want != nil && got == nil:
t.Fatalf("want %v, got nil", *tc.want)
case tc.want != nil && *got != *tc.want:
t.Fatalf("want %v, got %v", *tc.want, *got)
}
})
}
}
func durPtr(d time.Duration) *time.Duration { return &d }

View File

@@ -34,6 +34,7 @@ import (
"github.com/netbirdio/netbird/client/internal/profilemanager"
"github.com/netbirdio/netbird/client/internal/statemanager"
"github.com/netbirdio/netbird/client/internal/stdnet"
"github.com/netbirdio/netbird/client/internal/tunnelnotifier"
"github.com/netbirdio/netbird/client/internal/updater"
"github.com/netbirdio/netbird/client/internal/updater/installer"
nbnet "github.com/netbirdio/netbird/client/net"
@@ -136,10 +137,13 @@ func (c *ConnectClient) RunOniOS(
// Set GC percent to 5% to reduce memory usage as iOS only allows 50MB of memory for the extension.
debug.SetGCPercent(5)
notifier := tunnelnotifier.New(networkChangeListener, dnsManager)
defer notifier.Close()
mobileDependency := MobileDependency{
FileDescriptor: fileDescriptor,
NetworkChangeListener: networkChangeListener,
DnsManager: dnsManager,
NetworkChangeListener: notifier,
DnsManager: notifier,
StateFilePath: stateFilePath,
TempDir: cacheDir,
}

View File

@@ -11,12 +11,14 @@ import (
// MobileDependency collect all dependencies for mobile platform
type MobileDependency struct {
// Android only
TunAdapter device.TunAdapter
IFaceDiscover stdnet.ExternalIFaceDiscover
// Android and iOS
NetworkChangeListener listener.NetworkChangeListener
HostDNSAddresses []netip.AddrPort
DnsReadyListener dns.ReadyListener
// Android only
TunAdapter device.TunAdapter
IFaceDiscover stdnet.ExternalIFaceDiscover
HostDNSAddresses []netip.AddrPort
DnsReadyListener dns.ReadyListener
// iOS only
DnsManager dns.IosDnsManager

View File

@@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error {
// AllowedIPs should use real IPs
if d.currentPeerKey != "" {
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
}
}
@@ -172,7 +172,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error {
}
// AllowedIPs use real IPs
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil {
return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err)
}
@@ -205,7 +205,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error {
for _, prefixes := range d.interceptedDomains {
for _, prefix := range prefixes {
// AllowedIPs use real IPs
if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil {
if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err))
}
}

View File

@@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error {
var merr *multierror.Error
for _, domainPrefixes := range r.dynamicDomains {
for _, prefix := range domainPrefixes {
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
}
}
@@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) {
merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err))
}
if r.currentPeerKey != "" {
if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil {
if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err))
}
}

View File

@@ -52,6 +52,10 @@ type Manager interface {
UpdateRoutes(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
ClassifyRoutes(newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
TriggerSelection(route.HAMap)
SelectRoutes(ids []route.NetID, appendRoute bool) error
DeselectRoutes(ids []route.NetID) error
SelectAllRoutes()
DeselectAllRoutes()
GetRouteSelector() *routeselector.RouteSelector
GetClientRoutes() route.HAMap
GetSelectedClientRoutes() route.HAMap
@@ -216,7 +220,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
)
}
m.allowedIPsRefCounter = refcounter.New(
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
func(prefix netip.Prefix, peerKey string) (string, error) {
// save peerKey to use it in the remove function
return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix)
@@ -800,7 +804,7 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
var info exitNodeInfo
for haID, routes := range clientRoutes {
if !m.isExitNodeRoute(routes) {
if !isExitNodeRoutes(routes) {
continue
}
@@ -820,13 +824,6 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
return info
}
func (m *DefaultManager) isExitNodeRoute(routes []*route.Route) bool {
if len(routes) == 0 {
return false
}
return route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network)
}
func (m *DefaultManager) categorizeUserSelection(netID route.NetID, info *exitNodeInfo) {
if m.routeSelector.IsSelected(netID) {
info.userSelected = append(info.userSelected, netID)

View File

@@ -16,6 +16,8 @@ type MockManager struct {
ClassifyRoutesFunc func(routes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
UpdateRoutesFunc func(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
TriggerSelectionFunc func(haMap route.HAMap)
SelectRoutesFunc func(ids []route.NetID, appendRoute bool) error
DeselectRoutesFunc func(ids []route.NetID) error
GetRouteSelectorFunc func() *routeselector.RouteSelector
GetClientRoutesFunc func() route.HAMap
GetSelectedClientRoutesFunc func() route.HAMap
@@ -55,6 +57,30 @@ func (m *MockManager) TriggerSelection(networks route.HAMap) {
}
}
// SelectRoutes mock implementation of SelectRoutes from Manager interface
func (m *MockManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
if m.SelectRoutesFunc != nil {
return m.SelectRoutesFunc(ids, appendRoute)
}
return nil
}
// DeselectRoutes mock implementation of DeselectRoutes from Manager interface
func (m *MockManager) DeselectRoutes(ids []route.NetID) error {
if m.DeselectRoutesFunc != nil {
return m.DeselectRoutesFunc(ids)
}
return nil
}
// SelectAllRoutes mock implementation of SelectAllRoutes from Manager interface
func (m *MockManager) SelectAllRoutes() {
}
// DeselectAllRoutes mock implementation of DeselectAllRoutes from Manager interface
func (m *MockManager) DeselectAllRoutes() {
}
// GetRouteSelector mock implementation of GetRouteSelector from Manager interface
func (m *MockManager) GetRouteSelector() *routeselector.RouteSelector {
if m.GetRouteSelectorFunc != nil {

View File

@@ -3,7 +3,6 @@
package notifier
import (
"container/list"
"net/netip"
"slices"
"sort"
@@ -16,20 +15,12 @@ import (
type Notifier struct {
mu sync.Mutex
cond *sync.Cond
currentPrefixes []string
listener listener.NetworkChangeListener
queue *list.List
closed bool
}
func NewNotifier() *Notifier {
n := &Notifier{
queue: list.New(),
}
n.cond = sync.NewCond(&n.mu)
go n.deliverLoop()
return n
return &Notifier{}
}
func (n *Notifier) SetListener(listener listener.NetworkChangeListener) {
@@ -59,44 +50,19 @@ func (n *Notifier) OnNewPrefixes(prefixes []netip.Prefix) {
sort.Strings(newNets)
n.mu.Lock()
defer n.mu.Unlock()
if slices.Equal(n.currentPrefixes, newNets) {
n.mu.Unlock()
return
}
n.currentPrefixes = newNets
routes := strings.Join(n.currentPrefixes, ",")
n.queue.PushBack(routes)
n.cond.Signal()
n.mu.Unlock()
if n.listener != nil {
n.listener.OnNetworkChanged(strings.Join(n.currentPrefixes, ","))
}
}
func (n *Notifier) Close() {
n.mu.Lock()
n.closed = true
n.cond.Signal()
n.mu.Unlock()
}
func (n *Notifier) GetInitialRouteRanges() []string {
return nil
}
func (n *Notifier) deliverLoop() {
for {
n.mu.Lock()
for n.queue.Len() == 0 && !n.closed {
n.cond.Wait()
}
if n.closed && n.queue.Len() == 0 {
n.mu.Unlock()
return
}
routes := n.queue.Remove(n.queue.Front()).(string)
l := n.listener
n.mu.Unlock()
if l != nil {
l.OnNetworkChanged(routes)
}
}
}

View File

@@ -54,7 +54,7 @@ func (m *reconcileWGMock) GetNet() *netstack.Net { return n
func TestReconcilePeerAllowedIPs(t *testing.T) {
wg := &reconcileWGMock{}
m := &DefaultManager{wgInterface: wg}
m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
m.allowedIPsRefCounter = refcounter.NewAllowedIPs(
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
func(netip.Prefix, string) error { return nil },
)

View File

@@ -0,0 +1,206 @@
package refcounter
import (
"errors"
"fmt"
"net/netip"
"sort"
"sync"
"github.com/hashicorp/go-multierror"
nberrors "github.com/netbirdio/netbird/client/errors"
)
// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is
// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most
// one peer is active at a time even when several peers reference the prefix.
type allowedIPsEntry struct {
// peers maps a peerKey to the number of references holding the prefix for that peer.
peers map[string]int
// active is the peerKey currently installed in WireGuard for this prefix ("" if none).
active string
// total is the sum of all per-peer reference counts (kept in sync with peers).
total int
}
// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs.
//
// The generic Counter keys only by prefix and remembers a single Out value set by the first
// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or
// multiple resolved domains) can reference the same prefix through different peers, and when the
// peer currently installed in WireGuard releases its last reference the prefix must be handed over
// to a surviving peer instead of being left pointing at the released one.
//
// It calls add/remove (which program WireGuard) only on the transitions that matter:
// - add on the first reference for a prefix, or when swapping the active peer;
// - remove on the last reference for a prefix, or on the old peer during a swap.
type AllowedIPsRefCounter struct {
mu sync.Mutex
entries map[netip.Prefix]*allowedIPsEntry
add AddFunc[netip.Prefix, string, string]
remove RemoveFunc[netip.Prefix, string]
}
// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter.
// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer.
// remove unprograms the prefix from the given peer.
func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter {
return &AllowedIPsRefCounter{
entries: map[netip.Prefix]*allowedIPsEntry{},
add: add,
remove: remove,
}
}
// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first
// reference to a prefix; while a different peer is already installed the prefix is left with it
// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept.
func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) {
rm.mu.Lock()
defer rm.mu.Unlock()
e, ok := rm.entries[prefix]
if !ok {
e = &allowedIPsEntry{peers: map[string]int{}}
rm.entries[prefix] = e
}
logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active)
// Program WireGuard only when nothing is installed yet for this prefix.
if e.active == "" {
out, err := rm.add(prefix, peerKey)
if errors.Is(err, ErrIgnore) {
if e.total == 0 {
delete(rm.entries, prefix)
}
return Ref[string]{Count: e.total, Out: e.active}, nil
}
if err != nil {
if e.total == 0 {
delete(rm.entries, prefix)
}
return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err)
}
e.active = out
}
e.peers[peerKey]++
e.total++
return Ref[string]{Count: e.total, Out: e.active}, nil
}
// Decrement removes a reference to prefix for peerKey. When the peer currently installed in
// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists,
// otherwise it is removed from WireGuard.
func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) {
rm.mu.Lock()
defer rm.mu.Unlock()
e, ok := rm.entries[prefix]
if !ok {
logCallerF("No allowed IP reference found for prefix %v", prefix)
return Ref[string]{}, nil
}
if e.peers[peerKey] > 0 {
logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]",
prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active)
e.peers[peerKey]--
e.total--
if e.peers[peerKey] == 0 {
delete(e.peers, peerKey)
}
} else {
logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey)
}
// If the peer currently installed in WireGuard still holds references, nothing to reprogram.
// Keying the check on the active peer (not the one just released) makes this self-healing:
// a prior swap whose remove/add failed leaves e.active pointing at a peer with no references,
// and this retries the hand-off on the next Decrement instead of getting stuck.
if e.active != "" && e.peers[e.active] > 0 {
return Ref[string]{Count: e.total, Out: e.active}, nil
}
// Detach the stale/gone active peer from WireGuard before reprogramming.
if e.active != "" {
if err := rm.remove(prefix, e.active); err != nil {
return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)
}
e.active = ""
}
// Hand the prefix over to a surviving peer, or drop the entry when none remain.
if survivor, ok := pickSurvivor(e.peers); ok {
out, err := rm.add(prefix, survivor)
if err != nil {
return Ref[string]{Count: e.total, Out: ""}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err)
}
e.active = out
return Ref[string]{Count: e.total, Out: e.active}, nil
}
delete(rm.entries, prefix)
return Ref[string]{Count: 0, Out: ""}, nil
}
// Flush removes all prefixes from WireGuard and clears the counter.
func (rm *AllowedIPsRefCounter) Flush() error {
rm.mu.Lock()
defer rm.mu.Unlock()
var merr *multierror.Error
for prefix, e := range rm.entries {
if e.active == "" {
continue
}
logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active)
if err := rm.remove(prefix, e.active); err != nil {
merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err))
}
}
clear(rm.entries)
return nberrors.FormatErrorOrNil(merr)
}
// ReapplyMatching calls apply for every prefix whose currently installed (active) peer satisfies
// pred, holding the lock for the whole pass. It is used to re-push allowed IPs onto a peer whose
// WireGuard entry was rebuilt (e.g. a lazy connection cycling idle->wake) without a matching
// refcounter change, which would otherwise leave the prefix installed in the counter but missing
// on the device. Only the active peer is considered — a prefix that lost its installed peer to a
// failed swap is skipped here and reconciled by the next Increment/Decrement.
func (rm *AllowedIPsRefCounter) ReapplyMatching(pred func(out string) bool, apply func(key netip.Prefix) error) error {
rm.mu.Lock()
defer rm.mu.Unlock()
var merr *multierror.Error
for prefix, e := range rm.entries {
if e.active != "" && pred(e.active) {
if err := apply(prefix); err != nil {
merr = multierror.Append(merr, err)
}
}
}
return nberrors.FormatErrorOrNil(merr)
}
// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do
// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable
// (lowest peerKey) for predictable behavior and testability.
func pickSurvivor(peers map[string]int) (string, bool) {
if len(peers) == 0 {
return "", false
}
keys := make([]string, 0, len(peers))
for k := range peers {
keys = append(keys, k)
}
sort.Strings(keys)
return keys[0], true
}

View File

@@ -0,0 +1,241 @@
package refcounter
import (
"errors"
"net/netip"
"testing"
)
// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer.
// failAdd/failRemove make the next add/remove fail once, to exercise the self-healing error paths.
type fakeWG struct {
installed map[netip.Prefix]string
adds int
removes int
failAdd bool
failRemove bool
}
func newFakeWG() *fakeWG {
return &fakeWG{installed: map[netip.Prefix]string{}}
}
func (f *fakeWG) counter() *AllowedIPsRefCounter {
return NewAllowedIPs(
func(prefix netip.Prefix, peerKey string) (string, error) {
if f.failAdd {
f.failAdd = false
return "", errors.New("add failed")
}
f.adds++
f.installed[prefix] = peerKey
return peerKey, nil
},
func(prefix netip.Prefix, peerKey string) error {
if f.failRemove {
f.failRemove = false
return errors.New("remove failed")
}
f.removes++
// only clear if this peer is the one installed, mirroring wg semantics
if f.installed[prefix] == peerKey {
delete(f.installed, prefix)
}
return nil
},
)
}
func mustPrefix(t *testing.T, s string) netip.Prefix {
t.Helper()
p, err := netip.ParsePrefix(s)
if err != nil {
t.Fatalf("parse prefix %q: %v", s, err)
}
return p
}
func mustIncrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
t.Helper()
ref, err := c.Increment(p, peer)
if err != nil {
t.Fatalf("Increment(%v, %s): %v", p, peer, err)
}
return ref
}
func mustDecrement(t *testing.T, c *AllowedIPsRefCounter, p netip.Prefix, peer string) Ref[string] {
t.Helper()
ref, err := c.Decrement(p, peer)
if err != nil {
t.Fatalf("Decrement(%v, %s): %v", p, peer, err)
}
return ref
}
// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same
// prefix routed by different peers. Removing the network whose peer is installed must hand the
// prefix over to the surviving peer instead of leaving it on the removed one.
func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
mustIncrement(t, c, p, "peerA")
mustIncrement(t, c, p, "peerB")
// First peer wins while both are present.
if got := f.installed[p]; got != "peerA" {
t.Fatalf("expected peerA installed, got %q", got)
}
// Remove the active peer's network -> must swap to peerB.
mustDecrement(t, c, p, "peerA")
if got := f.installed[p]; got != "peerB" {
t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got)
}
// Remove the last one -> prefix gone.
mustDecrement(t, c, p, "peerB")
if _, ok := f.installed[p]; ok {
t.Fatalf("expected prefix removed, still installed on %q", f.installed[p])
}
}
// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard.
func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
mustIncrement(t, c, p, "peerA")
mustIncrement(t, c, p, "peerB")
removesBefore := f.removes
mustDecrement(t, c, p, "peerB")
if f.installed[p] != "peerA" {
t.Fatalf("active peer must stay peerA, got %q", f.installed[p])
}
if f.removes != removesBefore {
t.Fatalf("removing a non-active peer must not call wg remove")
}
}
// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until
// the last reference is released (the reason the per-peer count must be an int, not a set).
func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
mustIncrement(t, c, p, "peerA")
mustIncrement(t, c, p, "peerA")
if f.adds != 1 {
t.Fatalf("expected a single wg add for the same peer, got %d", f.adds)
}
mustDecrement(t, c, p, "peerA")
if f.installed[p] != "peerA" {
t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p])
}
if f.removes != 0 {
t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes)
}
mustDecrement(t, c, p, "peerA")
if _, ok := f.installed[p]; ok {
t.Fatalf("prefix must be removed after last reference")
}
}
// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log).
func TestAllowedIPs_RefCountAndActive(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
ref := mustIncrement(t, c, p, "peerA")
if ref.Count != 1 || ref.Out != "peerA" {
t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out)
}
ref = mustIncrement(t, c, p, "peerB")
if ref.Count != 2 || ref.Out != "peerA" {
t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out)
}
}
// TestAllowedIPs_Flush removes everything installed and clears the counter.
func TestAllowedIPs_Flush(t *testing.T) {
f := newFakeWG()
c := f.counter()
p1 := mustPrefix(t, "10.44.8.0/24")
p2 := mustPrefix(t, "10.44.9.0/24")
mustIncrement(t, c, p1, "peerA")
mustIncrement(t, c, p2, "peerB")
if err := c.Flush(); err != nil {
t.Fatal(err)
}
if len(f.installed) != 0 {
t.Fatalf("expected all prefixes removed, got %v", f.installed)
}
// After flush, a fresh increment must add again.
mustIncrement(t, c, p1, "peerC")
if f.installed[p1] != "peerC" {
t.Fatalf("counter not reset after flush")
}
}
// TestAllowedIPs_SelfHealAfterSwapAddError ensures a failed add during a swap does not permanently
// strand the prefix: the next Decrement (or Increment) must retry and install a surviving peer.
func TestAllowedIPs_SelfHealAfterSwapAddError(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
mustIncrement(t, c, p, "peerA")
mustIncrement(t, c, p, "peerB")
mustIncrement(t, c, p, "peerC")
// Removing the active peerA triggers a swap to a survivor; make the add fail once.
f.failAdd = true
if _, err := c.Decrement(p, "peerA"); err == nil {
t.Fatalf("expected error from failed swap add")
}
if _, ok := f.installed[p]; ok {
t.Fatalf("nothing should be installed after a failed swap add, got %q", f.installed[p])
}
// A later Decrement of a non-active survivor must retry the hand-off (self-heal), not stay stuck.
ref := mustDecrement(t, c, p, "peerC")
if got := f.installed[p]; got == "" {
t.Fatalf("self-heal failed: prefix left unrouted after add recovered")
}
if ref.Out == "" {
t.Fatalf("expected an active peer after self-heal, got empty")
}
}
// TestAllowedIPs_SelfHealAfterRemoveError ensures a failed remove during a swap is retried instead
// of leaving e.active stuck on a peer that no longer holds references.
func TestAllowedIPs_SelfHealAfterRemoveError(t *testing.T) {
f := newFakeWG()
c := f.counter()
p := mustPrefix(t, "10.44.8.0/24")
mustIncrement(t, c, p, "peerA")
mustIncrement(t, c, p, "peerB")
// Releasing active peerA must detach it (remove) then add peerB; fail the remove once.
f.failRemove = true
if _, err := c.Decrement(p, "peerA"); err == nil {
t.Fatalf("expected error from failed remove")
}
// Next Decrement of the non-active survivor retries: removes stale peerA, installs peerB.
mustDecrement(t, c, p, "peerB")
// peerB had only one ref, so after retry the prefix is fully released.
if _, ok := f.installed[p]; ok {
t.Fatalf("expected prefix released after self-heal, still on %q", f.installed[p])
}
}

View File

@@ -5,5 +5,7 @@ import "net/netip"
// RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement
type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}]
// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement
type AllowedIPsRefCounter = Counter[netip.Prefix, string, string]
// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware:
// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer,
// so the counter records the per-peer reference count and swaps the installed peer when the active one is released.
// See allowedips.go.

View File

@@ -0,0 +1,138 @@
package routemanager
import (
"fmt"
"slices"
"github.com/hashicorp/go-multierror"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
nberrors "github.com/netbirdio/netbird/client/errors"
"github.com/netbirdio/netbird/route"
)
// SelectRoutes selects the routes with the given network IDs and applies the
// new selection. V4/v6 exit-node pairs are expanded automatically. Exit nodes
// are mutually exclusive: if the selection activates an exit node, every other
// available exit node is deselected so two can't be active at once. With
// appendRoute=false the previous selection is replaced instead of extended.
func (m *DefaultManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
if err := m.selectRoutes(ids, appendRoute); err != nil {
return err
}
m.TriggerSelection(m.GetClientRoutes())
return nil
}
// DeselectRoutes removes the routes with the given network IDs from the
// selection and applies the change. V4/v6 exit-node pairs are expanded
// automatically.
func (m *DefaultManager) DeselectRoutes(ids []route.NetID) error {
if err := m.deselectRoutes(ids); err != nil {
return err
}
m.TriggerSelection(m.GetClientRoutes())
return nil
}
func (m *DefaultManager) deselectRoutes(ids []route.NetID) error {
routesMap := m.GetClientRoutesWithNetID()
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
log.Debugf("deselecting routes with ids: %v", routes)
if err := m.routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
return fmt.Errorf("deselect routes: %w", err)
}
return nil
}
// SelectAllRoutes selects every available route and applies the selection.
// Exit nodes stay mutually exclusive: at most one remains active.
func (m *DefaultManager) SelectAllRoutes() {
m.selectAllRoutes()
m.TriggerSelection(m.GetClientRoutes())
}
func (m *DefaultManager) selectAllRoutes() {
m.routeSelector.SelectAllRoutes()
// Select-all wipes every explicit selection, so exit nodes fall back to
// management's auto-apply flags — which may mark several at once.
// Reconcile immediately so at most one exit node stays active instead of
// waiting for the next network map to enforce it.
m.mux.Lock()
defer m.mux.Unlock()
m.updateRouteSelectorFromManagement(m.clientRoutes)
}
// DeselectAllRoutes deselects every route and applies the change.
func (m *DefaultManager) DeselectAllRoutes() {
m.routeSelector.DeselectAllRoutes()
m.TriggerSelection(m.GetClientRoutes())
}
func (m *DefaultManager) selectRoutes(ids []route.NetID, appendRoute bool) error {
routesMap := m.GetClientRoutesWithNetID()
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
allIDs := maps.Keys(routesMap)
log.Debugf("selecting routes with ids: %v", routes)
// A partial failure (e.g. an unknown ID in the request) still selects the
// valid routes, so exclusivity below must run regardless of the error.
var merr *multierror.Error
if err := m.routeSelector.SelectRoutes(routes, appendRoute, allIDs); err != nil {
merr = multierror.Append(merr, fmt.Errorf("select routes: %w", err))
}
// Exit nodes are mutually exclusive: if this selection activates an
// exit node, deselect every other available exit node so two can't be
// selected at once. Non-exit route selections are left untouched.
if requestActivatesExitNode(routes, routesMap) {
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
if err := m.routeSelector.DeselectRoutes(others, allIDs); err != nil {
merr = multierror.Append(merr, fmt.Errorf("deselect sibling exit nodes: %w", err))
}
}
}
return nberrors.FormatErrorOrNil(merr)
}
func isExitNodeRoutes(routes []*route.Route) bool {
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
}
// requestActivatesExitNode reports whether any requested NetID maps to an exit
// node (default route) in the current route table.
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
for _, id := range requested {
if isExitNodeRoutes(routesMap[id]) {
return true
}
}
return false
}
// otherExitNodeIDs returns every available exit-node NetID that is not in the
// requested set — the siblings to deselect so a single exit node stays active.
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
keep := make(map[route.NetID]struct{}, len(requested))
for _, id := range requested {
keep[id] = struct{}{}
}
var others []route.NetID
for id, routes := range routesMap {
if !isExitNodeRoutes(routes) {
continue
}
if _, ok := keep[id]; ok {
continue
}
others = append(others, id)
}
return others
}

View File

@@ -0,0 +1,129 @@
package routemanager
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/client/internal/routeselector"
"github.com/netbirdio/netbird/route"
)
func v6ExitRoute(netID, peer string) *route.Route {
return &route.Route{
NetID: route.NetID(netID),
Network: netip.MustParsePrefix("::/0"),
Peer: peer,
}
}
func newSelectionTestManager() *DefaultManager {
return &DefaultManager{
routeSelector: routeselector.NewRouteSelector(),
clientRoutes: route.HAMap{
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", true)},
"exitA-v6|::/0": {v6ExitRoute("exitA-v6", "p1")},
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", true)},
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
},
}
}
func TestSelectRoutes_ExitNodeExclusivity(t *testing.T) {
m := newSelectionTestManager()
// Selecting an exit node selects its v6 pair and deselects the sibling.
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
assert.True(t, m.routeSelector.IsSelected("exitA"), "exitA should be selected")
assert.True(t, m.routeSelector.IsSelected("exitA-v6"), "the v6 pair follows its v4 base")
assert.False(t, m.routeSelector.IsSelected("exitB"), "the sibling exit node must be deselected")
// Switching to the sibling deselects the previous exit node and its v6 pair.
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
assert.True(t, m.routeSelector.IsSelected("exitB"), "exitB should now be selected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "the previous exit node must be deselected")
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "the previous exit node's v6 pair must be deselected")
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
// Selecting a non-exit route leaves the active exit node alone.
require.NoError(t, m.selectRoutes([]route.NetID{"lan"}, true))
assert.True(t, m.routeSelector.IsSelected("exitB"), "selecting a non-exit route keeps the exit node")
// Deselecting the active exit node turns every exit node off.
require.NoError(t, m.deselectRoutes([]route.NetID{"exitB"}))
assert.False(t, m.routeSelector.IsSelected("exitB"), "exitB should be deselected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "exitA stays deselected")
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
}
func TestSelectRoutes_PartialErrorStillEnforcesExclusivity(t *testing.T) {
// The unknown ID must be reported, but the valid exit node in the same
// request is still selected — so its sibling must still be deselected.
// Both orderings are covered: processing must continue past the invalid
// ID wherever it sits in the request.
requests := map[string][]route.NetID{
"invalid id first": {"missing", "exitB"},
"invalid id last": {"exitB", "missing"},
}
for name, ids := range requests {
t.Run(name, func(t *testing.T) {
m := newSelectionTestManager()
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
err := m.selectRoutes(ids, true)
assert.Error(t, err, "unknown id must be reported")
assert.True(t, m.routeSelector.IsSelected("exitB"), "valid exit node from the request is selected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "sibling exit node must be deselected despite the error")
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "sibling's v6 pair must be deselected too")
})
}
}
func TestSelectAllRoutes_KeepsSingleExitNode(t *testing.T) {
// Both exit nodes are marked for auto-apply by management
// (SkipAutoApply=false), the state where select-all could turn on two at
// once without the immediate reconciliation.
m := &DefaultManager{
routeSelector: routeselector.NewRouteSelector(),
clientRoutes: route.HAMap{
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", false)},
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", false)},
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
},
}
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
m.selectAllRoutes()
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit routes are all selected")
assert.True(t, m.routeSelector.IsSelected("exitA"), "the deterministic management pick stays active")
assert.False(t, m.routeSelector.IsSelected("exitB"), "select-all must not leave a second exit node active")
}
func TestSelectRoutes_UnknownRoute(t *testing.T) {
m := newSelectionTestManager()
assert.Error(t, m.selectRoutes([]route.NetID{"missing"}, true), "selecting an unavailable route must fail")
assert.Error(t, m.deselectRoutes([]route.NetID{"missing"}), "deselecting an unavailable route must fail")
}
func TestExitNodeSelectionHelpers(t *testing.T) {
routesMap := map[route.NetID][]*route.Route{
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
}
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
}

View File

@@ -15,6 +15,11 @@ type Route struct {
route *route.Route
routeRefCounter *refcounter.RouteRefCounter
allowedIPsRefcounter *refcounter.AllowedIPsRefCounter
// currentPeerKey is the routing peer this watcher currently has the prefix installed on
// (the HA winner elected by the watcher). It can differ from route.Peer and change on
// failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement
// the exact peer that was incremented.
currentPeerKey string
}
func NewRoute(params common.HandlerParams) *Route {
@@ -52,12 +57,15 @@ func (r *Route) AddAllowedIPs(peerKey string) error {
ref.Out,
)
}
r.currentPeerKey = peerKey
return nil
}
func (r *Route) RemoveAllowedIPs() error {
if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil {
return err
var err error
if _, decErr := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); decErr != nil {
err = fmt.Errorf("remove allowed IP %s: %w", r.route.Network, decErr)
}
return nil
r.currentPeerKey = ""
return err
}

View File

@@ -20,6 +20,8 @@ const (
rpFilterPath = "net.ipv4.conf.all.rp_filter"
rpFilterInterfacePath = "net.ipv4.conf.%s.rp_filter"
srcValidMarkPath = "net.ipv4.conf.all.src_valid_mark"
percentEscape = "%25"
dotEscape = "%2E"
)
type iface interface {
@@ -56,7 +58,11 @@ func Setup(wgIface iface) (map[string]int, error) {
continue
}
i := fmt.Sprintf(rpFilterInterfacePath, intf.Name)
// Escape '%' and '.' so they survive the dot-to-slash conversion in Set()
safeName := strings.ReplaceAll(intf.Name, "%", percentEscape)
safeName = strings.ReplaceAll(safeName, ".", dotEscape)
i := fmt.Sprintf(rpFilterInterfacePath, safeName)
oldVal, err := Set(i, 2, true)
if err != nil {
result = multierror.Append(result, err)
@@ -70,7 +76,11 @@ func Setup(wgIface iface) (map[string]int, error) {
// Set sets a sysctl configuration, if onlyIfOne is true it will only set the new value if it's set to 1
func Set(key string, desiredValue int, onlyIfOne bool) (int, error) {
path := fmt.Sprintf("/proc/sys/%s", strings.ReplaceAll(key, ".", "/"))
path := strings.ReplaceAll(key, ".", "/")
// Unescape interface dots and percent signs
path = strings.ReplaceAll(path, dotEscape, ".")
path = strings.ReplaceAll(path, percentEscape, "%")
path = fmt.Sprintf("/proc/sys/%s", path)
currentValue, err := os.ReadFile(path)
if err != nil {
return -1, fmt.Errorf("read sysctl %s: %w", key, err)

View File

@@ -0,0 +1,124 @@
package tunnelnotifier
import (
"container/list"
"sync"
"github.com/netbirdio/netbird/client/internal/dns"
"github.com/netbirdio/netbird/client/internal/listener"
)
type eventKind int
const (
eventRoutes eventKind = iota
eventIfaceIP
eventIfaceIPv6
eventDNS
)
var (
_ listener.NetworkChangeListener = (*Notifier)(nil)
_ dns.IosDnsManager = (*Notifier)(nil)
)
type event struct {
kind eventKind
payload string
}
type Notifier struct {
mu sync.Mutex
cond *sync.Cond
queue *list.List
closed bool
done chan struct{}
listener listener.NetworkChangeListener
dnsManager dns.IosDnsManager
}
func New(l listener.NetworkChangeListener, dm dns.IosDnsManager) *Notifier {
n := &Notifier{
queue: list.New(),
done: make(chan struct{}),
listener: l,
dnsManager: dm,
}
n.cond = sync.NewCond(&n.mu)
go n.deliverLoop()
return n
}
func (n *Notifier) OnNetworkChanged(routes string) {
n.enqueue(event{kind: eventRoutes, payload: routes})
}
func (n *Notifier) SetInterfaceIP(ip string) {
n.enqueue(event{kind: eventIfaceIP, payload: ip})
}
func (n *Notifier) SetInterfaceIPv6(ip string) {
n.enqueue(event{kind: eventIfaceIPv6, payload: ip})
}
func (n *Notifier) ApplyDns(config string) {
n.enqueue(event{kind: eventDNS, payload: config})
}
// Close stops accepting new events and blocks until the delivery loop has
// drained all queued events and exited.
func (n *Notifier) Close() {
n.mu.Lock()
n.closed = true
n.cond.Signal()
n.mu.Unlock()
<-n.done
}
func (n *Notifier) enqueue(ev event) {
n.mu.Lock()
defer n.mu.Unlock()
if n.closed {
return
}
n.queue.PushBack(ev)
n.cond.Signal()
}
func (n *Notifier) deliverLoop() {
defer close(n.done)
for {
n.mu.Lock()
for n.queue.Len() == 0 && !n.closed {
n.cond.Wait()
}
if n.closed && n.queue.Len() == 0 {
n.mu.Unlock()
return
}
ev := n.queue.Remove(n.queue.Front()).(event)
l := n.listener
dm := n.dnsManager
n.mu.Unlock()
switch ev.kind {
case eventRoutes:
if l != nil {
l.OnNetworkChanged(ev.payload)
}
case eventIfaceIP:
if l != nil {
l.SetInterfaceIP(ev.payload)
}
case eventIfaceIPv6:
if l != nil {
l.SetInterfaceIPv6(ev.payload)
}
case eventDNS:
if dm != nil {
dm.ApplyDns(ev.payload)
}
}
}
}

View File

@@ -0,0 +1,192 @@
package tunnelnotifier
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type call struct {
kind string
payload string
}
type recorder struct {
mu sync.Mutex
calls []call
inFlight atomic.Int32
overlap atomic.Bool
delay time.Duration
}
func (r *recorder) record(kind, payload string) {
if r.inFlight.Add(1) != 1 {
r.overlap.Store(true)
}
if r.delay > 0 {
time.Sleep(r.delay)
}
r.mu.Lock()
r.calls = append(r.calls, call{kind: kind, payload: payload})
r.mu.Unlock()
r.inFlight.Add(-1)
}
func (r *recorder) count() int {
r.mu.Lock()
defer r.mu.Unlock()
return len(r.calls)
}
func (r *recorder) snapshot() []call {
r.mu.Lock()
defer r.mu.Unlock()
out := make([]call, len(r.calls))
copy(out, r.calls)
return out
}
type fakeListener struct {
rec *recorder
}
func (f *fakeListener) OnNetworkChanged(routes string) {
f.rec.record("routes", routes)
}
func (f *fakeListener) SetInterfaceIP(ip string) {
f.rec.record("ip", ip)
}
func (f *fakeListener) SetInterfaceIPv6(ip string) {
f.rec.record("ipv6", ip)
}
type fakeDNSManager struct {
rec *recorder
}
func (f *fakeDNSManager) ApplyDns(config string) {
f.rec.record("dns", config)
}
func TestFIFOOrder(t *testing.T) {
rec := &recorder{}
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
defer n.Close()
n.SetInterfaceIP("10.0.0.1")
n.SetInterfaceIPv6("fd00::1")
n.ApplyDns(`{"domains":[]}`)
n.OnNetworkChanged("10.0.0.0/8,192.168.0.0/16")
n.ApplyDns(`{"domains":["example.com"]}`)
require.Eventually(t, func() bool { return rec.count() == 5 }, time.Second, time.Millisecond)
expected := []call{
{kind: "ip", payload: "10.0.0.1"},
{kind: "ipv6", payload: "fd00::1"},
{kind: "dns", payload: `{"domains":[]}`},
{kind: "routes", payload: "10.0.0.0/8,192.168.0.0/16"},
{kind: "dns", payload: `{"domains":["example.com"]}`},
}
assert.Equal(t, expected, rec.snapshot())
}
func TestNoOverlappingCalls(t *testing.T) {
rec := &recorder{delay: 100 * time.Microsecond}
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
defer n.Close()
const producers = 8
const perProducer = 25
var wg sync.WaitGroup
for i := 0; i < producers; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
for j := 0; j < perProducer; j++ {
payload := fmt.Sprintf("%d-%d", id, j)
switch j % 4 {
case 0:
n.OnNetworkChanged(payload)
case 1:
n.SetInterfaceIP(payload)
case 2:
n.SetInterfaceIPv6(payload)
case 3:
n.ApplyDns(payload)
}
}
}(i)
}
wg.Wait()
require.Eventually(t, func() bool { return rec.count() == producers*perProducer }, 5*time.Second, time.Millisecond)
assert.False(t, rec.overlap.Load())
}
func TestDNSAndRoutesInterleaved(t *testing.T) {
rec := &recorder{delay: 100 * time.Microsecond}
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
defer n.Close()
const events = 50
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
for i := 0; i < events; i++ {
n.ApplyDns(fmt.Sprintf("dns-%d", i))
}
}()
go func() {
defer wg.Done()
for i := 0; i < events; i++ {
n.OnNetworkChanged(fmt.Sprintf("routes-%d", i))
}
}()
wg.Wait()
require.Eventually(t, func() bool { return rec.count() == 2*events }, 5*time.Second, time.Millisecond)
assert.False(t, rec.overlap.Load())
var dnsSeen, routesSeen int
for _, c := range rec.snapshot() {
switch c.kind {
case "dns":
assert.Equal(t, fmt.Sprintf("dns-%d", dnsSeen), c.payload)
dnsSeen++
case "routes":
assert.Equal(t, fmt.Sprintf("routes-%d", routesSeen), c.payload)
routesSeen++
}
}
assert.Equal(t, events, dnsSeen)
assert.Equal(t, events, routesSeen)
}
func TestCloseDrainsQueue(t *testing.T) {
rec := &recorder{delay: time.Millisecond}
n := New(&fakeListener{rec: rec}, &fakeDNSManager{rec: rec})
const events = 20
for i := 0; i < events; i++ {
n.OnNetworkChanged(fmt.Sprintf("routes-%d", i))
}
n.Close()
require.Equal(t, events, rec.count(), "Close must not return before all queued events are delivered")
n.OnNetworkChanged("after-close")
n.ApplyDns("after-close")
time.Sleep(50 * time.Millisecond)
assert.Equal(t, events, rec.count())
}

View File

@@ -13,7 +13,6 @@ import (
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
"github.com/netbirdio/netbird/client/internal"
"github.com/netbirdio/netbird/client/internal/auth"
@@ -637,23 +636,18 @@ func (c *Client) SelectRoute(id string) error {
}
routeManager := engine.GetRouteManager()
routeSelector := routeManager.GetRouteSelector()
if id == "All" {
log.Debugf("select all routes")
routeSelector.SelectAllRoutes()
} else {
log.Debugf("select route with id: %s", id)
routes := toNetIDs([]string{id})
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
if err := routeSelector.SelectRoutes(routes, true, maps.Keys(routesMap)); err != nil {
log.Debugf("error when selecting routes: %s", err)
return fmt.Errorf("select routes: %w", err)
}
routeManager.SelectAllRoutes()
return nil
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
return nil
log.Debugf("select route with id: %s", id)
if err := routeManager.SelectRoutes(toNetIDs([]string{id}), true); err != nil {
log.Debugf("error when selecting routes: %s", err)
return err
}
return nil
}
func (c *Client) DeselectRoute(id string) error {
@@ -667,21 +661,17 @@ func (c *Client) DeselectRoute(id string) error {
}
routeManager := engine.GetRouteManager()
routeSelector := routeManager.GetRouteSelector()
if id == "All" {
log.Debugf("deselect all routes")
routeSelector.DeselectAllRoutes()
} else {
log.Debugf("deselect route with id: %s", id)
routes := toNetIDs([]string{id})
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
if err := routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
log.Debugf("error when deselecting routes: %s", err)
return fmt.Errorf("deselect routes: %w", err)
}
routeManager.DeselectAllRoutes()
return nil
}
log.Debugf("deselect route with id: %s", id)
if err := routeManager.DeselectRoutes(toNetIDs([]string{id})); err != nil {
log.Debugf("error when deselecting routes: %s", err)
return err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
return nil
}

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

@@ -8,7 +8,6 @@ import (
"sort"
"strings"
"golang.org/x/exp/maps"
"google.golang.org/grpc/codes"
gstatus "google.golang.org/grpc/status"
@@ -161,30 +160,11 @@ func (s *Server) SelectNetworks(_ context.Context, req *proto.SelectNetworksRequ
return nil, fmt.Errorf("no route manager")
}
routeSelector := routeManager.GetRouteSelector()
if req.GetAll() {
routeSelector.SelectAllRoutes()
} else {
routes := toNetIDs(req.GetNetworkIDs())
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
netIdRoutes := maps.Keys(routesMap)
if err := routeSelector.SelectRoutes(routes, req.GetAppend(), netIdRoutes); err != nil {
return nil, fmt.Errorf("select routes: %w", err)
}
// Exit nodes are mutually exclusive: if this selection activates an
// exit node, deselect every other available exit node so two can't be
// selected at once. Non-exit route selections are left untouched.
if requestActivatesExitNode(routes, routesMap) {
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
if err := routeSelector.DeselectRoutes(others, netIdRoutes); err != nil {
return nil, fmt.Errorf("deselect sibling exit nodes: %w", err)
}
}
}
routeManager.SelectAllRoutes()
} else if err := routeManager.SelectRoutes(toNetIDs(req.GetNetworkIDs()), req.GetAppend()); err != nil {
return nil, err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
s.statusRecorder.PublishEvent(
proto.SystemEvent_INFO,
@@ -224,19 +204,11 @@ func (s *Server) DeselectNetworks(_ context.Context, req *proto.SelectNetworksRe
return nil, fmt.Errorf("no route manager")
}
routeSelector := routeManager.GetRouteSelector()
if req.GetAll() {
routeSelector.DeselectAllRoutes()
} else {
routes := toNetIDs(req.GetNetworkIDs())
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
netIdRoutes := maps.Keys(routesMap)
if err := routeSelector.DeselectRoutes(routes, netIdRoutes); err != nil {
return nil, fmt.Errorf("deselect routes: %w", err)
}
routeManager.DeselectAllRoutes()
} else if err := routeManager.DeselectRoutes(toNetIDs(req.GetNetworkIDs())); err != nil {
return nil, err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
s.statusRecorder.PublishEvent(
proto.SystemEvent_INFO,
@@ -261,37 +233,3 @@ func toNetIDs(routes []string) []route.NetID {
return netIDs
}
func isExitNodeRoutes(routes []*route.Route) bool {
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
}
// requestActivatesExitNode reports whether any requested NetID maps to an exit
// node (default route) in the current route table.
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
for _, id := range requested {
if isExitNodeRoutes(routesMap[id]) {
return true
}
}
return false
}
// otherExitNodeIDs returns every available exit-node NetID that is not in the
// requested set — the siblings to deselect so a single exit node stays active.
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
keep := make(map[route.NetID]struct{}, len(requested))
for _, id := range requested {
keep[id] = struct{}{}
}
var others []route.NetID
for id, routes := range routesMap {
if !isExitNodeRoutes(routes) {
continue
}
if _, ok := keep[id]; ok {
continue
}
others = append(others, id)
}
return others
}

View File

@@ -1,26 +0,0 @@
package server
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/route"
)
func TestExitNodeSelectionHelpers(t *testing.T) {
routesMap := map[route.NetID][]*route.Route{
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
}
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
}

View File

@@ -2,6 +2,7 @@ import { type ReactNode } from "react";
import { useTranslation } from "react-i18next";
import { Browser } from "@wailsio/runtime";
import { DownloadIcon, NotepadText } from "lucide-react";
import { Update as UpdateSvc } from "@bindings/services";
import { Button } from "@/components/buttons/Button";
import { useClientVersion } from "@/contexts/ClientVersionContext";
import { cn } from "@/lib/cn";
@@ -14,6 +15,12 @@ function openUrl(url: string) {
});
}
function openInstallerDownload() {
UpdateSvc.DownloadURL()
.then(openUrl)
.catch(() => openUrl(GITHUB_RELEASES));
}
export function UpdateVersionCard() {
const { t } = useTranslation();
const { updateVersion, enforced, triggerUpdate } = useClientVersion();
@@ -37,11 +44,7 @@ export function UpdateVersionCard() {
{t("update.card.installNow")}
</Button>
) : (
<Button
variant={"primary"}
size={"xs"}
onClick={() => openUrl(GITHUB_RELEASES)}
>
<Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
<DownloadIcon size={14} />
{t("update.card.getInstaller")}
</Button>

View File

@@ -281,6 +281,9 @@ func newApplication(onSecondInstance func()) *application.App {
Linux: application.LinuxOptions{
ProgramName: "netbird",
},
Windows: application.WindowsOptions{
WndProcInterceptor: endSessionInterceptor(),
},
SingleInstance: &application.SingleInstanceOptions{
UniqueID: "io.netbird.ui",
OnSecondInstanceLaunch: func(_ application.SecondInstanceData) {
@@ -367,6 +370,9 @@ func newMainWindow(app *application.App, prefStore *preferences.Store) *applicat
// Hide instead of quit on close; "really quit" is reached via tray -> Quit.
window.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
if services.ShuttingDown() {
return
}
e.Cancel()
window.Hide()
})

View File

@@ -0,0 +1,24 @@
package services
import "sync/atomic"
var (
sessionEnding atomic.Bool
quitting atomic.Bool
)
func BeginSessionEnd() {
sessionEnding.Store(true)
}
func AbortSessionEnd() {
sessionEnding.Store(false)
}
func BeginShutdown() {
quitting.Store(true)
}
func ShuttingDown() bool {
return sessionEnding.Load() || quitting.Load()
}

View File

@@ -10,6 +10,7 @@ import (
"github.com/netbirdio/netbird/client/proto"
"github.com/netbirdio/netbird/client/ui/updater"
"github.com/netbirdio/netbird/version"
)
// UpdateResult mirrors TriggerUpdateResponse.
@@ -33,6 +34,12 @@ func (s *Update) GetState() updater.State {
return s.holder.Get()
}
// DownloadURL returns the platform-appropriate installer download link for
// manual (non-enforced) updates.
func (s *Update) DownloadURL() string {
return version.DownloadUrl()
}
// Quit exits the app. Scheduled off the calling goroutine so the JS caller's
// response returns before the runtime tears down.
func (s *Update) Quit() {

View File

@@ -154,6 +154,9 @@ func NewWindowManager(app *application.App, mainWindow *application.WebviewWindo
})
// Hide (not destroy) on close to keep React state; reset to General for a flash-free reopen.
s.settings.RegisterHook(events.Common.WindowClosing, func(e *application.WindowEvent) {
if ShuttingDown() {
return
}
e.Cancel()
s.app.Event.Emit(EventSettingsOpen, "general")
s.settings.Hide()
@@ -393,6 +396,9 @@ func (s *WindowManager) CloseWelcome() {
// OpenError shows the custom error dialog; title/message are pre-localised and ride in the
// start URL. A second error replaces the open one via SetURL. Singleton, destroyed on close.
func (s *WindowManager) OpenError(title, message string) {
if ShuttingDown() {
return
}
s.mu.Lock()
defer s.mu.Unlock()
startURL := errorDialogURL(title, message)

View File

@@ -0,0 +1,7 @@
//go:build !windows && !android && !ios && !freebsd && !js
package main
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
return nil
}

View File

@@ -0,0 +1,36 @@
//go:build windows
package main
import (
"os"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/client/ui/services"
)
const (
wmQueryEndSession = 0x0011
wmEndSession = 0x0016
)
func endSessionInterceptor() func(hwnd uintptr, msg uint32, wParam, lParam uintptr) (uintptr, bool) {
return func(_ uintptr, msg uint32, wParam, _ uintptr) (uintptr, bool) {
switch msg {
case wmQueryEndSession:
services.BeginSessionEnd()
return 1, true
case wmEndSession:
if wParam == 0 {
services.AbortSessionEnd()
return 0, true
}
log.Info("windows session is ending; exiting immediately")
os.Exit(0)
return 0, true
default:
return 0, false
}
}
}

View File

@@ -32,9 +32,8 @@ const (
quitDownTimeout = 5 * time.Second
urlGitHubRepo = "https://github.com/netbirdio/netbird"
urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
urlDocs = "https://docs.netbird.io"
urlGitHubRepo = "https://github.com/netbirdio/netbird"
urlDocs = "https://docs.netbird.io"
)
// TrayServices bundles the services the tray menu needs, grouped so NewTray
@@ -453,6 +452,7 @@ func (t *Tray) buildMenu() *application.Menu {
}
func (t *Tray) handleQuit() {
services.BeginShutdown()
t.profileMu.Lock()
if t.switchCancel != nil {
t.switchCancel()

View File

@@ -25,6 +25,9 @@ type sendFn func(notifications.NotificationOptions) error
// event-dispatch goroutine that panic is fatal process-wide; recover() turns
// it into a logged no-op.
func safeSendNotification(send sendFn, what string, opts notifications.NotificationOptions) (err error) {
if services.ShuttingDown() {
return nil
}
defer func() {
if r := recover(); r != nil {
log.Errorf("notify %s: recovered from panic (notification bus unavailable): %v", what, r)

View File

@@ -13,6 +13,7 @@ import (
"github.com/netbirdio/netbird/client/ui/services"
"github.com/netbirdio/netbird/client/ui/updater"
"github.com/netbirdio/netbird/version"
)
// trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
@@ -76,15 +77,15 @@ func (u *trayUpdater) applyLanguage() {
u.refreshMenuItem(state)
}
// handleClick opens the GitHub releases page when not Enforced, otherwise shows
// the progress page and asks the daemon to start the installer.
// handleClick opens the installer download link when not Enforced, otherwise
// shows the progress page and asks the daemon to start the installer.
func (u *trayUpdater) handleClick() {
u.mu.Lock()
state := u.state
u.mu.Unlock()
if !state.Enforced {
_ = u.app.Browser.OpenURL(urlGitHubReleases)
_ = u.app.Browser.OpenURL(version.DownloadUrl())
return
}

View File

@@ -109,7 +109,7 @@ sequenceDiagram
Chk->>Inj: continue
Inj->>Inj: inject NetBird identity headers per provider config
Inj->>Grd: continue
Grd->>Grd: enforce model allowlist
Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
Grd->>Up: forward (over WireGuard)
Up-->>Resp: response (JSON or SSE stream)
Resp->>Resp: parse usage tokens, completion
@@ -135,6 +135,21 @@ sequenceDiagram
(`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
PII names — see `redact.go` for the full set. See
[`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
is authoritative: it resolves the policy that governs this
(provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
when no applicable policy permits the model — so an allowlist scoped to
one group/provider never leaks to another, and an un-guardrailed policy
is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
backstop: it only carries an allowlist for a provider every authorising
policy restricts, and blocks unknown/undetermined models even when
management is unreachable. Because that backstop allowlist is the UNION
of every restricting policy's models, per-group narrowing lives only in
the authoritative check: during a `CheckLLMPolicyLimits` outage
`llm_limit_check` fails open, so a caller can reach any model in the
provider's union — a group scoped to model A could reach model B if
another group restricts the same provider to B. This is the documented
fail-open trade-off; a future flag may switch it to fail-closed.
- SSE streaming requires special handling on the response side; the
parser must handle partial chunks without buffering the whole
stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).

View File

@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
| on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
| on_request | 2 | `llm_limit_check` | `{}` | |
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
| on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | |
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | |
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | |
| on_response | 6 | `cost_meter` | `{}` | |
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | |

View File

@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
| `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
| `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |

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

@@ -0,0 +1,209 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
// provider, its own guardrail whose allowlist holds ONLY that provider's
// allowed model, and its own policy. Each case runs in isolation (its own
// proxy + client), so the guardrail the proxy enforces contains exactly this
// provider's model — never a mixed cross-provider list.
type pathRoutedGuardrailCase struct {
name string
catalogID string // agent-network catalog provider id
wire string // harness.WireVertex | harness.WireBedrock
allowEntry string // the single model id put on the guardrail allowlist
allowModel string // model id sent that MUST be served (200)
blockModel string // model id sent that MUST be denied (403 model_blocked)
}
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
// guard for the customer report that a model-allowlist guardrail attached to a
// policy has no effect for PATH-ROUTED providers — where the model travels in
// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
// AWS Bedrock (/model/{id}/invoke).
//
// Each provider is tested in isolation with a guardrail allowlisting a single
// model of its own: the allowed model (in the URL path) is served (200) and an
// unselected model (in the URL path) is denied 403 by the guardrail
// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
// customer verbatim — allow Sonnet, and the unselected model is the exact
// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
// sends a region-prefixed, versioned inference-profile id so URL-path model
// normalization is exercised too.
//
// The provider is catch-all (no models), so the router forwards any model and a
// 403 can only come from the guardrail, never model_not_routable. Only the
// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
// synth/reconcile, the proxy middleware chain (URL-path model extraction,
// router, guardrail) and the tunnel are all real, and the guardrail denies
// before the upstream is dialed so the mock cannot influence the block. A
// static bearer api key is used so the router injects a static Authorization
// header instead of minting a GCP token — the only reason path-routed providers
// normally need live credentials — so the test runs with none and is always on.
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
cases := []pathRoutedGuardrailCase{
{
name: "vertex",
catalogID: "vertex_ai_api",
wire: harness.WireVertex,
allowEntry: "claude-sonnet-4-5",
allowModel: "claude-sonnet-4-5",
blockModel: "claude-opus-4-6", // the customer-reported model
},
{
name: "bedrock",
catalogID: "bedrock_api",
wire: harness.WireBedrock,
allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
blockModel: "us.anthropic.claude-opus-4-8-v1:0",
},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
runPathRoutedGuardrailCase(t, tc)
})
}
}
func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
t.Helper()
const (
vertexProject = "e2e-project"
vertexRegion = "global"
)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
require.NoError(t, err, "create group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-guardrail-" + tc.name + "-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grp.Id},
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
// Catch-all provider (no models) so the router forwards any model; a static
// bearer key means the router injects a static auth header instead of minting
// a GCP token. Bootstraps the cluster if it isn't already.
staticKey := "static-e2e-token"
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: tc.name,
ProviderId: tc.catalogID,
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create %s provider", tc.name)
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
// Guardrail allowlisting ONLY this provider's allowed model.
var gr api.AgentNetworkGuardrailRequest
gr.Name = "e2e-guardrail-" + tc.name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
guard, err := srv.CreateGuardrail(ctx, gr)
require.NoError(t, err, "create guardrail")
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
enabled := true
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-" + tc.name,
Enabled: &enabled,
SourceGroups: []string{grp.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{guard.Id},
})
require.NoError(t, err, "create policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Probe first: the GET resolves the endpoint and its first packet wakes the
// lazy proxy peer, so WaitProxyPeer then observes it connected.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
var code int
var body string
var cerr error
switch tc.wire {
case harness.WireVertex:
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
case harness.WireBedrock:
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
default:
t.Fatalf("unsupported wire %q", tc.wire)
}
require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
return code, body
}
// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
// jitter on the first call over the freshly warmed tunnel.
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(tc.allowModel)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
assert.Equal(t, 200, code,
"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
// Unselected model (in the URL path) must be blocked by the guardrail.
code, body = send(tc.blockModel)
assert.Equal(t, 403, code,
"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
}

View File

@@ -0,0 +1,205 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
// the proxy caches; the switch is invisible until that cache expires. The proxy
// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
// instead of the 5-minute default.
//
// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
// served and modelB denied; after switching the client grpA -> grpB, modelB is
// served and modelA denied. The cross-group deny comes from management's
// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
modelA = "e2e-model-a"
modelB = "e2e-model-b"
// Short tunnel-cache TTL so a group switch propagates in seconds.
// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
cacheTTL = 3 * time.Second
)
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
require.NoError(t, err, "create group A")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
require.NoError(t, err, "create group B")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-gswitch-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpA.Id}, // client starts in group A
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "gswitch",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: &[]api.AgentNetworkProviderModel{
{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
},
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
mkGuard := func(name, model string) api.AgentNetworkGuardrail {
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, gerr := srv.CreateGuardrail(ctx, gr)
require.NoError(t, gerr, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}
gA := mkGuard("e2e-gswitch-a", modelA)
gB := mkGuard("e2e-gswitch-b", modelB)
enabled := true
polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-gswitch-a",
Enabled: &enabled,
SourceGroups: []string{grpA.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gA.Id},
})
require.NoError(t, err, "create policy A")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-gswitch-b",
Enabled: &enabled,
SourceGroups: []string{grpB.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gB.Id},
})
require.NoError(t, err, "create policy B")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
})
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
var code int
var body string
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
code, body = send(model)
if code == want {
return code, body
}
time.Sleep(2 * time.Second)
}
return code, body
}
// Phase 1 — client is in group A: modelA served, modelB denied.
code, body := sendUntil(modelA, 200, 90*time.Second)
assert.Equal(t, 200, code,
"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = send(modelB)
assert.Equal(t, 403, code,
"group-B model must be denied while the client is in group A; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
// Switch the client peer from group A to group B.
peerID := clientPeerInGroup(t, ctx, grpA.Id)
_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
Name: grpB.Name,
Peers: &[]string{peerID},
})
require.NoError(t, err, "add peer to group B")
_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
Name: grpA.Name,
Peers: &[]string{},
})
require.NoError(t, err, "remove peer from group A")
// Phase 2 — after the short TTL expires the proxy re-validates the peer,
// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
code, body = sendUntil(modelB, 200, 60*time.Second)
assert.Equal(t, 200, code,
"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = send(modelA)
assert.Equal(t, 403, code,
"after the switch, the old group-A model must be denied; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
}
// clientPeerInGroup returns the id of the single peer that is a member of the
// given group — the test client. The proxy peer is never added to test groups.
func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
t.Helper()
peers, err := srv.API().Peers.List(ctx)
require.NoError(t, err, "list peers")
for _, p := range peers {
for _, g := range p.Groups {
if g.Id == groupID {
return p.Id
}
}
}
t.Fatalf("no peer found in group %s", groupID)
return ""
}

View File

@@ -0,0 +1,201 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
// modelOther denied for the grpMain client (403 model_blocked, no cross-group
// leak), and openModel on the un-guardrailed policy's provider served (200).
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
modelSelected = "e2e-selected"
modelOther = "e2e-other"
openModel = "e2e-open"
)
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
require.NoError(t, err, "create main group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
require.NoError(t, err, "create other group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-guardrail-mp-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
// pRestricted declares the two guardrailed models so routing is deterministic
// (model -> provider). Created first, so it carries the bootstrap cluster.
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "restricted",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: models(modelSelected, modelOther),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create restricted provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "open",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: models(openModel),
})
require.NoError(t, err, "create open provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, gerr := srv.CreateGuardrail(ctx, gr)
require.NoError(t, gerr, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
enabled := true
// polMain: grpMain restricted to modelSelected on pRestricted.
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-main",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{pRestricted.Id},
GuardrailIds: &[]string{gMain.Id},
})
require.NoError(t, err, "create main policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
// polOther: grpOther restricted to modelOther on the SAME provider. The
// client is not in grpOther, so modelOther must never be usable by it.
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-other",
Enabled: &enabled,
SourceGroups: []string{grpOther.Id},
DestinationProviderIds: []string{pRestricted.Id},
GuardrailIds: &[]string{gOther.Id},
})
require.NoError(t, err, "create other policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-open",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{pOpen.Id},
})
require.NoError(t, err, "create open policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
sendUntil200 := func(model string) (int, string) {
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(model)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
return code, body
}
t.Run("selected model allowed for its group", func(t *testing.T) {
code, body := sendUntil200(modelSelected)
assert.Equal(t, 200, code,
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
t.Run("other group's model does not leak", func(t *testing.T) {
// modelOther is allowlisted only for grpOther. The grpMain client must be
// denied by management's per-policy/group check — not waved through by an
// account-wide union. This is the security-critical wrong-ALLOW guard.
code, body := send(modelOther)
assert.Equal(t, 403, code,
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"denial must be a model-allowlist decision; body: %s", body)
})
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
// old account-wide union would have blocked openModel (it is on no
// allowlist); it must now be served — the false-DENY guard.
code, body := sendUntil200(openModel)
assert.Equal(t, 200, code,
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
}

View File

@@ -0,0 +1,422 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// pergroupCase describes one provider surface for the per-group allowlist matrix.
// selectedReq/otherReq are the model identifiers as they travel in the request
// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
// otherAllow are the (normalized) forms the guardrail allowlist holds — for
// Bedrock these differ from the request form so path normalization is exercised.
type pergroupCase struct {
name string
catalogID string
wire string // "chat", "messages", "vertex", "bedrock"
models *[]api.AgentNetworkProviderModel
selectedReq string
selectedAllow string
otherReq string
otherAllow string
providerID string // filled during setup
}
// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
// allowlist end to end across every always-on provider surface, including the
// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
//
// For each provider two policies target it: grpMain (the client) is allowed only
// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
// The client must get selectedReq served (200) and otherReq denied (403,
// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
// authoritative per-policy/group decision from management (the proxy per-provider
// backstop carries the union of both models), so this also confirms management
// receives the correct normalized model for path-routed providers.
func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
defer cancel()
const (
vertexProject = "e2e-project"
vertexRegion = "global"
)
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
cases := []*pergroupCase{
{
name: "openai", catalogID: "openai_api", wire: harness.WireChat,
models: priced("oai-model-a", "oai-model-b"),
selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
otherReq: "oai-model-b", otherAllow: "oai-model-b",
},
{
name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
models: priced("ant-model-a", "ant-model-b"),
selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
otherReq: "ant-model-b", otherAllow: "ant-model-b",
},
{
// Vertex catalog ids travel bare in the rawPredict path.
name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
},
{
// Bedrock request ids are region-prefixed/versioned; the parser
// normalizes them to the catalog key the allowlist holds.
name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
},
}
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
require.NoError(t, err, "create main group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
require.NoError(t, err, "create other group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-pergroup-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpMain.Id},
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
enabled := true
for i, c := range cases {
req := api.AgentNetworkProviderRequest{
Name: "e2e-pergroup-" + c.name,
ProviderId: c.catalogID,
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: c.models,
}
if i == 0 {
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
}
prov, perr := srv.CreateProvider(ctx, req)
require.NoError(t, perr, "create provider %s", c.name)
c.providerID = prov.Id
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-pergroup-" + c.name + "-main",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gSel.Id},
})
require.NoError(t, merr, "create main policy %s", c.name)
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-pergroup-" + c.name + "-other",
Enabled: &enabled,
SourceGroups: []string{grpOther.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gOth.Id},
})
require.NoError(t, oerr, "create other policy %s", c.name)
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
}
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(c *pergroupCase, model string) (int, string) {
var code int
var body string
var cerr error
switch c.wire {
case "vertex":
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
case "bedrock":
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
default:
code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
}
require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
return code, body
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
// the first call over the freshly warmed tunnel.
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(c, c.selectedReq)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
assert.Equal(t, 200, code,
"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
// grpOther's model must NOT leak to the grpMain client.
code, body = send(c, c.otherReq)
assert.Equal(t, 403, code,
"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
})
}
}
// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
//
// - union across the user's groups: the client is in gUX and gUY, each with
// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
// client may use BOTH models (the union of its groups' allowlists) while a
// third, un-allowlisted model is denied.
// - an un-guardrailed group lifts the restriction: the client is in gMP and
// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
// NO guardrail. Because one applicable policy is unrestricted, the client may
// use a model on no allowlist (mix-z) as well as mix-a.
func TestGuardrailMultiGroupUser(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
unionA = "mg-union-a"
unionB = "mg-union-b"
unionC = "mg-union-c" // allowlisted by neither group
mixA = "mg-mix-a"
mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
)
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
mkGroup := func(name string) *api.Group {
g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
require.NoError(t, gerr, "create group %s", name)
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
return g
}
gUX := mkGroup("e2e-mg-union-x")
gUY := mkGroup("e2e-mg-union-y")
gMP := mkGroup("e2e-mg-mix-p")
gMQ := mkGroup("e2e-mg-mix-q")
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-mg-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
enabled := true
// P1 — union scenario: two restricting policies, one per group.
p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "e2e-mg-union",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: priced(unionA, unionB, unionC),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create union provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-union-x",
Enabled: &enabled,
SourceGroups: []string{gUX.Id},
DestinationProviderIds: []string{p1.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
})
require.NoError(t, err, "create union policy X")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-union-y",
Enabled: &enabled,
SourceGroups: []string{gUY.Id},
DestinationProviderIds: []string{p1.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
})
require.NoError(t, err, "create union policy Y")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "e2e-mg-mix",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: priced(mixA, mixZ),
})
require.NoError(t, err, "create mix provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-mix-p",
Enabled: &enabled,
SourceGroups: []string{gMP.Id},
DestinationProviderIds: []string{p2.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
})
require.NoError(t, err, "create mix policy P")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-mix-q",
Enabled: &enabled,
SourceGroups: []string{gMQ.Id},
DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
})
require.NoError(t, err, "create mix policy Q")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
sendUntil200 := func(model string) (int, string) {
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(model)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
return code, body
}
t.Run("union across the user's groups", func(t *testing.T) {
code, body := sendUntil200(unionA)
assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = sendUntil200(unionB)
assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
code, body = send(unionC)
assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
})
t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
code, body := sendUntil200(mixA)
assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = sendUntil200(mixZ)
assert.Equal(t, 200, code,
"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
}
// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
// exactly the given model, registering cleanup.
func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
t.Helper()
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, err := srv.CreateGuardrail(ctx, gr)
require.NoError(t, err, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}

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

@@ -38,7 +38,11 @@ type Proxy struct {
// network, registered via the given account proxy token and serving the
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
// peer connectivity — callers poll management for the proxy peer.
func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
// StartProxy launches the reverse-proxy container. Optional envOverrides are
// merged into the container environment after the defaults, so callers can set
// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
// need a short authorization-cache window).
func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
root, err := repoRoot()
if err != nil {
return nil, err
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
}
for _, ov := range envOverrides {
for k, v := range ov {
req.Env[k] = v
}
}
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
ContainerRequest: req,
Started: true,

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

@@ -86,12 +86,17 @@ type Manager interface {
// PolicySelectionInput is the per-request selection envelope. The
// proxy populates it from CapturedData (account, user, groups) plus
// the provider llm_router resolved.
// the provider llm_router resolved and the model it extracted.
type PolicySelectionInput struct {
AccountID string
UserID string
GroupIDs []string
ProviderID string
// Model is the already-normalised upstream model id the proxy extracted
// (parser strips Bedrock region/version, Vertex @version), so a
// case-insensitive compare suffices. Empty = undetermined → not permitted
// (fail closed).
Model string
}
// PolicySelectionResult names the policy that "pays" for this request

View File

@@ -5,6 +5,7 @@ import (
"fmt"
"math"
"sort"
"strings"
"time"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
@@ -35,6 +36,10 @@ const (
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
//nolint:gosec // account deny code label, not a credential
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
// denyCodeModelBlocked is returned when policies govern the request's
// (provider, caller-groups) but none permits the model. Matches the proxy
// guardrail's code so both layers surface the same label.
denyCodeModelBlocked = "llm_policy.model_blocked"
)
// consumptionCache holds the consumption counters prefetched for one
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
}
candidates := filterApplicablePolicies(policies, in)
// Model-allowlist gate scoped to the matched policies: keep candidates whose
// guardrails permit the model (none enabled = unrestricted), deny when
// policies apply but none permits it. Skip the load when none has a guardrail.
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
if gErr != nil {
return nil, gErr
}
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
if len(permitted) == 0 {
return &PolicySelectionResult{
Allow: false,
DenyCode: denyCodeModelBlocked,
DenyReason: modelBlockedReason(in.Model),
}, nil
}
candidates = permitted
}
// Prefetch every consumption counter the ceiling + candidate policies will
// read, in a single store round-trip, then score against the cache.
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
return out
}
// anyPolicyHasGuardrails reports whether any policy references at least one
// guardrail, so the selector can skip loading guardrails when none do.
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
for _, p := range policies {
if p != nil && len(p.GuardrailIDs) > 0 {
return true
}
}
return false
}
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
// model-allowlist gate to resolve each candidate policy's attached guardrails.
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return nil, fmt.Errorf("list account guardrails: %w", err)
}
byID := make(map[string]*types.Guardrail, len(guardrails))
for _, g := range guardrails {
if g != nil {
byID[g.ID] = g
}
}
return byID, nil
}
// filterModelPermittedPolicies returns the subset of policies whose guardrails
// permit the model. Order is preserved so downstream scoring is unaffected.
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
out := make([]*types.Policy, 0, len(policies))
for _, p := range policies {
if policyPermitsModel(p, byID, model) {
out = append(out, p)
}
}
return out
}
// policyPermitsModel reports whether a policy permits the model. No
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
// otherwise the model must be in the union of its allowlists, so an
// empty/undetermined model fails closed.
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
if p == nil {
return false
}
wanted := normaliseModelID(model)
restricted := false
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
if wanted == "" {
continue
}
for _, allowed := range g.Checks.ModelAllowlist.Models {
if normaliseModelID(allowed) == wanted {
return true
}
}
}
return !restricted
}
// normaliseModelID lowercases and trims a model identifier so the allowlist
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
// normaliseModel so both layers agree on what "same model" means.
func normaliseModelID(model string) string {
return strings.ToLower(strings.TrimSpace(model))
}
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
// rejection. The model is quoted when known; an undetermined model is reported
// as such so the access log distinguishes "wrong model" from "no model".
func modelBlockedReason(model string) string {
if normaliseModelID(model) == "" {
return "request model could not be determined for the policy allowlist"
}
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
}
// candidate is the per-policy intermediate the selector ranks. A
// policy that's been exhausted on any enabled cap never makes it
// into this slice; the selector's deny envelope carries the latest

View File

@@ -0,0 +1,329 @@
package agentnetwork
import (
"context"
"errors"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/store"
)
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
// to reach providerID under the given guardrails. Uncapped keeps the selector's
// headroom scoring trivial so these tests isolate the model-allowlist gate.
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
AccountID: account,
Enabled: true,
SourceGroups: sourceGroups,
DestinationProviderIDs: []string{providerID},
GuardrailIDs: guardrailIDs,
CreatedAt: time.Now().UTC(),
}
}
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
// carries the given models.
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
return &types.Guardrail{
ID: id,
AccountID: account,
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
},
}
}
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
mockStore.EXPECT().
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
Return(policies, nil)
}
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
Return(guardrails, nil)
}
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
// decision: a policy authorises the (provider, group) but restricts the model,
// and the requested model isn't on the list, so the request is denied.
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
}
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
// is on the applicable policy's allowlist, so selection proceeds normally.
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
// and surrounding whitespace, matching the proxy guardrail's normalisation.
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.NoError(t, err)
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
}
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
// policies authorise the same (provider, group) and one has no guardrail, that
// policy makes the request unrestricted — not caught by the other's allowlist.
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
expectPolicies(mockStore, "acc-1", restricted, open)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
}
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
// model allowlisted only for grp-b must not be usable by a grp-a caller. The
// selector considers only policies applicable to the caller's groups.
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
expectPolicies(mockStore, "acc-1", polA, polB)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-a"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only allowed for grp-b
})
require.NoError(t, err)
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
// mirrors the proxy: with a restricted applicable policy and an empty model
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "", // undetermined
})
require.NoError(t, err)
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
// model allowlist is disabled imposes no model restriction, even though the
// policy references it.
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
disabled := &types.Guardrail{
ID: "g-1",
AccountID: "acc-1",
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
},
}
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", disabled)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
// allowlist guardrails permits the union of their models (not just the first).
func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only in the second guardrail's list
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
// resolving the candidate policies' guardrails surfaces as an error, not a
// silent allow/deny.
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
Return(nil, errors.New("store unavailable"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
assert.Nil(t, res)
}
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
// policy referencing a guardrail ID absent from the account's set (a stale
// reference) imposes no model restriction — same as no guardrail.
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1")
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
// gate narrows candidates before cap scoring: the permitting policy is selected
// even though the blocked one has a larger, more attractive cap.
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
polBig.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
}
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
polSmall.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
}
expectPolicies(mockStore, "acc-1", polBig, polSmall)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only pol-small's guardrail permits this
})
require.NoError(t, err)
assert.True(t, res.Allow)
assert.Equal(t, "pol-small", res.SelectedPolicyID,
"the model filter must exclude pol-big before cap scoring")
}

View File

@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
applyAccountCollectionControls(&mergedGuardrails, settings)
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
// The proxy guardrail is a per-provider fail-closed backstop; the
// authoritative per-policy/group decision is management's
// SelectPolicyForRequest. A provider lands in this map only when every
// authorising policy restricts models.
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
if err != nil {
return nil, err
}
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
// middleware expects. Mirrors the proxy registration documented in
// the management→proxy contract.
// the management→proxy contract. provider_allowlists is keyed by the
// resolved provider id llm_router stamps; a provider absent from the map is
// unrestricted at the proxy layer.
type guardrailConfig struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
}
type guardrailPromptCapture struct {
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
}
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
cfg := guardrailConfig{
ModelAllowlist: merged.ModelAllowlist,
PromptCapture: guardrailPromptCapture{
Enabled: merged.PromptCapture.Enabled,
RedactPii: merged.PromptCapture.RedactPii,
},
ProviderAllowlists: providerAllowlists,
PromptCapture: guardrailPromptCapture(capture),
}
out, err := json.Marshal(cfg)
if err != nil {
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
return out, nil
}
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
// is included only when every authorising policy restricts models (their union);
// if any leaves it unrestricted it is omitted, so management decides per group.
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
type providerAcc struct {
models map[string]struct{}
anyUnrestricted bool
}
accs := make(map[string]*providerAcc)
for _, p := range policies {
if p == nil {
continue
}
restricted, models := policyModelAllowlist(p, byID)
for _, providerID := range p.DestinationProviderIDs {
if providerID == "" {
continue
}
acc, ok := accs[providerID]
if !ok {
acc = &providerAcc{models: make(map[string]struct{})}
accs[providerID] = acc
}
if !restricted {
acc.anyUnrestricted = true
continue
}
for _, m := range models {
acc.models[m] = struct{}{}
}
}
}
out := make(map[string][]string, len(accs))
for providerID, acc := range accs {
if acc.anyUnrestricted {
continue
}
models := make([]string, 0, len(acc.models))
for m := range acc.models {
models = append(models, m)
}
sort.Strings(models)
out[providerID] = models
}
return out
}
// policyModelAllowlist reports whether a policy restricts models (has an
// allowlist-enabled guardrail) and the union of allowed models. Models are
// verbatim; the proxy factory lowercases/trims them at decode time.
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
restricted := false
var models []string
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
models = append(models, m)
}
}
}
return restricted, models
}
// buildAccountService composes the per-account gateway Service. The
// target carries the noop placeholder URL — the router middleware
// rewrites every request to the matched provider's upstream before the
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
return out
}
// MergedGuardrails is the JSON shape passed to the proxy via the
// guardrail middleware's config_json. Mirrors the proxy-side
// expectations and is intentionally distinct from
// types.GuardrailChecks so we can evolve either side independently.
// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
// merged here — the model allowlist is emitted per-provider, and
// token/budget/retention moved onto Policy.Limits and account Settings.
type MergedGuardrails struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
TokenLimits MergedTokenLimits `json:"token_limits"`
Budget MergedBudget `json:"budget"`
PromptCapture MergedPromptCapture `json:"prompt_capture"`
Retention MergedRetention `json:"retention"`
}
type MergedTokenLimits struct {
Hourly *MergedTokenWindow `json:"hourly,omitempty"`
Daily *MergedTokenWindow `json:"daily,omitempty"`
Monthly *MergedTokenWindow `json:"monthly,omitempty"`
}
type MergedTokenWindow struct {
MaxInputTokens int `json:"max_input_tokens,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
}
type MergedBudget struct {
Hourly *MergedBudgetWindow `json:"hourly,omitempty"`
Daily *MergedBudgetWindow `json:"daily,omitempty"`
Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
}
type MergedBudgetWindow struct {
SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
PromptCapture MergedPromptCapture
}
type MergedPromptCapture struct {
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
RedactPii bool `json:"redact_pii"`
}
type MergedRetention struct {
Enabled bool `json:"enabled"`
Days int `json:"days"`
}
// mergeGuardrails computes the effective guardrail spec applied at the
// proxy, given the referencing policies and the account's guardrail
// catalogue. Policy enabled-ness is the caller's responsibility — only
// enabled policies should be passed in.
// mergeGuardrails folds the referencing policies' guardrails into the
// prompt-capture decision only. The model allowlist is enforced per-policy/group
// in management and shipped per-provider; token/budget/retention live off
// guardrails now.
//
// Merge rules:
// - Model allowlist: union of allowlists across policies that enable it.
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
// if any enabling policy turns it on.
// - Retention: enabled if any enables it; smallest non-zero days wins.
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
// sticks if any enabling policy turns it on.
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
merged := MergedGuardrails{}
allowlist := make(map[string]struct{})
allowlistEnabled := false
for _, policy := range policies {
for _, gID := range policy.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil {
continue
}
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
mergeGuardrail(g, &merged)
}
}
if allowlistEnabled {
merged.ModelAllowlist = make([]string, 0, len(allowlist))
for m := range allowlist {
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
}
sort.Strings(merged.ModelAllowlist)
}
return merged
}
// mergeGuardrail folds a single guardrail's enabled checks into the
// running merge: model-allowlist models join the shared set (and flip
// allowlistEnabled), and prompt-capture / redact-pii stick once any
// enabling guardrail turns them on.
//
// TokenLimits, Budget, and Retention have moved off guardrails — token
// and budget caps now live on the Policy itself (Policy.Limits) and
// retention moves to account-level Settings — so they are not merged here.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
if g.Checks.ModelAllowlist.Enabled {
*allowlistEnabled = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
allowlist[m] = struct{}{}
}
}
}
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
// running merge: enabled / redact-pii stick once any enabling guardrail turns
// them on.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
if g.Checks.PromptCapture.Enabled {
merged.PromptCapture.Enabled = true
if g.Checks.PromptCapture.RedactPii {

View File

@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
require.Len(t, services, 1)
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
"model allowlist is a pure policy guardrail and must always reach the config")
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
"model allowlist is a pure policy guardrail and must reach the per-provider config")
assert.False(t, cfg.PromptCapture.Enabled,
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
}
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
require.Len(t, services, 1, "exactly one synth service expected")
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
}

View File

@@ -0,0 +1,95 @@
package agentnetwork
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// policyForProviders builds an enabled policy authorising the given providers
// under the given guardrails (both optional). Groups are irrelevant to
// buildProviderAllowlists, which keys purely on destination provider.
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
Enabled: true,
DestinationProviderIDs: providerIDs,
GuardrailIDs: guardrailIDs,
}
}
func TestBuildProviderAllowlists(t *testing.T) {
byID := map[string]*types.Guardrail{
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
}
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
"a provider every policy restricts carries the sorted union of their models")
})
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", nil, "prov-x"), // no guardrail
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
})
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a policy whose only guardrail has a disabled allowlist is unrestricted")
})
t.Run("providers are isolated from one another", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
})
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
})
t.Run("union across a single policy's guardrails", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
"a policy's own multiple allowlist guardrails union together")
})
t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
got := buildProviderAllowlists([]*types.Policy{
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
}, empty)
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
})
}

View File

@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
var cfg guardrailConfig
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
"model allowlist union must keep both models")
// Both policies restrict the same provider, so the per-provider backstop
// carries the union of their models — a coarse gate that management's
// per-policy/group check narrows; it only blocks models outside the union
// when management is down.
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
"per-provider allowlist union must keep both models")
}
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {

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

@@ -50,7 +50,7 @@ func ToComponentSyncResponse(
// TODO (dmitri) consider using invariants?
//
enableSSH := computeSSHEnabledForPeer(components, peer)
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH)
peerConfig := toPeerConfig(peer, components.Network, dnsName, settings, httpConfig, deviceFlowConfig, enableSSH, components.ForceRoutingPeerDNSResolution)
includeIPv6 := peer.SupportsIPv6() && peer.IPv6.IsValid()
useSourcePrefixes := peer.SupportsSourcePrefixes()

View File

@@ -119,7 +119,7 @@ func toNetbirdConfig(config *nbconfig.Config, turnCredentials *Token, relayToken
return nbConfig
}
func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool) *proto.PeerConfig {
func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, settings *types.Settings, httpConfig *nbconfig.HttpServerConfig, deviceFlowConfig *nbconfig.DeviceAuthorizationFlow, enableSSH bool, forceRoutingPeerDNS bool) *proto.PeerConfig {
netmask, _ := network.Net.Mask.Size()
fqdn := peer.FQDN(dnsName)
@@ -135,7 +135,7 @@ func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, set
Address: fmt.Sprintf("%s/%d", peer.IP.String(), netmask),
SshConfig: sshConfig,
Fqdn: fqdn,
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled,
RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled || peer.ProxyMeta.Embedded || forceRoutingPeerDNS,
LazyConnectionEnabled: settings.LazyConnectionEnabled,
AutoUpdate: &proto.AutoUpdateSettings{
Version: settings.AutoUpdateVersion,
@@ -162,12 +162,12 @@ func ToSyncResponse(ctx context.Context, config *nbconfig.Config, httpConfig *nb
useSourcePrefixes := peer.SupportsSourcePrefixes()
response := &proto.SyncResponse{
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
NetworkMap: &proto.NetworkMap{
Serial: networkMap.Network.CurrentSerial(),
Routes: networkmap.ToProtocolRoutes(networkMap.Routes),
DNSConfig: networkmap.ToProtocolDNSConfig(networkMap.DNSConfig, dnsCache, dnsFwdPort),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH),
PeerConfig: toPeerConfig(peer, networkMap.Network, dnsName, settings, httpConfig, deviceFlowConfig, networkMap.EnableSSH, networkMap.ForceRoutingPeerDNSResolution),
},
Checks: toProtocolChecks(ctx, checks),
}

View File

@@ -2,6 +2,7 @@ package grpc
import (
"fmt"
"net"
"net/netip"
"reflect"
"testing"
@@ -14,6 +15,7 @@ import (
"github.com/netbirdio/netbird/management/internals/controllers/network_map"
"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller/cache"
nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/shared/management/networkmap"
)
@@ -301,3 +303,35 @@ func TestToNetbirdConfig_RelayInvariant(t *testing.T) {
assert.True(t, nbCfg.Metrics.Enabled, "metrics flag should carry the settings value")
})
}
func TestToPeerConfig_RoutingPeerDNSResolution(t *testing.T) {
network := &types.Network{Net: net.IPNet{IP: net.IPv4(100, 0, 0, 0), Mask: net.CIDRMask(8, 32)}}
newPeer := func(embedded bool) *nbpeer.Peer {
p := &nbpeer.Peer{IP: netip.MustParseAddr("100.0.0.1")}
p.ProxyMeta.Embedded = embedded
return p
}
tests := []struct {
name string
globalFlag bool
embedded bool
forceParam bool
wantEnabled bool
}{
{name: "global off, regular peer, no force", wantEnabled: false},
{name: "global on wins", globalFlag: true, wantEnabled: true},
{name: "embedded proxy peer forced", embedded: true, wantEnabled: true},
{name: "routing peer forced via param", forceParam: true, wantEnabled: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
settings := &types.Settings{RoutingPeerDNSResolutionEnabled: tt.globalFlag}
cfg := toPeerConfig(newPeer(tt.embedded), network, "netbird.selfhosted", settings, nil, nil, false, tt.forceParam)
assert.Equal(t, tt.wantEnabled, cfg.RoutingPeerDnsResolutionEnabled,
"RoutingPeerDnsResolutionEnabled should reflect global || embedded || forced")
})
}
}

View File

@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
UserID: req.GetUserId(),
GroupIDs: req.GetGroupIds(),
ProviderID: req.GetProviderId(),
Model: req.GetModel(),
})
if err != nil {
log.WithContext(ctx).Errorf("select policy for request: %v", err)

View File

@@ -0,0 +1,138 @@
package grpc
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/codes"
grpcstatus "google.golang.org/grpc/status"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/shared/management/proto"
)
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
// and returns a pre-programmed result, so tests can assert what the handler
// forwards to the selector.
type fakeAgentNetworkLimits struct {
gotInput agentnetwork.PolicySelectionInput
result *agentnetwork.PolicySelectionResult
err error
}
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
f.gotInput = in
if f.err != nil {
return nil, f.err
}
return f.result, nil
}
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
return nil
}
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
// the model the proxy extracted must reach the selector's Model unchanged,
// alongside the account/user/group/provider fields.
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
UserId: "user-1",
GroupIds: []string{"grp-a", "grp-b"},
ProviderId: "prov-1",
Model: "claude-opus-4",
}
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
assert.Equal(t, "user-1", fake.gotInput.UserID)
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
"the request's model must be forwarded to the selector")
}
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
// model (empty string) is forwarded as-is; the selector decides how to treat it.
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
}
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
}
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
// envelope surfaces the model-allowlist deny code + reason through the response.
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
Allow: false,
DenyCode: "llm_policy.model_blocked",
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "deny", resp.Decision)
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
assert.NotEmpty(t, resp.DenyReason)
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
}
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
// failure surfaces as an Internal gRPC error rather than a silent allow.
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
}
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
// fallback: with no limits service wired the RPC returns Unimplemented.
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
s := &ProxyServiceServer{}
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Unimplemented, st.Code())
}

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, network, s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH),
PeerConfig: toPeerConfig(peer, network, s.networkMapController.GetDNSDomain(settings), settings, s.config.HttpConfig, s.config.DeviceAuthorizationFlow, enableSSH, false),
Checks: toProtocolChecks(ctx, postureChecks),
}

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

@@ -6,6 +6,7 @@ import (
"encoding/json"
"errors"
"fmt"
"math"
"net"
"net/netip"
"net/url"
@@ -2258,117 +2259,30 @@ func (s *SqlStore) getPostureChecks(ctx context.Context, accountID string) ([]*p
return checks, nil
}
func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpservice.Service, error) {
const serviceQuery = `SELECT id, account_id, name, domain, enabled, auth,
meta_created_at, meta_certificate_issued_at, meta_status, proxy_cluster,
pass_host_header, rewrite_redirects, session_private_key, session_public_key,
mode, listen_port, port_auto_assigned, source, source_peer, terminated,
private, access_groups
FROM services WHERE account_id = $1`
// serviceSelectColumns and targetSelectColumns are the column lists the Postgres
// pgx read path scans. They must stay in sync with the rpservice.Service and
// rpservice.Target gorm models; TestPgxServiceColumnsMatchGorm enforces this.
const serviceSelectColumns = `id, account_id, name, domain, enabled, auth, restrictions,
meta_created_at, meta_certificate_issued_at, meta_last_renewed_at, meta_status, proxy_cluster,
pass_host_header, rewrite_redirects, session_private_key, session_public_key,
mode, listen_port, port_auto_assigned, source, source_peer, terminated,
private, access_groups`
const targetsQuery = `SELECT id, account_id, service_id, path, host, port, protocol,
target_id, target_type, enabled
FROM targets WHERE service_id = ANY($1)`
const targetSelectColumns = `id, account_id, service_id, path, host, port, protocol,
target_id, target_type, enabled, proxy_protocol,
skip_tls_verify, request_timeout, session_idle_timeout, path_rewrite, custom_headers,
direct_upstream, middlewares, capture_max_request_bytes, capture_max_response_bytes,
capture_content_types, agent_network, disable_access_log`
func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpservice.Service, error) {
const serviceQuery = `SELECT ` + serviceSelectColumns + ` FROM services WHERE account_id = $1`
serviceRows, err := s.pool.Query(ctx, serviceQuery, accountID)
if err != nil {
return nil, err
}
services, err := pgx.CollectRows(serviceRows, func(row pgx.CollectableRow) (*rpservice.Service, error) {
var s rpservice.Service
var auth []byte
var accessGroups []byte
var createdAt, certIssuedAt sql.NullTime
var status, proxyCluster, sessionPrivateKey, sessionPublicKey sql.NullString
var mode, source, sourcePeer sql.NullString
var terminated, portAutoAssigned, private sql.NullBool
var listenPort sql.NullInt64
err := row.Scan(
&s.ID,
&s.AccountID,
&s.Name,
&s.Domain,
&s.Enabled,
&auth,
&createdAt,
&certIssuedAt,
&status,
&proxyCluster,
&s.PassHostHeader,
&s.RewriteRedirects,
&sessionPrivateKey,
&sessionPublicKey,
&mode,
&listenPort,
&portAutoAssigned,
&source,
&sourcePeer,
&terminated,
&private,
&accessGroups,
)
if err != nil {
return nil, err
}
if auth != nil {
if err := json.Unmarshal(auth, &s.Auth); err != nil {
return nil, err
}
}
if len(accessGroups) > 0 {
if err := json.Unmarshal(accessGroups, &s.AccessGroups); err != nil {
return nil, fmt.Errorf("unmarshal access_groups: %w", err)
}
}
if private.Valid {
s.Private = private.Bool
}
s.Meta = rpservice.Meta{}
if createdAt.Valid {
s.Meta.CreatedAt = createdAt.Time
}
if certIssuedAt.Valid {
t := certIssuedAt.Time
s.Meta.CertificateIssuedAt = &t
}
if status.Valid {
s.Meta.Status = status.String
}
if proxyCluster.Valid {
s.ProxyCluster = proxyCluster.String
}
if sessionPrivateKey.Valid {
s.SessionPrivateKey = sessionPrivateKey.String
}
if sessionPublicKey.Valid {
s.SessionPublicKey = sessionPublicKey.String
}
if mode.Valid {
s.Mode = mode.String
}
if source.Valid {
s.Source = source.String
}
if sourcePeer.Valid {
s.SourcePeer = sourcePeer.String
}
if terminated.Valid {
s.Terminated = terminated.Bool
}
if portAutoAssigned.Valid {
s.PortAutoAssigned = portAutoAssigned.Bool
}
if listenPort.Valid {
s.ListenPort = uint16(listenPort.Int64)
}
s.Targets = []*rpservice.Target{}
return &s, nil
})
services, err := pgx.CollectRows(serviceRows, scanService)
if err != nil {
return nil, err
}
@@ -2379,39 +2293,12 @@ func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpserv
serviceIDs := make([]string, len(services))
serviceMap := make(map[string]*rpservice.Service)
for i, s := range services {
serviceIDs[i] = s.ID
serviceMap[s.ID] = s
for i, svc := range services {
serviceIDs[i] = svc.ID
serviceMap[svc.ID] = svc
}
targetRows, err := s.pool.Query(ctx, targetsQuery, serviceIDs)
if err != nil {
return nil, err
}
targets, err := pgx.CollectRows(targetRows, func(row pgx.CollectableRow) (*rpservice.Target, error) {
var t rpservice.Target
var path sql.NullString
err := row.Scan(
&t.ID,
&t.AccountID,
&t.ServiceID,
&path,
&t.Host,
&t.Port,
&t.Protocol,
&t.TargetId,
&t.TargetType,
&t.Enabled,
)
if err != nil {
return nil, err
}
if path.Valid {
t.Path = &path.String
}
return &t, nil
})
targets, err := s.getServiceTargets(ctx, serviceIDs)
if err != nil {
return nil, err
}
@@ -2425,6 +2312,201 @@ func (s *SqlStore) getServices(ctx context.Context, accountID string) ([]*rpserv
return services, nil
}
func scanService(row pgx.CollectableRow) (*rpservice.Service, error) {
var s rpservice.Service
var auth []byte
var restrictions []byte
var accessGroups []byte
var createdAt, certIssuedAt, lastRenewedAt sql.NullTime
var status, proxyCluster, sessionPrivateKey, sessionPublicKey sql.NullString
var mode, source, sourcePeer sql.NullString
var terminated, portAutoAssigned, private sql.NullBool
var listenPort sql.NullInt64
err := row.Scan(
&s.ID,
&s.AccountID,
&s.Name,
&s.Domain,
&s.Enabled,
&auth,
&restrictions,
&createdAt,
&certIssuedAt,
&lastRenewedAt,
&status,
&proxyCluster,
&s.PassHostHeader,
&s.RewriteRedirects,
&sessionPrivateKey,
&sessionPublicKey,
&mode,
&listenPort,
&portAutoAssigned,
&source,
&sourcePeer,
&terminated,
&private,
&accessGroups,
)
if err != nil {
return nil, err
}
if auth != nil {
if err := json.Unmarshal(auth, &s.Auth); err != nil {
return nil, err
}
}
if len(restrictions) > 0 {
if err := json.Unmarshal(restrictions, &s.Restrictions); err != nil {
return nil, fmt.Errorf("unmarshal restrictions: %w", err)
}
}
if len(accessGroups) > 0 {
if err := json.Unmarshal(accessGroups, &s.AccessGroups); err != nil {
return nil, fmt.Errorf("unmarshal access_groups: %w", err)
}
}
if private.Valid {
s.Private = private.Bool
}
s.Meta = serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt, status)
if proxyCluster.Valid {
s.ProxyCluster = proxyCluster.String
}
if sessionPrivateKey.Valid {
s.SessionPrivateKey = sessionPrivateKey.String
}
if sessionPublicKey.Valid {
s.SessionPublicKey = sessionPublicKey.String
}
if mode.Valid {
s.Mode = mode.String
}
if source.Valid {
s.Source = source.String
}
if sourcePeer.Valid {
s.SourcePeer = sourcePeer.String
}
if terminated.Valid {
s.Terminated = terminated.Bool
}
if portAutoAssigned.Valid {
s.PortAutoAssigned = portAutoAssigned.Bool
}
if listenPort.Valid {
if listenPort.Int64 < 0 || listenPort.Int64 > math.MaxUint16 {
return nil, fmt.Errorf("listen_port %d out of range", listenPort.Int64)
}
s.ListenPort = uint16(listenPort.Int64)
}
s.Targets = []*rpservice.Target{}
return &s, nil
}
func serviceMetaFromRow(createdAt, certIssuedAt, lastRenewedAt sql.NullTime, status sql.NullString) rpservice.Meta {
meta := rpservice.Meta{}
if createdAt.Valid {
meta.CreatedAt = createdAt.Time
}
if certIssuedAt.Valid {
t := certIssuedAt.Time
meta.CertificateIssuedAt = &t
}
if lastRenewedAt.Valid {
t := lastRenewedAt.Time
meta.LastRenewedAt = &t
}
if status.Valid {
meta.Status = status.String
}
return meta
}
func (s *SqlStore) getServiceTargets(ctx context.Context, serviceIDs []string) ([]*rpservice.Target, error) {
const targetsQuery = `SELECT ` + targetSelectColumns + ` FROM targets WHERE service_id = ANY($1)`
rows, err := s.pool.Query(ctx, targetsQuery, serviceIDs)
if err != nil {
return nil, err
}
return pgx.CollectRows(rows, scanTarget)
}
func scanTarget(row pgx.CollectableRow) (*rpservice.Target, error) {
var t rpservice.Target
var path sql.NullString
var pathRewrite sql.NullString
var proxyProtocol, skipTLSVerify, directUpstream, agentNetwork, disableAccessLog sql.NullBool
var requestTimeout, sessionIdleTimeout, captureMaxRequestBytes, captureMaxResponseBytes sql.NullInt64
var customHeaders, middlewares, captureContentTypes []byte
err := row.Scan(
&t.ID,
&t.AccountID,
&t.ServiceID,
&path,
&t.Host,
&t.Port,
&t.Protocol,
&t.TargetId,
&t.TargetType,
&t.Enabled,
&proxyProtocol,
&skipTLSVerify,
&requestTimeout,
&sessionIdleTimeout,
&pathRewrite,
&customHeaders,
&directUpstream,
&middlewares,
&captureMaxRequestBytes,
&captureMaxResponseBytes,
&captureContentTypes,
&agentNetwork,
&disableAccessLog,
)
if err != nil {
return nil, err
}
if path.Valid {
t.Path = &path.String
}
t.ProxyProtocol = proxyProtocol.Bool
t.Options.SkipTLSVerify = skipTLSVerify.Bool
t.Options.RequestTimeout = time.Duration(requestTimeout.Int64)
t.Options.SessionIdleTimeout = time.Duration(sessionIdleTimeout.Int64)
t.Options.PathRewrite = rpservice.PathRewriteMode(pathRewrite.String)
t.Options.DirectUpstream = directUpstream.Bool
t.Options.CaptureMaxRequestBytes = captureMaxRequestBytes.Int64
t.Options.CaptureMaxResponseBytes = captureMaxResponseBytes.Int64
t.Options.AgentNetwork = agentNetwork.Bool
t.Options.DisableAccessLog = disableAccessLog.Bool
if len(customHeaders) > 0 {
if err := json.Unmarshal(customHeaders, &t.Options.CustomHeaders); err != nil {
return nil, fmt.Errorf("unmarshal custom_headers: %w", err)
}
}
if len(middlewares) > 0 {
if err := json.Unmarshal(middlewares, &t.Options.Middlewares); err != nil {
return nil, fmt.Errorf("unmarshal middlewares: %w", err)
}
}
if len(captureContentTypes) > 0 {
if err := json.Unmarshal(captureContentTypes, &t.Options.CaptureContentTypes); err != nil {
return nil, fmt.Errorf("unmarshal capture_content_types: %w", err)
}
}
return &t, nil
}
func (s *SqlStore) getNetworks(ctx context.Context, accountID string) ([]*networkTypes.Network, error) {
const query = `SELECT id, account_id, public_id, name, description FROM networks WHERE account_id = $1`
rows, err := s.pool.Query(ctx, query, accountID)

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

@@ -0,0 +1,74 @@
package store
import (
"strings"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm/schema"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
)
// TestPgxServiceColumnsMatchGorm guards the Postgres pgx read path against
// drifting from the gorm model. The SQLite/MySQL gorm path loads rows by struct,
// so a new column on a model is picked up automatically, but the hand-written
// pgx SELECT in sql_store.go must be updated by hand. This test fails when a
// gorm column is missing from the pgx column list, which otherwise silently
// returns zero-valued on Postgres with no compile error.
func TestPgxServiceColumnsMatchGorm(t *testing.T) {
tests := []struct {
name string
model any
selectColumns string
// excluded lists gorm columns intentionally not loaded by the pgx path.
excluded map[string]struct{}
}{
{
name: "service",
model: &rpservice.Service{},
selectColumns: serviceSelectColumns,
},
{
name: "target",
model: &rpservice.Target{},
selectColumns: targetSelectColumns,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
selected := parseColumnList(tc.selectColumns)
for _, col := range gormColumnNames(t, tc.model) {
if _, ok := tc.excluded[col]; ok {
continue
}
_, ok := selected[col]
assert.Truef(t, ok,
"gorm column %q is not read by the Postgres pgx SELECT; add it to %sSelectColumns in sql_store.go (or to the test's excluded set if it is intentionally not loaded)",
col, tc.name)
}
})
}
}
func parseColumnList(cols string) map[string]struct{} {
set := make(map[string]struct{})
for _, c := range strings.Split(cols, ",") {
if c = strings.TrimSpace(c); c != "" {
set[c] = struct{}{}
}
}
return set
}
// gormColumnNames returns the DB column names gorm would migrate for the model,
// using the same default naming strategy the store configures.
func gormColumnNames(t *testing.T, model any) []string {
t.Helper()
sch, err := schema.Parse(model, &sync.Map{}, schema.NamingStrategy{})
require.NoError(t, err)
return sch.DBNames
}

View File

@@ -5,6 +5,7 @@ import (
"os"
"runtime"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -44,3 +45,91 @@ func TestSqlStore_GetAccount_PrivateServiceRoundtrip(t *testing.T) {
assert.Equal(t, []string{"grp-admins", "grp-ops"}, got.AccessGroups)
})
}
// TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip guards the Postgres pgx
// read path (getServices) against silently dropping columns present on the gorm
// model. Before the fix these fields loaded correctly on SQLite but came back
// zero-valued on Postgres because the hand-written SELECT and scan omitted them.
func TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip(t *testing.T) {
if os.Getenv("CI") == "true" && (runtime.GOOS == "darwin" || runtime.GOOS == "windows") {
t.Skip("skip CI tests on darwin and windows")
}
runTestForAllEngines(t, "", func(t *testing.T, store Store) {
ctx := context.Background()
account := newAccountWithId(ctx, "account_svc_opts", "testuser", "")
require.NoError(t, store.SaveAccount(ctx, account))
renewedAt := time.Now().UTC().Truncate(time.Second)
targetPath := "/api"
svc := &rpservice.Service{
ID: "svc-opts",
AccountID: account.Id,
Name: "opts-svc",
Domain: "opts.example",
Enabled: true,
Mode: rpservice.ModeHTTP,
Restrictions: rpservice.AccessRestrictions{
AllowedCIDRs: []string{"10.0.0.0/8"},
BlockedCountries: []string{"XX"},
CrowdSecMode: "block",
},
Meta: rpservice.Meta{
LastRenewedAt: &renewedAt,
},
Targets: []*rpservice.Target{
{
AccountID: account.Id,
ServiceID: "svc-opts",
Path: &targetPath,
Host: "backend.internal",
Port: 8080,
Protocol: "http",
TargetId: "tgt-1",
Enabled: true,
ProxyProtocol: true,
Options: rpservice.TargetOptions{
SkipTLSVerify: true,
RequestTimeout: 30 * time.Second,
SessionIdleTimeout: 5 * time.Minute,
PathRewrite: rpservice.PathRewritePreserve,
CustomHeaders: map[string]string{"X-Foo": "bar"},
DirectUpstream: true,
CaptureMaxRequestBytes: 1024,
CaptureMaxResponseBytes: 2048,
CaptureContentTypes: []string{"application/json"},
AgentNetwork: true,
DisableAccessLog: true,
},
},
},
}
require.NoError(t, store.CreateService(ctx, svc))
loaded, err := store.GetAccount(ctx, account.Id)
require.NoError(t, err)
require.Len(t, loaded.Services, 1)
got := loaded.Services[0]
assert.Equal(t, []string{"10.0.0.0/8"}, got.Restrictions.AllowedCIDRs, "restrictions allowed CIDRs")
assert.Equal(t, []string{"XX"}, got.Restrictions.BlockedCountries, "restrictions blocked countries")
assert.Equal(t, "block", got.Restrictions.CrowdSecMode, "restrictions crowdsec mode")
require.NotNil(t, got.Meta.LastRenewedAt, "meta last renewed at")
assert.WithinDuration(t, renewedAt, *got.Meta.LastRenewedAt, time.Second, "meta last renewed at")
require.Len(t, got.Targets, 1)
tg := got.Targets[0]
assert.True(t, tg.ProxyProtocol, "target proxy protocol")
assert.True(t, tg.Options.SkipTLSVerify, "options skip TLS verify")
assert.Equal(t, 30*time.Second, tg.Options.RequestTimeout, "options request timeout")
assert.Equal(t, 5*time.Minute, tg.Options.SessionIdleTimeout, "options session idle timeout")
assert.Equal(t, rpservice.PathRewritePreserve, tg.Options.PathRewrite, "options path rewrite")
assert.Equal(t, map[string]string{"X-Foo": "bar"}, tg.Options.CustomHeaders, "options custom headers")
assert.True(t, tg.Options.DirectUpstream, "options direct upstream")
assert.Equal(t, int64(1024), tg.Options.CaptureMaxRequestBytes, "options capture max request bytes")
assert.Equal(t, int64(2048), tg.Options.CaptureMaxResponseBytes, "options capture max response bytes")
assert.Equal(t, []string{"application/json"}, tg.Options.CaptureContentTypes, "options capture content types")
assert.True(t, tg.Options.AgentNetwork, "options agent network")
assert.True(t, tg.Options.DisableAccessLog, "options disable access log")
})
}

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

@@ -1517,6 +1517,54 @@ func (a *Account) GetResourceRoutersMap() map[string]map[string]*routerTypes.Net
return routers
}
// forcesRoutingPeerDNSResolution reports whether the given peer must run
// routing-peer DNS resolution regardless of the account-global
// RoutingPeerDNSResolutionEnabled setting. It returns true when the peer is a
// router for a domain network resource that is targeted by an enabled
// reverse-proxy service, so the peer's DNS forwarder starts and can resolve
// the target for the embedded proxy peers. Embedded proxy peers themselves are
// handled at PeerConfig build time.
func (a *Account) forcesRoutingPeerDNSResolution(peerID string, routers map[string]map[string]*routerTypes.NetworkRouter) bool {
targeted := a.proxyTargetedDomainResourceIDs()
if len(targeted) == 0 {
return false
}
for _, resource := range a.NetworkResources {
if resource == nil || !resource.Enabled || resource.Type != resourceTypes.Domain {
continue
}
if _, ok := targeted[resource.ID]; !ok {
continue
}
if _, isRouter := routers[resource.NetworkID][peerID]; isRouter {
return true
}
}
return false
}
// proxyTargetedDomainResourceIDs returns the set of domain network resource IDs
// targeted by an enabled, non-terminated reverse-proxy service.
func (a *Account) proxyTargetedDomainResourceIDs() map[string]struct{} {
ids := make(map[string]struct{})
for _, svc := range a.Services {
if svc == nil || !svc.Enabled || svc.Terminated {
continue
}
for _, target := range svc.Targets {
if target == nil || !target.Enabled {
continue
}
if target.TargetType == service.TargetTypeDomain {
ids[target.TargetId] = struct{}{}
}
}
}
return ids
}
// getPoliciesSourcePeers collects all unique peers from the source groups defined in the given policies.
func getPoliciesSourcePeers(policies []*Policy, groups map[string]*Group) map[string]struct{} {
sourcePeers := make(map[string]struct{})

View File

@@ -140,6 +140,8 @@ func (a *Account) GetPeerNetworkMapComponents(
RouterPeers: make(map[string]*ComponentPeer),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
ForceRoutingPeerDNSResolution: a.forcesRoutingPeerDNSResolution(peerID, routers),
}
for _, n := range a.Networks {
if n != nil {

View File

@@ -1751,3 +1751,71 @@ func hasPrivateAccessPolicy(account *Account, serviceID string) bool {
}
return false
}
func TestForcesRoutingPeerDNSResolution(t *testing.T) {
buildAccountRes := func(serviceEnabled, targetEnabled, resourceEnabled bool, targetType service.TargetType, resType resourceTypes.NetworkResourceType) *Account {
return &Account{
Id: "accountID",
Groups: map[string]*Group{
"router-group": {ID: "router-group", Peers: []string{"router-peer-grp"}},
},
NetworkRouters: []*routerTypes.NetworkRouter{
{ID: "r1", NetworkID: "net-1", AccountID: "accountID", Peer: "router-peer", Enabled: true},
{ID: "r2", NetworkID: "net-1", AccountID: "accountID", PeerGroups: []string{"router-group"}, Enabled: true},
},
NetworkResources: []*resourceTypes.NetworkResource{
{ID: "res-domain", AccountID: "accountID", NetworkID: "net-1", Type: resType, Domain: "example.org", Enabled: resourceEnabled},
},
Services: []*service.Service{
{
ID: "svc-1", AccountID: "accountID", Enabled: serviceEnabled,
Targets: []*service.Target{
{TargetId: "res-domain", TargetType: targetType, Enabled: targetEnabled},
},
},
},
}
}
buildAccount := func(serviceEnabled, targetEnabled, resourceEnabled bool, targetType service.TargetType) *Account {
return buildAccountRes(serviceEnabled, targetEnabled, resourceEnabled, targetType, resourceTypes.Domain)
}
t.Run("router peer for RP-targeted domain resource is forced", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypeDomain)
routers := account.GetResourceRoutersMap()
assert.True(t, account.forcesRoutingPeerDNSResolution("router-peer", routers), "direct router peer should be forced")
assert.True(t, account.forcesRoutingPeerDNSResolution("router-peer-grp", routers), "group-member router peer should be forced")
})
t.Run("non-router peer is not forced", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("other-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when service disabled", func(t *testing.T) {
account := buildAccount(false, true, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when target disabled", func(t *testing.T) {
account := buildAccount(true, false, true, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when resource disabled", func(t *testing.T) {
account := buildAccount(true, true, false, service.TargetTypeDomain)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced for non-domain target type", func(t *testing.T) {
account := buildAccount(true, true, true, service.TargetTypePeer)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()))
})
t.Run("not forced when targeted resource is not a domain", func(t *testing.T) {
account := buildAccountRes(true, true, true, service.TargetTypeDomain, resourceTypes.Host)
assert.False(t, account.forcesRoutingPeerDNSResolution("router-peer", account.GetResourceRoutersMap()),
"a domain target pointing at a non-domain resource must not force resolution")
})
}

View File

@@ -321,6 +321,10 @@ func (am *DefaultAccountManager) DeleteUser(ctx context.Context, accountID, init
return err
}
if targetUser.AccountID != accountID {
return status.NewUserNotFoundError(targetUserID)
}
if targetUser.Role == types.UserRoleOwner {
return status.NewOwnerDeletePermissionError()
}

View File

@@ -802,6 +802,52 @@ func TestUser_DeleteUser_SelfDelete(t *testing.T) {
}
}
func TestUser_DeleteUser_OtherAccount(t *testing.T) {
testStore, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
if err != nil {
t.Fatalf("Error when creating store: %s", err)
}
t.Cleanup(cleanup)
account := newAccountWithId(context.Background(), mockAccountID, mockUserID, "", "", "", false)
if err = testStore.SaveAccount(context.Background(), account); err != nil {
t.Fatalf("Error when saving account: %s", err)
}
otherAccount := newAccountWithId(context.Background(), "otherAccount", "otherOwner", "", "", "", false)
otherAccount.Users["otherRegularUser"] = &types.User{
Id: "otherRegularUser",
AccountID: "otherAccount",
Role: types.UserRoleUser,
}
otherAccount.Users["otherServiceUser"] = &types.User{
Id: "otherServiceUser",
AccountID: "otherAccount",
Role: types.UserRoleUser,
IsServiceUser: true,
ServiceUserName: "otherServiceUser",
}
if err = testStore.SaveAccount(context.Background(), otherAccount); err != nil {
t.Fatalf("Error when saving other account: %s", err)
}
am := DefaultAccountManager{
Store: testStore,
eventStore: &activity.InMemoryEventStore{},
permissionsManager: permissions.NewManager(testStore),
}
for _, targetUserID := range []string{"otherRegularUser", "otherServiceUser"} {
t.Run(targetUserID, func(t *testing.T) {
err := am.DeleteUser(context.Background(), mockAccountID, mockUserID, targetUserID)
assert.Equal(t, status.NewUserNotFoundError(targetUserID), err)
_, err = testStore.GetUserByUserID(context.Background(), store.LockingStrengthNone, targetUserID)
assert.NoError(t, err, "user of another account must not be deleted")
})
}
}
func TestUser_DeleteUser_regularUser(t *testing.T) {
store, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
if err != nil {

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

@@ -3,20 +3,30 @@ package auth
import (
"context"
"net/netip"
"os"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/sync/singleflight"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
// reused before re-fetching from management. 5 minutes balances freshness
// against management load on busy mesh networks.
// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer
// result is reused before re-fetching from management. 5 minutes balances
// freshness against management load on busy mesh networks. Override it with
// envTunnelCacheTTL when an account needs authorization changes to take effect
// sooner (at the cost of more ValidateTunnelPeer RPCs).
const tunnelCacheTTL = 300 * time.Second
// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string
// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the
// default.
const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL"
// tunnelCachePerAccount caps the number of cached identities per account.
// Bounded eviction avoids memory growth in pathological cases (huge peer
// roster, brief request bursts) while staying generous for normal use.
@@ -60,16 +70,35 @@ type accountBucket struct {
order []tunnelCacheKey
}
// newTunnelValidationCache constructs a cache with default TTL and bounds.
// newTunnelValidationCache constructs a cache with the configured TTL
// (envTunnelCacheTTL override or default) and default bounds.
func newTunnelValidationCache() *tunnelValidationCache {
return &tunnelValidationCache{
entries: make(map[types.AccountID]*accountBucket),
ttl: tunnelCacheTTL,
ttl: tunnelCacheTTLFromEnv(),
maxSize: tunnelCachePerAccount,
now: time.Now,
}
}
// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the
// envTunnelCacheTTL override. The override must be a positive Go duration
// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive
// falls back to tunnelCacheTTL.
func tunnelCacheTTLFromEnv() time.Duration {
raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL))
if raw == "" {
return tunnelCacheTTL
}
d, err := time.ParseDuration(raw)
if err != nil || d <= 0 {
log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
envTunnelCacheTTL, raw, tunnelCacheTTL)
return tunnelCacheTTL
}
return d
}
// get returns a cached response for the key, or nil when missing or
// expired. Expired entries are evicted lazily on read.
func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {

View File

@@ -169,3 +169,32 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
}
func TestTunnelCacheTTLFromEnv(t *testing.T) {
t.Run("unset uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("valid duration overrides", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "45s")
assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv())
})
t.Run("whitespace trimmed", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, " 2m ")
assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv())
})
t.Run("unparseable uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "nonsense")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("non-positive uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "0s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
t.Setenv(envTunnelCacheTTL, "-30s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("constructor honors override", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "90s")
assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl)
})
}

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

@@ -10,11 +10,15 @@ import (
)
// Config is the JSON-decoded shape accepted by the factory. The
// runtime path consumes the normalised allowlist; raw config is not
// runtime path consumes the normalised allowlists; raw config is not
// retained beyond construction.
type Config struct {
ModelAllowlist []string `json:"model_allowlist"`
PromptCapture PromptCapture `json:"prompt_capture"`
// ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to
// its model allowlist. A provider present is restricted to those models; one
// absent is unrestricted. Kept per-provider so one provider's list can't leak
// onto another.
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture PromptCapture `json:"prompt_capture"`
}
// PromptCapture toggles the optional prompt capture + redaction step
@@ -54,21 +58,28 @@ func isEmptyJSON(raw []byte) bool {
return false
}
// normaliseConfig lowercases and trims allowlist entries so the runtime
// match is case-insensitive. Empty entries are dropped.
// normaliseConfig lowercases and trims allowlist entries for case-insensitive
// matching; empty entries drop. A provider whose entries all drop keeps an empty
// (non-nil) list — "deny every model" — distinct from an absent provider
// (unrestricted).
func normaliseConfig(cfg Config) Config {
if len(cfg.ModelAllowlist) == 0 {
if len(cfg.ProviderAllowlists) == 0 {
cfg.ProviderAllowlists = nil
return cfg
}
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
for _, entry := range cfg.ModelAllowlist {
n := normaliseModel(entry)
if n == "" {
continue
cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
for provider, models := range cfg.ProviderAllowlists {
list := make([]string, 0, len(models))
for _, entry := range models {
n := normaliseModel(entry)
if n == "" {
continue
}
list = append(list, n)
}
cleaned = append(cleaned, n)
cleaned[provider] = list
}
cfg.ModelAllowlist = cleaned
cfg.ProviderAllowlists = cleaned
return cfg
}

View File

@@ -83,8 +83,9 @@ func (m *Middleware) MutationsSupported() bool { return false }
// prompt capture only affects the metadata emitted alongside an allow.
func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil {
return denial, nil
}
@@ -110,20 +111,32 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
// is a no-op.
func (m *Middleware) Close() error { return nil }
// evaluateAllowlist returns a deny Output when the configured allowlist
// rejects the model. A nil return means the request should proceed.
func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output {
if len(m.cfg.ModelAllowlist) == 0 {
// evaluateAllowlist denies when the resolved provider's allowlist rejects the
// model; nil means proceed. Scoped to the provider llm_router resolved, so an
// unrestricted provider (absent from config) is never caught by another's list.
func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output {
if len(m.cfg.ProviderAllowlists) == 0 {
return nil
}
// Fail closed: with an allowlist configured, a request whose model the
// upstream parser could not extract (absent or empty) must be denied rather
// than allowed. This is what enforces the allowlist for URL/path-routed
// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
// Restrictions exist but the resolved provider is unknown, so we can't tell
// if this request targets a restricted provider — fail closed. llm_router
// normally stamps the provider first, so this is a defensive guard.
if providerID == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
allowlist, restricted := m.cfg.ProviderAllowlists[providerID]
if !restricted {
// This provider has no allowlist (some authorising policy left it
// unrestricted); management owns any per-policy/group decision.
return nil
}
// Fail closed: with an allowlist in effect for this provider, a request whose
// model the parser couldn't extract (absent/empty) is denied. This enforces
// the allowlist for path-routed providers (Bedrock, Vertex) with no body model.
if !modelPresent || normaliseModel(model) == "" {
return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
}
if m.modelInAllowlist(model) {
if modelInAllowlist(allowlist, model) {
return nil
}
return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
@@ -151,14 +164,15 @@ func denyModel(model, code, message, reason string) *middleware.Output {
}
}
// modelInAllowlist reports whether the model matches any allowlist
// entry under the case-insensitive, trim-tolerant comparison rule.
func (m *Middleware) modelInAllowlist(model string) bool {
// modelInAllowlist reports whether the model matches any entry in the supplied
// (already-normalised) allowlist under the case-insensitive, trim-tolerant
// comparison rule.
func modelInAllowlist(allowlist []string, model string) bool {
normalised := normaliseModel(model)
if normalised == "" {
return false
}
for _, allowed := range m.cfg.ModelAllowlist {
for _, allowed := range allowlist {
if allowed == normalised {
return true
}

View File

@@ -26,6 +26,25 @@ func newInput(meta ...middleware.KV) *middleware.Input {
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
}
const (
testProvider = "prov-1"
otherProvider = "prov-2"
)
// providerCfg builds a Config restricting testProvider to the given models.
func providerCfg(models ...string) Config {
return Config{ProviderAllowlists: map[string][]string{testProvider: models}}
}
// newInputProvider builds an input that carries a resolved provider id (as
// llm_router would stamp) plus any extra metadata.
func newInputProvider(provider string, meta ...middleware.KV) *middleware.Input {
all := make([]middleware.KV, 0, len(meta)+1)
all = append(all, middleware.KV{Key: middleware.KeyLLMResolvedProviderID, Value: provider})
all = append(all, meta...)
return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: all}
}
func TestMiddlewareIdentity(t *testing.T) {
mw := New(Config{})
assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
@@ -47,12 +66,12 @@ func TestMiddlewareIdentity(t *testing.T) {
func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model")
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no provider allowlists must allow any model")
v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
require.True(t, ok, "decision metadata must be emitted")
assert.Equal(t, "allow", v, "decision must be allow")
@@ -62,8 +81,8 @@ func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
}
func TestAllowlistMatchAllows(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o", "claude-opus-4"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
@@ -71,8 +90,8 @@ func TestAllowlistMatchAllows(t *testing.T) {
}
func TestAllowlistMissDenies(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
@@ -91,10 +110,10 @@ func TestAllowlistMissDenies(t *testing.T) {
}
func TestAllowlistCaseInsensitive(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{" GPT-4o ", "Claude-OPUS-4"}})
mw := New(providerCfg(" GPT-4o ", "Claude-OPUS-4"))
cases := []string{"gpt-4o", "GPT-4O", " claude-opus-4 "}
for _, model := range cases {
out, err := mw.Invoke(context.Background(), newInput(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: model},
))
require.NoError(t, err)
@@ -103,14 +122,15 @@ func TestAllowlistCaseInsensitive(t *testing.T) {
}
func TestAllowlistMissingModelKeyDenies(t *testing.T) {
// Fail closed: with an allowlist configured, a request whose model the
// parser could not extract (URL/path-routed providers such as Bedrock or
// Vertex whose shape wasn't recognised) must be denied, not allowed.
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput())
// Fail closed: with an allowlist in effect for the resolved provider, a
// request whose model the parser could not extract (URL/path-routed
// providers such as Bedrock or Vertex whose shape wasn't recognised) must be
// denied, not allowed.
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when the provider is restricted")
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
require.NotNil(t, out.DenyReason, "deny reason must be populated")
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
@@ -122,26 +142,101 @@ func TestAllowlistMissingModelKeyDenies(t *testing.T) {
func TestAllowlistEmptyModelValueDenies(t *testing.T) {
// A present-but-empty model is as undeterminable as an absent one.
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput(
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: " "},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when the provider is restricted")
require.NotNil(t, out.DenyReason, "deny reason must be populated")
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
}
func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
// Without an allowlist there is nothing to enforce, so a missing model is
// still allowed — the fail-closed rule only applies when a list is set.
// Without any provider allowlists there is nothing to enforce, so a missing
// model is still allowed — the fail-closed rule only applies when a
// restriction is in effect.
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput())
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
}
func TestUnrestrictedProviderAllowsAnyModel(t *testing.T) {
// The request resolved to otherProvider, which has no allowlist, so its
// traffic must not be caught by testProvider's restriction — the
// cross-provider-leak / false-deny guard.
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInputProvider(otherProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "an unrestricted provider must not inherit another provider's allowlist")
}
func TestPerProviderAllowlistsAreIsolated(t *testing.T) {
// gpt-4o is allowed only on testProvider; claude-opus-4 only on
// otherProvider. A model allowlisted for one provider must not be usable on
// the other — the fail-closed layer never unions allowlists across providers.
mw := New(Config{ProviderAllowlists: map[string][]string{
testProvider: {"gpt-4o"},
otherProvider: {"claude-opus-4"},
}})
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "claude-opus-4 is allowed only on otherProvider, not testProvider")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "cross-provider model must be blocked, not model_unknown")
}
func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
// Restrictions exist for the account but the resolved provider id is absent,
// so the request cannot be scoped to a provider. Fail closed rather than
// wave it through.
mw := New(providerCfg("gpt-4o"))
out, err := mw.Invoke(context.Background(), newInput(
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "missing resolved provider must fail closed when restrictions exist")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
}
func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) {
// An allowlist-enabled provider with zero models is distinct from an
// unrestricted (absent) provider: it must deny every model.
mw := New(providerCfg())
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown")
}
func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) {
// All the provider's entries are blank; they collapse to a non-nil empty
// list (deny everything for that provider), not "no restriction".
raw := []byte(`{"provider_allowlists":{"prov-1":[""," "]}}`)
mw, err := Factory{}.New(raw)
require.NoError(t, err)
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked")
}
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput(
@@ -217,8 +312,8 @@ func TestFactoryAcceptsZeroConfigs(t *testing.T) {
func TestFactoryDecodesValidConfig(t *testing.T) {
cfg := Config{
ModelAllowlist: []string{"gpt-4o"},
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
ProviderAllowlists: map[string][]string{testProvider: {"gpt-4o"}},
PromptCapture: PromptCapture{Enabled: true, RedactPii: true},
}
raw, err := json.Marshal(cfg)
require.NoError(t, err, "marshalling test config must succeed")
@@ -234,15 +329,15 @@ func TestFactoryRejectsMalformedJSON(t *testing.T) {
}
func TestFactoryNormalisesAllowlist(t *testing.T) {
raw := []byte(`{"model_allowlist":[" GPT-4o ","",""," Claude-3 "]}`)
raw := []byte(`{"provider_allowlists":{"prov-1":[" GPT-4o ","",""," Claude-3 "]}}`)
mw, err := Factory{}.New(raw)
require.NoError(t, err)
out, err := mw.Invoke(context.Background(), newInput(
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries")
out2, err := mw.Invoke(context.Background(), newInput(
out2, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
))
require.NoError(t, err)

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