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) } }