ModelRouter/internal/lua/plugin_login_test.go

package lua

import (
	"testing"
	"time"
)

// 本文件验证 R4-e：插件驱动的登录会话 + 登录 host 授权。
//
// 用户决策是「多流程抽象」（link-poll / callback / paste / VNC 逃生舱）。
// 这里断言的是**框架**：会话生命周期、跨插件隔离、nonce 校验、host 授权
// 与出站白名单的联动。各平台的具体流程属于插件业务。

func loginPlugins(t *testing.T) (*VM, *Plugins, string) {
	return runtimePlugins(t)
}

const loginProbe = `
local plugin = { name = "loginprobe", version = "1.0" }
plugin.exports = { "begin_link", "poll", "start_callback", "cancel_one", "resolve_ok", "read_other" }
plugin.auth_hosts = { "auth.example.com", "Accounts.Google.com:443" }

function plugin.begin_link()
  local s, err = __llmsproxy_login.begin({
    platform = "example", flow = "link-poll",
    url = "https://auth.example.com/device", user_code = "WDJB-MJHT",
  })
  return { s = s, err = err }
end

function plugin.poll(args)
  local s, err = __llmsproxy_login.get(args.id)
  return { s = s, err = err }
end

function plugin.start_callback()
  local s, err = __llmsproxy_login.begin({ platform = "example", flow = "callback" })
  return { s = s, err = err }
end

function plugin.cancel_one(args)
  local ok = __llmsproxy_login.cancel(args.id)
  return { ok = ok }
end

function plugin.resolve_ok(args)
  local ok, err = __llmsproxy_login.resolve(args.id, "completed", "alice", "")
  return { ok = ok, err = err }
end

function plugin.read_other(args)
  local s = __llmsproxy_login.get(args.id)
  return { got = s }
end
return plugin
`

// TestLoginSessionLifecycle：创建 → 轮询 → 完成。
func TestLoginSessionLifecycle(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	out, err := ps.Call("loginprobe", "begin_link", map[string]interface{}{})
	if err != nil {
		t.Fatal(err)
	}
	m := out.(map[string]interface{})
	if m["err"] != nil {
		t.Fatalf("begin 失败: %v", m["err"])
	}
	if m["s"] == nil {
		t.Fatalf("begin 返回 nil session, err=%v", m["err"])
	}
	s := m["s"].(map[string]interface{})
	if s["state"] != "pending" {
		t.Errorf("新会话应为 pending, 实得 %v", s["state"])
	}
	if s["url"] != "https://auth.example.com/device" {
		t.Errorf("URL 未保留: %v", s["url"])
	}
	if s["user_code"] != "WDJB-MJHT" {
		t.Errorf("设备码未保留: %v", s["user_code"])
	}
	id := s["id"].(string)

	// 轮询到 pending
	out, _ = ps.Call("loginprobe", "poll", map[string]interface{}{"id": id})
	got := out.(map[string]interface{})["s"].(map[string]interface{})
	if got["state"] != "pending" {
		t.Errorf("轮询应为 pending, 实得 %v", got["state"])
	}
	// 后续读取不得带 nonce（只在 begin 时给）
	if _, hasNonce := got["nonce"]; hasNonce {
		t.Error("后续读取不应暴露 nonce（CSRF 令牌）")
	}

	// 完成
	out, _ = ps.Call("loginprobe", "resolve_ok", map[string]interface{}{"id": id})
	if out.(map[string]interface{})["ok"] != true {
		t.Fatalf("resolve 失败: %#v", out)
	}
	out, _ = ps.Call("loginprobe", "poll", map[string]interface{}{"id": id})
	fin := out.(map[string]interface{})["s"].(map[string]interface{})
	if fin["state"] != "completed" {
		t.Errorf("应为 completed, 实得 %v", fin["state"])
	}
	if fin["account"] != "alice" {
		t.Errorf("账号标签应保留: %v", fin["account"])
	}
}

// TestLoginSessionIsolatedPerPlugin：插件读不到别人的会话（且与"不存在"不可区分）。
func TestLoginSessionIsolatedPerPlugin(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	other := `
local plugin = { name = "other", version = "1.0" }
plugin.exports = { "read_other" }
function plugin.read_other(args)
  local s = __llmsproxy_login.get(args.id)
  return { got = s }
end
return plugin
`
	if err := ps.LoadSource("other", other); err != nil {
		t.Fatal(err)
	}
	out, _ := ps.Call("loginprobe", "begin_link", map[string]interface{}{})
	id := out.(map[string]interface{})["s"].(map[string]interface{})["id"].(string)

	out, err := ps.Call("other", "read_other", map[string]interface{}{"id": id})
	if err != nil {
		t.Fatal(err)
	}
	if out.(map[string]interface{})["got"] != nil {
		t.Fatal("★ 插件 B 读到了插件 A 的登录会话")
	}
}

