ModelRouter/internal/core/channels_callback_loopback_test.go

package core

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"llmsproxy/internal/lua"
)

// 三个 callback 平台的回调路径**各是注册值**，且校验规则不同：
//
//	Trae    /authorize      不带 state（token 直接放 query 里）
//	有道    /auth/callback  必须带 state
//	Gemini  /oauth-callback 必须带 state
//
// 路径写错 → 上游 redirect_uri_mismatch；校验写错 → 要么接不住回调、
// 要么被伪造回调塞入凭据。本测试打真上游形态（本地 loopback 接收器）。
func TestCallbackPlatformsLoopback(t *testing.T) {
	c := kernelTestCore(t)
	ps := c.Plugins()
	src, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}

	start := func(t *testing.T, plat string) (m map[string]interface{}, nonce, sid string) {
		t.Helper()
		out, err := ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
			Method: "POST", Body: `{"platform":"` + plat + `"}`, Role: "admin",
		})
		if err != nil {
			t.Fatalf("%s start_login: %v", plat, err)
		}
		m = out.JSON.(map[string]interface{})
		s := m["session"].(map[string]interface{})
		return m, s["nonce"].(string), s["id"].(string)
	}

	poll := func(t *testing.T, sid string) map[string]interface{} {
		t.Helper()
		o, err := ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
			Method: "GET", Query: map[string]string{"id": sid}, Role: "admin",
		})
		if err != nil {
			t.Fatal(err)
		}
		return o.JSON.(map[string]interface{})
	}

	t.Run("行1 Trae 回调路径是 /authorize 且无 state 也能接住", func(t *testing.T) {
		_, _, sid := start(t, "trae")
		// 交换端点指到假上游
		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			w.Write([]byte(`{"Token":"ACCESS-NEW"}`))
		}))
		defer ts.Close()
		if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
			"platform": "trae", "callback": map[string]interface{}{"exchange_url": ts.URL},
		}); err != nil {
			t.Fatal(err)
		}
		// 重新起会话，拿到带 loopback 的 redirect_uri
		m, nonce, sid := start(t, "trae")
		ru, _ := m["redirect_uri"].(string)
		if !strings.HasSuffix(ru, "/authorize") {
			t.Fatalf("trae 回调路径必须是 /authorize，实得 %q", ru)
		}
		// Trae 的回调**不带 state**，而是把 token 放 query 里
		q := url.Values{}
		q.Set("refreshToken", "RT-1")
		q.Set("userInfo", `{"UserID":"u1","ScreenName":"张三"}`)
		resp, err := http.Get(ru + "?" + q.Encode())
		if err != nil {
			t.Fatalf("打 loopback 接收器失败: %v", err)
		}
		resp.Body.Close()
		if resp.StatusCode != 200 {
			t.Fatalf("Trae 回调（无 state）应被接住，实得 %d", resp.StatusCode)
		}
		st := waitState(t, poll, sid, "completed")
		if st["state"] != "completed" {
			t.Fatalf("应完成登录，实得 %v", st)
		}
		raw, ok, _ := ps.GetSecret("channels", "channels:trae")
		if !ok {
			t.Fatal("凭据未落库")
		}
		var cred map[string]interface{}
		if err := json.Unmarshal([]byte(raw), &cred); err != nil {
			t.Fatal(err)
		}
		if cred["access_token"] != "ACCESS-NEW" {
			t.Fatalf("凭据应是交换后的新 access，实得 %v", cred["access_token"])
		}
		if cred["uid"] != "u1" {
			t.Fatalf("uid 应从 userInfo.UserID 取出（签到设备身份的种子）, 实得 %v", cred["uid"])
		}
		_ = nonce
	})

	t.Run("行2 Gemini 路径是 /oauth-callback 且错 state 被拒", func(t *testing.T) {
		m, nonce, sid := start(t, "gemini")
		ru, _ := m["redirect_uri"].(string)
		if !strings.HasSuffix(ru, "/oauth-callback") {
			t.Fatalf("gemini 回调路径必须是 /oauth-callback（参考 GEMINI_CALLBACK_PATH），实得 %q", ru)
		}
		if !strings.HasPrefix(ru, "http://localhost:") {
			t.Fatalf("gemini 回调主机必须是 localhost（参考注释明确），实得 %q", ru)
		}
		// 错 state 必须被接收器拒（400），且不推进会话
		bad, err := http.Get(ru + "?code=C-1&state=WRONG")
		if err != nil {
			t.Fatal(err)
		}
		bad.Body.Close()
		if bad.StatusCode != 400 {
			t.Fatalf("state 不匹配应 400（防伪造回调），实得 %d", bad.StatusCode)
		}
		if st := poll(t, sid); st["state"] != "pending" {
			t.Fatalf("错 state 不应推进会话，实得 %v", st)
		}
		// 正确 state → 完成（交换端点指到假上游）
		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_ = r.ParseForm()
			// 必须带 client_secret：只发 client_id 会 invalid_request，
			// 症状正是「浏览器显示授权成功但账号不出现」
			if r.PostFormValue("client_secret") == "" {
				t.Error("token 交换必须带 client_secret")
			}
			if r.PostFormValue("grant_type") != "authorization_code" {
				t.Errorf("grant_type = %q", r.PostFormValue("grant_type"))
			}
			w.Header().Set("Content-Type", "application/json")
			w.Write([]byte(`{"access_token":"GA-1","refresh_token":"GR-1","expires_in":3600}`))
		}))
		defer ts.Close()
		if _, err := ps.Call("channels", "set_endpoints", map[string]interface{}{
			"platform": "gemini",
			"callback": map[string]interface{}{"token_url": ts.URL, "mode": "code"},
		}); err != nil {
			t.Fatal(err)
		}
		good, err := http.Get(ru + "?code=C-1&state=" + nonce)
		if err != nil {
			t.Fatal(err)
		}
		good.Body.Close()
		st := waitState(t, poll, sid, "completed")
		if st["state"] != "completed" {
			t.Fatalf("应完成登录，实得 %v", st)
		}
		raw, ok, _ := ps.GetSecret("channels", "channels:gemini")
		if !ok {
			t.Fatal("凭据未落库")
		}
		var cred map[string]interface{}
		if err := json.Unmarshal([]byte(raw), &cred); err != nil {
			t.Fatal(err)
		}
		if cred["access_token"] != "GA-1" {
			t.Fatalf("access_token = %v", cred["access_token"])
		}
	})
}

