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

package handler

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"go-backend/internal/http/response"
	"go-backend/internal/store/model"
	"go-backend/internal/ws"
)

func TestUserUpdateDisablesPermissionsAndRunningForwards(t *testing.T) {
	for _, password := range []string{"", "replacement-password"} {
		t.Run(fmt.Sprintf("password_changed=%t", password != ""), func(t *testing.T) {
			a := newCleanupAgent(t)
			now := time.Now().UnixMilli()
			for _, userID := range []int64{2, 3} {
				if err := a.h.repo.DB().Create(&model.User{ID: userID, User: fmt.Sprintf("user%d", userID), Pwd: "pwd", RoleID: 1, Status: 1, ExpTime: now + 86400000, Flow: 100, Num: 10}).Error; err != nil {
					t.Fatal(err)
				}
			}
			if err := a.h.repo.DB().Create(&model.Tunnel{ID: 1, Name: "test", Type: 1, Protocol: "tls", Status: 1, Flow: 1}).Error; err != nil {
				t.Fatal(err)
			}
			for _, permission := range []model.UserTunnel{
				{ID: 10, UserID: 2, TunnelID: 1, Status: 1, Flow: 100, Num: 10},
				{ID: 11, UserID: 2, TunnelID: 2, Status: 0, Flow: 100, Num: 10},
				{ID: 12, UserID: 3, TunnelID: 1, Status: 1, Flow: 100, Num: 10},
			} {
				if err := a.h.repo.DB().Create(&permission).Error; err != nil {
					t.Fatal(err)
				}
			}
			for _, forward := range []model.Forward{
				{ID: 20, UserID: 2, UserName: "user2", Name: "running", TunnelID: 1, RemoteAddr: "127.0.0.1:80", Status: 1},
				{ID: 21, UserID: 2, UserName: "user2", Name: "paused", TunnelID: 1, RemoteAddr: "127.0.0.1:81", Status: 0},
				{ID: 22, UserID: 3, UserName: "user3", Name: "other-user", TunnelID: 1, RemoteAddr: "127.0.0.1:82", Status: 1},
			} {
				if err := a.h.repo.DB().Create(&forward).Error; err != nil {
					t.Fatal(err)
				}
				if err := a.h.repo.DB().Create(&model.ForwardPort{ForwardID: forward.ID, NodeID: 1, Port: 31000 + int(forward.ID)}).Error; err != nil {
					t.Fatal(err)
				}
			}

			update := func(status int) {
				t.Helper()
				payload, err := json.Marshal(map[string]interface{}{"id": 2, "user": "user2", "status": status, "pwd": password, "flow": 100, "num": 10, "expTime": now + 86400000})
				if err != nil {
					t.Fatal(err)
				}
				res := httptest.NewRecorder()
				a.h.userUpdate(res, httptest.NewRequest(http.MethodPost, "/api/v1/user/update", bytes.NewReader(payload)))
				var result response.R
				if err := json.Unmarshal(res.Body.Bytes(), &result); err != nil {
					t.Fatal(err)
				}
				if result.Code != 0 {
					t.Fatalf("update user: %d %s", result.Code, result.Msg)
				}
			}

			update(1)
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM user_tunnel WHERE id = 10"); got != 1 {
				t.Errorf("ordinary account update disabled an enabled permission: %d", got)
			}
			if len(a.commandsOfType("PauseService")) != 0 {
				t.Error("ordinary account update paused services")
			}
			update(0)
			for _, check := range []struct {
				table string
				id    int64
				want  int
			}{{"user", 2, 0}, {"user_tunnel", 10, 0}, {"user_tunnel", 11, 0}, {"user_tunnel", 12, 1}, {"forward", 20, 0}, {"forward", 21, 0}, {"forward", 22, 1}} {
				if got := mustQueryInt(t, a.h.repo, "SELECT status FROM "+check.table+" WHERE id = ?", check.id); got != check.want {
					t.Errorf("%s %d status = %d, want %d", check.table, check.id, got, check.want)
				}
			}
			commands := a.commandsOfType("PauseService")
			if len(commands) == 0 {
				t.Error("running forward never received PauseService")
			}
			for _, command := range commands {
				var data struct {
					Services []string `json:"services"`
				}
				if err := json.Unmarshal(command.Data, &data); err != nil {
					t.Fatal(err)
				}
				for _, name := range data.Services {
					if !strings.HasPrefix(name, "20_2_") {
						t.Errorf("paused unrelated service %q", name)
					}
				}
			}
			update(0)
			if len(a.commandsOfType("PauseService")) != len(commands) {
				t.Error("repeated disable replayed pauses for already paused forwards")
			}

			// Account edits and re-enabling must not grant disabled permissions or
			// resume forwards that an administrator previously paused.
			update(1)
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM user_tunnel WHERE id = 10"); got != 0 {
				t.Errorf("account re-enable unexpectedly granted a disabled permission: %d", got)
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM user_tunnel WHERE id = 11"); got != 0 {
				t.Errorf("re-enabled independently disabled permission: %d", got)
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 21"); got != 0 {
				t.Errorf("resumed previously paused forward: %d", got)
			}
			if len(a.commandsOfType("ResumeService")) != 0 {
				t.Error("account re-enable unexpectedly resumed services")
			}
		})
	}
}

