myxce/src/frontend/apiserver/store/apikey_test.go

package store

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"
)

func resetAPIKeyState() {
	SetAPIKeyStatePath("")
	InitStores()
}

func TestHashAPIKey_Deterministic(t *testing.T) {
	h1 := HashAPIKey("xce_ak_abc123")
	h2 := HashAPIKey("xce_ak_abc123")
	if h1 != h2 {
		t.Errorf("HashAPIKey should be deterministic: %q != %q", h1, h2)
	}
	h3 := HashAPIKey("xce_ak_different")
	if h1 == h3 {
		t.Error("HashAPIKey should differ for different inputs")
	}
}

func TestHashAPIKey_NeverPlaintext(t *testing.T) {
	raw := "xce_ak_secret-value"
	hashed := HashAPIKey(raw)
	if hashed == raw {
		t.Error("HashAPIKey must not return the raw key unchanged")
	}
}

func TestAddAPIKey_ValidateAPIKey_RoundTrip(t *testing.T) {
	resetAPIKeyState()

	raw := "xce_ak_roundtrip"
	entry := APIKeyEntry{
		Username:   "mcp-bot",
		RoleName:   "mcp_readonly",
		CreateTime: time.Now(),
		ExpiryTime: time.Now().Add(24 * time.Hour),
	}
	AddAPIKey(raw, entry)

	got, ok := ValidateAPIKey(raw)
	if !ok {
		t.Fatal("expected ValidateAPIKey to succeed for a freshly-added key")
	}
	if got.Username != "mcp-bot" {
		t.Errorf("Username = %q, want mcp-bot", got.Username)
	}
	if got.RoleName != "mcp_readonly" {
		t.Errorf("RoleName = %q, want mcp_readonly", got.RoleName)
	}
}

func TestValidateAPIKey_Unknown(t *testing.T) {
	resetAPIKeyState()

	_, ok := ValidateAPIKey("xce_ak_never-issued")
	if ok {
		t.Error("expected ValidateAPIKey to fail for a key that was never added")
	}
}

func TestValidateAPIKey_Expired(t *testing.T) {
	resetAPIKeyState()

	raw := "xce_ak_expired"
	AddAPIKey(raw, APIKeyEntry{
		Username:   "mcp-bot",
		RoleName:   "mcp_readonly",
		CreateTime: time.Now().Add(-48 * time.Hour),
		ExpiryTime: time.Now().Add(-1 * time.Hour), // already expired
	})

	_, ok := ValidateAPIKey(raw)
	if ok {
		t.Error("expected ValidateAPIKey to fail for an expired key")
	}
}

func TestValidateAPIKey_StoredByHashNotPlaintext(t *testing.T) {
	resetAPIKeyState()

	raw := "xce_ak_hashed-storage"
	AddAPIKey(raw, APIKeyEntry{Username: "mcp-bot", ExpiryTime: time.Now().Add(time.Hour)})

	if APIKeys.Has(raw) {
		t.Error("store must not contain the raw key as a map key")
	}
	if !APIKeys.Has(HashAPIKey(raw)) {
		t.Error("store should contain the hashed key")
	}
}

// ========== Persistence across restarts ==========

// TestAPIKeyPersistence_ReloadSurvivesRestart simulates an apiserver restart:
// issue keys against one store instance with a state file, then re-run
// InitStores (fresh memory) and verify every key still validates and its
// user→role mapping was restored — the exact failure mode that made every
// apiserver restart silently revoke all service-account keys before.
func TestAPIKeyPersistence_ReloadSurvivesRestart(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "apikeys.json")
	SetAPIKeyStatePath(path)
	defer resetAPIKeyState()

	InitStores()
	raw := "xce_ak_survives-restart"
	AddAPIKey(raw, APIKeyEntry{
		Username:   "advisor-svc",
		RoleName:   "advisor_agent",
		CreateTime: time.Now(),
		ExpiryTime: time.Now().Add(90 * 24 * time.Hour),
	})

	// Restart: brand-new store over the same state path.
	InitStores()

	entry, ok := ValidateAPIKey(raw)
	if !ok {
		t.Fatal("key lost across restart — state file was not reloaded")
	}
	if entry.Username != "advisor-svc" || entry.RoleName != "advisor_agent" {
		t.Errorf("entry = %+v, want advisor-svc/advisor_agent", entry)
	}
	// Role mapping must be restored too, or the key authenticates but is
	// denied every permission check after restart.
	roles := GetUserRoleNames("advisor-svc")
	if len(roles) != 1 || roles[0] != "advisor_agent" {
		t.Errorf("GetUserRoleNames(advisor-svc) = %v, want [advisor_agent]", roles)
	}
}

