llmgw/internal/secrets/secrets.go

// Package secrets reads the secrets llmgw is configured with from where
// they are kept: a file, as always, an environment variable, or HashiCorp
// Vault (or OpenBao) — and seals the store's own secrets under a key that
// stays outside it.
//
// A reference is one of:
//
//	/etc/llmgw/secrets/store-dsn         a file, trimmed
//	env:LLMGW_STORE_DSN                  an environment variable
//	vault:secret/llmgw/store#dsn         a field of a KV secret in Vault
//
// A Vault path's first segment is its KV v2 mount; secret/data/… is taken
// as written. KV v1 mounts work too.
package secrets

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

// Vault is how Vault is reached.
type Vault struct {
	// Addr is Vault's address, like https://vault.internal:8200; empty
	// takes VAULT_ADDR.
	Addr string `koanf:"addr"`
	// TokenFile holds a token; empty takes VAULT_TOKEN, or logs in with
	// AppRole when RoleIDFile and SecretIDFile are set.
	TokenFile    string `koanf:"token_file"`
	RoleIDFile   string `koanf:"role_id_file"`
	SecretIDFile string `koanf:"secret_id_file"`
	// CAFile verifies Vault's certificate; empty uses the system's.
	CAFile string `koanf:"ca_file"`
	// Namespace is Vault Enterprise's or OpenBao's namespace.
	Namespace string `koanf:"namespace"`
}

var (
	mu     sync.Mutex
	vault  Vault
	client = &http.Client{Timeout: 10 * time.Second}
	token  string
	expiry time.Time
	cache  = map[string]cached{}
)

type cached struct {
	b  []byte
	at time.Time
}

// cacheFor is how long a value read from Vault is used before it is read
// again: a secret rotated there reaches llmgw within it.
const cacheFor = time.Minute

// Configure says how to reach Vault. Called once configuration is read.
func Configure(v Vault) error {
	mu.Lock()
	defer mu.Unlock()
	vault, token, expiry, cache = v, "", time.Time{}, map[string]cached{}
	if v.CAFile != "" {
		pem, err := os.ReadFile(v.CAFile)
		if err != nil {
			return fmt.Errorf("secrets.vault.ca_file: %w", err)
		}
		pool := x509.NewCertPool()
		if !pool.AppendCertsFromPEM(pem) {
			return fmt.Errorf("secrets.vault.ca_file %s holds no certificate", v.CAFile)
		}
		client = &http.Client{Timeout: 10 * time.Second,
			Transport: &http.Transport{TLSClientConfig: &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12}}}
	}
	return nil
}

// IsReference says whether s names a secret somewhere other than a file.
func IsReference(s string) bool {
	return strings.HasPrefix(s, "env:") || strings.HasPrefix(s, "vault:")
}

// Read reads a secret by reference, trimmed of surrounding space.
func Read(ctx context.Context, ref string) ([]byte, error) {
	switch {
	case strings.HasPrefix(ref, "env:"):
		name := strings.TrimPrefix(ref, "env:")
		v, ok := os.LookupEnv(name)
		if !ok {
			return nil, fmt.Errorf("the environment has no %s", name)
		}
		return []byte(strings.TrimSpace(v)), nil
	case strings.HasPrefix(ref, "vault:"):
		return readVault(ctx, strings.TrimPrefix(ref, "vault:"))
	}
	b, err := os.ReadFile(ref) //nolint:gosec // a path the operator configured
	if err != nil {
		return nil, err
	}
	return bytes.TrimSpace(b), nil
}

// ReadString is Read as a string.
func ReadString(ctx context.Context, ref string) (string, error) {
	b, err := Read(ctx, ref)
	return string(b), err
}

