perf(hdwallet): speed up ExportHDAddresses ~10-15x

hdwallet/hdwallet.go

@@ -5,8 +5,11 @@ import (
 	"crypto/sha256"
 	"crypto/sha512"
 	"encoding/hex"
+	"errors"
 	"fmt"
+	"math"
 	"regexp"
+	"runtime"
 	"strconv"
 	"strings"
 
@@ -22,6 +25,7 @@ import (
 	"golang.org/x/crypto/pbkdf2"
 	"golang.org/x/crypto/ripemd160" //nolint:staticcheck
 	"golang.org/x/crypto/sha3"
+	"golang.org/x/sync/errgroup"
 )
 
 type (
@@ -228,13 +232,14 @@ func (p *PolyWallet) ExportRootAddress() (*PolyWalletExport, error) {
 	}
 
 	pwe.RootKey = rootKey.String()
-	pwe.HexPublicKey = hex.EncodeToString(rootKey.PublicKey().Key)
+	rootPubKey := rootKey.PublicKey()
+	rootUncompressedPubKey := toUncompressedPubKey(rootKey)
+	pwe.HexPublicKey = hex.EncodeToString(rootPubKey.Key)
 	pwe.HexPrivateKey = hex.EncodeToString(rootKey.Key)
 	pwe.WIF = toWIF(rootKey)
-	pwe.BTCAddress = toBTCAddress(rootKey)
-	rootEthAddress := toETHAddress(rootKey)
-	pwe.ETHAddress = rootEthAddress.String()
-	pwe.HexFullPublicKey = hex.EncodeToString(toUncompressedPubKey(rootKey))
+	pwe.BTCAddress = toBTCAddress(rootPubKey)
+	pwe.ETHAddress = RawPubKeyToETHAddress(rootUncompressedPubKey).String()
+	pwe.HexFullPublicKey = hex.EncodeToString(rootUncompressedPubKey)
 	addr, err := GetPublicKeyFromSeed(p.rawSeed, SignatureSecp256k1, true)
 	if err != nil {
 		return nil, err
@@ -322,41 +327,79 @@ func (p *PolyWallet) ExportHDAddresses(count int) (*PolyWalletExport, error) {
 
 		lastDerivationPathPart := derivationPathParts[len(derivationPathParts)-1]
 		lastDerivationPathPart = strings.ReplaceAll(lastDerivationPathPart, "'", "")
-		idx, err := strconv.Atoi(lastDerivationPathPart)
-		if err == nil {
-			firstIndex = idx
-		} else {
+		// Non-hardened bip32 child indexes are bounded by 2^31-1, so parse
+		// with an explicit 31-bit bound, which also guarantees the value
+		// fits in an int on 32-bit platforms.
+		idx, idxErr := strconv.ParseUint(lastDerivationPathPart, 10, 31)
+		switch {
+		case idxErr == nil:
+			firstIndex = int(idx)
+		case errors.Is(idxErr, strconv.ErrRange):
+			return nil, fmt.Errorf("address index %s in derivation path exceeds the maximum non-hardened bip32 child index: %w", lastDerivationPathPart, idxErr)
+		default:
 			log.Warn().Msg("Failed to identify the index of the address in the derivation path, starting at 0")
 		}
 	}
 	lastIndex := firstIndex + count
 
-	for i := firstIndex; i < lastIndex; i = i + 1 {
-		// TODO if we want to provide support for hardened addresses it would need to be accommodated here
-		currentPath := p.derivationPath
-		if lastIndex-firstIndex > 1 {
-			currentPath = strings.Join(derivationPathParts[:len(derivationPathParts)-1], "/") + "/" + strconv.Itoa(i)
+	if count == 1 {
+		k, keyErr := p.GetKeyForPath(p.derivationPath)
+		if keyErr != nil {
+			return nil, keyErr
 		}
+		pwe.Addresses = append(pwe.Addresses, exportAddress(p.derivationPath, k))
+		return pwe, nil
+	}
+	if count < 1 {
+		return pwe, nil
+	}
+	if int64(lastIndex-1) > math.MaxUint32 {
+		return nil, fmt.Errorf("address index %d exceeds the maximum bip32 child index %d", lastIndex-1, uint32(math.MaxUint32))
+	}
 
-		k, err := p.GetKeyForPath(currentPath)
-		if err != nil {
-			return nil, err
-		}
+	// TODO if we want to provide support for hardened addresses it would need to be accommodated here
+	// Derive the shared parent key once, then derive a single child per address
+	// rather than re-deriving the full path from the master key every time.
+	parentPath := strings.Join(derivationPathParts[:len(derivationPathParts)-1], "/")
+	parentKey, err := p.GetKeyForPath(parentPath)
+	if err != nil {
+		return nil, err
+	}
 
-		pae := new(PolyAddressExport)
-		pae.Path = currentPath
-		pae.HexPublicKey = hex.EncodeToString(k.PublicKey().Key)
-		pae.HexPrivateKey = hex.EncodeToString(k.Key)
-		pae.WIF = toWIF(k)
-		pae.BTCAddress = toBTCAddress(k)
-		ethAddress := toETHAddress(k)
-		pae.ETHAddress = ethAddress.String()
-		pae.HexFullPublicKey = hex.EncodeToString(toUncompressedPubKey(k))
-		pwe.Addresses = append(pwe.Addresses, pae)
+	addresses := make([]*PolyAddressExport, count)
+	g := new(errgroup.Group)
+	g.SetLimit(runtime.NumCPU())
+	for i := firstIndex; i < lastIndex; i = i + 1 {
+		g.Go(func() error {
+			k, err := parentKey.NewChildKey(uint32(i))
+			if err != nil {
+				return fmt.Errorf("failed to derive child key %d of %s: %w", i, parentPath, err)
+			}
+			addresses[i-firstIndex] = exportAddress(parentPath+"/"+strconv.Itoa(i), k)
+			return nil
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return nil, err
 	}
+	pwe.Addresses = addresses
 	return pwe, nil
 }
 
+func exportAddress(path string, k *bip32.Key) *PolyAddressExport {
+	pubKey := k.PublicKey()
+	uncompressedPubKey := toUncompressedPubKey(k)
+	pae := new(PolyAddressExport)
+	pae.Path = path
+	pae.HexPublicKey = hex.EncodeToString(pubKey.Key)
+	pae.HexPrivateKey = hex.EncodeToString(k.Key)
+	pae.WIF = toWIF(k)
+	pae.BTCAddress = toBTCAddress(pubKey)
+	pae.ETHAddress = RawPubKeyToETHAddress(uncompressedPubKey).String()
+	pae.HexFullPublicKey = hex.EncodeToString(uncompressedPubKey)
+	return pae
+}
+
 // https://en.bitcoin.it/wiki/Wallet_import_format
 func toWIF(prvKey *bip32.Key) string {
 	mainnet := []byte{0x80}
@@ -369,11 +412,6 @@ func toWIF(prvKey *bip32.Key) string {
 	return base58.Encode(h3)
 }
 
-func toETHAddress(prvKey *bip32.Key) common.Address {
-	concat := toUncompressedPubKey(prvKey)
-	return RawPubKeyToETHAddress(concat)
-
-}
 func RawPubKeyToETHAddress(concat []byte) common.Address {
 	h := sha3.NewLegacyKeccak256()
 	h.Write(concat)
@@ -401,8 +439,7 @@ func toUncompressedPubKey(prvKey *bip32.Key) []byte {
 }
 
 // https://en.bitcoin.it/wiki/Technical_background_of_version_1_Bitcoin_addresses
-func toBTCAddress(prvKey *bip32.Key) string {
-	publicKey := prvKey.PublicKey()
+func toBTCAddress(publicKey *bip32.Key) string {
 	h := sha256.Sum256(publicKey.Key)
 	ripe160 := ripemd160.New()
 	ripe160.Write(h[:])

hdwallet/hdwallet_test.go

@@ -348,6 +348,69 @@ func TestPaddedPublicKey(t *testing.T) {
 	}
 }
 
+// TestExportHDAddressesMatchesPerPathDerivation verifies that the batched,
+// parallel derivation in ExportHDAddresses produces exactly the same keys as
+// deriving each address independently from the master key via GetKeyForPath.
+// Golden values were captured from the output of the pre-optimization
+// implementation.
+func TestExportHDAddressesMatchesPerPathDerivation(t *testing.T) {
+	const mnemonic = "code code code code code code code code code code code quality"
+	const count = 500
+
+	pw, err := NewPolyWallet(mnemonic, "")
+	require.NoError(t, err)
+
+	export, err := pw.ExportHDAddresses(count)
+	require.NoError(t, err)
+	require.Len(t, export.Addresses, count)
+
+	// Golden values from the pre-optimization implementation.
+	first := export.Addresses[0]
+	assert.Equal(t, "m/44'/60'/0'/0/0", first.Path)
+	assert.Equal(t, "0x85dA99c8a7C2C95964c8EfD687E95E632Fc533D6", first.ETHAddress)
+	assert.Equal(t, "42b6e34dc21598a807dc19d7784c71b2a7a01f6480dc6f58258f78e539f1a1fa", first.HexPrivateKey)
+	assert.Equal(t, "03507cf9a75e053cda6922467721ddb10412da9bec30620347d9529cc77fca2433", first.HexPublicKey)
+	assert.Equal(t, "1HdqWQqsVD41pKNHVrpFGNHqW6t3fuAfkh", first.BTCAddress)
+	assert.Equal(t, "KyTPrvjtqbyu9J4bRAvJgnBeYAdrdvCAToceM1RwDGFjEAdra6Fa", first.WIF)
+	last := export.Addresses[count-1]
+	assert.Equal(t, "m/44'/60'/0'/0/499", last.Path)
+	assert.Equal(t, "0x927C2d5aEab7BEFfb2e61C151f2524D00169d146", last.ETHAddress)
+	assert.Equal(t, "c29ee1a8886b2f1f592f3e4156d2d57842f0674f05b5e831b49c99b4ab233019", last.HexPrivateKey)
+
+	// Re-derive every address independently on a fresh wallet and compare.
+	reference, err := NewPolyWallet(mnemonic, "")
+	require.NoError(t, err)
+
+	for i, addr := range export.Addresses {
+		expectedPath := fmt.Sprintf("m/44'/60'/0'/0/%d", i)
+		require.Equal(t, expectedPath, addr.Path)
+
+		k, err := reference.GetKeyForPath(expectedPath)
+		require.NoError(t, err)
+
+		pubKey := k.PublicKey()
+		uncompressedPubKey := toUncompressedPubKey(k)
+		assert.Equal(t, hex.EncodeToString(k.Key), addr.HexPrivateKey, "path %s", expectedPath)
+		assert.Equal(t, hex.EncodeToString(pubKey.Key), addr.HexPublicKey, "path %s", expectedPath)
+		assert.Equal(t, hex.EncodeToString(uncompressedPubKey), addr.HexFullPublicKey, "path %s", expectedPath)
+		assert.Equal(t, RawPubKeyToETHAddress(uncompressedPubKey).String(), addr.ETHAddress, "path %s", expectedPath)
+		assert.Equal(t, toWIF(k), addr.WIF, "path %s", expectedPath)
+		assert.Equal(t, toBTCAddress(pubKey), addr.BTCAddress, "path %s", expectedPath)
+	}
+}
+
+func BenchmarkExportHDAddresses(b *testing.B) {
+	const mnemonic = "code code code code code code code code code code code quality"
+	pw, err := NewPolyWallet(mnemonic, "")
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for b.Loop() {
+		_, err := pw.ExportHDAddresses(500)
+		require.NoError(b, err)
+	}
+}
+
 func TestDerivationPath(t *testing.T) {
 	type testCase struct {
 		derivationPathInput  string