Merge branch 'luc/ic-keep-main-memory-option' into 'master'

feat(IC-1635): Keep main memory on upgrade option

Introducing a new upgrade option `wasm_memory_persistence` to retain the main memory on a canister upgrade.
This is part of the required IC extension to support Motoko's enhanced orthogonal persistence (IC-1635, cf. dfinity/motoko#4225) and requires support for passive data segments (cf. https://gitlab.com/dfinity-lab/public/ic/-/merge_requests/17892). See dfinity/interface-spec#281 for the corresponding IC specification extension.

Specific design aspects:
* The option `wasm_memory_persistence` can be omitted for backwards compatibility.
* A variant type has been chosen for `wasm_memory_persistence` to allow extendibility in the future, e.g. if multiple Wasm memories would be supported.
* When set to `opt keep`, the upgrade retains both main memory and stable memory. Otherwise, the existing behavior of only retaining the stable memory and reinitializing the main memory is applied.
* As a safety guard, programs with enhanced orthogonal persistence (i.e. Motoko) can add a custom section `enhanced-orthogonal-persistence` into the Wasm. When this is specified, the runtime system checks that the `wasm_memory_persistence` option is not omitted in error. This is to prevent accidental data loss for enhanced orthogonal persistence. Moreover, the system checks that `wasm_memory_persistence: opt keep` is only applied to canisters with enhanced orthogonal persistence.
* For convenience, `drun` is adjusted to automatically apply `wasm_memory_persistence: opt keep` for canisters that have the `enhanced-orthogonal-persistence` section specified. 

See merge request dfinity-lab/public/ic!17893

hs/spec_compliance/ic.did

@@ -46,6 +46,18 @@ type change = record {
   details : change_details;
 };
 
+type canister_install_mode = variant {
+  install;
+  reinstall;
+  upgrade : opt record {
+    skip_pre_upgrade: opt bool;
+    wasm_memory_persistence : opt variant {
+      keep;
+      replace;
+    };
+  };
+};
+
 type chunk_hash = blob;
 
 type http_header = record { name: text; value: text };