func TestUserDisableMultiEntryNeverRecreatesMissingServices(t *testing.T) {
	for _, secondNodeState := range []string{"missing", "offline", "timeout"} {
		t.Run(secondNodeState, func(t *testing.T) {
			a := newCleanupAgent(t)
			var remoteMu sync.Mutex
			remoteServices := make(map[string]bool)
			var remoteCommands []string
			remote := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				var command struct {
					Type string          `json:"commandType"`
					Data json.RawMessage `json:"data"`
				}
				if err := json.NewDecoder(r.Body).Decode(&command); err != nil {
					t.Error(err)
					return
				}
				remoteMu.Lock()
				defer remoteMu.Unlock()
				remoteCommands = append(remoteCommands, command.Type)
				switch command.Type {
				case "UpdateService", "AddService":
					var services []struct {
						Name string `json:"name"`
					}
					if err := json.Unmarshal(command.Data, &services); err != nil {
						t.Error(err)
						return
					}
					for _, service := range services {
						remoteServices[service.Name] = true
					}
				case "PauseService":
					if secondNodeState != "missing" {
						message := "节点不在线"
						if secondNodeState == "timeout" {
							message = "等待节点响应超时"
						}
						response.WriteJSON(w, response.ErrDefault(message))
						return
					}
					var data struct {
						Services []string `json:"services"`
					}
					if err := json.Unmarshal(command.Data, &data); err != nil {
						t.Error(err)
						return
					}
					for _, name := range data.Services {
						if _, exists := remoteServices[name]; !exists {
							response.WriteJSON(w, response.OK(map[string]interface{}{"success": false, "message": "service not found"}))
							return
						}
						remoteServices[name] = false
					}
				}
				response.WriteJSON(w, response.OK(map[string]interface{}{"success": true, "type": command.Type}))
			}))
			defer remote.Close()
			now := time.Now().UnixMilli()
			userID, err := a.h.repo.CreateUser("multi-entry-disable", "pwd", 1, now+86400000, 100, 1, 10, 1, 0, now)
			if err != nil {
				t.Fatal(err)
			}
			for _, value := range []interface{}{
				&model.Node{ID: 2, Name: "second-entry", Secret: "test", ServerIP: "127.0.0.1", Status: 1, IsRemote: 1, RemoteURL: sql.NullString{String: remote.URL, Valid: true}, RemoteToken: sql.NullString{String: "test-token", Valid: true}},
				&model.Tunnel{ID: 1, Name: "multi-entry", Type: 1, Status: 1, Protocol: "tls", Flow: 1},
				&model.ChainTunnel{TunnelID: 1, NodeID: 1, ChainType: "1"},
				&model.ChainTunnel{TunnelID: 1, NodeID: 2, ChainType: "1"},
				&model.UserTunnel{ID: 10, UserID: userID, TunnelID: 1, Status: 1, Flow: 100, Num: 10},
				&model.Forward{ID: 20, UserID: userID, UserName: "multi-entry-disable", Name: "running", TunnelID: 1, RemoteAddr: "127.0.0.1:80", Status: 1},
				&model.ForwardPort{ForwardID: 20, NodeID: 1, Port: 31020},
				&model.ForwardPort{ForwardID: 20, NodeID: 2, Port: 31020},
			} {
				if err := a.h.repo.DB().Create(value).Error; err != nil {
					t.Fatal(err)
				}
			}
			res := postJSONToHandler(t, a.h.userUpdate, map[string]interface{}{"id": userID, "user": "multi-entry-disable", "status": 0})
			var out response.R
			if err := json.Unmarshal(res.Body.Bytes(), &out); err != nil {
				t.Fatal(err)
			}
			wantStatus := 1
			if secondNodeState == "missing" {
				wantStatus = 0
				if out.Code != 0 {
					t.Errorf("missing service should count as stopped: %s", out.Msg)
				}
			} else if out.Code == 0 {
				t.Error("unconfirmed second entry pause returned success")
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 20"); got != wantStatus {
				t.Errorf("forward status = %d, want %d", got, wantStatus)
			}
			a.mu.Lock()
			commands := append([]cleanupAgentCommand(nil), a.commands...)
			a.mu.Unlock()
			firstEntryRunning := true
			for _, command := range commands {
				switch command.Type {
				case "PauseService":
					firstEntryRunning = false
				case "UpdateService", "AddService", "ResumeService":
					firstEntryRunning = true
					t.Errorf("disable sent %s to already paused entry A", command.Type)
				}
			}
			if firstEntryRunning {
				t.Error("entry A is running after disable")
			}
			remoteMu.Lock()
			defer remoteMu.Unlock()
			if len(remoteCommands) == 0 {
				t.Error("entry B was never checked")
			}
			for _, command := range remoteCommands {
				if command != "PauseService" {
					t.Errorf("disable sent %s to entry B", command)
				}
			}
			for name, running := range remoteServices {
				if running {
					t.Errorf("entry B service %s is still running", name)
				}
			}
		})
	}
}

