flvx/go-backend/internal/http/handler/flow_cleanup_regression_test.go

package handler

import (
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"reflect"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"gorm.io/gorm"

	"go-backend/internal/security"
	"go-backend/internal/store/model"
	"go-backend/internal/store/repo"
	"go-backend/internal/ws"
)

type cleanupAgentCommand struct {
	Type      string          `json:"type"`
	Data      json.RawMessage `json:"data"`
	RequestID string          `json:"requestId"`
}

// cleanupAgent exercises the actual command transport. Each command is recorded
// before its ACK, so synchronous handler calls need no sleeps to inspect it.
type cleanupAgent struct {
	h        *Handler
	mu       sync.Mutex
	commands []cleanupAgentCommand
}

func newCleanupAgent(t *testing.T) *cleanupAgent {
	t.Helper()
	r, err := repo.Open(filepath.Join(t.TempDir(), "cleanup.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = r.Close() })
	const secret = "cleanup-regression-node"
	if err := r.DB().Create(&model.Node{ID: 1, Name: "relay", Secret: secret, ServerIP: "127.0.0.1", Port: "30000", Status: 1}).Error; err != nil {
		t.Fatal(err)
	}
	server := ws.NewServer(r, "test-jwt-secret")
	online := make(chan struct{})
	server.SetNodeOnlineHook(func(int64) { close(online) })
	serverDone := make(chan struct{})
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		defer close(serverDone)
		server.ServeHTTP(w, req)
	}))
	t.Cleanup(ts.Close)
	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(ts.URL, "http")+"?type=1&secret="+secret, nil)
	if err != nil {
		t.Fatal(err)
	}
	crypto, err := security.NewAESCrypto(secret)
	if err != nil {
		t.Fatal(err)
	}
	a := &cleanupAgent{h: &Handler{repo: r, wsServer: server}}
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for {
			_, payload, err := conn.ReadMessage()
			if err != nil {
				return
			}
			var envelope struct {
				Encrypted bool   `json:"encrypted"`
				Data      string `json:"data"`
			}
			if err := json.Unmarshal(payload, &envelope); err != nil {
				t.Errorf("decode command envelope: %v", err)
				return
			}
			if envelope.Encrypted {
				payload, err = crypto.Decrypt(envelope.Data)
				if err != nil {
					t.Errorf("decrypt command: %v", err)
					return
				}
			}
			var command cleanupAgentCommand
			if err := json.Unmarshal(payload, &command); err != nil {
				t.Errorf("decode command: %v", err)
				return
			}
			a.mu.Lock()
			a.commands = append(a.commands, command)
			a.mu.Unlock()
			if err := conn.WriteJSON(map[string]interface{}{"type": command.Type, "requestId": command.RequestID, "success": true}); err != nil {
				return
			}
		}
	}()
	t.Cleanup(func() {
		_ = conn.Close()
		<-readerDone
		<-serverDone
	})
	select {
	case <-online:
	case <-time.After(3 * time.Second):
		t.Fatal("mock node did not come online")
	}
	a.probe(t)
	return a
}

func (a *cleanupAgent) probe(t *testing.T) {
	t.Helper()
	if _, err := a.h.wsServer.SendCommand(1, "CleanupTestProbe", nil, time.Second); err != nil {
		t.Fatalf("mock agent transport unavailable: %v", err)
	}
}

func (a *cleanupAgent) commandsOfType(commandType string) []cleanupAgentCommand {
	a.mu.Lock()
	defer a.mu.Unlock()
	var commands []cleanupAgentCommand
	for _, command := range a.commands {
		if command.Type == commandType {
			commands = append(commands, command)
		}
	}
	return commands
}

func (a *cleanupAgent) addRuntime(t *testing.T, name string, nodeID int64, status, applied int, updated time.Time) int64 {
	t.Helper()
	now := time.Now().UnixMilli()
	share := &repo.PeerShare{Name: "shared", NodeID: nodeID, Token: "cleanup-share-token", PortRangeStart: 31000, PortRangeEnd: 31010, IsActive: 1, CreatedTime: now, UpdatedTime: now}
	if err := a.h.repo.CreatePeerShare(share); err != nil {
		t.Fatal(err)
	}
	stored, err := a.h.repo.GetPeerShareByToken(share.Token)
	if err != nil || stored == nil {
		t.Fatalf("load share: %v", err)
	}
	// SQL preserves status=0; GORM's default tag would replace that with 1.
	if err := a.h.repo.DB().Exec(`INSERT INTO peer_share_runtime
		(share_id, node_id, reservation_id, resource_key, role, service_name, applied, status, created_time, updated_time)
		VALUES (?, ?, 'cleanup-reservation', 'cleanup-resource', 'forward', ?, ?, ?, ?, ?)`,
		stored.ID, nodeID, name, applied, status, now, updated.UnixMilli()).Error; err != nil {
		t.Fatal(err)
	}
	return stored.ID
}

