ModelRouter/internal/lua/secrets_crypto_test.go

package lua

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
)

// ---- 测试用加密桩：与 config.SecretBox 同语义（AEAD + enc:v1 前缀）----

const testEncPrefix = "enc:v1:"

// realBox：真 AEAD-GCM（与 config.SecretBox 同原语），行为一致。
type realBox struct {
	aead cipher.AEAD
}

func newRealBox(t *testing.T) *realBox {
	t.Helper()
	digest := sha256.Sum256([]byte("test-master-key"))
	block, err := aes.NewCipher(digest[:])
	if err != nil {
		t.Fatal(err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		t.Fatal(err)
	}
	return &realBox{aead: aead}
}

func (b *realBox) Encrypt(plain string) (string, error) {
	nonce := make([]byte, b.aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	sealed := b.aead.Seal(nonce, nonce, []byte(plain), nil)
	return testEncPrefix + b64RawURL(sealed), nil
}

func (b *realBox) Decrypt(v string) (string, error) {
	if !strings.HasPrefix(v, testEncPrefix) {
		return v, nil
	}
	raw, err := b64RawURLDecode(strings.TrimPrefix(v, testEncPrefix))
	if err != nil {
		return v, err
	}
	n := b.aead.NonceSize()
	if len(raw) < n {
		return v, os.ErrInvalid
	}
	plain, err := b.aead.Open(nil, raw[:n], raw[n:], nil)
	if err != nil {
		return v, err
	}
	return string(plain), nil
}

// failBox 永远解密失败（模拟 master.key 被换/丢失）。
type failBox struct{}

func (failBox) Encrypt(plain string) (string, error) { return testEncPrefix + "AAAA", nil }
func (failBox) Decrypt(v string) (string, error) {
	if !strings.HasPrefix(v, testEncPrefix) {
		return v, nil
	}
	return v, os.ErrPermission
}

func b64RawURL(b []byte) string { return base64.RawURLEncoding.EncodeToString(b) }
func b64RawURLDecode(s string) ([]byte, error) {
	return base64.RawURLEncoding.DecodeString(s)
}

// TestSecretVaultEncryptedAtRest 验证保险库落盘是密文（enc:v1 前缀），
// 且读回时解密为原值。明文 token 绝不能出现在磁盘字节里。
func TestSecretVaultEncryptedAtRest(t *testing.T) {
	dir := t.TempDir()
	ps := NewPlugins(nil, dir)
	ps.SetSecretBox(newRealBox(t))

	if err := ps.SetSecret("channels", "channels:minimax", "SUPER-TOKEN-123"); err != nil {
		t.Fatalf("SetSecret: %v", err)
	}

	// 磁盘字节必须不含明文
	raw, err := os.ReadFile(filepath.Join(dir, ".secrets", "channels.json"))
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(raw), "SUPER-TOKEN-123") {
		t.Fatalf("保险库明文落盘! 文件内容含 token:\n%s", raw)
	}
	if !strings.Contains(string(raw), testEncPrefix) {
		t.Fatalf("落盘值应带 %s 前缀（统一密文形态）:\n%s", testEncPrefix, raw)
	}

	// 读回是原值
	v, ok, err := ps.GetSecret("channels", "channels:minimax")
	if err != nil || !ok || v != "SUPER-TOKEN-123" {
		t.Fatalf("round-trip 失败: v=%q ok=%v err=%v", v, ok, err)
	}

	// 文件权限 0600、目录 0700
	if fi, err := os.Stat(filepath.Join(dir, ".secrets", "channels.json")); err != nil || fi.Mode().Perm() != 0o600 {
		t.Fatalf("vault 文件权限应 0600, 实得 %v (err=%v)", fi.Mode(), err)
	}
	if fi, err := os.Stat(filepath.Join(dir, ".secrets")); err != nil || fi.Mode().Perm() != 0o700 {
		t.Fatalf("vault 目录权限应 0700, 实得 %v (err=%v)", fi.Mode(), err)
	}
}

// TestSecretVaultPlaintextFallbackAndUpgrade：旧明文库可读（向后兼容），
// 且下一次写会整体升级为密文。
func TestSecretVaultPlaintextFallbackAndUpgrade(t *testing.T) {
	dir := t.TempDir()
	ps := NewPlugins(nil, dir)
	// 1) 无 box 时写入 = 旧版本形态（明文）
	if err := ps.SetSecret("channels", "k", "OLD-PLAIN"); err != nil {
		t.Fatal(err)
	}
	// 2) 换上有 box 的新版本：旧明文必须能读出来
	ps.SetSecretBox(newRealBox(t))
	v, ok, err := ps.GetSecret("channels", "k")
	if err != nil || !ok || v != "OLD-PLAIN" {
		t.Fatalf("旧明文读回应原值: v=%q ok=%v err=%v", v, ok, err)
	}
	// 3) 任意一次写（哪怕写别的键）会把整个文件升级为密文
	if err := ps.SetSecret("channels", "k2", "NEW"); err != nil {
		t.Fatal(err)
	}
	raw, _ := os.ReadFile(filepath.Join(dir, ".secrets", "channels.json"))
	if strings.Contains(string(raw), "OLD-PLAIN") {
		t.Fatalf("升级写后旧明文值仍在盘上:\n%s", raw)
	}
	var m map[string]string
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(m["k"], testEncPrefix) {
		t.Fatalf("旧值应被重封为密文: %v", m["k"])
	}
	// 4) 两个值都还能读回
	for key, want := range map[string]string{"k": "OLD-PLAIN", "k2": "NEW"} {
		if v, _, _ := ps.GetSecret("channels", key); v != want {
			t.Fatalf("升级后 %q: got %q want %q", key, v, want)
		}
	}
}

// TestSecretVaultWrongMasterKey：master.key 换掉后读库必须**报错**而不是静默
// 返回空 —— 空会让插件重新登录并覆盖文件，把可恢复问题变成凭据丢失。
func TestSecretVaultWrongMasterKey(t *testing.T) {
	dir := t.TempDir()
	ps := NewPlugins(nil, dir)
	ps.SetSecretBox(newRealBox(t))
	if err := ps.SetSecret("channels", "k", "SECRET"); err != nil {
		t.Fatal(err)
	}
	// 换一把钥匙（模拟 master.key 丢失重建）
	ps.SetSecretBox(failBox{})
	if _, _, err := ps.GetSecret("channels", "k"); err == nil {
		t.Fatal("换 master key 后读库未报错 —— 会静默吞凭据")
	}
}

// ---- 并发压力：多插件同时读写各自的库 ----

// TestSecretVaultConcurrentPlugins 压力：8 个插件 × 每插件 2 goroutine 并发
// set/get/delete 自己的键。断言：无数据串台（跨插件不可见）、无丢写、文件数正确。
func TestSecretVaultConcurrentPlugins(t *testing.T) {
	dir := t.TempDir()
	ps := NewPlugins(nil, dir)
	ps.SetSecretBox(newRealBox(t))

	const plugins = 8
	const perG = 50
	var wg sync.WaitGroup
	for p := 0; p < plugins; p++ {
		name := "plug" + string(rune('A'+p))
		for g := 0; g < 2; g++ {
			wg.Add(1)
			go func(name string, g int) {
				defer wg.Done()
				for i := 0; i < perG; i++ {
					key := name + ":k" + string(rune('a'+g))
					if err := ps.SetSecret(name, key, name+"-token"); err != nil {
						t.Errorf("SetSecret(%s,%s): %v", name, key, err)
						return
					}
					v, ok, err := ps.GetSecret(name, key)
					if err != nil || !ok || v != name+"-token" {
						t.Errorf("GetSecret(%s,%s)=%q ok=%v err=%v", name, key, v, ok, err)
						return
					}
					// 跨插件必须不可见
					other := "plug" + string(rune('A'+(p+1)%plugins))
					if _, ok, _ := ps.GetSecret(other, key); ok {
						t.Errorf("凭据串台: %s 的键在 %s 可见", name, other)
						return
					}
					if g == 1 && i%10 == 9 {
						_ = ps.DeleteSecret(name, key)
					}
				}
			}(name, g)
		}
	}
	wg.Wait()

	entries, err := os.ReadDir(filepath.Join(dir, ".secrets"))
	if err != nil {
		t.Fatal(err)
	}
	if len(entries) > plugins {
		t.Errorf("vault 文件数 %d 超过插件数 %d（应一插件一文件）", len(entries), plugins)
	}
}