Merge pull request #2766 from TheBlueMatt/2023-12-2314-cleanups-2
[rust-lightning] / lightning / src / util / test_utils.rs
1 // This file is Copyright its original authors, visible in version control
2 // history.
3 //
4 // This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
5 // or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
6 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
7 // You may not use this file except in accordance with one or both of these
8 // licenses.
9
10 use crate::chain;
11 use crate::chain::WatchedOutput;
12 use crate::chain::chaininterface;
13 use crate::chain::chaininterface::ConfirmationTarget;
14 use crate::chain::chaininterface::FEERATE_FLOOR_SATS_PER_KW;
15 use crate::chain::chainmonitor;
16 use crate::chain::chainmonitor::{MonitorUpdateId, UpdateOrigin};
17 use crate::chain::channelmonitor;
18 use crate::chain::channelmonitor::MonitorEvent;
19 use crate::chain::transaction::OutPoint;
20 use crate::sign;
21 use crate::events;
22 use crate::events::bump_transaction::{WalletSource, Utxo};
23 use crate::ln::ChannelId;
24 use crate::ln::channelmanager;
25 use crate::ln::chan_utils::CommitmentTransaction;
26 use crate::ln::features::{ChannelFeatures, InitFeatures, NodeFeatures};
27 use crate::ln::{msgs, wire};
28 use crate::ln::msgs::LightningError;
29 use crate::ln::script::ShutdownScript;
30 use crate::offers::invoice::UnsignedBolt12Invoice;
31 use crate::offers::invoice_request::UnsignedInvoiceRequest;
32 use crate::routing::gossip::{EffectiveCapacity, NetworkGraph, NodeId};
33 use crate::routing::utxo::{UtxoLookup, UtxoLookupError, UtxoResult};
34 use crate::routing::router::{find_route, InFlightHtlcs, Path, Route, RouteParameters, Router, ScorerAccountingForInFlightHtlcs};
35 use crate::routing::scoring::{ChannelUsage, ScoreUpdate, ScoreLookUp};
36 use crate::sync::RwLock;
37 use crate::util::config::UserConfig;
38 use crate::util::test_channel_signer::{TestChannelSigner, EnforcementState};
39 use crate::util::logger::{Logger, Level, Record};
40 use crate::util::ser::{Readable, ReadableArgs, Writer, Writeable};
41 use crate::util::persist::KVStore;
42
43 use bitcoin::blockdata::constants::ChainHash;
44 use bitcoin::blockdata::constants::genesis_block;
45 use bitcoin::blockdata::transaction::{Transaction, TxOut};
46 use bitcoin::blockdata::script::{Builder, Script, ScriptBuf};
47 use bitcoin::blockdata::opcodes;
48 use bitcoin::blockdata::block::Block;
49 use bitcoin::network::constants::Network;
50 use bitcoin::hash_types::{BlockHash, Txid};
51 use bitcoin::sighash::{SighashCache, EcdsaSighashType};
52
53 use bitcoin::secp256k1::{PublicKey, Scalar, Secp256k1, SecretKey};
54 use bitcoin::secp256k1::ecdh::SharedSecret;
55 use bitcoin::secp256k1::ecdsa::{RecoverableSignature, Signature};
56 use bitcoin::secp256k1::schnorr;
57
58 #[cfg(any(test, feature = "_test_utils"))]
59 use regex;
60
61 use crate::io;
62 use crate::prelude::*;
63 use core::cell::RefCell;
64 use core::time::Duration;
65 use crate::sync::{Mutex, Arc};
66 use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
67 use core::mem;
68 use bitcoin::bech32::u5;
69 use crate::sign::{InMemorySigner, Recipient, EntropySource, NodeSigner, SignerProvider};
70
71 #[cfg(feature = "std")]
72 use std::time::{SystemTime, UNIX_EPOCH};
73 use bitcoin::Sequence;
74
75 pub fn pubkey(byte: u8) -> PublicKey {
76         let secp_ctx = Secp256k1::new();
77         PublicKey::from_secret_key(&secp_ctx, &privkey(byte))
78 }
79
80 pub fn privkey(byte: u8) -> SecretKey {
81         SecretKey::from_slice(&[byte; 32]).unwrap()
82 }
83
84 pub struct TestVecWriter(pub Vec<u8>);
85 impl Writer for TestVecWriter {
86         fn write_all(&mut self, buf: &[u8]) -> Result<(), io::Error> {
87                 self.0.extend_from_slice(buf);
88                 Ok(())
89         }
90 }
91
92 pub struct TestFeeEstimator {
93         pub sat_per_kw: Mutex<u32>,
94 }
95 impl chaininterface::FeeEstimator for TestFeeEstimator {
96         fn get_est_sat_per_1000_weight(&self, _confirmation_target: ConfirmationTarget) -> u32 {
97                 *self.sat_per_kw.lock().unwrap()
98         }
99 }
100
101 pub struct TestRouter<'a> {
102         pub network_graph: Arc<NetworkGraph<&'a TestLogger>>,
103         pub next_routes: Mutex<VecDeque<(RouteParameters, Result<Route, LightningError>)>>,
104         pub scorer: &'a RwLock<TestScorer>,
105 }
106
107 impl<'a> TestRouter<'a> {
108         pub fn new(network_graph: Arc<NetworkGraph<&'a TestLogger>>, scorer: &'a RwLock<TestScorer>) -> Self {
109                 Self { network_graph, next_routes: Mutex::new(VecDeque::new()), scorer }
110         }
111
112         pub fn expect_find_route(&self, query: RouteParameters, result: Result<Route, LightningError>) {
113                 let mut expected_routes = self.next_routes.lock().unwrap();
114                 expected_routes.push_back((query, result));
115         }
116 }
117
118 impl<'a> Router for TestRouter<'a> {
119         fn find_route(
120                 &self, payer: &PublicKey, params: &RouteParameters, first_hops: Option<&[&channelmanager::ChannelDetails]>,
121                 inflight_htlcs: InFlightHtlcs
122         ) -> Result<Route, msgs::LightningError> {
123                 if let Some((find_route_query, find_route_res)) = self.next_routes.lock().unwrap().pop_front() {
124                         assert_eq!(find_route_query, *params);
125                         if let Ok(ref route) = find_route_res {
126                                 assert_eq!(route.route_params, Some(find_route_query));
127                                 let scorer = self.scorer.read().unwrap();
128                                 let scorer = ScorerAccountingForInFlightHtlcs::new(scorer, &inflight_htlcs);
129                                 for path in &route.paths {
130                                         let mut aggregate_msat = 0u64;
131                                         for (idx, hop) in path.hops.iter().rev().enumerate() {
132                                                 aggregate_msat += hop.fee_msat;
133                                                 let usage = ChannelUsage {
134                                                         amount_msat: aggregate_msat,
135                                                         inflight_htlc_msat: 0,
136                                                         effective_capacity: EffectiveCapacity::Unknown,
137                                                 };
138
139                                                 // Since the path is reversed, the last element in our iteration is the first
140                                                 // hop.
141                                                 if idx == path.hops.len() - 1 {
142                                                         scorer.channel_penalty_msat(hop.short_channel_id, &NodeId::from_pubkey(payer), &NodeId::from_pubkey(&hop.pubkey), usage, &Default::default());
143                                                 } else {
144                                                         let curr_hop_path_idx = path.hops.len() - 1 - idx;
145                                                         scorer.channel_penalty_msat(hop.short_channel_id, &NodeId::from_pubkey(&path.hops[curr_hop_path_idx - 1].pubkey), &NodeId::from_pubkey(&hop.pubkey), usage, &Default::default());
146                                                 }
147                                         }
148                                 }
149                         }
150                         return find_route_res;
151                 }
152                 let logger = TestLogger::new();
153                 find_route(
154                         payer, params, &self.network_graph, first_hops, &logger,
155                         &ScorerAccountingForInFlightHtlcs::new(self.scorer.read().unwrap(), &inflight_htlcs), &Default::default(),
156                         &[42; 32]
157                 )
158         }
159 }
160
161 impl<'a> Drop for TestRouter<'a> {
162         fn drop(&mut self) {
163                 #[cfg(feature = "std")] {
164                         if std::thread::panicking() {
165                                 return;
166                         }
167                 }
168                 assert!(self.next_routes.lock().unwrap().is_empty());
169         }
170 }
171
172 pub struct OnlyReadsKeysInterface {}
173
174 impl EntropySource for OnlyReadsKeysInterface {
175         fn get_secure_random_bytes(&self) -> [u8; 32] { [0; 32] }}
176
177 impl SignerProvider for OnlyReadsKeysInterface {
178         type EcdsaSigner = TestChannelSigner;
179         #[cfg(taproot)]
180         type TaprootSigner = TestChannelSigner;
181
182         fn generate_channel_keys_id(&self, _inbound: bool, _channel_value_satoshis: u64, _user_channel_id: u128) -> [u8; 32] { unreachable!(); }
183
184         fn derive_channel_signer(&self, _channel_value_satoshis: u64, _channel_keys_id: [u8; 32]) -> Self::EcdsaSigner { unreachable!(); }
185
186         fn read_chan_signer(&self, mut reader: &[u8]) -> Result<Self::EcdsaSigner, msgs::DecodeError> {
187                 let inner: InMemorySigner = ReadableArgs::read(&mut reader, self)?;
188                 let state = Arc::new(Mutex::new(EnforcementState::new()));
189
190                 Ok(TestChannelSigner::new_with_revoked(
191                         inner,
192                         state,
193                         false
194                 ))
195         }
196
197         fn get_destination_script(&self, _channel_keys_id: [u8; 32]) -> Result<ScriptBuf, ()> { Err(()) }
198         fn get_shutdown_scriptpubkey(&self) -> Result<ShutdownScript, ()> { Err(()) }
199 }
200
201 pub struct TestChainMonitor<'a> {
202         pub added_monitors: Mutex<Vec<(OutPoint, channelmonitor::ChannelMonitor<TestChannelSigner>)>>,
203         pub monitor_updates: Mutex<HashMap<ChannelId, Vec<channelmonitor::ChannelMonitorUpdate>>>,
204         pub latest_monitor_update_id: Mutex<HashMap<ChannelId, (OutPoint, u64, MonitorUpdateId)>>,
205         pub chain_monitor: chainmonitor::ChainMonitor<TestChannelSigner, &'a TestChainSource, &'a chaininterface::BroadcasterInterface, &'a TestFeeEstimator, &'a TestLogger, &'a chainmonitor::Persist<TestChannelSigner>>,
206         pub keys_manager: &'a TestKeysInterface,
207         /// If this is set to Some(), the next update_channel call (not watch_channel) must be a
208         /// ChannelForceClosed event for the given channel_id with should_broadcast set to the given
209         /// boolean.
210         pub expect_channel_force_closed: Mutex<Option<(ChannelId, bool)>>,
211         /// If this is set to Some(), the next round trip serialization check will not hold after an
212         /// update_channel call (not watch_channel) for the given channel_id.
213         pub expect_monitor_round_trip_fail: Mutex<Option<ChannelId>>,
214 }
215 impl<'a> TestChainMonitor<'a> {
216         pub fn new(chain_source: Option<&'a TestChainSource>, broadcaster: &'a chaininterface::BroadcasterInterface, logger: &'a TestLogger, fee_estimator: &'a TestFeeEstimator, persister: &'a chainmonitor::Persist<TestChannelSigner>, keys_manager: &'a TestKeysInterface) -> Self {
217                 Self {
218                         added_monitors: Mutex::new(Vec::new()),
219                         monitor_updates: Mutex::new(HashMap::new()),
220                         latest_monitor_update_id: Mutex::new(HashMap::new()),
221                         chain_monitor: chainmonitor::ChainMonitor::new(chain_source, broadcaster, logger, fee_estimator, persister),
222                         keys_manager,
223                         expect_channel_force_closed: Mutex::new(None),
224                         expect_monitor_round_trip_fail: Mutex::new(None),
225                 }
226         }
227
228         pub fn complete_sole_pending_chan_update(&self, channel_id: &ChannelId) {
229                 let (outpoint, _, latest_update) = self.latest_monitor_update_id.lock().unwrap().get(channel_id).unwrap().clone();
230                 self.chain_monitor.channel_monitor_updated(outpoint, latest_update).unwrap();
231         }
232 }
233 impl<'a> chain::Watch<TestChannelSigner> for TestChainMonitor<'a> {
234         fn watch_channel(&self, funding_txo: OutPoint, monitor: channelmonitor::ChannelMonitor<TestChannelSigner>) -> Result<chain::ChannelMonitorUpdateStatus, ()> {
235                 // At every point where we get a monitor update, we should be able to send a useful monitor
236                 // to a watchtower and disk...
237                 let mut w = TestVecWriter(Vec::new());
238                 monitor.write(&mut w).unwrap();
239                 let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<TestChannelSigner>)>::read(
240                         &mut io::Cursor::new(&w.0), (self.keys_manager, self.keys_manager)).unwrap().1;
241                 assert!(new_monitor == monitor);
242                 self.latest_monitor_update_id.lock().unwrap().insert(funding_txo.to_channel_id(),
243                         (funding_txo, monitor.get_latest_update_id(), MonitorUpdateId::from_new_monitor(&monitor)));
244                 self.added_monitors.lock().unwrap().push((funding_txo, monitor));
245                 self.chain_monitor.watch_channel(funding_txo, new_monitor)
246         }
247
248         fn update_channel(&self, funding_txo: OutPoint, update: &channelmonitor::ChannelMonitorUpdate) -> chain::ChannelMonitorUpdateStatus {
249                 // Every monitor update should survive roundtrip
250                 let mut w = TestVecWriter(Vec::new());
251                 update.write(&mut w).unwrap();
252                 assert!(channelmonitor::ChannelMonitorUpdate::read(
253                                 &mut io::Cursor::new(&w.0)).unwrap() == *update);
254
255                 self.monitor_updates.lock().unwrap().entry(funding_txo.to_channel_id()).or_insert(Vec::new()).push(update.clone());
256
257                 if let Some(exp) = self.expect_channel_force_closed.lock().unwrap().take() {
258                         assert_eq!(funding_txo.to_channel_id(), exp.0);
259                         assert_eq!(update.updates.len(), 1);
260                         if let channelmonitor::ChannelMonitorUpdateStep::ChannelForceClosed { should_broadcast } = update.updates[0] {
261                                 assert_eq!(should_broadcast, exp.1);
262                         } else { panic!(); }
263                 }
264
265                 self.latest_monitor_update_id.lock().unwrap().insert(funding_txo.to_channel_id(),
266                         (funding_txo, update.update_id, MonitorUpdateId::from_monitor_update(update)));
267                 let update_res = self.chain_monitor.update_channel(funding_txo, update);
268                 // At every point where we get a monitor update, we should be able to send a useful monitor
269                 // to a watchtower and disk...
270                 let monitor = self.chain_monitor.get_monitor(funding_txo).unwrap();
271                 w.0.clear();
272                 monitor.write(&mut w).unwrap();
273                 let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<TestChannelSigner>)>::read(
274                         &mut io::Cursor::new(&w.0), (self.keys_manager, self.keys_manager)).unwrap().1;
275                 if let Some(chan_id) = self.expect_monitor_round_trip_fail.lock().unwrap().take() {
276                         assert_eq!(chan_id, funding_txo.to_channel_id());
277                         assert!(new_monitor != *monitor);
278                 } else {
279                         assert!(new_monitor == *monitor);
280                 }
281                 self.added_monitors.lock().unwrap().push((funding_txo, new_monitor));
282                 update_res
283         }
284
285         fn release_pending_monitor_events(&self) -> Vec<(OutPoint, Vec<MonitorEvent>, Option<PublicKey>)> {
286                 return self.chain_monitor.release_pending_monitor_events();
287         }
288 }
289
290 struct JusticeTxData {
291         justice_tx: Transaction,
292         value: u64,
293         commitment_number: u64,
294 }
295
296 pub(crate) struct WatchtowerPersister {
297         persister: TestPersister,
298         /// Upon a new commitment_signed, we'll get a
299         /// ChannelMonitorUpdateStep::LatestCounterpartyCommitmentTxInfo. We'll store the justice tx
300         /// amount, and commitment number so we can build the justice tx after our counterparty
301         /// revokes it.
302         unsigned_justice_tx_data: Mutex<HashMap<OutPoint, VecDeque<JusticeTxData>>>,
303         /// After receiving a revoke_and_ack for a commitment number, we'll form and store the justice
304         /// tx which would be used to provide a watchtower with the data it needs.
305         watchtower_state: Mutex<HashMap<OutPoint, HashMap<Txid, Transaction>>>,
306         destination_script: ScriptBuf,
307 }
308
309 impl WatchtowerPersister {
310         #[cfg(test)]
311         pub(crate) fn new(destination_script: ScriptBuf) -> Self {
312                 WatchtowerPersister {
313                         persister: TestPersister::new(),
314                         unsigned_justice_tx_data: Mutex::new(HashMap::new()),
315                         watchtower_state: Mutex::new(HashMap::new()),
316                         destination_script,
317                 }
318         }
319
320         #[cfg(test)]
321         pub(crate) fn justice_tx(&self, funding_txo: OutPoint, commitment_txid: &Txid)
322         -> Option<Transaction> {
323                 self.watchtower_state.lock().unwrap().get(&funding_txo).unwrap().get(commitment_txid).cloned()
324         }
325
326         fn form_justice_data_from_commitment(&self, counterparty_commitment_tx: &CommitmentTransaction)
327         -> Option<JusticeTxData> {
328                 let trusted_tx = counterparty_commitment_tx.trust();
329                 let output_idx = trusted_tx.revokeable_output_index()?;
330                 let built_tx = trusted_tx.built_transaction();
331                 let value = built_tx.transaction.output[output_idx as usize].value;
332                 let justice_tx = trusted_tx.build_to_local_justice_tx(
333                         FEERATE_FLOOR_SATS_PER_KW as u64, self.destination_script.clone()).ok()?;
334                 let commitment_number = counterparty_commitment_tx.commitment_number();
335                 Some(JusticeTxData { justice_tx, value, commitment_number })
336         }
337 }
338
339 impl<Signer: sign::ecdsa::WriteableEcdsaChannelSigner> chainmonitor::Persist<Signer> for WatchtowerPersister {
340         fn persist_new_channel(&self, funding_txo: OutPoint,
341                 data: &channelmonitor::ChannelMonitor<Signer>, id: MonitorUpdateId
342         ) -> chain::ChannelMonitorUpdateStatus {
343                 let res = self.persister.persist_new_channel(funding_txo, data, id);
344
345                 assert!(self.unsigned_justice_tx_data.lock().unwrap()
346                         .insert(funding_txo, VecDeque::new()).is_none());
347                 assert!(self.watchtower_state.lock().unwrap()
348                         .insert(funding_txo, HashMap::new()).is_none());
349
350                 let initial_counterparty_commitment_tx = data.initial_counterparty_commitment_tx()
351                         .expect("First and only call expects Some");
352                 if let Some(justice_data)
353                         = self.form_justice_data_from_commitment(&initial_counterparty_commitment_tx) {
354                         self.unsigned_justice_tx_data.lock().unwrap()
355                                 .get_mut(&funding_txo).unwrap()
356                                 .push_back(justice_data);
357                 }
358                 res
359         }
360
361         fn update_persisted_channel(
362                 &self, funding_txo: OutPoint, update: Option<&channelmonitor::ChannelMonitorUpdate>,
363                 data: &channelmonitor::ChannelMonitor<Signer>, update_id: MonitorUpdateId
364         ) -> chain::ChannelMonitorUpdateStatus {
365                 let res = self.persister.update_persisted_channel(funding_txo, update, data, update_id);
366
367                 if let Some(update) = update {
368                         let commitment_txs = data.counterparty_commitment_txs_from_update(update);
369                         let justice_datas = commitment_txs.into_iter()
370                                 .filter_map(|commitment_tx| self.form_justice_data_from_commitment(&commitment_tx));
371                         let mut channels_justice_txs = self.unsigned_justice_tx_data.lock().unwrap();
372                         let channel_state = channels_justice_txs.get_mut(&funding_txo).unwrap();
373                         channel_state.extend(justice_datas);
374
375                         while let Some(JusticeTxData { justice_tx, value, commitment_number }) = channel_state.front() {
376                                 let input_idx = 0;
377                                 let commitment_txid = justice_tx.input[input_idx].previous_output.txid;
378                                 match data.sign_to_local_justice_tx(justice_tx.clone(), input_idx, *value, *commitment_number) {
379                                         Ok(signed_justice_tx) => {
380                                                 let dup = self.watchtower_state.lock().unwrap()
381                                                         .get_mut(&funding_txo).unwrap()
382                                                         .insert(commitment_txid, signed_justice_tx);
383                                                 assert!(dup.is_none());
384                                                 channel_state.pop_front();
385                                         },
386                                         Err(_) => break,
387                                 }
388                         }
389                 }
390                 res
391         }
392 }
393
394 pub struct TestPersister {
395         /// The queue of update statuses we'll return. If none are queued, ::Completed will always be
396         /// returned.
397         pub update_rets: Mutex<VecDeque<chain::ChannelMonitorUpdateStatus>>,
398         /// When we get an update_persisted_channel call with no ChannelMonitorUpdate, we insert the
399         /// MonitorUpdateId here.
400         pub chain_sync_monitor_persistences: Mutex<HashMap<OutPoint, HashSet<MonitorUpdateId>>>,
401         /// When we get an update_persisted_channel call *with* a ChannelMonitorUpdate, we insert the
402         /// MonitorUpdateId here.
403         pub offchain_monitor_updates: Mutex<HashMap<OutPoint, HashSet<MonitorUpdateId>>>,
404 }
405 impl TestPersister {
406         pub fn new() -> Self {
407                 Self {
408                         update_rets: Mutex::new(VecDeque::new()),
409                         chain_sync_monitor_persistences: Mutex::new(HashMap::new()),
410                         offchain_monitor_updates: Mutex::new(HashMap::new()),
411                 }
412         }
413
414         /// Queue an update status to return.
415         pub fn set_update_ret(&self, next_ret: chain::ChannelMonitorUpdateStatus) {
416                 self.update_rets.lock().unwrap().push_back(next_ret);
417         }
418 }
419 impl<Signer: sign::ecdsa::WriteableEcdsaChannelSigner> chainmonitor::Persist<Signer> for TestPersister {
420         fn persist_new_channel(&self, _funding_txo: OutPoint, _data: &channelmonitor::ChannelMonitor<Signer>, _id: MonitorUpdateId) -> chain::ChannelMonitorUpdateStatus {
421                 if let Some(update_ret) = self.update_rets.lock().unwrap().pop_front() {
422                         return update_ret
423                 }
424                 chain::ChannelMonitorUpdateStatus::Completed
425         }
426
427         fn update_persisted_channel(&self, funding_txo: OutPoint, _update: Option<&channelmonitor::ChannelMonitorUpdate>, _data: &channelmonitor::ChannelMonitor<Signer>, update_id: MonitorUpdateId) -> chain::ChannelMonitorUpdateStatus {
428                 let mut ret = chain::ChannelMonitorUpdateStatus::Completed;
429                 if let Some(update_ret) = self.update_rets.lock().unwrap().pop_front() {
430                         ret = update_ret;
431                 }
432                 let is_chain_sync = if let UpdateOrigin::ChainSync(_) = update_id.contents { true } else { false };
433                 if is_chain_sync {
434                         self.chain_sync_monitor_persistences.lock().unwrap().entry(funding_txo).or_insert(HashSet::new()).insert(update_id);
435                 } else {
436                         self.offchain_monitor_updates.lock().unwrap().entry(funding_txo).or_insert(HashSet::new()).insert(update_id);
437                 }
438                 ret
439         }
440 }
441
442 pub struct TestStore {
443         persisted_bytes: Mutex<HashMap<String, HashMap<String, Vec<u8>>>>,
444         read_only: bool,
445 }
446
447 impl TestStore {
448         pub fn new(read_only: bool) -> Self {
449                 let persisted_bytes = Mutex::new(HashMap::new());
450                 Self { persisted_bytes, read_only }
451         }
452 }
453
454 impl KVStore for TestStore {
455         fn read(&self, primary_namespace: &str, secondary_namespace: &str, key: &str) -> io::Result<Vec<u8>> {
456                 let persisted_lock = self.persisted_bytes.lock().unwrap();
457                 let prefixed = if secondary_namespace.is_empty() {
458                         primary_namespace.to_string()
459                 } else {
460                         format!("{}/{}", primary_namespace, secondary_namespace)
461                 };
462
463                 if let Some(outer_ref) = persisted_lock.get(&prefixed) {
464                         if let Some(inner_ref) = outer_ref.get(key) {
465                                 let bytes = inner_ref.clone();
466                                 Ok(bytes)
467                         } else {
468                                 Err(io::Error::new(io::ErrorKind::NotFound, "Key not found"))
469                         }
470                 } else {
471                         Err(io::Error::new(io::ErrorKind::NotFound, "Namespace not found"))
472                 }
473         }
474
475         fn write(&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: &[u8]) -> io::Result<()> {
476                 if self.read_only {
477                         return Err(io::Error::new(
478                                 io::ErrorKind::PermissionDenied,
479                                 "Cannot modify read-only store",
480                         ));
481                 }
482                 let mut persisted_lock = self.persisted_bytes.lock().unwrap();
483
484                 let prefixed = if secondary_namespace.is_empty() {
485                         primary_namespace.to_string()
486                 } else {
487                         format!("{}/{}", primary_namespace, secondary_namespace)
488                 };
489                 let outer_e = persisted_lock.entry(prefixed).or_insert(HashMap::new());
490                 let mut bytes = Vec::new();
491                 bytes.write_all(buf)?;
492                 outer_e.insert(key.to_string(), bytes);
493                 Ok(())
494         }
495
496         fn remove(&self, primary_namespace: &str, secondary_namespace: &str, key: &str, _lazy: bool) -> io::Result<()> {
497                 if self.read_only {
498                         return Err(io::Error::new(
499                                 io::ErrorKind::PermissionDenied,
500                                 "Cannot modify read-only store",
501                         ));
502                 }
503
504                 let mut persisted_lock = self.persisted_bytes.lock().unwrap();
505
506                 let prefixed = if secondary_namespace.is_empty() {
507                         primary_namespace.to_string()
508                 } else {
509                         format!("{}/{}", primary_namespace, secondary_namespace)
510                 };
511                 if let Some(outer_ref) = persisted_lock.get_mut(&prefixed) {
512                                 outer_ref.remove(&key.to_string());
513                 }
514
515                 Ok(())
516         }
517
518         fn list(&self, primary_namespace: &str, secondary_namespace: &str) -> io::Result<Vec<String>> {
519                 let mut persisted_lock = self.persisted_bytes.lock().unwrap();
520
521                 let prefixed = if secondary_namespace.is_empty() {
522                         primary_namespace.to_string()
523                 } else {
524                         format!("{}/{}", primary_namespace, secondary_namespace)
525                 };
526                 match persisted_lock.entry(prefixed) {
527                         hash_map::Entry::Occupied(e) => Ok(e.get().keys().cloned().collect()),
528                         hash_map::Entry::Vacant(_) => Ok(Vec::new()),
529                 }
530         }
531 }
532
533 pub struct TestBroadcaster {
534         pub txn_broadcasted: Mutex<Vec<Transaction>>,
535         pub blocks: Arc<Mutex<Vec<(Block, u32)>>>,
536 }
537
538 impl TestBroadcaster {
539         pub fn new(network: Network) -> Self {
540                 Self {
541                         txn_broadcasted: Mutex::new(Vec::new()),
542                         blocks: Arc::new(Mutex::new(vec![(genesis_block(network), 0)])),
543                 }
544         }
545
546         pub fn with_blocks(blocks: Arc<Mutex<Vec<(Block, u32)>>>) -> Self {
547                 Self { txn_broadcasted: Mutex::new(Vec::new()), blocks }
548         }
549
550         pub fn txn_broadcast(&self) -> Vec<Transaction> {
551                 self.txn_broadcasted.lock().unwrap().split_off(0)
552         }
553
554         pub fn unique_txn_broadcast(&self) -> Vec<Transaction> {
555                 let mut txn = self.txn_broadcasted.lock().unwrap().split_off(0);
556                 let mut seen = HashSet::new();
557                 txn.retain(|tx| seen.insert(tx.txid()));
558                 txn
559         }
560 }
561
562 impl chaininterface::BroadcasterInterface for TestBroadcaster {
563         fn broadcast_transactions(&self, txs: &[&Transaction]) {
564                 for tx in txs {
565                         let lock_time = tx.lock_time.to_consensus_u32();
566                         assert!(lock_time < 1_500_000_000);
567                         if tx.lock_time.is_block_height() && lock_time > self.blocks.lock().unwrap().last().unwrap().1 {
568                                 for inp in tx.input.iter() {
569                                         if inp.sequence != Sequence::MAX {
570                                                 panic!("We should never broadcast a transaction before its locktime ({})!", tx.lock_time);
571                                         }
572                                 }
573                         }
574                 }
575                 let owned_txs: Vec<Transaction> = txs.iter().map(|tx| (*tx).clone()).collect();
576                 self.txn_broadcasted.lock().unwrap().extend(owned_txs);
577         }
578 }
579
580 pub struct TestChannelMessageHandler {
581         pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
582         expected_recv_msgs: Mutex<Option<Vec<wire::Message<()>>>>,
583         connected_peers: Mutex<HashSet<PublicKey>>,
584         pub message_fetch_counter: AtomicUsize,
585         chain_hash: ChainHash,
586 }
587
588 impl TestChannelMessageHandler {
589         pub fn new(chain_hash: ChainHash) -> Self {
590                 TestChannelMessageHandler {
591                         pending_events: Mutex::new(Vec::new()),
592                         expected_recv_msgs: Mutex::new(None),
593                         connected_peers: Mutex::new(HashSet::new()),
594                         message_fetch_counter: AtomicUsize::new(0),
595                         chain_hash,
596                 }
597         }
598
599         #[cfg(test)]
600         pub(crate) fn expect_receive_msg(&self, ev: wire::Message<()>) {
601                 let mut expected_msgs = self.expected_recv_msgs.lock().unwrap();
602                 if expected_msgs.is_none() { *expected_msgs = Some(Vec::new()); }
603                 expected_msgs.as_mut().unwrap().push(ev);
604         }
605
606         fn received_msg(&self, _ev: wire::Message<()>) {
607                 let mut msgs = self.expected_recv_msgs.lock().unwrap();
608                 if msgs.is_none() { return; }
609                 assert!(!msgs.as_ref().unwrap().is_empty(), "Received message when we weren't expecting one");
610                 #[cfg(test)]
611                 assert_eq!(msgs.as_ref().unwrap()[0], _ev);
612                 msgs.as_mut().unwrap().remove(0);
613         }
614 }
615
616 impl Drop for TestChannelMessageHandler {
617         fn drop(&mut self) {
618                 #[cfg(feature = "std")]
619                 {
620                         let l = self.expected_recv_msgs.lock().unwrap();
621                         if !std::thread::panicking() {
622                                 assert!(l.is_none() || l.as_ref().unwrap().is_empty());
623                         }
624                 }
625         }
626 }
627
628 impl msgs::ChannelMessageHandler for TestChannelMessageHandler {
629         fn handle_open_channel(&self, _their_node_id: &PublicKey, msg: &msgs::OpenChannel) {
630                 self.received_msg(wire::Message::OpenChannel(msg.clone()));
631         }
632         fn handle_accept_channel(&self, _their_node_id: &PublicKey, msg: &msgs::AcceptChannel) {
633                 self.received_msg(wire::Message::AcceptChannel(msg.clone()));
634         }
635         fn handle_funding_created(&self, _their_node_id: &PublicKey, msg: &msgs::FundingCreated) {
636                 self.received_msg(wire::Message::FundingCreated(msg.clone()));
637         }
638         fn handle_funding_signed(&self, _their_node_id: &PublicKey, msg: &msgs::FundingSigned) {
639                 self.received_msg(wire::Message::FundingSigned(msg.clone()));
640         }
641         fn handle_channel_ready(&self, _their_node_id: &PublicKey, msg: &msgs::ChannelReady) {
642                 self.received_msg(wire::Message::ChannelReady(msg.clone()));
643         }
644         fn handle_shutdown(&self, _their_node_id: &PublicKey, msg: &msgs::Shutdown) {
645                 self.received_msg(wire::Message::Shutdown(msg.clone()));
646         }
647         fn handle_closing_signed(&self, _their_node_id: &PublicKey, msg: &msgs::ClosingSigned) {
648                 self.received_msg(wire::Message::ClosingSigned(msg.clone()));
649         }
650         fn handle_stfu(&self, _their_node_id: &PublicKey, msg: &msgs::Stfu) {
651                 self.received_msg(wire::Message::Stfu(msg.clone()));
652         }
653         fn handle_splice(&self, _their_node_id: &PublicKey, msg: &msgs::Splice) {
654                 self.received_msg(wire::Message::Splice(msg.clone()));
655         }
656         fn handle_splice_ack(&self, _their_node_id: &PublicKey, msg: &msgs::SpliceAck) {
657                 self.received_msg(wire::Message::SpliceAck(msg.clone()));
658         }
659         fn handle_splice_locked(&self, _their_node_id: &PublicKey, msg: &msgs::SpliceLocked) {
660                 self.received_msg(wire::Message::SpliceLocked(msg.clone()));
661         }
662         fn handle_update_add_htlc(&self, _their_node_id: &PublicKey, msg: &msgs::UpdateAddHTLC) {
663                 self.received_msg(wire::Message::UpdateAddHTLC(msg.clone()));
664         }
665         fn handle_update_fulfill_htlc(&self, _their_node_id: &PublicKey, msg: &msgs::UpdateFulfillHTLC) {
666                 self.received_msg(wire::Message::UpdateFulfillHTLC(msg.clone()));
667         }
668         fn handle_update_fail_htlc(&self, _their_node_id: &PublicKey, msg: &msgs::UpdateFailHTLC) {
669                 self.received_msg(wire::Message::UpdateFailHTLC(msg.clone()));
670         }
671         fn handle_update_fail_malformed_htlc(&self, _their_node_id: &PublicKey, msg: &msgs::UpdateFailMalformedHTLC) {
672                 self.received_msg(wire::Message::UpdateFailMalformedHTLC(msg.clone()));
673         }
674         fn handle_commitment_signed(&self, _their_node_id: &PublicKey, msg: &msgs::CommitmentSigned) {
675                 self.received_msg(wire::Message::CommitmentSigned(msg.clone()));
676         }
677         fn handle_revoke_and_ack(&self, _their_node_id: &PublicKey, msg: &msgs::RevokeAndACK) {
678                 self.received_msg(wire::Message::RevokeAndACK(msg.clone()));
679         }
680         fn handle_update_fee(&self, _their_node_id: &PublicKey, msg: &msgs::UpdateFee) {
681                 self.received_msg(wire::Message::UpdateFee(msg.clone()));
682         }
683         fn handle_channel_update(&self, _their_node_id: &PublicKey, _msg: &msgs::ChannelUpdate) {
684                 // Don't call `received_msg` here as `TestRoutingMessageHandler` generates these sometimes
685         }
686         fn handle_announcement_signatures(&self, _their_node_id: &PublicKey, msg: &msgs::AnnouncementSignatures) {
687                 self.received_msg(wire::Message::AnnouncementSignatures(msg.clone()));
688         }
689         fn handle_channel_reestablish(&self, _their_node_id: &PublicKey, msg: &msgs::ChannelReestablish) {
690                 self.received_msg(wire::Message::ChannelReestablish(msg.clone()));
691         }
692         fn peer_disconnected(&self, their_node_id: &PublicKey) {
693                 assert!(self.connected_peers.lock().unwrap().remove(their_node_id));
694         }
695         fn peer_connected(&self, their_node_id: &PublicKey, _msg: &msgs::Init, _inbound: bool) -> Result<(), ()> {
696                 assert!(self.connected_peers.lock().unwrap().insert(their_node_id.clone()));
697                 // Don't bother with `received_msg` for Init as its auto-generated and we don't want to
698                 // bother re-generating the expected Init message in all tests.
699                 Ok(())
700         }
701         fn handle_error(&self, _their_node_id: &PublicKey, msg: &msgs::ErrorMessage) {
702                 self.received_msg(wire::Message::Error(msg.clone()));
703         }
704         fn provided_node_features(&self) -> NodeFeatures {
705                 channelmanager::provided_node_features(&UserConfig::default())
706         }
707         fn provided_init_features(&self, _their_init_features: &PublicKey) -> InitFeatures {
708                 channelmanager::provided_init_features(&UserConfig::default())
709         }
710
711         fn get_chain_hashes(&self) -> Option<Vec<ChainHash>> {
712                 Some(vec![self.chain_hash])
713         }
714
715         fn handle_open_channel_v2(&self, _their_node_id: &PublicKey, msg: &msgs::OpenChannelV2) {
716                 self.received_msg(wire::Message::OpenChannelV2(msg.clone()));
717         }
718
719         fn handle_accept_channel_v2(&self, _their_node_id: &PublicKey, msg: &msgs::AcceptChannelV2) {
720                 self.received_msg(wire::Message::AcceptChannelV2(msg.clone()));
721         }
722
723         fn handle_tx_add_input(&self, _their_node_id: &PublicKey, msg: &msgs::TxAddInput) {
724                 self.received_msg(wire::Message::TxAddInput(msg.clone()));
725         }
726
727         fn handle_tx_add_output(&self, _their_node_id: &PublicKey, msg: &msgs::TxAddOutput) {
728                 self.received_msg(wire::Message::TxAddOutput(msg.clone()));
729         }
730
731         fn handle_tx_remove_input(&self, _their_node_id: &PublicKey, msg: &msgs::TxRemoveInput) {
732                 self.received_msg(wire::Message::TxRemoveInput(msg.clone()));
733         }
734
735         fn handle_tx_remove_output(&self, _their_node_id: &PublicKey, msg: &msgs::TxRemoveOutput) {
736                 self.received_msg(wire::Message::TxRemoveOutput(msg.clone()));
737         }
738
739         fn handle_tx_complete(&self, _their_node_id: &PublicKey, msg: &msgs::TxComplete) {
740                 self.received_msg(wire::Message::TxComplete(msg.clone()));
741         }
742
743         fn handle_tx_signatures(&self, _their_node_id: &PublicKey, msg: &msgs::TxSignatures) {
744                 self.received_msg(wire::Message::TxSignatures(msg.clone()));
745         }
746
747         fn handle_tx_init_rbf(&self, _their_node_id: &PublicKey, msg: &msgs::TxInitRbf) {
748                 self.received_msg(wire::Message::TxInitRbf(msg.clone()));
749         }
750
751         fn handle_tx_ack_rbf(&self, _their_node_id: &PublicKey, msg: &msgs::TxAckRbf) {
752                 self.received_msg(wire::Message::TxAckRbf(msg.clone()));
753         }
754
755         fn handle_tx_abort(&self, _their_node_id: &PublicKey, msg: &msgs::TxAbort) {
756                 self.received_msg(wire::Message::TxAbort(msg.clone()));
757         }
758 }
759
760 impl events::MessageSendEventsProvider for TestChannelMessageHandler {
761         fn get_and_clear_pending_msg_events(&self) -> Vec<events::MessageSendEvent> {
762                 self.message_fetch_counter.fetch_add(1, Ordering::AcqRel);
763                 let mut pending_events = self.pending_events.lock().unwrap();
764                 let mut ret = Vec::new();
765                 mem::swap(&mut ret, &mut *pending_events);
766                 ret
767         }
768 }
769
770 fn get_dummy_channel_announcement(short_chan_id: u64) -> msgs::ChannelAnnouncement {
771         use bitcoin::secp256k1::ffi::Signature as FFISignature;
772         let secp_ctx = Secp256k1::new();
773         let network = Network::Testnet;
774         let node_1_privkey = SecretKey::from_slice(&[42; 32]).unwrap();
775         let node_2_privkey = SecretKey::from_slice(&[41; 32]).unwrap();
776         let node_1_btckey = SecretKey::from_slice(&[40; 32]).unwrap();
777         let node_2_btckey = SecretKey::from_slice(&[39; 32]).unwrap();
778         let unsigned_ann = msgs::UnsignedChannelAnnouncement {
779                 features: ChannelFeatures::empty(),
780                 chain_hash: ChainHash::using_genesis_block(network),
781                 short_channel_id: short_chan_id,
782                 node_id_1: NodeId::from_pubkey(&PublicKey::from_secret_key(&secp_ctx, &node_1_privkey)),
783                 node_id_2: NodeId::from_pubkey(&PublicKey::from_secret_key(&secp_ctx, &node_2_privkey)),
784                 bitcoin_key_1: NodeId::from_pubkey(&PublicKey::from_secret_key(&secp_ctx, &node_1_btckey)),
785                 bitcoin_key_2: NodeId::from_pubkey(&PublicKey::from_secret_key(&secp_ctx, &node_2_btckey)),
786                 excess_data: Vec::new(),
787         };
788
789         unsafe {
790                 msgs::ChannelAnnouncement {
791                         node_signature_1: Signature::from(FFISignature::new()),
792                         node_signature_2: Signature::from(FFISignature::new()),
793                         bitcoin_signature_1: Signature::from(FFISignature::new()),
794                         bitcoin_signature_2: Signature::from(FFISignature::new()),
795                         contents: unsigned_ann,
796                 }
797         }
798 }
799
800 fn get_dummy_channel_update(short_chan_id: u64) -> msgs::ChannelUpdate {
801         use bitcoin::secp256k1::ffi::Signature as FFISignature;
802         let network = Network::Testnet;
803         msgs::ChannelUpdate {
804                 signature: Signature::from(unsafe { FFISignature::new() }),
805                 contents: msgs::UnsignedChannelUpdate {
806                         chain_hash: ChainHash::using_genesis_block(network),
807                         short_channel_id: short_chan_id,
808                         timestamp: 0,
809                         flags: 0,
810                         cltv_expiry_delta: 0,
811                         htlc_minimum_msat: 0,
812                         htlc_maximum_msat: msgs::MAX_VALUE_MSAT,
813                         fee_base_msat: 0,
814                         fee_proportional_millionths: 0,
815                         excess_data: vec![],
816                 }
817         }
818 }
819
820 pub struct TestRoutingMessageHandler {
821         pub chan_upds_recvd: AtomicUsize,
822         pub chan_anns_recvd: AtomicUsize,
823         pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
824         pub request_full_sync: AtomicBool,
825 }
826
827 impl TestRoutingMessageHandler {
828         pub fn new() -> Self {
829                 TestRoutingMessageHandler {
830                         chan_upds_recvd: AtomicUsize::new(0),
831                         chan_anns_recvd: AtomicUsize::new(0),
832                         pending_events: Mutex::new(vec![]),
833                         request_full_sync: AtomicBool::new(false),
834                 }
835         }
836 }
837 impl msgs::RoutingMessageHandler for TestRoutingMessageHandler {
838         fn handle_node_announcement(&self, _msg: &msgs::NodeAnnouncement) -> Result<bool, msgs::LightningError> {
839                 Err(msgs::LightningError { err: "".to_owned(), action: msgs::ErrorAction::IgnoreError })
840         }
841         fn handle_channel_announcement(&self, _msg: &msgs::ChannelAnnouncement) -> Result<bool, msgs::LightningError> {
842                 self.chan_anns_recvd.fetch_add(1, Ordering::AcqRel);
843                 Err(msgs::LightningError { err: "".to_owned(), action: msgs::ErrorAction::IgnoreError })
844         }
845         fn handle_channel_update(&self, _msg: &msgs::ChannelUpdate) -> Result<bool, msgs::LightningError> {
846                 self.chan_upds_recvd.fetch_add(1, Ordering::AcqRel);
847                 Err(msgs::LightningError { err: "".to_owned(), action: msgs::ErrorAction::IgnoreError })
848         }
849         fn get_next_channel_announcement(&self, starting_point: u64) -> Option<(msgs::ChannelAnnouncement, Option<msgs::ChannelUpdate>, Option<msgs::ChannelUpdate>)> {
850                 let chan_upd_1 = get_dummy_channel_update(starting_point);
851                 let chan_upd_2 = get_dummy_channel_update(starting_point);
852                 let chan_ann = get_dummy_channel_announcement(starting_point);
853
854                 Some((chan_ann, Some(chan_upd_1), Some(chan_upd_2)))
855         }
856
857         fn get_next_node_announcement(&self, _starting_point: Option<&NodeId>) -> Option<msgs::NodeAnnouncement> {
858                 None
859         }
860
861         fn peer_connected(&self, their_node_id: &PublicKey, init_msg: &msgs::Init, _inbound: bool) -> Result<(), ()> {
862                 if !init_msg.features.supports_gossip_queries() {
863                         return Ok(());
864                 }
865
866                 #[allow(unused_mut, unused_assignments)]
867                 let mut gossip_start_time = 0;
868                 #[cfg(feature = "std")]
869                 {
870                         gossip_start_time = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time must be > 1970").as_secs();
871                         if self.request_full_sync.load(Ordering::Acquire) {
872                                 gossip_start_time -= 60 * 60 * 24 * 7 * 2; // 2 weeks ago
873                         } else {
874                                 gossip_start_time -= 60 * 60; // an hour ago
875                         }
876                 }
877
878                 let mut pending_events = self.pending_events.lock().unwrap();
879                 pending_events.push(events::MessageSendEvent::SendGossipTimestampFilter {
880                         node_id: their_node_id.clone(),
881                         msg: msgs::GossipTimestampFilter {
882                                 chain_hash: ChainHash::using_genesis_block(Network::Testnet),
883                                 first_timestamp: gossip_start_time as u32,
884                                 timestamp_range: u32::max_value(),
885                         },
886                 });
887                 Ok(())
888         }
889
890         fn handle_reply_channel_range(&self, _their_node_id: &PublicKey, _msg: msgs::ReplyChannelRange) -> Result<(), msgs::LightningError> {
891                 Ok(())
892         }
893
894         fn handle_reply_short_channel_ids_end(&self, _their_node_id: &PublicKey, _msg: msgs::ReplyShortChannelIdsEnd) -> Result<(), msgs::LightningError> {
895                 Ok(())
896         }
897
898         fn handle_query_channel_range(&self, _their_node_id: &PublicKey, _msg: msgs::QueryChannelRange) -> Result<(), msgs::LightningError> {
899                 Ok(())
900         }
901
902         fn handle_query_short_channel_ids(&self, _their_node_id: &PublicKey, _msg: msgs::QueryShortChannelIds) -> Result<(), msgs::LightningError> {
903                 Ok(())
904         }
905
906         fn provided_node_features(&self) -> NodeFeatures {
907                 let mut features = NodeFeatures::empty();
908                 features.set_gossip_queries_optional();
909                 features
910         }
911
912         fn provided_init_features(&self, _their_init_features: &PublicKey) -> InitFeatures {
913                 let mut features = InitFeatures::empty();
914                 features.set_gossip_queries_optional();
915                 features
916         }
917
918         fn processing_queue_high(&self) -> bool { false }
919 }
920
921 impl events::MessageSendEventsProvider for TestRoutingMessageHandler {
922         fn get_and_clear_pending_msg_events(&self) -> Vec<events::MessageSendEvent> {
923                 let mut ret = Vec::new();
924                 let mut pending_events = self.pending_events.lock().unwrap();
925                 core::mem::swap(&mut ret, &mut pending_events);
926                 ret
927         }
928 }
929
930 pub struct TestLogger {
931         level: Level,
932         pub(crate) id: String,
933         pub lines: Mutex<HashMap<(&'static str, String), usize>>,
934         pub context: Mutex<HashMap<(&'static str, Option<PublicKey>, Option<ChannelId>), usize>>,
935 }
936
937 impl TestLogger {
938         pub fn new() -> TestLogger {
939                 Self::with_id("".to_owned())
940         }
941         pub fn with_id(id: String) -> TestLogger {
942                 TestLogger {
943                         level: Level::Trace,
944                         id,
945                         lines: Mutex::new(HashMap::new()),
946                         context: Mutex::new(HashMap::new()),
947                 }
948         }
949         pub fn enable(&mut self, level: Level) {
950                 self.level = level;
951         }
952         pub fn assert_log(&self, module: &str, line: String, count: usize) {
953                 let log_entries = self.lines.lock().unwrap();
954                 assert_eq!(log_entries.get(&(module, line)), Some(&count));
955         }
956
957         /// Search for the number of occurrence of the logged lines which
958         /// 1. belongs to the specified module and
959         /// 2. contains `line` in it.
960         /// And asserts if the number of occurrences is the same with the given `count`
961         pub fn assert_log_contains(&self, module: &str, line: &str, count: usize) {
962                 let log_entries = self.lines.lock().unwrap();
963                 let l: usize = log_entries.iter().filter(|&(&(ref m, ref l), _c)| {
964                         *m == module && l.contains(line)
965                 }).map(|(_, c) | { c }).sum();
966                 assert_eq!(l, count)
967         }
968
969         /// Search for the number of occurrences of logged lines which
970         /// 1. belong to the specified module and
971         /// 2. match the given regex pattern.
972         /// Assert that the number of occurrences equals the given `count`
973         #[cfg(any(test, feature = "_test_utils"))]
974         pub fn assert_log_regex(&self, module: &str, pattern: regex::Regex, count: usize) {
975                 let log_entries = self.lines.lock().unwrap();
976                 let l: usize = log_entries.iter().filter(|&(&(ref m, ref l), _c)| {
977                         *m == module && pattern.is_match(&l)
978                 }).map(|(_, c) | { c }).sum();
979                 assert_eq!(l, count)
980         }
981
982         pub fn assert_log_context_contains(
983                 &self, module: &str, peer_id: Option<PublicKey>, channel_id: Option<ChannelId>, count: usize
984         ) {
985                 let context_entries = self.context.lock().unwrap();
986                 let l = context_entries.get(&(module, peer_id, channel_id)).unwrap();
987                 assert_eq!(*l, count)
988         }
989 }
990
991 impl Logger for TestLogger {
992         fn log(&self, record: Record) {
993                 *self.lines.lock().unwrap().entry((record.module_path, format!("{}", record.args))).or_insert(0) += 1;
994                 *self.context.lock().unwrap().entry((record.module_path, record.peer_id, record.channel_id)).or_insert(0) += 1;
995                 if record.level >= self.level {
996                         #[cfg(all(not(ldk_bench), feature = "std"))] {
997                                 let pfx = format!("{} {} [{}:{}]", self.id, record.level.to_string(), record.module_path, record.line);
998                                 println!("{:<55}{}", pfx, record.args);
999                         }
1000                 }
1001         }
1002 }
1003
1004 pub struct TestNodeSigner {
1005         node_secret: SecretKey,
1006 }
1007
1008 impl TestNodeSigner {
1009         pub fn new(node_secret: SecretKey) -> Self {
1010                 Self { node_secret }
1011         }
1012 }
1013
1014 impl NodeSigner for TestNodeSigner {
1015         fn get_inbound_payment_key_material(&self) -> crate::sign::KeyMaterial {
1016                 unreachable!()
1017         }
1018
1019         fn get_node_id(&self, recipient: Recipient) -> Result<PublicKey, ()> {
1020                 let node_secret = match recipient {
1021                         Recipient::Node => Ok(&self.node_secret),
1022                         Recipient::PhantomNode => Err(())
1023                 }?;
1024                 Ok(PublicKey::from_secret_key(&Secp256k1::signing_only(), node_secret))
1025         }
1026
1027         fn ecdh(&self, recipient: Recipient, other_key: &PublicKey, tweak: Option<&bitcoin::secp256k1::Scalar>) -> Result<SharedSecret, ()> {
1028                 let mut node_secret = match recipient {
1029                         Recipient::Node => Ok(self.node_secret.clone()),
1030                         Recipient::PhantomNode => Err(())
1031                 }?;
1032                 if let Some(tweak) = tweak {
1033                         node_secret = node_secret.mul_tweak(tweak).map_err(|_| ())?;
1034                 }
1035                 Ok(SharedSecret::new(other_key, &node_secret))
1036         }
1037
1038         fn sign_invoice(&self, _: &[u8], _: &[bitcoin::bech32::u5], _: Recipient) -> Result<bitcoin::secp256k1::ecdsa::RecoverableSignature, ()> {
1039                 unreachable!()
1040         }
1041
1042         fn sign_bolt12_invoice_request(
1043                 &self, _invoice_request: &UnsignedInvoiceRequest
1044         ) -> Result<schnorr::Signature, ()> {
1045                 unreachable!()
1046         }
1047
1048         fn sign_bolt12_invoice(
1049                 &self, _invoice: &UnsignedBolt12Invoice,
1050         ) -> Result<schnorr::Signature, ()> {
1051                 unreachable!()
1052         }
1053
1054         fn sign_gossip_message(&self, _msg: msgs::UnsignedGossipMessage) -> Result<Signature, ()> {
1055                 unreachable!()
1056         }
1057 }
1058
1059 pub struct TestKeysInterface {
1060         pub backing: sign::PhantomKeysManager,
1061         pub override_random_bytes: Mutex<Option<[u8; 32]>>,
1062         pub disable_revocation_policy_check: bool,
1063         enforcement_states: Mutex<HashMap<[u8;32], Arc<Mutex<EnforcementState>>>>,
1064         expectations: Mutex<Option<VecDeque<OnGetShutdownScriptpubkey>>>,
1065 }
1066
1067 impl EntropySource for TestKeysInterface {
1068         fn get_secure_random_bytes(&self) -> [u8; 32] {
1069                 let override_random_bytes = self.override_random_bytes.lock().unwrap();
1070                 if let Some(bytes) = &*override_random_bytes {
1071                         return *bytes;
1072                 }
1073                 self.backing.get_secure_random_bytes()
1074         }
1075 }
1076
1077 impl NodeSigner for TestKeysInterface {
1078         fn get_node_id(&self, recipient: Recipient) -> Result<PublicKey, ()> {
1079                 self.backing.get_node_id(recipient)
1080         }
1081
1082         fn ecdh(&self, recipient: Recipient, other_key: &PublicKey, tweak: Option<&Scalar>) -> Result<SharedSecret, ()> {
1083                 self.backing.ecdh(recipient, other_key, tweak)
1084         }
1085
1086         fn get_inbound_payment_key_material(&self) -> sign::KeyMaterial {
1087                 self.backing.get_inbound_payment_key_material()
1088         }
1089
1090         fn sign_invoice(&self, hrp_bytes: &[u8], invoice_data: &[u5], recipient: Recipient) -> Result<RecoverableSignature, ()> {
1091                 self.backing.sign_invoice(hrp_bytes, invoice_data, recipient)
1092         }
1093
1094         fn sign_bolt12_invoice_request(
1095                 &self, invoice_request: &UnsignedInvoiceRequest
1096         ) -> Result<schnorr::Signature, ()> {
1097                 self.backing.sign_bolt12_invoice_request(invoice_request)
1098         }
1099
1100         fn sign_bolt12_invoice(
1101                 &self, invoice: &UnsignedBolt12Invoice,
1102         ) -> Result<schnorr::Signature, ()> {
1103                 self.backing.sign_bolt12_invoice(invoice)
1104         }
1105
1106         fn sign_gossip_message(&self, msg: msgs::UnsignedGossipMessage) -> Result<Signature, ()> {
1107                 self.backing.sign_gossip_message(msg)
1108         }
1109 }
1110
1111 impl SignerProvider for TestKeysInterface {
1112         type EcdsaSigner = TestChannelSigner;
1113         #[cfg(taproot)]
1114         type TaprootSigner = TestChannelSigner;
1115
1116         fn generate_channel_keys_id(&self, inbound: bool, channel_value_satoshis: u64, user_channel_id: u128) -> [u8; 32] {
1117                 self.backing.generate_channel_keys_id(inbound, channel_value_satoshis, user_channel_id)
1118         }
1119
1120         fn derive_channel_signer(&self, channel_value_satoshis: u64, channel_keys_id: [u8; 32]) -> TestChannelSigner {
1121                 let keys = self.backing.derive_channel_signer(channel_value_satoshis, channel_keys_id);
1122                 let state = self.make_enforcement_state_cell(keys.commitment_seed);
1123                 TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
1124         }
1125
1126         fn read_chan_signer(&self, buffer: &[u8]) -> Result<Self::EcdsaSigner, msgs::DecodeError> {
1127                 let mut reader = io::Cursor::new(buffer);
1128
1129                 let inner: InMemorySigner = ReadableArgs::read(&mut reader, self)?;
1130                 let state = self.make_enforcement_state_cell(inner.commitment_seed);
1131
1132                 Ok(TestChannelSigner::new_with_revoked(
1133                         inner,
1134                         state,
1135                         self.disable_revocation_policy_check
1136                 ))
1137         }
1138
1139         fn get_destination_script(&self, channel_keys_id: [u8; 32]) -> Result<ScriptBuf, ()> { self.backing.get_destination_script(channel_keys_id) }
1140
1141         fn get_shutdown_scriptpubkey(&self) -> Result<ShutdownScript, ()> {
1142                 match &mut *self.expectations.lock().unwrap() {
1143                         None => self.backing.get_shutdown_scriptpubkey(),
1144                         Some(expectations) => match expectations.pop_front() {
1145                                 None => panic!("Unexpected get_shutdown_scriptpubkey"),
1146                                 Some(expectation) => Ok(expectation.returns),
1147                         },
1148                 }
1149         }
1150 }
1151
1152 impl TestKeysInterface {
1153         pub fn new(seed: &[u8; 32], network: Network) -> Self {
1154                 let now = Duration::from_secs(genesis_block(network).header.time as u64);
1155                 Self {
1156                         backing: sign::PhantomKeysManager::new(seed, now.as_secs(), now.subsec_nanos(), seed),
1157                         override_random_bytes: Mutex::new(None),
1158                         disable_revocation_policy_check: false,
1159                         enforcement_states: Mutex::new(HashMap::new()),
1160                         expectations: Mutex::new(None),
1161                 }
1162         }
1163
1164         /// Sets an expectation that [`sign::SignerProvider::get_shutdown_scriptpubkey`] is
1165         /// called.
1166         pub fn expect(&self, expectation: OnGetShutdownScriptpubkey) -> &Self {
1167                 self.expectations.lock().unwrap()
1168                         .get_or_insert_with(|| VecDeque::new())
1169                         .push_back(expectation);
1170                 self
1171         }
1172
1173         pub fn derive_channel_keys(&self, channel_value_satoshis: u64, id: &[u8; 32]) -> TestChannelSigner {
1174                 let keys = self.backing.derive_channel_keys(channel_value_satoshis, id);
1175                 let state = self.make_enforcement_state_cell(keys.commitment_seed);
1176                 TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
1177         }
1178
1179         fn make_enforcement_state_cell(&self, commitment_seed: [u8; 32]) -> Arc<Mutex<EnforcementState>> {
1180                 let mut states = self.enforcement_states.lock().unwrap();
1181                 if !states.contains_key(&commitment_seed) {
1182                         let state = EnforcementState::new();
1183                         states.insert(commitment_seed, Arc::new(Mutex::new(state)));
1184                 }
1185                 let cell = states.get(&commitment_seed).unwrap();
1186                 Arc::clone(cell)
1187         }
1188 }
1189
1190 pub(crate) fn panicking() -> bool {
1191         #[cfg(feature = "std")]
1192         let panicking = ::std::thread::panicking();
1193         #[cfg(not(feature = "std"))]
1194         let panicking = false;
1195         return panicking;
1196 }
1197
1198 impl Drop for TestKeysInterface {
1199         fn drop(&mut self) {
1200                 if panicking() {
1201                         return;
1202                 }
1203
1204                 if let Some(expectations) = &*self.expectations.lock().unwrap() {
1205                         if !expectations.is_empty() {
1206                                 panic!("Unsatisfied expectations: {:?}", expectations);
1207                         }
1208                 }
1209         }
1210 }
1211
1212 /// An expectation that [`sign::SignerProvider::get_shutdown_scriptpubkey`] was called and
1213 /// returns a [`ShutdownScript`].
1214 pub struct OnGetShutdownScriptpubkey {
1215         /// A shutdown script used to close a channel.
1216         pub returns: ShutdownScript,
1217 }
1218
1219 impl core::fmt::Debug for OnGetShutdownScriptpubkey {
1220         fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1221                 f.debug_struct("OnGetShutdownScriptpubkey").finish()
1222         }
1223 }
1224
1225 pub struct TestChainSource {
1226         pub chain_hash: ChainHash,
1227         pub utxo_ret: Mutex<UtxoResult>,
1228         pub get_utxo_call_count: AtomicUsize,
1229         pub watched_txn: Mutex<HashSet<(Txid, ScriptBuf)>>,
1230         pub watched_outputs: Mutex<HashSet<(OutPoint, ScriptBuf)>>,
1231 }
1232
1233 impl TestChainSource {
1234         pub fn new(network: Network) -> Self {
1235                 let script_pubkey = Builder::new().push_opcode(opcodes::OP_TRUE).into_script();
1236                 Self {
1237                         chain_hash: ChainHash::using_genesis_block(network),
1238                         utxo_ret: Mutex::new(UtxoResult::Sync(Ok(TxOut { value: u64::max_value(), script_pubkey }))),
1239                         get_utxo_call_count: AtomicUsize::new(0),
1240                         watched_txn: Mutex::new(HashSet::new()),
1241                         watched_outputs: Mutex::new(HashSet::new()),
1242                 }
1243         }
1244 }
1245
1246 impl UtxoLookup for TestChainSource {
1247         fn get_utxo(&self, chain_hash: &ChainHash, _short_channel_id: u64) -> UtxoResult {
1248                 self.get_utxo_call_count.fetch_add(1, Ordering::Relaxed);
1249                 if self.chain_hash != *chain_hash {
1250                         return UtxoResult::Sync(Err(UtxoLookupError::UnknownChain));
1251                 }
1252
1253                 self.utxo_ret.lock().unwrap().clone()
1254         }
1255 }
1256
1257 impl chain::Filter for TestChainSource {
1258         fn register_tx(&self, txid: &Txid, script_pubkey: &Script) {
1259                 self.watched_txn.lock().unwrap().insert((*txid, script_pubkey.into()));
1260         }
1261
1262         fn register_output(&self, output: WatchedOutput) {
1263                 self.watched_outputs.lock().unwrap().insert((output.outpoint, output.script_pubkey));
1264         }
1265 }
1266
1267 impl Drop for TestChainSource {
1268         fn drop(&mut self) {
1269                 if panicking() {
1270                         return;
1271                 }
1272         }
1273 }
1274
1275 pub struct TestScorer {
1276         /// Stores a tuple of (scid, ChannelUsage)
1277         scorer_expectations: RefCell<Option<VecDeque<(u64, ChannelUsage)>>>,
1278 }
1279
1280 impl TestScorer {
1281         pub fn new() -> Self {
1282                 Self {
1283                         scorer_expectations: RefCell::new(None),
1284                 }
1285         }
1286
1287         pub fn expect_usage(&self, scid: u64, expectation: ChannelUsage) {
1288                 self.scorer_expectations.borrow_mut().get_or_insert_with(|| VecDeque::new()).push_back((scid, expectation));
1289         }
1290 }
1291
1292 #[cfg(c_bindings)]
1293 impl crate::util::ser::Writeable for TestScorer {
1294         fn write<W: crate::util::ser::Writer>(&self, _: &mut W) -> Result<(), crate::io::Error> { unreachable!(); }
1295 }
1296
1297 impl ScoreLookUp for TestScorer {
1298         type ScoreParams = ();
1299         fn channel_penalty_msat(
1300                 &self, short_channel_id: u64, _source: &NodeId, _target: &NodeId, usage: ChannelUsage, _score_params: &Self::ScoreParams
1301         ) -> u64 {
1302                 if let Some(scorer_expectations) = self.scorer_expectations.borrow_mut().as_mut() {
1303                         match scorer_expectations.pop_front() {
1304                                 Some((scid, expectation)) => {
1305                                         assert_eq!(expectation, usage);
1306                                         assert_eq!(scid, short_channel_id);
1307                                 },
1308                                 None => {},
1309                         }
1310                 }
1311                 0
1312         }
1313 }
1314
1315 impl ScoreUpdate for TestScorer {
1316         fn payment_path_failed(&mut self, _actual_path: &Path, _actual_short_channel_id: u64) {}
1317
1318         fn payment_path_successful(&mut self, _actual_path: &Path) {}
1319
1320         fn probe_failed(&mut self, _actual_path: &Path, _: u64) {}
1321
1322         fn probe_successful(&mut self, _actual_path: &Path) {}
1323 }
1324
1325 impl Drop for TestScorer {
1326         fn drop(&mut self) {
1327                 #[cfg(feature = "std")] {
1328                         if std::thread::panicking() {
1329                                 return;
1330                         }
1331                 }
1332
1333                 if let Some(scorer_expectations) = self.scorer_expectations.borrow().as_ref() {
1334                         if !scorer_expectations.is_empty() {
1335                                 panic!("Unsatisfied scorer expectations: {:?}", scorer_expectations)
1336                         }
1337                 }
1338         }
1339 }
1340
1341 pub struct TestWalletSource {
1342         secret_key: SecretKey,
1343         utxos: RefCell<Vec<Utxo>>,
1344         secp: Secp256k1<bitcoin::secp256k1::All>,
1345 }
1346
1347 impl TestWalletSource {
1348         pub fn new(secret_key: SecretKey) -> Self {
1349                 Self {
1350                         secret_key,
1351                         utxos: RefCell::new(Vec::new()),
1352                         secp: Secp256k1::new(),
1353                 }
1354         }
1355
1356         pub fn add_utxo(&self, outpoint: bitcoin::OutPoint, value: u64) -> TxOut {
1357                 let public_key = bitcoin::PublicKey::new(self.secret_key.public_key(&self.secp));
1358                 let utxo = Utxo::new_p2pkh(outpoint, value, &public_key.pubkey_hash());
1359                 self.utxos.borrow_mut().push(utxo.clone());
1360                 utxo.output
1361         }
1362
1363         pub fn add_custom_utxo(&self, utxo: Utxo) -> TxOut {
1364                 let output = utxo.output.clone();
1365                 self.utxos.borrow_mut().push(utxo);
1366                 output
1367         }
1368
1369         pub fn remove_utxo(&self, outpoint: bitcoin::OutPoint) {
1370                 self.utxos.borrow_mut().retain(|utxo| utxo.outpoint != outpoint);
1371         }
1372 }
1373
1374 impl WalletSource for TestWalletSource {
1375         fn list_confirmed_utxos(&self) -> Result<Vec<Utxo>, ()> {
1376                 Ok(self.utxos.borrow().clone())
1377         }
1378
1379         fn get_change_script(&self) -> Result<ScriptBuf, ()> {
1380                 let public_key = bitcoin::PublicKey::new(self.secret_key.public_key(&self.secp));
1381                 Ok(ScriptBuf::new_p2pkh(&public_key.pubkey_hash()))
1382         }
1383
1384         fn sign_tx(&self, mut tx: Transaction) -> Result<Transaction, ()> {
1385                 let utxos = self.utxos.borrow();
1386                 for i in 0..tx.input.len() {
1387                         if let Some(utxo) = utxos.iter().find(|utxo| utxo.outpoint == tx.input[i].previous_output) {
1388                                 let sighash = SighashCache::new(&tx)
1389                                         .legacy_signature_hash(i, &utxo.output.script_pubkey, EcdsaSighashType::All as u32)
1390                                         .map_err(|_| ())?;
1391                                 let sig = self.secp.sign_ecdsa(&(*sighash.as_raw_hash()).into(), &self.secret_key);
1392                                 let bitcoin_sig = bitcoin::ecdsa::Signature { sig, hash_ty: EcdsaSighashType::All };
1393                                 tx.input[i].script_sig = Builder::new()
1394                                         .push_slice(&bitcoin_sig.serialize())
1395                                         .push_slice(&self.secret_key.public_key(&self.secp).serialize())
1396                                         .into_script();
1397                         }
1398                 }
1399                 Ok(tx)
1400         }
1401 }