diff --git a/hdwallet/hdwallet.go b/hdwallet/hdwallet.go index 2e9d0d740..8540b3399 100644 --- a/hdwallet/hdwallet.go +++ b/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[:]) diff --git a/hdwallet/hdwallet_test.go b/hdwallet/hdwallet_test.go index de4866a63..42f9778e8 100644 --- a/hdwallet/hdwallet_test.go +++ b/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