Refactor archive deployment code (#1017)

This PR introduces a new type to express the verification status of
the wasm.
Additionally, it makes the deploy_archive more clear by splitting the big
match statement into smaller pieces.

Author: Frederik Rothenberger

src/internet_identity/src/archive.rs

@@ -19,7 +19,6 @@ use sha2::Sha256;
 use std::time::Duration;
 use ArchiveState::{Created, CreationInProgress, NotCreated};
 use CanisterInstallMode::Upgrade;
-use DeployArchiveResult::{Failed, Success};
 
 #[derive(Clone, Debug, Default, CandidType, Deserialize, Eq, PartialEq)]
 pub struct ArchiveInfo {
@@ -60,94 +59,86 @@ pub struct ArchiveStatusCache {
     pub status: CanisterStatusResponse,
 }
 
-/// Wrapper for wasm to carry verification state.
-struct VerifiableWasm {
-    wasm: Vec<u8>,
-    verified: bool,
-}
+struct VerifiedWasm(Vec<u8>);
 
-impl VerifiableWasm {
-    pub fn from_unverified_wasm(wasm: ByteBuf) -> Self {
-        VerifiableWasm {
-            wasm: wasm.into_vec(),
-            verified: false,
-        }
+pub async fn deploy_archive(wasm: ByteBuf) -> DeployArchiveResult {
+    // archive state without creation_in_progress
+    // if creation is in progress we exit early since we do not want to deploy the archive twice
+    enum ReducedArchiveState {
+        NotCreated,
+        Created(ArchiveData),
     }
 
-    pub fn verify_wasm_hash(&mut self) -> Result<(), DeployArchiveResult> {
-        let expected_hash =
-            state::expected_archive_hash().expect("bug: no wasm hash to check against");
+    let Some(expected_hash) = state::expected_archive_hash() else {
+        return DeployArchiveResult::Failed("archive deployment disabled".to_string());
+    };
 
-        let mut hasher = Sha256::new();
-        hasher.update(&self.wasm);
-        let wasm_hash: [u8; 32] = hasher.finalize().into();
+    unlock_archive_if_stuck();
+    // exit early if archive is being created by another call to deploy_archive
+    let reduced_state = match state::archive_state() {
+        NotCreated => ReducedArchiveState::NotCreated,
+        CreationInProgress(_) => return DeployArchiveResult::CreationInProgress,
+        Created(data) => ReducedArchiveState::Created(data),
+    };
 
-        if wasm_hash != expected_hash {
-            return Err(Failed("invalid wasm module".to_string()));
+    // exit early if the expected wasm module is already installed
+    if let ReducedArchiveState::Created(ref data) = reduced_state {
+        let status = archive_status(data.archive_canister).await;
+        match status.module_hash {
+            Some(hash) if hash == expected_hash.to_vec() => {
+                // we already have an archive with the expected module and don't need to do anything
+                return DeployArchiveResult::Success(data.archive_canister);
+            }
+            None | Some(_) => {}
         }
-
-        self.verified = true;
-        Ok(())
-    }
-
-    pub fn get_verified_wasm(mut self) -> Result<Vec<u8>, DeployArchiveResult> {
-        if !self.verified {
-            self.verify_wasm_hash()?;
-        };
-        Ok(self.wasm)
     }
-}
 
-pub async fn deploy_archive(wasm: ByteBuf) -> DeployArchiveResult {
-    let mut wasm = VerifiableWasm::from_unverified_wasm(wasm);
-    let expected_hash = if let Some(hash) = state::expected_archive_hash() {
-        hash
-    } else {
-        return Failed("archive deployment disabled".to_string());
+    // all early checks passed, we need to make changes to the archive --> verify wasm
+    let verified_wasm = match verify_wasm(wasm.into_vec()) {
+        Ok(verified_wasm) => verified_wasm,
+        Err(err) => return DeployArchiveResult::Failed(err),
     };
 
-    let (archive_canister, install_mode) = match state::archive_state() {
-        NotCreated => match create_archive(&mut wasm).await {
-            Ok(archive) => (archive, Install),
-            Err(result) => return result,
-        },
-        CreationInProgress(timestamp) => {
-            if time() - timestamp > Duration::from_secs(24 * 60 * 60).as_nanos() as u64 {
-                // The archive has been in creation for more than a day and the creation process
-                // has likely failed thus another attempt should be made.
-                match create_archive(&mut wasm).await {
-                    Ok(archive) => (archive, Install),
-                    Err(result) => return result,
-                }
-            } else {
-                return DeployArchiveResult::CreationInProgress;
-            }
-        }
-        Created(archive_data) => {
-            let status = match archive_status(archive_data.archive_canister).await {
-                Ok(status) => status,
-                Err(message) => return Failed(message),
+    // create if not exists and determine install mode
+    let (archive_canister, install_mode) = match reduced_state {
+        ReducedArchiveState::NotCreated => {
+            let archive = match create_archive().await {
+                Ok(archive) => archive,
+                Err(err) => return DeployArchiveResult::Failed(err),
             };
+            (archive, Install)
+        }
+        ReducedArchiveState::Created(data) => {
+            let status = archive_status(data.archive_canister).await;
             match status.module_hash {
-                None => (archive_data.archive_canister, Install),
-                Some(hash) if hash == expected_hash.to_vec() => {
-                    // we already have an archive with the expected module and don't need to do anything
-                    return Success(archive_data.archive_canister);
-                }
-                Some(_) => (archive_data.archive_canister, Upgrade),
+                None => (data.archive_canister, Install),
+                Some(_) => (data.archive_canister, Upgrade),
             }
         }
     };
 
-    match install_archive(archive_canister, wasm, install_mode).await {
-        Ok(()) => Success(archive_canister),
-        Err(err) => err,
+    match install_archive(archive_canister, verified_wasm, install_mode).await {
+        Ok(()) => DeployArchiveResult::Success(archive_canister),
+        Err(err) => DeployArchiveResult::Failed(err),
     }
 }
 
-async fn create_archive(wasm: &mut VerifiableWasm) -> Result<Principal, DeployArchiveResult> {
-    wasm.verify_wasm_hash()?;
+fn unlock_archive_if_stuck() {
+    match state::archive_state() {
+        NotCreated | Created(_) => {}
+        CreationInProgress(timestamp) => {
+            // The archive has been in creation for more than a day and the creation process
+            // has likely failed thus another attempt should be made.
+            if time() - timestamp > Duration::from_secs(24 * 60 * 60).as_nanos() as u64 {
+                state::persistent_state_mut(|persistent_state| {
+                    persistent_state.archive_info.state = NotCreated
+                })
+            }
+        }
+    }
+}
 
+async fn create_archive() -> Result<Principal, String> {
     // lock the archive
     state::persistent_state_mut(|persistent_state| {
         persistent_state.archive_info.state = CreationInProgress(time());
@@ -171,10 +162,10 @@ async fn create_archive(wasm: &mut VerifiableWasm) -> Result<Principal, DeployAr
             state::persistent_state_mut(|persistent_state| {
                 persistent_state.archive_info.state = NotCreated
             });
-            Err(Failed(format!(
+            Err(format!(
                 "failed to create archive! error code: {:?}, message: {}",
                 reject_code, message
-            )))
+            ))
         }
     }
 }
@@ -201,52 +192,65 @@ async fn create_canister(arg: CreateCanisterArgument) -> CallResult<(CanisterIdR
 
 async fn install_archive(
     archive_canister: Principal,
-    wasm: VerifiableWasm,
+    wasm: VerifiedWasm,
     install_mode: CanisterInstallMode,
-) -> Result<(), DeployArchiveResult> {
+) -> Result<(), String> {
     let settings = ArchiveInit {
         ii_canister: id(),
         max_entries_per_call: 1000,
     };
-    let encoded_arg = candid::encode_one(settings).map_err(|err| {
-        Failed(format!(
-            "failed to encode archive install argument: {:?}",
-            err
-        ))
-    })?;
+    let encoded_arg = candid::encode_one(settings)
+        .map_err(|err| format!("failed to encode archive install argument: {:?}", err))?;
 
     install_code(InstallCodeArgument {
         mode: install_mode,
         canister_id: archive_canister,
-        wasm_module: wasm.get_verified_wasm()?,
+        wasm_module: wasm.0,
         arg: encoded_arg,
     })
     .await
     .map_err(|(code, message)| {
-        Failed(format!(
+        format!(
             "failed to install archive canister! error code: {:?}, message: {}",
             code, message
-        ))
+        )
     })
 }
 
-async fn archive_status(archive_canister: Principal) -> Result<CanisterStatusResponse, String> {
+fn verify_wasm(wasm: Vec<u8>) -> Result<VerifiedWasm, String> {
+    let expected_hash = match state::expected_archive_hash() {
+        None => return Err("archive deployment disabled".to_string()),
+        Some(hash) => hash,
+    };
+
+    let mut hasher = Sha256::new();
+    hasher.update(&wasm);
+    let actual_hash: [u8; 32] = hasher.finalize().into();
+
+    if actual_hash != expected_hash {
+        return Err("invalid wasm module".to_string());
+    }
+
+    Ok(VerifiedWasm(wasm))
+}
+
+async fn archive_status(archive_canister: Principal) -> CanisterStatusResponse {
     let status_opt = state::cached_archive_status();
     match status_opt {
         None => {
             let (archive_status,) = canister_status(CanisterIdRecord {
                 canister_id: archive_canister,
             })
             .await
-            .map_err(|err| format!("failed to retrieve archive status: {:?}", err))?;
+            .expect(&format!("failed to retrieve archive status"));
 
             state::cache_archive_status(ArchiveStatusCache {
                 timestamp: time(),
                 status: archive_status.clone(),
             });
-            Ok(archive_status)
+            archive_status
         }
-        Some(status) => Ok(status),
+        Some(status) => status,
     }
 }