From f22b5018286b7c3739675359d14557fc180ecef9 Mon Sep 17 00:00:00 2001
From: David Burkett <davidburkett38@gmail.com>
Date: Sat, 29 Jan 2022 15:06:24 -0500
Subject: [PATCH] MWEB: Add DesitnationAddr wrapper

---
 src/Makefile.am                |  2 ++
 src/script/address.cpp         | 42 ++++++++++++++++++++++++++++++++++
 src/script/address.h           | 37 ++++++++++++++++++++++++++++++
 src/script/sign.cpp            |  8 ++++++-
 src/script/sign.h              |  3 ++-
 src/script/signingprovider.cpp |  3 +++
 src/script/standard.h          |  2 ++
 7 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 src/script/address.cpp
 create mode 100644 src/script/address.h

diff --git a/src/Makefile.am b/src/Makefile.am
index 2e2287b09..77820c6ad 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -242,6 +242,7 @@ BITCOIN_CORE_H = \
   rpc/server.h \
   rpc/util.h \
   scheduler.h \
+  script/address.h \
   script/descriptor.h \
   script/keyorigin.h \
   script/sigcache.h \
@@ -557,6 +558,7 @@ libbitcoin_common_a_SOURCES = \
   rpc/rawtransaction_util.cpp \
   rpc/util.cpp \
   scheduler.cpp \
+  script/address.cpp \
   script/descriptor.cpp \
   script/sign.cpp \
   script/signingprovider.cpp \
diff --git a/src/script/address.cpp b/src/script/address.cpp
new file mode 100644
index 000000000..3a88e135c
--- /dev/null
+++ b/src/script/address.cpp
@@ -0,0 +1,42 @@
+#include <script/address.h>
+#include <key_io.h>
+
+DestinationAddr::DestinationAddr(const CTxDestination& dest)
+{
+    if (dest.type() == typeid(StealthAddress)) {
+        m_script = boost::get<StealthAddress>(dest);
+    } else {
+        m_script = ::GetScriptForDestination(dest);
+    }
+}
+
+std::string DestinationAddr::Encode() const
+{
+    CTxDestination dest;
+    if (ExtractDestination(dest)) {
+        return ::EncodeDestination(dest);
+    }
+
+    return HexStr(GetScript());
+}
+
+const CScript& DestinationAddr::GetScript() const noexcept
+{
+    assert(m_script.type() == typeid(CScript));
+    return boost::get<CScript>(m_script);
+}
+
+const StealthAddress& DestinationAddr::GetMWEBAddress() const noexcept
+{
+    assert(m_script.type() == typeid(StealthAddress));
+    return boost::get<StealthAddress>(m_script);
+}
+
+bool DestinationAddr::ExtractDestination(CTxDestination& dest) const
+{
+    if (IsMWEB()) {
+        dest = GetMWEBAddress();
+        return true;
+    }
+    return ::ExtractDestination(GetScript(), dest);
+}
\ No newline at end of file
diff --git a/src/script/address.h b/src/script/address.h
new file mode 100644
index 000000000..ee7d443a4
--- /dev/null
+++ b/src/script/address.h
@@ -0,0 +1,37 @@
+#pragma once
+
+#include <script/standard.h>
+#include <boost/variant.hpp>
+
+//
+// A wrapper around a boost::variant<CScript, StealthAddress> to deal with cases where
+// a raw CScript is passed around as an address, rather than a CTxDestination.
+// Since StealthAddresses can't be represented as a CScript, this wrapper should be used
+// wherever the address could contain a StealthAddress instead of a CScript
+//
+class DestinationAddr
+{
+public:
+    DestinationAddr() = default;
+    DestinationAddr(CScript script)
+        : m_script(std::move(script)) {}
+    DestinationAddr(StealthAddress address)
+        : m_script(std::move(address)) {}
+    DestinationAddr(const CTxDestination& dest);
+
+    bool operator==(const DestinationAddr& rhs) const noexcept { return this->m_script == rhs.m_script; }
+    bool operator<(const DestinationAddr& rhs) const noexcept { return this->m_script < rhs.m_script; }
+    bool operator<=(const DestinationAddr& rhs) const noexcept { return this->m_script <= rhs.m_script; }
+
+    std::string Encode() const;
+
+    bool IsMWEB() const noexcept { return m_script.type() == typeid(StealthAddress); }
+    bool IsEmpty() const noexcept { return !IsMWEB() && GetScript().empty(); }
+
+    const CScript& GetScript() const noexcept;
+    const StealthAddress& GetMWEBAddress() const noexcept;
+    bool ExtractDestination(CTxDestination& dest) const;
+
+private:
+    boost::variant<CScript, StealthAddress> m_script;
+};
\ No newline at end of file
diff --git a/src/script/sign.cpp b/src/script/sign.cpp
index 903c95a5c..d8003880f 100644
--- a/src/script/sign.cpp
+++ b/src/script/sign.cpp
@@ -434,8 +434,14 @@ public:
 const BaseSignatureCreator& DUMMY_SIGNATURE_CREATOR = DummySignatureCreator(32, 32);
 const BaseSignatureCreator& DUMMY_MAXIMUM_SIGNATURE_CREATOR = DummySignatureCreator(33, 32);
 