func TestUserDisableContinuesAcrossNodesAndForwards(t *testing.T) {
	for _, failure := range []string{"节点不在线", "等待节点响应超时", "pause rejected"} {
		for _, failedNodeFirst := range []bool{false, true} {
			t.Run(fmt.Sprintf("%s/failed_node_first=%t", failure, failedNodeFirst), func(t *testing.T) {
				a := newCleanupAgent(t)
				var mu sync.Mutex
				recovered := false
				attempts := make(map[string]int)
				remote := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					var command struct {
						Type string `json:"commandType"`
					}
					if err := json.NewDecoder(r.Body).Decode(&command); err != nil {
						t.Error(err)
						return
					}
					if command.Type != "PauseService" {
						t.Errorf("disable sent %s", command.Type)
					}
					mu.Lock()
					attempts[r.Header.Get("Authorization")]++
					isRecovered := recovered
					mu.Unlock()
					if !isRecovered {
						response.WriteJSON(w, response.ErrDefault(failure))
						return
					}
					response.WriteJSON(w, response.OK(map[string]interface{}{"success": true, "type": command.Type}))
				}))
				defer remote.Close()
				now := time.Now().UnixMilli()
				userID, err := a.h.repo.CreateUser("continue-disable", "pwd", 1, now+86400000, 100, 1, 10, 1, 0, now)
				if err != nil {
					t.Fatal(err)
				}
				for _, value := range []interface{}{
					&model.Tunnel{ID: 1, Name: "multi-entry", Type: 2, Status: 1, Flow: 1},
					&model.Node{ID: 4, Name: "offline-middle", Secret: "test", ServerIP: "127.0.0.1", Status: 0},
					&model.ChainTunnel{TunnelID: 1, NodeID: 4, ChainType: "2"},
					&model.UserTunnel{ID: 10, UserID: userID, TunnelID: 1, Status: 1, Flow: 100, Num: 10},
					&model.Forward{ID: 20, UserID: userID, UserName: "continue-disable", Name: "mixed-nodes", TunnelID: 1, RemoteAddr: "127.0.0.1:80", Status: 1},
					&model.Forward{ID: 21, UserID: userID, UserName: "continue-disable", Name: "healthy-forward", TunnelID: 1, RemoteAddr: "127.0.0.1:81", Status: 1},
				} {
					if err := a.h.repo.DB().Create(value).Error; err != nil {
						t.Fatal(err)
					}
				}
				for _, nodeID := range []int64{2, 3} {
					node := model.Node{ID: nodeID, Name: fmt.Sprintf("failed-node-%d", nodeID), Secret: "test", ServerIP: "127.0.0.1", Status: 1, IsRemote: 1, RemoteURL: sql.NullString{String: remote.URL, Valid: true}, RemoteToken: sql.NullString{String: fmt.Sprintf("test-node-%d", nodeID), Valid: true}}
					if err := a.h.repo.DB().Create(&node).Error; err != nil {
						t.Fatal(err)
					}
				}
				nodeOrder := []int64{1, 2, 3}
				if failedNodeFirst {
					nodeOrder = []int64{2, 1, 3}
				}
				for _, nodeID := range nodeOrder {
					if err := a.h.repo.DB().Create(&model.ForwardPort{ForwardID: 20, NodeID: nodeID, Port: 31020}).Error; err != nil {
						t.Fatal(err)
					}
					if err := a.h.repo.DB().Create(&model.ChainTunnel{TunnelID: 1, NodeID: nodeID, ChainType: "1"}).Error; err != nil {
						t.Fatal(err)
					}
				}
				if err := a.h.repo.DB().Create(&model.ForwardPort{ForwardID: 21, NodeID: 1, Port: 31021}).Error; err != nil {
					t.Fatal(err)
				}
				update := func() response.R {
					t.Helper()
					res := postJSONToHandler(t, a.h.userUpdate, map[string]interface{}{"id": userID, "user": "continue-disable", "status": 0})
					var out response.R
					if err := json.Unmarshal(res.Body.Bytes(), &out); err != nil {
						t.Fatal(err)
					}
					return out
				}
				out := update()
				if out.Code == 0 {
					t.Error("partial node failure reported complete success")
				}
				state, ok := out.Data.(map[string]interface{})
				if !ok || state["userDisabled"] != true {
					t.Errorf("partial response did not tell the UI the account is disabled: %#v", out.Data)
				}
				for _, nodeID := range []int64{2, 3} {
					if !strings.Contains(out.Msg, fmt.Sprintf("failed-node-%d", nodeID)) || !strings.Contains(out.Msg, fmt.Sprintf("ID %d", nodeID)) {
						t.Errorf("failure response did not identify node %d: %s", nodeID, out.Msg)
					}
					mu.Lock()
					count := attempts[fmt.Sprintf("Bearer test-node-%d", nodeID)]
					mu.Unlock()
					if count == 0 {
						t.Errorf("failure on another node prevented node %d from being attempted", nodeID)
					}
				}
				paused := make(map[string]bool)
				for _, command := range a.commandsOfType("PauseService") {
					var data struct {
						Services []string `json:"services"`
					}
					if err := json.Unmarshal(command.Data, &data); err != nil {
						t.Fatal(err)
					}
					for _, name := range data.Services {
						paused[strings.Split(name, "_")[0]] = true
					}
				}
				for _, id := range []string{"20", "21"} {
					if !paused[id] {
						t.Errorf("healthy node never paused forward %s", id)
					}
				}
				if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 20"); got != 1 {
					t.Errorf("unconfirmed forward status = %d", got)
				}
				if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 21"); got != 0 {
					t.Errorf("healthy later forward was not paused: %d", got)
				}
				mu.Lock()
				recovered = true
				mu.Unlock()
				if retry := update(); retry.Code != 0 {
					t.Errorf("retry after recovery failed: %s", retry.Msg)
				}
				if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 20"); got != 0 {
					t.Errorf("retry did not confirm all nodes: %d", got)
				}
				for _, kind := range []string{"UpdateService", "AddService", "ResumeService"} {
					if len(a.commandsOfType(kind)) != 0 {
						t.Errorf("disable sent %s", kind)
					}
				}
			})
		}
	}
}

