ModelRouter/internal/core/kernel_actions_test.go

package core

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"

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

// 本文件验证 R4-c/R4-d 的闭环：**插件自己驱动外部源的链接认证，然后把源接进内核**。
//
// 这是用户点明的"同类插件"能力（对照 dsh-our-free-model）：插件主导 OAuth 登录、
// 存凭据、定时领额度、把拿到的端点注册成内核上游。此前 llmsproxy 的插件只能管
// 自己的适配器，够不到内核。

// kernelTestCore 起一个最小 core（一个源 + 一个 key）。
func kernelTestCore(t *testing.T) *Core {
	t.Helper()
	dir := t.TempDir()
	pluginDir := filepath.Join(dir, "plugins")
	// 插件的自有存储位置由插件自己决定（读 config.yaml 的 plugin_dir，或
	// LLMSPROXY_PLUGIN_DIR 覆盖）。测试**必须**指向临时目录 —— 否则插件会读到
	// /etc/llmsproxy/config.yaml 并写进**生产插件目录**（真实事故：一次全量测试
	// 把 26 个测试账号写进了生产 .channels-store.json）。
	// 调用方已显式设定时尊重它：多次「重启」要共用同一目录才能测持久化。
	if os.Getenv("LLMSPROXY_PLUGIN_DIR") == "" {
		t.Setenv("LLMSPROXY_PLUGIN_DIR", pluginDir)
	}
	cfg := &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		PluginDir:    pluginDir,
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-kernel-0001"},
		Sources: []config.Source{{
			Name: "primary", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
			APIKey: "secret-primary-key", MaxConcurrent: 1,
			Models: []config.Model{{ID: "m1", Kind: "chat"}},
		}},
	}
	c, err := NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	// 测试结束**必须**把插件拉起的服务端浏览器栈收掉。
	// 为何必需：callback 流（trae/gemini/有道）在 start_login 时就起一整套
	// Xvfb+chromium+x11vnc+websockify，而绝大多数测试不会显式 stop_browser、
	// 会话 TTL 也还没到 —— 于是一轮 core 全包跑完就攒出 7 套常驻进程
	// （实测 websockify×7 + x11vnc×7 + 7 个 profile 目录），把显示器号与端口
	// 段逐步耗光，后续测试报「chromium/websockify 端口未就绪」。
	// 在这一处统一收掉，比让每个测试各自记得更可靠。
	t.Cleanup(func() {
		ps := c.Plugins()
		if ps == nil {
			return
		}
		for _, name := range []string{"channels"} {
			_, _ = ps.InvokeHandler(name, "stop_browser", &lua.HandlerRequest{
				Method: "POST", Role: "admin",
			})
		}
	})
	t.Cleanup(c.Close)
	return c
}

// grant 给插件授权若干动作（模拟运维在设置页点选）。
func grant(t *testing.T, c *Core, plugin string, actions ...string) {
	t.Helper()
	if err := c.Plugins().SetPluginKernelPermission(plugin, actions, false); err != nil {
		t.Fatalf("grant %v to %s: %v", actions, plugin, err)
	}
}

const kernelProbePlugin = `
local plugin = { name = "kernelprobe", version = "1.0" }
plugin.exports = { "read_status", "read_models", "read_sources", "try_delete_key", "probe_count", "call_reload" }
plugin.kernel = { "status", "list_models", "reload" }

function plugin.read_status()
  local st, err = __llmsproxy_kernel.call("status", {})
  if err then return { err = err } end
  return { ready = st.ready, nsources = #(st.sources or {}) }
end

function plugin.read_models()
  local m, err = __llmsproxy_kernel.call("list_models", {})
  if err then return { err = err } end
  return { models = m }
end

-- sources 动作返回的密钥必须是掩码的
function plugin.read_sources()
  local s, err = __llmsproxy_kernel.call("list_sources", {})
  if err then return { err = err } end
  return { sources = s }
end

-- 未授权的破坏性动作必须被拒
function plugin.try_delete_key()
  local _, err = __llmsproxy_kernel.call("delete_key", { key = "sk-kernel-0001" })
  return { err = err }
end

-- 授权数量可见（插件据此优雅降级）
function plugin.probe_count()
  local acts = __llmsproxy_kernel.actions()
  local granted = 0
  for _, a in ipairs(acts) do
    if a.granted then granted = granted + 1 end
  end
  return { total = #acts, granted = granted }
end

function plugin.call_reload()
  local r, err = __llmsproxy_kernel.call("reload", {})
  return { ok = r and r.ok, err = err }
end
return plugin
`

