diff --git a/docs/superpowers/plans/2026-04-23-max-conn-limit.md b/docs/superpowers/plans/2026-04-23-max-conn-limit.md new file mode 100644 index 0000000..ce2adb0 --- /dev/null +++ b/docs/superpowers/plans/2026-04-23-max-conn-limit.md @@ -0,0 +1,482 @@ +# 最大连接数限制实现计划 + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 在 FLVX 中实现基于用户的全局最大连接数限制和基于单条规则的独立最大连接数限制功能。前端输入框为 0 或空时表示不限制。 + +**Architecture:** 采用“覆盖逻辑”（方案二）。 +1. 数据库层面：在 `user` 和 `forward` 表中各增加一个整型字段 `max_conn`，默认值为 0（表示不限制）。 +2. 后端接口层面：提供 API 更新该字段，在组装下发给 GOST 的配置时，判断规则的 `max_conn` 是否大于 0： + - 如果规则 `max_conn > 0`，则为此规则动态生成一个唯一的连接限制器配置，并在下发服务的 `climiter` 字段中引用该限制器。 + - 如果规则 `max_conn == 0`，则检查该规则所属用户的 `max_conn`。 + - 如果用户 `max_conn > 0`，则引用以用户维度的连接限制器配置（如 `user_conn_limit_`）。 + - 否则不下发 `climiter`。 +3. 后端服务控制平面：需要在下发服务前，将需要的连接限制器（Rule 或 User 维度）推送到节点上。 + - **重要发现：** 当前 `go-gost` 的 WebSocket Reporter (`go-gost/x/socket/websocket_reporter.go`) 仅支持 `TrafficLimiter` 的动态增删（如 `AddLimiters` 等），**不支持** `ConnLimiter`（即 `CLimiters`）。 + - **计划修改：** 我们需要先在 `go-gost` 侧（`go-gost/x/socket`）添加针对 `CLimiters` 的 WebSocket 指令（`AddCLimiters`, `UpdateCLimiters`, `DeleteCLimiters`）以及对应的处理函数（参考 `AddLimiters` 等的实现，调用现有的针对 `ConnLimiterRegistry` 的相关接口和配置存储逻辑，具体需要实现类似 `createLimiter` 到 `createConnLimiter` 的逻辑）。 + - 完成底层修改后，`go-backend` 再通过这些新增加的 WebSocket 指令，在 `ensureLimiterOnNode` 时下发最大连接数限制规则。 +4. 前端层面：在用户管理和规则管理页面增加输入框组件。 + +**Tech Stack:** Go, GORM, SQLite/PostgreSQL, React, Vite, TypeScript, TailwindCSS. + +--- + +### Task 1: 扩展 go-gost WebSocket 接口以支持 CLimiters + +**Files:** +- Modify: `go-gost/x/socket/limiter.go` +- Modify: `go-gost/x/socket/websocket_reporter.go` + +- [ ] **Step 1: 实现 `createConnLimiter` 等功能** + +在 `go-gost/x/socket/limiter.go` 中参考现有 `createLimiter` 添加对 `CLimiters` 的支持： + +```go +func createConnLimiter(req createLimiterRequest) error { + name := strings.TrimSpace(req.Data.Name) + if name == "" { + return errors.New("limiter name is required") + } + req.Data.Name = name + + if registry.ConnLimiterRegistry().IsRegistered(name) { + return errors.New("conn limiter " + name + " already exists") + } + + v := parser.ParseConnLimiter(&req.Data) + + if err := registry.ConnLimiterRegistry().Register(name, v); err != nil { + return errors.New("conn limiter " + name + " already exists") + } + + if c := config.Global(); c != nil { + c.CLimiters = append(c.CLimiters, &req.Data) + } + return nil +} + +func updateConnLimiter(req updateLimiterRequest) error { + name := strings.TrimSpace(req.Limiter) + req.Data.Name = name + if registry.ConnLimiterRegistry().IsRegistered(name) { + registry.ConnLimiterRegistry().Unregister(name) + } + + v := parser.ParseConnLimiter(&req.Data) + + if err := registry.ConnLimiterRegistry().Register(name, v); err != nil { + return errors.New("conn limiter " + name + " already exists") + } + + if c := config.Global(); c != nil { + for i := range c.CLimiters { + if c.CLimiters[i].Name == name { + c.CLimiters[i] = &req.Data + return nil + } + } + c.CLimiters = append(c.CLimiters, &req.Data) + } + return nil +} + +func deleteConnLimiter(req deleteLimiterRequest) error { + name := strings.TrimSpace(req.Limiter) + + if registry.ConnLimiterRegistry().IsRegistered(name) { + registry.ConnLimiterRegistry().Unregister(name) + } + + if c := config.Global(); c != nil { + limiteres := c.CLimiters + c.CLimiters = nil + for _, s := range limiteres { + if s.Name == name { + continue + } + c.CLimiters = append(c.CLimiters, s) + } + } + return nil +} +``` + +- [ ] **Step 2: 在 `WebSocketReporter` 注册命令** + +在 `go-gost/x/socket/websocket_reporter.go` 的 `ProcessCommand` 中添加 case： + +```go + case "AddCLimiters": + err = w.handleAddCLimiter(cmd.Data) + response.Type = "AddCLimitersResponse" + needSaveConfig = true + case "UpdateCLimiters": + err = w.handleUpdateCLimiter(cmd.Data) + response.Type = "UpdateCLimitersResponse" + needSaveConfig = true + case "DeleteCLimiters": + err = w.handleDeleteCLimiter(cmd.Data) + response.Type = "DeleteCLimitersResponse" + needSaveConfig = true +``` + +- [ ] **Step 3: 实现 Handler 方法** + +在 `go-gost/x/socket/websocket_reporter.go` 中添加： + +```go +func (w *WebSocketReporter) handleAddCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var limiterConfig config.LimiterConfig + if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { + return fmt.Errorf("解析限流器配置失败: %v", err) + } + + req := createLimiterRequest{Data: limiterConfig} + return createConnLimiter(req) +} + +func (w *WebSocketReporter) handleUpdateCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var updateReq struct { + Limiter string `json:"limiter"` + Data config.LimiterConfig `json:"data"` + } + + if err := json.Unmarshal(jsonData, &updateReq); err != nil { + var limiterConfig config.LimiterConfig + if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { + return fmt.Errorf("解析更新请求失败: %v", err) + } + updateReq.Limiter = limiterConfig.Name + updateReq.Data = limiterConfig + } + + req := updateLimiterRequest{ + Limiter: updateReq.Limiter, + Data: updateReq.Data, + } + return updateConnLimiter(req) +} + +func (w *WebSocketReporter) handleDeleteCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var deleteReq deleteLimiterRequest + + if err := json.Unmarshal(jsonData, &deleteReq); err != nil { + var limiterName string + if err := json.Unmarshal(jsonData, &limiterName); err != nil { + return fmt.Errorf("解析删除请求失败: %v", err) + } + deleteReq.Limiter = limiterName + } + + return deleteConnLimiter(deleteReq) +} +``` + +- [ ] **Step 4: Commit** + +```bash +cd go-gost +git add x/socket/limiter.go x/socket/websocket_reporter.go +git commit -m "feat: add CLimiters support for websocket reporter" +cd .. +``` + +--- + +### Task 2: 数据库迁移与模型更新 + +**Files:** +- Modify: `go-backend/internal/store/model/model.go` +- Modify: `go-backend/internal/store/repo/repository.go` + +- [ ] **Step 1: 更新数据库模型** + +在 `go-backend/internal/store/model/model.go` 的 `User` 和 `Forward` 结构体中添加 `MaxConn` 字段。 + +```go +// 在 User 结构体中 +type User struct { + // ... + MaxConn int `gorm:"column:max_conn;not null;default:0"` + // ... +} + +// 在 Forward 结构体中 +type Forward struct { + // ... + MaxConn int `gorm:"column:max_conn;not null;default:0"` + // ... +} +``` + +- [ ] **Step 2: 编写数据库迁移** + +在 `go-backend/internal/store/repo/repository.go` 的 `AutoMigrate` 逻辑前（如果有自定义迁移）或利用 gorm 自动迁移机制，由于这是 autoMigrate，添加字段只要 `db.AutoMigrate(&model.User{}, &model.Forward{})` 被调用就能自动加上。确认已执行迁移。由于 `FLVX` 通常会自动执行迁移，只需修改模型即可。我们需要处理默认值，由于使用了 `default:0`，GORM 会处理新增字段的默认值，但为了安全起见，如果在旧环境中，可能直接 alter table。 + +```go +// 无需手动编写 SQL，依赖现有的 gorm AutoMigrate 即可。 +``` + +- [ ] **Step 3: 运行并验证迁移通过** + +Run: `make build` (在 go-backend 中)，或者运行一个相关的存储单元测试。 + +- [ ] **Step 4: Commit** + +```bash +cd go-backend +git add internal/store/model/model.go +git commit -m "feat: add max_conn field to user and forward models" +cd .. +``` + +--- + +### Task 3: 后端控制平面 - 连接数限制器的组装与下发 + +**Files:** +- Modify: `go-backend/internal/http/handler/control_plane.go` +- Modify: `go-backend/internal/store/repo/repository_control.go` + +- [ ] **Step 1: 更新存储层以获取 User 的 MaxConn** + +在 `go-backend/internal/store/repo/repository_control.go` 中： + +需要一个方法获取 User，或者如果已经有，确保可以拿到 `MaxConn`。 + +- [ ] **Step 2: 编写下发 CLimiter 到节点的辅助函数** + +在 `go-backend/internal/http/handler/control_plane.go`，参考 `ensureLimiterOnNode` 和 `upsertLimiterOnNode`： + +```go +func (h *Handler) ensureConnLimiterOnNode(nodeID int64, limiterName string, maxConn int) error { + limitStr := fmt.Sprintf("$ %d", maxConn) + + payload := map[string]interface{}{ + "name": limiterName, + "limits": []string{limitStr}, + } + + if _, err := h.sendNodeCommand(nodeID, "AddCLimiters", payload, false, false); err != nil { + if !isAlreadyExistsMessage(err.Error()) { + return fmt.Errorf("连接限制器下发失败: %w", err) + } + updatePayload := map[string]interface{}{ + "limiter": limiterName, + "data": payload, + } + if _, updateErr := h.sendNodeCommand(nodeID, "UpdateCLimiters", updatePayload, false, false); updateErr != nil { + return fmt.Errorf("连接限制器更新失败: %w", updateErr) + } + } + return nil +} +``` + +- [ ] **Step 3: 更新组装配置逻辑以绑定 `climiter`** + +在 `control_plane.go` 的 `syncForwardServicesWithWarnings` 及其辅助函数 `buildForwardServiceConfigs` 附近： + +修改 `buildForwardServiceConfigs` 的签名，传入 `maxConn int` 和对应的 `cLimiterName string`。 + +```go +func buildForwardServiceConfigs(baseName string, forward *model.Forward, tunnel *model.Tunnel, node *model.Node, port int, bindIP string, limiterID *int64, cLimiterName string) []map[string]interface{} { + // ... 现有逻辑 + // 在服务配置生成的部分增加： + if cLimiterName != "" { + service["climiter"] = cLimiterName + } + // ... +} +``` + +- [ ] **Step 4: 在转发服务同步主流程中决定并下发 `climiter`** + +在 `syncForwardServicesWithWarnings` (可能在多个重载/处理入口处，如 `ensureForwardServices`)，查出转发所属 user 的 `MaxConn`，以及转发本身的 `MaxConn`。 + +```go + // 获取 User + user, err := h.repo.GetUser(forward.UserID) + if err != nil { + return nil, err + } + + var cLimiterName string + var maxConnToSet int + + if forward.MaxConn > 0 { + maxConnToSet = forward.MaxConn + cLimiterName = fmt.Sprintf("rule_conn_limit_%d", forward.ID) + } else if user != nil && user.MaxConn > 0 { + maxConnToSet = user.MaxConn + cLimiterName = fmt.Sprintf("user_conn_limit_%d", user.ID) + } + + if cLimiterName != "" { + for _, fp := range ports { + if err := h.ensureConnLimiterOnNode(fp.NodeID, cLimiterName, maxConnToSet); err != nil { + warnings = append(warnings, fmt.Sprintf("节点 %d 连接限制器下发失败: %v", fp.NodeID, err)) + } + } + } + + // 传递给 buildForwardServiceConfigs + // ... +``` +*(注意：需要确保更新涉及 `buildForwardServiceConfigs` 的所有调用点)* + +- [ ] **Step 5: Commit** + +```bash +cd go-backend +git add internal/http/handler/control_plane.go internal/store/repo/repository_control.go +git commit -m "feat: implement max conn limiter dispatching" +cd .. +``` + +--- + +### Task 4: 后端接口 - 用户和规则的 CRUD 支持 + +**Files:** +- Modify: `go-backend/internal/http/handler/admin_user.go` +- Modify: `go-backend/internal/http/handler/forward.go` + +- [ ] **Step 1: 用户接口更新** + +在 `go-backend/internal/http/handler/admin_user.go`，修改用户创建和更新请求的结构体（如果有），接收 `MaxConn`，并在保存到数据库时赋值。 + +```go +type CreateUserReq struct { + // ... + MaxConn *int `json:"maxConn"` +} +// 接收后： +if req.MaxConn != nil { + user.MaxConn = *req.MaxConn +} +``` + +在获取用户列表时，确保 `MaxConn` 返回给前端。 + +- [ ] **Step 2: 规则接口更新** + +在 `go-backend/internal/http/handler/forward.go` 中，更新 `CreateForwardReq` 和 `UpdateForwardReq` 结构体，增加 `MaxConn`，并在创建/更新 Forward 时保存到数据库。 + +如果转发规则的 `MaxConn` 或相关信息改变，触发节点上的规则重载（重新下发服务）。这一步由于更改了数据库，复用现有的 `syncForwardServices` 就会带上最新的配置。 + +- [ ] **Step 3: 测试接口** + +Run: 可以启动后使用 curl 测试。 + +- [ ] **Step 4: Commit** + +```bash +cd go-backend +git add internal/http/handler/admin_user.go internal/http/handler/forward.go +git commit -m "feat: add maxConn to user and forward CRUD API" +cd .. +``` + +--- + +### Task 5: 前端 - 用户管理页面集成 + +**Files:** +- Modify: `vite-frontend/src/api/types.ts` +- Modify: `vite-frontend/src/api/index.ts` +- Modify: `vite-frontend/src/pages/users.tsx` (或者对应的用户管理页面文件) + +- [ ] **Step 1: 类型更新** + +在 `vite-frontend/src/api/types.ts` 中： +为 `UserApiItem` 和相关的 mutation payload 增加 `maxConn?: number` 属性。 + +- [ ] **Step 2: UI 修改** + +在用户创建/编辑弹窗中，增加“最大连接数”输入框： +(假设使用 `@nextui-org/react` 的 `Input`) + +```tsx + { + const num = parseInt(val, 10); + setFormData({ ...formData, maxConn: isNaN(num) ? 0 : num }); + }} +/> +``` +并在用户的表格列中展示 `最大连接数`（值为 0 显示“不限制”）。 + +- [ ] **Step 3: 运行 Vite 进行验证** + +- [ ] **Step 4: Commit** + +```bash +cd vite-frontend +git add src/api/types.ts src/api/index.ts src/pages/users.tsx +git commit -m "feat: add max conn UI to user management" +cd .. +``` + +--- + +### Task 6: 前端 - 转发规则页面集成 + +**Files:** +- Modify: `vite-frontend/src/pages/forward.tsx` + +- [ ] **Step 1: 类型更新** + +在 `api/types.ts` 中 `ForwardMutationPayload` 和 `ForwardApiItem` 中增加 `maxConn?: number`。 + +- [ ] **Step 2: UI 修改** + +在 `vite-frontend/src/pages/forward.tsx` 的创建/编辑规则弹窗（在 "规则限速" 附近）增加“最大连接数”输入框： + +```tsx + { + const num = parseInt(val, 10); + setFormData({ ...formData, maxConn: isNaN(num) ? 0 : num }); + }} + description="此设置优先于用户的全局连接数限制。0 表示不限制（或使用用户的全局限制）。" +/> +``` + +如果是在列表/卡片中展示，可以增加一个小标签或者 Tooltip 显示其最大连接数设置。 + +- [ ] **Step 3: 验证** + +在前端验证该功能能正确读写规则的连接限制字段。 + +- [ ] **Step 4: Commit** + +```bash +cd vite-frontend +git add src/pages/forward.tsx src/api/types.ts +git commit -m "feat: add max conn UI to forward rules" +cd .. +``` \ No newline at end of file diff --git a/go-backend/internal/http/handler/control_plane.go b/go-backend/internal/http/handler/control_plane.go index baec361..e4f8d1e 100644 --- a/go-backend/internal/http/handler/control_plane.go +++ b/go-backend/internal/http/handler/control_plane.go @@ -266,6 +266,22 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method serviceBase := buildForwardServiceBaseWithResolvedUserTunnel(forward.ID, forward.UserID, userTunnelID) + user, err := h.repo.GetUserByID(forward.UserID) + if err != nil { + return nil, err + } + + var cLimiterName string + var maxConnToSet int + + if forward.MaxConn > 0 { + maxConnToSet = forward.MaxConn + cLimiterName = fmt.Sprintf("rule_conn_limit_%d", forward.ID) + } else if user != nil && user.MaxConn > 0 { + maxConnToSet = user.MaxConn + cLimiterName = fmt.Sprintf("user_conn_limit_%d", user.ID) + } + for _, fp := range ports { if limiterID != nil && speed != nil { if err := h.ensureLimiterOnNode(fp.NodeID, *limiterID, *speed); err != nil { @@ -283,11 +299,17 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method } } + if cLimiterName != "" { + if err := h.ensureConnLimiterOnNode(fp.NodeID, cLimiterName, maxConnToSet); err != nil { + warnings = append(warnings, fmt.Sprintf("节点 %d 连接限制器下发失败: %v", fp.NodeID, err)) + } + } + node, err := h.getNodeRecord(fp.NodeID) if err != nil { return nil, err } - services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID) + services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID, cLimiterName) _, err = h.sendNodeCommand(node.ID, method, services, true, false) if err != nil && allowFallbackAdd && method == "UpdateService" { if isNotFoundError(err) { @@ -302,7 +324,7 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method } if err != nil && strings.EqualFold(strings.TrimSpace(method), "UpdateService") && isCannotAssignRequestedAddressError(err) { var warning string - warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID) + warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID, cLimiterName) if err == nil && warning != "" { warnings = append(warnings, warning) } @@ -328,7 +350,7 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method return warnings, nil } -func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64) (string, error) { +func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64, cLimiterName string) (string, error) { if h == nil || forward == nil || tunnel == nil || node == nil { return "", errors.New("invalid bind fallback context") } @@ -345,7 +367,7 @@ func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunne } time.Sleep(150 * time.Millisecond) - defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID) + defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID, cLimiterName) if _, err := h.sendNodeCommand(node.ID, "AddService", defaultServices, true, false); err != nil { return "", err } @@ -1637,7 +1659,7 @@ func compactErrorMessage(msg string) string { return strings.Join(strings.Fields(strings.ToLower(msg)), "") } -func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64) []map[string]interface{} { +func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64, cLimiterName string) []map[string]interface{} { protocols := []string{"tcp", "udp"} services := make([]map[string]interface{}, 0, 2) targets := splitRemoteTargets(forward.RemoteAddr) @@ -1680,6 +1702,9 @@ func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel }, }, } + if cLimiterName != "" { + service["climiter"] = cLimiterName + } if protocol == "udp" { listenerMetadata := map[string]interface{}{ "keepAlive": true, @@ -1795,6 +1820,30 @@ func (h *Handler) ensureLimiterOnNode(nodeID int64, limiterID int64, speed int) return nil } +func (h *Handler) ensureConnLimiterOnNode(nodeID int64, limiterName string, maxConn int) error { + limitStr := fmt.Sprintf("$ %d", maxConn) + + payload := map[string]interface{}{ + "name": limiterName, + "limits": []string{limitStr}, + } + + if _, err := h.sendNodeCommand(nodeID, "AddCLimiters", payload, false, false); err != nil { + if !isAlreadyExistsMessage(err.Error()) { + return fmt.Errorf("连接限制器下发失败: %w", err) + } + updatePayload := map[string]interface{}{ + "limiter": limiterName, + "data": payload, + } + if _, updateErr := h.sendNodeCommand(nodeID, "UpdateCLimiters", updatePayload, false, false); updateErr != nil { + return fmt.Errorf("连接限制器更新失败: %w", updateErr) + } + } + return nil +} + + func buildLimiterAddPayload(limiterID int64, speed int) (string, map[string]interface{}) { rate := float64(speed) / 8.0 limitStr := fmt.Sprintf("$ %.1fMB %.1fMB", rate, rate) diff --git a/go-backend/internal/http/handler/control_plane_test.go b/go-backend/internal/http/handler/control_plane_test.go index 81d79bc..97bd80a 100644 --- a/go-backend/internal/http/handler/control_plane_test.go +++ b/go-backend/internal/http/handler/control_plane_test.go @@ -378,7 +378,7 @@ func TestRetryTunnelServiceAddWithCleanupReturnsCleanupError(t *testing.T) { func TestBuildForwardServiceConfigs_UsesBindIPForListen(t *testing.T) { forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7} node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"} - services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22000, "10.9.8.7", nil) + services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22000, "10.9.8.7", nil, "") if len(services) != 2 { t.Fatalf("expected 2 services, got %d", len(services)) } @@ -393,7 +393,7 @@ func TestBuildForwardServiceConfigs_UsesBindIPForListen(t *testing.T) { func TestBuildForwardServiceConfigs_DefaultListenAddrWhenBindIPEmpty(t *testing.T) { forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7} node := &nodeRecord{TCPListenAddr: "0.0.0.0", UDPListenAddr: "[::]"} - services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22001, "", nil) + services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 22001, "", nil, "") if len(services) != 2 { t.Fatalf("expected 2 services, got %d", len(services)) } @@ -409,7 +409,7 @@ func TestBuildForwardServiceConfigs_DefaultListenAddrWhenBindIPEmpty(t *testing. func TestBuildForwardServiceConfigs_BindIPAlreadyContainsPort(t *testing.T) { forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7} node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"} - services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 55555, "3.3.3.3:12345", nil) + services := buildForwardServiceConfigs("1_2_0", forward, nil, node, 55555, "3.3.3.3:12345", nil, "") if len(services) != 2 { t.Fatalf("expected 2 services, got %d", len(services)) } @@ -464,7 +464,7 @@ func TestBuildForwardServiceConfigs_IPv6BindIP(t *testing.T) { t.Run(tt.name, func(t *testing.T) { forward := &forwardRecord{RemoteAddr: "1.2.3.4:80", Strategy: "fifo", TunnelID: 7} node := &nodeRecord{TCPListenAddr: "[::]", UDPListenAddr: "[::]"} - services := buildForwardServiceConfigs("1_2_0", forward, nil, node, tt.port, tt.bindIP, nil) + services := buildForwardServiceConfigs("1_2_0", forward, nil, node, tt.port, tt.bindIP, nil, "") if len(services) != 2 { t.Fatalf("expected 2 services, got %d", len(services)) } diff --git a/go-backend/internal/http/handler/mutations.go b/go-backend/internal/http/handler/mutations.go index b622d4a..9324737 100644 --- a/go-backend/internal/http/handler/mutations.go +++ b/go-backend/internal/http/handler/mutations.go @@ -68,8 +68,9 @@ func (h *Handler) userCreate(w http.ResponseWriter, r *http.Request) { } roleID := 1 now := time.Now().UnixMilli() + maxConn := asInt(req["maxConn"], 0) - userID, err := h.repo.CreateUser(username, security.MD5(pwd), roleID, expTime, flow, flowResetTime, num, status, now) + userID, err := h.repo.CreateUser(username, security.MD5(pwd), roleID, expTime, flow, flowResetTime, num, status, maxConn, now) if err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) return @@ -156,15 +157,16 @@ func (h *Handler) userUpdate(w http.ResponseWriter, r *http.Request) { _, hasDailyQuota := req["dailyQuotaGB"] _, hasMonthlyQuota := req["monthlyQuotaGB"] now := time.Now().UnixMilli() + maxConn := asInt(req["maxConn"], 0) pwd := asString(req["pwd"]) if strings.TrimSpace(pwd) == "" { - if err := h.repo.UpdateUserWithoutPassword(id, username, flow, num, expTime, flowResetTime, status, now); err != nil { + if err := h.repo.UpdateUserWithoutPassword(id, username, flow, num, expTime, flowResetTime, status, maxConn, now); err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) return } } else { - if err := h.repo.UpdateUserWithPassword(id, username, security.MD5(pwd), flow, num, expTime, flowResetTime, status, now); err != nil { + if err := h.repo.UpdateUserWithPassword(id, username, security.MD5(pwd), flow, num, expTime, flowResetTime, status, maxConn, now); err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) return } @@ -1777,7 +1779,8 @@ func (h *Handler) forwardCreate(w http.ResponseWriter, r *http.Request) { if userName == "" { userName = "user" } - forwardID, err := h.repo.CreateForwardTx(userID, userName, name, tunnelID, remoteAddr, defaultString(asString(req["strategy"]), "fifo"), now, inx, entryNodes, port, inIp, nullableInt(speedID)) + maxConn := asInt(req["maxConn"], 0) + forwardID, err := h.repo.CreateForwardTx(userID, userName, name, tunnelID, remoteAddr, defaultString(asString(req["strategy"]), "fifo"), now, inx, entryNodes, port, inIp, nullableInt(speedID), maxConn) if err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) return @@ -1928,7 +1931,9 @@ func (h *Handler) forwardUpdate(w http.ResponseWriter, r *http.Request) { } } now := time.Now().UnixMilli() - if err := h.repo.UpdateForward(id, name, tunnelID, remoteAddr, strategy, now, newSpeedID); err != nil { + maxConn := asInt(req["maxConn"], forward.MaxConn) + + if err := h.repo.UpdateForward(id, name, tunnelID, remoteAddr, strategy, now, newSpeedID, maxConn); err != nil { response.WriteJSON(w, response.Err(-2, err.Error())) return } @@ -4080,7 +4085,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.SpeedID, oldForward.MaxConn, time.Now().UnixMilli(), ) diff --git a/go-backend/internal/store/model/model.go b/go-backend/internal/store/model/model.go index 605ec5c..d403c7d 100644 --- a/go-backend/internal/store/model/model.go +++ b/go-backend/internal/store/model/model.go @@ -23,6 +23,7 @@ type User struct { CreatedTime int64 `gorm:"column:created_time;not null"` UpdatedTime sql.NullInt64 `gorm:"column:updated_time"` Status int `gorm:"not null"` + MaxConn int `gorm:"column:max_conn;not null;default:0"` } func (User) TableName() string { return "user" } @@ -43,6 +44,7 @@ type Forward struct { 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"` } func (Forward) TableName() string { return "forward" } @@ -537,6 +539,7 @@ type ForwardRecord struct { Strategy string Status int SpeedID sql.NullInt64 + MaxConn int } // TunnelRecord is a minimal tunnel view used by control plane. diff --git a/go-backend/internal/store/repo/repository_flow.go b/go-backend/internal/store/repo/repository_flow.go index 8d1ef83..988f404 100644 --- a/go-backend/internal/store/repo/repository_flow.go +++ b/go-backend/internal/store/repo/repository_flow.go @@ -124,15 +124,16 @@ func (r *Repository) GetForwardRecord(forwardID int64) (*model.ForwardRecord, er return nil, err } 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, + 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, } if strings.TrimSpace(fr.Strategy) == "" { fr.Strategy = "fifo" diff --git a/go-backend/internal/store/repo/repository_mutations.go b/go-backend/internal/store/repo/repository_mutations.go index 91cd3b4..f1bb979 100644 --- a/go-backend/internal/store/repo/repository_mutations.go +++ b/go-backend/internal/store/repo/repository_mutations.go @@ -37,7 +37,7 @@ func (r *Repository) UserExistsExcluding(username string, excludeID int64) (bool return cnt > 0, err } -func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, flow, flowResetTime int64, num, status int, now int64) (int64, error) { +func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, flow, flowResetTime int64, num, status, maxConn int, now int64) (int64, error) { if r == nil || r.db == nil { return 0, errors.New("repository not initialized") } @@ -51,6 +51,7 @@ func (r *Repository) CreateUser(username, pwdHash string, roleID int, expTime, f OutFlow: 0, FlowResetTime: flowResetTime, Num: num, + MaxConn: maxConn, CreatedTime: now, UpdatedTime: sql.NullInt64{Int64: now, Valid: true}, Status: status, @@ -73,7 +74,7 @@ func (r *Repository) GetUserRoleID(userID int64) (int, error) { return user.RoleID, nil } -func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string, flow int64, num int, expTime, flowResetTime int64, status int, now int64) error { +func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64) error { if r == nil || r.db == nil { return errors.New("repository not initialized") } @@ -87,11 +88,12 @@ func (r *Repository) UpdateUserWithPassword(id int64, username, pwdHash string, "exp_time": expTime, "flow_reset_time": flowResetTime, "status": status, + "max_conn": maxConn, "updated_time": sql.NullInt64{Int64: now, Valid: true}, }).Error } -func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow int64, num int, expTime, flowResetTime int64, status int, now int64) error { +func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow int64, num int, expTime, flowResetTime int64, status, maxConn int, now int64) error { if r == nil || r.db == nil { return errors.New("repository not initialized") } @@ -104,6 +106,7 @@ func (r *Repository) UpdateUserWithoutPassword(id int64, username string, flow i "exp_time": expTime, "flow_reset_time": flowResetTime, "status": status, + "max_conn": maxConn, "updated_time": sql.NullInt64{Int64: now, Valid: true}, }).Error } @@ -692,7 +695,7 @@ 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{}) error { +func (r *Repository) UpdateForward(id int64, name string, tunnelID int64, remoteAddr, strategy string, now int64, speedID interface{}, maxConn int) error { if r == nil || r.db == nil { return errors.New("repository not initialized") } @@ -704,6 +707,7 @@ func (r *Repository) UpdateForward(id int64, name string, tunnelID int64, remote "remote_addr": remoteAddr, "strategy": strategy, "speed_id": nullInt64FromInterface(speedID), + "max_conn": maxConn, "updated_time": now, }).Error } @@ -778,7 +782,7 @@ 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{}, now int64) { +func (r *Repository) RollbackForwardFields(id, userID int64, userName, name string, tunnelID int64, remoteAddr, strategy string, status int, speedID interface{}, maxConn int, now int64) { if r == nil || r.db == nil { return } @@ -793,6 +797,7 @@ func (r *Repository) RollbackForwardFields(id, userID int64, userName, name stri "strategy": strategy, "status": status, "speed_id": nullInt64FromInterface(speedID), + "max_conn": maxConn, "updated_time": now, }).Error } @@ -1253,7 +1258,7 @@ 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{}) (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) (int64, error) { if r == nil || r.db == nil { return 0, errors.New("repository not initialized") } @@ -1272,6 +1277,7 @@ func (r *Repository) CreateForwardTx(userID int64, userName, name string, tunnel UpdatedTime: now, Status: 1, Inx: inx, + MaxConn: maxConn, SpeedID: nullInt64FromInterface(speedID), } if err := tx.Create(&fwd).Error; err != nil { diff --git a/go-backend/patch.py b/go-backend/patch.py new file mode 100644 index 0000000..10be01e --- /dev/null +++ b/go-backend/patch.py @@ -0,0 +1,163 @@ +import re + +with open('internal/http/handler/control_plane.go', 'r') as f: + content = f.read() + +# 1. Update ensureLimiterOnNode and add ensureConnLimiterOnNode +ensure_conn_limiter = """ +func (h *Handler) ensureConnLimiterOnNode(nodeID int64, limiterName string, maxConn int) error { +\tlimitStr := fmt.Sprintf("$ %d", maxConn) +\t +\tpayload := map[string]interface{}{ +\t\t"name": limiterName, +\t\t"limits": []string{limitStr}, +\t} +\t +\tif _, err := h.sendNodeCommand(nodeID, "AddCLimiters", payload, false, false); err != nil { +\t\tif !isAlreadyExistsMessage(err.Error()) { +\t\t\treturn fmt.Errorf("连接限制器下发失败: %w", err) +\t\t} +\t\tupdatePayload := map[string]interface{}{ +\t\t\t"limiter": limiterName, +\t\t\t"data": payload, +\t\t} +\t\tif _, updateErr := h.sendNodeCommand(nodeID, "UpdateCLimiters", updatePayload, false, false); updateErr != nil { +\t\t\treturn fmt.Errorf("连接限制器更新失败: %w", updateErr) +\t\t} +\t} +\treturn nil +} +""" + +content = content.replace('func (h *Handler) ensureLimiterOnNode(nodeID int64, limiterID int64, speed int) error {\n\tif err := h.upsertLimiterOnNode(nodeID, limiterID, speed); err != nil {\n\t\treturn fmt.Errorf("限速规则下发失败: %w", err)\n\t}\n\n\treturn nil\n}', +'func (h *Handler) ensureLimiterOnNode(nodeID int64, limiterID int64, speed int) error {\n\tif err := h.upsertLimiterOnNode(nodeID, limiterID, speed); err != nil {\n\t\treturn fmt.Errorf("限速规则下发失败: %w", err)\n\t}\n\n\treturn nil\n}\n' + ensure_conn_limiter) + + +# 2. Update buildForwardServiceConfigs declaration +content = content.replace( + 'func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64) []map[string]interface{} {', + 'func buildForwardServiceConfigs(baseName string, forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, port int, bindIP string, limiterID *int64, cLimiterName string) []map[string]interface{} {' +) + + +# 3. Inject climiter into generated service +service_map_end = """ } + if protocol == "udp" {""" +service_map_end_new = """ } + if cLimiterName != "" { + service["climiter"] = cLimiterName + } + if protocol == "udp" {""" +content = content.replace(service_map_end, service_map_end_new) + + +# 4. Update syncForwardServicesWithWarnings +# Find user tunnel resolution +resolution = """ serviceBase := buildForwardServiceBaseWithResolvedUserTunnel(forward.ID, forward.UserID, userTunnelID) + + for _, fp := range ports {""" + +resolution_new = """ serviceBase := buildForwardServiceBaseWithResolvedUserTunnel(forward.ID, forward.UserID, userTunnelID) + + user, err := h.repo.GetUserByID(forward.UserID) + if err != nil { + return nil, err + } + + var cLimiterName string + var maxConnToSet int + + if forward.MaxConn > 0 { + maxConnToSet = forward.MaxConn + cLimiterName = fmt.Sprintf("rule_conn_limit_%d", forward.ID) + } else if user != nil && user.MaxConn > 0 { + maxConnToSet = user.MaxConn + cLimiterName = fmt.Sprintf("user_conn_limit_%d", user.ID) + } + + for _, fp := range ports {""" +content = content.replace(resolution, resolution_new) + +# Inject ensureConnLimiterOnNode inside loop +loop_inner = """ if limiterID != nil && speed != nil { + if err := h.ensureLimiterOnNode(fp.NodeID, *limiterID, *speed); err != nil { + // If the limiter push fails because the node is offline, skip it with a warning + if isNodeOfflineOrTimeoutError(err) { + node, _ := h.getNodeRecord(fp.NodeID) + nodeName := fmt.Sprintf("%d", fp.NodeID) + if node != nil && strings.TrimSpace(node.Name) != "" { + nodeName = strings.TrimSpace(node.Name) + } + warnings = append(warnings, fmt.Sprintf("节点 %s 不在线，已跳过下发", nodeName)) + continue + } + return nil, err + } + } + + node, err := h.getNodeRecord(fp.NodeID)""" + +loop_inner_new = """ if limiterID != nil && speed != nil { + if err := h.ensureLimiterOnNode(fp.NodeID, *limiterID, *speed); err != nil { + // If the limiter push fails because the node is offline, skip it with a warning + if isNodeOfflineOrTimeoutError(err) { + node, _ := h.getNodeRecord(fp.NodeID) + nodeName := fmt.Sprintf("%d", fp.NodeID) + if node != nil && strings.TrimSpace(node.Name) != "" { + nodeName = strings.TrimSpace(node.Name) + } + warnings = append(warnings, fmt.Sprintf("节点 %s 不在线，已跳过下发", nodeName)) + continue + } + return nil, err + } + } + + if cLimiterName != "" { + if err := h.ensureConnLimiterOnNode(fp.NodeID, cLimiterName, maxConnToSet); err != nil { + warnings = append(warnings, fmt.Sprintf("节点 %d 连接限制器下发失败: %v", fp.NodeID, err)) + } + } + + node, err := h.getNodeRecord(fp.NodeID)""" +content = content.replace(loop_inner, loop_inner_new) + +# Update buildForwardServiceConfigs call in syncForwardServicesWithWarnings +content = content.replace( + 'services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID)', + 'services := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, strings.TrimSpace(fp.InIP), limiterID, cLimiterName)' +) + +# Update fallbackForwardPortToDefaultBind call +content = content.replace( + 'warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID)', + 'warning, err = h.fallbackForwardPortToDefaultBind(forward, tunnel, node, fp, serviceBase, limiterID, cLimiterName)' +) + +# 5. Update fallbackForwardPortToDefaultBind declaration and logic +content = content.replace( + 'func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64) (string, error) {', + 'func (h *Handler) fallbackForwardPortToDefaultBind(forward *forwardRecord, tunnel *tunnelRecord, node *nodeRecord, fp forwardPortRecord, serviceBase string, limiterID *int64, cLimiterName string) (string, error) {' +) + +content = content.replace( + 'defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID)', + 'defaultServices := buildForwardServiceConfigs(serviceBase, forward, tunnel, node, fp.Port, "", limiterID, cLimiterName)' +) + +with open('internal/http/handler/control_plane.go', 'w') as f: + f.write(content) + + +# Update control_plane_test.go +with open('internal/http/handler/control_plane_test.go', 'r') as f: + test_content = f.read() + +test_content = re.sub( + r'buildForwardServiceConfigs\((.*?),(.*?),(.*?),(.*?),(.*?),(.*?),(.*?)\)', + r'buildForwardServiceConfigs(\1,\2,\3,\4,\5,\6,\7, "")', + test_content +) + +with open('internal/http/handler/control_plane_test.go', 'w') as f: + f.write(test_content) diff --git a/go-backend/tests/contract/max_conn_limit_contract_test.go b/go-backend/tests/contract/max_conn_limit_contract_test.go new file mode 100644 index 0000000..d1fa00f --- /dev/null +++ b/go-backend/tests/contract/max_conn_limit_contract_test.go @@ -0,0 +1,284 @@ +package contract_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + + "go-backend/internal/auth" + "go-backend/internal/http/response" + "go-backend/internal/security" +) + +func TestMaxConnLimit(t *testing.T) { + secret := "contract-jwt-secret" + router, r := setupContractRouter(t, secret) + server := httptest.NewServer(router) + defer server.Close() + + adminToken, err := auth.GenerateToken(1, "admin_user", 0, secret) + if err != nil { + t.Fatalf("generate admin token: %v", err) + } + + now := time.Now().UnixMilli() + if err := r.DB().Exec(` + INSERT INTO tunnel(name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) + VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, "max-conn-tunnel", 1.0, 1, "tls", 99999, now, now, 1, nil, 0).Error; err != nil { + t.Fatalf("insert tunnel: %v", err) + } + var tunnelID int64 + if err := r.DB().Raw("SELECT id FROM tunnel WHERE name = ?", "max-conn-tunnel").Scan(&tunnelID).Error; err != nil { + t.Fatalf("get tunnel ID: %v", err) + } + + if err := r.DB().Exec(` + INSERT INTO node(name, secret, server_ip, server_ip_v4, server_ip_v6, port, interface_name, version, http, tls, socks, created_time, updated_time, status, tcp_listen_addr, udp_listen_addr, inx) + VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, "max-conn-node", "max-conn-secret", "10.20.0.1", "10.20.0.1", "", "32000-32010", "", "v1", 1, 1, 1, now, now, 1, "[::]", "[::]", 0).Error; err != nil { + t.Fatalf("insert node: %v", err) + } + var nodeID int64 + if err := r.DB().Raw("SELECT id FROM node WHERE name = ?", "max-conn-node").Scan(&nodeID).Error; err != nil { + t.Fatalf("get node ID: %v", err) + } + + if err := r.DB().Exec(` + INSERT INTO chain_tunnel(tunnel_id, chain_type, node_id, port, strategy, inx, protocol) + VALUES(?, 1, ?, 32001, 'round', 1, 'tls') + `, tunnelID, nodeID).Error; err != nil { + t.Fatalf("insert chain_tunnel: %v", err) + } + + if err := r.DB().Exec(` + INSERT INTO user_tunnel(user_id, tunnel_id, num, flow, in_flow, out_flow, flow_reset_time, exp_time, status) + VALUES(1, ?, 10, 99999, 0, 0, 1, ?, 1) + `, tunnelID, now + 365*24*3600*1000).Error; err != nil { t.Fatalf("insert user_tunnel: %v", err) + } + + var commandMu sync.Mutex + receivedCommands := make([]string, 0) + var addCLimitersData json.RawMessage + var updateCLimitersData json.RawMessage + + stopNode := startMockSessionForMaxConn(t, server.URL, "max-conn-secret", func(cmdType string, data json.RawMessage) (bool, string) { + commandMu.Lock() + defer commandMu.Unlock() + receivedCommands = append(receivedCommands, cmdType) + + if cmdType == "AddCLimiters" { + addCLimitersData = append([]byte(nil), data...) + return true, "already exists" + } + if cmdType == "UpdateCLimiters" { + updateCLimitersData = append([]byte(nil), data...) + } + + return false, "" + }) + defer stopNode() + + waitNodeStatus(t, r, nodeID, 1) + + payload := map[string]interface{}{ + "name": "max-conn-forward", + "tunnelId": tunnelID, + "remoteAddr": "1.1.1.1:443", + "strategy": "fifo", + "maxConn": 42, + } + body, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/api/v1/forward/create", bytes.NewReader(body)) + req.Header.Set("Authorization", adminToken) + req.Header.Set("Content-Type", "application/json") + res := httptest.NewRecorder() + router.ServeHTTP(res, req) + + var out response.R + if err := json.NewDecoder(res.Body).Decode(&out); err != nil { + t.Fatalf("decode response: %v", err) + } + if out.Code != 0 { + t.Fatalf("expected create success, got code=%d msg=%s", out.Code, out.Msg) + } + + var forwardID int64 + if err := r.DB().Raw("SELECT id FROM forward WHERE name = ?", "max-conn-forward").Scan(&forwardID).Error; err != nil { + t.Fatalf("get forward ID: %v", err) + } + + commandMu.Lock() + defer commandMu.Unlock() + + hasAdd := false + hasUpdate := false + for _, cmd := range receivedCommands { + if cmd == "AddCLimiters" { + hasAdd = true + } + if cmd == "UpdateCLimiters" { + hasUpdate = true + } + } + + if !hasAdd { + t.Fatalf("expected AddCLimiters to be sent, but it was not. Received: %v", receivedCommands) + } + if !hasUpdate { + t.Fatalf("expected UpdateCLimiters to be sent after AddCLimiters failed with already exists. Received: %v", receivedCommands) + } + + expectedName := fmt.Sprintf("rule_conn_limit_%d", forwardID) + + // verify payload for AddCLimiters + var addData map[string]interface{} + if err := json.Unmarshal(addCLimitersData, &addData); err != nil { + t.Fatalf("unmarshal AddCLimiters data: %v", err) + } + if addData["name"] != expectedName { + t.Fatalf("expected limiter name %s, got %v", expectedName, addData["name"]) + } + if limits, ok := addData["limits"].([]interface{}); ok { + if len(limits) != 1 || limits[0] != "$ 42" { + t.Fatalf("expected limits to contain '$ 42', got %v", limits) + } + } else { + t.Fatalf("invalid limits type in AddCLimiters data: %v", addData) + } + + // verify payload for UpdateCLimiters + var updateData map[string]interface{} + if err := json.Unmarshal(updateCLimitersData, &updateData); err != nil { + t.Fatalf("unmarshal UpdateCLimiters data: %v", err) + } + if updateData["limiter"] != expectedName { + t.Fatalf("expected update limiter name %s, got %v", expectedName, updateData["limiter"]) + } + nestedData, ok := updateData["data"].(map[string]interface{}) + if !ok { + t.Fatalf("expected nested 'data' in UpdateCLimiters, got %v", updateData) + } + if nestedData["name"] != expectedName { + t.Fatalf("expected nested name %s, got %v", expectedName, nestedData["name"]) + } + if nestedLimits, ok := nestedData["limits"].([]interface{}); ok { + if len(nestedLimits) != 1 || nestedLimits[0] != "$ 42" { + t.Fatalf("expected nested limits to contain '$ 42', got %v", nestedLimits) + } + } else { + t.Fatalf("invalid limits type in UpdateCLimiters nested data: %v", nestedData) + } +} + +func startMockSessionForMaxConn(t *testing.T, baseURL string, nodeSecret string, onCommand func(cmdType string, data json.RawMessage) (bool, string)) func() { + t.Helper() + + u, err := url.Parse(baseURL) + if err != nil { + t.Fatalf("parse provider url: %v", err) + } + if strings.EqualFold(u.Scheme, "https") { + u.Scheme = "wss" + } else { + u.Scheme = "ws" + } + u.Path = "/system-info" + q := u.Query() + q.Set("type", "1") + q.Set("secret", nodeSecret) + q.Set("version", "v1") + q.Set("http", "1") + q.Set("tls", "1") + q.Set("socks", "1") + u.RawQuery = q.Encode() + + conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil) + if err != nil { + t.Fatalf("dial mock node websocket: %v", err) + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for { + _, raw, readErr := conn.ReadMessage() + if readErr != nil { + return + } + + plain := raw + var wrap struct { + Encrypted bool `json:"encrypted"` + Data string `json:"data"` + } + if err := json.Unmarshal(raw, &wrap); err == nil && wrap.Encrypted && strings.TrimSpace(wrap.Data) != "" { + crypto, cryptoErr := security.NewAESCrypto(nodeSecret) + if cryptoErr == nil { + if dec, decErr := crypto.Decrypt(wrap.Data); decErr == nil { + plain = []byte(dec) + } + } + } + + var cmd struct { + Type string `json:"type"` + RequestID string `json:"requestId"` + Data json.RawMessage `json:"data"` + } + if err := json.Unmarshal(plain, &cmd); err != nil { + continue + } + if strings.TrimSpace(cmd.RequestID) == "" { + continue + } + + shouldFail := false + failMsg := "" + if onCommand != nil { + shouldFail, failMsg = onCommand(strings.TrimSpace(cmd.Type), cmd.Data) + } + + respType := fmt.Sprintf("%sResponse", cmd.Type) + respPayload := map[string]interface{}{ + "type": respType, + "success": !shouldFail, + "message": "OK", + "requestId": cmd.RequestID, + } + if shouldFail { + if strings.TrimSpace(failMsg) == "" { + failMsg = "mock command failed" + } + respPayload["message"] = failMsg + } + + respBytes, err := json.Marshal(respPayload) + if err != nil { + continue + } + _ = conn.WriteMessage(websocket.TextMessage, respBytes) + } + }() + + var stopOnce sync.Once + return func() { + stopOnce.Do(func() { + _ = conn.Close() + wg.Wait() + }) + } +} diff --git a/go-gost/x/socket/limiter.go b/go-gost/x/socket/limiter.go index 61ed9e5..e623f4e 100644 --- a/go-gost/x/socket/limiter.go +++ b/go-gost/x/socket/limiter.go @@ -102,3 +102,71 @@ type updateLimiterRequest struct { type deleteLimiterRequest struct { Limiter string `json:"limiter"` } + +func createConnLimiter(req createLimiterRequest) error { + name := strings.TrimSpace(req.Data.Name) + if name == "" { + return errors.New("limiter name is required") + } + req.Data.Name = name + + if registry.ConnLimiterRegistry().IsRegistered(name) { + return errors.New("conn limiter " + name + " already exists") + } + + v := parser.ParseConnLimiter(&req.Data) + + if err := registry.ConnLimiterRegistry().Register(name, v); err != nil { + return errors.New("conn limiter " + name + " already exists") + } + + if c := config.Global(); c != nil { + c.CLimiters = append(c.CLimiters, &req.Data) + } + return nil +} + +func updateConnLimiter(req updateLimiterRequest) error { + name := strings.TrimSpace(req.Limiter) + req.Data.Name = name + if registry.ConnLimiterRegistry().IsRegistered(name) { + registry.ConnLimiterRegistry().Unregister(name) + } + + v := parser.ParseConnLimiter(&req.Data) + + if err := registry.ConnLimiterRegistry().Register(name, v); err != nil { + return errors.New("conn limiter " + name + " already exists") + } + + if c := config.Global(); c != nil { + for i := range c.CLimiters { + if c.CLimiters[i].Name == name { + c.CLimiters[i] = &req.Data + return nil + } + } + c.CLimiters = append(c.CLimiters, &req.Data) + } + return nil +} + +func deleteConnLimiter(req deleteLimiterRequest) error { + name := strings.TrimSpace(req.Limiter) + + if registry.ConnLimiterRegistry().IsRegistered(name) { + registry.ConnLimiterRegistry().Unregister(name) + } + + if c := config.Global(); c != nil { + limiteres := c.CLimiters + c.CLimiters = nil + for _, s := range limiteres { + if s.Name == name { + continue + } + c.CLimiters = append(c.CLimiters, s) + } + } + return nil +} diff --git a/go-gost/x/socket/websocket_reporter.go b/go-gost/x/socket/websocket_reporter.go index 72100b8..c32a963 100644 --- a/go-gost/x/socket/websocket_reporter.go +++ b/go-gost/x/socket/websocket_reporter.go @@ -836,6 +836,18 @@ func (w *WebSocketReporter) routeCommand(cmd CommandMessage) { err = w.handleDeleteLimiter(cmd.Data) response.Type = "DeleteLimitersResponse" needSaveConfig = true + case "AddCLimiters": + err = w.handleAddCLimiter(cmd.Data) + response.Type = "AddCLimitersResponse" + needSaveConfig = true + case "UpdateCLimiters": + err = w.handleUpdateCLimiter(cmd.Data) + response.Type = "UpdateCLimitersResponse" + needSaveConfig = true + case "DeleteCLimiters": + err = w.handleDeleteCLimiter(cmd.Data) + response.Type = "DeleteCLimitersResponse" + needSaveConfig = true // TCP Ping 诊断命令（只读，不需要保存配置） case "TcpPing": @@ -1130,6 +1142,67 @@ func (w *WebSocketReporter) handleDeleteLimiter(data interface{}) error { return deleteLimiter(deleteReq) } +func (w *WebSocketReporter) handleAddCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var limiterConfig config.LimiterConfig + if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { + return fmt.Errorf("解析限流器配置失败: %v", err) + } + + req := createLimiterRequest{Data: limiterConfig} + return createConnLimiter(req) +} + +func (w *WebSocketReporter) handleUpdateCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var updateReq struct { + Limiter string `json:"limiter"` + Data config.LimiterConfig `json:"data"` + } + + if err := json.Unmarshal(jsonData, &updateReq); err != nil { + var limiterConfig config.LimiterConfig + if err := json.Unmarshal(jsonData, &limiterConfig); err != nil { + return fmt.Errorf("解析更新请求失败: %v", err) + } + updateReq.Limiter = limiterConfig.Name + updateReq.Data = limiterConfig + } + + req := updateLimiterRequest{ + Limiter: updateReq.Limiter, + Data: updateReq.Data, + } + return updateConnLimiter(req) +} + +func (w *WebSocketReporter) handleDeleteCLimiter(data interface{}) error { + jsonData, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("序列化数据失败: %v", err) + } + + var deleteReq deleteLimiterRequest + + if err := json.Unmarshal(jsonData, &deleteReq); err != nil { + var limiterName string + if err := json.Unmarshal(jsonData, &limiterName); err != nil { + return fmt.Errorf("解析删除请求失败: %v", err) + } + deleteReq.Limiter = limiterName + } + + return deleteConnLimiter(deleteReq) +} + // handleSetProtocol 处理设置屏蔽协议的命令 func (w *WebSocketReporter) handleSetProtocol(data interface{}) error { jsonData, err := json.Marshal(data) diff --git a/vite-frontend/src/api/types.ts b/vite-frontend/src/api/types.ts index 7e92b67..dc31da5 100644 --- a/vite-frontend/src/api/types.ts +++ b/vite-frontend/src/api/types.ts @@ -20,6 +20,7 @@ export interface UserApiItem { num: number; expTime?: number; flowResetTime?: number; + maxConn?: number; inFlow?: number; outFlow?: number; dailyQuotaGB?: number; @@ -200,6 +201,7 @@ export interface UserPackageInfoApiData { num: number; expTime?: string; flowResetTime?: number; + maxConn?: number; [key: string]: unknown; }; tunnelPermissions: UserTunnelPermissionApiItem[]; @@ -276,6 +278,7 @@ export interface UserMutationPayload { num?: number; expTime?: number | string; flowResetTime?: number; + maxConn?: number; dailyQuotaGB?: number; monthlyQuotaGB?: number; tunnelFlow?: number; @@ -338,6 +341,7 @@ export interface UserTunnelAssignPayload { num?: number; expTime?: number; flowResetTime?: number; + maxConn?: number; status?: number; speedId?: number | null; tunnels?: Array<{ tunnelId: number; speedId?: number | null }>; diff --git a/vite-frontend/src/pages/forward.tsx b/vite-frontend/src/pages/forward.tsx index c63d17b..95048ef 100644 --- a/vite-frontend/src/pages/forward.tsx +++ b/vite-frontend/src/pages/forward.tsx @@ -157,6 +157,7 @@ interface ForwardForm { interfaceName?: string; strategy: string; speedId: number | null; + maxConn?: number; } interface ForwardUserGroup { @@ -1306,6 +1307,7 @@ export default function ForwardPage() { interfaceName: "", strategy: "fifo", speedId: null, + maxConn: 0, }); const [inIpTouched, setInIpTouched] = useState(false); @@ -4723,6 +4725,20 @@ export default function ForwardPage() {
+ { + const value = Math.max(Number(e.target.value) || 0, 0); + setForm((prev) => ({ ...prev, maxConn: value })); + }} + description="此设置优先于用户的全局连接数限制。0 表示不限制。" + variant="bordered" + /> + ({ ...prev, flow: value })); }} /> + { + const value = Math.max(Number(e.target.value) || 0, 0); + setUserForm((prev) => ({ ...prev, maxConn: value })); + }} + /> + { + const value = Math.max(Number(e.target.value) || 0, 0); + setUserForm((prev) => ({ ...prev, maxConn: value })); + }} + /> 规则数量 状态 限速规则 + 最大连接 重置时间 到期时间 操作 diff --git a/vite-frontend/src/types/index.ts b/vite-frontend/src/types/index.ts index 83fc3a4..6784911 100644 --- a/vite-frontend/src/types/index.ts +++ b/vite-frontend/src/types/index.ts @@ -44,6 +44,7 @@ export interface UserForm { num: number; expTime: Date | null; flowResetTime: number; + maxConn?: number; groupIds?: number[]; } @@ -56,7 +57,8 @@ export interface UserTunnel { flow: number; // 流量限制(GB) num: number; // 转发数量 expTime: number; // 过期时间戳 - flowResetTime: number; // 流量重置日期 + flowResetTime: number; + maxConn?: number; // 流量重置日期 speedId?: number | null; // 限速规则ID speedLimitName?: string; // 限速规则名称 inFlow?: number; // 下载流量(字节) @@ -70,6 +72,7 @@ export interface UserTunnelForm { num: number; expTime: Date | null; flowResetTime: number; + maxConn?: number; speedId: number | null; }