feat(PocketIC): optionally check caller when retrieving ingress status (

#3341)

This PR extends the PocketIC functionality of retrieving ingress status
with an optional caller: if provided and if a corresponding read state
request for the status of the same update call signed by that specified
caller was rejected because the update call was submitted by a different
caller, then an error is returned. To distinguish between the ingress
status not being available and forbidden due to the new error condition
introduced by this PR, a new enumeration `IngressStatusResult` is
introduced in the PocketIC library.

packages/pocket-ic/CHANGELOG.md

@@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - The function `PocketIcBuilder::with_bitcoind_addrs` to specify multiple addresses and ports at which `bitcoind` processes are listening.
 - The function `PocketIc::query_call_with_effective_principal` for making generic query calls (including management canister query calls).
-- The function `PocketIc::ingress_status` to fetch the status of an update call submitted through an ingress message (`None` means that the status is unknown yet).
+- The function `PocketIc::ingress_status` to fetch the status of an update call submitted through an ingress message.
 - The function `PocketIc::await_call_no_ticks` to await the status of an update call (submitted through an ingress message) becoming known without triggering round execution
   (round execution must be triggered separarely, e.g., on a "live" instance or by separate PocketIC library calls).
 

packages/pocket-ic/src/common/rest.rs

@@ -122,6 +122,12 @@ pub struct RawMessageId {
     pub message_id: Vec<u8>,
 }
 
+#[derive(Clone, Serialize, Deserialize, Debug, JsonSchema)]
+pub struct RawIngressStatusArgs {
+    pub raw_message_id: RawMessageId,
+    pub raw_caller: Option<RawPrincipalId>,
+}
+
 #[derive(Clone, Serialize, Deserialize, Debug, JsonSchema)]
 pub enum RawSubmitIngressResult {
     Ok(RawMessageId),

packages/pocket-ic/src/lib.rs

@@ -681,9 +681,10 @@ impl PocketIc {
     pub fn ingress_status(
         &self,
         message_id: RawMessageId,
-    ) -> Option<Result<WasmResult, UserError>> {
+        caller: Option<Principal>,
+    ) -> IngressStatusResult {
         let runtime = self.runtime.clone();
-        runtime.block_on(async { self.pocket_ic.ingress_status(message_id).await })
+        runtime.block_on(async { self.pocket_ic.ingress_status(message_id, caller).await })
     }
 
     /// Await an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
@@ -1525,6 +1526,18 @@ pub enum CallError {
     UserError(UserError),
 }
 
+/// This enum describes the result of retrieving ingress status.
+/// The `IngressStatusResult::Forbidden` variant is produced
+/// if an optional caller is provided and a corresponding read state request
+/// for the status of the same update call signed by that specified caller
+/// was rejected because the update call was submitted by a different caller.
+#[derive(Debug, Serialize, Deserialize)]
+pub enum IngressStatusResult {
+    NotAvailable,
+    Forbidden(String),
+    Success(Result<WasmResult, UserError>),
+}
+
 /// This struct describes the different types that executing a WASM function in
 /// a canister can produce.
 #[derive(PartialOrd, Ord, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]

packages/pocket-ic/src/nonblocking.rs

@@ -3,9 +3,10 @@ use crate::common::rest::{
     CreateHttpGatewayResponse, CreateInstanceResponse, ExtendedSubnetConfigSet, HttpGatewayBackend,
     HttpGatewayConfig, HttpGatewayInfo, HttpsConfig, InstanceConfig, InstanceId,
     MockCanisterHttpResponse, RawAddCycles, RawCanisterCall, RawCanisterHttpRequest, RawCanisterId,
-    RawCanisterResult, RawCycles, RawEffectivePrincipal, RawMessageId, RawMockCanisterHttpResponse,
-    RawPrincipalId, RawSetStableMemory, RawStableMemory, RawSubmitIngressResult, RawSubnetId,
-    RawTime, RawVerifyCanisterSigArg, RawWasmResult, SubnetId, Topology,
+    RawCanisterResult, RawCycles, RawEffectivePrincipal, RawIngressStatusArgs, RawMessageId,
+    RawMockCanisterHttpResponse, RawPrincipalId, RawSetStableMemory, RawStableMemory,
+    RawSubmitIngressResult, RawSubnetId, RawTime, RawVerifyCanisterSigArg, RawWasmResult, SubnetId,
+    Topology,
 };
 use crate::management_canister::{
     CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterInstallModeUpgradeInner,
@@ -16,7 +17,7 @@ use crate::management_canister::{
     TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs, UploadChunkResult,
 };
 pub use crate::DefaultEffectiveCanisterIdError;
-use crate::{CallError, PocketIcBuilder, UserError, WasmResult};
+use crate::{CallError, IngressStatusResult, PocketIcBuilder, UserError, WasmResult};
 use backoff::backoff::Backoff;
 use backoff::{ExponentialBackoff, ExponentialBackoffBuilder};
 use candid::{
@@ -599,17 +600,31 @@ impl PocketIc {
     /// or a round has been executed due to a separate PocketIC library call.
     pub async fn ingress_status(
         &self,
-        message_id: RawMessageId,
-    ) -> Option<Result<WasmResult, UserError>> {
+        raw_message_id: RawMessageId,
+        caller: Option<Principal>,
+    ) -> IngressStatusResult {
         let endpoint = "read/ingress_status";
-        let result: Option<RawCanisterResult> = self.post(endpoint, message_id).await;
-        result.map(|result| match result {
-            RawCanisterResult::Ok(raw_wasm_result) => match raw_wasm_result {
-                RawWasmResult::Reply(data) => Ok(WasmResult::Reply(data)),
-                RawWasmResult::Reject(text) => Ok(WasmResult::Reject(text)),
-            },
-            RawCanisterResult::Err(user_error) => Err(user_error),
-        })
+        let raw_ingress_status_args = RawIngressStatusArgs {
+            raw_message_id,
+            raw_caller: caller.map(|caller| caller.into()),
+        };
+        match self.try_post(endpoint, raw_ingress_status_args).await {
+            Ok(None) => IngressStatusResult::NotAvailable,
+            Ok(Some(raw_result)) => {
+                let result = match raw_result {
+                    RawCanisterResult::Ok(raw_wasm_result) => match raw_wasm_result {
+                        RawWasmResult::Reply(data) => Ok(WasmResult::Reply(data)),
+                        RawWasmResult::Reject(text) => Ok(WasmResult::Reject(text)),
+                    },
+                    RawCanisterResult::Err(user_error) => Err(user_error),
+                };
+                IngressStatusResult::Success(result)
+            }
+            Err((status, message)) => {
+                assert_eq!(status, StatusCode::FORBIDDEN, "HTTP error code {} for PocketIc::ingress_status is not StatusCode::FORBIDDEN. This is a bug!", status);
+                IngressStatusResult::Forbidden(message)
+            }
+        }
     }
 
     /// Await an update call submitted previously by `submit_call` or `submit_call_with_effective_principal`.
@@ -625,7 +640,9 @@ impl PocketIc {
             .with_multiplier(2.0)
             .build();
         loop {
-            if let Some(ingress_status) = self.ingress_status(message_id.clone()).await {
+            if let IngressStatusResult::Success(ingress_status) =
+                self.ingress_status(message_id.clone(), None).await
+            {
                 break ingress_status;
             }
             tokio::time::sleep(retry_policy.next_backoff().unwrap()).await;
@@ -1361,12 +1378,29 @@ impl PocketIc {
         self.request(HttpMethod::Post, endpoint, body).await
     }
 
+    async fn try_post<T: DeserializeOwned, B: Serialize>(
+        &self,
+        endpoint: &str,
+        body: B,
+    ) -> Result<T, (StatusCode, String)> {
+        self.try_request(HttpMethod::Post, endpoint, body).await
+    }
+
     async fn request<T: DeserializeOwned, B: Serialize>(
         &self,
         http_method: HttpMethod,
         endpoint: &str,
         body: B,
     ) -> T {
+        self.try_request(http_method, endpoint, body).await.unwrap()
+    }
+
+    async fn try_request<T: DeserializeOwned, B: Serialize>(
+        &self,
+        http_method: HttpMethod,
+        endpoint: &str,
+        body: B,
+    ) -> Result<T, (StatusCode, String)> {
         // we may have to try several times if the instance is busy
         let start = std::time::SystemTime::now();
         loop {
@@ -1377,9 +1411,10 @@ impl PocketIc {
                 HttpMethod::Post => reqwest_client.post(url).json(&body),
             };
             let result = builder.send().await.expect("HTTP failure");
+            let status = result.status();
             match ApiResponse::<_>::from_response(result).await {
-                ApiResponse::Success(t) => break t,
-                ApiResponse::Error { message } => panic!("{}", message),
+                ApiResponse::Success(t) => break Ok(t),
+                ApiResponse::Error { message } => break Err((status, message)),
                 ApiResponse::Busy { state_label, op_id } => {
                     debug!(
                         "instance_id={} Instance is busy (with a different computation): state_label: {}, op_id: {}",
@@ -1404,24 +1439,31 @@ impl PocketIc {
                             .send()
                             .await
                             .expect("HTTP failure");
-                        match ApiResponse::<_>::from_response(result).await {
-                            ApiResponse::Error { message } => {
-                                debug!("Polling has not succeeded yet: {}", message)
-                            }
-                            ApiResponse::Success(t) => {
-                                return t;
-                            }
-                            ApiResponse::Started { state_label, op_id } => {
-                                warn!(
-                                    "instance_id={} unexpected Started({} {})",
-                                    self.instance_id, state_label, op_id
-                                );
-                            }
-                            ApiResponse::Busy { state_label, op_id } => {
-                                warn!(
-                                    "instance_id={} unexpected Busy({} {})",
-                                    self.instance_id, state_label, op_id
-                                );
+                        if result.status() == reqwest::StatusCode::NOT_FOUND {
+                            let message =
+                                String::from_utf8(result.bytes().await.unwrap().to_vec()).unwrap();
+                            debug!("Polling has not succeeded yet: {}", message);
+                        } else {
+                            let status = result.status();
+                            match ApiResponse::<_>::from_response(result).await {
+                                ApiResponse::Error { message } => {
+                                    return Err((status, message));
+                                }
+                                ApiResponse::Success(t) => {
+                                    return Ok(t);
+                                }
+                                ApiResponse::Started { state_label, op_id } => {
+                                    warn!(
+                                        "instance_id={} unexpected Started({} {})",
+                                        self.instance_id, state_label, op_id
+                                    );
+                                }
+                                ApiResponse::Busy { state_label, op_id } => {
+                                    warn!(
+                                        "instance_id={} unexpected Busy({} {})",
+                                        self.instance_id, state_label, op_id
+                                    );
+                                }
                             }
                         }
                         if let Some(max_request_time_ms) = self.max_request_time_ms {

packages/pocket-ic/tests/tests.rs

@@ -1,4 +1,7 @@
 use candid::{decode_one, encode_one, CandidType, Decode, Deserialize, Encode, Principal};
+use ic_certification::Label;
+use ic_transport_types::Envelope;
+use ic_transport_types::EnvelopeContent::ReadState;
 use pocket_ic::management_canister::{
     CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterSettings, EcdsaPublicKeyResult,
     HttpRequestResult, ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm,
@@ -9,11 +12,12 @@ use pocket_ic::{
         BlobCompression, CanisterHttpReply, CanisterHttpResponse, MockCanisterHttpResponse,
         RawEffectivePrincipal, SubnetKind,
     },
-    query_candid, update_candid, DefaultEffectiveCanisterIdError, ErrorCode, PocketIc,
-    PocketIcBuilder, WasmResult,
+    query_candid, update_candid, DefaultEffectiveCanisterIdError, ErrorCode, IngressStatusResult,
+    PocketIc, PocketIcBuilder, WasmResult,
 };
 #[cfg(unix)]
 use reqwest::blocking::Client;
+use serde::Serialize;
 use sha2::{Digest, Sha256};
 use std::{io::Read, time::SystemTime};
 
@@ -2009,25 +2013,91 @@ fn ingress_status() {
     pic.add_cycles(canister_id, INIT_CYCLES);
     pic.install_canister(canister_id, test_canister_wasm(), vec![], None);
 
+    let caller = Principal::from_slice(&[0xFF; 29]);
     let msg_id = pic
-        .submit_call(
-            canister_id,
-            Principal::anonymous(),
-            "whoami",
-            encode_one(()).unwrap(),
-        )
+        .submit_call(canister_id, caller, "whoami", encode_one(()).unwrap())
         .unwrap();
 
-    assert!(pic.ingress_status(msg_id.clone()).is_none());
+    match pic.ingress_status(msg_id.clone(), None) {
+        IngressStatusResult::NotAvailable => (),
+        status => panic!("Unexpected ingress status: {:?}", status),
+    }
+
+    // since the ingress status is not available, any caller can attempt to retrieve it
+    match pic.ingress_status(msg_id.clone(), Some(Principal::anonymous())) {
+        IngressStatusResult::NotAvailable => (),
+        status => panic!("Unexpected ingress status: {:?}", status),
+    }
 
     pic.tick();
 
-    let ingress_status = pic.ingress_status(msg_id).unwrap().unwrap();
-    let principal = match ingress_status {
+    let reply = match pic.ingress_status(msg_id.clone(), None) {
+        IngressStatusResult::Success(result) => result.unwrap(),
+        status => panic!("Unexpected ingress status: {:?}", status),
+    };
+    let principal = match reply {
         WasmResult::Reply(data) => Decode!(&data, String).unwrap(),
         WasmResult::Reject(err) => panic!("Unexpected reject: {}", err),
     };
     assert_eq!(principal, canister_id.to_string());
+
+    // now that the ingress status is available, the caller must match
+    let expected_err = "The user tries to access Request ID not signed by the caller.";
+    match pic.ingress_status(msg_id.clone(), Some(Principal::anonymous())) {
+        IngressStatusResult::Forbidden(msg) => assert_eq!(msg, expected_err,),
+        status => panic!("Unexpected ingress status: {:?}", status),
+    }
+
+    // confirm the behavior of read state requests
+    let resp = read_state_request_status(&pic, canister_id, msg_id.message_id.as_slice());
+    assert_eq!(resp.status(), reqwest::StatusCode::FORBIDDEN);
+    assert_eq!(
+        String::from_utf8(resp.bytes().unwrap().to_vec()).unwrap(),
+        expected_err
+    );
+}
+
+fn read_state_request_status(
+    pic: &PocketIc,
+    canister_id: Principal,
+    msg_id: &[u8],
+) -> reqwest::blocking::Response {
+    let path = vec!["request_status".into(), Label::from_bytes(msg_id)];
+    let paths = vec![path.clone()];
+    let content = ReadState {
+        ingress_expiry: pic
+            .get_time()
+            .duration_since(std::time::SystemTime::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64
+            + 240_000_000_000,
+        sender: Principal::anonymous(),
+        paths,
+    };
+    let envelope = Envelope {
+        content: std::borrow::Cow::Borrowed(&content),
+        sender_pubkey: None,
+        sender_sig: None,
+        sender_delegation: None,
+    };
+
+    let mut serialized_bytes = Vec::new();
+    let mut serializer = serde_cbor::Serializer::new(&mut serialized_bytes);
+    serializer.self_describe().unwrap();
+    envelope.serialize(&mut serializer).unwrap();
+
+    let endpoint = format!(
+        "instances/{}/api/v2/canister/{}/read_state",
+        pic.instance_id(),
+        canister_id.to_text()
+    );
+    let client = reqwest::blocking::Client::new();
+    client
+        .post(pic.get_server_url().join(&endpoint).unwrap())
+        .header(reqwest::header::CONTENT_TYPE, "application/cbor")
+        .body(serialized_bytes)
+        .send()
+        .unwrap()
 }
 
 #[test]

rs/pocket_ic_server/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - New endpoint `/instances/<instance_id>/read/ingress_status` to fetch the status of an update call submitted through an ingress message.
+  If an optional caller is provided and a corresponding read state request for the status of the same update call
+  signed by that specified caller was rejected because the update call was submitted by a different caller, then an error is returned.
 
 ### Fixed
 - Canisters created via `provisional_create_canister_with_cycles` with the management canister ID as the effective canister ID