1 files changed, 66 insertions, 5 deletions

diff --git a/src/keystore.cpp b/src/keystore.cpp
index 5718a14ca2..fab1b81c9a 100644
--- a/src/keystore.cpp
+++ b/src/keystore.cpp
@@ -1,18 +1,41 @@
 // Copyright (c) 2009-2010 Satoshi Nakamoto
-// Copyright (c) 2009-2016 The Bitcoin Core developers
+// Copyright (c) 2009-2017 The Bitcoin Core developers
 // Distributed under the MIT software license, see the accompanying
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
 #include <keystore.h>
 
-#include <key.h>
-#include <pubkey.h>
 #include <util.h>
 
 bool CKeyStore::AddKey(const CKey &key) {
     return AddKeyPubKey(key, key.GetPubKey());
 }
 
+void CBasicKeyStore::ImplicitlyLearnRelatedKeyScripts(const CPubKey& pubkey)
+{
+    AssertLockHeld(cs_KeyStore);
+    CKeyID key_id = pubkey.GetID();
+    // We must actually know about this key already.
+    assert(HaveKey(key_id) || mapWatchKeys.count(key_id));
+    // This adds the redeemscripts necessary to detect P2WPKH and P2SH-P2WPKH
+    // outputs. Technically P2WPKH outputs don't have a redeemscript to be
+    // spent. However, our current IsMine logic requires the corresponding
+    // P2SH-P2WPKH redeemscript to be present in the wallet in order to accept
+    // payment even to P2WPKH outputs.
+    // Also note that having superfluous scripts in the keystore never hurts.
+    // They're only used to guide recursion in signing and IsMine logic - if
+    // a script is present but we can't do anything with it, it has no effect.
+    // "Implicitly" refers to fact that scripts are derived automatically from
+    // existing keys, and are present in memory, even without being explicitly
+    // loaded (e.g. from a file).
+    if (pubkey.IsCompressed()) {
+        CScript script = GetScriptForDestination(WitnessV0KeyHash(key_id));
+        // This does not use AddCScript, as it may be overridden.
+        CScriptID id(script);
+        mapScripts[id] = std::move(script);
+    }
+}
+
 bool CBasicKeyStore::GetPubKey(const CKeyID &address, CPubKey &vchPubKeyOut) const
 {
     CKey key;
@@ -33,6 +56,7 @@ bool CBasicKeyStore::AddKeyPubKey(const CKey& key, const CPubKey &pubkey)
 {
     LOCK(cs_KeyStore);
     mapKeys[pubkey.GetID()] = key;
+    ImplicitlyLearnRelatedKeyScripts(pubkey);
     return true;
 }
 
@@ -79,6 +103,16 @@ bool CBasicKeyStore::HaveCScript(const CScriptID& hash) const
     return mapScripts.count(hash) > 0;
 }
 
+std::set<CScriptID> CBasicKeyStore::GetCScripts() const
+{
+    LOCK(cs_KeyStore);
+    std::set<CScriptID> set_script;
+    for (const auto& mi : mapScripts) {
+        set_script.insert(mi.first);
+    }
+    return set_script;
+}
+
 bool CBasicKeyStore::GetCScript(const CScriptID &hash, CScript& redeemScriptOut) const
 {
     LOCK(cs_KeyStore);
@@ -112,8 +146,10 @@ bool CBasicKeyStore::AddWatchOnly(const CScript &dest)
     LOCK(cs_KeyStore);
     setWatchOnly.insert(dest);
     CPubKey pubKey;
-    if (ExtractPubKey(dest, pubKey))
+    if (ExtractPubKey(dest, pubKey)) {
         mapWatchKeys[pubKey.GetID()] = pubKey;
+        ImplicitlyLearnRelatedKeyScripts(pubKey);
+    }
     return true;
 }
 
@@ -122,8 +158,11 @@ bool CBasicKeyStore::RemoveWatchOnly(const CScript &dest)
     LOCK(cs_KeyStore);
     setWatchOnly.erase(dest);
     CPubKey pubKey;
-    if (ExtractPubKey(dest, pubKey))
+    if (ExtractPubKey(dest, pubKey)) {
         mapWatchKeys.erase(pubKey.GetID());
+    }
+    // Related CScripts are not removed; having superfluous scripts around is
+    // harmless (see comment in ImplicitlyLearnRelatedKeyScripts).
     return true;
 }
 
@@ -138,3 +177,25 @@ bool CBasicKeyStore::HaveWatchOnly() const
     LOCK(cs_KeyStore);
     return (!setWatchOnly.empty());
 }
+
+CKeyID GetKeyForDestination(const CKeyStore& store, const CTxDestination& dest)
+{
+    // Only supports destinations which map to single public keys, i.e. P2PKH,
+    // P2WPKH, and P2SH-P2WPKH.
+    if (auto id = boost::get<CKeyID>(&dest)) {
+        return *id;
+    }
+    if (auto witness_id = boost::get<WitnessV0KeyHash>(&dest)) {
+        return CKeyID(*witness_id);
+    }
+    if (auto script_id = boost::get<CScriptID>(&dest)) {
+        CScript script;
+        CTxDestination inner_dest;
+        if (store.GetCScript(*script_id, script) && ExtractDestination(script, inner_dest)) {
+            if (auto inner_witness_id = boost::get<WitnessV0KeyHash>(&inner_dest)) {
+                return CKeyID(*inner_witness_id);
+            }
+        }
+    }
+    return CKeyID();
+}