feat: impersonating sender of requests to a local PocketIC instance

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 # UNRELEASED
 
+### feat: impersonating sender of requests to a local PocketIC instance
+
+`dfx canister call`, `dfx canister status`, and `dfx canister update-settings` take
+an additional CLI argument `--impersonate` to specify a principal
+on behalf of which requests to a local PocketIC instance are sent.
+
 ### fix: `dfx deploy --by-proposal` no longer sends chunk data in ProposeCommitBatch
 
 Recently we made `dfx deploy` include some chunk data in CommitBatch, in order to streamline

e2e/tests-dfx/call.bash

@@ -298,3 +298,82 @@ teardown() {
   assert_command dfx canister call inter2_mo read
   assert_match '(8 : nat)'
 }
+
+@test "impersonate sender" {
+    [[ ! "$USE_POCKETIC" ]] && skip "skipped for replica: impersonating sender is only supported for PocketIC"
+
+    dfx_start
+    assert_command dfx deploy hello_backend
+    CANISTER_ID="$(dfx canister id hello_backend)"
+
+    # set the management canister as the only controller
+    assert_command dfx canister update-settings hello_backend --set-controller aaaaa-aa --yes
+
+    # updating settings now fails because the default identity does not control the canister anymore
+    assert_command_fail dfx canister update-settings hello_backend --freezing-threshold 0
+    assert_contains "Only controllers of canister $CANISTER_ID can call ic00 method update_settings"
+
+    # updating settings succeeds when impersonating the management canister as the sender
+    assert_command dfx canister update-settings hello_backend --freezing-threshold 0 --impersonate aaaaa-aa
+
+    # test management canister call failure (setting memory allocation to a low value)
+    assert_command_fail dfx canister update-settings hello_backend --memory-allocation 1 --impersonate aaaaa-aa
+    assert_contains "Management canister call failed: IC0402: Canister was given 1 B memory allocation but at least"
+
+    # canister status fails because the default identity does not control the canister anymore
+    assert_command_fail dfx canister status hello_backend
+    assert_contains "Only controllers of canister $CANISTER_ID can call ic00 method canister_status"
+
+    # canister status succeeds when impersonating the management canister as the sender
+    assert_command dfx canister status hello_backend --impersonate aaaaa-aa
+    assert_contains "Controllers: aaaaa-aa"
+    assert_contains "Freezing threshold: 0"
+
+    # freeze the canister
+    assert_command dfx canister update-settings hello_backend --freezing-threshold 9223372036854775808 --confirm-very-long-freezing-threshold --impersonate aaaaa-aa
+
+    # test management canister call submission failure
+    assert_command_fail dfx canister status hello_backend --impersonate aaaaa-aa
+    assert_contains "Failed to submit management canister call: IC0207: Canister $CANISTER_ID is out of cycles"
+
+    # test update call submission failure
+    assert_command_fail dfx canister call aaaaa-aa canister_status "(record { canister_id=principal\"$CANISTER_ID\" })" --update --impersonate aaaaa-aa
+    assert_contains "Failed to submit canister call: IC0207: Canister $CANISTER_ID is out of cycles"
+
+    # test async call submission failure
+    assert_command_fail dfx canister call aaaaa-aa canister_status "(record { canister_id=principal\"$CANISTER_ID\" })" --async --impersonate aaaaa-aa
+    assert_contains "Failed to submit canister call: IC0207: Canister $CANISTER_ID is out of cycles"
+
+    # unfreeze the canister
+    assert_command dfx canister update-settings hello_backend --freezing-threshold 0 --impersonate aaaaa-aa
+
+    # test update call failure
+    assert_command_fail dfx canister call aaaaa-aa delete_canister "(record { canister_id=principal\"$CANISTER_ID\" })" --update --impersonate aaaaa-aa
+    assert_contains "Canister call failed: IC0510: Canister $CANISTER_ID must be stopped before it is deleted."
+
+    # test update call
+    assert_command dfx canister call aaaaa-aa start_canister "(record { canister_id=principal\"$CANISTER_ID\" })" --update --impersonate aaaaa-aa
+    assert_contains "()"
+
+    # test async call
+    assert_command dfx canister call aaaaa-aa canister_status "(record { canister_id=principal\"$CANISTER_ID\" })" --async --impersonate aaaaa-aa
+    assert_contains "Request ID:"
+
+    # test request status failure
+    RID=$(dfx canister call aaaaa-aa delete_canister "(record { canister_id=principal\"$CANISTER_ID\" })" --async --impersonate aaaaa-aa)
+    assert_command_fail dfx canister request-status "$RID" hello_backend
+    assert_contains "Canister call failed: IC0510: Canister $CANISTER_ID must be stopped before it is deleted."
+
+    # test request status
+    RID=$(dfx canister call aaaaa-aa canister_status "(record { canister_id=principal\"$CANISTER_ID\" })" --async --impersonate aaaaa-aa)
+    assert_command dfx canister request-status "$RID" hello_backend
+    assert_contains "record {"
+
+    # test query call failure
+    assert_command_fail dfx canister call aaaaa-aa fetch_canister_logs "(record { canister_id=principal\"$CANISTER_ID\" })" --query --impersonate "$CANISTER_ID"
+    assert_contains "Failed to perform query call: IC0406: Caller $CANISTER_ID is not allowed to query ic00 method fetch_canister_logs"
+
+    # test query call
+    assert_command dfx canister call aaaaa-aa fetch_canister_logs "(record { canister_id=principal\"$CANISTER_ID\" })" --query --impersonate aaaaa-aa
+    assert_contains "(record { 1_754_302_831 = vec {} })"
+}

