mcp-framework/secretstore/store.go

package secretstore

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"runtime"
	"slices"
	"strings"

	"github.com/99designs/keyring"
)

var ErrNotFound = errors.New("secret not found")
var ErrBackendUnavailable = errors.New("secret backend unavailable")
var ErrReadOnly = errors.New("secret backend is read-only")
var ErrInvalidBackendPolicy = errors.New("invalid secret backend policy")
var ErrBWNotLoggedIn = errors.New("bitwarden is not logged in")
var ErrBWLocked = errors.New("bitwarden vault is locked or BW_SESSION is missing")
var ErrBWUnavailable = errors.New("bitwarden CLI unavailable")

type BackendPolicy string

const (
	BackendAuto         BackendPolicy = "auto"
	BackendKWalletOnly  BackendPolicy = "kwallet-only"
	BackendKeyringAny   BackendPolicy = "keyring-any"
	BackendEnvOnly      BackendPolicy = "env-only"
	BackendBitwardenCLI BackendPolicy = "bitwarden-cli"
)

type Options struct {
	ServiceName      string
	BackendPolicy    BackendPolicy
	LookupEnv        func(string) (string, bool)
	KWalletAppID     string
	KWalletFolder    string
	BitwardenCommand string
}

type Store interface {
	SetSecret(name, label, secret string) error
	GetSecret(name string) (string, error)
	DeleteSecret(name string) error
}

type BackendUnavailableError struct {
	Policy    BackendPolicy
	Required  string
	Available []string
}

func (e *BackendUnavailableError) Error() string {
	if len(e.Available) == 0 {
		return fmt.Sprintf("secret backend policy %q requires %s, but no compatible backend is available", e.Policy, e.Required)
	}

	return fmt.Sprintf(
		"secret backend policy %q requires %s, but only [%s] are available",
		e.Policy,
		e.Required,
		strings.Join(e.Available, ", "),
	)
}

func (e *BackendUnavailableError) Unwrap() error {
	return ErrBackendUnavailable
}

type keyringStore struct {
	ring        keyring.Keyring
	serviceName string
}

type envStore struct {
	lookupEnv func(string) (string, bool)
}

var (
	openKeyring            = keyring.Open
	availableKeyringPolicy = keyring.AvailableBackends
)

func Open(options Options) (Store, error) {
	policy, err := normalizeBackendPolicy(options.BackendPolicy)
	if err != nil {
		return nil, err
	}

	if policy == BackendEnvOnly {
		return newEnvStore(options), nil
	}

	serviceName := strings.TrimSpace(options.ServiceName)
	if serviceName == "" {
		return nil, errors.New("service name must not be empty")
	}

	if policy == BackendBitwardenCLI {
		return newBitwardenStore(options, policy, serviceName)
	}

	available := availableKeyringPolicy()
	allowed, err := allowedBackends(policy, available)
	if err != nil {
		if errors.Is(err, ErrBackendUnavailable) && policy == BackendAuto && options.LookupEnv != nil {
			return newEnvStore(options), nil
		}
		return nil, err
	}

	ring, err := openKeyring(keyring.Config{
		ServiceName:     serviceName,
		AllowedBackends: allowed,
		KWalletAppID:    strings.TrimSpace(options.KWalletAppID),
		KWalletFolder:   strings.TrimSpace(options.KWalletFolder),
	})
	if err == nil {
		return &keyringStore{
			ring:        ring,
			serviceName: serviceName,
		}, nil
	}

	if policy == BackendAuto && options.LookupEnv != nil {
		return newEnvStore(options), nil
	}

	return nil, fmt.Errorf("open secret backend for service %q with policy %q: %w", serviceName, policy, err)
}

func BackendName() string {
	switch runtime.GOOS {
	case "darwin":
		return "macOS Keychain"
	case "windows":
		return "Windows Credential Manager"
	case "linux":
		return "Linux Secret Service or KWallet"
	default:
		return "system wallet"
	}
}

func EffectiveBackendPolicy(store Store) BackendPolicy {
	switch store.(type) {
	case *bitwardenStore:
		return BackendBitwardenCLI
	case *envStore:
		return BackendEnvOnly
	case *keyringStore:
		return BackendKeyringAny
	default:
		return ""
	}
}

