ModelRouter/internal/core/channels_stress_test.go

package core

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
)

// mustBox 构造一个真 SecretBox（与运行期同一实现，密钥落 runtime dir）。
func mustBox(t *testing.T, runtimeFile string) *config.SecretBox {
	t.Helper()
	box, err := config.NewSecretBox(runtimeFile)
	if err != nil {
		t.Fatal(err)
	}
	return box
}

// channelsCore 起一个带 channels 插件、可指定源码与 box 的核心（压力测试用）。
func channelsCore(t *testing.T, src string) (*Core, *lua.Plugins) {
	t.Helper()
	c := kernelTestCore(t)
	ps := c.Plugins()
	if err := ps.LoadSource("channels", src); err != nil {
		t.Fatalf("加载 channels: %v", err)
	}
	ps.SetSecretBox(mustBox(t, filepath.Join(t.TempDir(), "runtime.json")))
	_ = ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)
	return c, ps
}

// TestChannelsConcurrentLogins 压力：并发走 device-poll 登录。
//
// 插件是单 VM 单锁 —— 本测试验证的是**正确性下的并发**（会话互不串台、
// 凭据互不可见、最终都 completed），真吞吐压力在 e2e 层（真二进制+HTTP 并发）。
func TestChannelsConcurrentLogins(t *testing.T) {
	var hits int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/oauth2/device/code":
			n := atomic.AddInt32(&hits, 1)
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprintf(w, `{"device_code":"DEV-%d","user_code":"UC-%d","verification_uri":"https://example/activate","interval":1,"expires_in":600}`, n, n)
		case "/oauth2/token":
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprint(w, `{"access_token":"TOK","refresh_token":"RF","nickname":"u"}`)
		default:
			w.WriteHeader(404)
		}
	}))
	defer ts.Close()

	_, ps := channelsCore(t, mustChannelsSrc(t))
	// 把 minimax 端点指到假上游
	if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "minimax",
		"device": map[string]interface{}{
			"code_url":  ts.URL + "/oauth2/device/code",
			"token_url": ts.URL + "/oauth2/token",
		},
	}); err != nil {
		t.Fatal(err)
	}

	const rounds = 12
	var wg sync.WaitGroup
	var failures int32
	for i := 0; i < rounds; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			out, err := ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
				Method: "POST", Body: `{"platform":"minimax"}`,
			})
			if err != nil {
				t.Errorf("start_login #%d: %v", n, err)
				atomic.AddInt32(&failures, 1)
				return
			}
			j := out.JSON.(map[string]interface{})
			sess, _ := j["session"].(map[string]interface{})
			if sess == nil {
				t.Errorf("start_login #%d 无会话: %v", n, j)
				atomic.AddInt32(&failures, 1)
				return
			}
			sid := sess["id"].(string)
			for k := 0; k < 10; k++ {
				out2, _ := ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
					Method: "GET", Query: map[string]string{"id": sid},
				})
				m2 := out2.JSON.(map[string]interface{})
				if m2["state"] == "completed" {
					return
				}
				if m2["state"] == "failed" {
					t.Errorf("会话 #%d failed: %v", n, m2)
					atomic.AddInt32(&failures, 1)
					return
				}
			}
			t.Errorf("会话 #%d 10 次轮询未终态", n)
			atomic.AddInt32(&failures, 1)
		}(i)
	}
	wg.Wait()
	if failures != 0 {
		t.Fatalf("并发登录失败 %d 轮", failures)
	}
}

