sendico/api/chain/gateway/internal/keymanager/vault/manager.go

package vault

import (
	"context"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/hex"
	"math/big"
	"os"
	"path"
	"strings"

	"github.com/decred/dcrd/dcrec/secp256k1/v4"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/hashicorp/vault/api"
	"go.uber.org/zap"

	"github.com/tech/sendico/chain/gateway/internal/keymanager"
	"github.com/tech/sendico/pkg/merrors"
	"github.com/tech/sendico/pkg/mlogger"
)

// Config describes how to connect to Vault for managed wallet keys.
type Config struct {
	Address   string `mapstructure:"address"`
	TokenEnv  string `mapstructure:"token_env"`
	Namespace string `mapstructure:"namespace"`
	MountPath string `mapstructure:"mount_path"`
	KeyPrefix string `mapstructure:"key_prefix"`
}

// Manager implements the keymanager.Manager contract backed by HashiCorp Vault.
type Manager struct {
	logger    mlogger.Logger
	client    *api.Client
	store     *api.KVv2
	keyPrefix string
}

// New constructs a Vault-backed key manager.
func New(logger mlogger.Logger, cfg Config) (*Manager, error) {
	if logger == nil {
		return nil, merrors.InvalidArgument("vault key manager: logger is required")
	}
	address := strings.TrimSpace(cfg.Address)
	if address == "" {
		logger.Error("vault address missing")
		return nil, merrors.InvalidArgument("vault key manager: address is required")
	}
	tokenEnv := strings.TrimSpace(cfg.TokenEnv)
	if tokenEnv == "" {
		logger.Error("vault token env missing")
		return nil, merrors.InvalidArgument("vault key manager: token_env is required")
	}
	token := strings.TrimSpace(os.Getenv(tokenEnv))
	if token == "" {
		logger.Error("vault token missing; expected Vault Agent to export token", zap.String("env", tokenEnv))
		return nil, merrors.InvalidArgument("vault key manager: token env " + tokenEnv + " is not set (expected Vault Agent sink to populate it)")
	}
	mountPath := strings.Trim(strings.TrimSpace(cfg.MountPath), "/")
	if mountPath == "" {
		logger.Error("vault mount path missing")
		return nil, merrors.InvalidArgument("vault key manager: mount_path is required")
	}
	keyPrefix := strings.Trim(strings.TrimSpace(cfg.KeyPrefix), "/")
	if keyPrefix == "" {
		keyPrefix = "chain/gateway/wallets"
	}

	clientCfg := api.DefaultConfig()
	clientCfg.Address = address

	client, err := api.NewClient(clientCfg)
	if err != nil {
		logger.Error("failed to create vault client", zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to create client: " + err.Error())
	}
	client.SetToken(token)
	if ns := strings.TrimSpace(cfg.Namespace); ns != "" {
		client.SetNamespace(ns)
	}

	kv := client.KVv2(mountPath)

	return &Manager{
		logger:    logger.Named("vault"),
		client:    client,
		store:     kv,
		keyPrefix: keyPrefix,
	}, nil
}

// CreateManagedWalletKey creates a new managed wallet key and stores it in Vault.
func (m *Manager) CreateManagedWalletKey(ctx context.Context, walletRef string, network string) (*keymanager.ManagedWalletKey, error) {
	if strings.TrimSpace(walletRef) == "" {
		m.logger.Warn("walletRef missing for managed key creation", zap.String("network", network))
		return nil, merrors.InvalidArgument("vault key manager: walletRef is required")
	}
	if strings.TrimSpace(network) == "" {
		m.logger.Warn("network missing for managed key creation", zap.String("wallet_ref", walletRef))
		return nil, merrors.InvalidArgument("vault key manager: network is required")
	}

	privateKey, err := ecdsa.GenerateKey(secp256k1.S256(), rand.Reader)
	if err != nil {
		m.logger.Warn("failed to generate managed wallet key", zap.String("wallet_ref", walletRef), zap.String("network", network), zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to generate key: " + err.Error())
	}
	privateKeyBytes := crypto.FromECDSA(privateKey)
	publicKey := privateKey.PublicKey
	publicKeyBytes := crypto.FromECDSAPub(&publicKey)
	publicKeyHex := hex.EncodeToString(publicKeyBytes)
	address := crypto.PubkeyToAddress(publicKey).Hex()

	err = m.persistKey(ctx, walletRef, network, privateKeyBytes, publicKeyBytes, address)
	if err != nil {
		m.logger.Warn("failed to persist managed wallet key", zap.String("wallet_ref", walletRef), zap.String("network", network), zap.Error(err))
		zeroBytes(privateKeyBytes)
		zeroBytes(publicKeyBytes)
		return nil, err
	}
	zeroBytes(privateKeyBytes)
	zeroBytes(publicKeyBytes)

	m.logger.Info("managed wallet key created",
		zap.String("wallet_ref", walletRef),
		zap.String("network", network),
		zap.String("address", strings.ToLower(address)),
	)

	return &keymanager.ManagedWalletKey{
		KeyID:     m.buildKeyID(network, walletRef),
		Address:   strings.ToLower(address),
		PublicKey: publicKeyHex,
	}, nil
}

func (m *Manager) persistKey(ctx context.Context, walletRef, network string, privateKey, publicKey []byte, address string) error {
	secretPath := m.buildKeyID(network, walletRef)
	payload := map[string]interface{}{
		"private_key": hex.EncodeToString(privateKey),
		"public_key":  hex.EncodeToString(publicKey),
		"address":     strings.ToLower(address),
		"network":     strings.ToLower(network),
	}
	if _, err := m.store.Put(ctx, secretPath, payload); err != nil {
		return merrors.Internal("vault key manager: failed to write secret at " + secretPath + ": " + err.Error())
	}
	return nil
}

func (m *Manager) buildKeyID(network, walletRef string) string {
	net := strings.Trim(strings.ToLower(network), "/")
	return path.Join(m.keyPrefix, net, walletRef)
}

// SignTransaction loads the key material from Vault and signs the transaction.
func (m *Manager) SignTransaction(ctx context.Context, keyID string, tx *types.Transaction, chainID *big.Int) (*types.Transaction, error) {
	if strings.TrimSpace(keyID) == "" {
		m.logger.Warn("signing failed: empty key id")
		return nil, merrors.InvalidArgument("vault key manager: keyID is required")
	}
	if tx == nil {
		m.logger.Warn("signing failed: nil transaction", zap.String("key_id", keyID))
		return nil, merrors.InvalidArgument("vault key manager: transaction is nil")
	}
	if chainID == nil {
		m.logger.Warn("signing failed: nil chain id", zap.String("key_id", keyID))
		return nil, merrors.InvalidArgument("vault key manager: chainID is nil")
	}

	material, err := m.loadKey(ctx, keyID)
	if err != nil {
		m.logger.Warn("failed to load key material", zap.String("key_id", keyID), zap.Error(err))
		return nil, err
	}

	keyBytes, err := hex.DecodeString(material.PrivateKey)
	if err != nil {
		m.logger.Warn("invalid key material", zap.String("key_id", keyID), zap.Error(err))
		return nil, merrors.Internal("vault key manager: invalid key material: " + err.Error())
	}
	defer zeroBytes(keyBytes)

	privateKey, err := crypto.ToECDSA(keyBytes)
	if err != nil {
		m.logger.Warn("failed to construct private key", zap.String("key_id", keyID), zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to construct private key: " + err.Error())
	}

	signed, err := types.SignTx(tx, types.LatestSignerForChainID(chainID), privateKey)
	if err != nil {
		m.logger.Warn("failed to sign transaction", zap.String("key_id", keyID), zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to sign transaction: " + err.Error())
	}
	m.logger.Info("transaction signed with managed key",
		zap.String("key_id", keyID),
		zap.String("network", material.Network),
		zap.String("tx_hash", signed.Hash().Hex()),
	)
	return signed, nil
}

type keyMaterial struct {
	PrivateKey string
	PublicKey  string
	Address    string
	Network    string
}

func (m *Manager) loadKey(ctx context.Context, keyID string) (*keyMaterial, error) {
	secretPath := strings.Trim(strings.TrimPrefix(keyID, "/"), "/")
	secret, err := m.store.Get(ctx, secretPath)
	if err != nil {
		m.logger.Warn("failed to read secret", zap.String("path", secretPath), zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to read secret at " + secretPath + ": " + err.Error())
	}
	if secret == nil || secret.Data == nil {
		m.logger.Warn("secret not found", zap.String("path", secretPath))
		return nil, merrors.NoData("vault key manager: secret " + secretPath + " not found")
	}

	getString := func(key string) (string, error) {
		val, ok := secret.Data[key]
		if !ok {
			m.logger.Warn("secret missing field", zap.String("path", secretPath), zap.String("field", key))
			return "", merrors.Internal("vault key manager: secret " + secretPath + " missing " + key)
		}
		str, ok := val.(string)
		if !ok || strings.TrimSpace(str) == "" {
			m.logger.Warn("secret field invalid", zap.String("path", secretPath), zap.String("field", key))
			return "", merrors.Internal("vault key manager: secret " + secretPath + " invalid " + key)
		}
		return str, nil
	}

	privateKey, err := getString("private_key")
	if err != nil {
		return nil, err
	}
	publicKey, err := getString("public_key")
	if err != nil {
		return nil, err
	}
	address, err := getString("address")
	if err != nil {
		return nil, err
	}
	network, err := getString("network")
	if err != nil {
		return nil, err
	}

	return &keyMaterial{
		PrivateKey: privateKey,
		PublicKey:  publicKey,
		Address:    address,
		Network:    network,
	}, nil
}

func zeroBytes(data []byte) {
	for i := range data {
		data[i] = 0
	}
}

var _ keymanager.Manager = (*Manager)(nil)