test(sns-cli): Port UpgradeSnsControlledCanister with Large Wasm integration test to use SNS CLI with PocketIc (dfinity#3696)

This PR leverages the SNS CLI functionality to prepare and submit
UpgradeSnsControlledCanister proposals for large Wasms as part of the
existing integration test scenario. The goal is to have the CLI
functionality tested in CI. This entails the need to generalize the CLI
command upgrade-sns-controlled-canister to work with a CallCanister
instance which can be instantiated not only with ic-agent, but also with
PocketIc instances.

Author: aterga

Cargo.lock

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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).
 - The function `PocketIc::set_certified_time` to set the current certified time on all subnets of the PocketIC instance.
+- The function `PocketIc::update_call_with_effective_principal` is made public. It is helpful, e.g., for
+modeling management canister calls that need to be routed to the right subnet using effective principals.
 
 ### Changed
 - The response types `pocket_ic::WasmResult`, `pocket_ic::UserError`, and `pocket_ic::CallError` are replaced by a single reject response type `pocket_ic::RejectResponse`.

packages/pocket-ic/CHANGELOG.md

@@ -1129,7 +1129,7 @@ impl PocketIc {
         runtime.block_on(async { self.pocket_ic.get_subnet_metrics(subnet_id).await })
     }
 
-    fn update_call_with_effective_principal(
+    pub fn update_call_with_effective_principal(
         &self,
         canister_id: CanisterId,
         effective_principal: RawEffectivePrincipal,

packages/pocket-ic/src/lib.rs

@@ -1529,7 +1529,7 @@ impl PocketIc {
         result.into()
     }
 
-    pub(crate) async fn update_call_with_effective_principal(
+    pub async fn update_call_with_effective_principal(
         &self,
         canister_id: CanisterId,
         effective_principal: RawEffectivePrincipal,

packages/pocket-ic/src/nonblocking.rs

@@ -19,7 +19,9 @@ DEPENDENCIES = [
     "@crate_index//:anyhow",
     "@crate_index//:candid",
     "@crate_index//:ic-agent",
+    "@crate_index//:itertools",
     "@crate_index//:serde",
+    "@crate_index//:serde_cbor",
     "@crate_index//:tempfile",
     "@crate_index//:thiserror",
     "@crate_index//:tokio",

rs/nervous_system/agent/BUILD.bazel

@@ -23,7 +23,9 @@ pocket-ic = { path = "../../../packages/pocket-ic" }
 registry-canister = { path = "../../registry/canister" }
 ic-sns-root = { path = "../../sns/root" }
 ic-sns-swap = { path = "../../sns/swap" }
+itertools = { workspace = true }
 serde = { workspace = true }
+serde_cbor = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }

rs/nervous_system/agent/Cargo.toml

@@ -1,24 +1,31 @@
-use crate::Request;
+use crate::CallCanisters;
+use crate::{CanisterInfo, Request};
 use candid::Principal;
 use ic_agent::Agent;
+use itertools::{Either, Itertools};
+use serde_cbor::Value;
+use std::collections::BTreeSet;
 use thiserror::Error;
 
-use crate::CallCanisters;
-
 #[derive(Error, Debug)]
 pub enum AgentCallError {
+    #[error("agent identity error: {0}")]
+    Identity(String),
     #[error("agent error: {0}")]
     Agent(#[from] ic_agent::AgentError),
     #[error("canister request could not be encoded: {0}")]
     CandidEncode(candid::Error),
     #[error("canister did not respond with the expected response type: {0}")]
     CandidDecode(candid::Error),
+    #[error("invalid canister controllers: {0}")]
+    CanisterControllers(String),
 }
 
 impl crate::sealed::Sealed for Agent {}
 
 impl CallCanisters for Agent {
     type Error = AgentCallError;
+
     async fn call<R: Request>(
         &self,
         canister_id: impl Into<Principal> + Send,
@@ -50,4 +57,86 @@ impl CallCanisters for Agent {
             candid::decode_one(response.as_slice()).map_err(AgentCallError::CandidDecode)?;
         Ok(response)
     }
+
+    async fn canister_info(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+    ) -> Result<CanisterInfo, Self::Error> {
+        let canister_id = canister_id.into();
+
+        let read_state_result = self
+            .read_state_canister_info(canister_id, "module_hash")
+            .await;
+
+        let module_hash = match read_state_result {
+            Ok(module_hash) => Some(module_hash),
+            Err(ic_agent::AgentError::LookupPathAbsent(_)) => {
+                // https://internetcomputer.org/docs/current/references/ic-interface-spec#state-tree-canister-information
+                None
+            }
+            Err(err) => {
+                return Err(Self::Error::Agent(err));
+            }
+        };
+
+        let controllers_blob = self
+            .read_state_canister_info(canister_id, "controllers")
+            .await
+            .map_err(AgentCallError::Agent)?;
+
+        let cbor: Value = serde_cbor::from_slice(&controllers_blob).map_err(|err| {
+            Self::Error::CanisterControllers(format!("Failed decoding CBOR data: {:?}", err))
+        })?;
+
+        let Value::Array(controllers) = cbor else {
+            return Err(Self::Error::CanisterControllers(format!(
+                "Expected controllers to be an array, but got {:?}",
+                cbor
+            )));
+        };
+
+        let (controllers, errors): (Vec<_>, Vec<_>) =
+            controllers.into_iter().partition_map(|value| {
+                let Value::Bytes(bytes) = value else {
+                    let err = format!(
+                        "Expected canister controller to be of type bytes, got {:?}",
+                        value
+                    );
+                    return Either::Right(err);
+                };
+                match Principal::try_from(&bytes) {
+                    Err(err) => {
+                        let err =
+                            format!("Cannot interpret canister controller principal: {}", err);
+                        Either::Right(err)
+                    }
+                    Ok(principal) => Either::Left(principal),
+                }
+            });
+
+        if !errors.is_empty() {
+            return Err(Self::Error::CanisterControllers(format!(
+                "\n  - {}",
+                errors.join("\n  - ")
+            )));
+        }
+
+        let unique_controllers = BTreeSet::from_iter(controllers.iter().copied());
+
+        if unique_controllers.len() != controllers.len() {
+            return Err(Self::Error::CanisterControllers(format!(
+                "Canister controllers have duplicates: {}",
+                controllers.into_iter().join(", ")
+            )));
+        }
+
+        Ok(CanisterInfo {
+            module_hash,
+            controllers: unique_controllers,
+        })
+    }
+
+    fn caller(&self) -> Result<Principal, Self::Error> {
+        self.get_principal().map_err(Self::Error::Identity)
+    }
 }

rs/nervous_system/agent/src/agent_impl.rs

@@ -1,14 +1,15 @@
+use candid::{CandidType, Principal};
+use serde::de::DeserializeOwned;
+use std::collections::BTreeSet;
+use std::fmt::Display;
+
 pub mod agent_impl;
 pub mod management_canister;
 pub mod nns;
 mod null_request;
 pub mod pocketic_impl;
 pub mod sns;
 
-use candid::{CandidType, Principal};
-use serde::de::DeserializeOwned;
-use std::fmt::Display;
-
 // This is used to "seal" the CallCanisters trait so that it cannot be implemented outside of this crate.
 // This is useful because it means we can modify the trait in the future without worrying about
 // breaking backwards compatibility with implementations outside of this crate.
@@ -27,11 +28,24 @@ pub trait Request: Send {
     type Response: CandidType + DeserializeOwned;
 }
 
+pub struct CanisterInfo {
+    pub module_hash: Option<Vec<u8>>,
+    pub controllers: BTreeSet<Principal>,
+}
+
 pub trait CallCanisters: sealed::Sealed {
     type Error: Display + Send + std::error::Error + 'static;
+
+    fn caller(&self) -> Result<Principal, Self::Error>;
+
     fn call<R: Request>(
         &self,
         canister_id: impl Into<Principal> + Send,
         request: R,
     ) -> impl std::future::Future<Output = Result<R::Response, Self::Error>> + Send;
+
+    fn canister_info(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+    ) -> impl std::future::Future<Output = Result<CanisterInfo, Self::Error>> + Send;
 }