@@ -130,26 +142,14 @@ service ic : {
   clear_chunk_store: (record {canister_id : canister_id}) -> ();
   stored_chunks: (record {canister_id : canister_id}) -> (vec chunk_hash);
   install_code : (record {
-    mode : variant {
-      install;
-      reinstall;
-      upgrade : opt record {
-        skip_pre_upgrade: opt bool;
-      }
-    };
+    mode : canister_install_mode;
     canister_id : canister_id;
     wasm_module : wasm_module;
     arg : blob;
     sender_canister_version : opt nat64;
   }) -> ();
   install_chunked_code: (record {
-    mode : variant {
-      install;
-      reinstall;
-      upgrade : opt record {
-        skip_pre_upgrade: opt bool;
-      };
-    };
+    mode : canister_install_mode;
     target_canister: canister_id;
     storage_canister: opt canister_id;
     chunk_hashes_list: vec chunk_hash;

hs/spec_compliance/src/IC/Management.hs

@@ -36,9 +36,17 @@ type SenderCanisterVersion =
 
 type InstallMode =
   [candidType|
-    variant {install : null; reinstall : null; upgrade : opt record {
+    variant {
+      install : null; 
+      reinstall : null; 
+      upgrade : opt record {
         skip_pre_upgrade : opt bool;
-      }}
+        wasm_memory_persistence : opt variant {
+          keep;
+          replace;
+        };
+      };
+    }
   |]
 
 type RunState =

rs/drun/BUILD.bazel

@@ -36,6 +36,7 @@ DEPENDENCIES = [
     "@crate_index//:tokio",
     "@crate_index//:tower",
     "@crate_index//:rand",
+    "@crate_index//:wasmparser",
 ]
 
 rust_library(

rs/drun/Cargo.toml

@@ -45,6 +45,7 @@ tokio = { workspace = true }
 rand = "0.8"
 tower = { workspace = true }
 futures.workspace = true
+wasmparser = "0.115.0"
 
 [[bin]]
 name = "drun"

rs/drun/README.adoc

@@ -45,6 +45,7 @@ Code installation messages have the following format:
 ----
 
 * `<mode>` is one of `install`, `reinstall` or `upgrade`
+- `drun` automatically detects whether the Wasm binary expects enhanced orthogonal persistence (as used in Motoko) and if so, sets the `wasm_memory_persistence` upgrade option.
 
 * `<canister_id>` is the desired ID for the canister to be installed, given in textual
 representation (e.g. `rwlgt-iiaaa-aaaaa-aaaaa-cai`) as specified in https://internetcomputer.org/docs/current/references/ic-interface-spec#textual-ids.

rs/drun/src/message.rs

@@ -1,15 +1,17 @@
 use super::CanisterId;
 
 use hex::decode;
-use ic_management_canister_types::{self as ic00, CanisterInstallMode, Payload};
+use ic_execution_environment::execution::upgrade::ENHANCED_ORTHOGONAL_PERSISTENCE_SECTION;
+use ic_management_canister_types::{
+    self as ic00, CanisterInstallModeV2, CanisterUpgradeOptions, Payload, WasmMemoryPersistence,
+};
 use ic_types::{
     messages::{SignedIngress, UserQuery},
     time::expiry_time_from_now,
     PrincipalId, UserId,
 };
 
 use std::{
-    convert::TryFrom,
     fmt,
     fs::File,
     io::{self, Read},
@@ -216,6 +218,21 @@ fn parse_create(nonce: u64) -> Result<Message, String> {
     Ok(Message::Create(signed_ingress))
 }
 
+fn contains_icp_private_custom_section(wasm_binary: &[u8], name: &str) -> Result<bool, String> {
+    use wasmparser::{Parser, Payload::CustomSection};
+
+    let icp_section_name = format!("icp:private {name}");
+    let parser = Parser::new(0);
+    for payload in parser.parse_all(wasm_binary) {
+        if let CustomSection(reader) = payload.map_err(|e| format!("Wasm parsing error: {}", e))? {
+            if reader.name() == icp_section_name {
+                return Ok(true);
+            }
+        }
+    }
+    Ok(false)
+}
+
 fn parse_install(
     nonce: u64,
     canister_id: &str,
@@ -235,14 +252,36 @@ fn parse_install(
     let canister_id = parse_canister_id(canister_id)?;
     let payload = parse_octet_string(payload)?;
 
+    let install_mode = match mode {
+        "install" => CanisterInstallModeV2::Install,
+        "reinstall" => CanisterInstallModeV2::Reinstall,
+        "upgrade" => {
+            let wasm_memory_persistence = if contains_icp_private_custom_section(
+                wasm_data.as_ref(),
+                ENHANCED_ORTHOGONAL_PERSISTENCE_SECTION,
+            )? {
+                Some(WasmMemoryPersistence::Keep)
+            } else {
+                None
+            };
+            CanisterInstallModeV2::Upgrade(Some(CanisterUpgradeOptions {
+                skip_pre_upgrade: None,
+                wasm_memory_persistence,
+            }))
+        }
+        _ => {
+            return Err(String::from("Unsupported install mode: {mode}"));
+        }
+    };
+
     let signed_ingress = SignedIngressBuilder::new()
         // `source` should become a self-authenticating id according
         // to https://internetcomputer.org/docs/current/references/ic-interface-spec#id-classes
         .canister_id(ic00::IC_00)
         .method_name(ic00::Method::InstallCode)
         .method_payload(
-            ic00::InstallCodeArgs::new(
-                CanisterInstallMode::try_from(mode.to_string()).unwrap(),
+            ic00::InstallCodeArgsV2::new(
+                install_mode,
                 canister_id,
                 wasm_data,
                 payload,

rs/execution_environment/src/canister_manager.rs

@@ -881,14 +881,14 @@ impl CanisterManager {
     ///
     /// There are three modes of installation that are supported:
     ///
-    /// 1. `CanisterInstallMode::Install`
+    /// 1. `CanisterInstallModeV2::Install`
     ///    Used for installing code on an empty canister.
     ///
-    /// 2. `CanisterInstallMode::Reinstall`
+    /// 2. `CanisterInstallModeV2::Reinstall`
     ///    Used for installing code on a _non-empty_ canister. All existing
     ///    state in the canister is cleared.
     ///
-    /// 3. `CanisterInstallMode::Upgrade`
+    /// 3. `CanisterInstallModeV2::Upgrade`
     ///    Used for upgrading a canister while providing a mechanism to
     ///    preserve its state.
     ///
@@ -2097,6 +2097,12 @@ pub(crate) enum CanisterManagerError {
         canister_id: CanisterId,
         snapshot_id: SnapshotId,
     },
+    MissingUpgradeOptionError {
+        message: String,
+    },
+    InvalidUpgradeOptionError {
+        message: String,
+    },
 }
 
 impl AsErrorHelp for CanisterManagerError {
@@ -2132,12 +2138,12 @@ impl AsErrorHelp for CanisterManagerError {
             | CanisterManagerError::WasmChunkStoreError { .. }
             | CanisterManagerError::CanisterSnapshotNotFound { .. }
             | CanisterManagerError::CanisterHeapDeltaRateLimited { .. }
-            | CanisterManagerError::CanisterSnapshotInvalidOwnership { .. } => {
-                ErrorHelp::UserError {
-                    suggestion: "".to_string(),
-                    doc_link: "".to_string(),
-                }
-            }
+            | CanisterManagerError::CanisterSnapshotInvalidOwnership { .. }
+            | CanisterManagerError::MissingUpgradeOptionError { .. }
+            | CanisterManagerError::InvalidUpgradeOptionError { .. } => ErrorHelp::UserError {
+                suggestion: "".to_string(),
+                doc_link: "".to_string(),
+            },
         }
     }
 }
@@ -2405,6 +2411,22 @@ impl From<CanisterManagerError> for UserError {
                     )
                 )
             }
+            MissingUpgradeOptionError { message } => {
+                Self::new(
+                    ErrorCode::CanisterContractViolation,
+                    format!(
+                        "Missing upgrade option: {}", message
+                    )
+                )
+            }
+            InvalidUpgradeOptionError { message } => {
+                Self::new(
+                    ErrorCode::CanisterContractViolation,
+                    format!(
+                        "Invalid upgrade option: {}", message
+                    )
+                )
+            }
         }
     }
 }