Compare commits

..

16 Commits

Author SHA1 Message Date
sagit 8e513a1bae Fix nftables node updates after panel restart (#525)
## Summary
- Skip agent protocol SetProtocol commands when updating nftables nodes.
- Preserve existing nftables SSH credentials when edit forms omit secret
fields.
- Add regression coverage for nftables updates and SSH config
persistence across repository reopen.

## Test Plan
- rtk go test ./...
2026-07-02 10:28:00 +08:00
sagitchu f98be845d3 fix: skip agent protocol updates for nftables nodes 2026-07-02 10:26:10 +08:00
sagit 0c2acfdd8a Fix nftables recovery and diagnostics (#524)
## Summary
- Reconcile nftables nodes when background jobs start so rules are
restored after server reboot.
- Collect nftables traffic immediately at startup and every 30 seconds
by default.
- Return nftables rule binding status in forward diagnostics and cover
it with regression tests.

## Test Plan
- `cd go-backend && go test ./...`
- `cd go-backend && make build`
2026-06-30 17:11:10 +08:00
sagitchu 777db8767f fix nftables recovery and diagnostics 2026-06-30 17:07:50 +08:00
sagit 82f6047506 fix(forward): split proxy protocol directions
Closes #520
2026-06-21 21:03:47 +08:00
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
32 changed files with 4991 additions and 367 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 后续补齐。
@@ -14,6 +14,7 @@ import (
"time"
"go-backend/internal/http/client"
runtimenft "go-backend/internal/runtime/nftables"
"go-backend/internal/store/model"
"go-backend/internal/ws"
)
@@ -772,6 +773,9 @@ func (h *Handler) diagnoseForwardRuntime(ctx context.Context, forward *forwardRe
if ctx == nil {
ctx = context.Background()
}
if payload, handled, err := h.diagnoseNftablesForwardRuntime(forward); handled || err != nil {
return payload, err
}
forwardName, workItems, err := h.prepareForwardDiagnosis(forward)
if err != nil {
return nil, err
@@ -787,6 +791,111 @@ func (h *Handler) diagnoseForwardRuntime(ctx context.Context, forward *forwardRe
return payload, nil
}
func (h *Handler) diagnoseNftablesForwardRuntime(forward *forwardRecord) (map[string]interface{}, bool, error) {
if forward == nil {
return nil, false, errForwardNotFound
}
nftMode, entryNodeIDs, err := h.tunnelUsesNftables(forward.TunnelID)
if err != nil {
return nil, false, err
}
if !nftMode {
return nil, false, nil
}
if len(entryNodeIDs) == 0 {
return nil, true, errors.New("nftables 转发缺少入口节点")
}
targets, err := resolveDiagnosisTargets(forward.RemoteAddr)
if err != nil {
return nil, true, err
}
results, err := h.buildNftablesForwardDiagnosisResults(forward, entryNodeIDs[0], targets)
if err != nil {
return nil, true, err
}
payload := map[string]interface{}{
"forwardName": forward.Name,
"timestamp": time.Now().UnixMilli(),
"results": results,
}
return payload, true, nil
}
func (h *Handler) buildNftablesForwardDiagnosisResults(forward *forwardRecord, nodeID int64, targets []diagnosisTarget) ([]map[string]interface{}, error) {
if h == nil || h.repo == nil {
return nil, errors.New("handler not initialized")
}
node, err := h.getNodeRecord(nodeID)
if err != nil {
return nil, err
}
bindings, err := h.repo.ListNftRuleBindingsByNode(nodeID)
if err != nil {
return nil, err
}
var binding *model.NftRuleBinding
for i := range bindings {
if bindings[i].ForwardID == forward.ID {
binding = &bindings[i]
break
}
}
target := diagnosisTarget{}
if len(targets) > 0 {
target = targets[0]
}
status := "missing"
message := "nftables 规则未下发"
success := false
inPort := 0
protocols := ""
targetAddr := strings.TrimSpace(forward.RemoteAddr)
ruleHash := ""
if binding != nil {
status = strings.ToLower(strings.TrimSpace(binding.Status))
inPort = binding.InPort
protocols = strings.TrimSpace(binding.Protocols)
targetAddr = strings.TrimSpace(binding.TargetAddr)
ruleHash = strings.TrimSpace(binding.RuleHash)
if status == "" {
status = "pending"
}
if status == runtimenft.StatusApplied {
success = true
message = "nftables 规则已下发"
} else if strings.TrimSpace(binding.LastError) != "" {
message = binding.LastError
} else {
message = "nftables 规则未完成下发"
}
}
packetLoss := 100
if success {
packetLoss = 0
}
result := map[string]interface{}{
"success": success,
"nodeName": node.Name,
"nodeId": strconv.FormatInt(nodeID, 10),
"targetIp": target.IP,
"targetPort": target.Port,
"description": fmt.Sprintf("nftables规则(%s)->目标(%s)", node.Name, defaultString(target.Address, targetAddr)),
"averageTime": 0,
"packetLoss": packetLoss,
"message": message,
"fromChainType": 1,
"forwardMode": "nftables",
"nftRuleStatus": status,
"nftRuleHash": ruleHash,
"inPort": inPort,
"protocols": protocols,
"targetAddr": targetAddr,
}
return []map[string]interface{}{result}, nil
}
func (h *Handler) prepareForwardDiagnosis(forward *forwardRecord) (string, []diagnosisWorkItem, error) {
if forward == nil {
return "", nil, errForwardNotFound
@@ -1757,6 +1866,7 @@ func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel
services := make([]map[string]interface{}, 0, 2)
targets := splitRemoteTargets(forward.RemoteAddr)
strategy := strings.TrimSpace(forward.Strategy)
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(forward.ProxyProtocol, forward.ProxyProtocolReceive, forward.ProxyProtocolSend)
if strategy == "" {
strategy = "fifo"
}
@@ -1801,12 +1911,16 @@ func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel
if runtimeLimiters.TrafficLimiter != "" {
service["limiter"] = runtimeLimiters.TrafficLimiter
}
if forward.ProxyProtocol > 0 {
if proxyProtocolReceive > 0 {
serviceMetadata := ensureServiceMetadata(service)
serviceMetadata["proxyProtocol"] = proxyProtocolReceive
}
if proxyProtocolSend > 0 {
handlerConfig := service["handler"].(map[string]interface{})
if handlerConfig["metadata"] == nil {
handlerConfig["metadata"] = map[string]interface{}{}
}
handlerConfig["metadata"].(map[string]interface{})["proxyProtocol"] = forward.ProxyProtocol
handlerConfig["metadata"].(map[string]interface{})["proxyProtocol"] = proxyProtocolSend
}
if protocol == "udp" {
listenerMetadata := map[string]interface{}{
@@ -1819,10 +1933,8 @@ func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel
service["handler"].(map[string]interface{})["chain"] = fmt.Sprintf("chains_%d", forward.TunnelID)
}
if tunnel != nil && tunnel.Type == 1 && strings.TrimSpace(node.InterfaceName) != "" {
if service["metadata"] == nil {
service["metadata"] = map[string]interface{}{}
}
service["metadata"].(map[string]interface{})["interface"] = node.InterfaceName
serviceMetadata := ensureServiceMetadata(service)
serviceMetadata["interface"] = node.InterfaceName
}
services = append(services, service)
}
@@ -1841,6 +1953,25 @@ func buildForwarderNodes(targets []string) []map[string]interface{} {
return nodes
}
func ensureServiceMetadata(service map[string]interface{}) map[string]interface{} {
if service["metadata"] == nil {
service["metadata"] = map[string]interface{}{}
}
metadata, ok := service["metadata"].(map[string]interface{})
if !ok {
metadata = map[string]interface{}{}
service["metadata"] = metadata
}
return metadata
}
func normalizeForwardProxyProtocol(legacy, receive, send int) (int, int) {
if send == 0 && legacy > 0 {
send = legacy
}
return receive, send
}
func processServerAddress(serverAddr string) string {
serverAddr = normalizeServerAddressInput(serverAddr)
if serverAddr == "" {
@@ -9,14 +9,15 @@ import (
"go-backend/internal/store/repo"
)
func TestBuildForwardServiceConfigsSendsProxyProtocolToForwardHandler(t *testing.T) {
func TestBuildForwardServiceConfigsAppliesProxyProtocolReceiveAndSendIndependently(t *testing.T) {
forward := &forwardRecord{
ID: 1,
UserID: 2,
TunnelID: 3,
RemoteAddr: "1.1.1.1:443",
Strategy: "fifo",
ProxyProtocol: 2,
ID: 1,
UserID: 2,
TunnelID: 3,
RemoteAddr: "1.1.1.1:443",
Strategy: "fifo",
ProxyProtocolReceive: 1,
ProxyProtocolSend: 2,
}
tunnel := &tunnelRecord{Type: 1}
node := &nodeRecord{
@@ -38,8 +39,8 @@ func TestBuildForwardServiceConfigsSendsProxyProtocolToForwardHandler(t *testing
if serviceMetadata["interface"] != "eth0" {
t.Fatalf("expected interface metadata eth0, got %v", serviceMetadata["interface"])
}
if _, ok := serviceMetadata["proxyProtocol"]; ok {
t.Fatalf("proxyProtocol should not be listener metadata: %v", serviceMetadata)
if serviceMetadata["proxyProtocol"] != 1 {
t.Fatalf("expected service proxyProtocol 1 for receive mode, got %v", serviceMetadata["proxyProtocol"])
}
handlerConfig, ok := service["handler"].(map[string]interface{})
@@ -56,6 +57,46 @@ func TestBuildForwardServiceConfigsSendsProxyProtocolToForwardHandler(t *testing
}
}
func TestBuildForwardServiceConfigsKeepsLegacyProxyProtocolAsSend(t *testing.T) {
forward := &forwardRecord{
ID: 1,
UserID: 2,
TunnelID: 3,
RemoteAddr: "1.1.1.1:443",
Strategy: "fifo",
ProxyProtocol: 2,
}
tunnel := &tunnelRecord{Type: 1}
node := &nodeRecord{
TCPListenAddr: "0.0.0.0",
UDPListenAddr: "0.0.0.0",
}
services := buildForwardServiceConfigs("1_2_3", forward, tunnel, node, 4001, "", forwardRuntimeLimiters{})
if len(services) != 2 {
t.Fatalf("expected 2 services, got %d", len(services))
}
for _, service := range services {
serviceMetadata, _ := service["metadata"].(map[string]interface{})
if _, ok := serviceMetadata["proxyProtocol"]; ok {
t.Fatalf("legacy proxyProtocol should not enable receive mode: %v", serviceMetadata)
}
handlerConfig, ok := service["handler"].(map[string]interface{})
if !ok {
t.Fatalf("expected handler config map, got %T", service["handler"])
}
handlerMetadata, ok := handlerConfig["metadata"].(map[string]interface{})
if !ok {
t.Fatalf("expected handler metadata map, got %T", handlerConfig["metadata"])
}
if handlerMetadata["proxyProtocol"] != 2 {
t.Fatalf("expected legacy proxyProtocol to send version 2, got %v", handlerMetadata["proxyProtocol"])
}
}
}
func TestRollbackForwardMutationRestoresProxyProtocol(t *testing.T) {
r, err := repo.Open(":memory:")
if err != nil {
+54 -2
View File
@@ -2,11 +2,14 @@ package handler
import (
"context"
"log"
"time"
"go-backend/internal/license"
)
var nftablesTrafficCollectInterval = 30 * time.Second
func (h *Handler) StartBackgroundJobs() {
if h == nil || h.repo == nil {
return
@@ -20,7 +23,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 +33,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 +68,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 +132,54 @@ func (h *Handler) runTunnelQualityProber(ctx context.Context) {
h.qualityProber.Start(ctx)
}
func (h *Handler) runNftablesTrafficCollectLoop(ctx context.Context) {
defer h.jobsWG.Done()
h.runNftablesStartupReconcile(ctx)
select {
case <-ctx.Done():
return
default:
h.runNftablesTrafficCollectJob(time.Now())
}
interval := nftablesTrafficCollectInterval
if interval <= 0 {
interval = 30 * time.Second
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
h.runNftablesTrafficCollectJob(time.Now())
}
}
}
func (h *Handler) runNftablesStartupReconcile(ctx context.Context) {
if h == nil || h.repo == nil {
return
}
nodes, err := h.repo.ListNftablesNodesForCollection()
if err != nil {
log.Printf("nftables startup reconcile failed op=list_nodes err=%v", err)
return
}
for _, node := range nodes {
select {
case <-ctx.Done():
return
default:
}
if err := h.syncNftablesNode(node.NodeID); err != nil {
log.Printf("nftables startup reconcile failed node_id=%d err=%v", node.NodeID, err)
}
}
}
func (h *Handler) runHourlyStatsLoop(ctx context.Context) {
defer h.jobsWG.Done()
+51 -15
View File
@@ -417,6 +417,12 @@ func (h *Handler) nodeUpdate(w http.ResponseWriter, r *http.Request) {
newHTTP := asInt(req["http"], currentHTTP)
newTLS := asInt(req["tls"], currentTLS)
newSocks := asInt(req["socks"], currentSocks)
forwardMode := defaultNodeForwardMode(strings.TrimSpace(asString(req["forwardMode"])))
currentForwardMode, err := h.repo.GetNodeForwardMode(id)
if err != nil {
response.WriteJSON(w, response.Err(-2, err.Error()))
return
}
serverIP := asString(req["serverIp"])
if serverIP != "" {
if err := IsValidNodeAddress(serverIP); err != nil {
@@ -424,7 +430,8 @@ func (h *Handler) nodeUpdate(w http.ResponseWriter, r *http.Request) {
return
}
}
if currentStatus == 1 && (newHTTP != currentHTTP || newTLS != currentTLS || newSocks != currentSocks) {
usesNftablesRuntime := forwardMode == "nftables" || defaultNodeForwardMode(currentForwardMode) == "nftables"
if currentStatus == 1 && !usesNftablesRuntime && (newHTTP != currentHTTP || newTLS != currentTLS || newSocks != currentSocks) {
if err := h.applyNodeProtocolChange(id, newHTTP, newTLS, newSocks); err != nil {
response.WriteJSON(w, response.ErrDefault(err.Error()))
return
@@ -432,7 +439,6 @@ func (h *Handler) nodeUpdate(w http.ResponseWriter, r *http.Request) {
}
now := time.Now().UnixMilli()
forwardMode := defaultNodeForwardMode(strings.TrimSpace(asString(req["forwardMode"])))
if err := h.repo.UpdateNode(id,
asString(req["name"]),
serverIP,
@@ -2143,8 +2149,10 @@ func (h *Handler) forwardCreate(w http.ResponseWriter, r *http.Request) {
ipMaxConn = 0
}
proxyProtocol := asInt(req["proxyProtocol"], 0)
proxyProtocolReceive := asInt(req["proxyProtocolReceive"], 0)
proxyProtocolSend := asInt(req["proxyProtocolSend"], proxyProtocol)
forwardID, err := h.repo.CreateForwardTx(userID, userName, name, tunnelID, remoteAddr, defaultString(asString(req["strategy"]), "fifo"), now, inx, entryNodes, port, inIp, nullableInt(speedID), maxConn, ipMaxConn, nullableInt(ipSpeedID), proxyProtocol)
forwardID, err := h.repo.CreateForwardTx(userID, userName, name, tunnelID, remoteAddr, defaultString(asString(req["strategy"]), "fifo"), now, inx, entryNodes, port, inIp, nullableInt(speedID), maxConn, ipMaxConn, nullableInt(ipSpeedID), proxyProtocol, proxyProtocolReceive, proxyProtocolSend)
if err != nil {
response.WriteJSON(w, response.Err(-2, err.Error()))
return
@@ -2333,8 +2341,10 @@ func (h *Handler) forwardUpdate(w http.ResponseWriter, r *http.Request) {
ipMaxConn = 0
}
proxyProtocol := asInt(req["proxyProtocol"], forward.ProxyProtocol)
proxyProtocolReceive := asInt(req["proxyProtocolReceive"], forward.ProxyProtocolReceive)
proxyProtocolSend := asInt(req["proxyProtocolSend"], forward.ProxyProtocolSend)
if err := h.repo.UpdateForward(id, name, tunnelID, remoteAddr, strategy, now, newSpeedID, maxConn, ipMaxConn, newIPSpeedID, proxyProtocol); err != nil {
if err := h.repo.UpdateForward(id, name, tunnelID, remoteAddr, strategy, now, newSpeedID, maxConn, ipMaxConn, newIPSpeedID, proxyProtocol, proxyProtocolReceive, proxyProtocolSend); err != nil {
response.WriteJSON(w, response.Err(-2, err.Error()))
return
}
@@ -2412,23 +2422,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 +2594,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 +2614,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}))
}
@@ -4696,7 +4732,7 @@ func (h *Handler) rollbackForwardMutation(oldForward *forwardRecord, oldPorts []
h.repo.RollbackForwardFields(
oldForward.ID, oldForward.UserID, oldForward.UserName, oldForward.Name,
oldForward.TunnelID, oldForward.RemoteAddr, oldForward.Strategy, oldForward.Status,
oldForward.SpeedID, oldForward.MaxConn, oldForward.IPMaxConn, oldForward.IPSpeedID, oldForward.ProxyProtocol,
oldForward.SpeedID, oldForward.MaxConn, oldForward.IPMaxConn, oldForward.IPSpeedID, oldForward.ProxyProtocol, oldForward.ProxyProtocolReceive, oldForward.ProxyProtocolSend,
time.Now().UnixMilli(),
)
@@ -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
}
@@ -10,6 +10,7 @@ import (
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"testing"
"time"
@@ -20,21 +21,29 @@ import (
)
type fakeNftablesManager struct {
testErr error
reconcileErr error
reconcileHit int
clearErr error
clearHit int
lastConfig runtimenft.SSHConfig
lastPlan runtimenft.NodePlan
mu sync.Mutex
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 {
f.mu.Lock()
defer f.mu.Unlock()
f.lastConfig = cfg
return f.testErr
}
func (f *fakeNftablesManager) Reconcile(_ context.Context, cfg runtimenft.SSHConfig, plan runtimenft.NodePlan) (runtimenft.ApplyResult, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.reconcileHit++
f.lastConfig = cfg
f.lastPlan = plan
@@ -44,15 +53,40 @@ func (f *fakeNftablesManager) Reconcile(_ context.Context, cfg runtimenft.SSHCon
return runtimenft.ApplyResult{
NodeID: plan.NodeID,
Script: "table inet flvx {}",
Hashes: map[int64]string{plan.NodeID: "hash"},
Hashes: runtimenft.PlanHashes(plan),
}, nil
}
func (f *fakeNftablesManager) Clear(context.Context, runtimenft.SSHConfig) error {
f.mu.Lock()
defer f.mu.Unlock()
f.clearHit++
return f.clearErr
}
func (f *fakeNftablesManager) CollectCounters(_ context.Context, cfg runtimenft.SSHConfig) ([]runtimenft.CounterSample, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.collectHit++
f.lastConfig = cfg
if f.collectErr != nil {
return nil, f.collectErr
}
return f.counterSamples, nil
}
func (f *fakeNftablesManager) reconcileCount() int {
f.mu.Lock()
defer f.mu.Unlock()
return f.reconcileHit
}
func (f *fakeNftablesManager) collectCount() int {
f.mu.Lock()
defer f.mu.Unlock()
return f.collectHit
}
type nftablesTestFixture struct {
handler *Handler
nodeID int64
@@ -146,6 +180,49 @@ func TestNodeNftablesReconcileEndpointPersistsBindings(t *testing.T) {
}
}
func TestStartBackgroundJobsReconcilesNftablesRulesAtStartup(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
seedNftablesSSHConfig(t, h, fixture.nodeID)
tunnelID := seedTunnelForNftables(t, h, "nft-startup-tunnel", fixture.nodeID)
seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
manager := &fakeNftablesManager{}
h.nftablesManager = manager
h.StartBackgroundJobs()
t.Cleanup(h.StopBackgroundJobs)
waitForCondition(t, time.Second, func() bool {
return manager.reconcileCount() > 0
}, "nftables startup reconcile")
}
func TestStartBackgroundJobsCollectsNftablesTrafficImmediatelyAndUsesFastInterval(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},
}}
h.nftablesManager = manager
oldInterval := nftablesTrafficCollectInterval
nftablesTrafficCollectInterval = 20 * time.Millisecond
t.Cleanup(func() { nftablesTrafficCollectInterval = oldInterval })
h.StartBackgroundJobs()
t.Cleanup(h.StopBackgroundJobs)
waitForCondition(t, time.Second, func() bool {
return manager.collectCount() >= 2
}, "immediate and repeated nftables traffic collection")
}
func TestNftablesTrafficCollectIntervalDefaultsToThirtySeconds(t *testing.T) {
if nftablesTrafficCollectInterval != 30*time.Second {
t.Fatalf("expected default nftables traffic collection interval 30s, got %s", nftablesTrafficCollectInterval)
}
}
func TestNodeNftablesClearEndpointClearsBindings(t *testing.T) {
fixture := setupNftablesHandler(t)
seedNftablesSSHConfig(t, fixture.handler, fixture.nodeID)
@@ -268,6 +345,83 @@ func TestNodeUpdatePreservesExistingNftablesSecretsWhenFieldsOmitted(t *testing.
}
}
func TestNodeUpdateSkipsAgentProtocolCommandForNftablesNode(t *testing.T) {
fixture := setupNftablesHandler(t)
seedNftablesSSHConfig(t, fixture.handler, fixture.nodeID)
req := newAuthenticatedJSONRequest(t, map[string]interface{}{
"id": fixture.nodeID,
"name": "nft-node-updated",
"serverIp": "198.51.100.10",
"serverIpV4": "198.51.100.10",
"port": "1000-65535",
"forwardMode": "nftables",
"http": 1,
"tls": 1,
"socks": 1,
"sshConfig": map[string]interface{}{
"host": "203.0.113.30",
"port": 22,
"username": "admin",
"authType": "password",
"sudoMode": "none",
},
})
res := httptest.NewRecorder()
fixture.handler.nodeUpdate(res, req)
assertNftablesSuccessWithBody(t, res)
cfg, err := fixture.handler.repo.GetNodeSSHConfig(fixture.nodeID)
if err != nil {
t.Fatalf("load ssh config: %v", err)
}
if !cfg.Password.Valid || cfg.Password.String != "secret" {
t.Fatalf("expected password secret to be preserved, got %+v", cfg)
}
}
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 +461,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
@@ -343,6 +521,61 @@ func TestTunnelBatchRedeployUsesNftablesReconcile(t *testing.T) {
}
}
func TestDiagnoseForwardRuntimeReturnsNftablesRuleStatus(t *testing.T) {
fixture := setupNftablesHandler(t)
h := fixture.handler
tunnelID := seedTunnelForNftables(t, h, "nft-diagnose-tunnel", fixture.nodeID)
forward := seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
now := time.Now().UnixMilli()
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)
}
payload, err := h.diagnoseForwardRuntime(context.Background(), &forwardRecord{
ID: forward.ID,
Name: forward.Name,
TunnelID: tunnelID,
RemoteAddr: "203.0.113.9:8080",
})
if err != nil {
t.Fatalf("diagnose forward: %v", err)
}
results, ok := payload["results"].([]map[string]interface{})
if !ok || len(results) != 1 {
t.Fatalf("expected one nftables diagnosis result, got %#v", payload["results"])
}
result := results[0]
if result["forwardMode"] != "nftables" || result["nftRuleStatus"] != runtimenft.StatusApplied {
t.Fatalf("expected nftables applied result, got %#v", result)
}
if result["success"] != true {
t.Fatalf("expected nftables diagnosis success, got %#v", result)
}
if !strings.Contains(asString(result["message"]), "已下发") {
t.Fatalf("expected applied message, got %#v", result["message"])
}
}
func waitForCondition(t *testing.T, timeout time.Duration, condition func() bool, description string) {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
if condition() {
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("timed out waiting for %s", description)
}
func setupNftablesHandler(t *testing.T) nftablesTestFixture {
t.Helper()
@@ -412,7 +645,7 @@ func seedForwardForNftables(t *testing.T, h *Handler, tunnelID, nodeID int64, re
now := time.Now().UnixMilli()
forwardID, err := h.repo.CreateForwardTx(
1, "admin", "nft-forward", tunnelID, remoteAddr, "fifo", now, 1,
[]int64{nodeID}, 20000, "", nil, 0, 0, nil, 0,
[]int64{nodeID}, 20000, "", nil, 0, 0, nil, 0, 0, 0,
)
if err != nil {
t.Fatalf("create forward: %v", err)
@@ -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 {
+71 -49
View File
@@ -31,24 +31,26 @@ func (User) TableName() string { return "user" }
// Forward maps to the "forward" table.
type Forward struct {
ID int64 `gorm:"primaryKey;autoIncrement"`
UserID int64 `gorm:"column:user_id;not null"`
UserName string `gorm:"column:user_name;type:varchar(100);not null"`
Name string `gorm:"type:varchar(100);not null"`
TunnelID int64 `gorm:"column:tunnel_id;not null"`
RemoteAddr string `gorm:"column:remote_addr;type:text;not null"`
Strategy string `gorm:"type:varchar(100);not null;default:'fifo'"`
InFlow int64 `gorm:"not null;default:0"`
OutFlow int64 `gorm:"column:out_flow;not null;default:0"`
CreatedTime int64 `gorm:"column:created_time;not null"`
UpdatedTime int64 `gorm:"column:updated_time;not null"`
Status int `gorm:"not null"`
Inx int `gorm:"not null;default:0"`
SpeedID sql.NullInt64 `gorm:"column:speed_id"`
MaxConn int `gorm:"column:max_conn;not null;default:0"`
IPMaxConn int `gorm:"column:ip_max_conn;not null;default:0"`
IPSpeedID sql.NullInt64 `gorm:"column:ip_speed_id"`
ProxyProtocol int `gorm:"column:proxy_protocol;not null;default:0"`
ID int64 `gorm:"primaryKey;autoIncrement"`
UserID int64 `gorm:"column:user_id;not null"`
UserName string `gorm:"column:user_name;type:varchar(100);not null"`
Name string `gorm:"type:varchar(100);not null"`
TunnelID int64 `gorm:"column:tunnel_id;not null"`
RemoteAddr string `gorm:"column:remote_addr;type:text;not null"`
Strategy string `gorm:"type:varchar(100);not null;default:'fifo'"`
InFlow int64 `gorm:"not null;default:0"`
OutFlow int64 `gorm:"column:out_flow;not null;default:0"`
CreatedTime int64 `gorm:"column:created_time;not null"`
UpdatedTime int64 `gorm:"column:updated_time;not null"`
Status int `gorm:"not null"`
Inx int `gorm:"not null;default:0"`
SpeedID sql.NullInt64 `gorm:"column:speed_id"`
MaxConn int `gorm:"column:max_conn;not null;default:0"`
IPMaxConn int `gorm:"column:ip_max_conn;not null;default:0"`
IPSpeedID sql.NullInt64 `gorm:"column:ip_speed_id"`
ProxyProtocol int `gorm:"column:proxy_protocol;not null;default:0"`
ProxyProtocolReceive int `gorm:"column:proxy_protocol_receive;not null;default:0"`
ProxyProtocolSend int `gorm:"column:proxy_protocol_send;not null;default:0"`
}
func (Forward) TableName() string { return "forward" }
@@ -131,6 +133,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"`
@@ -471,24 +489,26 @@ type ChainTunnelBackup struct {
}
type ForwardBackup struct {
ID int64 `json:"id"`
UserID int64 `json:"userId"`
UserName string `json:"userName"`
Name string `json:"name"`
TunnelID int64 `json:"tunnelId"`
RemoteAddr string `json:"remoteAddr"`
Strategy string `json:"strategy"`
InFlow int64 `json:"inFlow"`
OutFlow int64 `json:"outFlow"`
CreatedTime int64 `json:"createdTime"`
UpdatedTime int64 `json:"updatedTime"`
Status int `json:"status"`
Inx int `json:"inx"`
SpeedID *int64 `json:"speedId,omitempty"`
IPMaxConn int `json:"ipMaxConn,omitempty"`
IPSpeedID *int64 `json:"ipSpeedId,omitempty"`
ForwardPorts *[]ForwardPortBackup `json:"forwardPorts,omitempty"`
ProxyProtocol int `json:"proxyProtocol"`
ID int64 `json:"id"`
UserID int64 `json:"userId"`
UserName string `json:"userName"`
Name string `json:"name"`
TunnelID int64 `json:"tunnelId"`
RemoteAddr string `json:"remoteAddr"`
Strategy string `json:"strategy"`
InFlow int64 `json:"inFlow"`
OutFlow int64 `json:"outFlow"`
CreatedTime int64 `json:"createdTime"`
UpdatedTime int64 `json:"updatedTime"`
Status int `json:"status"`
Inx int `json:"inx"`
SpeedID *int64 `json:"speedId,omitempty"`
IPMaxConn int `json:"ipMaxConn,omitempty"`
IPSpeedID *int64 `json:"ipSpeedId,omitempty"`
ForwardPorts *[]ForwardPortBackup `json:"forwardPorts,omitempty"`
ProxyProtocol int `json:"proxyProtocol"`
ProxyProtocolReceive int `json:"proxyProtocolReceive,omitempty"`
ProxyProtocolSend int `json:"proxyProtocolSend,omitempty"`
}
type ForwardPortBackup struct {
@@ -577,19 +597,21 @@ type ImportResult struct {
// ForwardRecord is a minimal forward view used by control plane and flow policy.
type ForwardRecord struct {
ID int64
UserID int64
UserName string
Name string
TunnelID int64
RemoteAddr string
Strategy string
Status int
SpeedID sql.NullInt64
MaxConn int
IPMaxConn int
IPSpeedID sql.NullInt64
ProxyProtocol int
ID int64
UserID int64
UserName string
Name string
TunnelID int64
RemoteAddr string
Strategy string
Status int
SpeedID sql.NullInt64
MaxConn int
IPMaxConn int
IPSpeedID sql.NullInt64
ProxyProtocol int
ProxyProtocolReceive int
ProxyProtocolSend int
}
// TunnelRecord is a minimal tunnel view used by control plane.
+93 -72
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{},
@@ -412,7 +424,7 @@ func prepareSQLiteLegacyColumns(db *gorm.DB) error {
}
if m.HasTable(&model.Forward{}) {
for _, field := range []string{"MaxConn", "IPMaxConn", "IPSpeedID", "ProxyProtocol"} {
for _, field := range []string{"MaxConn", "IPMaxConn", "IPSpeedID", "ProxyProtocol", "ProxyProtocolReceive", "ProxyProtocolSend"} {
if m.HasColumn(&model.Forward{}, field) {
continue
}
@@ -932,31 +944,33 @@ func (r *Repository) ListForwards() ([]map[string]interface{}, error) {
}
type fwdRow struct {
ID int64
UserID int64
UserName string
Name string
TunnelID int64
TunnelName string
TrafficRatio float64
RemoteAddr string
Strategy string
InFlow int64
OutFlow int64
CreatedTime int64
Status int
Inx int
SpeedID sql.NullInt64
MaxConn int
IPMaxConn int
IPSpeedID sql.NullInt64
IPSpeedLimitName string
ProxyProtocol int
ID int64
UserID int64
UserName string
Name string
TunnelID int64
TunnelName string
TrafficRatio float64
RemoteAddr string
Strategy string
InFlow int64
OutFlow int64
CreatedTime int64
Status int
Inx int
SpeedID sql.NullInt64
MaxConn int
IPMaxConn int
IPSpeedID sql.NullInt64
IPSpeedLimitName string
ProxyProtocol int
ProxyProtocolReceive int
ProxyProtocolSend int
}
var rows []fwdRow
err := r.db.Model(&model.Forward{}).
Select("forward.id, forward.user_id, forward.user_name, forward.name, forward.tunnel_id, COALESCE(tunnel.name, '') AS tunnel_name, COALESCE(tunnel.traffic_ratio, 1.0) AS traffic_ratio, forward.remote_addr, COALESCE(forward.strategy, 'fifo') AS strategy, forward.in_flow, forward.out_flow, forward.created_time, forward.status, forward.inx, forward.speed_id, forward.max_conn, forward.ip_max_conn, forward.ip_speed_id, COALESCE(ip_speed_limit.name, '') AS ip_speed_limit_name, forward.proxy_protocol").
Select("forward.id, forward.user_id, forward.user_name, forward.name, forward.tunnel_id, COALESCE(tunnel.name, '') AS tunnel_name, COALESCE(tunnel.traffic_ratio, 1.0) AS traffic_ratio, forward.remote_addr, COALESCE(forward.strategy, 'fifo') AS strategy, forward.in_flow, forward.out_flow, forward.created_time, forward.status, forward.inx, forward.speed_id, forward.max_conn, forward.ip_max_conn, forward.ip_speed_id, COALESCE(ip_speed_limit.name, '') AS ip_speed_limit_name, forward.proxy_protocol, forward.proxy_protocol_receive, forward.proxy_protocol_send").
Joins("LEFT JOIN tunnel ON tunnel.id = forward.tunnel_id").
Joins("LEFT JOIN speed_limit AS ip_speed_limit ON ip_speed_limit.id = forward.ip_speed_id").
Order("forward.inx ASC, forward.id ASC").
@@ -967,6 +981,7 @@ func (r *Repository) ListForwards() ([]map[string]interface{}, error) {
items := make([]map[string]interface{}, 0, len(rows))
for _, row := range rows {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(row.ProxyProtocol, row.ProxyProtocolReceive, row.ProxyProtocolSend)
inIP, inPort, err := resolveForwardIngress(r.db, row.ID, row.TunnelID)
if err != nil {
return nil, err
@@ -979,9 +994,11 @@ func (r *Repository) ListForwards() ([]map[string]interface{}, error) {
"remoteAddr": row.RemoteAddr, "strategy": row.Strategy,
"inFlow": row.InFlow, "outFlow": row.OutFlow,
"createdTime": row.CreatedTime, "status": row.Status, "inx": int64(row.Inx),
"maxConn": row.MaxConn,
"ipMaxConn": row.IPMaxConn,
"proxyProtocol": row.ProxyProtocol,
"maxConn": row.MaxConn,
"ipMaxConn": row.IPMaxConn,
"proxyProtocol": row.ProxyProtocol,
"proxyProtocolReceive": proxyProtocolReceive,
"proxyProtocolSend": proxyProtocolSend,
}
if row.SpeedID.Valid {
item["speedId"] = row.SpeedID.Int64
@@ -2183,8 +2200,10 @@ func (r *Repository) exportForwards() ([]model.ForwardBackup, error) {
TunnelID: f.TunnelID, RemoteAddr: f.RemoteAddr, Strategy: f.Strategy,
InFlow: f.InFlow, OutFlow: f.OutFlow, CreatedTime: f.CreatedTime,
UpdatedTime: f.UpdatedTime, Status: f.Status, Inx: f.Inx,
IPMaxConn: f.IPMaxConn,
ProxyProtocol: f.ProxyProtocol,
IPMaxConn: f.IPMaxConn,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: f.ProxyProtocolReceive,
ProxyProtocolSend: f.ProxyProtocolSend,
}
if f.SpeedID.Valid {
v := f.SpeedID.Int64
@@ -2618,29 +2637,31 @@ func importForwards(tx *gorm.DB, forwards []model.ForwardBackup, now int64) (int
count := 0
for _, f := range forwards {
item := model.Forward{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
InFlow: f.InFlow,
OutFlow: f.OutFlow,
CreatedTime: f.CreatedTime,
UpdatedTime: now,
Status: f.Status,
Inx: f.Inx,
SpeedID: sql.NullInt64{Int64: nullableBackupInt64(f.SpeedID), Valid: f.SpeedID != nil && *f.SpeedID > 0},
IPMaxConn: f.IPMaxConn,
IPSpeedID: sql.NullInt64{Int64: nullableBackupInt64(f.IPSpeedID), Valid: f.IPSpeedID != nil && *f.IPSpeedID > 0},
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
InFlow: f.InFlow,
OutFlow: f.OutFlow,
CreatedTime: f.CreatedTime,
UpdatedTime: now,
Status: f.Status,
Inx: f.Inx,
SpeedID: sql.NullInt64{Int64: nullableBackupInt64(f.SpeedID), Valid: f.SpeedID != nil && *f.SpeedID > 0},
IPMaxConn: f.IPMaxConn,
IPSpeedID: sql.NullInt64{Int64: nullableBackupInt64(f.IPSpeedID), Valid: f.IPSpeedID != nil && *f.IPSpeedID > 0},
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: f.ProxyProtocolReceive,
ProxyProtocolSend: f.ProxyProtocolSend,
}
err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
DoUpdates: clause.AssignmentColumns([]string{
"user_id", "user_name", "name", "tunnel_id", "remote_addr", "strategy",
"in_flow", "out_flow", "updated_time", "status", "inx", "speed_id", "ip_max_conn", "ip_speed_id", "proxy_protocol",
"in_flow", "out_flow", "updated_time", "status", "inx", "speed_id", "ip_max_conn", "ip_speed_id", "proxy_protocol", "proxy_protocol_receive", "proxy_protocol_send",
}),
}).Create(&item).Error
if err != nil {
@@ -44,20 +44,23 @@ func (r *Repository) ListForwardsByTunnelTx(tx *gorm.DB, tunnelID int64) ([]mode
}
rows := make([]model.ForwardRecord, 0, len(forwards))
for _, f := range forwards {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
rows = append(rows, model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
})
}
for i := range rows {
@@ -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,
@@ -113,20 +120,23 @@ func (r *Repository) ListActiveForwardsByUser(userID int64) ([]model.ForwardReco
}
rows := make([]model.ForwardRecord, 0, len(forwards))
for _, f := range forwards {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
rows = append(rows, model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
})
}
for i := range rows {
@@ -148,20 +158,23 @@ func (r *Repository) ListActiveForwardsByUserTunnel(userID, tunnelID int64) ([]m
}
rows := make([]model.ForwardRecord, 0, len(forwards))
for _, f := range forwards {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
rows = append(rows, model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
})
}
for i := range rows {
@@ -183,20 +196,23 @@ func (r *Repository) ListForwardsByUserAndTunnel(userID, tunnelID int64) ([]mode
}
rows := make([]model.ForwardRecord, 0, len(forwards))
for _, f := range forwards {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
rows = append(rows, model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
})
}
for i := range rows {
@@ -219,20 +235,23 @@ func (r *Repository) GetForwardRecord(forwardID int64) (*model.ForwardRecord, er
}
return nil, err
}
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
fr := model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
}
if strings.TrimSpace(fr.Strategy) == "" {
fr.Strategy = "fifo"
@@ -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)
}
}
@@ -160,7 +160,7 @@ func TestForwardRepositoryPersistsPerIPLimits(t *testing.T) {
defer r.Close()
now := time.Now().UnixMilli()
forwardID, err := r.CreateForwardTx(1, "admin", "per-ip-forward", 2, "1.1.1.1:443", "fifo", now, 1, []int64{3}, 24000, "", nil, 0, 5, int64(21), 0)
forwardID, err := r.CreateForwardTx(1, "admin", "per-ip-forward", 2, "1.1.1.1:443", "fifo", now, 1, []int64{3}, 24000, "", nil, 0, 5, int64(21), 0, 0, 0)
if err != nil {
t.Fatalf("CreateForwardTx: %v", err)
}
@@ -175,7 +175,7 @@ func TestForwardRepositoryPersistsPerIPLimits(t *testing.T) {
t.Fatalf("expected created ipSpeedId 21, got %+v", record.IPSpeedID)
}
if err := r.UpdateForward(forwardID, "per-ip-forward", 2, "2.2.2.2:443", "fifo", now+1, nil, 0, 9, int64(22), 0); err != nil {
if err := r.UpdateForward(forwardID, "per-ip-forward", 2, "2.2.2.2:443", "fifo", now+1, nil, 0, 9, int64(22), 0, 0, 0); err != nil {
t.Fatalf("UpdateForward: %v", err)
}
record, err = r.GetForwardRecord(forwardID)
@@ -216,6 +216,83 @@ func TestForwardRepositoryPersistsPerIPLimits(t *testing.T) {
}
}
func TestForwardRepositoryPersistsProxyProtocolReceiveAndSend(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now().UnixMilli()
forwardID, err := r.CreateForwardTx(1, "admin", "proxy-protocol-forward", 2, "1.1.1.1:443", "fifo", now, 1, []int64{3}, 24000, "", nil, 0, 0, nil, 0, 1, 2)
if err != nil {
t.Fatalf("CreateForwardTx: %v", err)
}
record, err := r.GetForwardRecord(forwardID)
if err != nil {
t.Fatalf("GetForwardRecord after create: %v", err)
}
if record.ProxyProtocolReceive != 1 || record.ProxyProtocolSend != 2 {
t.Fatalf("expected proxyProtocol receive/send 1/2 after create, got %d/%d", record.ProxyProtocolReceive, record.ProxyProtocolSend)
}
if err := r.UpdateForward(forwardID, "proxy-protocol-forward", 2, "2.2.2.2:443", "fifo", now+1, nil, 0, 0, nil, 0, 2, 1); err != nil {
t.Fatalf("UpdateForward: %v", err)
}
record, err = r.GetForwardRecord(forwardID)
if err != nil {
t.Fatalf("GetForwardRecord after update: %v", err)
}
if record.ProxyProtocolReceive != 2 || record.ProxyProtocolSend != 1 {
t.Fatalf("expected proxyProtocol receive/send 2/1 after update, got %d/%d", record.ProxyProtocolReceive, record.ProxyProtocolSend)
}
records, err := r.ListForwardsByTunnel(2)
if err != nil {
t.Fatalf("ListForwardsByTunnel: %v", err)
}
if len(records) != 1 {
t.Fatalf("expected 1 listed record, got %d", len(records))
}
if records[0].ProxyProtocolReceive != 2 || records[0].ProxyProtocolSend != 1 {
t.Fatalf("expected listed proxyProtocol receive/send 2/1, got %d/%d", records[0].ProxyProtocolReceive, records[0].ProxyProtocolSend)
}
}
func TestForwardRepositoryMapsLegacyProxyProtocolToSend(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
t.Fatalf("open repo: %v", err)
}
defer r.Close()
now := time.Now().UnixMilli()
if err := r.DB().Create(&model.Forward{
UserID: 2,
UserName: "user",
Name: "legacy-proxy-protocol-forward",
TunnelID: 9,
RemoteAddr: "1.1.1.1:443",
Strategy: "fifo",
CreatedTime: now,
UpdatedTime: now,
Status: 1,
ProxyProtocol: 2,
}).Error; err != nil {
t.Fatalf("create legacy forward: %v", err)
}
forwardID := mustRepoLastInsertID(t, r)
record, err := r.GetForwardRecord(forwardID)
if err != nil {
t.Fatalf("GetForwardRecord: %v", err)
}
if record.ProxyProtocolReceive != 0 || record.ProxyProtocolSend != 2 {
t.Fatalf("expected legacy proxyProtocol to map to receive/send 0/2, got %d/%d", record.ProxyProtocolReceive, record.ProxyProtocolSend)
}
}
func TestRollbackForwardFieldsRestoresPerIPLimits(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
@@ -224,15 +301,15 @@ func TestRollbackForwardFieldsRestoresPerIPLimits(t *testing.T) {
defer r.Close()
now := time.Now().UnixMilli()
forwardID, err := r.CreateForwardTx(1, "admin", "rollback-per-ip-forward", 2, "1.1.1.1:443", "fifo", now, 1, nil, 0, "", nil, 7, 5, int64(21), 2)
forwardID, err := r.CreateForwardTx(1, "admin", "rollback-per-ip-forward", 2, "1.1.1.1:443", "fifo", now, 1, nil, 0, "", nil, 7, 5, int64(21), 2, 0, 2)
if err != nil {
t.Fatalf("CreateForwardTx: %v", err)
}
if err := r.UpdateForward(forwardID, "rollback-per-ip-forward", 2, "2.2.2.2:443", "fifo", now+1, nil, 0, 0, nil, 0); err != nil {
if err := r.UpdateForward(forwardID, "rollback-per-ip-forward", 2, "2.2.2.2:443", "fifo", now+1, nil, 0, 0, nil, 0, 0, 0); err != nil {
t.Fatalf("UpdateForward: %v", err)
}
r.RollbackForwardFields(forwardID, 1, "admin", "rollback-per-ip-forward", 2, "1.1.1.1:443", "fifo", 1, nil, 7, 5, int64(21), 2, now+2)
r.RollbackForwardFields(forwardID, 1, "admin", "rollback-per-ip-forward", 2, "1.1.1.1:443", "fifo", 1, nil, 7, 5, int64(21), 2, 0, 2, now+2)
record, err := r.GetForwardRecord(forwardID)
if err != nil {
@@ -726,23 +726,26 @@ func (r *Repository) GetMinForwardPort(forwardID int64) sql.NullInt64 {
return p
}
func (r *Repository) UpdateForward(id int64, name string, tunnelID int64, remoteAddr, strategy string, now int64, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int) error {
func (r *Repository) UpdateForward(id int64, name string, tunnelID int64, remoteAddr, strategy string, now int64, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int, proxyProtocolReceive int, proxyProtocolSend int) error {
if r == nil || r.db == nil {
return errors.New("repository not initialized")
}
_, proxyProtocolSend = normalizeForwardProxyProtocol(proxyProtocol, proxyProtocolReceive, proxyProtocolSend)
return r.db.Model(&model.Forward{}).
Where("id = ?", id).
Updates(map[string]interface{}{
"name": name,
"tunnel_id": tunnelID,
"remote_addr": remoteAddr,
"strategy": strategy,
"speed_id": nullInt64FromInterface(speedID),
"max_conn": maxConn,
"ip_max_conn": ipMaxConn,
"ip_speed_id": nullInt64FromInterface(ipSpeedID),
"proxy_protocol": proxyProtocol,
"updated_time": now,
"name": name,
"tunnel_id": tunnelID,
"remote_addr": remoteAddr,
"strategy": strategy,
"speed_id": nullInt64FromInterface(speedID),
"max_conn": maxConn,
"ip_max_conn": ipMaxConn,
"ip_speed_id": nullInt64FromInterface(ipSpeedID),
"proxy_protocol": proxyProtocol,
"proxy_protocol_receive": proxyProtocolReceive,
"proxy_protocol_send": proxyProtocolSend,
"updated_time": now,
}).Error
}
@@ -769,6 +772,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
}
@@ -816,26 +822,29 @@ func (r *Repository) UpdateForwardPortBindIP(forwardID, nodeID int64, port int,
Update("in_ip", sql.NullString{String: inIP, Valid: strings.TrimSpace(inIP) != ""}).Error
}
func (r *Repository) RollbackForwardFields(id, userID int64, userName, name string, tunnelID int64, remoteAddr, strategy string, status int, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int, now int64) {
func (r *Repository) RollbackForwardFields(id, userID int64, userName, name string, tunnelID int64, remoteAddr, strategy string, status int, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int, proxyProtocolReceive int, proxyProtocolSend int, now int64) {
if r == nil || r.db == nil {
return
}
_, proxyProtocolSend = normalizeForwardProxyProtocol(proxyProtocol, proxyProtocolReceive, proxyProtocolSend)
_ = r.db.Model(&model.Forward{}).
Where("id = ?", id).
Updates(map[string]interface{}{
"user_id": userID,
"user_name": userName,
"name": name,
"tunnel_id": tunnelID,
"remote_addr": remoteAddr,
"strategy": strategy,
"status": status,
"speed_id": nullInt64FromInterface(speedID),
"max_conn": maxConn,
"ip_max_conn": ipMaxConn,
"ip_speed_id": nullInt64FromInterface(ipSpeedID),
"proxy_protocol": proxyProtocol,
"updated_time": now,
"user_id": userID,
"user_name": userName,
"name": name,
"tunnel_id": tunnelID,
"remote_addr": remoteAddr,
"strategy": strategy,
"status": status,
"speed_id": nullInt64FromInterface(speedID),
"max_conn": maxConn,
"ip_max_conn": ipMaxConn,
"ip_speed_id": nullInt64FromInterface(ipSpeedID),
"proxy_protocol": proxyProtocol,
"proxy_protocol_receive": proxyProtocolReceive,
"proxy_protocol_send": proxyProtocolSend,
"updated_time": now,
}).Error
}
@@ -1295,30 +1304,33 @@ func (r *Repository) EnsureUserTunnelGrant(userID, tunnelID int64) (int64, bool,
return ut.ID, true, nil
}
func (r *Repository) CreateForwardTx(userID int64, userName, name string, tunnelID int64, remoteAddr, strategy string, now int64, inx int, entryNodeIDs []int64, port int, inIp string, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int) (int64, error) {
func (r *Repository) CreateForwardTx(userID int64, userName, name string, tunnelID int64, remoteAddr, strategy string, now int64, inx int, entryNodeIDs []int64, port int, inIp string, speedID interface{}, maxConn int, ipMaxConn int, ipSpeedID interface{}, proxyProtocol int, proxyProtocolReceive int, proxyProtocolSend int) (int64, error) {
if r == nil || r.db == nil {
return 0, errors.New("repository not initialized")
}
_, proxyProtocolSend = normalizeForwardProxyProtocol(proxyProtocol, proxyProtocolReceive, proxyProtocolSend)
var forwardID int64
err := r.db.Transaction(func(tx *gorm.DB) error {
fwd := model.Forward{
UserID: userID,
UserName: userName,
Name: name,
TunnelID: tunnelID,
RemoteAddr: remoteAddr,
Strategy: strategy,
InFlow: 0,
OutFlow: 0,
CreatedTime: now,
UpdatedTime: now,
Status: 1,
Inx: inx,
MaxConn: maxConn,
SpeedID: nullInt64FromInterface(speedID),
IPMaxConn: ipMaxConn,
IPSpeedID: nullInt64FromInterface(ipSpeedID),
ProxyProtocol: proxyProtocol,
UserID: userID,
UserName: userName,
Name: name,
TunnelID: tunnelID,
RemoteAddr: remoteAddr,
Strategy: strategy,
InFlow: 0,
OutFlow: 0,
CreatedTime: now,
UpdatedTime: now,
Status: 1,
Inx: inx,
MaxConn: maxConn,
SpeedID: nullInt64FromInterface(speedID),
IPMaxConn: ipMaxConn,
IPSpeedID: nullInt64FromInterface(ipSpeedID),
ProxyProtocol: proxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
}
if err := tx.Create(&fwd).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)
}
}
@@ -207,20 +207,23 @@ func (r *Repository) ListActiveForwardsByNode(nodeID int64) ([]model.ForwardReco
}
rows := make([]model.ForwardRecord, 0, len(forwards))
for _, f := range forwards {
proxyProtocolReceive, proxyProtocolSend := normalizeForwardProxyProtocol(f.ProxyProtocol, f.ProxyProtocolReceive, f.ProxyProtocolSend)
rows = append(rows, model.ForwardRecord{
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ID: f.ID,
UserID: f.UserID,
UserName: f.UserName,
Name: f.Name,
TunnelID: f.TunnelID,
RemoteAddr: f.RemoteAddr,
Strategy: f.Strategy,
Status: f.Status,
SpeedID: f.SpeedID,
MaxConn: f.MaxConn,
IPMaxConn: f.IPMaxConn,
IPSpeedID: f.IPSpeedID,
ProxyProtocol: f.ProxyProtocol,
ProxyProtocolReceive: proxyProtocolReceive,
ProxyProtocolSend: proxyProtocolSend,
})
}
for i := range rows {
@@ -237,3 +240,10 @@ func defaultString(value, fallback string) string {
}
return value
}
func normalizeForwardProxyProtocol(legacy, receive, send int) (int, int) {
if send == 0 && legacy > 0 {
send = legacy
}
return receive, send
}
@@ -1,6 +1,7 @@
package repo
import (
"path/filepath"
"strings"
"testing"
"time"
@@ -123,6 +124,114 @@ func TestNftablesNodeModeSSHConfigAndBindingPersistence(t *testing.T) {
}
}
func TestNftablesNodeSSHConfigSurvivesRepositoryReopen(t *testing.T) {
tests := []struct {
name string
cfg NftSSHConfigInput
}{
{
name: "password",
cfg: NftSSHConfigInput{
Host: "203.0.113.10",
Port: 2222,
Username: "root",
AuthType: "password",
Password: "ssh-password",
Passphrase: "key-passphrase",
SudoMode: "sudo",
},
},
{
name: "private-key",
cfg: NftSSHConfigInput{
Host: "203.0.113.11",
Port: 2223,
Username: "admin",
AuthType: "private_key",
PrivateKey: "PRIVATE-KEY-SHOULD-PERSIST",
Passphrase: "key-passphrase",
SudoMode: "sudo_su",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "nftables-ssh.sqlite")
r, err := Open(dbPath)
if err != nil {
t.Fatalf("open repo: %v", err)
}
now := time.Now().UnixMilli()
if err := r.CreateNode(
"nft-node",
"secret",
tt.cfg.Host,
nil,
nil,
"10000-20000",
nil,
nil,
nil,
nil,
nil,
0,
0,
0,
now,
1,
"[::]",
"[::]",
1,
0,
nil,
nil,
nil,
nil,
"nftables",
); err != nil {
t.Fatalf("CreateNode: %v", err)
}
nodes, err := r.ListNodes()
if err != nil {
t.Fatalf("ListNodes: %v", err)
}
nodeID := nodes[0]["id"].(int64)
if err := r.UpsertNodeSSHConfig(nodeID, tt.cfg, now); err != nil {
t.Fatalf("UpsertNodeSSHConfig: %v", err)
}
if err := r.Close(); err != nil {
t.Fatalf("close repo: %v", err)
}
reopened, err := Open(dbPath)
if err != nil {
t.Fatalf("reopen repo: %v", err)
}
defer reopened.Close()
loaded, err := reopened.GetNodeSSHConfig(nodeID)
if err != nil {
t.Fatalf("GetNodeSSHConfig after reopen: %v", err)
}
if loaded.Host != tt.cfg.Host || loaded.Port != tt.cfg.Port || loaded.Username != tt.cfg.Username || loaded.AuthType != tt.cfg.AuthType || loaded.SudoMode != tt.cfg.SudoMode {
t.Fatalf("unexpected ssh config after reopen: %+v", loaded)
}
if tt.cfg.Password != "" && (!loaded.Password.Valid || loaded.Password.String != tt.cfg.Password) {
t.Fatalf("expected password to persist after reopen, got %+v", loaded.Password)
}
if tt.cfg.PrivateKey != "" && (!loaded.PrivateKey.Valid || loaded.PrivateKey.String != tt.cfg.PrivateKey) {
t.Fatalf("expected private key to persist after reopen, got %+v", loaded.PrivateKey)
}
if tt.cfg.Passphrase != "" && (!loaded.Passphrase.Valid || loaded.Passphrase.String != tt.cfg.Passphrase) {
t.Fatalf("expected passphrase to persist after reopen, got %+v", loaded.Passphrase)
}
})
}
}
func TestUpdateNodeWithoutForwardModePreservesExistingMode(t *testing.T) {
r, err := Open(":memory:")
if err != nil {
@@ -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")
+4
View File
@@ -82,6 +82,8 @@ export interface ForwardApiItem {
ipSpeedLimitName?: string;
maxConn?: number;
proxyProtocol?: number;
proxyProtocolReceive?: number;
proxyProtocolSend?: number;
inx?: number;
[key: string]: unknown;
}
@@ -422,6 +424,8 @@ export interface ForwardMutationPayload {
ipSpeedId?: number | null;
maxConn?: number;
proxyProtocol?: number;
proxyProtocolReceive?: number;
proxyProtocolSend?: number;
}
export interface SpeedLimitMutationPayload {
+69 -20
View File
@@ -130,6 +130,8 @@ interface Forward {
ipSpeedId?: number | null;
ipSpeedLimitName?: string;
proxyProtocol?: number;
proxyProtocolReceive?: number;
proxyProtocolSend?: number;
}
interface Tunnel {
@@ -169,6 +171,8 @@ interface ForwardForm {
ipSpeedId: number | null;
maxConn?: number;
proxyProtocol?: number;
proxyProtocolReceive?: number;
proxyProtocolSend?: number;
}
interface ForwardUserGroup {
@@ -600,6 +604,16 @@ const mapForwardApiItems = (items: ForwardApiItem[]): Forward[] => {
typeof forward.proxyProtocol === "number"
? forward.proxyProtocol
: undefined,
proxyProtocolReceive:
typeof forward.proxyProtocolReceive === "number"
? forward.proxyProtocolReceive
: 0,
proxyProtocolSend:
typeof forward.proxyProtocolSend === "number"
? forward.proxyProtocolSend
: typeof forward.proxyProtocol === "number"
? forward.proxyProtocol
: 0,
serviceRunning: forward.status === 1,
}));
};
@@ -1337,6 +1351,8 @@ export default function ForwardPage() {
ipSpeedId: null,
maxConn: 0,
proxyProtocol: 0,
proxyProtocolReceive: 0,
proxyProtocolSend: 0,
});
const [inIpTouched, setInIpTouched] = useState(false);
@@ -2128,6 +2144,8 @@ export default function ForwardPage() {
ipMaxConn: 0,
ipSpeedId: null,
proxyProtocol: 0,
proxyProtocolReceive: 0,
proxyProtocolSend: 0,
});
setErrors({});
setModalOpen(true);
@@ -2152,6 +2170,8 @@ export default function ForwardPage() {
ipSpeedId: normalizeSpeedId(forward.ipSpeedId),
maxConn: forward.maxConn ?? 0,
proxyProtocol: forward.proxyProtocol ?? 0,
proxyProtocolReceive: forward.proxyProtocolReceive ?? 0,
proxyProtocolSend: forward.proxyProtocolSend ?? forward.proxyProtocol ?? 0,
});
setErrors({});
setModalOpen(true);
@@ -2285,7 +2305,9 @@ export default function ForwardPage() {
ipMaxConn: form.ipMaxConn,
...(isAdmin ? { ipSpeedId: normalizedIPSpeedId } : {}),
maxConn: form.maxConn,
proxyProtocol: form.proxyProtocol,
proxyProtocol: form.proxyProtocolSend,
proxyProtocolReceive: form.proxyProtocolReceive,
proxyProtocolSend: form.proxyProtocolSend,
};
res = await updateForward(updateData);
@@ -2301,7 +2323,9 @@ export default function ForwardPage() {
ipMaxConn: form.ipMaxConn,
...(isAdmin ? { ipSpeedId: normalizedIPSpeedId } : {}),
maxConn: form.maxConn,
proxyProtocol: form.proxyProtocol,
proxyProtocol: form.proxyProtocolSend,
proxyProtocolReceive: form.proxyProtocolReceive,
proxyProtocolSend: form.proxyProtocolSend,
};
res = await createForward(createData);
@@ -4968,25 +4992,50 @@ export default function ForwardPage() {
setForm((prev) => ({ ...prev, ipMaxConn: value }));
}}
/>
<Select
description="启用 PROXY protocol,用于透传客户端真实 IP"
label="Proxy Protocol"
placeholder="禁用"
selectedKeys={[String(form.proxyProtocol || 0)]}
variant="bordered"
onSelectionChange={(keys) => {
const selectedKey = Array.from(keys)[0] as string;
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
<Select
description="入口监听接收 PROXY protocol,用于读取上游传入的真实客户端 IP。"
label="Proxy Protocol 接收"
placeholder="禁用"
selectedKeys={[
String(form.proxyProtocolReceive || 0),
]}
variant="bordered"
onSelectionChange={(keys) => {
const selectedKey = Array.from(keys)[0] as string;
setForm((prev) => ({
...prev,
proxyProtocol: Number(selectedKey),
}));
}}
>
<SelectItem key="0">禁用</SelectItem>
<SelectItem key="1">Version 1</SelectItem>
<SelectItem key="2">Version 2</SelectItem>
</Select>
setForm((prev) => ({
...prev,
proxyProtocolReceive: Number(selectedKey),
}));
}}
>
<SelectItem key="0">禁用</SelectItem>
<SelectItem key="1">Version 1</SelectItem>
<SelectItem key="2">Version 2</SelectItem>
</Select>
<Select
description="连接目标地址时发送 PROXY protocol,用于向下游透传客户端真实 IP。"
label="Proxy Protocol 发送"
placeholder="禁用"
selectedKeys={[String(form.proxyProtocolSend || 0)]}
variant="bordered"
onSelectionChange={(keys) => {
const selectedKey = Array.from(keys)[0] as string;
const proxyProtocolSend = Number(selectedKey);
setForm((prev) => ({
...prev,
proxyProtocol: proxyProtocolSend,
proxyProtocolSend,
}));
}}
>
<SelectItem key="0">禁用</SelectItem>
<SelectItem key="1">Version 1</SelectItem>
<SelectItem key="2">Version 2</SelectItem>
</Select>
</div>
{isAdmin && (
<Select
label="规则限速"