ai-gateway-go/internal/platform/cryptox/keyring.go

package cryptox

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
)

type Cipher interface {
	Encrypt([]byte) ([]byte, int, error)
	Decrypt([]byte, int) ([]byte, error)
}

type Keyring struct {
	activeVersion int
	ciphers       map[int]*AESGCM
}

func NewKeyring(activeKey string, activeVersion int, encodedKeyring, purpose string) (*Keyring, error) {
	if activeVersion < 1 || purpose == "" {
		return nil, errors.New("encryption version and purpose are required")
	}
	encoded := make(map[int]string)
	if strings.TrimSpace(encodedKeyring) != "" {
		var values map[string]string
		if err := json.Unmarshal([]byte(encodedKeyring), &values); err != nil {
			return nil, fmt.Errorf("encryption keyring must be a JSON object: %w", err)
		}
		for rawVersion, key := range values {
			version, err := strconv.Atoi(rawVersion)
			if err != nil || version < 1 || strings.TrimSpace(key) == "" {
				return nil, fmt.Errorf("invalid encryption keyring version %q", rawVersion)
			}
			encoded[version] = strings.TrimSpace(key)
		}
	}
	if strings.TrimSpace(activeKey) != "" {
		if existing, ok := encoded[activeVersion]; ok && existing != strings.TrimSpace(activeKey) {
			return nil, fmt.Errorf("active encryption key version %d is configured twice with different values", activeVersion)
		}
		encoded[activeVersion] = strings.TrimSpace(activeKey)
	}
	if len(encoded) == 0 {
		return nil, nil
	}
	if _, ok := encoded[activeVersion]; !ok {
		return nil, fmt.Errorf("active encryption key version %d is not loaded", activeVersion)
	}
	keyring := &Keyring{activeVersion: activeVersion, ciphers: make(map[int]*AESGCM, len(encoded))}
	for version, key := range encoded {
		cipher, err := NewAESGCM(key, version, purpose)
		if err != nil {
			return nil, fmt.Errorf("encryption key version %d: %w", version, err)
		}
		keyring.ciphers[version] = cipher
	}
	return keyring, nil
}

func (k *Keyring) Encrypt(plaintext []byte) ([]byte, int, error) {
	if k == nil {
		return nil, 0, ErrKeyUnavailable
	}
	return k.ciphers[k.activeVersion].Encrypt(plaintext)
}

func (k *Keyring) Decrypt(encrypted []byte, version int) ([]byte, error) {
	if k == nil {
		return nil, ErrKeyUnavailable
	}
	cipher, ok := k.ciphers[version]
	if !ok {
		return nil, fmt.Errorf("encryption key version %d is not loaded", version)
	}
	return cipher.Decrypt(encrypted, version)
}

func (k *Keyring) ActiveVersion() int {
	if k == nil {
		return 0
	}
	return k.activeVersion
}

func (k *Keyring) Versions() []int {
	if k == nil {
		return nil
	}
	versions := make([]int, 0, len(k.ciphers))
	for version := range k.ciphers {
		versions = append(versions, version)
	}
	sort.Ints(versions)
	return versions
}

var _ Cipher = (*Keyring)(nil)