fix workload shieldsync

workload/src/build_checks/utils.rs

@@ -188,6 +188,8 @@ pub async fn shield_sync(sdk: &Sdk) -> Result<(), StepError> {
         .await
         .map_err(|e| StepError::ShieldedSync(e.to_string()))?;
 
+    shielded_ctx.save().await.map_err(|e| StepError::ShieldedSync(e.to_string()))?;
+
     tracing::info!("Done shieldsync (took {}s)!", now.elapsed().as_secs());
 
     Ok(())

workload/src/main.rs

@@ -89,8 +89,13 @@ async fn inner_main() -> i32 {
 
         // Setup shielded context storage
         let shielded_ctx_path = state.base_dir.join(format!("masp-{}", config.id));
-        let shielded_ctx = ShieldedContext::new(FsShieldedUtils::new(shielded_ctx_path));
-        shielded_ctx.save().await.unwrap();
+
+        let mut shielded_ctx = ShieldedContext::new(FsShieldedUtils::new(shielded_ctx_path.clone()));
+        if shielded_ctx_path.join("shielded.dat").exists() {
+            shielded_ctx.load().await.expect("Should be able to load shielded context");
+        } else {
+            shielded_ctx.save().await.unwrap();
+        }
 
         let io = NullIo;