1  // Copyright (c) 2021-2022 The Bitcoin Core developers
  2  // Distributed under the MIT software license, see the accompanying
  3  // file COPYING or http://www.opensource.org/licenses/mit-license.php.
  4  
  5  #include <wallet/test/util.h>
  6  
  7  #include <chain.h>
  8  #include <key.h>
  9  #include <key_io.h>
 10  #include <streams.h>
 11  #include <test/util/setup_common.h>
 12  #include <validationinterface.h>
 13  #include <wallet/context.h>
 14  #include <wallet/wallet.h>
 15  #include <wallet/walletdb.h>
 16  
 17  #include <memory>
 18  
 19  namespace wallet {
 20  std::unique_ptr<CWallet> CreateSyncedWallet(interfaces::Chain& chain, CChain& cchain, const CKey& key)
 21  {
 22      auto wallet = std::make_unique<CWallet>(&chain, "", CreateMockableWalletDatabase());
 23      {
 24          LOCK2(wallet->cs_wallet, ::cs_main);
 25          wallet->SetLastBlockProcessed(cchain.Height(), cchain.Tip()->GetBlockHash());
 26      }
 27      {
 28          LOCK(wallet->cs_wallet);
 29          wallet->SetWalletFlag(WALLET_FLAG_DESCRIPTORS);
 30          wallet->SetupDescriptorScriptPubKeyMans();
 31  
 32          FlatSigningProvider provider;
 33          std::string error;
 34          std::unique_ptr<Descriptor> desc = Parse("combo(" + EncodeSecret(key) + ")", provider, error, /* require_checksum=*/ false);
 35          assert(desc);
 36          WalletDescriptor w_desc(std::move(desc), 0, 0, 1, 1);
 37          if (!wallet->AddWalletDescriptor(w_desc, provider, "", false)) assert(false);
 38      }
 39      WalletRescanReserver reserver(*wallet);
 40      reserver.reserve();
 41      CWallet::ScanResult result = wallet->ScanForWalletTransactions(cchain.Genesis()->GetBlockHash(), /*start_height=*/0, /*max_height=*/{}, reserver, /*fUpdate=*/false, /*save_progress=*/false);
 42      assert(result.status == CWallet::ScanResult::SUCCESS);
 43      assert(result.last_scanned_block == cchain.Tip()->GetBlockHash());
 44      assert(*result.last_scanned_height == cchain.Height());
 45      assert(result.last_failed_block.IsNull());
 46      return wallet;
 47  }
 48  
 49  std::shared_ptr<CWallet> TestLoadWallet(std::unique_ptr<WalletDatabase> database, WalletContext& context, uint64_t create_flags)
 50  {
 51      bilingual_str error;
 52      std::vector<bilingual_str> warnings;
 53      auto wallet = CWallet::Create(context, "", std::move(database), create_flags, error, warnings);
 54      NotifyWalletLoaded(context, wallet);
 55      if (context.chain) {
 56          wallet->postInitProcess();
 57      }
 58      return wallet;
 59  }
 60  
 61  std::shared_ptr<CWallet> TestLoadWallet(WalletContext& context)
 62  {
 63      DatabaseOptions options;
 64      options.create_flags = WALLET_FLAG_DESCRIPTORS;
 65      DatabaseStatus status;
 66      bilingual_str error;
 67      std::vector<bilingual_str> warnings;
 68      auto database = MakeWalletDatabase("", options, status, error);
 69      return TestLoadWallet(std::move(database), context, options.create_flags);
 70  }
 71  
 72  void TestUnloadWallet(std::shared_ptr<CWallet>&& wallet)
 73  {
 74      // Calls SyncWithValidationInterfaceQueue
 75      wallet->chain().waitForNotificationsIfTipChanged({});
 76      wallet->m_chain_notifications_handler.reset();
 77      UnloadWallet(std::move(wallet));
 78  }
 79  
 80  std::unique_ptr<WalletDatabase> DuplicateMockDatabase(WalletDatabase& database)
 81  {
 82      return std::make_unique<MockableDatabase>(dynamic_cast<MockableDatabase&>(database).m_records);
 83  }
 84  
 85  std::string getnewaddress(CWallet& w)
 86  {
 87      constexpr auto output_type = OutputType::BECH32;
 88      return EncodeDestination(getNewDestination(w, output_type));
 89  }
 90  
 91  CTxDestination getNewDestination(CWallet& w, OutputType output_type)
 92  {
 93      return *Assert(w.GetNewDestination(output_type, ""));
 94  }
 95  
 96  // BytePrefix compares equality with other byte spans that begin with the same prefix.
 97  struct BytePrefix { Span<const std::byte> prefix; };
 98  bool operator<(BytePrefix a, Span<const std::byte> b) { return a.prefix < b.subspan(0, std::min(a.prefix.size(), b.size())); }
 99  bool operator<(Span<const std::byte> a, BytePrefix b) { return a.subspan(0, std::min(a.size(), b.prefix.size())) < b.prefix; }
100  
101  MockableCursor::MockableCursor(const MockableData& records, bool pass, Span<const std::byte> prefix)
102  {
103      m_pass = pass;
104      std::tie(m_cursor, m_cursor_end) = records.equal_range(BytePrefix{prefix});
105  }
106  
107  DatabaseCursor::Status MockableCursor::Next(DataStream& key, DataStream& value)
108  {
109      if (!m_pass) {
110          return Status::FAIL;
111      }
112      if (m_cursor == m_cursor_end) {
113          return Status::DONE;
114      }
115      key.clear();
116      value.clear();
117      const auto& [key_data, value_data] = *m_cursor;
118      key.write(key_data);
119      value.write(value_data);
120      m_cursor++;
121      return Status::MORE;
122  }
123  
124  bool MockableBatch::ReadKey(DataStream&& key, DataStream& value)
125  {
126      if (!m_pass) {
127          return false;
128      }
129      SerializeData key_data{key.begin(), key.end()};
130      const auto& it = m_records.find(key_data);
131      if (it == m_records.end()) {
132          return false;
133      }
134      value.clear();
135      value.write(it->second);
136      return true;
137  }
138  
139  bool MockableBatch::WriteKey(DataStream&& key, DataStream&& value, bool overwrite)
140  {
141      if (!m_pass) {
142          return false;
143      }
144      SerializeData key_data{key.begin(), key.end()};
145      SerializeData value_data{value.begin(), value.end()};
146      auto [it, inserted] = m_records.emplace(key_data, value_data);
147      if (!inserted && overwrite) { // Overwrite if requested
148          it->second = value_data;
149          inserted = true;
150      }
151      return inserted;
152  }
153  
154  bool MockableBatch::EraseKey(DataStream&& key)
155  {
156      if (!m_pass) {
157          return false;
158      }
159      SerializeData key_data{key.begin(), key.end()};
160      m_records.erase(key_data);
161      return true;
162  }
163  
164  bool MockableBatch::HasKey(DataStream&& key)
165  {
166      if (!m_pass) {
167          return false;
168      }
169      SerializeData key_data{key.begin(), key.end()};
170      return m_records.count(key_data) > 0;
171  }
172  
173  bool MockableBatch::ErasePrefix(Span<const std::byte> prefix)
174  {
175      if (!m_pass) {
176          return false;
177      }
178      auto it = m_records.begin();
179      while (it != m_records.end()) {
180          auto& key = it->first;
181          if (key.size() < prefix.size() || std::search(key.begin(), key.end(), prefix.begin(), prefix.end()) != key.begin()) {
182              it++;
183              continue;
184          }
185          it = m_records.erase(it);
186      }
187      return true;
188  }
189  
190  std::unique_ptr<WalletDatabase> CreateMockableWalletDatabase(MockableData records)
191  {
192      return std::make_unique<MockableDatabase>(records);
193  }
194  
195  MockableDatabase& GetMockableDatabase(CWallet& wallet)
196  {
197      return dynamic_cast<MockableDatabase&>(wallet.GetDatabase());
198  }
199  } // namespace wallet