feat: add set controller to the management canister

ic-utils/src/interfaces/management_canister.rs

@@ -326,4 +326,24 @@ impl<'agent> Canister<'agent, ManagementCanister> {
     ) -> InstallCodeBuilder<'agent, 'canister, ManagementCanister> {
         InstallCodeBuilder::builder(self, canister_id, wasm)
     }
+
+    /// Set controller for a canister.
+    pub fn set_controller<'canister: 'agent>(
+        &'canister self,
+        canister_id: &Principal,
+        new_controller: &Principal,
+    ) -> impl 'agent + AsyncCall<()> {
+        #[derive(CandidType)]
+        struct Argument {
+            canister_id: Principal,
+            new_controller: Principal,
+        }
+
+        self.update_("set_controller")
+            .with_arg(Argument {
+                canister_id: canister_id.clone(),
+                new_controller: new_controller.clone(),
+            })
+            .build()
+    }
 }

ref-tests/src/utils.rs

@@ -25,7 +25,7 @@ pub async fn create_identity() -> Result<BasicIdentity, String> {
     ))
 }
 
-pub async fn create_agent() -> Result<Agent, String> {
+pub async fn create_agent(identity: BasicIdentity) -> Result<Agent, String> {
     let port_env = std::env::var("IC_REF_PORT")
         .expect("Need to specify the IC_REF_PORT environment variable.");
     let port = port_env
@@ -34,7 +34,7 @@ pub async fn create_agent() -> Result<Agent, String> {
 
     Agent::builder()
         .with_url(format!("http://127.0.0.1:{}", port))
-        .with_identity(create_identity().await?)
+        .with_identity(identity)
         .build()
         .map_err(|e| format!("{:?}", e))
 }
@@ -46,7 +46,12 @@ where
 {
     let mut runtime = tokio::runtime::Runtime::new().expect("Could not create tokio runtime.");
     runtime.block_on(async {
-        let agent = create_agent().await.expect("Could not create an agent.");
+        let agent_identity = create_identity()
+            .await
+            .expect("Could not create an identity.");
+        let agent = create_agent(agent_identity)
+            .await
+            .expect("Could not create an agent.");
         match f(agent).await {
             Ok(_) => {}
             Err(e) => assert!(false, "{:?}", e),
@@ -63,7 +68,12 @@ where
     match runtime.block_on(async {
         let canister_wasm = universal_canister::wasm();
 
-        let agent = create_agent().await.expect("Could not create an agent.");
+        let agent_identity = create_identity()
+            .await
+            .expect("Could not create an identity.");
+        let agent = create_agent(agent_identity)
+            .await
+            .expect("Could not create an agent.");
         let ic00 = ManagementCanister::create(&agent);
 
         let (canister_id,) = ic00

ref-tests/tests/ic-ref.rs

@@ -38,10 +38,11 @@ fn spec_compliance_claimed() {
 
 mod management_canister {
     use ic_agent::AgentError;
+    use ic_agent::Identity;
     use ic_utils::call::AsyncCall;
     use ic_utils::interfaces::management_canister::{CanisterStatus, InstallMode};
     use ic_utils::interfaces::ManagementCanister;
-    use ref_tests::{create_agent, create_waiter, with_agent};
+    use ref_tests::{create_agent, create_identity, create_waiter, with_agent};
 
     mod create_canister {
         use super::{create_waiter, with_agent};
@@ -130,7 +131,9 @@ mod management_canister {
                 .await?;
 
             // Each agent has their own identity.
-            let other_agent = create_agent().await?;
+            let other_agent_identity = create_identity().await?;
+            let other_agent_principal = other_agent_identity.sender()?;
+            let other_agent = create_agent(other_agent_identity).await?;
             let other_ic00 = ManagementCanister::create(&other_agent);
 
             // Reinstall with another agent should fail.
@@ -162,7 +165,29 @@ mod management_canister {
             });
 
             // Change controller.
-            // TODO: set controller tests.
+            ic00.set_controller(&canister_id, &other_agent_principal)
+                .call_and_wait(create_waiter())
+                .await?;
+
+            // Change controller with wrong controller should fail
+            let result = ic00
+                .set_controller(&canister_id, &other_agent_principal)
+                .call_and_wait(create_waiter())
+                .await;
+            assert!(match result {
+                Err(AgentError::ReplicaError {
+                    reject_code: 5,
+                    reject_message,
+                }) if reject_message.contains("is not authorized to manage canister") => true,
+                _ => false,
+            });
+
+            // Reinstall as new controller
+            other_ic00
+                .install_code(&canister_id, &canister_wasm)
+                .with_mode(InstallMode::Reinstall)
+                .call_and_wait(create_waiter())
+                .await?;
 
             // Reinstall on empty should succeed.
             let (canister_id_2,) = ic00
@@ -378,6 +403,84 @@ mod management_canister {
             Ok(())
         })
     }
+
+    #[ignore]
+    #[test]
+    fn canister_lifecycle_as_wrong_controller() {
+        with_agent(|agent| async move {
+            let ic00 = ManagementCanister::create(&agent);
+            let (canister_id,) = ic00
+                .create_canister()
+                .call_and_wait(create_waiter())
+                .await?;
+            let canister_wasm = b"\0asm\x01\0\0\0".to_vec();
+
+            // Install once.
+            ic00.install_code(&canister_id, &canister_wasm)
+                .with_mode(InstallMode::Install)
+                .call_and_wait(create_waiter())
+                .await?;
+
+            // Create another agent with different identity.
+            let other_agent_identity = create_identity().await?;
+            let other_agent = create_agent(other_agent_identity).await?;
+            let other_ic00 = ManagementCanister::create(&other_agent);
+
+            // Start as a wrong controller should fail.
+            let result = other_ic00
+                .start_canister(&canister_id)
+                .call_and_wait(create_waiter())
+                .await;
+            assert!(match result {
+                Err(AgentError::ReplicaError {
+                    reject_code: 5,
+                    reject_message,
+                }) if reject_message.contains("is not authorized to manage canister") => true,
+                _ => false,
+            });
+
+            // Stop as a wrong controller should fail.
+            let result = other_ic00
+                .stop_canister(&canister_id)
+                .call_and_wait(create_waiter())
+                .await;
+            assert!(match result {
+                Err(AgentError::ReplicaError {
+                    reject_code: 5,
+                    reject_message,
+                }) if reject_message.contains("is not authorized to manage canister") => true,
+                _ => false,
+            });
+
+            // Get canister status as a wrong controller should fail.
+            let result = other_ic00
+                .canister_status(&canister_id)
+                .call_and_wait(create_waiter())
+                .await;
+            assert!(match result {
+                Err(AgentError::ReplicaError {
+                    reject_code: 5,
+                    reject_message,
+                }) if reject_message.contains("is not authorized to manage canister") => true,
+                _ => false,
+            });
+
+            // Delete as a wrong controller should fail.
+            let result = other_ic00
+                .delete_canister(&canister_id)
+                .call_and_wait(create_waiter())
+                .await;
+            assert!(match result {
+                Err(AgentError::ReplicaError {
+                    reject_code: 5,
+                    reject_message,
+                }) if reject_message.contains("is not authorized to manage canister") => true,
+                _ => false,
+            });
+
+            Ok(())
+        })
+    }
 }
 
 mod simple_calls {