From 6a9c429084b40356aa36aa67992da35f61c2f6a2 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Thu, 13 Feb 2020 17:09:12 -0500
Subject: [PATCH] Move direct calls to MessageSign into new SignMessage
 functions in CWallet and ScriptPubKeyMan

Instead of getting a SigningProvider and then going to MessageSign,
have ScriptPubKeyMan handle the message signing internally.
---
 src/interfaces/wallet.cpp          |  4 ++++
 src/interfaces/wallet.h            |  4 ++++
 src/qt/signverifymessagedialog.cpp | 27 +++++++++++++++++----------
 src/util/message.cpp               | 14 ++++++++++++++
 src/util/message.h                 |  8 ++++++++
 src/wallet/rpcwallet.cpp           | 20 +++++---------------
 src/wallet/scriptpubkeyman.cpp     | 14 ++++++++++++++
 src/wallet/scriptpubkeyman.h       |  4 ++++
 src/wallet/wallet.cpp              | 12 ++++++++++++
 src/wallet/wallet.h                |  2 ++
 10 files changed, 84 insertions(+), 25 deletions(-)

diff --git a/src/interfaces/wallet.cpp b/src/interfaces/wallet.cpp
index 83ce9aba740..7b3ab27b503 100644
--- a/src/interfaces/wallet.cpp
+++ b/src/interfaces/wallet.cpp
@@ -132,6 +132,10 @@ public:
         }
         return false;
     }
+    SigningResult signMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) override
+    {
+        return m_wallet->SignMessage(message, pkhash, str_sig);
+    }
     bool isSpendable(const CTxDestination& dest) override { return m_wallet->IsMine(dest) & ISMINE_SPENDABLE; }
     bool haveWatchOnly() override
     {
diff --git a/src/interfaces/wallet.h b/src/interfaces/wallet.h
index 964608ff827..6f27ee126d8 100644
--- a/src/interfaces/wallet.h
+++ b/src/interfaces/wallet.h
@@ -10,6 +10,7 @@
 #include <script/standard.h>           // For CTxDestination
 #include <support/allocators/secure.h> // For SecureString
 #include <ui_interface.h>              // For ChangeType
+#include <util/message.h>
 
 #include <functional>
 #include <map>
@@ -87,6 +88,9 @@ public:
     //! Get private key.
     virtual bool getPrivKey(const CScript& script, const CKeyID& address, CKey& key) = 0;
 
+    //! Sign message
+    virtual SigningResult signMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) = 0;
+
     //! Return whether wallet has private key.
     virtual bool isSpendable(const CTxDestination& dest) = 0;
 
diff --git a/src/qt/signverifymessagedialog.cpp b/src/qt/signverifymessagedialog.cpp
index 883dcecf9a1..4552753bf6d 100644
--- a/src/qt/signverifymessagedialog.cpp
+++ b/src/qt/signverifymessagedialog.cpp
@@ -133,20 +133,27 @@ void SignVerifyMessageDialog::on_signMessageButton_SM_clicked()
         return;
     }
 
-    CKey key;
-    if (!model->wallet().getPrivKey(GetScriptForDestination(destination), CKeyID(*pkhash), key))
-    {
-        ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
-        ui->statusLabel_SM->setText(tr("Private key for the entered address is not available."));
-        return;
-    }
-
     const std::string& message = ui->messageIn_SM->document()->toPlainText().toStdString();
     std::string signature;
+    SigningResult res = model->wallet().signMessage(message, *pkhash, signature);
 
