feat: add upgrade arg to asset canister to overwrite permissions

CHANGELOG.md

@@ -51,6 +51,8 @@ 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.
 
+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: 965c8899f0a033593dc9b1634b2ab4e0f3fd28c1cfa06993069be2040a2f700e
 - https://github.com/dfinity/sdk/pull/3429
 - https://github.com/dfinity/sdk/pull/3428

docs/design/asset-canister-interface.md

@@ -159,6 +159,30 @@ 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: opt vec principal;
+      commit: opt vec principal;
+      manage_permissions: opt 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.
+Any of the permissions in `set_permissions` that are not `null` are set to the newly provided list of principals.
+The previous list of principals in the given permission is discarded.
+
 ### Method: `get`
 
 ```candid

e2e/tests-dfx/assetscanister.bash

@@ -1732,3 +1732,61 @@ 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_match '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_match 'vec {}'
+
+  dfx identity new alice --storage-mode plaintext
+  ALICE="$(dfx --identity alice identity get-principal)"
+
+  # set new lists
+  dfx canister install e2e_project_frontend --upgrade-unchanged --mode upgrade --argument "(opt variant {
+    Upgrade = record {
+      set_permissions = opt record {
+        prepare = opt vec {
+          principal \""${ALICE}"\";
+        };
+        commit = opt vec {
+          principal \""$(dfx identity get-principal)"\";
+          principal \"aaaaa-aa\";
+        };
+        manage_permissions = opt 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)"
+
+  # set only some lists - others are not touched
+  dfx canister install e2e_project_frontend --upgrade-unchanged --mode upgrade --argument "(opt variant {
+    Upgrade = record {
+      set_permissions = opt record {
+        prepare = null;
+        commit = null;
+        manage_permissions = opt vec {};
+      }
+    }
+  })"
+  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 "vec {}"
+}

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 with a certain permission for every permission that is set.
+type SetPermissions = record {
+  prepare: opt vec principal;
+  commit: opt vec principal;
+  manage_permissions: opt 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,26 @@ impl State {
             .ok_or_else(|| "asset not found".to_string())
     }
 
+    pub fn set_permissions(
+        &mut self,
+        SetPermissions {
+            prepare,
+            commit,
+            manage_permissions,
+        }: SetPermissions,
+    ) {
+        if let Some(prepare) = prepare {
+            *self.get_mut_permission_list(&Permission::Prepare) = prepare.into_iter().collect();
+        }
+        if let Some(commit) = commit {
+            *self.get_mut_permission_list(&Permission::Commit) = commit.into_iter().collect();
+        }
+        if let Some(manage_permissions) = manage_permissions {
+            *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: Option<Vec<Principal>>,
+    pub commit: Option<Vec<Principal>>,
+    pub manage_permissions: Option<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/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 with a certain permission for every permission that is set.
+type SetPermissions = record {
+  prepare: opt vec principal;
+  commit: opt vec principal;
+  manage_permissions: opt vec principal;
+};
+
+service: (asset_canister_args: opt AssetCanisterArgs) -> {
   api_version: () -> (nat16) query;
 
   get: (record {