Add SegWit support to importmulti with some ProcessImport cleanup

2 files changed, 124 insertions, 157 deletions

diff --git a/src/wallet/rpcdump.cpp b/src/wallet/rpcdump.cpp
index c97bc38e6f..188c2246f8 100644
--- a/src/wallet/rpcdump.cpp
+++ b/src/wallet/rpcdump.cpp
@@ -809,29 +809,24 @@ UniValue dumpwallet(const JSONRPCRequest& request)
 static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, const int64_t timestamp) EXCLUSIVE_LOCKS_REQUIRED(pwallet->cs_wallet)
 {
     try {
-        bool success = false;
-
-        // Required fields.
+        // First ensure scriptPubKey has either a script or JSON with "address" string
         const UniValue& scriptPubKey = data["scriptPubKey"];
-
-        // Should have script or JSON with "address".
-        if (!(scriptPubKey.getType() == UniValue::VOBJ && scriptPubKey.exists("address")) && !(scriptPubKey.getType() == UniValue::VSTR)) {
-            throw JSONRPCError(RPC_INVALID_PARAMETER, "Invalid scriptPubKey");
+        bool isScript = scriptPubKey.getType() == UniValue::VSTR;
+        if (!isScript && !(scriptPubKey.getType() == UniValue::VOBJ && scriptPubKey.exists("address"))) {
+            throw JSONRPCError(RPC_INVALID_PARAMETER, "scriptPubKey must be string with script or JSON with address string");
         }
+        const std::string& output = isScript ? scriptPubKey.get_str() : scriptPubKey["address"].get_str();
 
         // Optional fields.
         const std::string& strRedeemScript = data.exists("redeemscript") ? data["redeemscript"].get_str() : "";
+        const std::string& witness_script_hex = data.exists("witnessscript") ? data["witnessscript"].get_str() : "";
         const UniValue& pubKeys = data.exists("pubkeys") ? data["pubkeys"].get_array() : UniValue();
         const UniValue& keys = data.exists("keys") ? data["keys"].get_array() : UniValue();
         const bool internal = data.exists("internal") ? data["internal"].get_bool() : false;
         const bool watchOnly = data.exists("watchonly") ? data["watchonly"].get_bool() : false;
-        const std::string& label = data.exists("label") && !internal ? data["label"].get_str() : "";
-
-        bool isScript = scriptPubKey.getType() == UniValue::VSTR;
-        bool isP2SH = strRedeemScript.length() > 0;
-        const std::string& output = isScript ? scriptPubKey.get_str() : scriptPubKey["address"].get_str();
+        const std::string& label = data.exists("label") ? data["label"].get_str() : "";
 
-        // Parse the output.
+        // Generate the script and destination for the scriptPubKey provided
         CScript script;
         CTxDestination dest;
 
@@ -855,35 +850,38 @@ static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, con
 
         // Watchonly and private keys
         if (watchOnly && keys.size()) {
-            throw JSONRPCError(RPC_INVALID_PARAMETER, "Incompatibility found between watchonly and keys");
+            throw JSONRPCError(RPC_INVALID_PARAMETER, "Watch-only addresses should not include private keys");
         }
 
-        // Internal + Label
+        // Internal addresses should not have a label
         if (internal && data.exists("label")) {
-            throw JSONRPCError(RPC_INVALID_PARAMETER, "Incompatibility found between internal and label");
+            throw JSONRPCError(RPC_INVALID_PARAMETER, "Internal addresses should not have a label");
         }
 
-        // Keys / PubKeys size check.
-        if (!isP2SH && (keys.size() > 1 || pubKeys.size() > 1)) { // Address / scriptPubKey
-            throw JSONRPCError(RPC_INVALID_PARAMETER, "More than private key given for one address");
+        // Force users to provide the witness script in its field rather than redeemscript
+        if (!strRedeemScript.empty() && script.IsPayToWitnessScriptHash()) {
+            throw JSONRPCError(RPC_INVALID_PARAMETER, "P2WSH addresses have an empty redeemscript. Please provide the witnessscript instead.");
         }
 
-        // Invalid P2SH redeemScript
-        if (isP2SH && !IsHex(strRedeemScript)) {
-            throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid redeem script");
-        }
-
-        // Process. //
+        CScript scriptpubkey_script = script;
+        CTxDestination scriptpubkey_dest = dest;
+        bool allow_p2wpkh = true;
 
         // P2SH
-        if (isP2SH) {
+        if (!strRedeemScript.empty() && script.IsPayToScriptHash()) {
+            // Check the redeemScript is valid
+            if (!IsHex(strRedeemScript)) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid redeem script: must be hex string");
+            }
+
             // Import redeem script.
             std::vector<unsigned char> vData(ParseHex(strRedeemScript));
             CScript redeemScript = CScript(vData.begin(), vData.end());
+            CScriptID redeem_id(redeemScript);
 
-            // Invalid P2SH address
-            if (!script.IsPayToScriptHash()) {
-                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid P2SH address / script");
+            // Check that the redeemScript and scriptPubKey match
+            if (GetScriptForDestination(redeem_id) != script) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "The redeemScript does not match the scriptPubKey");
             }
 
             pwallet->MarkDirty();
@@ -892,103 +890,83 @@ static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, con
                 throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
             }
 
-            CScriptID redeem_id(redeemScript);
             if (!pwallet->HaveCScript(redeem_id) && !pwallet->AddCScript(redeemScript)) {
                 throw JSONRPCError(RPC_WALLET_ERROR, "Error adding p2sh redeemScript to wallet");
             }
 
-            CScript redeemDestination = GetScriptForDestination(redeem_id);
+            // Now set script to the redeemScript so we parse the inner script as P2WSH or P2WPKH below
+            script = redeemScript;
+            ExtractDestination(script, dest);
+        }
 
-            if (::IsMine(*pwallet, redeemDestination) == ISMINE_SPENDABLE) {
-                throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
+        // (P2SH-)P2WSH
+        if (!witness_script_hex.empty() && script.IsPayToWitnessScriptHash()) {
+            if (!IsHex(witness_script_hex)) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid witness script: must be hex string");
             }
 
-            pwallet->MarkDirty();
+            // Generate the scripts
+            std::vector<unsigned char> witness_script_parsed(ParseHex(witness_script_hex));
+            CScript witness_script = CScript(witness_script_parsed.begin(), witness_script_parsed.end());
+            CScriptID witness_id(witness_script);
 
-            if (!pwallet->AddWatchOnly(redeemDestination, timestamp)) {
-                throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
+            // Check that the witnessScript and scriptPubKey match
+            if (GetScriptForDestination(WitnessV0ScriptHash(witness_script)) != script) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "The witnessScript does not match the scriptPubKey or redeemScript");
             }
 
-            // add to address book or update label
-            if (IsValidDestination(dest)) {
-                pwallet->SetAddressBook(dest, label, "receive");
+            // Add the witness script as watch only only if it is not for P2SH-P2WSH
+            if (!scriptpubkey_script.IsPayToScriptHash() && !pwallet->AddWatchOnly(witness_script, timestamp)) {
+                throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
             }
 
-            // Import private keys.
-            if (keys.size()) {
-                for (size_t i = 0; i < keys.size(); i++) {
-                    const std::string& privkey = keys[i].get_str();
-
-                    CKey key = DecodeSecret(privkey);
-
-                    if (!key.IsValid()) {
-                        throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid private key encoding");
-                    }
-
-                    CPubKey pubkey = key.GetPubKey();
-                    assert(key.VerifyPubKey(pubkey));
-
-                    CKeyID vchAddress = pubkey.GetID();
-                    pwallet->MarkDirty();
-                    pwallet->SetAddressBook(vchAddress, label, "receive");
-
-                    if (pwallet->HaveKey(vchAddress)) {
-                        throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Already have this key");
-                    }
-
-                    pwallet->mapKeyMetadata[vchAddress].nCreateTime = timestamp;
+            if (!pwallet->HaveCScript(witness_id) && !pwallet->AddCScript(witness_script)) {
+                throw JSONRPCError(RPC_WALLET_ERROR, "Error adding p2wsh witnessScript to wallet");
+            }
 
-                    if (!pwallet->AddKeyPubKey(key, pubkey)) {
-                        throw JSONRPCError(RPC_WALLET_ERROR, "Error adding key to wallet");
-                    }
+            // Now set script to the witnessScript so we parse the inner script as P2PK or P2PKH below
+            script = witness_script;
+            ExtractDestination(script, dest);
+            allow_p2wpkh = false; // P2WPKH cannot be embedded in P2WSH
+        }
 
-                    pwallet->UpdateTimeFirstKey(timestamp);
-                }
+        // (P2SH-)P2PK/P2PKH/P2WPKH
+        if (dest.type() == typeid(CKeyID) || dest.type() == typeid(WitnessV0KeyHash)) {
+            if (!allow_p2wpkh && dest.type() == typeid(WitnessV0KeyHash)) {
+                throw JSONRPCError(RPC_INVALID_PARAMETER, "P2WPKH cannot be embedded in P2WSH");
             }
-
-            success = true;
-        } else {
-            // Import public keys.
-            if (pubKeys.size() && keys.size() == 0) {
+            if (keys.size() > 1 || pubKeys.size() > 1) {
+                throw JSONRPCError(RPC_INVALID_PARAMETER, "More than one key given for one single-key address");
+            }
+            CPubKey pubkey;
+            if (keys.size()) {
+                pubkey = DecodeSecret(keys[0].get_str()).GetPubKey();
+            }
+            if (pubKeys.size()) {
                 const std::string& strPubKey = pubKeys[0].get_str();
-
                 if (!IsHex(strPubKey)) {
                     throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey must be a hex string");
                 }
-
-                std::vector<unsigned char> vData(ParseHex(strPubKey));
-                CPubKey pubKey(vData.begin(), vData.end());
-
-                if (!pubKey.IsFullyValid()) {
-                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey is not a valid public key");
-                }
-
-                CTxDestination pubkey_dest = pubKey.GetID();
-
-                // Consistency check.
-                if (!(pubkey_dest == dest)) {
-                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Consistency check failed");
-                }
-
-                CScript pubKeyScript = GetScriptForDestination(pubkey_dest);
-
-                if (::IsMine(*pwallet, pubKeyScript) == ISMINE_SPENDABLE) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
+                std::vector<unsigned char> vData(ParseHex(pubKeys[0].get_str()));
+                CPubKey pubkey_temp(vData.begin(), vData.end());
+                if (pubkey.size() && pubkey_temp != pubkey) {
+                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Private key does not match public key for address");
                 }
-
-                pwallet->MarkDirty();
-
-                if (!pwallet->AddWatchOnly(pubKeyScript, timestamp)) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
+                pubkey = pubkey_temp;
+            }
+            if (pubkey.size() > 0) {
+                if (!pubkey.IsFullyValid()) {
+                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey is not a valid public key");
                 }
 
-                // add to address book or update label
-                if (IsValidDestination(pubkey_dest)) {
-                    pwallet->SetAddressBook(pubkey_dest, label, "receive");
+                // Check the key corresponds to the destination given
+                std::vector<CTxDestination> destinations = GetAllDestinationsForKey(pubkey);
+                if (std::find(destinations.begin(), destinations.end(), dest) == destinations.end()) {
+                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Key does not match address destination");
                 }
 
-                // TODO Is this necessary?
-                CScript scriptRawPubKey = GetScriptForRawPubKey(pubKey);
+                // This is necessary to force the wallet to import the pubKey
+                CScript scriptRawPubKey = GetScriptForRawPubKey(pubkey);
 
                 if (::IsMine(*pwallet, scriptRawPubKey) == ISMINE_SPENDABLE) {
                     throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
@@ -999,73 +977,61 @@ static UniValue ProcessImport(CWallet * const pwallet, const UniValue& data, con
                 if (!pwallet->AddWatchOnly(scriptRawPubKey, timestamp)) {
                     throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
                 }
-
-                success = true;
             }
+        }
 
-            // Import private keys.
-            if (keys.size()) {
-                const std::string& strPrivkey = keys[0].get_str();
-
-                // Checks.
-                CKey key = DecodeSecret(strPrivkey);
-
-                if (!key.IsValid()) {
-                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid private key encoding");
-                }
-
-                CPubKey pubKey = key.GetPubKey();
-                assert(key.VerifyPubKey(pubKey));
-
-                CTxDestination pubkey_dest = pubKey.GetID();
+        // Import the address
+        if (::IsMine(*pwallet, scriptpubkey_script) == ISMINE_SPENDABLE) {
+            throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
+        }
 
-                // Consistency check.
-                if (!(pubkey_dest == dest)) {
-                    throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Consistency check failed");
-                }
+        pwallet->MarkDirty();
 
-                CKeyID vchAddress = pubKey.GetID();
-                pwallet->MarkDirty();
-                pwallet->SetAddressBook(vchAddress, label, "receive");
+        if (!pwallet->AddWatchOnly(scriptpubkey_script, timestamp)) {
+            throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
+        }
 
-                if (pwallet->HaveKey(vchAddress)) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
-                }
+        if (!watchOnly && !pwallet->HaveCScript(CScriptID(scriptpubkey_script)) && !pwallet->AddCScript(scriptpubkey_script)) {
+            throw JSONRPCError(RPC_WALLET_ERROR, "Error adding scriptPubKey script to wallet");
+        }
 
-                pwallet->mapKeyMetadata[vchAddress].nCreateTime = timestamp;
+        // add to address book or update label
+        if (IsValidDestination(scriptpubkey_dest)) {
+            pwallet->SetAddressBook(scriptpubkey_dest, label, "receive");
+        }
 
-                if (!pwallet->AddKeyPubKey(key, pubKey)) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "Error adding key to wallet");
-                }
+        // Import private keys.
+        for (size_t i = 0; i < keys.size(); i++) {
+            const std::string& strPrivkey = keys[i].get_str();
 
-                pwallet->UpdateTimeFirstKey(timestamp);
+            // Checks.
+            CKey key = DecodeSecret(strPrivkey);
 
-                success = true;
+            if (!key.IsValid()) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid private key encoding");
             }
 
-            // Import scriptPubKey only.
-            if (pubKeys.size() == 0 && keys.size() == 0) {
-                if (::IsMine(*pwallet, script) == ISMINE_SPENDABLE) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "The wallet already contains the private key for this address or script");
-                }
+            CPubKey pubKey = key.GetPubKey();
+            assert(key.VerifyPubKey(pubKey));
 
-                pwallet->MarkDirty();
+            CKeyID vchAddress = pubKey.GetID();
+            pwallet->MarkDirty();
 
-                if (!pwallet->AddWatchOnly(script, timestamp)) {
-                    throw JSONRPCError(RPC_WALLET_ERROR, "Error adding address to wallet");
-                }
+            if (pwallet->HaveKey(vchAddress)) {
+                throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Already have this key");
+            }
 
-                // add to address book or update label
-                if (IsValidDestination(dest)) {
-                    pwallet->SetAddressBook(dest, label, "receive");
-                }
+            pwallet->mapKeyMetadata[vchAddress].nCreateTime = timestamp;
 
-                success = true;
+            if (!pwallet->AddKeyPubKey(key, pubKey)) {
+                throw JSONRPCError(RPC_WALLET_ERROR, "Error adding key to wallet");
             }
+
+            pwallet->UpdateTimeFirstKey(timestamp);
         }
 
         UniValue result = UniValue(UniValue::VOBJ);
-        result.pushKV("success", UniValue(success));
+        result.pushKV("success", UniValue(true));
         return result;
     } catch (const UniValue& e) {
         UniValue result = UniValue(UniValue::VOBJ);
@@ -1118,7 +1084,8 @@ UniValue importmulti(const JSONRPCRequest& mainRequest)
             "                                                              \"now\" can be specified to bypass scanning, for keys which are known to never have been used, and\n"
             "                                                              0 can be specified to scan the entire blockchain. Blocks up to 2 hours before the earliest key\n"
             "                                                              creation time of all keys being imported by the importmulti call will be scanned.\n"
-            "      \"redeemscript\": \"<script>\"                            , (string, optional) Allowed only if the scriptPubKey is a P2SH address or a P2SH scriptPubKey\n"
+            "      \"redeemscript\": \"<script>\"                            , (string, optional) Allowed only if the scriptPubKey is a P2SH or P2SH-P2WSH address/scriptPubKey\n"
+            "      \"witnessscript\": \"<script>\"                           , (string, optional) Allowed only if the scriptPubKey is a P2SH-P2WSH or P2WSH address/scriptPubKey\n"
             "      \"pubkeys\": [\"<pubKey>\", ... ]                         , (array, optional) Array of strings giving pubkeys that must occur in the output or redeemscript\n"
             "      \"keys\": [\"<key>\", ... ]                               , (array, optional) Array of strings giving private keys whose corresponding public keys must occur in the output or redeemscript\n"
             "      \"internal\": <true>                                    , (boolean, optional, default: false) Stating whether matching outputs should be treated as not incoming payments\n"
diff --git a/test/functional/wallet_importmulti.py b/test/functional/wallet_importmulti.py
index f78ff19791..888e84f7d8 100755
--- a/test/functional/wallet_importmulti.py
+++ b/test/functional/wallet_importmulti.py
@@ -198,7 +198,7 @@ class ImportMultiTest(BitcoinTestFramework):
         }])
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -8)
-        assert_equal(result[0]['error']['message'], 'Incompatibility found between watchonly and keys')
+        assert_equal(result[0]['error']['message'], 'Watch-only addresses should not include private keys')
         address_assert = self.nodes[1].getaddressinfo(address['address'])
         assert_equal(address_assert['iswatchonly'], False)
         assert_equal(address_assert['ismine'], False)