// TestAPIKeyPersistence_FileNeverContainsRawKey verifies the state file holds
// only hashes — reading it must not leak usable credentials.
func TestAPIKeyPersistence_FileNeverContainsRawKey(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "apikeys.json")
	SetAPIKeyStatePath(path)
	defer resetAPIKeyState()

	InitStores()
	raw := "xce_ak_secret-never-on-disk"
	AddAPIKey(raw, APIKeyEntry{Username: "mcp-bot", ExpiryTime: time.Now().Add(time.Hour)})

	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("state file missing: %v", err)
	}
	if string(data) != "" && containsPlain(data, raw) {
		t.Error("state file contains the raw key — must store hashes only")
	}
	var state persistedAPIKeyState
	if err := json.Unmarshal(data, &state); err != nil {
		t.Fatalf("state file is not valid JSON: %v", err)
	}
	if len(state.Keys) != 1 {
		t.Fatalf("state.Keys = %d entries, want 1", len(state.Keys))
	}
	for hash := range state.Keys {
		if hash != HashAPIKey(raw) {
			t.Errorf("stored key %q is not the expected SHA-256 hash", hash)
		}
	}
}

// TestAPIKeyPersistence_CorruptFileFailsFast verifies a corrupt state file
// makes InitStores panic instead of silently starting with zero keys (which
// would turn every service account into a mystery 401).
func TestAPIKeyPersistence_CorruptFileFailsFast(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "apikeys.json")
	SetAPIKeyStatePath(path)

	if err := os.WriteFile(path, []byte("{not json"), 0o600); err != nil {
		t.Fatal(err)
	}
	defer resetAPIKeyState()

	defer func() {
		if recover() == nil {
			t.Error("InitStores did not fail on a corrupt state file")
		}
	}()
	InitStores()
}

// TestAPIKeyPersistence_MissingFileIsFreshInstall verifies a missing state
// file is the normal fresh-install case, not an error.
func TestAPIKeyPersistence_MissingFileIsFreshInstall(t *testing.T) {
	SetAPIKeyStatePath(filepath.Join(t.TempDir(), "does-not-exist.json"))
	defer resetAPIKeyState()

	InitStores() // must not panic
}

// TestAPIKeyPersistence_SaveFailurePropagates verifies a state file that
// cannot be written makes AddAPIKey return an error (the controller turns
// this into a failed request rather than issuing a key that dies on restart).
func TestAPIKeyPersistence_SaveFailurePropagates(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "apikeys.json")
	SetAPIKeyStatePath(path)
	defer resetAPIKeyState()

	InitStores()
	// A directory where the state file lives makes saveAPIKeyState's rename
	// fail — os.Rename(file, dir) fails with EISDIR/ENOTDIR on Linux.
	if err := os.Mkdir(path, 0o700); err != nil {
		t.Fatal(err)
	}
	err := AddAPIKey("xce_ak_unwritable", APIKeyEntry{Username: "x"})
	if err == nil {
		t.Error("AddAPIKey did not report the failed state save")
	}
}

func containsPlain(data []byte, needle string) bool {
	return len(needle) > 0 && string(data) != "" && jsonContains(data, []byte(needle))
}

func jsonContains(haystack, needle []byte) bool {
	for i := 0; i+len(needle) <= len(haystack); i++ {
		if string(haystack[i:i+len(needle)]) == string(needle) {
			return true
		}
	}
	return false
}