From f64adb57371e00f6af11b574baebc0d0c8734dd5 Mon Sep 17 00:00:00 2001
From: Adam R <13562139+catenocrypt@users.noreply.github.com>
Date: Sat, 7 Dec 2019 08:03:45 +0100
Subject: [PATCH] HDWallet: New method to derive address directly from wallet.
 (#760)

* HDWallet: New method to derive address directly from wallet.
* Test update: preserve original test case intact.
---
 include/TrustWalletCore/TWHDWallet.h |  4 ++++
 src/Bitcoin/SegwitAddress.cpp        |  2 +-
 src/interface/TWHDWallet.cpp         |  7 +++++++
 tests/interface/TWHDWalletTests.cpp  | 27 ++++++++++++++++++++++++++-
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/include/TrustWalletCore/TWHDWallet.h b/include/TrustWalletCore/TWHDWallet.h
index 1e4e178116e..73e45cac44c 100644
--- a/include/TrustWalletCore/TWHDWallet.h
+++ b/include/TrustWalletCore/TWHDWallet.h
@@ -57,6 +57,10 @@ struct TWPrivateKey *_Nonnull TWHDWalletGetMasterKey(struct TWHDWallet *_Nonnull
 TW_EXPORT_METHOD
 struct TWPrivateKey *_Nonnull TWHDWalletGetKeyForCoin(struct TWHDWallet *_Nonnull wallet, enum TWCoinType coin);
 
+/// Generates the default address for the specified coin (without exposing intermediary private key).
+TW_EXPORT_METHOD
+TWString *_Nonnull TWHDWalletGetAddressForCoin(struct TWHDWallet *_Nonnull wallet, enum TWCoinType coin);
+
 /// Generates the private key for the specified derivation path.
 TW_EXPORT_METHOD
 struct TWPrivateKey *_Nonnull TWHDWalletGetKey(struct TWHDWallet *_Nonnull wallet, TWString *_Nonnull derivationPath);
diff --git a/src/Bitcoin/SegwitAddress.cpp b/src/Bitcoin/SegwitAddress.cpp
index 40a7120b575..fafbd57c67b 100644
--- a/src/Bitcoin/SegwitAddress.cpp
+++ b/src/Bitcoin/SegwitAddress.cpp
@@ -51,7 +51,7 @@ bool SegwitAddress::isValid(const std::string& string, const std::string& hrp) {
 SegwitAddress::SegwitAddress(const PublicKey& publicKey, int witver, std::string hrp)
     : hrp(std::move(hrp)), witnessVersion(witver), witnessProgram() {
     if (publicKey.type != TWPublicKeyTypeSECP256k1) {
-        throw std::invalid_argument("SegwitAddressneeds a compressed SECP256k1 public key.");
+        throw std::invalid_argument("SegwitAddress needs a compressed SECP256k1 public key.");
     }
     witnessProgram.resize(20);
     ecdsa_get_pubkeyhash(publicKey.compressed().bytes.data(), HASHER_SHA2_RIPEMD,
diff --git a/src/interface/TWHDWallet.cpp b/src/interface/TWHDWallet.cpp
index c8d58cc0930..a6d5b7e95a1 100644
--- a/src/interface/TWHDWallet.cpp
+++ b/src/interface/TWHDWallet.cpp
@@ -49,6 +49,13 @@ struct TWPrivateKey *_Nonnull TWHDWalletGetKeyForCoin(struct TWHDWallet *wallet,
     return new TWPrivateKey{ wallet->impl.getKey(derivationPath) };
 }
 
+TWString *_Nonnull TWHDWalletGetAddressForCoin(struct TWHDWallet *wallet, TWCoinType coin) {
+    auto derivationPath = TW::derivationPath(coin);
+    PrivateKey privateKey = wallet->impl.getKey(derivationPath);
+    std::string address = deriveAddress(coin, privateKey);
+    return TWStringCreateWithUTF8Bytes(address.c_str());
+}
+
 struct TWPrivateKey *_Nonnull TWHDWalletGetKey(struct TWHDWallet *_Nonnull wallet, TWString *_Nonnull derivationPath) {
     auto& s = *reinterpret_cast<const std::string*>(derivationPath);
     return new TWPrivateKey{ wallet->impl.getKey( TW::DerivationPath(s)) };
diff --git a/tests/interface/TWHDWalletTests.cpp b/tests/interface/TWHDWalletTests.cpp
index e5b4cb011e4..0b1b8e1b22a 100644
--- a/tests/interface/TWHDWalletTests.cpp
+++ b/tests/interface/TWHDWalletTests.cpp
@@ -95,13 +95,32 @@ TEST(HDWallet, Derive) {
     assertHexEqual(publicKey0Data, "0414acbe5a06c68210fcbb77763f9612e45a526990aeb69d692d705f276f558a5ae68268e9389bb099ed5ac84d8d6861110f63644f6e5b447e3f86b4bab5dee011");
 }
 
-TEST(HDWallet, DeriveBitcoin) {
+TEST(HDWallet, DeriveBitcoinNonextended) {
     auto wallet = WRAP(TWHDWallet, TWHDWalletCreateWithMnemonic(words.get(), passphrase.get()));
     auto key = WRAP(TWPrivateKey, TWHDWalletGetKeyForCoin(wallet.get(), TWCoinTypeBitcoin));
     auto publicKey = TWPrivateKeyGetPublicKeySecp256k1(key.get(), false);
     auto publicKeyData = WRAPD(TWPublicKeyData(publicKey));
 
     assertHexEqual(publicKeyData, "047ea5dff03f677502c4a1d73c5ac897200e56b155e876774c8fba0cc22f80b9414ec07cda7b1c9a84c2e04ea2746c21afacc5e91b47427c453c3f1a4a3e983ce5");
+    // Note: address derivation does not work with nonextended public key
+}
+
+TEST(HDWallet, DeriveBitcoinExtended) {
+    auto wallet = WRAP(TWHDWallet, TWHDWalletCreateWithMnemonic(words.get(), passphrase.get()));
+    auto key = WRAP(TWPrivateKey, TWHDWalletGetKeyForCoin(wallet.get(), TWCoinTypeBitcoin));
+    auto publicKey = TWPrivateKeyGetPublicKeySecp256k1(key.get(), true);
+    auto publicKeyData = WRAPD(TWPublicKeyData(publicKey));
+
+    assertHexEqual(publicKeyData, "037ea5dff03f677502c4a1d73c5ac897200e56b155e876774c8fba0cc22f80b941");
+
+    auto address = WRAPS(TWCoinTypeDeriveAddressFromPublicKey(TWCoinTypeBitcoin, publicKey));
+    assertStringsEqual(address, "bc1qumwjg8danv2vm29lp5swdux4r60ezptzz7ce85");
+}
+
+TEST(HDWallet, DeriveAddressBitcoin) {
+    auto wallet = WRAP(TWHDWallet, TWHDWalletCreateWithMnemonic(words.get(), passphrase.get()));
+    auto address = WRAP(TWString, TWHDWalletGetAddressForCoin(wallet.get(), TWCoinTypeBitcoin));
+    assertStringsEqual(address, "bc1qumwjg8danv2vm29lp5swdux4r60ezptzz7ce85");
 }
 
 TEST(HDWallet, DeriveEthereum) {
@@ -116,6 +135,12 @@ TEST(HDWallet, DeriveEthereum) {
     assertStringsEqual(address, "0x27Ef5cDBe01777D62438AfFeb695e33fC2335979");
 }
 
+TEST(HDWallet, DeriveAddressEthereum) {
+    auto wallet = WRAP(TWHDWallet, TWHDWalletCreateWithMnemonic(words.get(), passphrase.get()));
+    auto address = WRAP(TWString, TWHDWalletGetAddressForCoin(wallet.get(), TWCoinTypeEthereum));
+    assertStringsEqual(address, "0x27Ef5cDBe01777D62438AfFeb695e33fC2335979");
+}
+
 TEST(HDWallet, DeriveCosmos) {
     // use `gaiacli keys add key_name` to generate mnemonic words and private key
     auto words = STRING("attract term foster morning tail foam excite copper disease measure cheese camera rug enroll cause flip sword waste try local purchase between idea thank");