diff --git a/go-backend/internal/http/handler/flow_policy.go b/go-backend/internal/http/handler/flow_policy.go index 3f8abee..81c641c 100644 --- a/go-backend/internal/http/handler/flow_policy.go +++ b/go-backend/internal/http/handler/flow_policy.go @@ -40,6 +40,7 @@ func (h *Handler) processFlowItem(item flowItem) { if ok { inFlow, outFlow := h.scaleFlowByTunnel(forwardID, item.D, item.U) _ = h.repo.AddFlow(forwardID, userID, userTunnelID, inFlow, outFlow) + h.processPeerShareFlowFromForward(forwardID, item) if userTunnelID > 0 { h.enforceFlowPolicies(userID, userTunnelID) @@ -87,6 +88,23 @@ func parsePeerShareRuntimeServiceID(serviceName string) (int64, bool) { return runtimeID, true } +func parsePeerShareIDFromFederationTunnelName(tunnelName string) (int64, bool) { + tunnelName = strings.TrimSpace(tunnelName) + if !strings.HasPrefix(tunnelName, "Share-") { + return 0, false + } + raw := strings.TrimPrefix(tunnelName, "Share-") + idx := strings.Index(raw, "-Port-") + if idx <= 0 { + return 0, false + } + shareID, err := strconv.ParseInt(raw[:idx], 10, 64) + if err != nil || shareID <= 0 { + return 0, false + } + return shareID, true +} + func (h *Handler) processPeerShareFlow(runtimeID int64, item flowItem) { if h == nil || h.repo == nil || runtimeID <= 0 { return @@ -113,6 +131,44 @@ func (h *Handler) processPeerShareFlow(runtimeID int64, item flowItem) { h.enforcePeerShareFlowLimit(share.ID) } +func (h *Handler) processPeerShareFlowFromForward(forwardID int64, item flowItem) { + if h == nil || h.repo == nil || forwardID <= 0 { + return + } + forward, err := h.getForwardRecord(forwardID) + if err != nil || forward == nil { + return + } + tunnel, err := h.getTunnelRecord(forward.TunnelID) + if err != nil || tunnel == nil { + return + } + tunnelName, err := h.repo.GetTunnelName(forward.TunnelID) + if err != nil { + return + } + shareID, ok := parsePeerShareIDFromFederationTunnelName(tunnelName) + if !ok { + return + } + + delta := item.D + item.U + if delta <= 0 { + return + } + + _ = h.repo.AddPeerShareCurrentFlow(shareID, delta) + + share, err := h.repo.GetPeerShare(shareID) + if err != nil || share == nil { + return + } + if !isPeerShareFlowExceeded(share) { + return + } + h.enforcePeerShareFlowLimit(share.ID) +} + func (h *Handler) enforcePeerShareFlowLimit(shareID int64) { if h == nil || h.repo == nil || shareID <= 0 { return diff --git a/go-backend/internal/http/handler/flow_policy_federation_test.go b/go-backend/internal/http/handler/flow_policy_federation_test.go index b9ca83a..81e8dcf 100644 --- a/go-backend/internal/http/handler/flow_policy_federation_test.go +++ b/go-backend/internal/http/handler/flow_policy_federation_test.go @@ -2,6 +2,7 @@ package handler import ( "path/filepath" + "strconv" "testing" "time" @@ -61,3 +62,71 @@ func TestProcessFlowItemTracksPeerShareFlowAndEnforcesLimit(t *testing.T) { t.Fatalf("expected runtime status=0 after limit enforcement, got %d", runtime.Status) } } + +func TestProcessFlowItemTracksPeerShareFlowForFederationPortForward(t *testing.T) { + r, err := repo.Open(filepath.Join(t.TempDir(), "panel-forward.db")) + if err != nil { + t.Fatalf("open repo: %v", err) + } + defer r.Close() + + now := time.Now().UnixMilli() + if err := r.CreatePeerShare(&repo.PeerShare{ + Name: "forward-share", + NodeID: 1, + Token: "forward-share-token", + MaxBandwidth: 0, + CurrentFlow: 0, + PortRangeStart: 30000, + PortRangeEnd: 30010, + IsActive: 1, + CreatedTime: now, + UpdatedTime: now, + }); err != nil { + t.Fatalf("create peer share: %v", err) + } + share, err := r.GetPeerShareByToken("forward-share-token") + if err != nil || share == nil { + t.Fatalf("load peer share: %v", err) + } + + 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', 'x', 1, ?, 99999, 0, 0, 1, 1, ?, ?, 1) + `, now+24*60*60*1000, now, now).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + + tunnelName := "Share-" + strconv.FormatInt(share.ID, 10) + "-Port-30001" + 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, ?, 1.0, 1, 'tls', 1, ?, ?, 1, NULL, 0) + `, tunnelName, 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, 1, 99999, 0, 0, 1, ?, 1) + `, now+24*60*60*1000).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:443', 'fifo', 0, 0, ?, ?, 1, 0) + `, now, now).Error; err != nil { + t.Fatalf("insert forward: %v", err) + } + + h := &Handler{repo: r} + h.processFlowItem(flowItem{N: "20_2_10", U: 120, D: 80}) + + updatedShare, err := r.GetPeerShare(share.ID) + if err != nil || updatedShare == nil { + t.Fatalf("reload share: %v", err) + } + if updatedShare.CurrentFlow != 200 { + t.Fatalf("expected current_flow=200, got %d", updatedShare.CurrentFlow) + } +}