func TestUserDisableRemovesNftablesRuleOnlyAfterConfirmedApply(t *testing.T) {
	fixture := setupNftablesHandler(t)
	h := fixture.handler
	t.Cleanup(func() { _ = h.repo.Close() })
	seedNftablesSSHConfig(t, h, fixture.nodeID)
	tunnelID := seedTunnelForNftables(t, h, "disable-nft", fixture.nodeID)
	forward := seedForwardForNftables(t, h, tunnelID, fixture.nodeID, "203.0.113.9:8080")
	now := time.Now().UnixMilli()
	userID, err := h.repo.CreateUser("disable-nft", "pwd", 1, now+86400000, 100, 1, 10, 1, 0, now)
	if err != nil {
		t.Fatal(err)
	}
	if err := h.repo.DB().Model(&model.Forward{}).Where("id = ?", forward.ID).Update("user_id", userID).Error; err != nil {
		t.Fatal(err)
	}
	if err := h.repo.DB().Create(&model.UserTunnel{UserID: userID, TunnelID: tunnelID, Status: 1, Flow: 100, Num: 10}).Error; err != nil {
		t.Fatal(err)
	}
	otherForwardID := forward.ID + 100
	for _, value := range []interface{}{
		&model.Forward{ID: otherForwardID, UserID: 1, UserName: "admin", Name: "other-user", TunnelID: tunnelID, RemoteAddr: "203.0.113.9:8081", Status: 1},
		&model.ForwardPort{ForwardID: otherForwardID, NodeID: fixture.nodeID, Port: 20001},
	} {
		if err := h.repo.DB().Create(value).Error; err != nil {
			t.Fatal(err)
		}
	}
	manager := &fakeNftablesManager{reconcileErr: errors.New("ssh unavailable")}
	h.nftablesManager = manager
	update := func() response.R {
		t.Helper()
		res := postJSONToHandler(t, h.userUpdate, map[string]interface{}{"id": userID, "user": "disable-nft", "status": 0})
		var out response.R
		if err := json.Unmarshal(res.Body.Bytes(), &out); err != nil {
			t.Fatal(err)
		}
		return out
	}
	if out := update(); out.Code == 0 {
		t.Error("failed SSH apply returned success")
	}
	if got := mustQueryInt(t, h.repo, "SELECT status FROM forward WHERE id = ?", forward.ID); got != 1 {
		t.Errorf("unconfirmed nftables pause marked completed: %d", got)
	}
	if manager.reconcileCount() != 1 || len(manager.lastPlan.Rules) != 1 || manager.lastPlan.Rules[0].ForwardID != otherForwardID {
		t.Fatalf("did not attempt to remove disabled user's nftables rule: %+v", manager.lastPlan)
	}
	manager.reconcileErr = nil
	if out := update(); out.Code != 0 {
		t.Fatalf("retry nftables disable: %s", out.Msg)
	}
	if got := mustQueryInt(t, h.repo, "SELECT status FROM forward WHERE id = ?", forward.ID); got != 0 {
		t.Errorf("confirmed nftables pause was not stored: %d", got)
	}
	if manager.reconcileCount() != 2 || len(manager.lastPlan.Rules) != 1 || manager.lastPlan.Rules[0].ForwardID != otherForwardID {
		t.Errorf("retry did not apply a plan without disabled forward: %+v", manager.lastPlan)
	}
	if got := mustQueryInt(t, h.repo, "SELECT status FROM forward WHERE id = ?", otherForwardID); got != 1 {
		t.Errorf("disabled another user's nftables forward: %d", got)
	}
}

