Merge pull request #2521 from TheBlueMatt/2023-08-one-less-write
[rust-lightning] / lightning / src / util / test_utils.rs
index 8cd95f62a930243adfb6d330fb28e6536651b8aa..09c99815bb70077ad19355eb7683a50a9f31526e 100644 (file)
@@ -35,9 +35,10 @@ use crate::routing::router::{find_route, InFlightHtlcs, Path, Route, RouteParame
 use crate::routing::scoring::{ChannelUsage, ScoreUpdate, ScoreLookUp};
 use crate::sync::RwLock;
 use crate::util::config::UserConfig;
-use crate::util::enforcing_trait_impls::{EnforcingSigner, EnforcementState};
+use crate::util::test_channel_signer::{TestChannelSigner, EnforcementState};
 use crate::util::logger::{Logger, Level, Record};
 use crate::util::ser::{Readable, ReadableArgs, Writer, Writeable};
+use crate::util::persist::KVStore;
 
 use bitcoin::EcdsaSighashType;
 use bitcoin::blockdata::constants::ChainHash;
@@ -61,7 +62,6 @@ use regex;
 use crate::io;
 use crate::prelude::*;
 use core::cell::RefCell;
-use core::ops::Deref;
 use core::time::Duration;
 use crate::sync::{Mutex, Arc};
 use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
@@ -175,7 +175,7 @@ impl EntropySource for OnlyReadsKeysInterface {
        fn get_secure_random_bytes(&self) -> [u8; 32] { [0; 32] }}
 
 impl SignerProvider for OnlyReadsKeysInterface {
-       type Signer = EnforcingSigner;
+       type Signer = TestChannelSigner;
 
        fn generate_channel_keys_id(&self, _inbound: bool, _channel_value_satoshis: u64, _user_channel_id: u128) -> [u8; 32] { unreachable!(); }
 
@@ -185,7 +185,7 @@ impl SignerProvider for OnlyReadsKeysInterface {
                let inner: InMemorySigner = ReadableArgs::read(&mut reader, self)?;
                let state = Arc::new(Mutex::new(EnforcementState::new()));
 
-               Ok(EnforcingSigner::new_with_revoked(
+               Ok(TestChannelSigner::new_with_revoked(
                        inner,
                        state,
                        false
@@ -197,10 +197,10 @@ impl SignerProvider for OnlyReadsKeysInterface {
 }
 
 pub struct TestChainMonitor<'a> {
-       pub added_monitors: Mutex<Vec<(OutPoint, channelmonitor::ChannelMonitor<EnforcingSigner>)>>,
+       pub added_monitors: Mutex<Vec<(OutPoint, channelmonitor::ChannelMonitor<TestChannelSigner>)>>,
        pub monitor_updates: Mutex<HashMap<ChannelId, Vec<channelmonitor::ChannelMonitorUpdate>>>,
        pub latest_monitor_update_id: Mutex<HashMap<ChannelId, (OutPoint, u64, MonitorUpdateId)>>,
-       pub chain_monitor: chainmonitor::ChainMonitor<EnforcingSigner, &'a TestChainSource, &'a chaininterface::BroadcasterInterface, &'a TestFeeEstimator, &'a TestLogger, &'a chainmonitor::Persist<EnforcingSigner>>,
+       pub chain_monitor: chainmonitor::ChainMonitor<TestChannelSigner, &'a TestChainSource, &'a chaininterface::BroadcasterInterface, &'a TestFeeEstimator, &'a TestLogger, &'a chainmonitor::Persist<TestChannelSigner>>,
        pub keys_manager: &'a TestKeysInterface,
        /// If this is set to Some(), the next update_channel call (not watch_channel) must be a
        /// ChannelForceClosed event for the given channel_id with should_broadcast set to the given
@@ -208,7 +208,7 @@ pub struct TestChainMonitor<'a> {
        pub expect_channel_force_closed: Mutex<Option<(ChannelId, bool)>>,
 }
 impl<'a> TestChainMonitor<'a> {
-       pub fn new(chain_source: Option<&'a TestChainSource>, broadcaster: &'a chaininterface::BroadcasterInterface, logger: &'a TestLogger, fee_estimator: &'a TestFeeEstimator, persister: &'a chainmonitor::Persist<EnforcingSigner>, keys_manager: &'a TestKeysInterface) -> Self {
+       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 {
                Self {
                        added_monitors: Mutex::new(Vec::new()),
                        monitor_updates: Mutex::new(HashMap::new()),
@@ -224,13 +224,13 @@ impl<'a> TestChainMonitor<'a> {
                self.chain_monitor.channel_monitor_updated(outpoint, latest_update).unwrap();
        }
 }
-impl<'a> chain::Watch<EnforcingSigner> for TestChainMonitor<'a> {
-       fn watch_channel(&self, funding_txo: OutPoint, monitor: channelmonitor::ChannelMonitor<EnforcingSigner>) -> chain::ChannelMonitorUpdateStatus {
+impl<'a> chain::Watch<TestChannelSigner> for TestChainMonitor<'a> {
+       fn watch_channel(&self, funding_txo: OutPoint, monitor: channelmonitor::ChannelMonitor<TestChannelSigner>) -> chain::ChannelMonitorUpdateStatus {
                // At every point where we get a monitor update, we should be able to send a useful monitor
                // to a watchtower and disk...
                let mut w = TestVecWriter(Vec::new());
                monitor.write(&mut w).unwrap();
-               let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<EnforcingSigner>)>::read(
+               let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<TestChannelSigner>)>::read(
                        &mut io::Cursor::new(&w.0), (self.keys_manager, self.keys_manager)).unwrap().1;
                assert!(new_monitor == monitor);
                self.latest_monitor_update_id.lock().unwrap().insert(funding_txo.to_channel_id(),
@@ -264,7 +264,7 @@ impl<'a> chain::Watch<EnforcingSigner> for TestChainMonitor<'a> {
                let monitor = self.chain_monitor.get_monitor(funding_txo).unwrap();
                w.0.clear();
                monitor.write(&mut w).unwrap();
-               let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<EnforcingSigner>)>::read(
+               let new_monitor = <(BlockHash, channelmonitor::ChannelMonitor<TestChannelSigner>)>::read(
                        &mut io::Cursor::new(&w.0), (self.keys_manager, self.keys_manager)).unwrap().1;
                assert!(new_monitor == *monitor);
                self.added_monitors.lock().unwrap().push((funding_txo, new_monitor));
@@ -296,6 +296,7 @@ pub(crate) struct WatchtowerPersister {
 }
 
 impl WatchtowerPersister {
+       #[cfg(test)]
        pub(crate) fn new(destination_script: Script) -> Self {
                WatchtowerPersister {
                        persister: TestPersister::new(),
@@ -305,6 +306,7 @@ impl WatchtowerPersister {
                }
        }
 
+       #[cfg(test)]
        pub(crate) fn justice_tx(&self, funding_txo: OutPoint, commitment_txid: &Txid)
        -> Option<Transaction> {
                self.watchtower_state.lock().unwrap().get(&funding_txo).unwrap().get(commitment_txid).cloned()
@@ -425,6 +427,97 @@ impl<Signer: sign::WriteableEcdsaChannelSigner> chainmonitor::Persist<Signer> fo
        }
 }
 
+pub(crate) struct TestStore {
+       persisted_bytes: Mutex<HashMap<String, HashMap<String, Vec<u8>>>>,
+       read_only: bool,
+}
+
+impl TestStore {
+       pub fn new(read_only: bool) -> Self {
+               let persisted_bytes = Mutex::new(HashMap::new());
+               Self { persisted_bytes, read_only }
+       }
+}
+
+impl KVStore for TestStore {
+       fn read(&self, namespace: &str, sub_namespace: &str, key: &str) -> io::Result<Vec<u8>> {
+               let persisted_lock = self.persisted_bytes.lock().unwrap();
+               let prefixed = if sub_namespace.is_empty() {
+                       namespace.to_string()
+               } else {
+                       format!("{}/{}", namespace, sub_namespace)
+               };
+
+               if let Some(outer_ref) = persisted_lock.get(&prefixed) {
+                       if let Some(inner_ref) = outer_ref.get(key) {
+                               let bytes = inner_ref.clone();
+                               Ok(bytes)
+                       } else {
+                               Err(io::Error::new(io::ErrorKind::NotFound, "Key not found"))
+                       }
+               } else {
+                       Err(io::Error::new(io::ErrorKind::NotFound, "Namespace not found"))
+               }
+       }
+
+       fn write(&self, namespace: &str, sub_namespace: &str, key: &str, buf: &[u8]) -> io::Result<()> {
+               if self.read_only {
+                       return Err(io::Error::new(
+                               io::ErrorKind::PermissionDenied,
+                               "Cannot modify read-only store",
+                       ));
+               }
+               let mut persisted_lock = self.persisted_bytes.lock().unwrap();
+
+               let prefixed = if sub_namespace.is_empty() {
+                       namespace.to_string()
+               } else {
+                       format!("{}/{}", namespace, sub_namespace)
+               };
+               let outer_e = persisted_lock.entry(prefixed).or_insert(HashMap::new());
+               let mut bytes = Vec::new();
+               bytes.write_all(buf)?;
+               outer_e.insert(key.to_string(), bytes);
+               Ok(())
+       }
+
+       fn remove(&self, namespace: &str, sub_namespace: &str, key: &str, _lazy: bool) -> io::Result<()> {
+               if self.read_only {
+                       return Err(io::Error::new(
+                               io::ErrorKind::PermissionDenied,
+                               "Cannot modify read-only store",
+                       ));
+               }
+
+               let mut persisted_lock = self.persisted_bytes.lock().unwrap();
+
+               let prefixed = if sub_namespace.is_empty() {
+                       namespace.to_string()
+               } else {
+                       format!("{}/{}", namespace, sub_namespace)
+               };
+               if let Some(outer_ref) = persisted_lock.get_mut(&prefixed) {
+                               outer_ref.remove(&key.to_string());
+               }
+
+               Ok(())
+       }
+
+       fn list(&self, namespace: &str, sub_namespace: &str) -> io::Result<Vec<String>> {
+               let mut persisted_lock = self.persisted_bytes.lock().unwrap();
+
+               let prefixed = if sub_namespace.is_empty() {
+                       namespace.to_string()
+               } else {
+                       format!("{}/{}", namespace, sub_namespace)
+               };
+               match persisted_lock.entry(prefixed) {
+                       hash_map::Entry::Occupied(e) => Ok(e.get().keys().cloned().collect()),
+                       hash_map::Entry::Vacant(_) => Ok(Vec::new()),
+               }
+       }
+}
+
 pub struct TestBroadcaster {
        pub txn_broadcasted: Mutex<Vec<Transaction>>,
        pub blocks: Arc<Mutex<Vec<(Block, u32)>>>,
@@ -979,16 +1072,16 @@ impl NodeSigner for TestKeysInterface {
 }
 
 impl SignerProvider for TestKeysInterface {
-       type Signer = EnforcingSigner;
+       type Signer = TestChannelSigner;
 
        fn generate_channel_keys_id(&self, inbound: bool, channel_value_satoshis: u64, user_channel_id: u128) -> [u8; 32] {
                self.backing.generate_channel_keys_id(inbound, channel_value_satoshis, user_channel_id)
        }
 
-       fn derive_channel_signer(&self, channel_value_satoshis: u64, channel_keys_id: [u8; 32]) -> EnforcingSigner {
+       fn derive_channel_signer(&self, channel_value_satoshis: u64, channel_keys_id: [u8; 32]) -> TestChannelSigner {
                let keys = self.backing.derive_channel_signer(channel_value_satoshis, channel_keys_id);
                let state = self.make_enforcement_state_cell(keys.commitment_seed);
-               EnforcingSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
+               TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
        }
 
        fn read_chan_signer(&self, buffer: &[u8]) -> Result<Self::Signer, msgs::DecodeError> {
@@ -997,7 +1090,7 @@ impl SignerProvider for TestKeysInterface {
                let inner: InMemorySigner = ReadableArgs::read(&mut reader, self)?;
                let state = self.make_enforcement_state_cell(inner.commitment_seed);
 
-               Ok(EnforcingSigner::new_with_revoked(
+               Ok(TestChannelSigner::new_with_revoked(
                        inner,
                        state,
                        self.disable_revocation_policy_check
@@ -1038,10 +1131,10 @@ impl TestKeysInterface {
                self
        }
 
-       pub fn derive_channel_keys(&self, channel_value_satoshis: u64, id: &[u8; 32]) -> EnforcingSigner {
+       pub fn derive_channel_keys(&self, channel_value_satoshis: u64, id: &[u8; 32]) -> TestChannelSigner {
                let keys = self.backing.derive_channel_keys(channel_value_satoshis, id);
                let state = self.make_enforcement_state_cell(keys.commitment_seed);
-               EnforcingSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
+               TestChannelSigner::new_with_revoked(keys, state, self.disable_revocation_policy_check)
        }
 
        fn make_enforcement_state_cell(&self, commitment_seed: [u8; 32]) -> Arc<Mutex<EnforcementState>> {