@@ -339,7 +339,7 @@ class ImportMultiTest(BitcoinTestFramework):
         }])
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -8)
-        assert_equal(result[0]['error']['message'], 'Incompatibility found between watchonly and keys')
+        assert_equal(result[0]['error']['message'], 'Watch-only addresses should not include private keys')
 
 
         # Address + Public key + !Internal + Wrong pubkey
@@ -355,7 +355,7 @@ class ImportMultiTest(BitcoinTestFramework):
         }])
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -5)
-        assert_equal(result[0]['error']['message'], 'Consistency check failed')
+        assert_equal(result[0]['error']['message'], 'Key does not match address destination')
         address_assert = self.nodes[1].getaddressinfo(address['address'])
         assert_equal(address_assert['iswatchonly'], False)
         assert_equal(address_assert['ismine'], False)
@@ -375,7 +375,7 @@ class ImportMultiTest(BitcoinTestFramework):
         result = self.nodes[1].importmulti(request)
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -5)
-        assert_equal(result[0]['error']['message'], 'Consistency check failed')
+        assert_equal(result[0]['error']['message'], 'Key does not match address destination')
         address_assert = self.nodes[1].getaddressinfo(address['address'])
         assert_equal(address_assert['iswatchonly'], False)
         assert_equal(address_assert['ismine'], False)
@@ -395,7 +395,7 @@ class ImportMultiTest(BitcoinTestFramework):
         }])
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -5)
-        assert_equal(result[0]['error']['message'], 'Consistency check failed')
+        assert_equal(result[0]['error']['message'], 'Key does not match address destination')
         address_assert = self.nodes[1].getaddressinfo(address['address'])
         assert_equal(address_assert['iswatchonly'], False)
         assert_equal(address_assert['ismine'], False)
@@ -414,7 +414,7 @@ class ImportMultiTest(BitcoinTestFramework):
         }])
         assert_equal(result[0]['success'], False)
         assert_equal(result[0]['error']['code'], -5)
-        assert_equal(result[0]['error']['message'], 'Consistency check failed')
+        assert_equal(result[0]['error']['message'], 'Key does not match address destination')
         address_assert = self.nodes[1].getaddressinfo(address['address'])
         assert_equal(address_assert['iswatchonly'], False)
         assert_equal(address_assert['ismine'], False)