func readVault(ctx context.Context, ref string) ([]byte, error) {
	path, field, ok := strings.Cut(ref, "#")
	if !ok || path == "" || field == "" {
		return nil, fmt.Errorf("vault:%s: write it as vault:mount/path#field", ref)
	}
	mu.Lock()
	if c, ok := cache[ref]; ok && time.Since(c.at) < cacheFor {
		mu.Unlock()
		return c.b, nil
	}
	mu.Unlock()

	api := path
	if mount, rest, ok := strings.Cut(path, "/"); ok && !strings.HasPrefix(rest, "data/") {
		api = mount + "/data/" + rest
	}
	var out struct {
		Data map[string]any `json:"data"`
	}
	if err := call(ctx, http.MethodGet, "/v1/"+api, nil, &out); err != nil {
		// A KV v1 mount has no /data/ in its path.
		if api != path && errors.Is(err, errNotFound) {
			err = call(ctx, http.MethodGet, "/v1/"+path, nil, &out)
		}
		if err != nil {
			return nil, fmt.Errorf("vault:%s: %w", ref, err)
		}
	}
	data := out.Data
	if inner, ok := data["data"].(map[string]any); ok {
		data = inner // KV v2
	}
	v, ok := data[field]
	if !ok {
		return nil, fmt.Errorf("vault:%s: the secret has no field %q", ref, field)
	}
	s, ok := v.(string)
	if !ok {
		return nil, fmt.Errorf("vault:%s: field %q is not a string", ref, field)
	}
	b := []byte(strings.TrimSpace(s))
	mu.Lock()
	cache[ref] = cached{b: b, at: time.Now()}
	mu.Unlock()
	return b, nil
}

var errNotFound = errors.New("not found")

// call makes one request of Vault's API with the token, logging in first
// when there is none.
func call(ctx context.Context, method, path string, body, out any) error {
	tok, addr, err := currentToken(ctx)
	if err != nil {
		return err
	}
	return do(ctx, addr, tok, method, path, body, out)
}

func do(ctx context.Context, addr, tok, method, path string, body, out any) error {
	var rd io.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		rd = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(addr, "/")+path, rd)
	if err != nil {
		return err
	}
	if tok != "" {
		req.Header.Set("X-Vault-Token", tok)
	}
	mu.Lock()
	ns, c := vault.Namespace, client
	mu.Unlock()
	if ns != "" {
		req.Header.Set("X-Vault-Namespace", ns)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	switch {
	case resp.StatusCode == http.StatusNotFound:
		return errNotFound
	case resp.StatusCode >= 300:
		var e struct {
			Errors []string `json:"errors"`
		}
		_ = json.Unmarshal(b, &e)
		return fmt.Errorf("vault answered %s: %s", resp.Status, strings.Join(e.Errors, "; "))
	}
	if out != nil {
		return json.Unmarshal(b, out)
	}
	return nil
}

// currentToken is the token to send, logging in with AppRole when that is
// how it is configured and the last login has expired.
func currentToken(ctx context.Context) (string, string, error) {
	mu.Lock()
	v := vault
	tok, exp := token, expiry
	mu.Unlock()
	addr := v.Addr
	if addr == "" {
		addr = os.Getenv("VAULT_ADDR")
	}
	if addr == "" {
		return "", "", errors.New("no Vault to ask: set secrets.vault.addr or VAULT_ADDR")
	}
	if tok != "" && (exp.IsZero() || time.Now().Before(exp)) {
		return tok, addr, nil
	}
	switch {
	case v.TokenFile != "":
		b, err := os.ReadFile(v.TokenFile)
		if err != nil {
			return "", "", fmt.Errorf("secrets.vault.token_file: %w", err)
		}
		tok, exp = strings.TrimSpace(string(b)), time.Time{}
	case v.RoleIDFile != "" && v.SecretIDFile != "":
		role, err := os.ReadFile(v.RoleIDFile)
		if err != nil {
			return "", "", err
		}
		sec, err := os.ReadFile(v.SecretIDFile)
		if err != nil {
			return "", "", err
		}
		var out struct {
			Auth struct {
				ClientToken   string `json:"client_token"`
				LeaseDuration int    `json:"lease_duration"`
			} `json:"auth"`
		}
		if err := do(ctx, addr, "", http.MethodPost, "/v1/auth/approle/login", map[string]string{
			"role_id": strings.TrimSpace(string(role)), "secret_id": strings.TrimSpace(string(sec))}, &out); err != nil {
			return "", "", fmt.Errorf("logging in to Vault with AppRole: %w", err)
		}
		tok = out.Auth.ClientToken
		if d := out.Auth.LeaseDuration; d > 0 {
			// Renewed by logging in again, a little before it runs out.
			exp = time.Now().Add(time.Duration(d) * time.Second * 9 / 10)
		}
	default:
		tok, exp = os.Getenv("VAULT_TOKEN"), time.Time{}
	}
	if tok == "" {
		return "", "", errors.New("no Vault token: set secrets.vault.token_file, AppRole's files, or VAULT_TOKEN")
	}
	mu.Lock()
	token, expiry = tok, exp
	mu.Unlock()
	return tok, addr, nil
}