motr-enclave/internal/crypto/bip44/bip44.go

// Package bip44 provides BIP44 address derivation using MPC public keys.
package bip44

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"

	"golang.org/x/crypto/ripemd160"
	"golang.org/x/crypto/sha3"
)

type CoinType uint32

const (
	CoinTypeBitcoin  CoinType = 0
	CoinTypeEthereum CoinType = 60
	CoinTypeCosmos   CoinType = 118
	CoinTypeSonr     CoinType = 703
)

type ChainConfig struct {
	CoinType CoinType
	Prefix   string
	ChainID  string
}

var DefaultChains = map[string]ChainConfig{
	"bitcoin": {
		CoinType: CoinTypeBitcoin,
		Prefix:   "bc1",
		ChainID:  "bitcoin-mainnet",
	},
	"ethereum": {
		CoinType: CoinTypeEthereum,
		Prefix:   "0x",
		ChainID:  "1",
	},
	"sonr": {
		CoinType: CoinTypeSonr,
		Prefix:   "snr",
		ChainID:  "sonr-testnet-1",
	},
	"cosmos": {
		CoinType: CoinTypeCosmos,
		Prefix:   "cosmos",
		ChainID:  "cosmoshub-4",
	},
}

func DeriveAddress(pubKeyBytes []byte, chain string) (string, error) {
	config, ok := DefaultChains[chain]
	if !ok {
		return "", fmt.Errorf("bip44: unknown chain: %s", chain)
	}
	return DeriveAddressWithConfig(pubKeyBytes, config)
}

func DeriveAddressWithConfig(pubKeyBytes []byte, config ChainConfig) (string, error) {
	if len(pubKeyBytes) == 0 {
		return "", fmt.Errorf("bip44: public key cannot be empty")
	}

	switch config.CoinType {
	case CoinTypeBitcoin:
		return deriveBitcoinAddress(pubKeyBytes, config.Prefix)
	case CoinTypeEthereum:
		return deriveEthereumAddress(pubKeyBytes)
	case CoinTypeCosmos, CoinTypeSonr:
		return deriveCosmosAddress(pubKeyBytes, config.Prefix)
	default:
		return "", fmt.Errorf("bip44: unsupported coin type: %d", config.CoinType)
	}
}

// deriveBitcoinAddress creates P2WPKH (native SegWit) address using SHA256+RIPEMD160
func deriveBitcoinAddress(pubKeyBytes []byte, prefix string) (string, error) {
	compressed := ensureCompressed(pubKeyBytes)
	if compressed == nil {
		return "", fmt.Errorf("bip44: invalid public key length: %d", len(pubKeyBytes))
	}

	sha256Hash := sha256.Sum256(compressed)
	ripemd := ripemd160.New()
	ripemd.Write(sha256Hash[:])
	pubKeyHash := ripemd.Sum(nil)

	return bech32Encode(prefix, pubKeyHash)
}

// deriveEthereumAddress creates address using Keccak256 of uncompressed pubkey
func deriveEthereumAddress(pubKeyBytes []byte) (string, error) {
	var keyData []byte

	switch len(pubKeyBytes) {
	case 65:
		keyData = pubKeyBytes[1:] // strip 0x04 prefix
	case 64:
		keyData = pubKeyBytes
	default:
		return "", fmt.Errorf("bip44: ethereum requires uncompressed key (64 or 65 bytes), got %d", len(pubKeyBytes))
	}

	hash := sha3.NewLegacyKeccak256()
	hash.Write(keyData)
	hashBytes := hash.Sum(nil)

	return "0x" + hex.EncodeToString(hashBytes[12:]), nil
}

// deriveCosmosAddress creates bech32 address using SHA256+RIPEMD160
func deriveCosmosAddress(pubKeyBytes []byte, prefix string) (string, error) {
	compressed := ensureCompressed(pubKeyBytes)
	if compressed == nil {
		return "", fmt.Errorf("bip44: invalid public key length: %d", len(pubKeyBytes))
	}

	sha256Hash := sha256.Sum256(compressed)
	ripemd := ripemd160.New()
	ripemd.Write(sha256Hash[:])
	addressBytes := ripemd.Sum(nil)

	return bech32Encode(prefix, addressBytes)
}

