ModelRouter/internal/gateway/mount_test.go

package gateway

import (
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"modelrouter/internal/config"
	"modelrouter/internal/core"
)

// 本文件锁定 R3 的两条不变量：
//  1. 自定义挂载只能落在允许的命名空间内
//  2. 任何内核路由都不能被插件声明覆盖 —— 且这条在新增内核路由时**不会失效**

// TestValidateCustomMountNamespace 只有 /plugins/ 与 /v1/plugins/ 可被占用。
func TestValidateCustomMountNamespace(t *testing.T) {
	ok := []string{
		"/plugins/auth/token",
		"/plugins/auth",
		"/v1/plugins/callback",
		"/v1/plugins/callback/recv",
		"/plugins/nested/deep/path",
	}
	for _, p := range ok {
		if _, err := ValidateCustomMount(p); err != nil {
			t.Errorf("合法挂载 %q 被拒: %v", p, err)
		}
	}
	bad := []string{
		"",
		"plugins/x",            // 没有前导斜杠
		"/v1/chat/completions", // 影子内核路由（且不在允许命名空间）
		"/api/anything",        // 内核命名空间
		"/callback",            // 任意前缀
		"/plugins/../etc",      // 逃逸
		"/plugins//double",     // 空段
		"/v1/plugins/../../x",
	}
	for _, p := range bad {
		if _, err := ValidateCustomMount(p); err == nil {
			t.Errorf("非法挂载 %q 被接受", p)
		}
	}
}

// TestCustomMountCannotShadowKernelRoute 逐条核对内核保留前缀。
func TestCustomMountCannotShadowKernelRoute(t *testing.T) {
	// 这些是要保护的公开接口；任何能覆盖它们的挂载都必须被拒。
	for _, kernel := range []string{"/v1/chat/completions", "/v1/messages", "/v1/models"} {
		// 尝试用允许命名空间里的路径去"贴近"内核路由 —— 例如
		// /v1/plugins/../../v1/models（逃逸），必须被拒。
		cand := "/v1/plugins/../../" + strings.TrimPrefix(kernel, "/")
		if _, err := ValidateCustomMount(cand); err == nil {
			t.Errorf("逃逸路径 %q 应被拒", cand)
		}
	}
}

// TestKernelRoutesAreAllCoveredByReservedPrefixes 是这道防线**不腐坏**的保证。
//
// 它从 server.go 里把路由 switch 的分支字符串刮出来，逐个检查是否被
// kernelReservedPrefixes 覆盖。将来有人新增一条内核路由却忘了更新保留表时，
// 这条测试会失败 —— 否则新路由会静默变成可被插件覆盖的。
func TestKernelRoutesAreAllCoveredByReservedPrefixes(t *testing.T) {
	src, err := os.ReadFile("server.go")
	if err != nil {
		t.Fatal(err)
	}
	// 抓 switch 里的字面量路径：case r.URL.Path == "/x" / HasPrefix(..., "/x")
	re := regexp.MustCompile(`r\.URL\.Path\s*==\s*"(/[^"]*)"`)
	matches := re.FindAllStringSubmatch(string(src), -1)
	if len(matches) == 0 {
		t.Fatal("没有从 server.go 抓到任何路由字面量 —— 测试自身失效")
	}
	covered := func(p string) bool {
		// "/" is the UI catch-all (serveUI). It is the default branch, not a
		// path a plugin could shadow, and it is covered by the namespace rule:
		// a custom mount must start with /plugins/ or /v1/plugins/, neither of
		// which is "/".
		if p == "/" {
			return true
		}
		for _, res := range kernelReservedPrefixes {
			if res == "/" {
				continue
			}
			rp := strings.TrimRight(res, "/")
			if p == rp || strings.HasPrefix(p, rp+"/") {
				return true
			}
		}
		// /api/plugins/... 由 pluginHTTPPrefix 与 /api/ 覆盖
		if strings.HasPrefix(p, "/api/") {
			return true
		}
		return false
	}
	for _, m := range matches {
		p := m[1]
		if !covered(p) {
			t.Errorf("内核路由 %q 未被 kernelReservedPrefixes 覆盖 —— "+
				"新增路由时必须同步更新保留表，否则插件可影子它", p)
		}
	}
}

