ModelRouter/internal/config/secret.go

package config

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var encPrefix = "enc:v1:"

// NewGatewayKey generates a fresh random gateway admin key with the
// "sk-gw-" prefix, used when bootstrapping a default config on first run.
func NewGatewayKey() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return "sk-gw-" + hex.EncodeToString(b), nil
}

// SecretBox encrypts and decrypts sensitive values (upstream API keys,
// custom header values, gateway keys) for persist-time protection. The
// master key comes from the MODELROUTER_MASTER_KEY environment variable
// (hex, 64 chars; the pre-rename LLMS_PROXY_MASTER_KEY is still accepted) or
// from a generated <runtime-dir>/master.key file.
type SecretBox struct {
	aead cipher.AEAD
}

// NewSecretBox builds a box from the env var, falling back to a master.key
// file next to the runtime file. If neither exists, a fresh random key is
// generated and written to master.key with mode 0600.
func NewSecretBox(runtimeFile string) (*SecretBox, error) {
	// New name first, then the pre-rename name: an existing unit file that still
	// exports LLMS_PROXY_MASTER_KEY must keep starting after an upgrade.
	envKey := os.Getenv("MODELROUTER_MASTER_KEY")
	envName := "MODELROUTER_MASTER_KEY"
	if envKey == "" {
		envKey = os.Getenv("LLMS_PROXY_MASTER_KEY")
		envName = "LLMS_PROXY_MASTER_KEY"
	}
	if envKey != "" {
		key, err := hex.DecodeString(strings.TrimSpace(envKey))
		if err != nil || len(key) != 32 {
			return nil, fmt.Errorf("%s: expected 64 hex chars", envName)
		}
		return newBox(key)
	}
	path := filepath.Join(filepath.Dir(runtimeFile), "master.key")
	key, err := os.ReadFile(path)
	if os.IsNotExist(err) {
		key = make([]byte, 32)
		if _, err := rand.Read(key); err != nil {
			return nil, err
		}
		if err := os.WriteFile(path, []byte(hex.EncodeToString(key)), 0600); err != nil {
			return nil, fmt.Errorf("write %s: %w", path, err)
		}
		log.Printf("[config] generated master key at %s (mode 0600, keep it safe)", path)
	} else if err != nil {
		return nil, err
	} else {
		key, err = hex.DecodeString(strings.TrimSpace(string(key)))
		if err != nil || len(key) != 32 {
			return nil, fmt.Errorf("master.key %s: expected 64 hex chars (regenerate if empty)", path)
		}
	}
	return newBox(key)
}

func newBox(key []byte) (*SecretBox, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &SecretBox{aead: aead}, nil
}

// Encrypt returns "" for empty input, otherwise "enc:v1:<base64>".
func (b *SecretBox) Encrypt(plain string) (string, error) {
	if plain == "" {
		return "", nil
	}
	nonce := make([]byte, b.aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	sealed := b.aead.Seal(nonce, nonce, []byte(plain), nil)
	return encPrefix + base64.RawURLEncoding.EncodeToString(sealed), nil
}

// Decrypt reverses Encrypt. Plain values (no prefix) are returned unchanged
// so old unencrypted files keep loading. Bad ciphertext is returned as-is
// with an error, never silently rewritten.
func (b *SecretBox) Decrypt(v string) (string, error) {
	if !strings.HasPrefix(v, encPrefix) {
		return v, nil
	}
	raw, err := base64.RawURLEncoding.DecodeString(strings.TrimPrefix(v, encPrefix))
	if err != nil {
		return v, err
	}
	n := b.aead.NonceSize()
	if len(raw) < n {
		return v, fmt.Errorf("ciphertext too short")
	}
	nonce, sealed := raw[:n], raw[n:]
	plain, err := b.aead.Open(nil, nonce, sealed, nil)
	if err != nil {
		return v, fmt.Errorf("decrypt failed (master key changed?)")
	}
	return string(plain), nil
}

// MustDecrypt is a convenience wrapper that logs and returns the input on
// failure so a bad value never silently becomes empty.
func (b *SecretBox) MustDecrypt(v string) string {
	out, err := b.Decrypt(v)
	if err != nil {
		head := v
		if len(head) > 8 {
			head = head[:8]
		}
		log.Printf("[config] secret decrypt failed for value %q…: %v", head, err)
		return v
	}
	return out
}