fix: handle the return value of exports::storage_remove() (#712)

* fix: handle the return value of `exports::storage_remove()` * chore: minor documentation update * Update engine-sdk/src/near_runtime.rs Co-authored-by: Michael Birch <[email protected]> * fix: clippy error --------- Co-authored-by: Michael Birch <[email protected]>
aurora-is-near · Apr 5, 2023 · b9d5da0 · b9d5da0
1 parent 2edf523
commit b9d5da0
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/engine-sdk/src/near_runtime.rs b/engine-sdk/src/near_runtime.rs
@@ -57,15 +57,20 @@ impl Runtime {
             // Use register 0 as the destination for the promise.
             let promise_id = exports::promise_batch_create(u64::MAX as _, 0);
             // Remove code from storage and store it in register 1.
-            exports::storage_remove(code_key.len() as _, code_key.as_ptr() as _, 1);
-            exports::promise_batch_action_deploy_contract(promise_id, u64::MAX, 1);
-            Self::promise_batch_action_function_call(
-                promise_id,
-                b"state_migration",
-                &[],
-                0,
-                Self::GAS_FOR_STATE_MIGRATION.as_u64(),
-            );
+            // if the code_key exists call `promise_batch_action_deploy_contract`, otherwise panic with error.
+            match exports::storage_remove(code_key.len() as _, code_key.as_ptr() as _, 1) {
+                1 => {
+                    exports::promise_batch_action_deploy_contract(promise_id, u64::MAX, 1);
+                    Self::promise_batch_action_function_call(
+                        promise_id,
+                        b"state_migration",
+                        &[],
+                        0,
+                        Self::GAS_FOR_STATE_MIGRATION.as_u64(),
+                    );
+                }
+                _ => panic!("ERR_NO_CODE_STAGED"),
+            };
         }
     }