Merge pull request #9 from yangby-cryptape/feature/dummy-can-rollback

feat: dummy service can rollback to an old client

Author: chaoticlonghair

prover/src/dummy_service.rs

@@ -8,7 +8,7 @@ use ckb_bitcoin_spv_verifier::{
     utilities::{bitcoin::calculate_next_target, mmr},
 };
 
-use crate::result::Result;
+use crate::result::{Error, Result};
 
 /// A dummy service for testing the SPV client cells's bootstrap and update.
 pub struct DummyService {
@@ -20,7 +20,11 @@ pub struct DummyService {
 impl DummyService {
     pub fn bootstrap(height: u32, header: core::Header) -> Result<Self> {
         if height % DIFFCHANGE_INTERVAL != 0 {
-            panic!("bad height");
+            let msg = format!(
+                "bad bootstrap height, expected multiples of \
+                {DIFFCHANGE_INTERVAL} but got {height}"
+            );
+            return Err(Error::other(msg));
         }
         let mut headers = HashMap::new();
         let store = mmr::lib::util::MemStore::default();
@@ -117,6 +121,26 @@ impl DummyService {
             .build())
     }
 
+    // The `prev_client` is not checked, since this is just a dummy service for testing purpose only.
+    pub fn rollback_to(&mut self, prev_client: core::SpvClient) -> Result<()> {
+        let prev_height = prev_client.headers_mmr_root.max_height;
+        if prev_height < self.client.headers_mmr_root.min_height
+            || self.client.headers_mmr_root.max_height < prev_height
+        {
+            let msg = format!(
+                "the previous header (height: {prev_height}) is not found (current: [{}, {}])",
+                self.client.headers_mmr_root.min_height, self.client.headers_mmr_root.max_height
+            );
+            return Err(Error::other(msg));
+        }
+        let curr_height = self.client.headers_mmr_root.max_height;
+        for h in (prev_height + 1)..=curr_height {
+            self.headers.remove(&h);
+        }
+        self.client = prev_client;
+        Ok(())
+    }
+
     pub fn tip_client(&self) -> core::SpvClient {
         self.client.clone()
     }

prover/src/result.rs

@@ -16,7 +16,7 @@ pub enum Error {
 pub type Result<T> = ::std::result::Result<T, Error>;
 
 impl Error {
-    pub(crate) fn other<S: ToString>(arg: S) -> Self {
+    pub fn other<S: ToString>(arg: S) -> Self {
         Self::Other(arg.to_string())
     }
 }

prover/src/tests/service.rs

@@ -66,7 +66,17 @@ fn test_spv_client(
         }
 
         log::trace!("process {} headers at one time", headers.len());
-        let update = service.update(mem::take(&mut headers)).unwrap();
+        let update = if headers_group_size % 5 == 0 {
+            let tmp_headers = mem::take(&mut headers);
+            let _update = service.update(tmp_headers.clone()).unwrap();
+            log::trace!("rollback to previous client (for test): {old_client}");
+            service.rollback_to(old_client.unpack()).unwrap();
+            log::trace!("process {} headers again", headers.len());
+            service.update(tmp_headers)
+        } else {
+            service.update(mem::take(&mut headers))
+        }
+        .unwrap();
         if verify_tx_range.0 <= height + 1 && height <= verify_tx_range.1 {
             headers_group_size = 1;
         } else {