func TestSharedForwardCleanupRegression(t *testing.T) {
	for _, mode := range []string{"single", "batch", "config"} {
		for _, runtimeName := range []string{"70_1_0", "70_1_0_tcp", "70_1_0_udp"} {
			t.Run(mode+"/active/"+runtimeName, func(t *testing.T) {
				a := newCleanupAgent(t)
				shareID := a.addRuntime(t, runtimeName, 1, 1, 1, time.Now())
				runCleanupPath(a.h, mode, []string{"70_1_0", "70_1_0_tcp", "70_1_0_udp"})
				a.probe(t)
				if commands := a.commandsOfType("DeleteService"); len(commands) != 0 {
					t.Fatalf("active shared service family was deleted: %+v", commands)
				}
				if mode != "config" && runtimeName == "70_1_0" {
					share, err := a.h.repo.GetPeerShare(shareID)
					if err != nil || share == nil || share.CurrentFlow != 600 {
						t.Fatalf("shared traffic should accumulate 600 bytes: share=%+v err=%v", share, err)
					}
				}
			})
		}
		for _, tc := range []struct {
			name        string
			serviceName string
			nodeID      int64
			status      int
			applied     int
			age         time.Duration
			wantDelete  bool
		}{
			{name: "recent-unbound", nodeID: 1, status: 1, age: time.Minute},
			{name: "stale-unbound", nodeID: 1, status: 1, age: 11 * time.Minute, wantDelete: true},
			{name: "released", serviceName: "70_1_0", nodeID: 1, applied: 1, wantDelete: true},
			{name: "other-node-bound", serviceName: "70_1_0", nodeID: 2, status: 1, applied: 1, wantDelete: true},
			{name: "other-node-unbound", nodeID: 2, status: 1, wantDelete: true},
			{name: "unrelated-active-family", serviceName: "71_1_0", nodeID: 1, status: 1, applied: 1, wantDelete: true},
		} {
			t.Run(mode+"/"+tc.name, func(t *testing.T) {
				a := newCleanupAgent(t)
				a.addRuntime(t, tc.serviceName, tc.nodeID, tc.status, tc.applied, time.Now().Add(-tc.age))
				runCleanupPath(a.h, mode, []string{"70_1_0_tcp"})
				a.probe(t)
				commands := a.commandsOfType("DeleteService")
				if !tc.wantDelete {
					if len(commands) != 0 {
						t.Fatalf("recent unbound shared runtime was deleted: %+v", commands)
					}
					return
				}
				if len(commands) != 1 {
					t.Fatalf("expected one orphan cleanup command, got %+v", commands)
				}
				var data struct {
					Services []string `json:"services"`
				}
				if err := json.Unmarshal(commands[0].Data, &data); err != nil {
					t.Fatal(err)
				}
				sort.Strings(data.Services)
				if want := []string{"70_1_0", "70_1_0_tcp", "70_1_0_udp"}; !reflect.DeepEqual(data.Services, want) {
					t.Fatalf("orphan cleanup must delete complete family: got %v, want %v", data.Services, want)
				}
			})
		}
	}
}

func runCleanupPath(h *Handler, mode string, names []string) {
	items := make([]flowItem, 0, len(names))
	services := make([]namedConfigItem, 0, len(names))
	for _, name := range names {
		items = append(items, flowItem{N: name, U: 120, D: 80})
		services = append(services, namedConfigItem{Name: name})
	}
	switch mode {
	case "single":
		for _, item := range items {
			h.processFlowItem(1, item)
		}
	case "batch":
		h.applyFlowUploadBatch(1, h.buildNodeFlowUploadBatch(1, items, nil), time.Now())
	case "config":
		h.cleanOrphanedServices(1, services)
	}
}

func TestSharedForwardCleanupMissingMetadataPreservesLocalForward(t *testing.T) {
	a := newCleanupAgent(t)
	if err := a.h.repo.DB().Create(&model.Forward{ID: 70, UserID: 1, UserName: "local", Name: "local", TunnelID: 1, RemoteAddr: "127.0.0.1:80", Status: 1}).Error; err != nil {
		t.Fatal(err)
	}
	runCleanupPath(a.h, "batch", []string{"70_1_0", "70_1_0_tcp", "70_1_0_udp"})
	a.probe(t)
	if commands := a.commandsOfType("DeleteService"); len(commands) != 0 {
		t.Fatalf("missing batch metadata must not delete a local forward: %+v", commands)
	}
}

func TestSharedForwardCleanupChecksActualDeleteFamily(t *testing.T) {
	for _, mode := range []string{"single", "batch", "config"} {
		t.Run(mode, func(t *testing.T) {
			a := newCleanupAgent(t)
			a.addRuntime(t, "70_1_0", 1, 1, 1, time.Now())
			// Legacy parsing accepts additional suffixes. Cleanup must check
			// the family it would delete, not just the reported service name.
			runCleanupPath(a.h, mode, []string{"70_1_0_old", "70_1_0_old_tcp"})
			a.probe(t)
			if commands := a.commandsOfType("DeleteService"); len(commands) != 0 {
				t.Fatalf("suffix variation must not delete a protected shared family: %+v", commands)
			}
		})
	}
}

func TestSharedForwardCleanupQueryFailurePreservesServices(t *testing.T) {
	for _, mode := range []string{"single", "batch", "config"} {
		for _, failure := range []string{"forward", "peer_share_runtime"} {
			t.Run(mode+"/"+failure, func(t *testing.T) {
				a := newCleanupAgent(t)
				// Fail only the ownership lookup; leave node lookup and WebSocket
				// delivery functional so an erroneous delete remains observable.
				callback := "test:cleanup-query-failure"
				var injected atomic.Int32
				if err := a.h.repo.DB().Callback().Query().Before("gorm:query").Register(callback, func(tx *gorm.DB) {
					if tx.Statement.Table == failure {
						injected.Add(1)
						tx.AddError(errors.New("injected ownership query failure"))
					}
				}); err != nil {
					t.Fatal(err)
				}
				t.Cleanup(func() { _ = a.h.repo.DB().Callback().Query().Remove(callback) })
				runCleanupPath(a.h, mode, []string{"70_1_0_tcp"})
				a.probe(t)
				if injected.Load() == 0 {
					t.Fatal("test did not inject the expected query failure")
				}
				if commands := a.commandsOfType("DeleteService"); len(commands) != 0 {
					t.Fatalf("ownership query failure must preserve services: %+v", commands)
				}
			})
		}
	}
}