-    if (!MessageSign(key, message, signature)) {
+    QString error;
+    switch (res) {
+        case SigningResult::OK:
+            error = tr("No error");
+            break;
+        case SigningResult::PRIVATE_KEY_NOT_AVAILABLE:
+            error = tr("Private key for the entered address is not available.");
+            break;
+        case SigningResult::SIGNING_FAILED:
+            error = tr("Message signing failed.");
+            break;
+        // no default case, so the compiler can warn about missing cases
+    }
+
+    if (res != SigningResult::OK) {
         ui->statusLabel_SM->setStyleSheet("QLabel { color: red; }");
-        ui->statusLabel_SM->setText(QString("<nobr>") + tr("Message signing failed.") + QString("</nobr>"));
+        ui->statusLabel_SM->setText(QString("<nobr>") + error + QString("</nobr>"));
         return;
     }
 
diff --git a/src/util/message.cpp b/src/util/message.cpp
index 17603a43d2d..1e7128d2252 100644
--- a/src/util/message.cpp
+++ b/src/util/message.cpp
@@ -76,3 +76,17 @@ uint256 MessageHash(const std::string& message)
 
     return hasher.GetHash();
 }
+
+std::string SigningResultString(const SigningResult res)
+{
+    switch (res) {
+        case SigningResult::OK:
+            return "No error";
+        case SigningResult::PRIVATE_KEY_NOT_AVAILABLE:
+            return "Private key not available";
+        case SigningResult::SIGNING_FAILED:
+            return "Sign failed";
+        // no default case, so the compiler can warn about missing cases
+    }
+    assert(false);
+}
diff --git a/src/util/message.h b/src/util/message.h
index 01fd14ce2dc..b31c5f5761b 100644
--- a/src/util/message.h
+++ b/src/util/message.h
@@ -39,6 +39,12 @@ enum class MessageVerificationResult {
     OK
 };
 
+enum class SigningResult {
+    OK, //!< No error
+    PRIVATE_KEY_NOT_AVAILABLE,
+    SIGNING_FAILED,
+};
+
 /** Verify a signed message.
  * @param[in] address Signer's bitcoin address, it must refer to a public key.
  * @param[in] signature The signature in base64 format.
@@ -65,4 +71,6 @@ bool MessageSign(
  */
 uint256 MessageHash(const std::string& message);
 
+std::string SigningResultString(const SigningResult res);
+
 #endif // BITCOIN_UTIL_MESSAGE_H
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
index 760bc1c408a..90fd3e43738 100644
--- a/src/wallet/rpcwallet.cpp
+++ b/src/wallet/rpcwallet.cpp
@@ -565,22 +565,12 @@ static UniValue signmessage(const JSONRPCRequest& request)
         throw JSONRPCError(RPC_TYPE_ERROR, "Address does not refer to key");
     }
 
-    CScript script_pub_key = GetScriptForDestination(*pkhash);
-    std::unique_ptr<SigningProvider> provider = pwallet->GetSigningProvider(script_pub_key);
-    if (!provider) {
-        throw JSONRPCError(RPC_WALLET_ERROR, "Private key not available");
-    }
-
-    CKey key;
-    CKeyID keyID(*pkhash);
-    if (!provider->GetKey(keyID, key)) {
-        throw JSONRPCError(RPC_WALLET_ERROR, "Private key not available");
-    }
-
     std::string signature;
-
-    if (!MessageSign(key, strMessage, signature)) {
-        throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Sign failed");
+    SigningResult err = pwallet->SignMessage(strMessage, *pkhash, signature);
+    if (err == SigningResult::SIGNING_FAILED) {
+        throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, SigningResultString(err));
+    } else if (err != SigningResult::OK){
+        throw JSONRPCError(RPC_WALLET_ERROR, SigningResultString(err));
     }
 
     return signature;
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index cfb8184c3b2..bec4abf8572 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -511,6 +511,20 @@ bool LegacyScriptPubKeyMan::SignTransaction(CMutableTransaction& tx, const std::
     return ::SignTransaction(tx, this, coins, sighash, input_errors);
 }
 
