fix

dfinity · Oct 22, 2024 · 9f760ed · 9f760ed
1 parent 2e6fb7a
commit 9f760ed
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/service/pool/Main.mo b/service/pool/Main.mo
@@ -125,6 +125,19 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
                 };
                 if (need_uninstall) {
                     await* pool_uninstall_code(cid.canister_id);
+                    switch (params.stored_module_hash) {
+                    case null {};
+                    case (?stored) {
+                             await IC.install_chunked_code {
+                                 arg = "DIDL\00\00";
+                                 target_canister = cid.canister_id;
+                                 store_canister = ?(Principal.fromActor this);
+                                 chunk_hashes_list = [{ hash = stored }];
+                                 wasm_module_hash = stored;
+                                 mode = #install;
+                             }
+                         };
+                    };
                 };
                 switch (status.status) {
                     case (#stopped or #stopping) {
@@ -134,7 +147,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
                 };
                 stats := Logs.updateStats(stats, #getId topUpCycles);
                 statsByOrigin.addCanister(origin);
-                let mode = if (need_uninstall) { #install } else { #reinstall };
+                let mode = if (need_uninstall and Option.isNull(params.stored_module_hash)) { #install } else { #reinstall };
                 (info, mode);
             };
             case (#outOfCapacity time) {
@@ -249,6 +262,9 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
     };
 
     public shared ({ caller }) func getCanisterId(nonce : PoW.Nonce, origin : Logs.Origin) : async Types.CanisterInfo {
+        if (Option.get(params.admin_only, false)) {
+            throw Error.reject "Cannot call this endpoint when admin_only is true";
+        };
         if (not validateOrigin(origin)) {
             throw Error.reject "Please specify a valid origin";
         };

diff --git a/service/pool/Types.mo b/service/pool/Types.mo
@@ -22,6 +22,7 @@ module {
         max_family_tree_size: Nat;
         // Used for installing asset canister. If set, will not use timer to kill expired canisters, and will not uninstall code when fetching an expired canister (unless the module hash changed).
         stored_module_hash: ?Blob;
+        admin_only: ?Bool;
         wasm_utils_principal: ?Text;
     };
     public let defaultParams : InitParams = {
@@ -31,6 +32,7 @@ module {
         nonce_time_to_live = 300_000_000_000;
         max_family_tree_size = 5;
         stored_module_hash = null;
+        admin_only = null;
         wasm_utils_principal = ?"ozk6r-tyaaa-aaaab-qab4a-cai";
     };
     public type InstallArgs = {