ModelRouter/internal/core/startup_secrets_test.go

package core

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

	"llmsproxy/internal/config"
)

// Startup must be idempotent, and every credential consumer must see PLAINTEXT.
//
// Two bugs shared one root cause: NewFromConfig unsealed the config at the END,
// after the steps that read credentials had already run.
//
//  1. seedKeys() deduped cfg.Keys[i].Key against the plaintext gateway_keys
//     entries. With ciphertext keys the comparison never matched, so every
//     restart appended another copy of the same admin key — production reached
//     four identical admin keys.
//  2. rebuildRegistry() handed cfg.Sources[i].APIKey to provider.New, so with
//     ciphertext it built every upstream client with "enc:v1:..." as its bearer
//     token and every upstream call would 401.
//
// They masked each other: seedKeys' Save() unsealed memory as a side effect, so
// bug 2 was invisible until the redundant saves were removed. These tests drive
// the real entry point (core.New -> config.Load off disk), because building a
// fresh in-memory config per attempt hides both bugs — that is exactly how the
// first draft of this test failed to reproduce anything.

// writeSealedInstall creates a config on disk the way a real install looks
// after its first run: credentials already sealed, gateway_keys still plaintext.
func writeSealedInstall(t *testing.T, dir string) string {
	t.Helper()
	path := filepath.Join(dir, "config.yaml")
	cfg := &config.Config{
		Path:         path,
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "AUTO",
		GatewayKeys:  []string{"sk-gw-seed-me"},
		Sources: []config.Source{{
			Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
			Adapter: "openai", Models: []config.Model{{ID: "gpt-4o"}},
		}},
	}
	c, err := NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("initial install: %v", err)
	}
	c.Close()

	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(raw), "enc:v1:") {
		t.Fatalf("fixture is wrong: on-disk config holds no sealed credential, "+
			"so this test cannot detect ciphertext leaking into consumers:\n%s", raw)
	}
	return path
}

// writeSealedInstallNoSeed is writeSealedInstall without gateway_keys, so no
// key is ever seeded and no startup Save() can unseal the config as a side
// effect. That isolation is what makes the ciphertext leak observable.
func writeSealedInstallNoSeed(t *testing.T, dir string) string {
	t.Helper()
	path := filepath.Join(dir, "config.yaml")
	cfg := &config.Config{
		Path:         path,
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "AUTO",
		Sources: []config.Source{{
			Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
			Adapter: "openai", Models: []config.Model{{ID: "gpt-4o"}},
		}},
	}
	c, err := NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("initial install: %v", err)
	}
	c.Close()
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(raw), "enc:v1:") {
		t.Fatalf("fixture is wrong: no sealed credential on disk:\n%s", raw)
	}
	return path
}

// countKeysWithSecret loads the config, unseals it, and counts keys whose
// plaintext equals one of the gateway_keys entries.
func countKeysWithSecret(t *testing.T, path string) (total, matching int) {
	t.Helper()
	cfg, err := config.Load(path)
	if err != nil {
		t.Fatal(err)
	}
	box, err := config.NewSecretBox(cfg.RuntimeFile)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := cfg.UnsealSecrets(box); err != nil {
		t.Fatal(err)
	}
	want := map[string]bool{}
	for _, g := range cfg.GatewayKeys {
		want[g] = true
	}
	for _, k := range cfg.Keys {
		total++
		if want[k.Key] {
			matching++
		}
	}
	return total, matching
}

// Restarting must not grow the key list. Before the fix this went 2 -> 3 -> 4.
func TestRestartDoesNotDuplicateSeededKeys(t *testing.T) {
	dir := t.TempDir()
	path := writeSealedInstall(t, dir)

	if _, n := countKeysWithSecret(t, path); n != 1 {
		t.Fatalf("after install: %d keys carry the seeded secret, want 1", n)
	}

	for i := 2; i <= 4; i++ {
		c, err := New(path) // production path: reads the file
		if err != nil {
			t.Fatalf("restart %d: %v", i, err)
		}
		c.Close()
		total, n := countKeysWithSecret(t, path)
		if n != 1 {
			t.Fatalf("restart %d: %d keys carry the seeded secret (total %d), want "+
				"exactly 1 — seedKeys is comparing ciphertext against plaintext and "+
				"appending a duplicate on every start", i, n, total)
		}
	}
}

// Every provider must hold the real upstream credential, not its ciphertext.
//
// This deliberately runs with gateway_keys EMPTY. The buggy order was masked
// whenever seedKeys had work to do, because its Save() unsealed memory as a
// side effect — providers then happened to see plaintext. With no key to seed,
// no Save runs, and the ciphertext reaches the registry directly. An earlier
// version of this test kept gateway_keys populated and was insensitive: it
// passed even with the bug present.
func TestProvidersGetPlaintextCredentials(t *testing.T) {
	dir := t.TempDir()
	path := writeSealedInstallNoSeed(t, dir)

	c, err := New(path)
	if err != nil {
		t.Fatalf("restart: %v", err)
	}
	defer c.Close()

	for _, p := range c.registry.Providers() {
		key := p.Config().APIKey
		if strings.HasPrefix(key, "enc:v1:") {
			t.Errorf("provider %q holds CIPHERTEXT api_key %q — every upstream call "+
				"would 401", p.Name(), key)
		}
		if key != "sk-upstream-secret" {
			t.Errorf("provider %q api_key = %q, want the plaintext upstream secret",
				p.Name(), key)
		}
	}
}