src/dfx-core/src/canister/mod.rs

@@ -71,6 +71,9 @@ YOU WILL LOSE ALL DATA IN THE CANISTER.
                 .await
                 .map_err(CanisterInstallError::InstallWasmError)
         }
+        CallSender::Impersonate(_) => {
+            unreachable!("Impersonating sender when installing canisters is not supported.")
+        }
         CallSender::Wallet(wallet_id) => {
             let wallet = build_wallet_canister(*wallet_id, agent).await?;
             let install_args = CanisterInstall {

src/dfx-core/src/config/model/local_server_descriptor.rs

@@ -360,23 +360,38 @@ impl LocalServerDescriptor {
         logger: Option<&Logger>,
     ) -> Result<Option<u16>, NetworkConfigError> {
         let replica_port_path = self.replica_port_path();
-        let pocketic_port_path = self.pocketic_port_path();
         match read_port_from(&replica_port_path)? {
             Some(port) => {
                 if let Some(logger) = logger {
                     info!(logger, "Found local replica running on port {}", port);
                 }
                 Ok(Some(port))
             }
-            None => match read_port_from(&pocketic_port_path)? {
-                Some(port) => {
-                    if let Some(logger) = logger {
-                        info!(logger, "Found local PocketIC running on port {}", port);
-                    }
-                    Ok(Some(port))
+            None => {
+                let port = self
+                    .get_running_pocketic_port(logger)?
+                    .or(self.replica.port);
+                Ok(port)
+            }
+        }
+    }
+
+    /// Gets the port of a local PocketIC instance.
+    ///
+    /// # Prerequisites
+    /// - A local PocketIC instance needs to be running, e.g. with `dfx start --pocketic`.
+    pub fn get_running_pocketic_port(
+        &self,
+        logger: Option<&Logger>,
+    ) -> Result<Option<u16>, NetworkConfigError> {
+        match read_port_from(&self.pocketic_port_path())? {
+            Some(port) => {
+                if let Some(logger) = logger {
+                    info!(logger, "Found local PocketIC running on port {}", port);
                 }
-                None => Ok(self.replica.port),
-            },
+                Ok(Some(port))
+            }
+            None => Ok(None),
         }
     }
 }

src/dfx-core/src/identity/mod.rs

@@ -307,6 +307,7 @@ impl AsRef<Identity> for Identity {
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub enum CallSender {
     SelectedId,
+    Impersonate(Principal),
     Wallet(Principal),
 }
 

src/dfx/Cargo.toml

@@ -86,6 +86,7 @@ os_str_bytes = { version = "6.3.0", features = ["conversions"] }
 patch = "0.7.0"
 pem.workspace = true
 petgraph = "0.6.0"
+pocket-ic = { git = "https://github.com/dfinity/ic", rev = "3e24396441e4c7380928d4e8b4ccff7de77d0e7e" }
 rand = "0.8.5"
 regex = "1.5.5"
 reqwest = { workspace = true, features = ["blocking", "json"] }
@@ -125,9 +126,6 @@ ci_info = "0.14"
 [target.'cfg(windows)'.dependencies]
 junction = "1.0.0"
 
-[target.'cfg(unix)'.dependencies]
-pocket-ic = { git = "https://github.com/dfinity/ic", rev = "3e24396441e4c7380928d4e8b4ccff7de77d0e7e" }
-
 [dev-dependencies]
 env_logger = "0.10"
 proptest = "1.0"

src/dfx/src/commands/canister/call.rs

@@ -6,6 +6,7 @@ use crate::lib::root_key::fetch_root_key_if_needed;
 use crate::util::clap::argument_from_cli::ArgumentFromCliPositionalOpt;
 use crate::util::clap::parsers::cycle_amount_parser;
 use crate::util::{blob_from_arguments, fetch_remote_did_file, get_candid_type, print_idl_blob};
+use anyhow::bail;
 use anyhow::{anyhow, Context};
 use candid::Principal as CanisterId;
 use candid::{CandidType, Decode, Deserialize, Principal};
@@ -14,11 +15,14 @@ use clap::Parser;
 use dfx_core::canister::build_wallet_canister;
 use dfx_core::identity::CallSender;
 use ic_agent::agent::CallResponse;
+use ic_agent::RequestId;
 use ic_utils::canister::Argument;
 use ic_utils::interfaces::management_canister::builders::{CanisterInstall, CanisterSettings};
 use ic_utils::interfaces::management_canister::MgmtMethod;
 use ic_utils::interfaces::wallet::{CallForwarder, CallResult};
 use ic_utils::interfaces::WalletCanister;
+use pocket_ic::common::rest::RawEffectivePrincipal;
+use pocket_ic::WasmResult;
 use slog::warn;
 use std::option::Option;
 use std::path::PathBuf;
@@ -83,6 +87,11 @@ pub struct CanisterCallOpts {
         conflicts_with("random")
     )]
     always_assist: bool,
+
+    /// Send request on behalf of the specified principal.
+    /// This option only works for a local PocketIC instance.
+    #[arg(long)]
+    impersonate: Option<Principal>,
 }
 
 #[derive(Clone, CandidType, Deserialize, Debug)]
@@ -205,8 +214,13 @@ pub fn get_effective_canister_id(
 pub async fn exec(
     env: &dyn Environment,
     opts: CanisterCallOpts,
-    call_sender: &CallSender,
+    mut call_sender: &CallSender,
 ) -> DfxResult {
+    let call_sender_override = opts.impersonate.map(CallSender::Impersonate);
+    if let Some(ref call_sender_override) = call_sender_override {
+        call_sender = call_sender_override;
+    };
+
     let agent = env.get_agent();
     fetch_root_key_if_needed(env).await?;
 
@@ -334,6 +348,29 @@ To figure out the id of your wallet, run 'dfx identity get-wallet (--network ic)
                     .with_arg(arg_value);
                 query_builder.call().await.context("Failed query call.")?
             }
+            CallSender::Impersonate(sender) => {
+                let pocketic = env.get_pocketic();
+                if let Some(pocketic) = pocketic {
+                    let res = pocketic
+                        .query_call_with_effective_principal(
+                            canister_id,
+                            RawEffectivePrincipal::CanisterId(
+                                effective_canister_id.as_slice().to_vec(),
+                            ),
+                            *sender,
+                            method_name,
+                            arg_value,
+                        )
+                        .await
+                        .map_err(|err| anyhow!("Failed to perform query call: {}", err))?;
+                    match res {
+                        WasmResult::Reply(data) => data,
+                        WasmResult::Reject(err) => bail!("Canister rejected: {}", err),
+                    }
+                } else {
+                    bail!("Impersonating sender is only supported for a local PocketIC instance.")
+                }
+            }
             CallSender::Wallet(wallet_id) => {
                 let wallet = build_wallet_canister(*wallet_id, agent).await?;
                 do_wallet_call(
@@ -360,6 +397,27 @@ To figure out the id of your wallet, run 'dfx identity get-wallet (--network ic)
                 .await
                 .context("Failed update call.")?
                 .map(|(res, _)| res),
+            CallSender::Impersonate(sender) => {
+                let pocketic = env.get_pocketic();
+                if let Some(pocketic) = pocketic {
+                    let msg_id = pocketic
+                        .submit_call_with_effective_principal(
+                            canister_id,
+                            RawEffectivePrincipal::CanisterId(
+                                effective_canister_id.as_slice().to_vec(),
+                            ),
+                            *sender,
+                            method_name,
+                            arg_value,
+                        )
+                        .await
+                        .map_err(|err| anyhow!("Failed to submit canister call: {}", err))?
+                        .message_id;
+                    CallResponse::Poll(RequestId::new(msg_id.as_slice().try_into().unwrap()))
+                } else {
+                    bail!("Impersonating sender is only supported for a local PocketIC instance.")
+                }
+            }
             CallSender::Wallet(wallet_id) => {
                 let wallet = build_wallet_canister(*wallet_id, agent).await?;
                 let mut args = Argument::default();
@@ -388,6 +446,33 @@ To figure out the id of your wallet, run 'dfx identity get-wallet (--network ic)
                 .with_arg(arg_value)
                 .await
                 .context("Failed update call.")?,
+            CallSender::Impersonate(sender) => {
+                let pocketic = env.get_pocketic();
+                if let Some(pocketic) = pocketic {
+                    let msg_id = pocketic
+                        .submit_call_with_effective_principal(
+                            canister_id,
+                            RawEffectivePrincipal::CanisterId(
+                                effective_canister_id.as_slice().to_vec(),
+                            ),
+                            *sender,
+                            method_name,
+                            arg_value,
+                        )
+                        .await
+                        .map_err(|err| anyhow!("Failed to submit canister call: {}", err))?;
+                    let res = pocketic
+                        .await_call_no_ticks(msg_id)
+                        .await
+                        .map_err(|err| anyhow!("Canister call failed: {}", err))?;
+                    match res {
+                        WasmResult::Reply(data) => data,
+                        WasmResult::Reject(err) => bail!("Canister rejected: {}", err),
+                    }
+                } else {
+                    bail!("Impersonating sender is only supported for a local PocketIC instance.")
+                }
+            }
             CallSender::Wallet(wallet_id) => {
                 let wallet = build_wallet_canister(*wallet_id, agent).await?;
                 do_wallet_call(

src/dfx/src/commands/canister/delete.rs

@@ -127,6 +127,11 @@ async fn delete_canister(
                     CallSender::Wallet(wallet_id) => WithdrawTarget::Canister {
                         canister_id: *wallet_id,
                     },
+                    CallSender::Impersonate(_) => {
+                        unreachable!(
+                            "Impersonating sender when deleting canisters is not supported."
+                        )
+                    }
                     CallSender::SelectedId => {
                         let network = env.get_network_descriptor();
                         let identity_name = env

src/dfx/src/commands/canister/deposit_cycles.rs

@@ -67,6 +67,9 @@ async fn deposit_cycles(
             )
             .await?;
         }
+        CallSender::Impersonate(_) => {
+            unreachable!("Impersonating sender when depositing cycles is not supported.")
+        }
         CallSender::Wallet(_) => {
             canister::deposit_cycles(env, canister_id, call_sender, cycles).await?
         }