ModelRouter/internal/config/secret_config_test.go

package config

import (
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// writeMasterKey plants a deterministic master.key next to the given runtime
// file so tests exercise the real box instead of the env-var shortcut.
func writeMasterKey(t *testing.T, runtimeFile string) *SecretBox {
	t.Helper()
	if err := os.WriteFile(filepath.Join(filepath.Dir(runtimeFile), "master.key"),
		[]byte(strings.Repeat("ab", 32)), 0600); err != nil {
		t.Fatal(err)
	}
	box, err := NewSecretBox(runtimeFile)
	if err != nil {
		t.Fatalf("NewSecretBox: %v", err)
	}
	return box
}

func TestSaveSealsCredentialsAndLoadUnseals(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")
	runtime := filepath.Join(dir, "runtime.json")
	box := writeMasterKey(t, runtime)

	cfg := &Config{
		Path:   path,
		Listen: "127.0.0.1:0",
		Sources: []Source{{
			Name:    "up",
			BaseURL: "http://up/v1",
			APIKey:  "sk-plaintext-secret",
			Adapter: "openai",
			Headers: map[string]string{"X-Extra": "header-secret"},
			Models:  []Model{{ID: "m", Kind: "chat"}},
		}},
		Keys: []GWKey{{Key: "sk-gw-plain", Role: "admin", Name: "admin"}},
	}
	cfg.AttachSecretBox(box)
	if err := cfg.Save(); err != nil {
		t.Fatalf("Save: %v", err)
	}

	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	text := string(raw)
	if strings.Contains(text, "sk-plaintext-secret") {
		t.Error("source api_key written in plaintext")
	}
	if strings.Contains(text, "header-secret") {
		t.Error("source header written in plaintext")
	}
	if strings.Contains(text, "sk-gw-plain") {
		t.Error("gateway key written in plaintext")
	}
	if n := strings.Count(text, encPrefix); n != 3 {
		t.Errorf("expected 3 sealed values, found %d", n)
	}

	// The live in-memory config must still hold plaintext after Save.
	if cfg.Sources[0].APIKey != "sk-plaintext-secret" {
		t.Errorf("in-memory api_key mutated by Save: %q", cfg.Sources[0].APIKey)
	}
	if cfg.Keys[0].Key != "sk-gw-plain" {
		t.Errorf("in-memory gateway key mutated by Save: %q", cfg.Keys[0].Key)
	}

	// A fresh load must hand back plaintext again.
	loaded, err := Load(path)
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	loaded.AttachSecretBox(box)
	loaded.normalizeSecrets(box)
	if loaded.Sources[0].APIKey != "sk-plaintext-secret" {
		t.Errorf("loaded api_key = %q, want plaintext", loaded.Sources[0].APIKey)
	}
	if loaded.Sources[0].Headers["X-Extra"] != "header-secret" {
		t.Errorf("loaded header = %q, want plaintext", loaded.Sources[0].Headers["X-Extra"])
	}
	if loaded.Keys[0].Key != "sk-gw-plain" {
		t.Errorf("loaded gateway key = %q, want plaintext", loaded.Keys[0].Key)
	}
}

func TestPlaintextConfigStillLoads(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")
	// A hand-written, pre-encryption config: no enc:v1: anywhere.
	body := `listen: 127.0.0.1:0
sources:
  - name: legacy
    base_url: http://legacy/v1
    api_key: sk-written-by-hand
    adapter: openai
    models:
      - id: m
        kind: chat
`
	if err := os.WriteFile(path, []byte(body), 0644); err != nil {
		t.Fatal(err)
	}
	cfg, err := Load(path)
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))
	cfg.AttachSecretBox(box)
	cfg.normalizeSecrets(box)
	if cfg.Sources[0].APIKey != "sk-written-by-hand" {
		t.Errorf("plaintext config value changed: %q", cfg.Sources[0].APIKey)
	}
}

// Sealing a config is a one-time migration: the FIRST run writes, every later
// start must leave the file alone.
//
// This test previously called migratePlaintextSecrets() twice and compared
// ModTime. That assertion was both weak and wrong:
//
//   - migratePlaintextSecrets decides from IN-MEMORY state, and Save() ends by
//     unsealing memory so the running process keeps working. So every call sees
//     "plaintext present" and rewrites the file. Comparing ModTime hid this
//     because two writes inside one filesystem timestamp tick look identical —
//     the test flaked (~1 in 6) instead of failing.
//   - it also exercised a function production no longer calls on the startup
//     path, which now uses UnsealSecrets + SealIfNeeded (they decide from the
//     ON-DISK state).
//
// So the test targets that real path and compares file CONTENT, which cannot be
// fooled by timestamp granularity.
func TestSealingIsIdempotentAcrossStarts(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")
	runtime := filepath.Join(dir, "runtime.json")
	box := writeMasterKey(t, runtime)

	cfg := &Config{
		Path:    path,
		Listen:  "127.0.0.1:0",
		Sources: []Source{{Name: "up", BaseURL: "http://up/v1", APIKey: "sk-clear", Adapter: "openai", Models: []Model{{ID: "m"}}}},
	}
	cfg.AttachSecretBox(box)

	hadPlaintext, err := cfg.UnsealSecrets(box)
	if err != nil {
		t.Fatalf("unseal: %v", err)
	}
	if !hadPlaintext {
		t.Fatal("a config written with a clear-text credential must report plaintext")
	}
	if err := cfg.SealIfNeeded(hadPlaintext); err != nil {
		t.Fatalf("seal: %v", err)
	}
	first, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(first), encPrefix) {
		t.Fatal("sealing did not encrypt the value on disk")
	}
	// In-memory must be plaintext so the running process keeps working.
	if cfg.Sources[0].APIKey != "sk-clear" {
		t.Errorf("in-memory api_key = %q, want plaintext", cfg.Sources[0].APIKey)
	}

	// Simulate the next start: load from disk, unseal, seal-if-needed. The file
	// was already sealed, so nothing may be written.
	cfg2, err := Load(path)
	if err != nil {
		t.Fatal(err)
	}
	cfg2.AttachSecretBox(box)
	hadPlaintext2, err := cfg2.UnsealSecrets(box)
	if err != nil {
		t.Fatalf("second unseal: %v", err)
	}
	if hadPlaintext2 {
		t.Error("a sealed config reported plaintext — SealIfNeeded would rewrite " +
			"the file on every start")
	}
	if err := cfg2.SealIfNeeded(hadPlaintext2); err != nil {
		t.Fatalf("second seal: %v", err)
	}
	second, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if string(first) != string(second) {
		t.Error("the second start rewrote an already-sealed config; startup must be a no-op")
	}
}

