Compare commits

..

11 Commits

Author SHA1 Message Date
sagitchu 3ce320da5a fix(nftables): harden traffic accounting edges 2026-06-07 11:55:59 +08:00
sagitchu 7ab0db29ae fix(nftables): clean counter state on forward delete 2026-06-07 11:55:59 +08:00
sagitchu 006ea97200 feat(nftables): ingest traffic counters 2026-06-07 11:55:59 +08:00
sagitchu e569aedd3e feat(nftables): calculate traffic deltas 2026-06-07 11:55:59 +08:00
sagitchu 35080aea2d feat(flow): expose forward owner metadata 2026-06-07 11:55:59 +08:00
sagitchu 85e588ffe9 feat(nftables): persist counter state 2026-06-07 11:55:59 +08:00
sagitchu 03524f4a65 feat(nftables): collect counters over ssh 2026-06-07 11:55:59 +08:00
sagitchu 9e69e020ab feat(nftables): parse traffic counters 2026-06-07 11:55:59 +08:00
sagitchu 14bbd3907d feat(nftables): render traffic counters 2026-06-07 11:55:59 +08:00
sagitchu ca8d8e92ba docs: plan nftables traffic stats 2026-06-07 11:55:59 +08:00
sagitchu 079474fa06 docs: design nftables traffic stats 2026-06-07 11:55:59 +08:00
24 changed files with 4068 additions and 105 deletions
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,284 @@
# nftables 流量统计设计
**日期**: 2026-06-06
**状态**: 待审核
**作者**: Codex
## 概述
为 FLVX 的 `nftables` 转发模式补齐流量统计。当前 nftables 模式由面板通过 SSH 全量维护 `table inet flvx`,但没有 agent,因此不能复用 WebSocket 运行时上报。新方案由面板定时通过 SSH 拉取远端 nftables counter,计算增量后写入现有流量账本。
目标是让 nftables 转发在用户可见口径上尽量接近 agent 模式:
- forward 列表显示 `inFlow` / `outFlow`。
- 用户、用户隧道、配额和流量策略继续生效。
- 隧道监控继续获得分钟级 `tunnel_metric`。
- 节点不需要安装新的 agent 或常驻进程。
## 背景
现有 agent 模式通过 `/flow/upload` 接收加密上报,handler 会把服务名解析为 `forward_id/user_id/user_tunnel_id`,再复用以下路径:
- `ApplyFlowUploadDeltasBatch` 更新 `forward`、`user`、`user_tunnel`。
- `AddUserQuotaUsageBatch` 更新用户配额窗口。
- `enforceUserQuotaIfNeeded` 和 `enforceFlowPolicies` 做约束 enforcement。
- `recordTunnelMetricsFromForwardBatch` 写入分钟级隧道监控。
nftables 模式已经有 `nft_rule_binding` 记录规则应用状态,规则 comment 里包含 `forward_id`。这给 counter 到业务实体的映射提供了稳定锚点。
## 推荐方案
采用“面板 SSH 轮询 nftables counter”的方案:
1. 渲染 nftables 规则时,为每个 forward、协议和方向写入稳定 comment 和 `counter`。
2. 后端定时扫描 `forward_mode = nftables` 的节点。
3. 对每个节点通过 SSH 执行 `nft -j list table inet flvx`。
4. 解析 JSON 规则,按 comment 得到 `forward_id/protocol/direction/bytes/packets`。
5. 用数据库中的上次采样值计算 delta。
6. 将 delta 转成现有 flow upload 内部结构,复用既有入账、配额、策略和监控逻辑。
不采用节点 crontab 或 systemd timer 回推。它会重新引入节点侧组件,削弱 nftables 模式“不安装 agent”的产品边界。
## 统计口径
正式入账使用 `forward` filter chain 的计数,不使用 NAT chain 的 DNAT 命中计数作为主口径。
原因:
- DNAT counter 表示规则命中,不一定代表后续转发成功。
- filter forward chain 更接近实际经过内核转发的数据。
- SNAT/masquerade 会改变包头,入账规则应在可稳定匹配目标服务地址和端口的位置统计。
方向定义:
| direction | nft 匹配 | 写入字段 |
|-----------|----------|----------|
| `to-target` | 外部客户端到目标服务 | `in_flow` |
| `from-target` | 目标服务返回外部客户端 | `out_flow` |
用户总用量和配额仍按 `in_flow + out_flow` 计算。隧道 `traffic_ratio` 和 `flow` 倍率继续沿用 agent 模式逻辑,保证不同运行时模式的账单口径一致。
## nftables 规则设计
继续只维护 `table inet flvx`,避免触碰用户已有规则。每条 forward 对 TCP 和 UDP 各生成一组 DNAT 和统计规则。
示例:
```nft
table inet flvx {
chain prerouting {
type nat hook prerouting priority dstnat; policy accept;
tcp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat tcp"
udp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat udp"
}
chain postrouting {
type nat hook postrouting priority srcnat; policy accept;
masquerade comment "flvx masquerade"
}
chain forward {
type filter hook forward priority filter; policy accept;
ip daddr 198.51.100.20 tcp dport 443 counter comment "flvx forward:42 to-target tcp"
ip saddr 198.51.100.20 tcp sport 443 counter comment "flvx forward:42 from-target tcp"
ip daddr 198.51.100.20 udp dport 443 counter comment "flvx forward:42 to-target udp"
ip saddr 198.51.100.20 udp sport 443 counter comment "flvx forward:42 from-target udp"
}
}
```
IPv6 目标使用 `ip6`:
```nft
ip6 daddr 2001:db8::20 tcp dport 443 counter comment "flvx forward:42 to-target tcp"
ip6 saddr 2001:db8::20 tcp sport 443 counter comment "flvx forward:42 from-target tcp"
```
域名目标无法在 nftables 规则中动态匹配返回方向。统计第一阶段要求 nftables forward 的 `remoteAddr` host 必须是 IP 地址;如果当前纯转发实现允许域名,开启统计时应同步收紧校验。后续若要支持域名,应在规则同步时解析并固化 IP,同时明确 DNS 变化后的重建策略。
## Comment 格式
正式统计规则使用固定格式:
```text
flvx forward:<forward_id> <direction> <protocol>
```
字段:
- `forward_id`: 十进制整数。
- `direction`: `to-target` 或 `from-target`。
- `protocol`: `tcp` 或 `udp`。
DNAT 调试规则可使用 `dnat` direction,但 collector 不入账 `dnat`。后端只依赖 comment 解析,不依赖 nft handle,因为全量重建 table 会改变 handle。
## 数据模型
新增 `nft_counter_state` 表保存上次采样基线。
| 字段 | 说明 |
|------|------|
| `id` | 主键 |
| `node_id` | nftables 节点 ID |
| `forward_id` | 转发规则 ID |
| `protocol` | `tcp` / `udp` |
| `direction` | `to-target` / `from-target` |
| `rule_hash` | 当前规则 hash |
| `bytes` | 上次采样绝对字节数 |
| `packets` | 上次采样绝对包数 |
| `collected_time` | 上次采样时间 |
| `created_time` | 创建时间 |
| `updated_time` | 更新时间 |
唯一索引:
```text
node_id, forward_id, protocol, direction
```
GORM 模型必须定义 `TableName()`,字段 tag 保持 SQLite/PostgreSQL 兼容,不使用 `jsonb`、`serial` 等数据库专属类型。
## 后端组件
扩展 `go-backend/internal/runtime/nftables`:
| 组件 | 职责 |
|------|------|
| `CounterSample` | 表达单条 nft counter 采样 |
| `Collector` | 对外提供 `Collect(ctx, cfg)` |
| `SSHRunner.ListTableJSON` | 远端执行 `nft -j list table inet flvx` |
| `ParseCounterSamples` | 解析 nft JSON 和 FLVX comment |
扩展 repository:
| 方法 | 职责 |
|------|------|
| `ListNftablesNodesForCollection` | 找到启用 nftables 且有 SSH 配置的节点 |
| `GetNftCounterStatesByNode` | 读取节点上次 counter 基线 |
| `UpsertNftCounterStates` | 批量刷新基线 |
| `DeleteNftCounterStatesByForward` | forward 删除时清理状态 |
扩展 handler/job:
- 新增 `runNftablesTrafficCollectJob(now time.Time)`。
- 默认每 60 秒运行一次。
- 对节点采集设置并发上限,建议 3 到 5。
- 单节点失败只记录日志和节点采集状态,不影响其他节点。
## 增量算法
collector 返回的是 nftables 的绝对 counter。入账前必须和上次基线做差。
规则:
- 无旧状态:只保存当前值作为基线,不入账。
- `rule_hash` 变化:只刷新基线,不入账,避免新旧规则混算。
- 新 bytes 大于等于旧 bytes:`delta = new - old`。
- 新 bytes 小于旧 bytes:认为远端 table 重建、counter reset 或系统重启,只刷新基线,不入账。
- delta 为 0:刷新采集时间,不入账。
- 样本无法映射到有效 forward:忽略并记录 debug 日志。
同一 forward 的 TCP/UDP delta 要先聚合,再转换成现有账本:
- `to-target` bytes 聚合为原始 `bytesIn`。
- `from-target` bytes 聚合为原始 `bytesOut`。
- 入账时按 `traffic_ratio` 和 `tunnel.flow` 计算 scaled `InFlow` / `OutFlow`。
- 配额使用 scaled 后的 `InFlow + OutFlow`。
- `tunnel_metric` 使用原始 `bytesIn` / `bytesOut`。
## 入账路径
新增一个 nftables 专用的 batch builder,但输出沿用现有结构:
```go
type nftTrafficDelta struct {
ForwardID int64
BytesIn int64
BytesOut int64
}
```
处理流程:
1. 收集本轮所有 `forward_id`。
2. 调用 `GetFlowUploadForwardMetas` 获取 `user_id/user_tunnel_id/tunnel_id/traffic_ratio/tunnel_flow`。
3. 构造 `repo.FlowUploadCounterDelta`。
4. 调用 `recordTunnelMetricsFromForwardBatch` 写监控。
5. 抽出共享入账 helper,复用 `applyFlowDeltasWithFallback`、`applyQuotaUsageWithFallback`、`enforceUserQuotaIfNeeded` 和 `enforceFlowPolicies`。不要通过伪造 agent service name 去调用 agent 专用 builder。
不新增独立的 nftables 流量字段。`forward.in_flow/out_flow`、`user.in_flow/out_flow`、`user_tunnel.in_flow/out_flow` 仍是统一事实来源。
## 错误处理
采集错误分为三类:
| 类型 | 行为 |
|------|------|
| SSH 连接或认证失败 | 记录日志,保留下次继续采集 |
| 远端无 `table inet flvx` | 视为规则未应用或被清理,记录 warning,不清空账本 |
| JSON 解析失败 | 记录原始错误摘要,不入账 |
不要因为采集失败禁用 forward。流量统计失败和转发运行失败不是同一件事。
可在后续 UI 增加节点级采集状态,例如最近成功时间、最近错误。但第一步只要求后端具备日志和数据库状态即可。
## 与现有行为的关系
- agent 模式 `/flow/upload` 不变。
- nftables 模式不新增节点侧 HTTP 回调。
- 现有 `nft_rule_binding.rule_hash` 继续表示规则期望状态;counter state 用它判断采样是否跨规则版本。
- `statistics_flow` 小时统计 job 不需要改,它基于用户总流量快照自然包含 nftables 入账结果。
- 用户重置流量时不需要清空 nftables counter。重置只清业务账本;下一轮采集继续从 counter state 差值入账。
## 测试计划
后端单元测试:
- renderer 为 TCP/UDP、IPv4/IPv6 目标生成 `counter` 和稳定 comment。
- comment parser 能识别合法格式,拒绝未知 direction/protocol。
- nft JSON parser 能从 `nft -j list table` 输出中提取 bytes/packets。
- delta 算法覆盖首次基线、正常增长、counter reset、rule_hash 变化和零增量。
- batch builder 正确应用 `traffic_ratio` 和 `tunnel.flow`。
repository 测试:
- `nft_counter_state` 自动迁移。
- upsert 在 SQLite 下可重复刷新。
- forward 删除时清理 counter state。
handler/job 测试:
- 单节点采集成功会调用现有流量入账路径。
- 单节点 SSH 失败不影响其他节点。
- 无旧状态时不会误把历史 counter 入账。
验证命令:
```bash
(cd go-backend && go test ./...)
```
## 分阶段落地
第一阶段:
- 规则渲染加入 filter chain counter。
- 实现 SSH collector、JSON parser、counter state 和后台 job。
- 入账到现有账本和 tunnel metric。
第二阶段:
- UI 展示 nftables 采集状态。
- 节点详情显示最近采集时间和最近错误。
- 提供手动“采集一次”诊断按钮。
第三阶段:
- 探索域名目标的解析和重建策略。
- 优化大量节点下的采集调度、退避和超时配置。
## 开放问题
- 采集周期默认 60 秒是否满足产品预期;如果需要更实时,可以降到 30 秒,但 SSH 压力会增加。
- nftables 模式是否继续允许域名 remoteAddr。如果允许,需要先定义 DNS 固化和统计匹配规则。
- 是否要在第一阶段暴露采集状态 API。推荐后端先记录,UI 后续补齐。
+18 -2
View File
@@ -20,7 +20,7 @@ func (h *Handler) StartBackgroundJobs() {
ctx, cancel := context.WithCancel(context.Background())
h.jobsCancel = cancel
h.jobsStarted = true
h.jobsWG.Add(7)
h.jobsWG.Add(8)
h.jobsMu.Unlock()
go h.runHourlyStatsLoop(ctx)
@@ -30,6 +30,7 @@ func (h *Handler) StartBackgroundJobs() {
go h.runHealthChecks(ctx)
go h.runTunnelQualityProber(ctx)
go h.runValidateLicenseJob(ctx)
go h.runNftablesTrafficCollectLoop(ctx)
}
func (h *Handler) runValidateLicenseJob(ctx context.Context) {
@@ -64,7 +65,7 @@ func (h *Handler) validateLicenseJob() {
fingerprint, _ := h.repo.GetViteConfigValue("machine_fingerprint")
client := license.NewKeygenClient(accountID, "")
valResp, err := client.ValidateKeyWithFingerprint(key, fingerprint)
if err != nil {
// Network error or timeout. Grace period by not revoking immediately here.
return
@@ -128,6 +129,21 @@ func (h *Handler) runTunnelQualityProber(ctx context.Context) {
h.qualityProber.Start(ctx)
}
func (h *Handler) runNftablesTrafficCollectLoop(ctx context.Context) {
defer h.jobsWG.Done()
ticker := time.NewTicker(time.Minute)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
h.runNftablesTrafficCollectJob(time.Now())
}
}
}
func (h *Handler) runHourlyStatsLoop(ctx context.Context) {
defer h.jobsWG.Done()
+36 -10
View File
@@ -2412,23 +2412,30 @@ func (h *Handler) forwardDelete(w http.ResponseWriter, r *http.Request) {
response.WriteJSON(w, response.Err(-2, err.Error()))
return
}
if err := h.controlForwardServices(forward, "DeleteService", true); err != nil {
response.WriteJSON(w, response.ErrDefault(err.Error()))
return
}
var nftNodeID int64
if nftMode, entryNodeIDs, modeErr := h.tunnelUsesNftables(forward.TunnelID); modeErr != nil {
response.WriteJSON(w, response.Err(-2, modeErr.Error()))
return
} else if nftMode && len(entryNodeIDs) > 0 {
if err := h.reconcileNftablesNodeByRequest(entryNodeIDs[0]); err != nil {
response.WriteJSON(w, response.ErrDefault(err.Error()))
} else if nftMode {
if len(entryNodeIDs) == 0 {
response.WriteJSON(w, response.ErrDefault("nftables 转发缺少入口节点"))
return
}
nftNodeID = entryNodeIDs[0]
} else if err := h.controlForwardServices(forward, "DeleteService", true); err != nil {
response.WriteJSON(w, response.ErrDefault(err.Error()))
return
}
if err := h.deleteForwardByID(id); err != nil {
response.WriteJSON(w, response.Err(-2, err.Error()))
return
}
if nftNodeID > 0 {
if err := h.reconcileNftablesNodeByRequest(nftNodeID); err != nil {
response.WriteJSON(w, response.ErrDefault(err.Error()))
return
}
}
response.WriteJSON(w, response.OKEmpty())
}
@@ -2577,7 +2584,19 @@ func (h *Handler) forwardBatchDelete(w http.ResponseWriter, r *http.Request) {
failures = appendBatchFailure(failures, id, "", accessErr)
continue
}
if err := h.controlForwardServices(forward, "DeleteService", true); err != nil {
var nftNodeID int64
if nftMode, entryNodeIDs, modeErr := h.tunnelUsesNftables(forward.TunnelID); modeErr != nil {
f++
failures = appendBatchFailure(failures, id, forward.Name, modeErr)
continue
} else if nftMode {
if len(entryNodeIDs) == 0 {
f++
failures = appendBatchFailureReason(failures, id, forward.Name, "nftables 转发缺少入口节点")
continue
}
nftNodeID = entryNodeIDs[0]
} else if err := h.controlForwardServices(forward, "DeleteService", true); err != nil {
f++
failures = appendBatchFailure(failures, id, forward.Name, err)
continue
@@ -2585,9 +2604,16 @@ func (h *Handler) forwardBatchDelete(w http.ResponseWriter, r *http.Request) {
if err := h.deleteForwardByID(id); err != nil {
f++
failures = appendBatchFailure(failures, id, forward.Name, err)
} else {
s++
continue
}
if nftNodeID > 0 {
if err := h.reconcileNftablesNodeByRequest(nftNodeID); err != nil {
f++
failures = appendBatchFailure(failures, id, forward.Name, err)
continue
}
}
s++
}
response.WriteJSON(w, response.OK(batchOperationResult{SuccessCount: s, FailCount: f, Failures: failures}))
}
@@ -5,6 +5,7 @@ import (
"database/sql"
"errors"
"fmt"
"net"
"net/http"
"strings"
"time"
@@ -21,6 +22,7 @@ type nftablesRuntimeManager interface {
Test(ctx context.Context, cfg runtimenft.SSHConfig) error
Reconcile(ctx context.Context, cfg runtimenft.SSHConfig, plan runtimenft.NodePlan) (runtimenft.ApplyResult, error)
Clear(ctx context.Context, cfg runtimenft.SSHConfig) error
CollectCounters(ctx context.Context, cfg runtimenft.SSHConfig) ([]runtimenft.CounterSample, error)
}
func isNftablesForwardMode(mode string) bool {
@@ -77,9 +79,13 @@ func (h *Handler) validateNftablesForwardRequest(tunnel *tunnelRecord, remoteAdd
if len(entryNodeIDs) != 1 {
return errors.New("nftables 节点仅支持单入口隧道")
}
if _, err := runtimenft.ParseSingleTarget(remoteAddr); err != nil {
target, err := runtimenft.ParseSingleTarget(remoteAddr)
if err != nil {
return err
}
if net.ParseIP(strings.Trim(strings.TrimSpace(target.Host), "[]")) == nil {
return errors.New("nftables 节点仅支持 IP 目标地址")
}
return nil
}
@@ -20,13 +20,16 @@ import (
)
type fakeNftablesManager struct {
testErr error
reconcileErr error
reconcileHit int
clearErr error
clearHit int
lastConfig runtimenft.SSHConfig
lastPlan runtimenft.NodePlan
testErr error
reconcileErr error
reconcileHit int
clearErr error
clearHit int
collectErr error
collectHit int
counterSamples []runtimenft.CounterSample
lastConfig runtimenft.SSHConfig
lastPlan runtimenft.NodePlan
}
func (f *fakeNftablesManager) Test(_ context.Context, cfg runtimenft.SSHConfig) error {
@@ -53,6 +56,15 @@ func (f *fakeNftablesManager) Clear(context.Context, runtimenft.SSHConfig) error
return f.clearErr
}
func (f *fakeNftablesManager) CollectCounters(_ context.Context, cfg runtimenft.SSHConfig) ([]runtimenft.CounterSample, error) {
f.collectHit++
f.lastConfig = cfg
if f.collectErr != nil {
return nil, f.collectErr
}
return f.counterSamples, nil
}
type nftablesTestFixture struct {
handler *Handler
nodeID int64
@@ -268,6 +280,48 @@ func TestNodeUpdatePreservesExistingNftablesSecretsWhenFieldsOmitted(t *testing.
}
}
func TestValidateNftablesForwardRequestRejectsHostnameTarget(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
tunnelID := seedTunnelForNftables(t, h, "nft-tunnel", fixture.nodeID)
tunnel, err := h.getTunnelRecord(tunnelID)
if err != nil {
t.Fatalf("load tunnel: %v", err)
}
err = h.validateNftablesForwardRequest(tunnel, "example.com:443", []int64{fixture.nodeID})
if err == nil {
t.Fatalf("expected hostname target to be rejected")
}
if !strings.Contains(err.Error(), "IP") {
t.Fatalf("expected IP literal validation error, got %q", err)
}
}
func TestForwardDeleteReconcilesNftablesAfterDBDelete(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
seedNftablesSSHConfig(t, h, fixture.nodeID)
tunnelID := seedTunnelForNftables(t, h, "nft-tunnel", fixture.nodeID)
forward := seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
manager := &fakeNftablesManager{}
h.nftablesManager = manager
req := newAuthenticatedJSONRequest(t, map[string]int64{"id": forward.ID})
res := httptest.NewRecorder()
h.forwardDelete(res, req)
assertNftablesSuccessWithBody(t, res)
if manager.reconcileHit != 1 {
t.Fatalf("expected reconcile once, got %d", manager.reconcileHit)
}
if len(manager.lastPlan.Rules) != 0 {
t.Fatalf("expected reconcile after DB delete to render no rules, got %+v", manager.lastPlan.Rules)
}
if _, err := h.getForwardRecord(forward.ID); !errors.Is(err, errForwardNotFound) {
t.Fatalf("expected forward to be deleted, got %v", err)
}
}
func TestForwardForceDeleteRemovesNftablesBindingAndReconciles(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
@@ -307,6 +361,30 @@ func TestForwardForceDeleteRemovesNftablesBindingAndReconciles(t *testing.T) {
}
}
func TestForwardBatchDeleteReconcilesNftablesAfterDBDelete(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
seedNftablesSSHConfig(t, h, fixture.nodeID)
tunnelID := seedTunnelForNftables(t, h, "nft-tunnel", fixture.nodeID)
forward := seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
manager := &fakeNftablesManager{}
h.nftablesManager = manager
req := newAuthenticatedJSONRequest(t, map[string][]int64{"ids": {forward.ID}})
res := httptest.NewRecorder()
h.forwardBatchDelete(res, req)
assertNftablesSuccessWithBody(t, res)
if manager.reconcileHit != 1 {
t.Fatalf("expected reconcile once, got %d", manager.reconcileHit)
}
if len(manager.lastPlan.Rules) != 0 {
t.Fatalf("expected reconcile after DB delete to render no rules, got %+v", manager.lastPlan.Rules)
}
if _, err := h.getForwardRecord(forward.ID); !errors.Is(err, errForwardNotFound) {
t.Fatalf("expected forward to be deleted, got %v", err)
}
}
func TestForwardBatchRedeployUsesNftablesReconcile(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
@@ -0,0 +1,400 @@
package handler
import (
"context"
"log"
"math"
"sort"
"strings"
"time"
runtimenft "go-backend/internal/runtime/nftables"
"go-backend/internal/store/model"
"go-backend/internal/store/repo"
)
type nftTrafficDelta struct {
ForwardID int64
BytesIn int64
BytesOut int64
}
type nftCounterStateKey struct {
forwardID int64
protocol string
direction string
}
func (h *Handler) runNftablesTrafficCollectJob(now time.Time) {
if h == nil || h.repo == nil {
return
}
nodes, err := h.repo.ListNftablesNodesForCollection()
if err != nil {
log.Printf("nftables traffic collection failed op=list_nodes err=%v", err)
return
}
for i := range nodes {
node := &nodes[i]
h.collectNftablesNodeTraffic(node.NodeID, &node.Config, now)
}
}
func (h *Handler) collectNftablesNodeTraffic(nodeID int64, cfgModel *model.NodeSSHConfig, now time.Time) {
if h == nil || h.repo == nil {
return
}
if h.nftablesManager == nil {
log.Printf("nftables traffic collection failed op=collect node_id=%d err=%v", nodeID, "nftables manager not initialized")
return
}
sshCfg, err := sshConfigFromModel(cfgModel)
if err != nil {
log.Printf("nftables traffic collection failed op=ssh_config node_id=%d err=%v", nodeID, err)
return
}
samples, err := h.nftablesManager.CollectCounters(context.Background(), sshCfg)
if err != nil {
log.Printf("nftables traffic collection failed op=collect node_id=%d err=%v", nodeID, err)
return
}
oldStates, err := h.repo.GetNftCounterStatesByNode(nodeID)
if err != nil {
log.Printf("nftables traffic collection failed op=list_states node_id=%d err=%v", nodeID, err)
return
}
bindings, err := h.repo.ListNftRuleBindingsByNode(nodeID)
if err != nil {
log.Printf("nftables traffic collection failed op=list_bindings node_id=%d err=%v", nodeID, err)
return
}
hashes := make(map[int64]string, len(bindings))
for _, binding := range bindings {
if strings.ToLower(strings.TrimSpace(binding.Status)) != runtimenft.StatusApplied {
continue
}
ruleHash := strings.TrimSpace(binding.RuleHash)
if ruleHash == "" {
continue
}
hashes[binding.ForwardID] = ruleHash
}
nowMs := now.UnixMilli()
boundSamples := filterNftCounterSamplesWithBinding(samples, hashes)
deltas, newStates := buildNftCounterDeltas(nodeID, boundSamples, oldStates, hashes, nowMs)
if len(newStates) == 0 {
if len(deltas) != 0 {
log.Printf("nftables traffic collection skipped suspicious deltas without states node_id=%d deltas=%d", nodeID, len(deltas))
}
return
}
var metas map[int64]repo.FlowUploadForwardMeta
forwardIDs := make([]int64, 0, len(deltas))
for _, delta := range deltas {
if delta.ForwardID > 0 {
forwardIDs = append(forwardIDs, delta.ForwardID)
}
}
if len(deltas) != 0 {
metas, err = h.repo.GetFlowUploadForwardMetas(forwardIDs)
if err != nil {
log.Printf("nftables traffic collection failed op=load_flow_metas node_id=%d err=%v", nodeID, err)
return
}
if missingForwardID, ok := firstNftDeltaMissingMeta(deltas, metas); ok {
log.Printf("nftables traffic collection skipped state advance op=missing_flow_meta node_id=%d forward_id=%d", nodeID, missingForwardID)
return
}
}
if len(deltas) == 0 {
if err := h.repo.UpsertNftCounterStates(newStates, nowMs); err != nil {
log.Printf("nftables traffic collection failed op=upsert_states node_id=%d err=%v", nodeID, err)
return
}
return
}
batch := buildNftFlowUploadBatch(deltas, metas)
if missingForwardID, ok := firstNftBatchMissingDelta(deltas, batch); ok {
log.Printf("nftables traffic collection skipped state advance op=unaccounted_delta node_id=%d forward_id=%d", nodeID, missingForwardID)
return
}
quotaViews, err := h.repo.ApplyNftTrafficAccounting(batch.flowDeltas, batch.quotaUsage, newStates, now)
if err != nil {
log.Printf("nftables traffic collection failed op=accounting node_id=%d err=%v", nodeID, err)
return
}
h.recordTunnelMetricsFromForwardBatch(nodeID, batch.forwardTraffic, metas, nowMs)
for userID, quota := range quotaViews {
h.enforceUserQuotaIfNeeded(userID, quota)
}
for _, target := range batch.policyTargets {
if target.UserID <= 0 || target.UserTunnelID <= 0 {
continue
}
h.enforceFlowPolicies(target.UserID, target.UserTunnelID)
}
}
func firstNftBatchMissingDelta(deltas []nftTrafficDelta, batch flowUploadBatch) (int64, bool) {
flowSeen := make(map[int64]struct{}, len(batch.flowDeltas))
for _, delta := range batch.flowDeltas {
flowSeen[delta.ForwardID] = struct{}{}
}
expectedRaw := make(map[int64]tunnelTrafficDelta, len(batch.forwardTraffic))
for _, delta := range deltas {
if delta.ForwardID <= 0 || (delta.BytesIn == 0 && delta.BytesOut == 0) {
continue
}
if delta.BytesIn < 0 || delta.BytesOut < 0 {
return delta.ForwardID, true
}
raw := expectedRaw[delta.ForwardID]
if raw.bytesIn > math.MaxInt64-delta.BytesIn || raw.bytesOut > math.MaxInt64-delta.BytesOut {
return delta.ForwardID, true
}
raw.bytesIn += delta.BytesIn
raw.bytesOut += delta.BytesOut
expectedRaw[delta.ForwardID] = raw
}
for forwardID, expected := range expectedRaw {
actual, ok := batch.forwardTraffic[forwardID]
if !ok || actual.bytesIn != expected.bytesIn || actual.bytesOut != expected.bytesOut {
return forwardID, true
}
if expected.bytesIn != 0 || expected.bytesOut != 0 {
if _, ok := flowSeen[forwardID]; !ok {
return forwardID, true
}
}
}
return 0, false
}
func firstNftDeltaMissingMeta(deltas []nftTrafficDelta, metas map[int64]repo.FlowUploadForwardMeta) (int64, bool) {
for _, delta := range deltas {
if delta.ForwardID <= 0 {
continue
}
if _, ok := metas[delta.ForwardID]; !ok {
return delta.ForwardID, true
}
}
return 0, false
}
func filterNftCounterSamplesWithBinding(samples []runtimenft.CounterSample, hashes map[int64]string) []runtimenft.CounterSample {
if len(samples) == 0 || len(hashes) == 0 {
return nil
}
filtered := make([]runtimenft.CounterSample, 0, len(samples))
for _, sample := range samples {
if _, ok := hashes[sample.ForwardID]; !ok {
continue
}
filtered = append(filtered, sample)
}
return filtered
}
func nftCounterKey(forwardID int64, protocol, direction string) nftCounterStateKey {
return nftCounterStateKey{
forwardID: forwardID,
protocol: strings.ToLower(strings.TrimSpace(protocol)),
direction: strings.ToLower(strings.TrimSpace(direction)),
}
}
func buildNftCounterDeltas(nodeID int64, samples []runtimenft.CounterSample, oldStates []model.NftCounterState, hashes map[int64]string, nowMs int64) ([]nftTrafficDelta, []repo.NftCounterStateInput) {
oldByKey := make(map[nftCounterStateKey]model.NftCounterState, len(oldStates))
for _, old := range oldStates {
if old.NodeID != nodeID {
continue
}
oldByKey[nftCounterKey(old.ForwardID, old.Protocol, old.Direction)] = old
}
stateInputs := make([]repo.NftCounterStateInput, 0, len(samples))
deltaByForward := make(map[int64]nftTrafficDelta)
for _, sample := range samples {
direction := strings.ToLower(strings.TrimSpace(sample.Direction))
if direction != runtimenft.CounterDirectionToTarget && direction != runtimenft.CounterDirectionFromTarget {
continue
}
protocol := strings.ToLower(strings.TrimSpace(sample.Protocol))
if protocol != "tcp" && protocol != "udp" {
continue
}
if sample.Bytes > uint64(math.MaxInt64) || sample.Packets > uint64(math.MaxInt64) {
continue
}
ruleHash := strings.TrimSpace(hashes[sample.ForwardID])
stateInput := repo.NftCounterStateInput{
NodeID: nodeID,
ForwardID: sample.ForwardID,
Protocol: protocol,
Direction: direction,
RuleHash: ruleHash,
Bytes: sample.Bytes,
Packets: sample.Packets,
CollectedTime: nowMs,
}
old, exists := oldByKey[nftCounterKey(sample.ForwardID, protocol, direction)]
if !exists || old.RuleHash != ruleHash {
stateInputs = append(stateInputs, stateInput)
continue
}
if old.Bytes < 0 {
stateInputs = append(stateInputs, stateInput)
continue
}
oldBytes := uint64(old.Bytes)
if sample.Bytes < oldBytes {
stateInputs = append(stateInputs, stateInput)
continue
}
rawDelta := sample.Bytes - oldBytes
if rawDelta == 0 {
stateInputs = append(stateInputs, stateInput)
continue
}
delta := deltaByForward[sample.ForwardID]
delta.ForwardID = sample.ForwardID
rawDeltaInt := int64(rawDelta)
if direction == runtimenft.CounterDirectionToTarget {
if delta.BytesIn > math.MaxInt64-rawDeltaInt {
continue
}
delta.BytesIn += rawDeltaInt
} else {
if delta.BytesOut > math.MaxInt64-rawDeltaInt {
continue
}
delta.BytesOut += rawDeltaInt
}
stateInputs = append(stateInputs, stateInput)
deltaByForward[sample.ForwardID] = delta
}
forwardIDs := make([]int64, 0, len(deltaByForward))
for forwardID := range deltaByForward {
forwardIDs = append(forwardIDs, forwardID)
}
sort.Slice(forwardIDs, func(i, j int) bool { return forwardIDs[i] < forwardIDs[j] })
deltas := make([]nftTrafficDelta, 0, len(forwardIDs))
for _, forwardID := range forwardIDs {
delta := deltaByForward[forwardID]
if delta.BytesIn == 0 && delta.BytesOut == 0 {
continue
}
deltas = append(deltas, delta)
}
return deltas, stateInputs
}
func buildNftFlowUploadBatch(deltas []nftTrafficDelta, metas map[int64]repo.FlowUploadForwardMeta) flowUploadBatch {
batch := flowUploadBatch{
quotaUsage: make(map[int64]int64),
forwardTraffic: make(map[int64]tunnelTrafficDelta),
orphanServices: make(map[string]struct{}),
peerShareForwardItems: make(map[string]flowItem),
peerShareRuntimeItems: make(map[int64]flowItem),
}
policySeen := map[flowPolicyTarget]struct{}{}
flowSeen := map[int64]int{}
for _, delta := range deltas {
meta, exists := metas[delta.ForwardID]
if !exists {
continue
}
raw := batch.forwardTraffic[delta.ForwardID]
if delta.BytesIn < 0 || delta.BytesOut < 0 || raw.bytesIn > math.MaxInt64-delta.BytesIn || raw.bytesOut > math.MaxInt64-delta.BytesOut {
continue
}
scaledIn, ok := scaleNftTrafficBytes(delta.BytesIn, meta.TrafficRatio, meta.TunnelFlow)
if !ok {
continue
}
scaledOut, ok := scaleNftTrafficBytes(delta.BytesOut, meta.TrafficRatio, meta.TunnelFlow)
if !ok {
continue
}
if scaledIn > math.MaxInt64-scaledOut {
continue
}
quotaDelta := scaledIn + scaledOut
if batch.quotaUsage[meta.UserID] > math.MaxInt64-quotaDelta {
continue
}
flowIdx, flowExists := flowSeen[delta.ForwardID]
if flowExists && (batch.flowDeltas[flowIdx].InFlow > math.MaxInt64-scaledIn || batch.flowDeltas[flowIdx].OutFlow > math.MaxInt64-scaledOut) {
continue
}
raw.bytesIn += delta.BytesIn
raw.bytesOut += delta.BytesOut
batch.forwardTraffic[delta.ForwardID] = raw
if flowExists {
batch.flowDeltas[flowIdx].InFlow += scaledIn
batch.flowDeltas[flowIdx].OutFlow += scaledOut
} else {
flowSeen[delta.ForwardID] = len(batch.flowDeltas)
batch.flowDeltas = append(batch.flowDeltas, repo.FlowUploadCounterDelta{
ForwardID: delta.ForwardID,
UserID: meta.UserID,
UserTunnelID: meta.UserTunnelID,
InFlow: scaledIn,
OutFlow: scaledOut,
})
}
batch.quotaUsage[meta.UserID] += quotaDelta
target := flowPolicyTarget{UserID: meta.UserID, UserTunnelID: meta.UserTunnelID}
if _, seen := policySeen[target]; !seen {
policySeen[target] = struct{}{}
batch.policyTargets = append(batch.policyTargets, target)
}
}
sort.Slice(batch.policyTargets, func(i, j int) bool {
if batch.policyTargets[i].UserID == batch.policyTargets[j].UserID {
return batch.policyTargets[i].UserTunnelID < batch.policyTargets[j].UserTunnelID
}
return batch.policyTargets[i].UserID < batch.policyTargets[j].UserID
})
return batch
}
func scaleNftTrafficBytes(bytes int64, ratio float64, tunnelFlow int64) (int64, bool) {
if bytes < 0 || ratio < 0 || tunnelFlow < 0 {
return 0, false
}
var scaled int64
if ratio == 1 {
scaled = bytes
} else {
scaledFloat := float64(bytes) * ratio
if math.IsNaN(scaledFloat) || math.IsInf(scaledFloat, 0) || scaledFloat < 0 || scaledFloat >= math.Pow(2, 63) {
return 0, false
}
scaled = int64(scaledFloat)
}
if tunnelFlow != 0 && scaled > math.MaxInt64/tunnelFlow {
return 0, false
}
return scaled * tunnelFlow, true
}
@@ -0,0 +1,680 @@
package handler
import (
"errors"
"math"
"testing"
"time"
runtimenft "go-backend/internal/runtime/nftables"
"go-backend/internal/store/model"
"go-backend/internal/store/repo"
)
func TestBuildNftCounterDeltasSavesFirstBaselineWithoutDelta(t *testing.T) {
nowMs := int64(1700000000123)
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
}, nil, map[int64]string{42: "hash-a"}, nowMs)
if len(deltas) != 0 {
t.Fatalf("expected no deltas for first baseline, got %#v", deltas)
}
if len(states) != 1 {
t.Fatalf("expected one state input, got %d", len(states))
}
state := states[0]
if state.NodeID != 11 || state.ForwardID != 42 || state.Protocol != "tcp" || state.Direction != runtimenft.CounterDirectionToTarget {
t.Fatalf("unexpected state identity: %#v", state)
}
if state.RuleHash != "hash-a" || state.Bytes != 1000 || state.Packets != 10 || state.CollectedTime != nowMs {
t.Fatalf("unexpected state values: %#v", state)
}
}
func TestBuildNftCounterDeltasNormalGrowthProducesDirectionalBytes(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1500, Packets: 15},
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 2600, Packets: 26},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, RuleHash: "hash-a", Bytes: 2000, Packets: 20},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 1 {
t.Fatalf("expected one aggregated delta, got %#v", deltas)
}
if deltas[0].ForwardID != 42 || deltas[0].BytesIn != 500 || deltas[0].BytesOut != 600 {
t.Fatalf("unexpected delta: %#v", deltas[0])
}
if len(states) != 2 {
t.Fatalf("expected two state inputs, got %d", len(states))
}
}
func TestBuildNftCounterDeltasResetRefreshesBaselineWithoutDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 25, Packets: 2},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected reset to produce no deltas, got %#v", deltas)
}
if len(states) != 1 || states[0].Bytes != 25 || states[0].RuleHash != "hash-a" {
t.Fatalf("expected refreshed baseline state, got %#v", states)
}
}
func TestBuildNftCounterDeltasRuleHashChangeRefreshesBaselineWithoutDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1500, Packets: 15},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-b"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected rule hash change to produce no deltas, got %#v", deltas)
}
if len(states) != 1 || states[0].Bytes != 1500 || states[0].RuleHash != "hash-b" {
t.Fatalf("expected refreshed hash baseline state, got %#v", states)
}
}
func TestBuildNftCounterDeltasEqualBytesRefreshesBaselineWithoutDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 11},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected equal bytes to produce no deltas, got %#v", deltas)
}
if len(states) != 1 || states[0].Bytes != 1000 || states[0].Packets != 11 || states[0].RuleHash != "hash-a" {
t.Fatalf("expected refreshed baseline state, got %#v", states)
}
}
func TestBuildNftCounterDeltasAggregatesProtocolsAndDirections(t *testing.T) {
deltas, _ := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1100, Packets: 11},
{ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 2200, Packets: 22},
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 3300, Packets: 33},
{ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 4400, Packets: 44},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000},
{NodeID: 11, ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 2000},
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, RuleHash: "hash-a", Bytes: 3000},
{NodeID: 11, ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionFromTarget, RuleHash: "hash-a", Bytes: 4000},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 1 {
t.Fatalf("expected one aggregated delta, got %#v", deltas)
}
if deltas[0].ForwardID != 42 || deltas[0].BytesIn != 300 || deltas[0].BytesOut != 700 {
t.Fatalf("unexpected aggregated delta: %#v", deltas[0])
}
}
func TestBuildNftCounterDeltasSkipsInvalidProtocolBeforeStateAndDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "icmp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1500, Packets: 15},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "icmp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected invalid protocol to produce no deltas, got %#v", deltas)
}
if len(states) != 0 {
t.Fatalf("expected invalid protocol to produce no state inputs, got %#v", states)
}
}
func TestBuildNftCounterDeltasSkipsOversizedPacketsBeforeStateAndDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1500, Packets: uint64(math.MaxInt64) + 1},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected oversized packets to produce no deltas, got %#v", deltas)
}
if len(states) != 0 {
t.Fatalf("expected oversized packets to produce no state inputs, got %#v", states)
}
}
func TestBuildNftCounterDeltasSkipsOversizedBytesBeforeStateAndDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: uint64(math.MaxInt64) + 1, Packets: 10},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1000, Packets: 10},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected oversized bytes to produce no deltas, got %#v", deltas)
}
if len(states) != 0 {
t.Fatalf("expected oversized bytes to produce no state inputs, got %#v", states)
}
}
func TestBuildNftCounterDeltasSkipsOverflowingAggregateSampleWithoutState(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: uint64(math.MaxInt64), Packets: 10},
{ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 10, Packets: 1},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1, Packets: 1},
{NodeID: 11, ForwardID: 42, Protocol: "udp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-a", Bytes: 1, Packets: 1},
}, map[int64]string{42: "hash-a"}, 2000)
if len(deltas) != 1 {
t.Fatalf("expected only non-overflowing aggregate delta, got %#v", deltas)
}
if deltas[0].ForwardID != 42 || deltas[0].BytesIn != math.MaxInt64-1 || deltas[0].BytesOut != 0 {
t.Fatalf("unexpected aggregate delta: %#v", deltas[0])
}
if len(states) != 1 {
t.Fatalf("expected only the accounted safe sample to advance baseline, got %#v", states)
}
if states[0].ForwardID != 42 || states[0].Protocol != "tcp" || states[0].Bytes != uint64(math.MaxInt64) {
t.Fatalf("expected safe sample state input to be preserved, got %#v", states)
}
}
func TestBuildNftCounterDeltasSkipsUnknownDirectionAndOversizedDelta(t *testing.T) {
deltas, states := buildNftCounterDeltas(11, []runtimenft.CounterSample{
{ForwardID: 42, Protocol: "tcp", Direction: "sideways", Bytes: 1500, Packets: 15},
{ForwardID: 43, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: uint64(math.MaxInt64) + 1, Packets: 1},
}, []model.NftCounterState{
{NodeID: 11, ForwardID: 43, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, RuleHash: "hash-b", Bytes: 100},
}, map[int64]string{42: "hash-a", 43: "hash-b"}, 2000)
if len(deltas) != 0 {
t.Fatalf("expected no delta for skipped/oversized samples, got %#v", deltas)
}
if len(states) != 0 {
t.Fatalf("expected no state inputs for skipped/oversized samples, got %#v", states)
}
}
func TestBuildNftFlowUploadBatchScalesFlowAndPreservesRawTunnelTraffic(t *testing.T) {
batch := buildNftFlowUploadBatch([]nftTrafficDelta{
{ForwardID: 20, BytesIn: 80, BytesOut: 110},
{ForwardID: 21, BytesIn: 7, BytesOut: 11},
{ForwardID: 20, BytesIn: 20, BytesOut: 10},
}, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TunnelID: 1, TrafficRatio: 2, TunnelFlow: 3},
21: {ForwardID: 21, UserID: 2, UserTunnelID: 10, TunnelID: 1, TrafficRatio: 1.5, TunnelFlow: 2},
})
if len(batch.flowDeltas) != 2 {
t.Fatalf("expected two flow deltas, got %#v", batch.flowDeltas)
}
if batch.flowDeltas[0].ForwardID != 20 || batch.flowDeltas[0].InFlow != 600 || batch.flowDeltas[0].OutFlow != 720 {
t.Fatalf("unexpected first flow delta: %#v", batch.flowDeltas[0])
}
if batch.flowDeltas[1].ForwardID != 21 || batch.flowDeltas[1].InFlow != 20 || batch.flowDeltas[1].OutFlow != 32 {
t.Fatalf("unexpected second flow delta: %#v", batch.flowDeltas[1])
}
if batch.quotaUsage[2] != 1372 {
t.Fatalf("expected quota usage 1372, got %d", batch.quotaUsage[2])
}
if len(batch.policyTargets) != 1 || batch.policyTargets[0].UserID != 2 || batch.policyTargets[0].UserTunnelID != 10 {
t.Fatalf("expected deduped policy target, got %#v", batch.policyTargets)
}
if traffic := batch.forwardTraffic[20]; traffic.bytesIn != 100 || traffic.bytesOut != 120 {
t.Fatalf("expected raw traffic for forward 20, got %#v", traffic)
}
if traffic := batch.forwardTraffic[21]; traffic.bytesIn != 7 || traffic.bytesOut != 11 {
t.Fatalf("expected raw traffic for forward 21, got %#v", traffic)
}
}
func TestBuildNftFlowUploadBatchSkipsOverflowingScaledFlow(t *testing.T) {
batch := buildNftFlowUploadBatch([]nftTrafficDelta{
{ForwardID: 20, BytesIn: math.MaxInt64, BytesOut: 0},
}, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TrafficRatio: 2, TunnelFlow: 2},
})
if len(batch.flowDeltas) != 0 {
t.Fatalf("expected overflowing scaled flow to be skipped, got %#v", batch.flowDeltas)
}
if len(batch.quotaUsage) != 0 {
t.Fatalf("expected no quota usage for overflowing scaled flow, got %#v", batch.quotaUsage)
}
if len(batch.policyTargets) != 0 {
t.Fatalf("expected no policy targets for overflowing scaled flow, got %#v", batch.policyTargets)
}
if len(batch.forwardTraffic) != 0 {
t.Fatalf("expected no raw traffic for overflowing scaled flow, got %#v", batch.forwardTraffic)
}
}
func TestBuildNftFlowUploadBatchSkipsRawForwardTrafficOverflow(t *testing.T) {
batch := buildNftFlowUploadBatch([]nftTrafficDelta{
{ForwardID: 20, BytesIn: math.MaxInt64, BytesOut: 0},
{ForwardID: 20, BytesIn: 1, BytesOut: 0},
}, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TrafficRatio: 0.5, TunnelFlow: 1},
})
traffic := batch.forwardTraffic[20]
if traffic.bytesIn != math.MaxInt64 || traffic.bytesOut != 0 {
t.Fatalf("expected overflowing raw delta to be skipped without negative traffic, got %#v", traffic)
}
if len(batch.flowDeltas) != 1 || batch.flowDeltas[0].ForwardID != 20 {
t.Fatalf("expected only the safe flow delta, got %#v", batch.flowDeltas)
}
if len(batch.policyTargets) != 1 || batch.policyTargets[0].UserID != 2 || batch.policyTargets[0].UserTunnelID != 10 {
t.Fatalf("expected policy target only from safe delta, got %#v", batch.policyTargets)
}
}
func TestBuildNftFlowUploadBatchSkipsQuotaOverflow(t *testing.T) {
batch := buildNftFlowUploadBatch([]nftTrafficDelta{
{ForwardID: 20, BytesIn: math.MaxInt64, BytesOut: 0},
{ForwardID: 21, BytesIn: 1, BytesOut: 0},
}, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TrafficRatio: 1, TunnelFlow: 1},
21: {ForwardID: 21, UserID: 2, UserTunnelID: 10, TrafficRatio: 1, TunnelFlow: 1},
})
if len(batch.flowDeltas) != 1 || batch.flowDeltas[0].ForwardID != 20 || batch.flowDeltas[0].InFlow != math.MaxInt64 {
t.Fatalf("expected only non-overflowing quota delta, got %#v", batch.flowDeltas)
}
if batch.quotaUsage[2] != math.MaxInt64 {
t.Fatalf("expected quota usage to remain at max int64, got %#v", batch.quotaUsage)
}
if len(batch.policyTargets) != 1 || batch.policyTargets[0].UserID != 2 || batch.policyTargets[0].UserTunnelID != 10 {
t.Fatalf("expected one policy target from non-overflowing delta, got %#v", batch.policyTargets)
}
if _, ok := batch.forwardTraffic[21]; ok {
t.Fatalf("expected quota-overflowing delta to be skipped from raw traffic")
}
}
func TestBuildNftFlowUploadBatchSkipsMissingMeta(t *testing.T) {
batch := buildNftFlowUploadBatch([]nftTrafficDelta{
{ForwardID: 20, BytesIn: 80, BytesOut: 110},
{ForwardID: 99, BytesIn: 1, BytesOut: 2},
}, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TrafficRatio: 1, TunnelFlow: 1},
})
if len(batch.flowDeltas) != 1 || batch.flowDeltas[0].ForwardID != 20 {
t.Fatalf("expected only forward 20 delta, got %#v", batch.flowDeltas)
}
if _, ok := batch.forwardTraffic[99]; ok {
t.Fatalf("expected missing meta forward to be skipped from raw traffic")
}
}
func TestNftBatchCoversDeltasRequiresRawAndFlowEntries(t *testing.T) {
deltas := []nftTrafficDelta{{ForwardID: 20, BytesIn: 1, BytesOut: 0}}
batch := flowUploadBatch{
forwardTraffic: map[int64]tunnelTrafficDelta{20: {bytesIn: 1}},
flowDeltas: []repo.FlowUploadCounterDelta{{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 1}},
}
if missing, ok := firstNftBatchMissingDelta(deltas, batch); ok || missing != 0 {
t.Fatalf("expected batch to cover delta, missing=%d ok=%v", missing, ok)
}
delete(batch.forwardTraffic, 20)
if missing, ok := firstNftBatchMissingDelta(deltas, batch); !ok || missing != 20 {
t.Fatalf("expected missing raw traffic for forward 20, got missing=%d ok=%v", missing, ok)
}
batch.forwardTraffic[20] = tunnelTrafficDelta{bytesIn: 1}
batch.flowDeltas = nil
if missing, ok := firstNftBatchMissingDelta(deltas, batch); !ok || missing != 20 {
t.Fatalf("expected missing flow delta for forward 20, got missing=%d ok=%v", missing, ok)
}
}
func TestNftBatchCoversDeltasRequiresAggregateRawTotals(t *testing.T) {
deltas := []nftTrafficDelta{
{ForwardID: 20, BytesIn: math.MaxInt64, BytesOut: 0},
{ForwardID: 20, BytesIn: 1, BytesOut: 0},
}
batch := buildNftFlowUploadBatch(deltas, map[int64]repo.FlowUploadForwardMeta{
20: {ForwardID: 20, UserID: 2, UserTunnelID: 10, TrafficRatio: 0.5, TunnelFlow: 1},
})
if missing, ok := firstNftBatchMissingDelta(deltas, batch); !ok || missing != 20 {
t.Fatalf("expected aggregate raw overflow/mismatch for forward 20, got missing=%d ok=%v", missing, ok)
}
}
func TestCollectNftablesNodeTrafficFirstBaselineSavesStateWithoutFlow(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
manager := &fakeNftablesManager{counterSamples: []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 2000, Packets: 20},
}}
h.nftablesManager = manager
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
if manager.collectHit != 1 {
t.Fatalf("expected one collection, got %d", manager.collectHit)
}
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 2 {
t.Fatalf("expected two baseline states, got %+v", states)
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no forward flow on baseline, got %d", got)
}
if got := mustHandlerCount(t, h, `SELECT out_flow FROM user WHERE id = 1`); got != 0 {
t.Fatalf("expected no user flow on baseline, got %d", got)
}
}
func TestCollectNftablesNodeTrafficGrowthAppliesFlowAndUpdatesState(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
manager := &fakeNftablesManager{}
h.nftablesManager = manager
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 2000, Packets: 20},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1400, Packets: 14},
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionFromTarget, Bytes: 2600, Packets: 26},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000060, 0))
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 400 {
t.Fatalf("expected forward in_flow=400, got %d", got)
}
if got := mustHandlerCount(t, h, `SELECT out_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 600 {
t.Fatalf("expected forward out_flow=600, got %d", got)
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM user WHERE id = 1`); got != 400 {
t.Fatalf("expected user in_flow=400, got %d", got)
}
if got := mustHandlerCount(t, h, `SELECT out_flow FROM user_tunnel WHERE id = ?`, fixture.userTunnelID); got != 600 {
t.Fatalf("expected user_tunnel out_flow=600, got %d", got)
}
if got := mustHandlerCount(t, h, `SELECT COALESCE((SELECT daily_used_bytes FROM user_quota WHERE user_id = 1), 0)`); got != 1000 {
t.Fatalf("expected daily quota usage=1000, got %d", got)
}
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 2 {
t.Fatalf("expected two states after growth, got %+v", states)
}
for _, state := range states {
if state.Direction == runtimenft.CounterDirectionToTarget && state.Bytes != 1400 {
t.Fatalf("expected to-target state bytes 1400, got %+v", state)
}
if state.Direction == runtimenft.CounterDirectionFromTarget && state.Bytes != 2600 {
t.Fatalf("expected from-target state bytes 2600, got %+v", state)
}
}
}
func TestCollectNftablesNodeTrafficSkippedBatchDeltaDoesNotAdvanceState(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
if err := h.repo.DB().Exec(`UPDATE tunnel SET traffic_ratio = 2 WHERE id = (SELECT tunnel_id FROM forward WHERE id = ?)`, fixture.forwardID).Error; err != nil {
t.Fatalf("update tunnel ratio: %v", err)
}
manager := &fakeNftablesManager{}
h.nftablesManager = manager
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 0, Packets: 0},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: uint64(math.MaxInt64), Packets: 1},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000060, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 1 {
t.Fatalf("expected one state, got %+v", states)
}
if states[0].Bytes != 0 || states[0].Packets != 0 {
t.Fatalf("expected state to remain at old baseline after skipped batch delta, got %+v", states[0])
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no forward flow for skipped batch delta, got %d", got)
}
}
func TestCollectNftablesNodeTrafficMetadataErrorDoesNotAdvanceState(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
manager := &fakeNftablesManager{}
h.nftablesManager = manager
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
if err := h.repo.DB().Exec(`DROP TABLE tunnel`).Error; err != nil {
t.Fatalf("drop tunnel table: %v", err)
}
manager.counterSamples = []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1400, Packets: 14},
}
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000060, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 1 {
t.Fatalf("expected one baseline state, got %+v", states)
}
if states[0].Bytes != 1000 || states[0].Packets != 10 {
t.Fatalf("expected state to remain at first baseline after metadata failure, got %+v", states[0])
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no flow after metadata failure, got %d", got)
}
}
func TestCollectNftablesNodeTrafficMissingMetaDoesNotAdvanceState(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
seedNftablesSSHConfig(t, h, fixture.nodeID)
forwardID := int64(4242)
nowMs := time.Now().UnixMilli()
if err := h.repo.UpsertNftRuleBinding(repo.NftRuleBindingInput{
ForwardID: forwardID,
NodeID: fixture.nodeID,
InPort: 20000,
Protocols: "tcp",
TargetAddr: "203.0.113.9:8080",
RuleHash: "hash-a",
Status: runtimenft.StatusApplied,
}, nowMs); err != nil {
t.Fatalf("seed stale applied binding: %v", err)
}
if err := h.repo.UpsertNftCounterStates([]repo.NftCounterStateInput{{
NodeID: fixture.nodeID,
ForwardID: forwardID,
Protocol: "tcp",
Direction: runtimenft.CounterDirectionToTarget,
RuleHash: "hash-a",
Bytes: 1000,
Packets: 10,
CollectedTime: nowMs,
}}, nowMs); err != nil {
t.Fatalf("seed counter state: %v", err)
}
h.nftablesManager = &fakeNftablesManager{counterSamples: []runtimenft.CounterSample{
{ForwardID: forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1400, Packets: 14},
}}
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000060, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 1 {
t.Fatalf("expected one state, got %+v", states)
}
if states[0].Bytes != 1000 || states[0].Packets != 10 {
t.Fatalf("expected state to remain at old baseline when meta is missing, got %+v", states[0])
}
}
func TestCollectNftablesNodeTrafficSkipsSamplesWithoutBinding(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
if err := h.repo.DeleteNftRuleBindingsByForward(fixture.forwardID); err != nil {
t.Fatalf("delete nft binding: %v", err)
}
h.nftablesManager = &fakeNftablesManager{counterSamples: []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
}}
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 0 {
t.Fatalf("expected no state for unbound sample, got %+v", states)
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no flow for unbound sample, got %d", got)
}
}
func TestCollectNftablesNodeTrafficSkipsNonAppliedBinding(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
if err := h.repo.MarkNftRuleBindingError(fixture.forwardID, fixture.nodeID, "apply failed", time.Now().UnixMilli()); err != nil {
t.Fatalf("mark binding error: %v", err)
}
h.nftablesManager = &fakeNftablesManager{counterSamples: []runtimenft.CounterSample{
{ForwardID: fixture.forwardID, Protocol: "tcp", Direction: runtimenft.CounterDirectionToTarget, Bytes: 1000, Packets: 10},
}}
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 0 {
t.Fatalf("expected no state for non-applied binding, got %+v", states)
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no flow for non-applied binding, got %d", got)
}
}
func TestCollectNftablesNodeTrafficCollectionErrorDoesNotWriteState(t *testing.T) {
fixture := setupNftablesCollectionFixture(t)
h := fixture.handler
h.nftablesManager = &fakeNftablesManager{collectErr: errors.New("ssh failed")}
cfg := mustCollectionSSHConfig(t, h, fixture.nodeID)
h.collectNftablesNodeTraffic(fixture.nodeID, cfg, time.Unix(1700000000, 0))
states, err := h.repo.GetNftCounterStatesByNode(fixture.nodeID)
if err != nil {
t.Fatalf("load states: %v", err)
}
if len(states) != 0 {
t.Fatalf("expected no state on collection error, got %+v", states)
}
if got := mustHandlerCount(t, h, `SELECT in_flow FROM forward WHERE id = ?`, fixture.forwardID); got != 0 {
t.Fatalf("expected no flow on collection error, got %d", got)
}
}
type nftablesCollectionFixture struct {
handler *Handler
nodeID int64
forwardID int64
userTunnelID int64
}
func setupNftablesCollectionFixture(t *testing.T) nftablesCollectionFixture {
t.Helper()
fixture := setupNftablesHandler(t)
h := fixture.handler
seedNftablesSSHConfig(t, h, fixture.nodeID)
tunnelID := seedTunnelForNftables(t, h, "nft-traffic-tunnel", fixture.nodeID)
now := time.Now().UnixMilli()
if err := h.repo.DB().Exec(`
INSERT INTO user_tunnel(user_id, tunnel_id, speed_id, num, flow, in_flow, out_flow, flow_reset_time, exp_time, status)
VALUES(1, ?, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1)
`, tunnelID).Error; err != nil {
t.Fatalf("seed user_tunnel: %v", err)
}
forward := seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
if err := h.repo.UpsertNftRuleBinding(repo.NftRuleBindingInput{
ForwardID: forward.ID,
NodeID: fixture.nodeID,
InPort: 20000,
Protocols: "tcp,udp",
TargetAddr: "203.0.113.9:8080",
RuleHash: "hash-a",
Status: runtimenft.StatusApplied,
}, now); err != nil {
t.Fatalf("seed nft binding: %v", err)
}
userTunnelID := mustHandlerCount(t, h, `SELECT id FROM user_tunnel WHERE user_id = 1 AND tunnel_id = ?`, tunnelID)
return nftablesCollectionFixture{
handler: h,
nodeID: fixture.nodeID,
forwardID: forward.ID,
userTunnelID: userTunnelID,
}
}
func mustCollectionSSHConfig(t *testing.T, h *Handler, nodeID int64) *model.NodeSSHConfig {
t.Helper()
cfg, err := h.repo.GetNodeSSHConfig(nodeID)
if err != nil {
t.Fatalf("load ssh config: %v", err)
}
return cfg
}
func mustHandlerCount(t *testing.T, h *Handler, query string, args ...interface{}) int64 {
t.Helper()
var value int64
if err := h.repo.DB().Raw(query, args...).Row().Scan(&value); err != nil {
t.Fatalf("query %q failed: %v", query, err)
}
return value
}
@@ -0,0 +1,125 @@
package nftables
import (
"encoding/json"
"strconv"
"strings"
)
type CounterSample struct {
ForwardID int64
Direction string
Protocol string
Bytes uint64
Packets uint64
}
func ParseCounterComment(comment string) (CounterSample, bool) {
parts := strings.Split(comment, " ")
if len(parts) != 4 || parts[0] != "flvx" {
return CounterSample{}, false
}
if !strings.HasPrefix(parts[1], "forward:") {
return CounterSample{}, false
}
forwardText := strings.TrimPrefix(parts[1], "forward:")
forwardID, err := strconv.ParseInt(forwardText, 10, 64)
if err != nil || forwardID <= 0 {
return CounterSample{}, false
}
direction := parts[2]
if direction != CounterDirectionToTarget && direction != CounterDirectionFromTarget {
return CounterSample{}, false
}
protocol := parts[3]
if protocol != "tcp" && protocol != "udp" {
return CounterSample{}, false
}
return CounterSample{
ForwardID: forwardID,
Direction: direction,
Protocol: protocol,
}, true
}
func ParseCounterSamples(raw []byte) ([]CounterSample, error) {
var doc nftListTable
if err := json.Unmarshal(raw, &doc); err != nil {
return nil, err
}
samples := make([]CounterSample, 0)
for _, item := range doc.Nftables {
ruleRaw, ok := item["rule"]
if !ok {
continue
}
var rule nftCounterRule
if err := json.Unmarshal(ruleRaw, &rule); err != nil {
return nil, err
}
if rule.Table != "flvx" || rule.Chain != "forward" {
continue
}
sample, ok, err := parseCounterRule(rule)
if err != nil {
return nil, err
}
if !ok {
continue
}
samples = append(samples, sample)
}
return samples, nil
}
type nftListTable struct {
Nftables []map[string]json.RawMessage `json:"nftables"`
}
type nftCounterRule struct {
Table string `json:"table"`
Chain string `json:"chain"`
Comment string `json:"comment"`
Expr []map[string]json.RawMessage `json:"expr"`
}
type nftCounter struct {
Bytes uint64 `json:"bytes"`
Packets uint64 `json:"packets"`
}
func parseCounterRule(rule nftCounterRule) (CounterSample, bool, error) {
var (
counter nftCounter
hasCounter bool
comment = rule.Comment
)
for _, expr := range rule.Expr {
if rawCounter, ok := expr["counter"]; ok {
if err := json.Unmarshal(rawCounter, &counter); err != nil {
return CounterSample{}, false, err
}
hasCounter = true
continue
}
if rawComment, ok := expr["comment"]; ok && strings.TrimSpace(comment) == "" {
if err := json.Unmarshal(rawComment, &comment); err != nil {
return CounterSample{}, false, err
}
}
}
if !hasCounter {
return CounterSample{}, false, nil
}
sample, ok := ParseCounterComment(comment)
if !ok {
return CounterSample{}, false, nil
}
sample.Bytes = counter.Bytes
sample.Packets = counter.Packets
return sample, true, nil
}
@@ -0,0 +1,168 @@
package nftables
import "testing"
func TestParseCounterCommentAcceptsValidToTargetTCP(t *testing.T) {
sample, ok := ParseCounterComment("flvx forward:42 to-target tcp")
if !ok {
t.Fatal("expected comment to parse")
}
if sample.ForwardID != 42 ||
sample.Direction != CounterDirectionToTarget ||
sample.Protocol != "tcp" {
t.Fatalf("unexpected sample: %+v", sample)
}
}
func TestParseCounterCommentRejectsDNAT(t *testing.T) {
if sample, ok := ParseCounterComment("flvx forward:42 dnat tcp"); ok {
t.Fatalf("expected dnat comment to be rejected, got %+v", sample)
}
}
func TestParseCounterSamplesParsesForwardBillableCounters(t *testing.T) {
raw := []byte(`{
"nftables": [
{"metainfo": {"json_schema_version": 1}},
{"rule": {
"family": "inet",
"table": "flvx",
"chain": "forward",
"handle": 10,
"comment": "flvx forward:42 to-target tcp",
"expr": [
{"match": {"left": {"payload": {"protocol": "ip", "field": "daddr"}}, "op": "==", "right": "198.51.100.20"}},
{"counter": {"packets": 7, "bytes": 4096}}
]
}},
{"rule": {
"family": "inet",
"table": "flvx",
"chain": "forward",
"handle": 11,
"comment": "flvx forward:42 from-target udp",
"expr": [
{"counter": {"packets": 9, "bytes": 8192}}
]
}},
{"rule": {
"family": "inet",
"table": "flvx",
"chain": "prerouting",
"handle": 12,
"comment": "flvx forward:42 dnat tcp",
"expr": [
{"counter": {"packets": 100, "bytes": 65536}}
]
}}
]
}`)
samples, err := ParseCounterSamples(raw)
if err != nil {
t.Fatalf("ParseCounterSamples: %v", err)
}
if len(samples) != 2 {
t.Fatalf("expected 2 samples, got %d: %+v", len(samples), samples)
}
want := []CounterSample{
{ForwardID: 42, Direction: CounterDirectionToTarget, Protocol: "tcp", Bytes: 4096, Packets: 7},
{ForwardID: 42, Direction: CounterDirectionFromTarget, Protocol: "udp", Bytes: 8192, Packets: 9},
}
for i := range want {
if samples[i] != want[i] {
t.Fatalf("sample %d: expected %+v, got %+v", i, want[i], samples[i])
}
}
}
func TestParseCounterSamplesUsesRuleLevelComment(t *testing.T) {
raw := []byte(`{
"nftables": [
{"rule": {
"table": "flvx",
"chain": "forward",
"comment": "flvx forward:77 to-target udp",
"expr": [
{"counter": {"packets": 3, "bytes": 2048}}
]
}}
]
}`)
samples, err := ParseCounterSamples(raw)
if err != nil {
t.Fatalf("ParseCounterSamples: %v", err)
}
if len(samples) != 1 {
t.Fatalf("expected 1 sample, got %d: %+v", len(samples), samples)
}
want := CounterSample{
ForwardID: 77,
Direction: CounterDirectionToTarget,
Protocol: "udp",
Bytes: 2048,
Packets: 3,
}
if samples[0] != want {
t.Fatalf("expected %+v, got %+v", want, samples[0])
}
}
func TestParseCounterSamplesUsesExprLevelComment(t *testing.T) {
raw := []byte(`{
"nftables": [
{"rule": {
"table": "flvx",
"chain": "forward",
"expr": [
{"counter": {"packets": 4, "bytes": 3072}},
{"comment": "flvx forward:78 from-target tcp"}
]
}}
]
}`)
samples, err := ParseCounterSamples(raw)
if err != nil {
t.Fatalf("ParseCounterSamples: %v", err)
}
if len(samples) != 1 {
t.Fatalf("expected 1 sample, got %d: %+v", len(samples), samples)
}
want := CounterSample{
ForwardID: 78,
Direction: CounterDirectionFromTarget,
Protocol: "tcp",
Bytes: 3072,
Packets: 4,
}
if samples[0] != want {
t.Fatalf("expected %+v, got %+v", want, samples[0])
}
}
func TestParseCounterSamplesMalformedJSONReturnsError(t *testing.T) {
if _, err := ParseCounterSamples([]byte(`{"nftables": [`)); err == nil {
t.Fatal("expected malformed JSON error")
}
}
func TestParseCounterSamplesMalformedRuleJSONReturnsError(t *testing.T) {
raw := []byte(`{
"nftables": [
{"rule": {
"table": "flvx",
"chain": "forward",
"comment": "flvx forward:42 to-target tcp",
"expr": [
{"counter": {"packets": "bad", "bytes": 4096}}
]
}}
]
}`)
if _, err := ParseCounterSamples(raw); err == nil {
t.Fatal("expected malformed rule JSON error")
}
}
@@ -46,6 +46,17 @@ func (m *Manager) Clear(ctx context.Context, cfg SSHConfig) error {
return m.runner.ApplyScript(ctx, cfg, script)
}
func (m *Manager) CollectCounters(ctx context.Context, cfg SSHConfig) ([]CounterSample, error) {
if err := m.ensureInitialized(); err != nil {
return nil, err
}
raw, err := m.runner.ListTableJSON(ctx, cfg)
if err != nil {
return nil, err
}
return ParseCounterSamples(raw)
}
func (m *Manager) ensureInitialized() error {
if m == nil || m.runner == nil {
return errors.New("nftables manager not initialized")
@@ -8,9 +8,11 @@ import (
)
type fakeRunner struct {
scripts []string
err error
testErr error
scripts []string
err error
testErr error
listJSON []byte
listJSONErr error
}
func (f *fakeRunner) ApplyScript(ctx context.Context, cfg SSHConfig, script string) error {
@@ -22,12 +24,16 @@ func (f *fakeRunner) Test(ctx context.Context, cfg SSHConfig) error {
return f.testErr
}
func (f *fakeRunner) ListTableJSON(ctx context.Context, cfg SSHConfig) ([]byte, error) {
return f.listJSON, f.listJSONErr
}
func TestManagerReconcileAppliesRenderedScript(t *testing.T) {
runner := &fakeRunner{}
manager := NewManager(runner)
plan := NodePlan{
NodeID: 7,
Rules: []Rule{{ForwardID: 42, InPort: 24000, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp", "udp"}}},
Rules: []Rule{{ForwardID: 42, InPort: 24000, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp", "udp"}}},
}
result, err := manager.Reconcile(context.Background(), SSHConfig{Host: "203.0.113.10", Port: 22, Username: "root"}, plan)
@@ -37,7 +43,7 @@ func TestManagerReconcileAppliesRenderedScript(t *testing.T) {
if len(runner.scripts) != 1 {
t.Fatalf("expected 1 script, got %d", len(runner.scripts))
}
if !strings.Contains(runner.scripts[0], "flvx forward:42 tcp") {
if !strings.Contains(runner.scripts[0], `flvx forward:42 dnat tcp`) {
t.Fatalf("script missing forward comment:\n%s", runner.scripts[0])
}
if result.NodeID != 7 || result.Hashes[42] == "" {
@@ -80,6 +86,41 @@ func TestManagerTestPassesThroughRunnerError(t *testing.T) {
}
}
func TestManagerCollectCountersParsesRunnerTableJSON(t *testing.T) {
runner := &fakeRunner{listJSON: []byte(`{
"nftables": [
{"rule": {
"family": "inet",
"table": "flvx",
"chain": "forward",
"comment": "flvx forward:77 to-target tcp",
"expr": [
{"counter": {"packets": 3, "bytes": 2048}}
]
}}
]
}`)}
manager := NewManager(runner)
samples, err := manager.CollectCounters(context.Background(), SSHConfig{Host: "203.0.113.10", Port: 22, Username: "root"})
if err != nil {
t.Fatalf("CollectCounters: %v", err)
}
if len(samples) != 1 {
t.Fatalf("expected 1 sample, got %d: %+v", len(samples), samples)
}
want := CounterSample{
ForwardID: 77,
Direction: CounterDirectionToTarget,
Protocol: "tcp",
Bytes: 2048,
Packets: 3,
}
if samples[0] != want {
t.Fatalf("expected %+v, got %+v", want, samples[0])
}
}
func TestManagerMethodsRequireInitializedRunner(t *testing.T) {
cfg := SSHConfig{Host: "203.0.113.10", Port: 22, Username: "root"}
plan := NodePlan{NodeID: 7}
@@ -98,6 +139,10 @@ func TestManagerMethodsRequireInitializedRunner(t *testing.T) {
t.Fatalf("expected not initialized error from nil manager Clear, got %v", err)
}
if _, err := nilManager.CollectCounters(context.Background(), cfg); err == nil || err.Error() != expected.Error() {
t.Fatalf("expected not initialized error from nil manager CollectCounters, got %v", err)
}
manager := &Manager{}
if err := manager.Test(context.Background(), cfg); err == nil || err.Error() != expected.Error() {
t.Fatalf("expected not initialized error from Test, got %v", err)
@@ -110,4 +155,8 @@ func TestManagerMethodsRequireInitializedRunner(t *testing.T) {
if err := manager.Clear(context.Background(), cfg); err == nil || err.Error() != expected.Error() {
t.Fatalf("expected not initialized error from Clear, got %v", err)
}
if _, err := manager.CollectCounters(context.Background(), cfg); err == nil || err.Error() != expected.Error() {
t.Fatalf("expected not initialized error from CollectCounters, got %v", err)
}
}
@@ -15,14 +15,18 @@ func RenderTable(plan NodePlan) string {
b.WriteString(" chain prerouting {\n")
b.WriteString(" type nat hook prerouting priority dstnat; policy accept;\n")
for _, rule := range sortedRules(plan.Rules) {
family := nftAddressFamily(rule.TargetHost)
dnatFamily := ""
if family != "" {
dnatFamily = family + " "
}
for _, protocol := range normalizedProtocols(rule.Protocols) {
b.WriteString(fmt.Sprintf(" %s dport %d dnat %s to %s comment \"flvx forward:%d %s\"\n",
b.WriteString(fmt.Sprintf(" %s dport %d counter dnat %sto %s comment %q\n",
protocol,
rule.InPort,
dnatFamilyPrefix(rule.TargetHost),
dnatFamily,
formatDNATTarget(rule.TargetHost, rule.TargetPort),
rule.ForwardID,
protocol,
counterComment(rule.ForwardID, CounterDirectionDNAT, protocol),
))
}
}
@@ -35,11 +39,40 @@ func RenderTable(plan NodePlan) string {
b.WriteString(" }\n\n")
b.WriteString(" chain forward {\n")
b.WriteString(" type filter hook forward priority filter; policy accept;\n")
for _, rule := range sortedRules(plan.Rules) {
family := nftAddressFamily(rule.TargetHost)
if family == "" {
continue
}
targetHost := strings.Trim(strings.TrimSpace(rule.TargetHost), "[]")
for _, protocol := range normalizedProtocols(rule.Protocols) {
b.WriteString(fmt.Sprintf(" ct original proto-dst %d %s daddr %s %s dport %d counter comment %q\n",
rule.InPort,
family,
targetHost,
protocol,
rule.TargetPort,
counterComment(rule.ForwardID, CounterDirectionToTarget, protocol),
))
b.WriteString(fmt.Sprintf(" ct original proto-dst %d %s saddr %s %s sport %d counter comment %q\n",
rule.InPort,
family,
targetHost,
protocol,
rule.TargetPort,
counterComment(rule.ForwardID, CounterDirectionFromTarget, protocol),
))
}
}
b.WriteString(" }\n")
b.WriteString("}\n")
return b.String()
}
func counterComment(forwardID int64, direction, protocol string) string {
return fmt.Sprintf("flvx forward:%d %s %s", forwardID, direction, protocol)
}
func RuleHash(rule Rule) string {
protocols := normalizedProtocols(rule.Protocols)
sum := sha256.Sum256([]byte(fmt.Sprintf("%d|%d|%s|%d|%s",
@@ -100,14 +133,14 @@ func formatDNATTarget(host string, port int) string {
return fmt.Sprintf("%s:%d", trimmed, port)
}
func dnatFamilyPrefix(host string) string {
func nftAddressFamily(host string) string {
trimmed := strings.Trim(strings.TrimSpace(host), "[]")
ip := net.ParseIP(trimmed)
if ip == nil {
return ""
}
if ip.To4() != nil {
return "ip"
if ip.To4() == nil {
return "ip6"
}
return "ip6"
return "ip"
}
@@ -23,8 +23,8 @@ func TestRenderTableIncludesDNATAndMasquerade(t *testing.T) {
"table inet flvx",
"type nat hook prerouting priority dstnat; policy accept;",
"type nat hook postrouting priority srcnat; policy accept;",
"tcp dport 24000 dnat ip to 198.51.100.20:443 comment \"flvx forward:42 tcp\"",
"udp dport 24000 dnat ip to 198.51.100.20:443 comment \"flvx forward:42 udp\"",
"tcp dport 24000 counter dnat ip to 198.51.100.20:443 comment \"flvx forward:42 dnat tcp\"",
"udp dport 24000 counter dnat ip to 198.51.100.20:443 comment \"flvx forward:42 dnat udp\"",
"masquerade comment \"flvx masquerade\"",
}
for _, part := range expectedParts {
@@ -46,6 +46,109 @@ func TestRenderTableBracketsIPv6Target(t *testing.T) {
}
}
func TestRenderTableIncludesForwardAccountingCounters(t *testing.T) {
plan := NodePlan{
NodeID: 7,
Rules: []Rule{{
ForwardID: 42,
InPort: 12345,
TargetHost: "198.51.100.20",
TargetPort: 443,
Protocols: []string{"tcp", "udp"},
}},
}
got := RenderTable(plan)
wantLines := []string{
`tcp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat tcp"`,
`udp dport 12345 counter dnat ip to 198.51.100.20:443 comment "flvx forward:42 dnat udp"`,
`ct original proto-dst 12345 ip daddr 198.51.100.20 tcp dport 443 counter comment "flvx forward:42 to-target tcp"`,
`ct original proto-dst 12345 ip saddr 198.51.100.20 tcp sport 443 counter comment "flvx forward:42 from-target tcp"`,
`ct original proto-dst 12345 ip daddr 198.51.100.20 udp dport 443 counter comment "flvx forward:42 to-target udp"`,
`ct original proto-dst 12345 ip saddr 198.51.100.20 udp sport 443 counter comment "flvx forward:42 from-target udp"`,
}
for _, want := range wantLines {
if !strings.Contains(got, want) {
t.Fatalf("RenderTable() missing %q\n%s", want, got)
}
}
}
func TestRenderTableIncludesIPv6ForwardAccountingCounters(t *testing.T) {
plan := NodePlan{
NodeID: 7,
Rules: []Rule{{
ForwardID: 43,
InPort: 12346,
TargetHost: "2001:db8::20",
TargetPort: 8443,
Protocols: []string{"tcp"},
}},
}
got := RenderTable(plan)
wantLines := []string{
`tcp dport 12346 counter dnat ip6 to [2001:db8::20]:8443 comment "flvx forward:43 dnat tcp"`,
`ct original proto-dst 12346 ip6 daddr 2001:db8::20 tcp dport 8443 counter comment "flvx forward:43 to-target tcp"`,
`ct original proto-dst 12346 ip6 saddr 2001:db8::20 tcp sport 8443 counter comment "flvx forward:43 from-target tcp"`,
}
for _, want := range wantLines {
if !strings.Contains(got, want) {
t.Fatalf("RenderTable() missing %q\n%s", want, got)
}
}
}
func TestRenderTableAccountingCountersIncludeOriginalPort(t *testing.T) {
plan := NodePlan{
NodeID: 7,
Rules: []Rule{
{ForwardID: 42, InPort: 12345, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp"}},
{ForwardID: 43, InPort: 12346, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp"}},
},
}
got := RenderTable(plan)
wantLines := []string{
`ct original proto-dst 12345 ip daddr 198.51.100.20 tcp dport 443 counter comment "flvx forward:42 to-target tcp"`,
`ct original proto-dst 12346 ip daddr 198.51.100.20 tcp dport 443 counter comment "flvx forward:43 to-target tcp"`,
}
for _, want := range wantLines {
if !strings.Contains(got, want) {
t.Fatalf("RenderTable() missing %q\n%s", want, got)
}
}
}
func TestRenderTablePreservesHostnameDNATAndSkipsAccountingCounters(t *testing.T) {
plan := NodePlan{
NodeID: 7,
Rules: []Rule{{
ForwardID: 44,
InPort: 12347,
TargetHost: "example.com",
TargetPort: 9443,
Protocols: []string{"tcp"},
}},
}
got := RenderTable(plan)
want := `tcp dport 12347 counter dnat to example.com:9443 comment "flvx forward:44 dnat tcp"`
if !strings.Contains(got, want) {
t.Fatalf("RenderTable() missing %q\n%s", want, got)
}
unwantedLines := []string{
`dnat ip to example.com`,
`ip daddr example.com`,
`ip saddr example.com`,
}
for _, unwanted := range unwantedLines {
if strings.Contains(got, unwanted) {
t.Fatalf("RenderTable() unexpectedly contains %q\n%s", unwanted, got)
}
}
}
func TestRuleHashIsStable(t *testing.T) {
rule := Rule{ForwardID: 42, InPort: 24000, TargetHost: "198.51.100.20", TargetPort: 443, Protocols: []string{"tcp", "udp"}}
if RuleHash(rule) != RuleHash(rule) {
+20 -8
View File
@@ -14,6 +14,7 @@ import (
type Runner interface {
ApplyScript(ctx context.Context, cfg SSHConfig, script string) error
Test(ctx context.Context, cfg SSHConfig) error
ListTableJSON(ctx context.Context, cfg SSHConfig) ([]byte, error)
}
type SSHRunner struct {
@@ -44,7 +45,16 @@ func (r *SSHRunner) ApplyScript(ctx context.Context, cfg SSHConfig, script strin
return r.run(ctx, cfg, command)
}
func (r *SSHRunner) ListTableJSON(ctx context.Context, cfg SSHConfig) ([]byte, error) {
return r.runOutput(ctx, cfg, nftBinary(cfg)+" -j list table inet flvx")
}
func (r *SSHRunner) run(ctx context.Context, cfg SSHConfig, command string) error {
_, err := r.runOutput(ctx, cfg, command)
return err
}
func (r *SSHRunner) runOutput(ctx context.Context, cfg SSHConfig, command string) ([]byte, error) {
timeout := r.Timeout
if timeout <= 0 {
timeout = 15 * time.Second
@@ -54,31 +64,33 @@ func (r *SSHRunner) run(ctx context.Context, cfg SSHConfig, command string) erro
clientConfig, err := buildSSHClientConfig(cfg)
if err != nil {
return err
return nil, err
}
addr := net.JoinHostPort(strings.TrimSpace(cfg.Host), fmt.Sprintf("%d", normalizedSSHPort(cfg.Port)))
dialer := net.Dialer{Timeout: timeout}
conn, err := dialer.DialContext(runCtx, "tcp", addr)
if err != nil {
return fmt.Errorf("SSH 连接失败: %w", err)
return nil, fmt.Errorf("SSH 连接失败: %w", err)
}
defer conn.Close()
sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, clientConfig)
if err != nil {
return fmt.Errorf("SSH 认证失败: %w", err)
return nil, fmt.Errorf("SSH 认证失败: %w", err)
}
client := ssh.NewClient(sshConn, chans, reqs)
defer client.Close()
session, err := client.NewSession()
if err != nil {
return fmt.Errorf("SSH 会话创建失败: %w", err)
return nil, fmt.Errorf("SSH 会话创建失败: %w", err)
}
defer session.Close()
var stdout bytes.Buffer
var stderr bytes.Buffer
session.Stdout = &stdout
session.Stderr = &stderr
done := make(chan error, 1)
@@ -89,16 +101,16 @@ func (r *SSHRunner) run(ctx context.Context, cfg SSHConfig, command string) erro
select {
case <-runCtx.Done():
_ = session.Close()
return fmt.Errorf("SSH 命令超时: %w", runCtx.Err())
return nil, fmt.Errorf("SSH 命令超时: %w", runCtx.Err())
case err := <-done:
if err != nil {
message := strings.TrimSpace(stderr.String())
if message != "" {
return fmt.Errorf("远程执行失败: %s: %w", message, err)
return nil, fmt.Errorf("远程执行失败: %s: %w", message, err)
}
return fmt.Errorf("远程执行失败: %w", err)
return nil, fmt.Errorf("远程执行失败: %w", err)
}
return nil
return stdout.Bytes(), nil
}
}
@@ -7,6 +7,10 @@ const (
StatusPending = "pending"
StatusApplied = "applied"
StatusError = "error"
CounterDirectionDNAT = "dnat"
CounterDirectionToTarget = "to-target"
CounterDirectionFromTarget = "from-target"
)
type Target struct {
+16
View File
@@ -131,6 +131,22 @@ type NftRuleBinding struct {
func (NftRuleBinding) TableName() string { return "nft_rule_binding" }
type NftCounterState struct {
ID int64 `gorm:"primaryKey;autoIncrement"`
NodeID int64 `gorm:"column:node_id;not null;uniqueIndex:idx_nft_counter_state_key;index"`
ForwardID int64 `gorm:"column:forward_id;not null;uniqueIndex:idx_nft_counter_state_key;index"`
Protocol string `gorm:"type:varchar(10);not null;uniqueIndex:idx_nft_counter_state_key"`
Direction string `gorm:"type:varchar(20);not null;uniqueIndex:idx_nft_counter_state_key"`
RuleHash string `gorm:"column:rule_hash;type:varchar(128);not null;default:''"`
Bytes int64 `gorm:"not null;default:0"`
Packets int64 `gorm:"not null;default:0"`
CollectedTime int64 `gorm:"column:collected_time;not null;default:0"`
CreatedTime int64 `gorm:"column:created_time;not null"`
UpdatedTime int64 `gorm:"column:updated_time;not null"`
}
func (NftCounterState) TableName() string { return "nft_counter_state" }
type SpeedLimit struct {
ID int64 `gorm:"primaryKey;autoIncrement"`
Name string `gorm:"type:varchar(100);not null"`
+39 -27
View File
@@ -104,6 +104,19 @@ func (r *Repository) ApplyFlowUploadDeltasBatch(deltas []FlowUploadCounterDelta)
return nil
}
return r.db.Transaction(func(tx *gorm.DB) error {
return applyFlowUploadDeltasTx(tx, deltas)
})
}
func applyFlowUploadDeltasTx(tx *gorm.DB, deltas []FlowUploadCounterDelta) error {
if tx == nil {
return errors.New("database unavailable")
}
if len(deltas) == 0 {
return nil
}
forwardTotals := make(map[int64][2]int64, len(deltas))
userTotals := make(map[int64][2]int64, len(deltas))
userTunnelTotals := make(map[int64][2]int64, len(deltas))
@@ -128,36 +141,34 @@ func (r *Repository) ApplyFlowUploadDeltasBatch(deltas []FlowUploadCounterDelta)
}
}
return r.db.Transaction(func(tx *gorm.DB) error {
for _, forwardID := range sortedFlowUploadTargetIDs(forwardTotals) {
total := forwardTotals[forwardID]
if err := tx.Model(&model.Forward{}).Where("id = ?", forwardID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
for _, forwardID := range sortedFlowUploadTargetIDs(forwardTotals) {
total := forwardTotals[forwardID]
if err := tx.Model(&model.Forward{}).Where("id = ?", forwardID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
for _, userID := range sortedFlowUploadTargetIDs(userTotals) {
total := userTotals[userID]
if err := tx.Model(&model.User{}).Where("id = ?", userID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
}
for _, userID := range sortedFlowUploadTargetIDs(userTotals) {
total := userTotals[userID]
if err := tx.Model(&model.User{}).Where("id = ?", userID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
for _, userTunnelID := range sortedFlowUploadTargetIDs(userTunnelTotals) {
total := userTunnelTotals[userTunnelID]
if err := tx.Model(&model.UserTunnel{}).Where("id = ?", userTunnelID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
}
for _, userTunnelID := range sortedFlowUploadTargetIDs(userTunnelTotals) {
total := userTunnelTotals[userTunnelID]
if err := tx.Model(&model.UserTunnel{}).Where("id = ?", userTunnelID).UpdateColumns(map[string]interface{}{
"in_flow": gorm.Expr("in_flow + ?", total[0]),
"out_flow": gorm.Expr("out_flow + ?", total[1]),
}).Error; err != nil {
return err
}
return nil
})
}
return nil
}
// ─── Open / Close ────────────────────────────────────────────────────
@@ -282,6 +293,7 @@ func autoMigrateAll(db *gorm.DB) error {
&model.Node{},
&model.NodeSSHConfig{},
&model.NftRuleBinding{},
&model.NftCounterState{},
&model.SpeedLimit{},
&model.StatisticsFlow{},
&model.Tunnel{},
@@ -11,6 +11,8 @@ import (
type FlowUploadForwardMeta struct {
ForwardID int64
UserID int64
UserTunnelID int64
TunnelID int64
TrafficRatio float64
TunnelFlow int64
@@ -59,6 +61,8 @@ func (r *Repository) GetFlowUploadForwardMetas(forwardIDs []int64) (map[int64]Fl
type row struct {
ForwardID int64 `gorm:"column:forward_id"`
UserID int64 `gorm:"column:user_id"`
UserTunnelID int64 `gorm:"column:user_tunnel_id"`
TunnelID int64 `gorm:"column:tunnel_id"`
TrafficRatio float64 `gorm:"column:traffic_ratio"`
TunnelFlow int64 `gorm:"column:tunnel_flow"`
@@ -68,8 +72,9 @@ func (r *Repository) GetFlowUploadForwardMetas(forwardIDs []int64) (map[int64]Fl
for _, chunk := range chunkFlowUploadForwardIDs(ids) {
var rows []row
err := r.db.Table("forward AS f").
Select("f.id AS forward_id, f.tunnel_id AS tunnel_id, t.traffic_ratio AS traffic_ratio, t.flow AS tunnel_flow").
Select("f.id AS forward_id, f.user_id AS user_id, COALESCE(ut.id, 0) AS user_tunnel_id, f.tunnel_id AS tunnel_id, t.traffic_ratio AS traffic_ratio, t.flow AS tunnel_flow").
Joins("LEFT JOIN tunnel t ON t.id = f.tunnel_id").
Joins("LEFT JOIN user_tunnel ut ON ut.user_id = f.user_id AND ut.tunnel_id = f.tunnel_id").
Where("f.id IN ?", chunk).
Scan(&rows).Error
if err != nil {
@@ -84,6 +89,8 @@ func (r *Repository) GetFlowUploadForwardMetas(forwardIDs []int64) (map[int64]Fl
}
out[row.ForwardID] = FlowUploadForwardMeta{
ForwardID: row.ForwardID,
UserID: row.UserID,
UserTunnelID: row.UserTunnelID,
TunnelID: row.TunnelID,
TrafficRatio: row.TrafficRatio,
TunnelFlow: row.TunnelFlow,
@@ -64,7 +64,7 @@ func TestGetFlowUploadForwardMetasAndApplyFlowUploadDeltasBatch(t *testing.T) {
if err != nil {
t.Fatalf("get metas: %v", err)
}
if metas[20].TunnelID != 1 || metas[20].TrafficRatio != 2 || metas[20].TunnelFlow != 3 {
if metas[20].UserID != 2 || metas[20].UserTunnelID != 10 || metas[20].TunnelID != 1 || metas[20].TrafficRatio != 2 || metas[20].TunnelFlow != 3 {
t.Fatalf("unexpected meta for forward 20: %#v", metas[20])
}
if _, ok := metas[99]; ok {
@@ -86,6 +86,99 @@ func TestGetFlowUploadForwardMetasAndApplyFlowUploadDeltasBatch(t *testing.T) {
}
}
func TestApplyNftTrafficAccountingAppliesFlowQuotaAndStates(t *testing.T) {
r, err := Open(filepath.Join(t.TempDir(), "nft-accounting.db"))
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now()
nowMs := now.UnixMilli()
seedFlowBatchRows(t, r, nowMs)
quotaViews, err := r.ApplyNftTrafficAccounting(
[]FlowUploadCounterDelta{{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 480, OutFlow: 660}},
map[int64]int64{2: 1140},
[]NftCounterStateInput{{
NodeID: 11,
ForwardID: 20,
Protocol: "tcp",
Direction: "to-target",
RuleHash: "hash-a",
Bytes: 1400,
Packets: 14,
CollectedTime: nowMs,
}},
now,
)
if err != nil {
t.Fatalf("ApplyNftTrafficAccounting: %v", err)
}
if got := mustFlowBatchCount(t, r, `SELECT in_flow FROM forward WHERE id = 20`); got != 480 {
t.Fatalf("expected forward in_flow=480, got %d", got)
}
if got := mustFlowBatchCount(t, r, `SELECT out_flow FROM user WHERE id = 2`); got != 660 {
t.Fatalf("expected user out_flow=660, got %d", got)
}
if got := mustFlowBatchCount(t, r, `SELECT in_flow FROM user_tunnel WHERE id = 10`); got != 480 {
t.Fatalf("expected user_tunnel in_flow=480, got %d", got)
}
if quotaViews[2] == nil || quotaViews[2].DailyUsedBytes != 1140 || quotaViews[2].MonthlyUsedBytes != 1140 {
t.Fatalf("unexpected quota view: %#v", quotaViews[2])
}
states, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode: %v", err)
}
if len(states) != 1 || states[0].ForwardID != 20 || states[0].Bytes != 1400 {
t.Fatalf("unexpected nft counter state: %+v", states)
}
}
func TestApplyNftTrafficAccountingRollsBackFlowAndQuotaWhenStateWriteFails(t *testing.T) {
r, err := Open(filepath.Join(t.TempDir(), "nft-accounting-rollback.db"))
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now()
nowMs := now.UnixMilli()
seedFlowBatchRows(t, r, nowMs)
if err := r.DB().Exec(`DROP TABLE nft_counter_state`).Error; err != nil {
t.Fatalf("drop nft_counter_state: %v", err)
}
_, err = r.ApplyNftTrafficAccounting(
[]FlowUploadCounterDelta{{ForwardID: 20, UserID: 2, UserTunnelID: 10, InFlow: 480, OutFlow: 660}},
map[int64]int64{2: 1140},
[]NftCounterStateInput{{
NodeID: 11,
ForwardID: 20,
Protocol: "tcp",
Direction: "to-target",
RuleHash: "hash-a",
Bytes: 1400,
Packets: 14,
CollectedTime: nowMs,
}},
now,
)
if err == nil {
t.Fatalf("expected ApplyNftTrafficAccounting to fail")
}
if got := mustFlowBatchCount(t, r, `SELECT in_flow FROM forward WHERE id = 20`); got != 0 {
t.Fatalf("expected forward flow rollback, got %d", got)
}
if got := mustFlowBatchCount(t, r, `SELECT out_flow FROM user WHERE id = 2`); got != 0 {
t.Fatalf("expected user flow rollback, got %d", got)
}
if got := mustFlowBatchCount(t, r, `SELECT COALESCE((SELECT daily_used_bytes FROM user_quota WHERE user_id = 2), 0)`); got != 0 {
t.Fatalf("expected quota rollback, got %d", got)
}
}
func TestGetFlowUploadForwardMetasKeepsForwardsWhenTunnelRowMissing(t *testing.T) {
r, err := Open(filepath.Join(t.TempDir(), "flow-batch-missing-tunnel.db"))
if err != nil {
@@ -106,7 +199,7 @@ func TestGetFlowUploadForwardMetasKeepsForwardsWhenTunnelRowMissing(t *testing.T
if !ok {
t.Fatalf("expected metadata for forward with missing tunnel row")
}
if meta.ForwardID != 25 || meta.TunnelID != 99 || meta.TrafficRatio != 1 || meta.TunnelFlow != 1 {
if meta.ForwardID != 25 || meta.UserID != 2 || meta.UserTunnelID != 0 || meta.TunnelID != 99 || meta.TrafficRatio != 1 || meta.TunnelFlow != 1 {
t.Fatalf("unexpected fallback meta: %#v", meta)
}
}
@@ -167,3 +260,19 @@ func mustFlowBatchCount(t *testing.T, r *Repository, query string, args ...inter
}
return value
}
func seedFlowBatchRows(t *testing.T, r *Repository, now int64) {
t.Helper()
if err := r.DB().Exec(`INSERT INTO user(id, user, pwd, role_id, exp_time, flow, in_flow, out_flow, flow_reset_time, num, created_time, updated_time, status) VALUES(2, 'u2', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1)`, now, now).Error; err != nil {
t.Fatalf("insert user: %v", err)
}
if err := r.DB().Exec(`INSERT INTO tunnel(id, name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) VALUES(1, 't1', 2.0, 1, 'tls', 3, ?, ?, 1, NULL, 0)`, now, now).Error; err != nil {
t.Fatalf("insert tunnel: %v", err)
}
if err := r.DB().Exec(`INSERT INTO user_tunnel(id, user_id, tunnel_id, speed_id, num, flow, in_flow, out_flow, flow_reset_time, exp_time, status) VALUES(10, 2, 1, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1)`).Error; err != nil {
t.Fatalf("insert user_tunnel: %v", err)
}
if err := r.DB().Exec(`INSERT INTO forward(id, user_id, user_name, name, tunnel_id, remote_addr, strategy, in_flow, out_flow, created_time, updated_time, status, inx) VALUES(20, 2, 'u2', 'f20', 1, '1.1.1.1:80', 'fifo', 0, 0, ?, ?, 1, 0)`, now, now).Error; err != nil {
t.Fatalf("insert forward: %v", err)
}
}
@@ -769,6 +769,9 @@ func (r *Repository) DeleteForwardCascade(forwardID int64) error {
return errors.New("repository not initialized")
}
return r.db.Transaction(func(tx *gorm.DB) error {
if err := tx.Where("forward_id = ?", forwardID).Delete(&model.NftCounterState{}).Error; err != nil {
return err
}
if err := tx.Where("forward_id = ?", forwardID).Delete(&model.ForwardPort{}).Error; err != nil {
return err
}
@@ -0,0 +1,205 @@
package repo
import (
"errors"
"math"
"strings"
"time"
"go-backend/internal/store/model"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
const (
nftCounterProtocolTCP = "tcp"
nftCounterProtocolUDP = "udp"
nftCounterDirectionToTarget = "to-target"
nftCounterDirectionFromTarget = "from-target"
)
type NftCounterStateInput struct {
NodeID int64
ForwardID int64
Protocol string
Direction string
RuleHash string
Bytes uint64
Packets uint64
CollectedTime int64
}
type NftablesCollectionNode struct {
NodeID int64
Config model.NodeSSHConfig
}
func (r *Repository) ListNftablesNodesForCollection() ([]NftablesCollectionNode, error) {
if r == nil || r.db == nil {
return nil, errors.New("repository not initialized")
}
type collectionRow struct {
NodeID int64 `gorm:"column:node_id"`
ConfigID int64 `gorm:"column:config_id"`
Host string `gorm:"column:host"`
Port int `gorm:"column:port"`
Username string `gorm:"column:username"`
AuthType string `gorm:"column:auth_type"`
Password string `gorm:"column:password"`
PrivateKey string `gorm:"column:private_key"`
Passphrase string `gorm:"column:passphrase"`
SudoMode string `gorm:"column:sudo_mode"`
CreatedTime int64 `gorm:"column:created_time"`
UpdatedTime int64 `gorm:"column:updated_time"`
}
var rows []collectionRow
if err := r.db.Table("node").
Select("node.id AS node_id, node_ssh_config.id AS config_id, node_ssh_config.host, node_ssh_config.port, node_ssh_config.username, node_ssh_config.auth_type, node_ssh_config.password, node_ssh_config.private_key, node_ssh_config.passphrase, node_ssh_config.sudo_mode, node_ssh_config.created_time, node_ssh_config.updated_time").
Joins("JOIN node_ssh_config ON node_ssh_config.node_id = node.id").
Where("node.status = ? AND LOWER(TRIM(node.forward_mode)) = ?", 1, "nftables").
Order("node.id ASC").
Scan(&rows).Error; err != nil {
return nil, err
}
nodes := make([]NftablesCollectionNode, 0, len(rows))
for _, row := range rows {
nodes = append(nodes, NftablesCollectionNode{
NodeID: row.NodeID,
Config: model.NodeSSHConfig{
ID: row.ConfigID,
NodeID: row.NodeID,
Host: row.Host,
Port: row.Port,
Username: row.Username,
AuthType: row.AuthType,
Password: nullStringFromInterface(row.Password),
PrivateKey: nullStringFromInterface(row.PrivateKey),
Passphrase: nullStringFromInterface(row.Passphrase),
SudoMode: row.SudoMode,
CreatedTime: row.CreatedTime,
UpdatedTime: row.UpdatedTime,
},
})
}
return nodes, nil
}
func (r *Repository) GetNftCounterStatesByNode(nodeID int64) ([]model.NftCounterState, error) {
if r == nil || r.db == nil {
return nil, errors.New("repository not initialized")
}
var rows []model.NftCounterState
err := r.db.Where("node_id = ?", nodeID).
Order("forward_id ASC, protocol ASC, direction ASC").
Find(&rows).Error
return rows, err
}
func (r *Repository) UpsertNftCounterStates(inputs []NftCounterStateInput, now int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
if len(inputs) == 0 {
return nil
}
return r.db.Transaction(func(tx *gorm.DB) error {
return upsertNftCounterStatesTx(tx, inputs, now)
})
}
func (r *Repository) ApplyNftTrafficAccounting(deltas []FlowUploadCounterDelta, quotaUsage map[int64]int64, states []NftCounterStateInput, now time.Time) (map[int64]*model.UserQuotaView, error) {
if r == nil || r.db == nil {
return nil, errors.New("repository not initialized")
}
quotaViews := map[int64]*model.UserQuotaView{}
err := r.db.Transaction(func(tx *gorm.DB) error {
if err := applyFlowUploadDeltasTx(tx, deltas); err != nil {
return err
}
var err error
quotaViews, err = r.addUserQuotaUsageBatchTx(tx, quotaUsage, now)
if err != nil {
return err
}
return upsertNftCounterStatesTx(tx, states, now.UnixMilli())
})
if err != nil {
return nil, err
}
return quotaViews, nil
}
func upsertNftCounterStatesTx(tx *gorm.DB, inputs []NftCounterStateInput, now int64) error {
if tx == nil {
return errors.New("database unavailable")
}
for _, input := range inputs {
row, ok := nftCounterStateFromInput(input, now)
if !ok {
continue
}
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{
{Name: "node_id"},
{Name: "forward_id"},
{Name: "protocol"},
{Name: "direction"},
},
DoUpdates: clause.Assignments(map[string]interface{}{
"rule_hash": row.RuleHash,
"bytes": row.Bytes,
"packets": row.Packets,
"collected_time": row.CollectedTime,
"updated_time": row.UpdatedTime,
}),
}).Create(&row).Error; err != nil {
return err
}
}
return nil
}
func (r *Repository) DeleteNftCounterStatesByForward(forwardID int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
return r.db.Where("forward_id = ?", forwardID).Delete(&model.NftCounterState{}).Error
}
func nftCounterStateFromInput(input NftCounterStateInput, now int64) (model.NftCounterState, bool) {
protocol := strings.ToLower(strings.TrimSpace(input.Protocol))
direction := strings.ToLower(strings.TrimSpace(input.Direction))
if input.NodeID <= 0 || input.ForwardID <= 0 || !isValidNftCounterProtocol(protocol) || !isValidNftCounterDirection(direction) {
return model.NftCounterState{}, false
}
if input.Bytes > uint64(math.MaxInt64) || input.Packets > uint64(math.MaxInt64) {
return model.NftCounterState{}, false
}
return model.NftCounterState{
NodeID: input.NodeID,
ForwardID: input.ForwardID,
Protocol: protocol,
Direction: direction,
RuleHash: strings.TrimSpace(input.RuleHash),
Bytes: int64(input.Bytes),
Packets: int64(input.Packets),
CollectedTime: input.CollectedTime,
CreatedTime: now,
UpdatedTime: now,
}, true
}
func isValidNftCounterProtocol(protocol string) bool {
return protocol == nftCounterProtocolTCP || protocol == nftCounterProtocolUDP
}
func isValidNftCounterDirection(direction string) bool {
return direction == nftCounterDirectionToTarget || direction == nftCounterDirectionFromTarget
}
@@ -0,0 +1,269 @@
package repo
import (
"math"
"path/filepath"
"testing"
"time"
"go-backend/internal/store/model"
)
func TestNftCounterStateUpsertUpdatesExistingKey(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
first := []NftCounterStateInput{
{
NodeID: 11,
ForwardID: 42,
Protocol: "tcp",
Direction: "to-target",
RuleHash: "hash-a",
Bytes: 100,
Packets: 10,
CollectedTime: 1000,
},
{
NodeID: 0,
ForwardID: 42,
Protocol: "tcp",
Direction: "to-target",
Bytes: 999,
},
}
if err := r.UpsertNftCounterStates(first, 2000); err != nil {
t.Fatalf("first UpsertNftCounterStates: %v", err)
}
second := []NftCounterStateInput{
{
NodeID: 11,
ForwardID: 42,
Protocol: "tcp",
Direction: "to-target",
RuleHash: "hash-b",
Bytes: 250,
Packets: 25,
CollectedTime: 3000,
},
}
if err := r.UpsertNftCounterStates(second, 4000); err != nil {
t.Fatalf("second UpsertNftCounterStates: %v", err)
}
rows, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode: %v", err)
}
if len(rows) != 1 {
t.Fatalf("expected one counter state row, got %d: %+v", len(rows), rows)
}
got := rows[0]
if got.ForwardID != 42 || got.Protocol != "tcp" || got.Direction != "to-target" {
t.Fatalf("unexpected counter state key: %+v", got)
}
if got.RuleHash != "hash-b" || got.Bytes != 250 || got.Packets != 25 || got.CollectedTime != 3000 {
t.Fatalf("counter state was not updated: %+v", got)
}
if got.CreatedTime != 2000 || got.UpdatedTime != 4000 {
t.Fatalf("unexpected timestamps after upsert: %+v", got)
}
}
func TestNftCounterStateDeleteByForwardRemovesOnlyMatchingRows(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
inputs := []NftCounterStateInput{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: "to-target", RuleHash: "a", Bytes: 100, Packets: 10, CollectedTime: 1000},
{NodeID: 11, ForwardID: 43, Protocol: "udp", Direction: "from-target", RuleHash: "b", Bytes: 200, Packets: 20, CollectedTime: 1000},
{NodeID: 12, ForwardID: 42, Protocol: "tcp", Direction: "to-target", RuleHash: "c", Bytes: 300, Packets: 30, CollectedTime: 1000},
}
if err := r.UpsertNftCounterStates(inputs, 2000); err != nil {
t.Fatalf("UpsertNftCounterStates: %v", err)
}
if err := r.DeleteNftCounterStatesByForward(42); err != nil {
t.Fatalf("DeleteNftCounterStatesByForward: %v", err)
}
node11, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode(11): %v", err)
}
if len(node11) != 1 || node11[0].ForwardID != 43 {
t.Fatalf("expected only forward 43 for node 11, got %+v", node11)
}
node12, err := r.GetNftCounterStatesByNode(12)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode(12): %v", err)
}
if len(node12) != 0 {
t.Fatalf("expected forward 42 state removed from node 12, got %+v", node12)
}
}
func TestDeleteForwardCascadeRemovesNftCounterStateOnlyForDeletedForward(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now().UnixMilli()
forwards := []model.Forward{
{ID: 42, UserID: 1, UserName: "admin", Name: "forward-a", TunnelID: 10, RemoteAddr: "203.0.113.1:80", Strategy: "fifo", CreatedTime: now, UpdatedTime: now, Status: 1},
{ID: 43, UserID: 1, UserName: "admin", Name: "forward-b", TunnelID: 10, RemoteAddr: "203.0.113.2:80", Strategy: "fifo", CreatedTime: now, UpdatedTime: now, Status: 1},
}
if err := r.DB().Create(&forwards).Error; err != nil {
t.Fatalf("seed forwards: %v", err)
}
if err := r.UpsertNftCounterStates([]NftCounterStateInput{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: "to-target", RuleHash: "a", Bytes: 100, Packets: 10, CollectedTime: now},
{NodeID: 11, ForwardID: 43, Protocol: "udp", Direction: "from-target", RuleHash: "b", Bytes: 200, Packets: 20, CollectedTime: now},
}, now); err != nil {
t.Fatalf("UpsertNftCounterStates: %v", err)
}
if err := r.DeleteForwardCascade(42); err != nil {
t.Fatalf("DeleteForwardCascade: %v", err)
}
rows, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode: %v", err)
}
if len(rows) != 1 || rows[0].ForwardID != 43 {
t.Fatalf("expected only forward 43 counter state to remain, got %+v", rows)
}
var deletedForwardCount int64
if err := r.DB().Model(&model.Forward{}).Where("id = ?", int64(42)).Count(&deletedForwardCount).Error; err != nil {
t.Fatalf("count deleted forward: %v", err)
}
if deletedForwardCount != 0 {
t.Fatalf("expected forward 42 deleted, count=%d", deletedForwardCount)
}
}
func TestNftCounterStateUpsertSkipsInvalidProtocolAndDirection(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
inputs := []NftCounterStateInput{
{NodeID: 11, ForwardID: 42, Protocol: "icmp", Direction: "to-target", RuleHash: "bad-protocol", Bytes: 100, Packets: 10, CollectedTime: 1000},
{NodeID: 11, ForwardID: 43, Protocol: "tcp", Direction: "sideways", RuleHash: "bad-direction", Bytes: 200, Packets: 20, CollectedTime: 1000},
{NodeID: 11, ForwardID: 44, Protocol: " UDP ", Direction: " FROM-TARGET ", RuleHash: "valid", Bytes: 300, Packets: 30, CollectedTime: 1000},
}
if err := r.UpsertNftCounterStates(inputs, 2000); err != nil {
t.Fatalf("UpsertNftCounterStates: %v", err)
}
rows, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode: %v", err)
}
if len(rows) != 1 {
t.Fatalf("expected only the valid counter state row, got %d: %+v", len(rows), rows)
}
if rows[0].ForwardID != 44 || rows[0].Protocol != "udp" || rows[0].Direction != "from-target" {
t.Fatalf("unexpected valid counter state row: %+v", rows[0])
}
}
func TestNftCounterStateUpsertSkipsCountersAboveInt64(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
inputs := []NftCounterStateInput{
{NodeID: 11, ForwardID: 42, Protocol: "tcp", Direction: "to-target", RuleHash: "too-large", Bytes: uint64(math.MaxInt64) + 1, Packets: 10, CollectedTime: 1000},
{NodeID: 11, ForwardID: 43, Protocol: "udp", Direction: "from-target", RuleHash: "valid", Bytes: 300, Packets: 30, CollectedTime: 1000},
}
if err := r.UpsertNftCounterStates(inputs, 2000); err != nil {
t.Fatalf("UpsertNftCounterStates: %v", err)
}
rows, err := r.GetNftCounterStatesByNode(11)
if err != nil {
t.Fatalf("GetNftCounterStatesByNode: %v", err)
}
if len(rows) != 1 {
t.Fatalf("expected only the valid counter state row, got %d: %+v", len(rows), rows)
}
if rows[0].ForwardID != 43 || rows[0].Bytes != 300 || rows[0].Packets != 30 {
t.Fatalf("unexpected valid counter state row: %+v", rows[0])
}
}
func TestListNftablesNodesForCollectionReturnsActiveNftablesWithSSHOrdered(t *testing.T) {
r, err := Open(filepath.Join(t.TempDir(), "nft-collection.db"))
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now().UnixMilli()
seedCollectionNode(t, r, 1, "agent", 1, now)
seedCollectionNode(t, r, 2, " nftables ", 1, now)
seedCollectionNode(t, r, 3, "NFTABLES", 0, now)
seedCollectionNode(t, r, 4, "nftables", 1, now)
seedCollectionNode(t, r, 5, "nftables", 1, now)
if err := r.UpsertNodeSSHConfig(4, NftSSHConfigInput{
Host: "203.0.113.4",
Port: 2222,
Username: "root",
AuthType: "password",
Password: "secret-4",
SudoMode: "none",
}, now); err != nil {
t.Fatalf("upsert ssh config 4: %v", err)
}
if err := r.UpsertNodeSSHConfig(2, NftSSHConfigInput{
Host: "203.0.113.2",
Port: 22,
Username: "admin",
AuthType: "private_key",
SudoMode: "sudo",
}, now); err != nil {
t.Fatalf("upsert ssh config 2: %v", err)
}
nodes, err := r.ListNftablesNodesForCollection()
if err != nil {
t.Fatalf("ListNftablesNodesForCollection: %v", err)
}
if len(nodes) != 2 {
t.Fatalf("expected 2 collection nodes, got %d: %+v", len(nodes), nodes)
}
if nodes[0].NodeID != 2 || nodes[1].NodeID != 4 {
t.Fatalf("expected nodes ordered by id [2 4], got [%d %d]", nodes[0].NodeID, nodes[1].NodeID)
}
if nodes[0].Config.NodeID != 2 || nodes[0].Config.Host != "203.0.113.2" || nodes[0].Config.Username != "admin" {
t.Fatalf("unexpected first config: %+v", nodes[0].Config)
}
if nodes[1].Config.NodeID != 4 || nodes[1].Config.Port != 2222 || nodes[1].Config.Password.String != "secret-4" {
t.Fatalf("unexpected second config: %+v", nodes[1].Config)
}
}
func seedCollectionNode(t *testing.T, r *Repository, id int64, forwardMode string, status int, now int64) {
t.Helper()
if err := r.DB().Exec(`
INSERT INTO node(id, name, secret, server_ip, port, created_time, updated_time, status, tcp_listen_addr, udp_listen_addr, inx, forward_mode)
VALUES(?, ?, 'secret', ?, '1000-2000', ?, ?, ?, '[::]', '[::]', 0, ?)
`, id, "node", "198.51.100.1", now, now, status, forwardMode).Error; err != nil {
t.Fatalf("insert node %d: %v", id, err)
}
}
@@ -264,39 +264,11 @@ func (r *Repository) AddUserQuotaUsageBatch(usages map[int64]int64, now time.Tim
return map[int64]*model.UserQuotaView{}, nil
}
result := make(map[int64]*model.UserQuotaView, len(usages))
var result map[int64]*model.UserQuotaView
err := r.db.Transaction(func(tx *gorm.DB) error {
userIDs := make([]int64, 0, len(usages))
for userID := range usages {
if userID > 0 {
userIDs = append(userIDs, userID)
}
}
sort.Slice(userIDs, func(i, j int) bool { return userIDs[i] < userIDs[j] })
for _, userID := range userIDs {
q, err := r.loadOrCreateUserQuotaTx(tx, userID, now)
if err != nil {
return err
}
applyUserQuotaWindowRoll(q, now)
if usages[userID] > 0 {
q.DailyUsedBytes += usages[userID]
q.MonthlyUsedBytes += usages[userID]
}
q.UpdatedTime = now.UnixMilli()
if err := tx.Model(&model.UserQuota{}).Where("user_id = ?", userID).Updates(map[string]interface{}{
"daily_used_bytes": q.DailyUsedBytes,
"monthly_used_bytes": q.MonthlyUsedBytes,
"day_key": q.DayKey,
"month_key": q.MonthKey,
"updated_time": q.UpdatedTime,
}).Error; err != nil {
return err
}
result[userID] = normalizeUserQuotaView(cloneUserQuotaView(*q), now)
}
return nil
var err error
result, err = r.addUserQuotaUsageBatchTx(tx, usages, now)
return err
})
if err != nil {
return nil, err
@@ -304,6 +276,48 @@ func (r *Repository) AddUserQuotaUsageBatch(usages map[int64]int64, now time.Tim
return result, nil
}
func (r *Repository) addUserQuotaUsageBatchTx(tx *gorm.DB, usages map[int64]int64, now time.Time) (map[int64]*model.UserQuotaView, error) {
if tx == nil {
return nil, errors.New("database unavailable")
}
if len(usages) == 0 {
return map[int64]*model.UserQuotaView{}, nil
}
result := make(map[int64]*model.UserQuotaView, len(usages))
userIDs := make([]int64, 0, len(usages))
for userID := range usages {
if userID > 0 {
userIDs = append(userIDs, userID)
}
}
sort.Slice(userIDs, func(i, j int) bool { return userIDs[i] < userIDs[j] })
for _, userID := range userIDs {
q, err := r.loadOrCreateUserQuotaTx(tx, userID, now)
if err != nil {
return nil, err
}
applyUserQuotaWindowRoll(q, now)
if usages[userID] > 0 {
q.DailyUsedBytes += usages[userID]
q.MonthlyUsedBytes += usages[userID]
}
q.UpdatedTime = now.UnixMilli()
if err := tx.Model(&model.UserQuota{}).Where("user_id = ?", userID).Updates(map[string]interface{}{
"daily_used_bytes": q.DailyUsedBytes,
"monthly_used_bytes": q.MonthlyUsedBytes,
"day_key": q.DayKey,
"month_key": q.MonthKey,
"updated_time": q.UpdatedTime,
}).Error; err != nil {
return nil, err
}
result[userID] = normalizeUserQuotaView(cloneUserQuotaView(*q), now)
}
return result, nil
}
func (r *Repository) MarkUserQuotaDisabled(userID int64, pausedForwardIDs []int64, now int64) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")