From 2e8c0530a98b537e8feb69adf59dc8aefc5347b7 Mon Sep 17 00:00:00 2001 From: sagitchu Date: Wed, 11 Mar 2026 09:34:36 +0800 Subject: [PATCH] fix(backend): block forwards when flow exceeded --- .../internal/http/handler/control_plane.go | 15 ++ .../internal/http/handler/flow_policy.go | 62 ++++++ go-backend/internal/http/handler/mutations.go | 18 +- ...flow_limit_forward_enable_contract_test.go | 201 ++++++++++++++++++ 4 files changed, 294 insertions(+), 2 deletions(-) create mode 100644 go-backend/tests/contract/flow_limit_forward_enable_contract_test.go diff --git a/go-backend/internal/http/handler/control_plane.go b/go-backend/internal/http/handler/control_plane.go index b226c30..53c5219 100644 --- a/go-backend/internal/http/handler/control_plane.go +++ b/go-backend/internal/http/handler/control_plane.go @@ -312,6 +312,14 @@ func (h *Handler) syncForwardServicesWithWarnings(forward *forwardRecord, method return warnings, fmt.Errorf("节点 %s 下发失败: %w", node.Name, err) } } + + // Keep paused forwards paused after UpdateService/AddService, since agent-side UpdateService + // always restarts services. + if forward.Status != 1 { + if err := h.controlForwardServices(forward, "PauseService", false); err != nil { + return warnings, err + } + } return warnings, nil } @@ -1620,6 +1628,13 @@ func processServerAddress(serverAddr string) string { if strings.HasPrefix(serverAddr, "[") { return serverAddr } + // If the input is a bare IPv6 host (no port), bracket it. + // IPv6-with-port must be provided in bracket form: [::1]:443. + if looksLikeIPv6(serverAddr) { + if ip := net.ParseIP(serverAddr); ip != nil && ip.To4() == nil { + return "[" + serverAddr + "]" + } + } idx := strings.LastIndex(serverAddr, ":") if idx < 0 { diff --git a/go-backend/internal/http/handler/flow_policy.go b/go-backend/internal/http/handler/flow_policy.go index 0850cf5..9691d17 100644 --- a/go-backend/internal/http/handler/flow_policy.go +++ b/go-backend/internal/http/handler/flow_policy.go @@ -2,6 +2,7 @@ package handler import ( "encoding/json" + "errors" "log" "strconv" "strings" @@ -327,6 +328,67 @@ func (h *Handler) enforceFlowPolicies(userID int64, userTunnelID int64) { } } +func (h *Handler) ensureUserTunnelForwardAllowed(userID int64, tunnelID int64, now int64) error { + if h == nil || h.repo == nil { + return errors.New("invalid flow policy context") + } + if userID <= 0 || tunnelID <= 0 { + return nil + } + + user, err := h.repo.GetUserByID(userID) + if err != nil { + return err + } + if user == nil { + return errors.New("用户不存在") + } + + if user.Status != 1 { + return errors.New("账号已禁用") + } + if user.ExpTime > 0 && user.ExpTime <= now { + return errors.New("账号已过期") + } + + flowLimit := user.Flow * bytesPerGB + current := user.InFlow + user.OutFlow + if flowLimit < current { + return errors.New("流量已超额，禁止开启转发") + } + + userTunnelID, _, _, err := h.resolveUserTunnelAndLimiter(userID, tunnelID) + if err != nil { + return err + } + if userTunnelID <= 0 { + return nil + } + + policy, err := h.getUserTunnelPolicy(userTunnelID) + if err != nil { + return err + } + if policy == nil { + return nil + } + + if policy.Status != 1 { + return errors.New("该隧道已禁用") + } + if policy.ExpTime > 0 && policy.ExpTime <= now { + return errors.New("该隧道已过期") + } + + utFlowLimit := policy.Flow * bytesPerGB + utCurrent := policy.InFlow + policy.OutFlow + if utCurrent >= utFlowLimit { + return errors.New("该隧道流量已超额，禁止开启转发") + } + + return nil +} + func (h *Handler) shouldPauseUser(userID int64, now int64) bool { user, err := h.repo.GetUserByID(userID) if err != nil || user == nil { diff --git a/go-backend/internal/http/handler/mutations.go b/go-backend/internal/http/handler/mutations.go index 65a0373..1f6bf0e 100644 --- a/go-backend/internal/http/handler/mutations.go +++ b/go-backend/internal/http/handler/mutations.go @@ -1155,6 +1155,10 @@ func (h *Handler) forwardCreate(w http.ResponseWriter, r *http.Request) { response.WriteJSON(w, response.ErrDefault("隧道已禁用，无法创建转发")) return } + if err := h.ensureUserTunnelForwardAllowed(userID, tunnelID, time.Now().UnixMilli()); err != nil { + response.WriteJSON(w, response.ErrDefault(err.Error())) + return + } name := asString(req["name"]) remoteAddr := asString(req["remoteAddr"]) if name == "" || remoteAddr == "" { @@ -1444,11 +1448,16 @@ func (h *Handler) forwardResume(w http.ResponseWriter, r *http.Request) { response.WriteJSON(w, response.Err(-2, err.Error())) return } + now := time.Now().UnixMilli() + if err := h.ensureUserTunnelForwardAllowed(forward.UserID, forward.TunnelID, now); err != nil { + response.WriteJSON(w, response.ErrDefault(err.Error())) + return + } if err := h.controlForwardServices(forward, "ResumeService", false); err != nil { response.WriteJSON(w, response.ErrDefault(err.Error())) return } - _ = h.repo.UpdateForwardStatus(id, 1, time.Now().UnixMilli()) + _ = h.repo.UpdateForwardStatus(id, 1, now) response.WriteJSON(w, response.OKEmpty()) } @@ -1571,17 +1580,22 @@ func (h *Handler) forwardBatchResume(w http.ResponseWriter, r *http.Request) { } s := 0 f := 0 + now := time.Now().UnixMilli() for _, id := range ids { forward, accessErr := h.ensureForwardAccessByActor(actorUserID, actorRole, id) if accessErr != nil { f++ continue } + if err := h.ensureUserTunnelForwardAllowed(forward.UserID, forward.TunnelID, now); err != nil { + f++ + continue + } if err := h.controlForwardServices(forward, "ResumeService", false); err != nil { f++ continue } - if err := h.repo.UpdateForwardStatus(id, 1, time.Now().UnixMilli()); err != nil { + if err := h.repo.UpdateForwardStatus(id, 1, now); err != nil { f++ } else { s++ diff --git a/go-backend/tests/contract/flow_limit_forward_enable_contract_test.go b/go-backend/tests/contract/flow_limit_forward_enable_contract_test.go new file mode 100644 index 0000000..5d27be0 --- /dev/null +++ b/go-backend/tests/contract/flow_limit_forward_enable_contract_test.go @@ -0,0 +1,201 @@ +package contract_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "go-backend/internal/auth" + "go-backend/internal/http/response" +) + +const contractBytesPerGB int64 = 1024 * 1024 * 1024 + +func TestForwardResumeBlockedWhenUserFlowExceeded(t *testing.T) { + secret := "contract-jwt-secret" + router, repo := setupContractRouter(t, secret) + now := time.Now().UnixMilli() + + userID := int64(2) + tunnelID := int64(1) + forwardID := int64(1) + + flowGB := int64(120) + used := flowGB*contractBytesPerGB + 1 + + if err := repo.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(?, 'flow_user', 'pwd', 1, 2727251700000, ?, ?, 0, 1, 99999, ?, ?, 1) + `, userID, flowGB, used, now, now).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + + if err := repo.DB().Exec(` + INSERT INTO tunnel(id, name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) + VALUES(?, 'flow_tunnel', 1.0, 1, 'tls', 99999, ?, ?, 1, NULL, 0) + `, tunnelID, now, now).Error; err != nil { + t.Fatalf("insert tunnel: %v", err) + } + if err := repo.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, ?, ?, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1) + `, userID, tunnelID).Error; err != nil { + t.Fatalf("insert user_tunnel: %v", err) + } + if err := repo.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(?, ?, 'flow_user', 'flow_forward', ?, '8.8.8.8:53', 'fifo', 0, 0, ?, ?, 0, 0) + `, forwardID, userID, tunnelID, now, now).Error; err != nil { + t.Fatalf("insert forward: %v", err) + } + + token, err := auth.GenerateToken(userID, "flow_user", 1, secret) + if err != nil { + t.Fatalf("generate token: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/api/v1/forward/resume", bytes.NewBufferString(`{"id":1}`)) + req.Header.Set("Authorization", token) + 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 non-zero code when flow exceeded") + } + if !strings.Contains(out.Msg, "流量") { + t.Fatalf("expected flow exceeded message, got %q", out.Msg) + } + + status := mustQueryInt(t, repo, `SELECT status FROM forward WHERE id = ?`, forwardID) + if status != 0 { + t.Fatalf("expected forward status to remain 0, got %d", status) + } +} + +func TestForwardResumeBlockedWhenUserTunnelFlowExceeded(t *testing.T) { + secret := "contract-jwt-secret" + router, repo := setupContractRouter(t, secret) + now := time.Now().UnixMilli() + + userID := int64(2) + tunnelID := int64(1) + forwardID := int64(1) + + if err := repo.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(?, 'ut_flow_user', 'pwd', 1, 2727251700000, 99999, 0, 0, 1, 99999, ?, ?, 1) + `, userID, now, now).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + + if err := repo.DB().Exec(` + INSERT INTO tunnel(id, name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) + VALUES(?, 'ut_flow_tunnel', 1.0, 1, 'tls', 99999, ?, ?, 1, NULL, 0) + `, tunnelID, now, now).Error; err != nil { + t.Fatalf("insert tunnel: %v", err) + } + + utFlowGB := int64(120) + utUsed := utFlowGB * contractBytesPerGB + if err := repo.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, ?, ?, NULL, 99999, ?, ?, 0, 1, 2727251700000, 1) + `, userID, tunnelID, utFlowGB, utUsed).Error; err != nil { + t.Fatalf("insert user_tunnel: %v", err) + } + if err := repo.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(?, ?, 'ut_flow_user', 'ut_flow_forward', ?, '8.8.8.8:53', 'fifo', 0, 0, ?, ?, 0, 0) + `, forwardID, userID, tunnelID, now, now).Error; err != nil { + t.Fatalf("insert forward: %v", err) + } + + token, err := auth.GenerateToken(userID, "ut_flow_user", 1, secret) + if err != nil { + t.Fatalf("generate token: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/api/v1/forward/resume", bytes.NewBufferString(`{"id":1}`)) + req.Header.Set("Authorization", token) + 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 non-zero code when tunnel flow exceeded") + } + if !strings.Contains(out.Msg, "隧道") || !strings.Contains(out.Msg, "流量") { + t.Fatalf("expected tunnel flow exceeded message, got %q", out.Msg) + } +} + +func TestForwardCreateBlockedWhenFlowExceeded(t *testing.T) { + secret := "contract-jwt-secret" + router, repo := setupContractRouter(t, secret) + now := time.Now().UnixMilli() + + userID := int64(2) + tunnelID := int64(1) + + flowGB := int64(120) + used := flowGB*contractBytesPerGB + 1 + + if err := repo.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(?, 'create_flow_user', 'pwd', 1, 2727251700000, ?, ?, 0, 1, 99999, ?, ?, 1) + `, userID, flowGB, used, now, now).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + + if err := repo.DB().Exec(` + INSERT INTO tunnel(id, name, traffic_ratio, type, protocol, flow, created_time, updated_time, status, in_ip, inx) + VALUES(?, 'create_flow_tunnel', 1.0, 1, 'tls', 99999, ?, ?, 1, NULL, 0) + `, tunnelID, now, now).Error; err != nil { + t.Fatalf("insert tunnel: %v", err) + } + if err := repo.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, ?, ?, NULL, 99999, 99999, 0, 0, 1, 2727251700000, 1) + `, userID, tunnelID).Error; err != nil { + t.Fatalf("insert user_tunnel: %v", err) + } + + token, err := auth.GenerateToken(userID, "create_flow_user", 1, secret) + if err != nil { + t.Fatalf("generate token: %v", err) + } + + payload := `{"tunnelId":1,"name":"n","remoteAddr":"1.1.1.1:53"}` + req := httptest.NewRequest(http.MethodPost, "/api/v1/forward/create", bytes.NewBufferString(payload)) + req.Header.Set("Authorization", token) + 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 non-zero code when flow exceeded") + } + if !strings.Contains(out.Msg, "流量") { + t.Fatalf("expected flow exceeded message, got %q", out.Msg) + } +}