-bool IsSolvable(const SigningProvider& provider, const CScript& script)
+bool IsSolvable(const SigningProvider& provider, const DestinationAddr& dest_addr)
 {
+    if (dest_addr.IsMWEB()) {
+        CPubKey pubkey(dest_addr.GetMWEBAddress().GetSpendPubKey().vec());
+        return provider.HaveKey(pubkey.GetID());
+    }
+
+    const CScript& script = dest_addr.GetScript();
     // This check is to make sure that the script we created can actually be solved for and signed by us
     // if we were to have the private keys. This is just to make sure that the script is valid and that,
     // if found in a transaction, we would still accept and relay that transaction. In particular,
diff --git a/src/script/sign.h b/src/script/sign.h
index a1cfe1574..f2e2f5836 100644
--- a/src/script/sign.h
+++ b/src/script/sign.h
@@ -9,6 +9,7 @@
 #include <coins.h>
 #include <hash.h>
 #include <pubkey.h>
+#include <script/address.h>
 #include <script/interpreter.h>
 #include <script/keyorigin.h>
 #include <span.h>
@@ -165,7 +166,7 @@ void UpdateInput(CTxIn& input, const SignatureData& data);
  * have all private keys. While this function does not need private keys, the passed
  * provider is used to look up public keys and redeemscripts by hash.
  * Solvability is unrelated to whether we consider this output to be ours. */
-bool IsSolvable(const SigningProvider& provider, const CScript& script);
+bool IsSolvable(const SigningProvider& provider, const DestinationAddr& dest_addr);
 
 /** Check whether a scriptPubKey is known to be segwit. */
 bool IsSegWitOutput(const SigningProvider& provider, const CScript& script);
diff --git a/src/script/signingprovider.cpp b/src/script/signingprovider.cpp
index 2d8dc7d47..354aceff3 100644
--- a/src/script/signingprovider.cpp
+++ b/src/script/signingprovider.cpp
@@ -195,5 +195,8 @@ CKeyID GetKeyForDestination(const SigningProvider& store, const CTxDestination&
             }
         }
     }
+    if (auto stealth_address = boost::get<StealthAddress>(&dest)) {
+        return CPubKey(stealth_address->B().vec()).GetID();
+    }
     return CKeyID();
 }
diff --git a/src/script/standard.h b/src/script/standard.h
index 2608ab1ff..5f4eff3da 100644
--- a/src/script/standard.h
+++ b/src/script/standard.h
@@ -273,4 +273,6 @@ bool IsPegInOutput(const CTxOutput& output);
 
 CScript GetScriptForPegin(const mw::Hash& kernel_id);
 
+#include <script/address.h>
+
 #endif // BITCOIN_SCRIPT_STANDARD_H