// waitState 轮询到目标状态或超时（loopback 写文件是异步的，需要一拍）。
func waitState(t *testing.T, poll func(*testing.T, string) map[string]interface{},
	sid, want string) map[string]interface{} {
	t.Helper()
	var st map[string]interface{}
	for i := 0; i < 30; i++ {
		st = poll(t, sid)
		if st["state"] != "pending" {
			return st
		}
		time.Sleep(100 * time.Millisecond)
	}
	return st
}

// 占位符未替换检查：授权 URL 里残留 {xxx} 会让上游收到字面量占位符
// （而且往往不报错，只是行为诡异）。三个 callback 平台逐个验。
func TestAuthorizeURLsHaveNoUnfilledPlaceholders(t *testing.T) {
	c := kernelTestCore(t)
	ps := c.Plugins()
	src, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	for _, plat := range []string{"trae", "gemini", "lobsterai"} {
		out, err := ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
			Method: "POST", Body: `{"platform":"` + plat + `"}`, Role: "admin",
		})
		if err != nil {
			t.Fatalf("%s: %v", plat, err)
		}
		m := out.JSON.(map[string]interface{})
		for _, key := range []string{"url", "redirect_uri"} {
			v, _ := m[key].(string)
			if v == "" {
				continue
			}
			if strings.Contains(v, "{") || strings.Contains(v, "nil") {
				t.Errorf("%s 的 %s 仍有未替换占位符: %s", plat, key, v)
			}
		}
	}
}