+SigningResult LegacyScriptPubKeyMan::SignMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) const
+{
+    CKeyID key_id(pkhash);
+    CKey key;
+    if (!GetKey(key_id, key)) {
+        return SigningResult::PRIVATE_KEY_NOT_AVAILABLE;
+    }
+
+    if (MessageSign(key, message, str_sig)) {
+        return SigningResult::OK;
+    }
+    return SigningResult::SIGNING_FAILED;
+}
+
 TransactionError LegacyScriptPubKeyMan::FillPSBT(PartiallySignedTransaction& psbtx, int sighash_type, bool sign, bool bip32derivs) const
 {
     for (unsigned int i = 0; i < psbtx.tx->vin.size(); ++i) {
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 1628001a744..ab0d1c37bd1 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -9,6 +9,7 @@
 #include <script/signingprovider.h>
 #include <script/standard.h>
 #include <util/error.h>
+#include <util/message.h>
 #include <wallet/crypter.h>
 #include <wallet/ismine.h>
 #include <wallet/walletdb.h>
@@ -214,6 +215,8 @@ public:
 
     /** Creates new signatures and adds them to the transaction. Returns whether all inputs were signed */
     virtual bool SignTransaction(CMutableTransaction& tx, const std::map<COutPoint, Coin>& coins, int sighash, std::map<int, std::string>& input_errors) const { return false; }
+    /** Sign a message with the given script */
+    virtual SigningResult SignMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) const { return SigningResult::SIGNING_FAILED; };
     /** Adds script and derivation path information to a PSBT, and optionally signs it. */
     virtual TransactionError FillPSBT(PartiallySignedTransaction& psbt, int sighash_type = 1 /* SIGHASH_ALL */, bool sign = true, bool bip32derivs = false) const { return TransactionError::INVALID_PSBT; }
 
@@ -358,6 +361,7 @@ public:
     bool CanProvide(const CScript& script, SignatureData& sigdata) override;
 
     bool SignTransaction(CMutableTransaction& tx, const std::map<COutPoint, Coin>& coins, int sighash, std::map<int, std::string>& input_errors) const override;
+    SigningResult SignMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) const override;
     TransactionError FillPSBT(PartiallySignedTransaction& psbt, int sighash_type = 1 /* SIGHASH_ALL */, bool sign = true, bool bip32derivs = false) const override;
 
     uint256 GetID() const override;
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index de09e60014f..058f6597ae4 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2566,6 +2566,18 @@ TransactionError CWallet::FillPSBT(PartiallySignedTransaction& psbtx, bool& comp
     return TransactionError::OK;
 }
 
+SigningResult CWallet::SignMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) const
+{
+    SignatureData sigdata;
+    CScript script_pub_key = GetScriptForDestination(pkhash);
+    for (const auto& spk_man_pair : m_spk_managers) {
+        if (spk_man_pair.second->CanProvide(script_pub_key, sigdata)) {
+            return spk_man_pair.second->SignMessage(message, pkhash, str_sig);
+        }
+    }
+    return SigningResult::PRIVATE_KEY_NOT_AVAILABLE;
+}
+
 bool CWallet::FundTransaction(CMutableTransaction& tx, CAmount& nFeeRet, int& nChangePosInOut, std::string& strFailReason, bool lockUnspents, const std::set<int>& setSubtractFeeFromOutputs, CCoinControl coinControl)
 {
     std::vector<CRecipient> vecSend;
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 15353de4ace..de37930f4e0 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -14,6 +14,7 @@
 #include <psbt.h>
 #include <tinyformat.h>
 #include <ui_interface.h>
+#include <util/message.h>
 #include <util/strencodings.h>
 #include <util/system.h>
 #include <validationinterface.h>
@@ -921,6 +922,7 @@ public:
     bool SignTransaction(CMutableTransaction& tx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     // Sign the tx given the input coins and sighash.
     bool SignTransaction(CMutableTransaction& tx, const std::map<COutPoint, Coin>& coins, int sighash, std::map<int, std::string>& input_errors) const;
+    SigningResult SignMessage(const std::string& message, const PKHash& pkhash, std::string& str_sig) const;
 
     /**
      * Fills out a PSBT with information from the wallet. Fills in UTXOs if we have