feat: add canister management to ic_oss_cluster

src/ic_oss_cluster/ic_oss_cluster.did

@@ -1,3 +1,4 @@
+type AddWasmInput = record { kind : nat8; wasm : blob; description : text };
 type ChainArgs = variant { Upgrade : UpgradeArgs; Init : InitArgs };
 type ClusterInfo = record {
   ecdsa_token_public_key : text;
@@ -6,6 +7,11 @@ type ClusterInfo = record {
   name : text;
   token_expiration : nat64;
 };
+type DeployWasmInput = record {
+  args : opt blob;
+  kind : nat8;
+  canister : principal;
+};
 type InitArgs = record {
   ecdsa_key_name : text;
   name : text;
@@ -22,10 +28,14 @@ type Token = record {
 type UpgradeArgs = record { name : opt text; token_expiration : opt nat64 };
 service : (opt ChainArgs) -> {
   access_token : (principal) -> (Result);
+  admin_add_wasm : (AddWasmInput, opt blob) -> (Result_1);
   admin_attach_policies : (Token) -> (Result_1);
   admin_detach_policies : (Token) -> (Result_1);
+  admin_install_wasm : (DeployWasmInput, opt bool) -> (Result_1);
   admin_set_managers : (vec principal) -> (Result_1);
   admin_sign_access_token : (Token) -> (Result);
   get_cluster_info : () -> (Result_2) query;
+  validate_admin_add_wasm : (AddWasmInput, opt blob) -> (Result_1);
+  validate_admin_install_wasm : (DeployWasmInput, opt bool) -> (Result_1);
   validate_admin_set_managers : (vec principal) -> (Result_1);
 }

src/ic_oss_cluster/src/api_admin.rs

@@ -1,13 +1,22 @@
 use candid::Principal;
+use ic_cdk::api::management_canister::main::*;
 use ic_oss_cose::{
     cose_sign1, coset::CborSerializable, sha256, Token as CoseToken, BUCKET_TOKEN_AAD,
     CLUSTER_TOKEN_AAD, ES256K,
 };
-use ic_oss_types::{bucket::Token, permission::Policies};
+use ic_oss_types::{
+    bucket::Token,
+    cluster::{AddWasmInput, DeployWasmInput},
+    format_error,
+    permission::Policies,
+    ByteN,
+};
 use serde_bytes::ByteBuf;
 use std::collections::BTreeSet;
 
-use crate::{ecdsa, is_controller, store, ANONYMOUS, SECONDS};
+use crate::{
+    ecdsa, is_controller, is_controller_or_manager, store, ANONYMOUS, MILLISECONDS, SECONDS,
+};
 
 #[ic_cdk::update(guard = "is_controller")]
 fn admin_set_managers(args: BTreeSet<Principal>) -> Result<(), String> {
@@ -29,12 +38,8 @@ fn validate_admin_set_managers(args: BTreeSet<Principal>) -> Result<(), String>
     Ok(())
 }
 
-#[ic_cdk::update]
+#[ic_cdk::update(guard = "is_controller_or_manager")]
 async fn admin_sign_access_token(token: Token) -> Result<ByteBuf, String> {
-    if !store::state::is_manager(&ic_cdk::caller()) {
-        Err("user is not a manager".to_string())?;
-    }
-
     let now_sec = ic_cdk::api::time() / SECONDS;
     let (ecdsa_key_name, token_expiration) =
         store::state::with(|r| (r.ecdsa_key_name.clone(), r.token_expiration));
@@ -55,24 +60,128 @@ async fn admin_sign_access_token(token: Token) -> Result<ByteBuf, String> {
     Ok(ByteBuf::from(token))
 }
 
-#[ic_cdk::update]
+#[ic_cdk::update(guard = "is_controller_or_manager")]
 async fn admin_attach_policies(args: Token) -> Result<(), String> {
-    if !store::state::is_manager(&ic_cdk::caller()) {
-        Err("user is not a manager".to_string())?;
-    }
-
     let policies = Policies::try_from(args.policies.as_str())?;
     store::auth::attach_policies(args.subject, args.audience, policies);
     Ok(())
 }
 
-#[ic_cdk::update]
+#[ic_cdk::update(guard = "is_controller_or_manager")]
 async fn admin_detach_policies(args: Token) -> Result<(), String> {
-    if !store::state::is_manager(&ic_cdk::caller()) {
-        Err("user is not a manager".to_string())?;
-    }
-
     let policies = Policies::try_from(args.policies.as_str())?;
     store::auth::detach_policies(args.subject, args.audience, policies);
     Ok(())
 }
+
+#[ic_cdk::update(guard = "is_controller_or_manager")]
+async fn admin_add_wasm(
+    args: AddWasmInput,
+    force_prev_hash: Option<ByteN<32>>,
+) -> Result<(), String> {
+    store::wasm::add_wasm(
+        ic_cdk::caller(),
+        ic_cdk::api::time() / MILLISECONDS,
+        args,
+        force_prev_hash,
+        false,
+    )
+}
+
+#[ic_cdk::update]
+async fn validate_admin_add_wasm(
+    args: AddWasmInput,
+    force_prev_hash: Option<ByteN<32>>,
+) -> Result<(), String> {
+    store::wasm::add_wasm(
+        ic_cdk::caller(),
+        ic_cdk::api::time() / MILLISECONDS,
+        args,
+        force_prev_hash,
+        true,
+    )
+}
+
+#[ic_cdk::update(guard = "is_controller_or_manager")]
+async fn admin_install_wasm(args: DeployWasmInput, reinstall: Option<bool>) -> Result<(), String> {
+    let (info,) = canister_info(CanisterInfoRequest {
+        canister_id: args.canister,
+        num_requested_changes: None,
+    })
+    .await
+    .map_err(format_error)?;
+    let id = ic_cdk::id();
+    if !info.controllers.contains(&id) {
+        Err(format!(
+            "{} is not a controller of the canister {}",
+            id.to_text(),
+            args.canister.to_text()
+        ))?;
+    }
+
+    let mode = if info.module_hash.is_none() {
+        CanisterInstallMode::Install
+    } else if reinstall.unwrap_or(false) {
+        CanisterInstallMode::Reinstall
+    } else {
+        CanisterInstallMode::Upgrade(None)
+    };
+
+    let prev_hash: [u8; 32] = if let Some(hash) = info.module_hash {
+        hash.try_into().map_err(format_error)?
+    } else {
+        Default::default()
+    };
+    let prev_hash = ByteN::from(prev_hash);
+    let (hash, wasm) = store::wasm::next_version(args.kind, prev_hash)?;
+    let arg = args.args.unwrap_or_default();
+    let res = install_code(InstallCodeArgument {
+        mode,
+        canister_id: args.canister,
+        wasm_module: wasm.wasm.into_vec(),
+        arg: arg.clone().into_vec(),
+    })
+    .await
+    .map_err(format_error);
+
+    store::wasm::add_log(store::InstallLog {
+        kind: args.kind,
+        install_at: ic_cdk::api::time() / MILLISECONDS,
+        canister: args.canister,
+        prev_hash,
+        wasm_hash: hash,
+        args: arg,
+        error: res.clone().err(),
+    });
+    res
+}
+
+#[ic_cdk::update]
+async fn validate_admin_install_wasm(
+    args: DeployWasmInput,
+    _reinstall: Option<bool>,
+) -> Result<(), String> {
+    let (info,) = canister_info(CanisterInfoRequest {
+        canister_id: args.canister,
+        num_requested_changes: None,
+    })
+    .await
+    .map_err(format_error)?;
+    let id = ic_cdk::id();
+    if !info.controllers.contains(&id) {
+        Err(format!(
+            "{} is not a controller of the canister {}",
+            id.to_text(),
+            args.canister.to_text()
+        ))?;
+    }
+
+    let prev_hash: [u8; 32] = if let Some(hash) = info.module_hash {
+        hash.try_into().map_err(format_error)?
+    } else {
+        Default::default()
+    };
+    let prev_hash = ByteN::from(prev_hash);
+    let _ = store::wasm::next_version(args.kind, prev_hash)?;
+    Ok(())
+}

src/ic_oss_cluster/src/lib.rs

@@ -1,5 +1,9 @@
 use candid::Principal;
-use ic_oss_types::{bucket::Token, cluster::ClusterInfo};
+use ic_oss_types::{
+    bucket::Token,
+    cluster::{AddWasmInput, ClusterInfo, DeployWasmInput},
+    ByteN,
+};
 use serde_bytes::ByteBuf;
 use std::collections::BTreeSet;
 
@@ -13,6 +17,7 @@ use crate::init::ChainArgs;
 
 static ANONYMOUS: Principal = Principal::anonymous();
 const SECONDS: u64 = 1_000_000_000;
+const MILLISECONDS: u64 = 1_000_000;
 
 #[ic_cdk::query]
 fn get_cluster_info() -> Result<ClusterInfo, String> {
@@ -34,6 +39,15 @@ fn is_controller() -> Result<(), String> {
     }
 }
 
+fn is_controller_or_manager() -> Result<(), String> {
+    let caller = ic_cdk::caller();
+    if ic_cdk::api::is_controller(&caller) || store::state::is_manager(&caller) {
+        Ok(())
+    } else {
+        Err("user is not a controller or manager".to_string())
+    }
+}
+
 #[cfg(all(
     target_arch = "wasm32",
     target_vendor = "unknown",