llmgw/internal/secrets/secrets_test.go

package secrets

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// fakeVault answers KV v2 reads, AppRole logins and transit, and counts
// what it was asked.
func fakeVault(t *testing.T) (*httptest.Server, *int) {
	t.Helper()
	reads := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/v1/auth/approle/login" {
			var in map[string]string
			_ = json.NewDecoder(r.Body).Decode(&in)
			if in["role_id"] != "role" || in["secret_id"] != "sid" {
				http.Error(w, `{"errors":["denied"]}`, 400)
				return
			}
			_, _ = w.Write([]byte(`{"auth":{"client_token":"approle-token","lease_duration":3600}}`))
			return
		}
		if tok := r.Header.Get("X-Vault-Token"); tok != "root" && tok != "approle-token" {
			http.Error(w, `{"errors":["permission denied"]}`, 403)
			return
		}
		switch {
		case r.URL.Path == "/v1/secret/data/llmgw/store":
			reads++
			_, _ = w.Write([]byte(`{"data":{"data":{"dsn":"postgres://u:p@h/db"}}}`))
		case strings.HasPrefix(r.URL.Path, "/v1/transit/encrypt/llmgw"):
			var in map[string]string
			_ = json.NewDecoder(r.Body).Decode(&in)
			_, _ = w.Write([]byte(`{"data":{"ciphertext":"vault:v1:` + in["plaintext"] + `"}}`))
		case strings.HasPrefix(r.URL.Path, "/v1/transit/decrypt/llmgw"):
			var in map[string]string
			_ = json.NewDecoder(r.Body).Decode(&in)
			_, _ = w.Write([]byte(`{"data":{"plaintext":"` + strings.TrimPrefix(in["ciphertext"], "vault:v1:") + `"}}`))
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return srv, &reads
}

func TestReferencesAreRead(t *testing.T) {
	ctx := t.Context()
	f := filepath.Join(t.TempDir(), "s")
	_ = os.WriteFile(f, []byte(" from-file\n"), 0o600)
	if v, err := ReadString(ctx, f); err != nil || v != "from-file" {
		t.Errorf("file: %q %v", v, err)
	}
	t.Setenv("LLMGW_TEST_SECRET", "from-env")
	if v, err := ReadString(ctx, "env:LLMGW_TEST_SECRET"); err != nil || v != "from-env" {
		t.Errorf("env: %q %v", v, err)
	}
	if _, err := Read(ctx, "env:LLMGW_NOT_SET_ANYWHERE"); err == nil {
		t.Error("an unset variable read as empty")
	}

	srv, reads := fakeVault(t)
	t.Setenv("VAULT_TOKEN", "root")
	if err := Configure(Vault{Addr: srv.URL}); err != nil {
		t.Fatal(err)
	}
	for range 3 {
		if v, err := ReadString(ctx, "vault:secret/llmgw/store#dsn"); err != nil || v != "postgres://u:p@h/db" {
			t.Fatalf("vault: %q %v", v, err)
		}
	}
	if *reads != 1 {
		t.Errorf("read from Vault %d times; cached for a minute", *reads)
	}
	if _, err := Read(ctx, "vault:secret/llmgw/store#nope"); err == nil {
		t.Error("a missing field read")
	}
	if _, err := Read(ctx, "vault:secret/llmgw/store"); err == nil {
		t.Error("a reference without a field read")
	}

	// AppRole, with no token given.
	t.Setenv("VAULT_TOKEN", "")
	dir := t.TempDir()
	_ = os.WriteFile(filepath.Join(dir, "role"), []byte("role\n"), 0o600)
	_ = os.WriteFile(filepath.Join(dir, "sid"), []byte("sid\n"), 0o600)
	_ = Configure(Vault{Addr: srv.URL, RoleIDFile: filepath.Join(dir, "role"), SecretIDFile: filepath.Join(dir, "sid")})
	if v, err := ReadString(ctx, "vault:secret/llmgw/store#dsn"); err != nil || v == "" {
		t.Errorf("approle: %q %v", v, err)
	}
}

func TestSealersSealAndOpen(t *testing.T) {
	ctx := t.Context()
	keyFile := filepath.Join(t.TempDir(), "master")
	_ = os.WriteFile(keyFile, []byte(strings.Repeat("ab", 32)), 0o600)
	l, err := NewSealer(ctx, keyFile)
	if err != nil {
		t.Fatal(err)
	}
	sealed, err := l.Seal(ctx, []byte("the engine key"))
	if err != nil || !IsSealed(sealed) || strings.Contains(sealed, "engine key") {
		t.Fatalf("%q %v", sealed, err)
	}
	if got, err := l.Open(ctx, sealed); err != nil || string(got) != "the engine key" {
		t.Errorf("open: %q %v", got, err)
	}
	other := filepath.Join(t.TempDir(), "other")
	_ = os.WriteFile(other, []byte(strings.Repeat("cd", 32)), 0o600)
	l2, _ := NewSealer(ctx, other)
	if _, err := l2.Open(ctx, sealed); err == nil {
		t.Error("another key opened it")
	}
	short := filepath.Join(t.TempDir(), "short")
	_ = os.WriteFile(short, []byte("abc"), 0o600)
	if _, err := NewSealer(ctx, short); err == nil {
		t.Error("a short key was taken")
	}

	srv, _ := fakeVault(t)
	t.Setenv("VAULT_TOKEN", "root")
	_ = Configure(Vault{Addr: srv.URL})
	tr, err := NewSealer(ctx, "vault-transit:llmgw")
	if err != nil {
		t.Fatal(err)
	}
	sealed, err = tr.Seal(ctx, []byte("secret"))
	if err != nil || !strings.HasPrefix(sealed, transitPrefix) ||
		!strings.Contains(sealed, base64.StdEncoding.EncodeToString([]byte("secret"))) {
		t.Fatalf("%q %v", sealed, err)
	}
	if got, err := tr.Open(ctx, sealed); err != nil || string(got) != "secret" {
		t.Errorf("transit open: %q %v", got, err)
	}
}