func TestUserDisableReportsPauseFailureAndCanRetry(t *testing.T) {
	for _, offline := range []bool{false, true} {
		t.Run(fmt.Sprintf("offline=%t", offline), func(t *testing.T) {
			a := newCleanupAgent(t)
			now := time.Now().UnixMilli()
			userID, err := a.h.repo.CreateUser("disable-retry", "pwd", 1, now+86400000, 100, 1, 10, 1, 0, now)
			if err != nil {
				t.Fatal(err)
			}
			for _, value := range []interface{}{
				&model.Tunnel{ID: 1, Name: "test", Type: 1, Status: 1, Flow: 1},
				&model.UserTunnel{ID: 10, UserID: userID, TunnelID: 1, Status: 1, Flow: 100, Num: 10},
				&model.Forward{ID: 20, UserID: userID, UserName: "disable-retry", Name: "running", TunnelID: 1, RemoteAddr: "127.0.0.1:80", Status: 1},
				&model.ForwardPort{ForwardID: 20, NodeID: 1, Port: 31020},
			} {
				if err := a.h.repo.DB().Create(value).Error; err != nil {
					t.Fatal(err)
				}
			}
			onlineServer := a.h.wsServer
			if offline {
				a.h.wsServer = ws.NewServer(a.h.repo, "test-jwt-secret")
			} else {
				ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
					response.WriteJSON(w, response.ErrDefault("pause rejected"))
				}))
				defer ts.Close()
				if err := a.h.repo.DB().Model(&model.Node{}).Where("id = 1").Updates(map[string]interface{}{"is_remote": 1, "remote_url": ts.URL, "remote_token": "test-token"}).Error; err != nil {
					t.Fatal(err)
				}
			}
			update := func() response.R {
				t.Helper()
				res := postJSONToHandler(t, a.h.userUpdate, map[string]interface{}{"id": userID, "user": "disable-retry", "status": 0})
				var out response.R
				if err := json.Unmarshal(res.Body.Bytes(), &out); err != nil {
					t.Fatal(err)
				}
				return out
			}
			if out := update(); out.Code == 0 {
				t.Error("pause failure returned success")
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM user WHERE id = ?", userID); got != 0 {
				t.Errorf("failed pause restored account access: %d", got)
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM user_tunnel WHERE id = 10"); got != 0 {
				t.Errorf("failed pause restored tunnel access: %d", got)
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 20"); got != 1 {
				t.Errorf("unconfirmed pause marked completed, status = %d", got)
			}
			a.h.wsServer = onlineServer
			if err := a.h.repo.DB().Model(&model.Node{}).Where("id = 1").Update("is_remote", 0).Error; err != nil {
				t.Fatal(err)
			}
			if out := update(); out.Code != 0 {
				t.Fatalf("retry failed: %s", out.Msg)
			}
			if got := mustQueryInt(t, a.h.repo, "SELECT status FROM forward WHERE id = 20"); got != 0 {
				t.Errorf("retry did not pause forward: %d", got)
			}
			if len(a.commandsOfType("PauseService")) == 0 {
				t.Error("retry never sent PauseService")
			}
		})
	}
}