HomeAgent/internal/agent/api/security_premise_test.go

package api

// 前提验证：**安全模块能否稳在 provider 层**。
//
// 设计依据：docs/zh/kernel-secret-firewall-design.md §5.1。
//
// 整件事成立的前提是两条：
//
//	P1. 所有发往 LLM 的调用都经 `Provider.Chat` / `Provider.ChatStream`
//	    （没有绕过接口的直发路径）
//	P2. `ProviderManager.Register` 是所有 provider 的**唯一入口**
//	    ⇒ 在 Register 处包一层装饰器，就能覆盖全部取用路径
//	       （Get / Default / OrderedProviders / ResolveForModel 都从同一 map 取）
//
// 若 P1 或 P2 不成立，安全模块必须下沉到更底层（HTTP 客户端），改动面大得多。
// ⇒ 先验它，再写实现（本仓纪律：前提不成立时不要继续设计上层）。

import (
	"context"
	"encoding/json"
	"path/filepath"
	"strings"
	"testing"

	"github.com/JianFeeeee/HologramAgent/internal/security"
)

// countingProvider 记录它被调用了几次、收到什么消息。
type countingProvider struct {
	calls  int
	seen   []string
	stream int
}

func (c *countingProvider) Name() string { return "counting" }

func (c *countingProvider) Chat(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	c.calls++
	for _, m := range req.Messages {
		c.seen = append(c.seen, m.Content)
	}
	return &CompletionResponse{Content: "ok"}, nil
}

func (c *countingProvider) ChatStream(ctx context.Context, req *CompletionRequest) (<-chan StreamChunk, error) {
	c.stream++
	ch := make(chan StreamChunk, 1)
	close(ch)
	return ch, nil
}

func (c *countingProvider) MaxContextTokens() int { return 8192 }

// ★★ P2：装饰 Register ⇒ **所有取用路径都拿到装饰后的实例**。
//
// 这是安全模块能"一处生效、处处覆盖"的机制保证。
func TestPremise_RegisterDecorateCoversAllAccessPaths(t *testing.T) {
	m := NewProviderManager()

	inner := &countingProvider{}
	var wrapped Provider
	// 模拟安全模块：在 Register 处包一层
	m.Register("p1", &proxyProvider{inner: inner, onCall: func() { wrapped = nil }})
	_ = wrapped

	// ① Get 拿到的是装饰后的
	got := m.Get("p1")
	if got == nil {
		t.Fatal("Get 返回 nil")
	}
	if _, isProxy := got.(*proxyProvider); !isProxy {
		t.Fatalf("Get 未拿到装饰后的实例，实际类型 %T", got)
	}
	// ② Default 也拿到装饰后的
	if _, isProxy := m.Default().(*proxyProvider); !isProxy {
		t.Fatal("Default 未拿到装饰后的实例")
	}
	// ③ OrderedProviders 也是
	ord := m.OrderedProviders()
	if len(ord) != 1 {
		t.Fatalf("OrderedProviders 数量 %d", len(ord))
	}
	if _, isProxy := ord[0].(*proxyProvider); !isProxy {
		t.Fatal("OrderedProviders 未拿到装饰后的实例")
	}
	// ④ ResolveForModel 也是
	res := m.ResolveForModel("")
	if len(res) == 0 {
		t.Fatal("ResolveForModel 为空")
	}
	if _, isProxy := res[0].(*proxyProvider); !isProxy {
		t.Fatal("ResolveForModel 未拿到装饰后的实例")
	}
}

// ★★ P1：经 manager 取到的 provider，其 Chat 必经装饰器。
//
// 这条证明"安全模块能看见每一条出网请求"。
func TestPremise_DecoratorSeesEveryChat(t *testing.T) {
	m := NewProviderManager()
	inner := &countingProvider{}
	var seen int
	m.Register("p1", &proxyProvider{inner: inner, onCall: func() { seen++ }})

	p := m.Get("p1")
	if _, err := p.Chat(context.Background(), &CompletionRequest{
		Messages: []Message{{Role: "user", Content: "hello"}},
	}); err != nil {
		t.Fatal(err)
	}
	if seen != 1 {
		t.Fatalf("装饰器应看到 1 次 Chat，实际 %d", seen)
	}
	if inner.calls != 1 {
		t.Fatalf("内层应被调 1 次，实际 %d", inner.calls)
	}
	// ChatStream 同样要经过
	if _, err := p.ChatStream(context.Background(), &CompletionRequest{}); err != nil {
		t.Fatal(err)
	}
	if seen != 2 {
		t.Fatalf("ChatStream 也应经装饰器，实际 seen=%d", seen)
	}
}