func SetSecretVerified(store Store, name, label, secret string) error {
	if store == nil {
		return errors.New("secret store must not be nil")
	}

	if err := store.SetSecret(name, label, secret); err != nil {
		return err
	}

	verified, err := store.GetSecret(name)
	if err != nil {
		return fmt.Errorf("verify secret %q after write: %w", name, err)
	}

	if verified != secret {
		return fmt.Errorf("verify secret %q after write: read-back mismatch", name)
	}

	return nil
}

func SetJSON[T any](store Store, name, label string, value T) error {
	data, err := json.Marshal(value)
	if err != nil {
		return fmt.Errorf("encode structured secret %q as JSON: %w", name, err)
	}

	return store.SetSecret(name, label, string(data))
}

func GetJSON[T any](store Store, name string) (T, error) {
	var value T
	if err := GetJSONInto(store, name, &value); err != nil {
		var zero T
		return zero, err
	}

	return value, nil
}

func GetJSONInto(store Store, name string, target any) error {
	secret, err := store.GetSecret(name)
	if err != nil {
		return err
	}

	if err := json.Unmarshal([]byte(secret), target); err != nil {
		return fmt.Errorf("decode structured secret %q from JSON: %w", name, err)
	}

	return nil
}

func (s *keyringStore) SetSecret(name, label, secret string) error {
	if err := s.ring.Set(keyring.Item{
		Key:   name,
		Label: label,
		Data:  []byte(secret),
	}); err != nil {
		return fmt.Errorf("save secret %q in OS wallet for service %q: %w", name, s.serviceName, err)
	}

	return nil
}

func (s *keyringStore) GetSecret(name string) (string, error) {
	item, err := s.ring.Get(name)
	if err != nil {
		if errors.Is(err, keyring.ErrKeyNotFound) {
			return "", ErrNotFound
		}
		return "", fmt.Errorf("read secret %q from OS wallet for service %q: %w", name, s.serviceName, err)
	}

	return string(item.Data), nil
}

func (s *keyringStore) DeleteSecret(name string) error {
	if err := s.ring.Remove(name); err != nil && !errors.Is(err, keyring.ErrKeyNotFound) {
		return fmt.Errorf("delete secret %q from OS wallet for service %q: %w", name, s.serviceName, err)
	}
	return nil
}

func (s *envStore) SetSecret(name, label, secret string) error {
	return fmt.Errorf("save secret %q in environment backend: %w", name, ErrReadOnly)
}

func (s *envStore) GetSecret(name string) (string, error) {
	value, ok := s.lookupEnv(name)
	if !ok {
		return "", ErrNotFound
	}

	return value, nil
}

func (s *envStore) DeleteSecret(name string) error {
	return fmt.Errorf("delete secret %q from environment backend: %w", name, ErrReadOnly)
}

func newEnvStore(options Options) Store {
	lookupEnv := options.LookupEnv
	if lookupEnv == nil {
		lookupEnv = os.LookupEnv
	}

	return &envStore{lookupEnv: lookupEnv}
}

func allowedBackends(policy BackendPolicy, available []keyring.BackendType) ([]keyring.BackendType, error) {
	switch policy {
	case BackendAuto, BackendKeyringAny:
		if len(available) == 0 {
			return nil, &BackendUnavailableError{
				Policy:    policy,
				Required:  "any keyring backend",
				Available: nil,
			}
		}
		return slices.Clone(available), nil
	case BackendKWalletOnly:
		if !slices.Contains(available, keyring.KWalletBackend) {
			return nil, &BackendUnavailableError{
				Policy:    policy,
				Required:  string(keyring.KWalletBackend),
				Available: backendNames(available),
			}
		}
		return []keyring.BackendType{keyring.KWalletBackend}, nil
	default:
		return nil, invalidBackendPolicyError(policy)
	}
}

func backendNames(backends []keyring.BackendType) []string {
	names := make([]string, 0, len(backends))
	for _, backend := range backends {
		names = append(names, string(backend))
	}
	return names
}

func normalizeBackendPolicy(policy BackendPolicy) (BackendPolicy, error) {
	trimmed := BackendPolicy(strings.TrimSpace(string(policy)))
	if trimmed == "" {
		return BackendAuto, nil
	}

	switch trimmed {
	case BackendAuto, BackendKWalletOnly, BackendKeyringAny, BackendEnvOnly, BackendBitwardenCLI:
		return trimmed, nil
	default:
		return "", invalidBackendPolicyError(trimmed)
	}
}

func invalidBackendPolicyError(policy BackendPolicy) error {
	return fmt.Errorf("%w %q", ErrInvalidBackendPolicy, policy)
}