whitelist certain wasms

service/pool/Main.mo

@@ -17,26 +17,32 @@ import PoW "./PoW";
 import Logs "./Logs";
 import Metrics "./Metrics";
 import Wasm "canister:wasm-utils";
+import Blob "mo:base/Blob";
+import Buffer "mo:base/Buffer";
+import Nat32 "mo:base/Nat32";
 
 shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
     let IC : ICType.Self = actor "aaaaa-aa";
     let params = Option.get(opt_params, Types.defaultParams);
     var pool = Types.CanisterPool(params.max_num_canisters, params.canister_time_to_live, params.max_family_tree_size);
     let nonceCache = PoW.NonceCache(params.nonce_time_to_live);
+    var whitelistedWasmHashes = Buffer.Buffer<Text>(4);
 
     stable let controller = creator.caller;
     stable var stats = Logs.defaultStats;
     stable var stablePool : [Types.CanisterInfo] = [];
     stable var stableMetadata : [(Principal, (Int, Bool))] = [];
     stable var stableChildren : [(Principal, [Principal])] = [];
     stable var previousParam : ?Types.InitParams = null;
+    stable var whitelistedHashes : [Text] = [];
 
     system func preupgrade() {
         let (tree, metadata, children) = pool.share();
         stablePool := tree;
         stableMetadata := metadata;
         stableChildren := children;
         previousParam := ?params;
+        whitelistedHashes := Buffer.toArray(whitelistedWasmHashes);
     };
 
     system func postupgrade() {
@@ -46,6 +52,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
             };
         };
         pool.unshare(stablePool, stableMetadata, stableChildren);
+        whitelistedWasmHashes := Buffer.fromArray(whitelistedHashes);
     };
 
     public query func getInitParams() : async Types.InitParams {
@@ -65,6 +72,46 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
         ignore Cycles.accept amount;
     };
 
+    public query func getWhitelistedWasmHashes() : async [Text] {
+        Buffer.toArray(whitelistedWasmHashes);
+    };
+
+    public shared ({ caller }) func whitelistWasmHash(hash : Text) : async () {
+        if (caller != controller) {
+            throw Error.reject "Only called by controller";
+        };
+        whitelistedWasmHashes.add(hash);
+    };
+
+    public shared ({ caller }) func whitelistWasm(wasm : Blob) : async () {
+        if (caller != controller) {
+            throw Error.reject "Only called by controller";
+        };
+        let hash = await Wasm.hash(wasm);
+        whitelistedWasmHashes.add(hash);
+    };
+
+    public shared ({ caller }) func removeWhitelistedWasmHash(hash : Text) : async () {
+        if (caller != controller) {
+            throw Error.reject "Only called by controller";
+        };
+        switch (Buffer.indexOf<Text>(hash, whitelistedWasmHashes, Text.equal)) {
+            case (?index) ignore whitelistedWasmHashes.remove(index);
+            case null ();
+        };
+    };
+
+    public shared ({ caller }) func removeWhitelistedWasm(wasm : Blob) : async () {
+        if (caller != controller) {
+            throw Error.reject "Only called by controller";
+        };
+        let hash = await Wasm.hash(wasm);
+        switch (Buffer.indexOf<Text>(hash, whitelistedWasmHashes, Text.equal)) {
+            case (?index) ignore whitelistedWasmHashes.remove(index);
+            case null ();
+        };
+    };
+
     private func getExpiredCanisterInfo() : async Types.CanisterInfo {
         switch (pool.getExpiredCanisterId()) {
             case (#newId) {
@@ -134,7 +181,12 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
                 limit_stable_memory_page = ?(16384 : Nat32); // Limit to 1G of stable memory
                 backend_canister_id = ?Principal.fromActor(this);
             };
-            let wasm = await Wasm.transform(args.wasm_module, config);
+            let wasm_hash = await Wasm.hash(args.wasm_module);
+            let wasm = if (Buffer.contains<Text>(whitelistedWasmHashes, wasm_hash, Text.equal)) {
+                args.wasm_module;
+            } else {
+                await Wasm.transform(args.wasm_module, config);
+            };
             let newArgs = {
                 arg = args.arg;
                 wasm_module = wasm;
@@ -359,7 +411,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
         msg : {
             #GCCanisters : Any;
             #balance : Any;
-            #callForward: Any;
+            #callForward : Any;
             #dump : Any;
             #getCanisterId : Any;
             #getSubtree : Any;
@@ -370,6 +422,11 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
             #removeCode : Any;
             #resetStats : Any;
             #wallet_receive : Any;
+            #getWhitelistedWasmHashes : Any;
+            #whitelistWasmHash : Any;
+            #whitelistWasm : Any;
+            #removeWhitelistedWasmHash : Any;
+            #removeWhitelistedWasm : Any;
 
             #create_canister : Any;
             #update_settings : Any;

service/wasm-utils/Cargo.toml

@@ -8,12 +8,14 @@ path = "lib.rs"
 crate-type = ["cdylib"]
 
 [dependencies]
+hex = "0.4.3"
 ic-cdk = "0.8.0-beta.0"
 ic-cdk-macros = "0.8.0-beta.0"
 serde = "1.0"
 serde_bytes = "0.11"
 candid = "0.9.0-beta.2"
 ic-wasm = { version = "0.3.7", default-features = false }
+sha2 = "0.10.6"
 walrus = "0.19"
 
 [profile.release]

service/wasm-utils/lib.rs

@@ -2,6 +2,7 @@ use candid::{CandidType, Deserialize};
 use serde_bytes::ByteBuf;
 
 use ic_wasm::*;
+use sha2::Digest;
 
 #[derive(CandidType, Deserialize)]
 struct Config {
@@ -11,6 +12,11 @@ struct Config {
     backend_canister_id: Option<candid::Principal>,
 }
 
+#[ic_cdk_macros::query]
+fn hash(wasm: ByteBuf) -> String {
+    hex::encode(sha2::Sha256::digest(wasm))
+}
+
 #[ic_cdk_macros::query]
 fn transform(wasm: ByteBuf, config: Config) -> ByteBuf {
     let mut m = walrus::Module::from_buffer(&wasm).unwrap();

service/wasm-utils/wasm-utils.did

@@ -1,5 +1,6 @@
 type Config = record { profiling: bool; remove_cycles_add: bool; limit_stable_memory_page: opt nat32; backend_canister_id: opt principal};
 
 service : {
-  transform : (blob, Config) -> (blob);
+  transform : (blob, Config) -> (blob) query;
+  hash : (blob) -> (text) query;
 }