func TestSealedValueDoesNotDoubleEncrypt(t *testing.T) {
	dir := t.TempDir()
	box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))
	cfg := &Config{Path: filepath.Join(dir, "config.yaml"), Sources: []Source{{Name: "a", BaseURL: "http://a"}}}
	cfg.AttachSecretBox(box)

	once, err := box.Encrypt("sk-x")
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources[0].APIKey = once
	if err := cfg.Save(); err != nil {
		t.Fatal(err)
	}
	got, err := box.Decrypt(cfg.Sources[0].APIKey)
	if err != nil {
		t.Fatalf("value became undecryptable: %v", err)
	}
	if got != "sk-x" {
		t.Errorf("round trip = %q, want sk-x", got)
	}
}

func TestUpsertSourceInYAMLSealsCredentials(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")
	box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))

	if err := os.WriteFile(path, []byte("listen: 127.0.0.1:0\nsources: []\n"), 0644); err != nil {
		t.Fatal(err)
	}
	src := Source{Name: "new", BaseURL: "http://new/v1", APIKey: "sk-fresh", Adapter: "openai", Models: []Model{{ID: "m"}}}
	if err := UpsertSourceInYAML(path, src.Name, src, box); err != nil {
		t.Fatalf("UpsertSourceInYAML: %v", err)
	}
	raw, _ := os.ReadFile(path)
	if strings.Contains(string(raw), "sk-fresh") {
		t.Error("newly added source stored its api_key in plaintext")
	}
	if !strings.Contains(string(raw), encPrefix) {
		t.Error("newly added source was not sealed")
	}
}

func TestCountPlaintextSecrets(t *testing.T) {
	cfg := &Config{
		Sources: []Source{
			{Name: "a", APIKey: encPrefix + "abc", Headers: map[string]string{"H": encPrefix + "x"}},
			{Name: "b", APIKey: "sk-in-clear", Headers: map[string]string{"H2": "clear-too"}},
		},
		Keys: []GWKey{{Key: "sk-gw-clear"}},
	}
	if got := cfg.countPlaintextSecrets(); got != 3 {
		t.Errorf("countPlaintextSecrets = %d, want 3", got)
	}
	if !cfg.hasPlaintextSecrets() {
		t.Error("hasPlaintextSecrets = false, want true")
	}
	cfg.Sources[1].APIKey = encPrefix + "y"
	cfg.Sources[1].Headers["H2"] = encPrefix + "z"
	cfg.Keys[0].Key = encPrefix + "k"
	if cfg.hasPlaintextSecrets() {
		t.Error("hasPlaintextSecrets = true after sealing everything")
	}
}

// TestNormalizeSecretsFailsLoudlyOnWrongMasterKey is the negative case that
// guards the migration: with a wrong key, a sealed value must NOT be handed
// back as ciphertext for a later re-seal (that compounds corruption and turns
// one lost key file into permanently unusable config).
func TestNormalizeSecretsFailsLoudlyOnWrongMasterKey(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")
	runtime := filepath.Join(dir, "runtime.json")
	good := writeMasterKey(t, runtime)

	cfg := &Config{
		Path:    path,
		Listen:  "127.0.0.1:0",
		Sources: []Source{{Name: "up", BaseURL: "http://up/v1", APIKey: "sk-secret", Adapter: "openai", Models: []Model{{ID: "m"}}}},
	}
	cfg.AttachSecretBox(good)
	if err := cfg.Save(); err != nil {
		t.Fatal(err)
	}
	sealedOnce, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}

	// A different master key must be rejected, loudly. NewSecretBox derives the
	// key path from the RUNTIME file's directory, so the wrong key has to live
	// in a different directory — otherwise it would read the good master.key.
	otherDir := t.TempDir()
	if err := os.WriteFile(filepath.Join(otherDir, "master.key"),
		[]byte(strings.Repeat("cd", 32)), 0600); err != nil {
		t.Fatal(err)
	}
	badBox, err := NewSecretBox(filepath.Join(otherDir, "runtime.json"))
	if err != nil {
		t.Fatal(err)
	}
	loaded, err := Load(path)
	if err != nil {
		t.Fatal(err)
	}
	loaded.AttachSecretBox(badBox)
	if err := loaded.NormalizeSecretsForRun(badBox); err == nil {
		t.Fatal("expected an error with the wrong master key, got nil")
	}

	// Crucially: the file must be untouched — no second seal on top.
	after, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if string(after) != string(sealedOnce) {
		t.Error("config file was rewritten despite the decrypt failure")
	}
	if n := strings.Count(string(after), encPrefix); n != strings.Count(string(sealedOnce), encPrefix) {
		t.Errorf("sealed count changed %d -> %d (double encryption)", strings.Count(string(sealedOnce), encPrefix), n)
	}
}