Add a utility to replay CAPE wallet events locally

wallet/src/backend.rs

@@ -662,7 +662,7 @@ impl<'a> CapeWalletBackend<'a> for CapeBackend<'a> {
     }
 }
 
-fn gen_proving_keys(srs: &UniversalParam) -> ProverKeySet<key_set::OrderByOutputs> {
+pub fn gen_proving_keys(srs: &UniversalParam) -> ProverKeySet<key_set::OrderByOutputs> {
     use jf_cap::proof::{freeze, mint, transfer};
 
     ProverKeySet {

wallet/src/bin/wallet-replay.rs

@@ -0,0 +1,44 @@
+// Copyright (c) 2022 Espresso Systems (espressosys.com)
+// This file is part of the Configurable Asset Privacy for Ethereum (CAPE) library.
+//
+// This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
+// This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use cape_wallet::mocks::{MockCapeWalletLoader, ReplayBackend};
+use seahorse::{
+    events::EventIndex,
+    hd::{KeyTree, Mnemonic},
+    Wallet,
+};
+use std::fs;
+use std::path::PathBuf;
+use structopt::StructOpt;
+use tempdir::TempDir;
+
+#[derive(StructOpt)]
+struct Options {
+    #[structopt(long, env = "CAPE_REPLAY_MNEMONIC")]
+    mnemonic: Mnemonic,
+    events: PathBuf,
+}
+
+#[async_std::main]
+async fn main() {
+    let opt = Options::from_args();
+    let storage = TempDir::new("cape-wallet-replay").unwrap();
+    let mut loader = MockCapeWalletLoader {
+        key: KeyTree::from_mnemonic(&opt.mnemonic),
+        path: storage.path().to_path_buf(),
+    };
+    let events_bytes = fs::read(opt.events).unwrap();
+    let events = serde_json::from_slice(&events_bytes).unwrap();
+    let backend = ReplayBackend::new(events, &mut loader);
+    let mut wallet = Wallet::new(backend).await.unwrap();
+
+    let key = wallet
+        .generate_user_key("key".into(), Some(EventIndex::default()))
+        .await
+        .unwrap();
+    wallet.await_key_scan(&key.address()).await.unwrap();
+}

wallet/src/mocks.rs

@@ -17,7 +17,7 @@ use cap_rust_sandbox::{
     deploy::EthMiddleware, ledger::*, model::*, universal_param::UNIVERSAL_PARAM,
 };
 use commit::Committable;
-use futures::stream::Stream;
+use futures::stream::{iter, pending, Stream, StreamExt};
 use itertools::izip;
 use jf_cap::{
     keys::{UserAddress, UserKeyPair, UserPubKey},
@@ -764,6 +764,165 @@ fn cape_to_wallet_err(err: CapeValidationError) -> WalletError<CapeLedger> {
     }
 }
 
+/// A mock CAPE wallet backend that can be used to replay a given list of events.
+pub struct ReplayBackend<'a, Meta: Send + DeserializeOwned + Serialize> {
+    events: Vec<LedgerEvent<CapeLedger>>,
+    storage: Arc<Mutex<AtomicWalletStorage<'a, CapeLedger, Meta>>>,
+    key_stream: KeyTree,
+}
+
+impl<'a, Meta: Clone + PartialEq + Send + DeserializeOwned + Serialize> ReplayBackend<'a, Meta> {
+    pub fn new(
+        events: Vec<LedgerEvent<CapeLedger>>,
+        loader: &mut impl WalletLoader<CapeLedger, Meta = Meta>,
+    ) -> Self {
+        let storage = AtomicWalletStorage::new(loader, 1024).unwrap();
+        let key_stream = storage.key_stream();
+        Self {
+            events,
+            storage: Arc::new(Mutex::new(storage)),
+            key_stream,
+        }
+    }
+}
+
+#[async_trait]
+impl<'a, Meta: Send + DeserializeOwned + Serialize> WalletBackend<'a, CapeLedger>
+    for ReplayBackend<'a, Meta>
+{
+    type EventStream = Pin<Box<dyn Stream<Item = (LedgerEvent<CapeLedger>, EventSource)> + Send>>;
+    type Storage = AtomicWalletStorage<'a, CapeLedger, Meta>;
+
+    async fn storage<'l>(&'l mut self) -> MutexGuard<'l, Self::Storage> {
+        self.storage.lock().await
+    }
+
+    async fn create(&mut self) -> Result<WalletState<'a, CapeLedger>, WalletError<CapeLedger>> {
+        // Get the state from the snapshotted time (i.e. after replaying all the events).
+        let mut record_mt = MerkleTree::new(CAPE_MERKLE_HEIGHT).unwrap();
+        let mut block_height = 0;
+        for event in &self.events {
+            if let LedgerEvent::Commit { block, .. } = event {
+                for txn in block.txns() {
+                    for comm in txn.output_commitments() {
+                        record_mt.push(comm.to_field_element());
+                    }
+                }
+                block_height += 1;
+            }
+        }
+
+        // `record_mt` should be completely sparse.
+        for uid in 0..record_mt.num_leaves() - 1 {
+            record_mt.forget(uid);
+        }
+        let merkle_leaf_to_forget = if record_mt.num_leaves() > 0 {
+            Some(record_mt.num_leaves() - 1)
+        } else {
+            None
+        };
+
+        Ok(WalletState {
+            proving_keys: Arc::new(crate::backend::gen_proving_keys(&*UNIVERSAL_PARAM)),
+            txn_state: TransactionState {
+                validator: CapeTruster::new(block_height, record_mt.num_leaves()),
+                now: EventIndex::from_source(EventSource::QueryService, self.events.len()),
+                // Completely sparse nullifier set.
+                nullifiers: Default::default(),
+                record_mt,
+                records: RecordDatabase::default(),
+                merkle_leaf_to_forget,
+                transactions: Default::default(),
+            },
+            key_state: Default::default(),
+            assets: Default::default(),
+            viewing_accounts: Default::default(),
+            freezing_accounts: Default::default(),
+            sending_accounts: Default::default(),
+        })
+    }
+
+    async fn subscribe(&self, from: EventIndex, to: Option<EventIndex>) -> Self::EventStream {
+        let from = from.index(EventSource::QueryService);
+        let to = to.map(|to| to.index(EventSource::QueryService));
+
+        println!(
+            "playing back {} events from {} to {:?}",
+            self.events.len(),
+            from,
+            to
+        );
+        let events = iter(self.events.clone())
+            .enumerate()
+            .map(|(i, event)| {
+                println!("replaying event {} {:?}", i, event);
+                (event, EventSource::QueryService)
+            })
+            .skip(from);
+        let events: Self::EventStream = if let Some(to) = to {
+            Box::pin(events.take(to - from))
+        } else {
+            Box::pin(events)
+        };
+        // Append a stream which blocks forever, since the event stream is not supposed to terminate.
+        Box::pin(events.chain(pending()))
+    }
+
+    async fn get_public_key(
+        &self,
+        _address: &UserAddress,
+    ) -> Result<UserPubKey, WalletError<CapeLedger>> {
+        // Since we're not generating transactions, we don't need to support address book queries.
+        Err(WalletError::Failed {
+            msg: "address book not supported".into(),
+        })
+    }
+
+    async fn register_user_key(
+        &mut self,
+        _key_pair: &UserKeyPair,
+    ) -> Result<(), WalletError<CapeLedger>> {
+        // The wallet calls this function when it generates a key, so it has to succeed, but since
+        // we don't support querying the address book, it doesn't have to do anything.
+        Ok(())
+    }
+
+    async fn get_initial_scan_state(
+        &self,
+        _from: EventIndex,
+    ) -> Result<(MerkleTree, EventIndex), CapeWalletError> {
+        Ok((
+            MerkleTree::new(CAPE_MERKLE_HEIGHT).unwrap(),
+            EventIndex::default(),
+        ))
+    }
+
+    async fn get_nullifier_proof(
+        &self,
+        _nullifiers: &mut CapeNullifierSet,
+        _nullifier: Nullifier,
+    ) -> Result<(bool, ()), WalletError<CapeLedger>> {
+        // Nullifier queries are not needed for event playback.
+        Err(WalletError::Failed {
+            msg: "nullifier queries not supported".into(),
+        })
+    }
+
+    fn key_stream(&self) -> KeyTree {
+        self.key_stream.clone()
+    }
+
+    async fn submit(
+        &mut self,
+        _txn: CapeTransition,
+        _info: TransactionInfo<CapeLedger>,
+    ) -> Result<(), WalletError<CapeLedger>> {
+        Err(WalletError::Failed {
+            msg: "transacting not supported".into(),
+        })
+    }
+}
+
 pub struct MockCapeWalletLoader {
     pub path: PathBuf,
     pub key: KeyTree,