// TestLoginCallbackNonce：callback 流程的 nonce 必须能匹配，且只在 begin 时给出。
func TestLoginCallbackNonce(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	out, _ := ps.Call("loginprobe", "start_callback", map[string]interface{}{})
	s := out.(map[string]interface{})["s"].(map[string]interface{})
	nonce, _ := s["nonce"].(string)
	if nonce == "" {
		t.Fatal("callback 流程必须在 begin 时给出 nonce")
	}
	// 内核侧能按 nonce 找到会话（回调路由用它）
	sess, ok := ps.MatchCallbackNonce(nonce)
	if !ok {
		t.Fatal("nonce 应能匹配到会话")
	}
	if sess.ID != s["id"].(string) {
		t.Errorf("匹配到错误的会话: %s vs %s", sess.ID, s["id"])
	}
	// 空 nonce 不得匹配任何东西（否则回调可绕过校验）
	if _, ok := ps.MatchCallbackNonce(""); ok {
		t.Error("空 nonce 不应匹配")
	}
}

// TestLoginHostsGateOutbound 是本次发现的**关键阻碍**的回归测试：
// 登录主机不是源，默认被出站白名单拒绝；只有声明 + 批准后才放行。
func TestLoginHostsGateOutbound(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	// 声明被读取（normalizeHost 会去端口、小写）
	declared := ps.PluginAuthHosts("loginprobe")
	if len(declared) != 2 {
		t.Fatalf("应读到 2 个声明 host, 实得 %#v", declared)
	}
	want := map[string]bool{"auth.example.com": true, "accounts.google.com": true}
	for _, h := range declared {
		if !want[h] {
			t.Errorf("声明 host 未规范化: %q", h)
		}
	}

	// 未批准 → 出站被拒
	if _, err := ps.vm.Transform("noop", "x", "{}"); err == nil {
		_ = err // 不关心，仅为占位
	}
	allow := ps.outboundAllowlist()
	if allow.allows("auth.example.com", "") {
		t.Fatal("★ 未批准的登录 host 不应在放行名单里")
	}

	// 批准后放行
	if err := ps.GrantAuthHosts("loginprobe", []string{"auth.example.com"}); err != nil {
		t.Fatalf("批准失败: %v", err)
	}
	allow = ps.outboundAllowlist()
	if !allow.allows("auth.example.com", "") {
		t.Fatal("批准后应放行")
	}
	if allow.allows("accounts.google.com", "") {
		t.Error("未逐个批准的 host 不应被连带放行")
	}
}

// TestGrantUndeclaredAuthHostRejected：只能批准插件声明过的 host。
func TestGrantUndeclaredAuthHostRejected(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	if err := ps.GrantAuthHosts("loginprobe", []string{"evil.example.net"}); err == nil {
		t.Fatal("★ 批准未声明的 host 应被拒（否则声明只是装饰）")
	}
}

// TestLoginExpiry：超时未完成的会话变 expired（读取时判定，无需后台定时器）。
func TestLoginExpiry(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	if err := ps.LoadSource("loginprobe", loginProbe); err != nil {
		t.Fatal(err)
	}
	out, _ := ps.Call("loginprobe", "begin_link", map[string]interface{}{})
	id := out.(map[string]interface{})["s"].(map[string]interface{})["id"].(string)

	// 人为把过期时间提前
	ps.logins().mu.Lock()
	ps.logins().byID[id].ExpiresAt = time.Now().Add(-time.Second)
	ps.logins().mu.Unlock()

	out, _ = ps.Call("loginprobe", "poll", map[string]interface{}{"id": id})
	s := out.(map[string]interface{})["s"].(map[string]interface{})
	if s["state"] != "expired" {
		t.Fatalf("过期会话应为 expired, 实得 %v", s["state"])
	}
}

// TestVNCRequiresOptIn：VNC 逃生舱默认关闭。
func TestVNCRequiresOptIn(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	code := `
local plugin = { name = "vncprobe", version = "1.0" }
plugin.exports = { "begin_vnc" }
function plugin.begin_vnc()
  local s, err = __llmsproxy_login.begin({ platform = "x", flow = "vnc" })
  return { s = s, err = err }
end
return plugin
`
	if err := ps.LoadSource("vncprobe", code); err != nil {
		t.Fatal(err)
	}
	out, err := ps.Call("vncprobe", "begin_vnc", map[string]interface{}{})
	if err != nil {
		t.Fatal(err)
	}
	m := out.(map[string]interface{})
	if m["err"] == nil || m["s"] != nil {
		t.Fatal("★ VNC 流程默认应被拒绝")
	}
	t.Logf("默认拒绝信息: %v", m["err"])

	// 运维显式开启后可用
	ps.SetVNCEnabled(true)
	out, _ = ps.Call("vncprobe", "begin_vnc", map[string]interface{}{})
	if out.(map[string]interface{})["s"] == nil {
		t.Fatal("开启后 VNC 应可用")
	}
}

// TestUnknownLoginFlowRejected：拼错的流程名必须被拒。
func TestUnknownLoginFlowRejected(t *testing.T) {
	_, ps, _ := loginPlugins(t)
	code := `
local plugin = { name = "badflow", version = "1.0" }
plugin.exports = { "go" }
function plugin.go()
  local s, err = __llmsproxy_login.begin({ platform = "x", flow = "carrier-pigeon" })
  return { s = s, err = err }
end
return plugin
`
	if err := ps.LoadSource("badflow", code); err != nil {
		t.Fatal(err)
	}
	out, _ := ps.Call("badflow", "go", map[string]interface{}{})
	if out.(map[string]interface{})["err"] == nil {
		t.Fatal("未知流程名应被拒")
	}
}