// TestChannelsClaimStress 压力：同一平台并发领签到（幂等 + 无登录态丢失）。
func TestChannelsClaimStress(t *testing.T) {
	var claims int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/v1/checkin" {
			atomic.AddInt32(&claims, 1)
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprint(w, `{"ok":true,"credits":10}`)
			return
		}
		w.WriteHeader(404)
		fmt.Fprint(w, `{}`)
	}))
	defer ts.Close()

	_, ps := channelsCore(t, mustChannelsSrc(t))
	if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "cookieish",
		"checkin":  map[string]interface{}{"url": ts.URL + "/v1/checkin"},
	}); err != nil {
		t.Fatal(err)
	}
	if err := ps.SetSecret("channels", "channels:cookieish", "PASTE-TOKEN"); err != nil {
		t.Fatal(err)
	}

	const n = 30
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(k int) {
			defer wg.Done()
			out, err := ps.InvokeHandler("channels", "claim_daily", &lua.HandlerRequest{
				Method: "POST", Body: `{"platform":"cookieish"}`,
			})
			if err != nil {
				t.Errorf("claim #%d: %v", k, err)
				return
			}
			j := out.JSON.(map[string]interface{})
			if j["ok"] != true {
				t.Errorf("claim #%d 未成功: %v", k, j)
			}
		}(i)
	}
	wg.Wait()
	if got := atomic.LoadInt32(&claims); got != n {
		t.Errorf("上游应收到 %d 次领取, 实得 %d", n, got)
	}
}

// TestLegacyVaultPlaintextLoad 旧通路：老版本（无加密）写的明文库在新版本
// （带 box）下必须能读 —— 凭据升级不丢。落盘则会在下次写时整体重封为密文。
func TestLegacyVaultPlaintextLoad(t *testing.T) {
	dir := t.TempDir()
	ps := lua.NewPlugins(nil, dir)
	if err := ps.SetSecret("legacy", "channels:old", "PLAIN-OLD-TOKEN"); err != nil {
		t.Fatal(err)
	}
	ps.SetSecretBox(mustBox(t, filepath.Join(dir, "runtime.json")))
	v, ok, err := ps.GetSecret("legacy", "channels:old")
	if err != nil || !ok || v != "PLAIN-OLD-TOKEN" {
		t.Fatalf("旧明文库读回失败: v=%q ok=%v err=%v", v, ok, err)
	}
	// 任意一次写触发整体重封
	if err := ps.SetSecret("legacy", "k2", "NEW"); err != nil {
		t.Fatal(err)
	}
	raw, err := os.ReadFile(filepath.Join(dir, ".secrets", "legacy.json"))
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(raw), "PLAIN-OLD-TOKEN") {
		t.Fatalf("★ 重封后旧明文仍在盘上:\n%s", raw)
	}
	if !strings.Contains(string(raw), "enc:v1:") {
		t.Fatalf("重封后应含密文:\n%s", raw)
	}
}

// TestLegacySourceSealedConfig 旧通路：config.yaml 里明文 api_key 的旧配置，
// 新版本加载后运行期拿到解密值、落盘被 seal（SealIfNeeded 语义）。
func TestLegacySourceSealedConfig(t *testing.T) {
	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	yaml := "listen: 127.0.0.1:0\n" +
		"gateway_keys:\n  - sk-legacy-0001\n" +
		"runtime_file: " + filepath.Join(dir, "runtime.json") + "\n" +
		"adapter_dir: " + filepath.Join(dir, "adapters") + "\n" +
		"sources:\n" +
		"  - name: legacy-src\n" +
		"    base_url: https://legacy.example.com/v1\n" +
		"    adapter: openai\n" +
		"    api_key: PLAIN-LEGACY-KEY\n" +
		"    max_concurrent: 2\n" +
		"    models:\n" +
		"      - id: m1\n" +
		"        kind: chat\n" +
		"        priority: 10\n"
	if err := os.WriteFile(cfgPath, []byte(yaml), 0o644); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	c, err := NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	got := c.Sources()
	if len(got) != 1 || got[0].APIKey != "PLAIN-LEGACY-KEY" {
		t.Fatalf("运行期 api_key 应解密: %#v", got)
	}
	raw, err := os.ReadFile(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(raw), "PLAIN-LEGACY-KEY") {
		t.Fatalf("★ 旧明文 api_key 升级后仍明文落盘:\n%s", raw)
	}
	if !strings.Contains(string(raw), "enc:v1:") {
		t.Fatalf("落盘应含 enc:v1 密文:\n%s", raw)
	}
	mk, err := os.ReadFile(filepath.Join(dir, "master.key"))
	if err != nil || len(mk) != 64 {
		t.Fatalf("master.key 应 64 hex, len=%d err=%v", len(mk), err)
	}
}