feat: add upgrade arg to asset canister to overwrite permissions (#3443)

CHANGELOG.md

@@ -57,7 +57,10 @@ Defining a custom `etag` header no longer breaks certification.
 
 Fixed a certification issue where under certain conditions the fallback file (`/index.html`) was served with an incomplete certificate tree, not proving sufficiently that the fallback file may be used as a replacement.
 
-- Module hash: 1621e9ead6463304ceb3a81b10577c61c9d24f6e70b2e275a10b3a9be982dfb4
+Add the option to (re)set all permissions using upgrade arguments. This is especially useful for SNSes that cannot make calls as the canister's controller.
+
+- Module hash: 657938477f1dee46db70b5a9f0bd167ec5ffcd2f930a1d96593c17dcddef61b3
+- https://github.com/dfinity/sdk/pull/3443
 - https://github.com/dfinity/sdk/pull/3451
 - https://github.com/dfinity/sdk/pull/3429
 - https://github.com/dfinity/sdk/pull/3428

docs/design/asset-canister-interface.md

@@ -159,6 +159,29 @@ The size of any chunk cannot exceed the message ingress limit.
 
 ## Method Reference
 
+### Method: `init` and `post_upgrade`
+
+```candid
+service: (asset_canister_args: variant {
+  Init: record {};
+  Upgrade: record {
+    set_permissions: opt record {
+      prepare: vec principal;
+      commit: vec principal;
+      manage_permissions: vec principal;
+    };
+  };
+})
+```
+
+The methods `init` and `post_upgrade` are called automatically by the system after new code is installed in the canister.
+
+Both methods take the same argument type by definition. Therefore, to be able to have different arguments for the two cases, an enum is used to make the distinction.
+If `init` is called with the `Upgrade` variant or if `post_upgrade` is called with the `Init` variant the asset canister traps and thereby reverts the code changes.
+
+In `Upgrade`, the field `set_permissions` can be used to (re)set the list of principals with the listed permissions.
+If `set_permissions` that is not `null`, then all permissions are set to the newly provided list of principals and the previous lists of principals are discarded.
+
 ### Method: `get`
 
 ```candid

e2e/tests-dfx/assetscanister.bash

@@ -1732,3 +1732,42 @@ WARN: {
   assert_command      dfx canister call e2e_project_frontend configure '(record { max_chunks=opt opt 3; max_bytes = opt opt 5500 })'
   assert_command      dfx deploy
 }
+
+@test "set permissions through upgrade argument" {
+  dfx_start
+  dfx deploy
+
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Prepare }; })'
+  assert_eq "(vec {})"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Commit }; })'
+  assert_match "$(dfx identity get-principal)"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { ManagePermissions }; })'
+  assert_eq "(vec {})"
+
+  dfx identity new alice --storage-mode plaintext
+  ALICE="$(dfx --identity alice identity get-principal)"
+
+  dfx canister install e2e_project_frontend --upgrade-unchanged --mode upgrade --argument "(opt variant {
+    Upgrade = record {
+      set_permissions = opt record {
+        prepare = vec {
+          principal \"${ALICE}\";
+        };
+        commit = vec {
+          principal \"$(dfx identity get-principal)\";
+          principal \"aaaaa-aa\";
+        };
+        manage_permissions = vec {
+          principal \"$(dfx identity get-principal)\";
+        };
+      }
+    }
+  })"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Prepare }; })'
+  assert_match "${ALICE}"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Commit }; })'
+  assert_match "$(dfx identity get-principal)"
+  assert_match '"aaaaa-aa"'
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { ManagePermissions }; })'
+  assert_match "$(dfx identity get-principal)"
+}

src/canisters/frontend/ic-certified-assets/assets.did

@@ -139,7 +139,25 @@ type ListPermitted = record { permission: Permission };
 
 type ValidationResult = variant { Ok : text; Err : text };
 
-service: {
+type AssetCanisterArgs = variant {
+  Init: InitArgs;
+  Upgrade: UpgradeArgs;
+};
+
+type InitArgs = record {};
+
+type UpgradeArgs = record {
+  set_permissions: opt SetPermissions;
+};
+
+/// Sets the list of principals granted each permission.
+type SetPermissions = record {
+  prepare: vec principal;
+  commit: vec principal;
+  manage_permissions: vec principal;
+};
+
+service: (asset_canister_args: opt AssetCanisterArgs) -> {
   api_version: () -> (nat16) query;
 
   get: (record {

src/canisters/frontend/ic-certified-assets/src/lib.rs

@@ -413,7 +413,10 @@ fn is_controller() -> Result<(), String> {
     }
 }
 
-pub fn init() {
+pub fn init(args: Option<AssetCanisterArgs>) {
+    if let Some(upgrade_arg) = args {
+        let AssetCanisterArgs::Init(InitArgs {}) = upgrade_arg else { ic_cdk::trap("Cannot initialize the canister with an Upgrade argument. Please provide an Init argument.")};
+    }
     STATE.with(|s| {
         let mut s = s.borrow_mut();
         s.clear();
@@ -425,10 +428,18 @@ pub fn pre_upgrade() -> StableState {
     STATE.with(|s| s.take().into())
 }
 
-pub fn post_upgrade(stable_state: StableState) {
+pub fn post_upgrade(stable_state: StableState, args: Option<AssetCanisterArgs>) {
+    let set_permissions = args.and_then(|args| {
+        let AssetCanisterArgs::Upgrade(UpgradeArgs { set_permissions }) = args else {ic_cdk::trap("Cannot upgrade the canister with an Init argument. Please provide an Upgrade argument.")};
+        set_permissions
+    });
+
     STATE.with(|s| {
         *s.borrow_mut() = State::from(stable_state);
         set_certified_data(&s.borrow().root_hash());
+        if let Some(set_permissions) = set_permissions {
+            s.borrow_mut().set_permissions(set_permissions);
+        }
     });
 }
 

src/canisters/frontend/ic-certified-assets/src/state_machine.rs

@@ -320,6 +320,20 @@ impl State {
             .ok_or_else(|| "asset not found".to_string())
     }
 
+    pub fn set_permissions(
+        &mut self,
+        SetPermissions {
+            prepare,
+            commit,
+            manage_permissions,
+        }: SetPermissions,
+    ) {
+        *self.get_mut_permission_list(&Permission::Prepare) = prepare.into_iter().collect();
+        *self.get_mut_permission_list(&Permission::Commit) = commit.into_iter().collect();
+        *self.get_mut_permission_list(&Permission::ManagePermissions) =
+            manage_permissions.into_iter().collect();
+    }
+
     pub fn grant_permission(&mut self, principal: Principal, permission: &Permission) {
         let permitted = self.get_mut_permission_list(permission);
         permitted.insert(principal);

src/canisters/frontend/ic-certified-assets/src/types.rs

@@ -184,3 +184,27 @@ pub struct RevokePermissionArguments {
 pub struct ListPermittedArguments {
     pub permission: Permission,
 }
+
+/// The argument to `init` and `post_upgrade` needs to have the same argument type by definition.
+/// `AssetCanisterArgs` is there so that the two functions can take different argument types.
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub enum AssetCanisterArgs {
+    Init(InitArgs),
+    Upgrade(UpgradeArgs),
+}
+
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub struct InitArgs {}
+
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub struct UpgradeArgs {
+    pub set_permissions: Option<SetPermissions>,
+}
+
+/// Sets the list of principals with a certain permission for every permission that is `Some`.
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub struct SetPermissions {
+    pub prepare: Vec<Principal>,
+    pub commit: Vec<Principal>,
+    pub manage_permissions: Vec<Principal>,
+}

src/canisters/frontend/ic-frontend-canister/src/lib.rs

@@ -1,8 +1,9 @@
 use ic_cdk::{init, post_upgrade, pre_upgrade};
+use ic_certified_assets::types::AssetCanisterArgs;
 
 #[init]
-fn init() {
-    ic_certified_assets::init();
+fn init(args: Option<AssetCanisterArgs>) {
+    ic_certified_assets::init(args);
 }
 
 #[pre_upgrade]
@@ -12,8 +13,8 @@ fn pre_upgrade() {
 }
 
 #[post_upgrade]
-fn post_upgrade() {
+fn post_upgrade(args: Option<AssetCanisterArgs>) {
     let (stable_state,): (ic_certified_assets::StableState,) =
         ic_cdk::storage::stable_restore().expect("failed to restore stable state");
-    ic_certified_assets::post_upgrade(stable_state);
+    ic_certified_assets::post_upgrade(stable_state, args);
 }

src/dfx/src/commands/canister/install.rs

@@ -203,7 +203,9 @@ pub async fn exec(
                     Principal::from_text(canister).or_else(|_| canister_id_store.get(canister))?;
                 let canister_info = CanisterInfo::load(&config, canister, Some(canister_id))?;
 
-                let install_args = || Ok(vec![]);
+                let idl_path = canister_info.get_constructor_idl_path();
+                let init_type = get_candid_init_type(&idl_path);
+                let install_args = || blob_from_arguments(None, None, None, &init_type);
 
                 install_canister(
                     env,

src/distributed/assetstorage.did

@@ -139,7 +139,25 @@ type ListPermitted = record { permission: Permission };
 
 type ValidationResult = variant { Ok : text; Err : text };
 
-service: {
+type AssetCanisterArgs = variant {
+  Init: InitArgs;
+  Upgrade: UpgradeArgs;
+};
+
+type InitArgs = record {};
+
+type UpgradeArgs = record {
+  set_permissions: opt SetPermissions;
+};
+
+/// Sets the list of principals granted each permission.
+type SetPermissions = record {
+  prepare: vec principal;
+  commit: vec principal;
+  manage_permissions: vec principal;
+};
+
+service: (asset_canister_args: opt AssetCanisterArgs) -> {
   api_version: () -> (nat16) query;
 
   get: (record {