// TestPluginCallsKernelActions 验证插件能调内核动作，且权限边界成立。
func TestPluginCallsKernelActions(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	grant(t, c, "kernelprobe", "status", "list_models", "list_sources")

	// 只读动作可用
	out, err := c.Plugins().Call("kernelprobe", "read_status", map[string]interface{}{})
	if err != nil {
		t.Fatalf("read_status: %v", err)
	}
	m := out.(map[string]interface{})
	if m["err"] != nil {
		t.Fatalf("status 调用失败: %v", m["err"])
	}
	if m["ready"] != true {
		t.Errorf("registry 应已就绪: %#v", m)
	}
	if n, _ := m["nsources"].(float64); n < 1 {
		t.Errorf("应至少有 1 个源, 实得 %v", m["nsources"])
	}

	// 模型列表
	out, err = c.Plugins().Call("kernelprobe", "read_models", map[string]interface{}{})
	if err != nil {
		t.Fatal(err)
	}
	if models, ok := out.(map[string]interface{})["models"].([]interface{}); !ok || len(models) == 0 {
		t.Errorf("模型列表应为非空数组: %#v", out)
	}
}

// TestKernelActionMasksSourceKeys 是安全断言：list_sources **不得**泄漏 api_key。
func TestKernelActionMasksSourceKeys(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	grant(t, c, "kernelprobe", "list_sources")

	out, err := c.Plugins().Call("kernelprobe", "read_sources", map[string]interface{}{})
	if err != nil {
		t.Fatal(err)
	}
	// 把整个返回值序列化后检查：明文密钥绝不能出现在里面
	b, _ := json.Marshal(out)
	if s := string(b); json.Valid(b) && contains(s, "secret-primary-key") {
		t.Fatalf("list_sources 泄漏了源的明文 api_key: %s", s)
	}
	// 但应有可区分性（has_key + 掩码后缀）
	if !contains(string(b), "***") {
		t.Errorf("应给出掩码提示而不是什么都没有: %s", string(b))
	}
}

// TestKernelActionDeniedWithoutGrant 是权限默认拒绝的断言。
func TestKernelActionDeniedWithoutGrant(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	// 故意不授权任何动作
	out, err := c.Plugins().Call("kernelprobe", "try_delete_key", map[string]interface{}{})
	if err != nil {
		t.Fatalf("调用本身不该报错（拒绝应作为返回值）: %v", err)
	}
	m := out.(map[string]interface{})
	if m["err"] == nil || m["err"] == "" {
		t.Fatal("未授权的 delete_key 必须被拒")
	}
	t.Logf("拒绝信息: %v", m["err"])
	// 拒绝信息要点名动作与插件（否则运维无法行动）
	msg, _ := m["err"].(string)
	if !contains(msg, "delete_key") || !contains(msg, "kernelprobe") {
		t.Errorf("拒绝信息应点名动作与插件, 实得 %q", msg)
	}
	// 且密钥必须还在
	if _, ok := c.FindKey("sk-kernel-0001"); !ok {
		t.Fatal("被拒的动作不应真的执行")
	}
}

// TestKernelActionsReportGranted 插件能看到自己拿到了哪些权限。
func TestKernelActionsReportGranted(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	grant(t, c, "kernelprobe", "status", "list_models")

	out, err := c.Plugins().Call("kernelprobe", "probe_count", map[string]interface{}{})
	if err != nil {
		t.Fatal(err)
	}
	m := out.(map[string]interface{})
	total, _ := m["total"].(float64)
	granted, _ := m["granted"].(float64)
	if total < 15 {
		t.Errorf("动作总数应 >=15, 实得 %v", total)
	}
	if granted != 2 {
		t.Errorf("授权数应为 2（status+list_models）, 实得 %v", granted)
	}
}