func ensureCompressed(pubKeyBytes []byte) []byte {
	switch len(pubKeyBytes) {
	case 33:
		return pubKeyBytes
	case 65:
		return compressPublicKey(pubKeyBytes)
	default:
		return nil
	}
}

func compressPublicKey(uncompressed []byte) []byte {
	if len(uncompressed) != 65 {
		return uncompressed
	}

	x := uncompressed[1:33]
	y := uncompressed[33:65]

	compressed := make([]byte, 33)
	if y[31]&1 == 0 {
		compressed[0] = 0x02
	} else {
		compressed[0] = 0x03
	}
	copy(compressed[1:], x)

	return compressed
}

func bech32Encode(hrp string, data []byte) (string, error) {
	converted, err := convertBits(data, 8, 5, true)
	if err != nil {
		return "", fmt.Errorf("bip44: convert bits: %w", err)
	}
	return encode(hrp, converted)
}

func convertBits(data []byte, fromBits, toBits uint8, pad bool) ([]byte, error) {
	acc := uint32(0)
	bits := uint8(0)
	maxv := uint32(1<<toBits) - 1
	result := make([]byte, 0, len(data)*int(fromBits)/int(toBits)+1)

	for _, b := range data {
		acc = (acc << fromBits) | uint32(b)
		bits += fromBits
		for bits >= toBits {
			bits -= toBits
			result = append(result, byte((acc>>bits)&maxv))
		}
	}

	if pad {
		if bits > 0 {
			result = append(result, byte((acc<<(toBits-bits))&maxv))
		}
	} else if bits >= fromBits {
		return nil, fmt.Errorf("illegal zero padding")
	} else if ((acc << (toBits - bits)) & maxv) != 0 {
		return nil, fmt.Errorf("non-zero padding")
	}

	return result, nil
}

const charset = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

func encode(hrp string, data []byte) (string, error) {
	checksum := createChecksum(hrp, data)
	combined := append(data, checksum...)

	result := make([]byte, len(hrp)+1+len(combined))
	copy(result, hrp)
	result[len(hrp)] = '1'

	for i, b := range combined {
		result[len(hrp)+1+i] = charset[b]
	}

	return string(result), nil
}

func createChecksum(hrp string, data []byte) []byte {
	values := append(hrpExpand(hrp), data...)
	values = append(values, []byte{0, 0, 0, 0, 0, 0}...)
	polymod := polymod(values) ^ 1

	checksum := make([]byte, 6)
	for i := 0; i < 6; i++ {
		checksum[i] = byte((polymod >> (5 * (5 - i))) & 31)
	}
	return checksum
}

func hrpExpand(hrp string) []byte {
	result := make([]byte, len(hrp)*2+1)
	for i, c := range hrp {
		result[i] = byte(c >> 5)
		result[i+len(hrp)+1] = byte(c & 31)
	}
	result[len(hrp)] = 0
	return result
}

func polymod(values []byte) uint32 {
	gen := []uint32{0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3}
	chk := uint32(1)
	for _, v := range values {
		top := chk >> 25
		chk = (chk&0x1ffffff)<<5 ^ uint32(v)
		for i := 0; i < 5; i++ {
			if (top>>i)&1 == 1 {
				chk ^= gen[i]
			}
		}
	}
	return chk
}

type Account struct {
	Address      string
	ChainID      string
	CoinType     CoinType
	AccountIndex uint32
	AddressIndex uint32
}

func DeriveAccounts(pubKeyBytes []byte, chains []string) ([]Account, error) {
	accounts := make([]Account, 0, len(chains))

	for _, chain := range chains {
		config, ok := DefaultChains[chain]
		if !ok {
			continue
		}

		addr, err := DeriveAddressWithConfig(pubKeyBytes, config)
		if err != nil {
			continue
		}

		accounts = append(accounts, Account{
			Address:      addr,
			ChainID:      config.ChainID,
			CoinType:     config.CoinType,
			AccountIndex: 0,
			AddressIndex: 0,
		})
	}

	return accounts, nil
}