feat(execution_environment): Allow msg_reject_code in call_on_cleanup (#6915)

This PR allows `msg_reject_code` system API to be accessible in the
cleanup callback. This allows users to know the status of the **inter
canister call (success / failure) in the cleanup (in case of traps)**
and can make data decisions based on that.

Author: venkkatesh-sekar

rs/embedders/src/wasmtime_embedder/system_api.rs

@@ -410,6 +410,7 @@ pub enum ApiType {
     Cleanup {
         caller: PrincipalId,
         time: Time,
+        reject_code: i32,
         /// The total number of instructions executed in the call context
         call_context_instructions_executed: NumInstructions,
     },
@@ -1445,7 +1446,6 @@ impl SystemApiImpl {
             ApiType::Start { .. }
             | ApiType::Init { .. }
             | ApiType::SystemTask { .. }
-            | ApiType::Cleanup { .. }
             | ApiType::CompositeCleanup { .. }
             | ApiType::ReplicatedQuery { .. }
             | ApiType::NonReplicatedQuery { .. }
@@ -1458,6 +1458,7 @@ impl SystemApiImpl {
             | ApiType::CompositeRejectCallback { reject_context, .. } => {
                 Some(reject_context.code() as i32)
             }
+            ApiType::Cleanup { reject_code, .. } => Some(*reject_code),
         }
     }
 

rs/embedders/tests/system_api.rs

@@ -157,6 +157,7 @@ fn cleanup_api() -> ApiType {
     ApiType::Cleanup {
         caller: PrincipalId::new_anonymous(),
         time: UNIX_EPOCH,
+        reject_code: 0,
         call_context_instructions_executed: 0.into(),
     }
 }
@@ -186,7 +187,7 @@ fn is_supported(api_type: SystemApiCallId, context: &str) -> bool {
         SystemApiCallId::MsgArgDataCopy => vec!["I", "U", "RQ", "NRQ", "CQ", "Ry", "CRy", "F"],
         SystemApiCallId::MsgCallerSize => vec!["*"],
         SystemApiCallId::MsgCallerCopy => vec!["*"],
-        SystemApiCallId::MsgRejectCode => vec!["Ry", "Rt", "CRy", "CRt"],
+        SystemApiCallId::MsgRejectCode => vec!["Ry", "Rt", "CRy", "CRt", "C"],
         SystemApiCallId::MsgRejectMsgSize => vec!["Rt", "CRt"],
         SystemApiCallId::MsgRejectMsgCopy => vec!["Rt", "CRt"],
         SystemApiCallId::MsgReplyDataAppend => vec!["U", "RQ", "NRQ", "CQ", "Ry", "Rt", "CRy", "CRt"],

rs/execution_environment/src/execution/response.rs

@@ -1075,10 +1075,17 @@ fn execute_response_cleanup(
             FuncRef::QueryClosure(cleanup_closure)
         }
     };
+
+    let reject_code = match &original.message.response_payload {
+        Payload::Data(_) => 0,
+        Payload::Reject(context) => context.code() as i32,
+    };
+
     let result = round.hypervisor.execute_dts(
         ApiType::Cleanup {
             caller: original.call_origin.get_principal(),
             time: original.time,
+            reject_code,
             call_context_instructions_executed: original.instructions_executed,
         },
         helper.canister().execution_state.as_ref().unwrap(),

rs/execution_environment/tests/hypervisor.rs

@@ -10518,3 +10518,164 @@ fn parallel_callbacks() {
     assert_eq!(err.code(), ErrorCode::CanisterCalledTrap);
     assert!(err.description().contains("trap in second callback"));
 }
+
+const REPLY_REJECT_CLEANUP_CALLBACK_WAT: &str = r#"
+    (module
+        (import "ic0" "call_new"
+            (func $ic0_call_new
+                (param i32 i32)
+                (param $method_name_src i32)    (param $method_name_len i32)
+                (param $reply_fun i32)          (param $reply_env i32)
+                (param $reject_fun i32)         (param $reject_env i32)
+            ))
+        (import "ic0" "call_perform" (func $ic0_call_perform (result i32)))
+
+        (import "ic0" "trap" (func $ic0_trap (param i32) (param i32)))
+        (import "ic0" "call_on_cleanup" (func $ic0_call_on_cleanup (param i32) (param i32)))
+        (import "ic0" "msg_reject_code" (func $ic0_msg_reject_code (result i32)))
+
+        (import "ic0" "msg_reply" (func $msg_reply))
+        (import "ic0" "msg_reply_data_append"
+            (func $msg_reply_data_append (param i32) (param i32)))
+        
+        (func $dummy
+            (call $msg_reply_data_append
+                (i32.const 300) (i32.const 1))  ;; refers to 9 on the heap
+            (call $msg_reply)
+        )   
+
+
+        (func $test
+            (call $ic0_call_new
+                (i32.const 100) (i32.const 10)   ;; callee canister id = 0
+                (i32.const 0) (i32.const 4)      ;; refers to "test" on the heap
+                (i32.const 0) (i32.const 200)    ;; on_reply closure at table index 0
+                (i32.const 1) (i32.const 200))   ;; on_reject closure at table index 1
+            (call $ic0_call_on_cleanup 
+                (i32.const 2) (i32.const 200))   ;; cleanup closure at table index 2
+            (drop (call $ic0_call_perform))
+        )
+
+        (func $on_reply (param i32)
+            (call $ic0_trap 
+                (i32.const 200) (i32.const 12))  ;; reply callback traps
+        )
+
+        (func $on_reject (param i32) 
+            (call $ic0_trap 
+                (i32.const 200) (i32.const 12))  ;; reject callback traps
+        )
+
+        (func $on_cleanup (param i32)
+            (i32.store8                                      ;; cleanup can't reply, so
+                (i32.const 300) (call $ic0_msg_reject_code)) ;; we write cleanup code to memory and retrieve it via $dummy                    
+        )
+
+        (export "canister_update test" (func $test))
+        (export "canister_update dummy" (func $dummy))
+        (memory $memory 1)
+        (export "memory" (memory $memory))
+        (data (i32.const 0) "test")
+        (data (i32.const 100) "\00\00\00\00\00\00\00\00\01\01") ;; cansister_id of the installed canister is 0
+        (data (i32.const 200) "trap message")
+        (data (i32.const 300) "\09")
+        (table 3 3 funcref)
+        (elem (i32.const 0) $on_reply $on_reject $on_cleanup)
+    )
+"#;
+
+#[test]
+fn can_access_reject_code_in_cleanup_call_rejected() {
+    const CANISTER_SIMPLE_WAT: &str = r#"
+        (module
+            (import "ic0" "msg_reject"
+            (func $msg_reject (param $msg_reject_src i32) (param $msg_reject_size i32)))
+            (func $test
+                (call $msg_reject
+                    (i32.const 0) (i32.const 26))  ;; refers to "explictly_rejected_message" on the heap
+            )
+            (export "canister_update test" (func $test))
+            (memory $memory 1)
+            (export "memory" (memory $memory))
+            (data (i32.const 0) "explictly_rejected_message")
+        )"#;
+
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.create_canister(Cycles::new(1_000_000_000_000));
+
+    // assert that the hardcoded canister_id in REPLY_REJECT_CLEANUP_CALLBACK_WAT
+    // matches the one from instantiation
+    assert_eq!(canister_id, CanisterId::from_u64(0));
+
+    test.install_canister(canister_id, wat::parse_str(CANISTER_SIMPLE_WAT).unwrap())
+        .unwrap();
+
+    let main_canister_id = test.create_canister(Cycles::new(1_000_000_000_000));
+    test.install_canister(
+        main_canister_id,
+        wat::parse_str(REPLY_REJECT_CLEANUP_CALLBACK_WAT).unwrap(),
+    )
+    .unwrap();
+
+    // Initial value at the cleanup reference is 9
+    let result = test.ingress(main_canister_id, "dummy", vec![]).unwrap();
+    assert_eq!(WasmResult::Reply(vec![9]), result);
+
+    // Reject callback traps
+    let result = test.ingress(main_canister_id, "test", vec![]);
+    assert!(result.is_err());
+
+    // msg_reject_code should be written to memory now via cleanup callback
+    // reject code is 4 when a canister explicitly rejects
+    let result = test.ingress(main_canister_id, "dummy", vec![]).unwrap();
+    assert_eq!(WasmResult::Reply(vec![4]), result);
+}
+
+#[test]
+fn can_access_reject_code_in_cleanup_call_replied() {
+    const CANISTER_SIMPLE_WAT: &str = r#"
+        (module
+            (import "ic0" "msg_reply" (func $msg_reply))
+            (import "ic0" "msg_reply_data_append"
+                (func $msg_reply_data_append (param i32) (param i32)))
+            (func $test
+                (call $msg_reply_data_append
+                    (i32.const 0) (i32.const 26))  ;; refers to "explictly_approved_message" on the heap
+                (call $msg_reply)
+            )
+            (export "canister_update test" (func $test))
+            (memory $memory 1)
+            (export "memory" (memory $memory))
+            (data (i32.const 0) "explictly_approved_message")
+        )"#;
+
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.create_canister(Cycles::new(1_000_000_000_000));
+
+    // assert that the hardcoded canister id in REPLY_REJECT_CLEANUP_CALLBACK_WAT
+    // matches the one from instantiation
+    assert_eq!(canister_id, CanisterId::from_u64(0));
+
+    test.install_canister(canister_id, wat::parse_str(CANISTER_SIMPLE_WAT).unwrap())
+        .unwrap();
+
+    let main_canister_id = test.create_canister(Cycles::new(1_000_000_000_000));
+    test.install_canister(
+        main_canister_id,
+        wat::parse_str(REPLY_REJECT_CLEANUP_CALLBACK_WAT).unwrap(),
+    )
+    .unwrap();
+
+    // Initial value at the cleanup reference is 9
+    let result = test.ingress(main_canister_id, "dummy", vec![]).unwrap();
+    assert_eq!(WasmResult::Reply(vec![9]), result);
+
+    // Reply callback traps
+    let result = test.ingress(main_canister_id, "test", vec![]);
+    assert!(result.is_err());
+
+    // msg_reject_code should be written to memory now via cleanup callback
+    // canister replied would be reject code 0
+    let result = test.ingress(main_canister_id, "dummy", vec![]).unwrap();
+    assert_eq!(WasmResult::Reply(vec![0]), result);
+}