// TestMountConflictDetectionBothDirections 冲突检测必须双向。
//
// 只查一个方向会让「先注册 /plugins/x，后注册 /plugins/x/sub」或反过来时的
// 一个被另一个静默吞掉。
func TestMountConflictDetectionBothDirections(t *testing.T) {
	table := []RouteOwner{{Path: "/plugins/x", Owner: "plugin:a"}}
	// 被包含
	if _, conflict := FindMountConflict("/plugins/x/sub", table); !conflict {
		t.Error("嵌套在已有挂载下的路径应判冲突")
	}
	// 包含
	if _, conflict := FindMountConflict("/plugins", table); !conflict {
		t.Error("包含已有挂载的路径应判冲突")
	}
	// 相同
	if _, conflict := FindMountConflict("/plugins/x", table); !conflict {
		t.Error("完全相同的路径应判冲突")
	}
	// 无关（前缀相似但不是路径段）
	if _, conflict := FindMountConflict("/plugins/xy", table); conflict {
		t.Error("/plugins/xy 与 /plugins/x 不是同一个路径段，不该判冲突")
	}
}

// TestMountConflictNamesTheOwner 冲突信息必须点名持有者（否则不可行动）。
func TestMountConflictNamesTheOwner(t *testing.T) {
	table := []RouteOwner{{Path: "/plugins/taken", Owner: "plugin:other"}}
	owner, conflict := FindMountConflict("/plugins/taken/deep", table)
	if !conflict {
		t.Fatal("应判冲突")
	}
	if owner.Owner != "plugin:other" {
		t.Errorf("应点名持有者 plugin:other, 实得 %q", owner.Owner)
	}
}

// TestParseMountRef 解析两种命名空间。
func TestParseMountRef(t *testing.T) {
	cases := []struct {
		in, owner, rest string
		ok              bool
	}{
		{"/plugins/auth/token", "auth", "token", true},
		{"/plugins/auth", "auth", "", true},
		{"/v1/plugins/cb/recv", "cb", "recv", true},
		{"/v1/plugins/cb", "cb", "", true},
		{"/plugins/", "", "", false},
		{"/plugins", "", "", false},
		{"/other/x", "", "", false},
	}
	for _, c := range cases {
		owner, rest, ok := parseMountRef(c.in)
		if ok != c.ok || owner != c.owner || rest != c.rest {
			t.Errorf("parseMountRef(%q) = (%q,%q,%v), want (%q,%q,%v)",
				c.in, owner, rest, ok, c.owner, c.rest, c.ok)
		}
	}
}

// TestBuildRouteTableIncludesKernel 路由表含内核保留项（供冲突提示用）。
func TestBuildRouteTableIncludesKernel(t *testing.T) {
	table := BuildRouteTable(
		map[string]string{"/plugins/a": "alpha"},
		map[string]string{"/v1/plugins/b": "beta"},
	)
	var hasKernel, hasPlugin, hasAdapter bool
	for _, r := range table {
		if r.Owner == "kernel" {
			hasKernel = true
		}
		if r.Owner == "plugin:alpha" {
			hasPlugin = true
		}
		if r.Owner == "adapter:beta" {
			hasAdapter = true
		}
	}
	if !hasKernel || !hasPlugin || !hasAdapter {
		t.Errorf("路由表应含 kernel/plugin/adapter 三类: %+v", table)
	}
}

