From 279595e616bd5c3c633bdb4a3340dab0968cc1ae Mon Sep 17 00:00:00 2001
From: Severin Siffert <severin.siffert@dfinity.org>
Date: Mon, 28 Oct 2024 10:25:58 +0100
Subject: [PATCH] test and canister interface change

---
 e2e/tests-dfx/assetscanister.bash             | 39 +++++++++++++++++++
 .../frontend/ic-certified-assets/src/lib.rs   | 17 +++++---
 .../frontend/ic-certified-assets/src/types.rs |  4 +-
 3 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/e2e/tests-dfx/assetscanister.bash b/e2e/tests-dfx/assetscanister.bash
index 5de6817e58..d4b3df692f 100644
--- a/e2e/tests-dfx/assetscanister.bash
+++ b/e2e/tests-dfx/assetscanister.bash
@@ -1988,6 +1988,45 @@ WARN: {
   assert_command      dfx deploy
 }
 
+@test "set permissions through init 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 --mode reinstall --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)"
+}
+
 @test "set permissions through upgrade argument" {
   dfx_start
   dfx deploy
diff --git a/src/canisters/frontend/ic-certified-assets/src/lib.rs b/src/canisters/frontend/ic-certified-assets/src/lib.rs
index 44847a1e1c..4b96ef1b44 100644
--- a/src/canisters/frontend/ic-certified-assets/src/lib.rs
+++ b/src/canisters/frontend/ic-certified-assets/src/lib.rs
@@ -423,16 +423,23 @@ fn is_controller() -> Result<(), String> {
 }
 
 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();
         s.grant_permission(caller(), &Permission::Commit);
     });
+
+    if let Some(upgrade_arg) = args {
+        let AssetCanisterArgs::Init(init_args) = upgrade_arg else {
+            ic_cdk::trap("Cannot initialize the canister with an Upgrade argument. Please provide an Init argument.")
+        };
+        STATE.with(|s| {
+            let mut state = s.borrow_mut();
+            if let Some(set_permissions) = init_args.set_permissions {
+                state.set_permissions(set_permissions);
+            }
+        });
+    }
 }
 
 pub fn pre_upgrade() -> StableState {
diff --git a/src/canisters/frontend/ic-certified-assets/src/types.rs b/src/canisters/frontend/ic-certified-assets/src/types.rs
index a164dc12c5..5274b9b06b 100644
--- a/src/canisters/frontend/ic-certified-assets/src/types.rs
+++ b/src/canisters/frontend/ic-certified-assets/src/types.rs
@@ -207,7 +207,9 @@ pub enum AssetCanisterArgs {
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
-pub struct InitArgs {}
+pub struct InitArgs {
+    pub set_permissions: Option<SetPermissions>,
+}
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub struct UpgradeArgs {