// TestPluginRegistersSourceIntoKernel 是"同类插件"的核心闭环：
// 插件通过内核动作把一个新的外部源接进网关。
func TestPluginRegistersSourceIntoKernel(t *testing.T) {
	c := kernelTestCore(t)
	code := `
local plugin = { name = "upstream", version = "1.0" }
plugin.exports = { "add_my_source" }
plugin.kernel = { "add_source", "list_sources" }

function plugin.add_my_source(args)
  local src = {
    name = "plugin-provided",
    base_url = "http://127.0.0.1:9/v1",
    adapter = "openai",
    api_key = "plugin-managed-token",
    max_concurrent = 2,
    models = { { id = "plugin-model", kind = "chat", priority = 10 } },
  }
  local r, err = __llmsproxy_kernel.call("add_source", { source = src })
  if err then return { err = err } end
  return { ok = r.ok, name = r.name }
end
return plugin
`
	if err := c.Plugins().LoadSource("upstream", code); err != nil {
		t.Fatal(err)
	}
	grant(t, c, "upstream", "add_source", "list_sources")

	out, err := c.Plugins().Call("upstream", "add_my_source", map[string]interface{}{})
	if err != nil {
		t.Fatalf("add_source: %v", err)
	}
	m := out.(map[string]interface{})
	if m["err"] != nil {
		t.Fatalf("插件注册源失败: %v", m["err"])
	}
	if m["name"] != "plugin-provided" {
		t.Fatalf("返回的源名不对: %#v", m)
	}

	// 内核侧必须真的有了这个源 —— 这是"接进网关"的定义
	found := false
	for _, s := range c.Sources() {
		if s.Name == "plugin-provided" {
			found = true
			if s.BaseURL != "http://127.0.0.1:9/v1" {
				t.Errorf("base_url 未落库: %q", s.BaseURL)
			}
		}
	}
	if !found {
		t.Fatal("插件注册的源没有出现在内核源列表里")
	}
	// 且模型已可路由
	for _, id := range c.Registry().ModelList() {
		if id == "plugin-model" {
			return // 通过
		}
	}
	t.Errorf("插件注册的模型未进入可路由列表: %v", c.Registry().ModelList())
}

// TestKernelPermissionPersists 授权必须落盘（否则重启即失效）。
func TestKernelPermissionPersists(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	grant(t, c, "kernelprobe", "status", "reload")

	pdir := c.Config().PluginDir
	path := filepath.Join(pdir, ".kernel-perms.json")
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("授权未落盘: %v", err)
	}
	if !contains(string(b), "reload") {
		t.Errorf("落盘内容缺少动作名: %s", b)
	}

	// 模拟重启：新建 Plugins、加载授权、权限仍在
	vm2 := lua.NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm2.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm2.Stop()
	ps2 := lua.NewPlugins(vm2, pdir)
	if err := ps2.LoadKernelPermissions(); err != nil {
		t.Fatal(err)
	}
	perm := ps2.PluginKernelPermission("kernelprobe")
	if len(perm.Actions) != 2 {
		t.Fatalf("重启后授权应保留 2 项, 实得 %#v", perm.Actions)
	}
}

// TestUnknownKernelActionRejected 拼错的动作名必须在授权时就被拒，
// 而不是静默授予"什么都不允许"。
func TestUnknownKernelActionRejected(t *testing.T) {
	c := kernelTestCore(t)
	if err := c.Plugins().LoadSource("kernelprobe", kernelProbePlugin); err != nil {
		t.Fatal(err)
	}
	if err := c.Plugins().SetPluginKernelPermission("kernelprobe", []string{"no_such_action"}, false); err == nil {
		t.Fatal("未知动作名应被拒")
	}
}

// TestKernelActionSurfaceExcludesPlaintextKeys 是凭据边界的回归断言：
// 动作面里绝不能出现返回明文密钥的动作（ListKeys/FindKey 被有意排除）。
func TestKernelActionSurfaceExcludesPlaintextKeys(t *testing.T) {
	for _, a := range lua.KernelActions() {
		switch a.Name {
		case "list_keys", "find_key", "get_key":
			t.Fatalf("动作面出现了明文密钥读取动作 %q —— 这会把部署的凭据面交给插件", a.Name)
		}
	}
}

func contains(haystack, needle string) bool {
	return len(needle) > 0 && len(haystack) >= len(needle) &&
		func() bool {
			for i := 0; i+len(needle) <= len(haystack); i++ {
				if haystack[i:i+len(needle)] == needle {
					return true
				}
			}
			return false
		}()
}