// TestPluginServesCustomMountEndToEnd 是 R3 的端到端验收：
// 插件声明一个自定义对外挂载点，客户端按该路径访问，得到插件的应答。
func TestPluginServesCustomMountEndToEnd(t *testing.T) {
	ts, key := messagesGateway(t)
	// messagesGateway 用的 core 没有插件目录配置，这里直接建一个带插件的。
	_ = ts
	_ = key

	dir := t.TempDir()
	cfg := &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		PluginDir:    filepath.Join(dir, "plugins"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-mount-0001"},
		Sources: []config.Source{{
			Name: "s", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
			APIKey: "k", MaxConcurrent: 1,
			Models: []config.Model{{ID: "m", Kind: "chat"}},
		}},
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	gw, err := New(c)
	if err != nil {
		t.Fatal(err)
	}
	ts2 := httptest.NewServer(gw.Handler())
	defer ts2.Close()

	// 装一个声明自定义挂载的插件
	code := `
local plugin = { name = "publicapi", version = "1.0" }
plugin.http = {
  { path = "publicapi/token", method = "GET", handler = "token",
    roles = { "admin", "user" }, description = "对外 token 端点" },
}
function plugin.token(req)
  return { status = 200, body = '{"token":"abc"}' }
end
return plugin
`
	if err := c.Plugins().LoadSource("publicapi", code); err != nil {
		t.Fatalf("load plugin: %v", err)
	}

	// 按自定义挂载访问
	req, _ := http.NewRequest("GET", ts2.URL+"/plugins/publicapi/token", nil)
	req.Header.Set("Authorization", "Bearer "+cfg.GatewayKeys[0])
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		t.Fatalf("自定义挂载 status = %d: %s", resp.StatusCode, body)
	}
	if !strings.Contains(string(body), "abc") {
		t.Errorf("自定义挂载未返回插件应答: %s", body)
	}

	// /v1/plugins/ 形式同样可用
	req2, _ := http.NewRequest("GET", ts2.URL+"/v1/plugins/publicapi/token", nil)
	req2.Header.Set("Authorization", "Bearer "+cfg.GatewayKeys[0])
	resp2, err := http.DefaultClient.Do(req2)
	if err != nil {
		t.Fatal(err)
	}
	defer resp2.Body.Close()
	if resp2.StatusCode != 200 {
		t.Errorf("/v1/plugins/ 形式 status = %d", resp2.StatusCode)
	}

	// 未声明的命名空间 404（不能因为落在允许命名空间就放行）
	req3, _ := http.NewRequest("GET", ts2.URL+"/plugins/nobody/here", nil)
	req3.Header.Set("Authorization", "Bearer "+cfg.GatewayKeys[0])
	resp3, err := http.DefaultClient.Do(req3)
	if err != nil {
		t.Fatal(err)
	}
	defer resp3.Body.Close()
	if resp3.StatusCode != 404 {
		t.Errorf("未注册的挂载应 404, 实得 %d", resp3.StatusCode)
	}

	// 鉴权仍生效：无 key 时，
	//   - API 路径（/v1/plugins/…）应 401（isAPIPath 判定）
	//   - 浏览器路径（/plugins/…）应 302 跳登录页
	req4, _ := http.NewRequest("GET", ts2.URL+"/v1/plugins/publicapi/token", nil)
	noRedirect := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error {
		return http.ErrUseLastResponse
	}}
	resp4, err := noRedirect.Do(req4)
	if err != nil {
		t.Fatal(err)
	}
	defer resp4.Body.Close()
	if resp4.StatusCode != 401 {
		t.Errorf("无 key 访问 /v1/plugins/ 应 401, 实得 %d", resp4.StatusCode)
	}

	req5, _ := http.NewRequest("GET", ts2.URL+"/plugins/publicapi/token", nil)
	resp5, err := noRedirect.Do(req5)
	if err != nil {
		t.Fatal(err)
	}
	defer resp5.Body.Close()
	if resp5.StatusCode != 302 {
		t.Errorf("无 key 访问 /plugins/ 应 302 跳登录, 实得 %d", resp5.StatusCode)
	}
	if loc := resp5.Header.Get("Location"); !strings.Contains(loc, "/login") {
		t.Errorf("应跳转到 /login, 实得 %q", loc)
	}
}

// TestCustomMountCannotInterceptKernelRoute 端到端确认内核路由不受影响：
// 即使有插件在用 /plugins/，/v1/chat/completions 与 /v1/models 仍归内核。
func TestCustomMountCannotInterceptKernelRoute(t *testing.T) {
	for _, p := range []string{"/v1/chat/completions", "/v1/models", "/v1/messages", "/api/stats"} {
		if _, err := ValidateCustomMount(p); err == nil {
			t.Errorf("内核路由 %q 必须不能被声明为自定义挂载", p)
		}
	}
}