diff --git a/rusk/src/lib/error.rs b/rusk/src/lib/error.rs index d5986e2e75..d6b715aa70 100644 --- a/rusk/src/lib/error.rs +++ b/rusk/src/lib/error.rs @@ -49,6 +49,8 @@ pub enum Error { InconsistentState(VerificationOutput), /// Other Other(Box), + /// Commit not found amongst existing commits + CommitNotFound([u8; 32]), } impl std::error::Error for Error {} @@ -140,6 +142,9 @@ impl fmt::Display for Error { "Inconsistent state verification data {verification_output}", ) } + Error::CommitNotFound(commit_id) => { + write!(f, "Commit not found, id = {}", hex::encode(commit_id),) + } } } } diff --git a/rusk/src/lib/lib.rs b/rusk/src/lib/lib.rs index 4b5775e321..e0a573f2f8 100644 --- a/rusk/src/lib/lib.rs +++ b/rusk/src/lib/lib.rs @@ -334,9 +334,15 @@ impl Rusk { Ok(()) } - pub fn revert(&self) -> Result<[u8; 32]> { + pub fn revert(&self, state_hash: [u8; 32]) -> Result<[u8; 32]> { let mut inner = self.inner.lock(); - inner.current_commit = inner.base_commit; + + let mut commits = &inner.vm.commits(); + if !commits.contains(&state_hash) { + return Err(Error::CommitNotFound(state_hash)); + } + + inner.current_commit = state_hash; Ok(inner.current_commit) } diff --git a/rusk/src/lib/vm.rs b/rusk/src/lib/vm.rs index 271a620584..58208ac1af 100644 --- a/rusk/src/lib/vm.rs +++ b/rusk/src/lib/vm.rs @@ -149,9 +149,13 @@ impl VMExecution for Rusk { Ok(self.state_root()) } - fn revert(&self) -> anyhow::Result<[u8; 32]> { + fn get_base_state_root(&self) -> anyhow::Result<[u8; 32]> { + Ok(self.base_root()) + } + + fn revert(&self, state_hash: [u8; 32]) -> anyhow::Result<[u8; 32]> { let state_hash = self - .revert() + .revert(state_hash) .map_err(|inner| anyhow::anyhow!("Cannot revert: {inner}"))?; Ok(state_hash)