// ★★ 关键能力验证：装饰器**能改写请求**（这正是安全模块要做的替换）。
func TestPremise_DecoratorCanRewriteMessages(t *testing.T) {
	m := NewProviderManager()
	inner := &countingProvider{}
	// 模拟安全模块：把明文换成别名
	m.Register("p1", &proxyProvider{inner: inner, onCall: func() {}})

	// 直接用能改写的装饰器
	m.Register("p2", &rewriteProvider{inner: inner})

	req := &CompletionRequest{Messages: []Message{
		{Role: "user", Content: "token=SECRET-XYZ"},
	}}
	if _, err := m.Get("p2").Chat(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	// 内层必须收到**改写后**的内容（证明改写真的生效、不是只读）
	found := false
	for _, s := range inner.seen {
		if strings.Contains(s, "SECRET-XYZ") {
			t.Fatal("明文被原样透传，改写未生效")
		}
		if strings.Contains(s, "@{sec_") {
			found = true
		}
	}
	if !found {
		t.Fatalf("内层未收到改写后的别名，实际收到: %v", inner.seen)
	}
}

// ★ 覆盖面：ChatStream 也必须能被装饰（流式是主路径之一）。
func TestPremise_StreamAlsoDecorated(t *testing.T) {
	m := NewProviderManager()
	inner := &countingProvider{}
	m.Register("p1", &proxyProvider{inner: inner, onCall: func() {}})
	if _, err := m.Get("p1").ChatStream(context.Background(), &CompletionRequest{}); err != nil {
		t.Fatal(err)
	}
	if inner.stream != 1 {
		t.Fatalf("内层 ChatStream 应被调 1 次，实际 %d", inner.stream)
	}
}

// ---------------- 测试替身 ----------------

// proxyProvider 是"只观察不修改"的装饰器。
type proxyProvider struct {
	inner  Provider
	onCall func()
}

func (p *proxyProvider) Name() string { return p.inner.Name() }
func (p *proxyProvider) Chat(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	if p.onCall != nil {
		p.onCall()
	}
	return p.inner.Chat(ctx, req)
}
func (p *proxyProvider) ChatStream(ctx context.Context, req *CompletionRequest) (<-chan StreamChunk, error) {
	if p.onCall != nil {
		p.onCall()
	}
	return p.inner.ChatStream(ctx, req)
}
func (p *proxyProvider) MaxContextTokens() int { return p.inner.MaxContextTokens() }

// rewriteProvider 是"会改写请求"的装饰器（模拟安全模块）。
type rewriteProvider struct{ inner Provider }

func (p *rewriteProvider) Name() string { return p.inner.Name() }
func (p *rewriteProvider) Chat(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	// 模拟：把明文换成别名（真实实现会用模式库 + HMAC 键）
	for i := range req.Messages {
		if strings.Contains(req.Messages[i].Content, "SECRET-XYZ") {
			req.Messages[i].Content = "token=@{sec_deadbeef}"
		}
	}
	return p.inner.Chat(ctx, req)
}
func (p *rewriteProvider) ChatStream(ctx context.Context, req *CompletionRequest) (<-chan StreamChunk, error) {
	return p.inner.ChatStream(ctx, req)
}
func (p *rewriteProvider) MaxContextTokens() int { return p.inner.MaxContextTokens() }

// ---------------- 接线判据（装饰器真的挂在 Register 上） ----------------

// ★★ 判据：经 Register 登记后，取出的 provider 就是**装饰过的**，
// 且出向明文被换成了别名。
//
// 这是端到端接线判据：它同时验证
//
//	① Register 挂上了装饰器
//	② 装饰器能拿到全局单例
//	③ 替换真的发生了
func TestWired_RegisterDecoratesAndScrubs(t *testing.T) {
	dir := t.TempDir()
	store, err := security.InitGlobal(security.Options{
		Path:    filepath.Join(dir, "vars.json"),
		KeyFn:   func() ([]byte, error) { return []byte("0123456789abcdef0123456789abcdef"), nil },
		Enabled: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { security.SetGlobal(nil) })

	const tok = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	var got string
	inner := &recordingProvider{name: "wired", capture: &got}
	m := NewProviderManager()
	m.Register("p", inner)

	// 取出（任何一条取用路径都应拿到装饰后的实例）
	p := m.Default()
	if p == nil {
		t.Fatal("Default 返回 nil")
	}
	req := &CompletionRequest{Messages: []Message{{Role: "user", Content: "token=" + tok}}}
	if _, err := p.Chat(context.Background(), req); err != nil {
		t.Fatal(err)
	}

	if strings.Contains(got, tok) {
		t.Fatalf("明文泄漏到 provider: %s", got)
	}
	if !strings.Contains(got, "安全模块自动拦截") {
		t.Errorf("未见拦截提示: %s", got)
	}
	if store.Count() != 1 {
		t.Errorf("应登记 1 条变量，实际 %d", store.Count())
	}
	// 记录里不得有明文
	b, _ := json.Marshal(store.List())
	if strings.Contains(string(b), tok) {
		t.Error("变量记录泄露明文")
	}
}

// ★ 判据：安全模块未初始化时，Register **不介入**（零成本）。
func TestWired_NoStoreMeansNoInterference(t *testing.T) {
	security.SetGlobal(nil)
	t.Cleanup(func() { security.SetGlobal(nil) })

	const tok = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	var got string
	m := NewProviderManager()
	m.Register("p", &recordingProvider{name: "plain", capture: &got})

	if _, err := m.Default().Chat(context.Background(),
		&CompletionRequest{Messages: []Message{{Role: "user", Content: "token=" + tok}}}); err != nil {
		t.Fatal(err)
	}
	if got != "token="+tok {
		t.Errorf("未启用安全模块时不得改动内容: %q", got)
	}
}

// ★ 判据：回填把别名换回明文（入向通道）。
func TestWired_FillToolArgs(t *testing.T) {
	dir := t.TempDir()
	store, err := security.InitGlobal(security.Options{
		Path:    filepath.Join(dir, "vars.json"),
		KeyFn:   func() ([]byte, error) { return []byte("0123456789abcdef0123456789abcdef"), nil },
		Enabled: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { security.SetGlobal(nil) })

	const tok = "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	alias, err := store.Submit(tok, "")
	if err != nil {
		t.Fatal(err)
	}

	args := map[string]interface{}{"token": "@{" + alias + "}", "other": "unchanged"}
	n, used := security.FillToolArgs("some_plugin_tool", args)
	if n != 1 {
		t.Fatalf("应回填 1 处，实际 %d", n)
	}
	if args["token"] != tok {
		t.Errorf("回填结果不对: %v", args["token"])
	}
	if args["other"] != "unchanged" {
		t.Error("无关字段不应被改动")
	}
	if len(used) != 1 || used[0] != alias {
		t.Errorf("应报告别名: %v", used)
	}
}

// recordingProvider 记录它收到的最后一条消息内容。
type recordingProvider struct {
	name    string
	capture *string
}

func (r *recordingProvider) Name() string { return r.name }
func (r *recordingProvider) Chat(_ context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	if req != nil && len(req.Messages) > 0 {
		*r.capture = req.Messages[0].Content
	}
	return &CompletionResponse{}, nil
}
func (r *recordingProvider) ChatStream(_ context.Context, req *CompletionRequest) (<-chan StreamChunk, error) {
	if req != nil && len(req.Messages) > 0 {
		*r.capture = req.Messages[0].Content
	}
	ch := make(chan StreamChunk)
	close(ch)
	return ch, nil
}
func (r *recordingProvider) MaxContextTokens() int { return 1000 }