diff --git a/apps/wallet/src/components/external-canisters/CanisterInstallForm.spec.ts b/apps/wallet/src/components/external-canisters/CanisterInstallForm.spec.ts index a622d17d0..10e44c556 100644 --- a/apps/wallet/src/components/external-canisters/CanisterInstallForm.spec.ts +++ b/apps/wallet/src/components/external-canisters/CanisterInstallForm.spec.ts @@ -1,5 +1,7 @@ import { Principal } from '@dfinity/principal'; import { describe, expect, it } from 'vitest'; +import { VCheckbox } from 'vuetify/components'; +import CanisterWasmMemoryPersistenceSelect from '~/components/inputs/CanisterWasmMemoryPersistenceSelect.vue'; import { mount } from '~/test.utils'; import CanisterInstallForm from './CanisterInstallForm.vue'; @@ -27,4 +29,93 @@ describe('CanisterInstallForm', () => { expect(canisterIdInput.exists()).toBe(true); }); + + it('hides the upgrade options unless the mode is upgrade', () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { mode: { install: null } }, + }, + }); + + expect(form.findComponent(CanisterWasmMemoryPersistenceSelect).exists()).toBe(false); + }); + + it('shows the upgrade options when the mode is upgrade', () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { mode: { upgrade: [] } }, + }, + }); + + expect(form.findComponent(CanisterWasmMemoryPersistenceSelect).exists()).toBe(true); + }); + + it('folds the selected wasm_memory_persistence into the upgrade mode', async () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { mode: { upgrade: [] } }, + }, + }); + + form.findComponent(CanisterWasmMemoryPersistenceSelect).vm.$emit('update:modelValue', { + keep: null, + }); + await form.vm.$nextTick(); + + const updates = form.emitted('update:modelValue'); + expect(updates).toBeTruthy(); + expect(updates?.at(-1)?.[0]).toEqual({ + mode: { upgrade: [{ wasm_memory_persistence: [{ keep: null }], skip_pre_upgrade: [] }] }, + }); + }); + + it('collapses back to a plain upgrade when the options are cleared', async () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { + mode: { upgrade: [{ wasm_memory_persistence: [{ keep: null }], skip_pre_upgrade: [] }] }, + }, + }, + }); + + form + .findComponent(CanisterWasmMemoryPersistenceSelect) + .vm.$emit('update:modelValue', undefined); + await form.vm.$nextTick(); + + const updates = form.emitted('update:modelValue'); + expect(updates?.at(-1)?.[0]).toEqual({ mode: { upgrade: [] } }); + }); + + it('folds the skip_pre_upgrade toggle into the upgrade mode', async () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { mode: { upgrade: [] } }, + }, + }); + + form.findComponent(VCheckbox).vm.$emit('update:modelValue', true); + await form.vm.$nextTick(); + + const updates = form.emitted('update:modelValue'); + expect(updates?.at(-1)?.[0]).toEqual({ + mode: { upgrade: [{ wasm_memory_persistence: [], skip_pre_upgrade: [true] }] }, + }); + }); + + it('collapses back to a plain upgrade when skip_pre_upgrade is disabled', async () => { + const form = mount(CanisterInstallForm, { + props: { + modelValue: { + mode: { upgrade: [{ wasm_memory_persistence: [], skip_pre_upgrade: [true] }] }, + }, + }, + }); + + form.findComponent(VCheckbox).vm.$emit('update:modelValue', false); + await form.vm.$nextTick(); + + const updates = form.emitted('update:modelValue'); + expect(updates?.at(-1)?.[0]).toEqual({ mode: { upgrade: [] } }); + }); }); diff --git a/apps/wallet/src/components/external-canisters/CanisterInstallForm.vue b/apps/wallet/src/components/external-canisters/CanisterInstallForm.vue index adb827375..b63ab9f92 100644 --- a/apps/wallet/src/components/external-canisters/CanisterInstallForm.vue +++ b/apps/wallet/src/components/external-canisters/CanisterInstallForm.vue @@ -15,6 +15,26 @@ + diff --git a/apps/wallet/src/locales/en.locale.ts b/apps/wallet/src/locales/en.locale.ts index c9ef374c3..6f18aa3e1 100644 --- a/apps/wallet/src/locales/en.locale.ts +++ b/apps/wallet/src/locales/en.locale.ts @@ -663,6 +663,17 @@ export default { upgrade: 'Upgrade', install: 'Install', }, + wasm_memory_persistence: { + label: 'Wasm Memory Persistence', + hint: 'Controls the canister main memory on upgrade. Motoko canisters using Enhanced Orthogonal Persistence require "Keep".', + default: 'Default (replace)', + keep: 'Keep', + replace: 'Replace', + }, + skip_pre_upgrade: { + label: 'Skip pre-upgrade hook', + hint: 'Skips the canister pre_upgrade hook during the upgrade. Useful for recovery when the hook traps.', + }, }, terms: { license: 'License', diff --git a/apps/wallet/src/locales/fr.locale.ts b/apps/wallet/src/locales/fr.locale.ts index 0eb98ca59..9e47eedba 100644 --- a/apps/wallet/src/locales/fr.locale.ts +++ b/apps/wallet/src/locales/fr.locale.ts @@ -673,6 +673,17 @@ export default { upgrade: 'Mettre à jour', install: 'Installer', }, + wasm_memory_persistence: { + label: 'Persistance de la mémoire Wasm', + hint: 'Contrôle la mémoire principale du canister lors de la mise à jour. Les canisters Motoko utilisant la persistance orthogonale améliorée nécessitent « Conserver ».', + default: 'Par défaut (remplacer)', + keep: 'Conserver', + replace: 'Remplacer', + }, + skip_pre_upgrade: { + label: 'Ignorer le hook pre-upgrade', + hint: 'Ignore le hook pre_upgrade du canister lors de la mise à jour. Utile pour la récupération lorsque le hook échoue.', + }, }, terms: { license: 'Licence', diff --git a/apps/wallet/src/locales/pt.locale.ts b/apps/wallet/src/locales/pt.locale.ts index e1816dff3..5af8d01a8 100644 --- a/apps/wallet/src/locales/pt.locale.ts +++ b/apps/wallet/src/locales/pt.locale.ts @@ -669,6 +669,17 @@ export default { upgrade: 'Atualizar', install: 'Instalar', }, + wasm_memory_persistence: { + label: 'Persistência da memória Wasm', + hint: 'Controla a memória principal do canister na atualização. Canisters Motoko que usam Persistência Ortogonal Aprimorada exigem "Manter".', + default: 'Padrão (substituir)', + keep: 'Manter', + replace: 'Substituir', + }, + skip_pre_upgrade: { + label: 'Ignorar hook de pré-atualização', + hint: 'Ignora o hook pre_upgrade do canister durante a atualização. Útil para recuperação quando o hook falha.', + }, }, terms: { license: 'Licença', diff --git a/core/station/api/src/common.rs b/core/station/api/src/common.rs index 831725fda..297578cf2 100644 --- a/core/station/api/src/common.rs +++ b/core/station/api/src/common.rs @@ -30,7 +30,7 @@ pub enum SortDirection { Desc, } -#[derive(CandidType, serde::Serialize, Deserialize, Debug, Clone)] +#[derive(CandidType, serde::Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] pub enum CanisterInstallMode { #[serde(rename = "install")] Install, diff --git a/tests/integration/src/dfx_orbit/install.rs b/tests/integration/src/dfx_orbit/install.rs index fc40fa73b..39494afcb 100644 --- a/tests/integration/src/dfx_orbit/install.rs +++ b/tests/integration/src/dfx_orbit/install.rs @@ -8,7 +8,9 @@ use crate::{ TestEnv, }; use candid::Encode; -use dfx_orbit::canister::{CanisterInstallModeArgs, RequestCanisterInstallArgs}; +use dfx_orbit::canister::{ + CanisterInstallModeArgs, RequestCanisterInstallArgs, WasmMemoryPersistenceArgs, +}; use dfx_orbit::{ args::{RequestArgs, RequestArgsActions, VerifyArgs, VerifyArgsAction}, canister::{ @@ -16,9 +18,10 @@ use dfx_orbit::{ VerifyCanisterArgs, }, }; +use flate2::read::GzDecoder; use sha2::{Digest, Sha256}; use station_api::{GetRequestInput, RequestApprovalStatusDTO}; -use std::io::Write; +use std::io::{Read, Write}; use tempfile::NamedTempFile; fn hash(data: &[u8]) -> Vec { @@ -27,6 +30,36 @@ fn hash(data: &[u8]) -> Vec { hasher.finalize().to_vec() } +/// Appends a custom section to a Wasm module. The IC decides whether a module +/// supports Enhanced Orthogonal Persistence by the presence of the +/// `icp:private enhanced-orthogonal-persistence` custom section, so appending +/// it turns the plain test canister into an EOP module from the IC's +/// perspective. +fn append_wasm_custom_section(module: &mut Vec, name: &str, payload: &[u8]) { + fn write_leb128(mut value: usize, out: &mut Vec) { + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + out.push(byte); + if value == 0 { + break; + } + } + } + + let mut contents = Vec::new(); + write_leb128(name.len(), &mut contents); + contents.extend_from_slice(name.as_bytes()); + contents.extend_from_slice(payload); + + module.push(0); // custom section id + write_leb128(contents.len(), module); + module.extend_from_slice(&contents); +} + /// Test installing a canister through orbit using the station agent #[test] fn canister_install_no_chunks() { @@ -85,6 +118,8 @@ fn canister_install(use_chunks: bool) { let inner_args = RequestCanisterInstallArgs { canister: String::from("test"), mode: CanisterInstallModeArgs::Install, + wasm_memory_persistence: None, + skip_pre_upgrade: false, wasm: wasm.path().as_os_str().to_str().unwrap().to_string(), argument: None, arg_file: None, @@ -152,3 +187,247 @@ fn canister_install(use_chunks: bool) { let status = canister_status(&env, Some(canister_ids.station), test_canister); assert_eq!(status.module_hash, Some(module_hash)); } + +/// Test that `--wasm-memory-persistence` and `--skip-pre-upgrade` are plumbed +/// through into the upgrade request, survive the round-trip through the +/// station (so `verify` accepts a matching request and rejects a mismatched +/// one), and are forwarded to the IC when the approved upgrade executes. +/// +/// The upgrade requests `wasm_memory_persistence = replace`: the test canister +/// is not built with Enhanced Orthogonal Persistence, and the IC rejects +/// `keep` for such modules. +#[test] +fn canister_upgrade_options_round_trip() { + let TestEnv { + mut env, + canister_ids, + .. + } = setup_new_env(); + + let (_dfx_user, _) = setup_dfx_user(&env, &canister_ids); + let other_user = user_test_id(1); + add_user(&env, other_user, vec![], canister_ids.station); + + // Create and install the test canister, so that it can be upgraded later. + let test_canister = create_canister(&env, canister_ids.station); + let module_bytes = get_canister_wasm("test_canister"); + let module_hash = hash(&module_bytes); + env.install_canister( + test_canister, + module_bytes.clone(), + vec![], + Some(canister_ids.station), + ); + + permit_change_operation(&env, &canister_ids); + set_four_eyes_on_change(&env, &canister_ids); + + let config = DfxOrbitTestConfig { + canister_ids: vec![(String::from("test"), test_canister)], + ..Default::default() + }; + + let mut wasm = NamedTempFile::new().unwrap(); + wasm.write_all(&module_bytes).unwrap(); + + let install_args = RequestCanisterInstallArgs { + canister: String::from("test"), + mode: CanisterInstallModeArgs::Upgrade, + wasm_memory_persistence: Some(WasmMemoryPersistenceArgs::Replace), + skip_pre_upgrade: true, + wasm: wasm.path().as_os_str().to_str().unwrap().to_string(), + argument: None, + arg_file: None, + asset_canister: None, + }; + + let request = dfx_orbit_test(&mut env, config, async { + let dfx_orbit = setup_dfx_orbit(canister_ids.station).await; + + let request = RequestArgs { + title: None, + summary: None, + action: RequestArgsActions::Canister(RequestCanisterArgs { + action: RequestCanisterActionArgs::Install(install_args.clone()), + }), + } + .into_request(&dfx_orbit) + .await + .unwrap(); + + let response = dfx_orbit.station.request(request.clone()).await.unwrap(); + + let review_response = dfx_orbit + .station + .review_id(GetRequestInput { + request_id: response.request.id.clone(), + with_full_info: Some(false), + }) + .await + .unwrap(); + + // The stored upgrade options match, so verification succeeds. + VerifyArgs { + request_id: response.request.id.clone(), + and_approve: false, + or_reject: false, + action: VerifyArgsAction::Canister(VerifyCanisterArgs { + action: VerifyCanisterActionArgs::Install(install_args.clone()), + }), + } + .verify(&dfx_orbit, &review_response) + .await + .unwrap(); + + // A mismatched `wasm_memory_persistence` must fail verification. + let mut mismatched_args = install_args.clone(); + mismatched_args.wasm_memory_persistence = Some(WasmMemoryPersistenceArgs::Keep); + VerifyArgs { + request_id: response.request.id.clone(), + and_approve: false, + or_reject: false, + action: VerifyArgsAction::Canister(VerifyCanisterArgs { + action: VerifyCanisterActionArgs::Install(mismatched_args), + }), + } + .verify(&dfx_orbit, &review_response) + .await + .expect_err("verification must reject a mismatched wasm_memory_persistence"); + + response.request + }); + + // The other user approves the request; the upgrade carrying the upgrade + // options executes successfully (a failure to forward the options to the + // IC would leave the request failed and `wait_for_request` erroring). + submit_request_approval( + &env, + other_user, + canister_ids.station, + request.clone(), + RequestApprovalStatusDTO::Approved, + ); + wait_for_request(&env, other_user, canister_ids.station, request).unwrap(); + + // The canister still runs the same module after the upgrade. + let status = canister_status(&env, Some(canister_ids.station), test_canister); + assert_eq!(status.module_hash, Some(module_hash)); +} + +/// Test upgrading a canister that runs a module declaring Enhanced Orthogonal +/// Persistence — the actual use case motivating `--wasm-memory-persistence`. +/// The IC refuses to upgrade such canisters unless +/// `wasm_memory_persistence = keep` is set (a safety check against +/// accidentally dropping their main memory), so the same upgrade request must +/// fail without the flag and execute successfully with it. +#[test] +fn canister_upgrade_with_wasm_memory_persistence_keep() { + let TestEnv { + mut env, + canister_ids, + .. + } = setup_new_env(); + + let (_dfx_user, _) = setup_dfx_user(&env, &canister_ids); + let other_user = user_test_id(1); + add_user(&env, other_user, vec![], canister_ids.station); + + // Mark the test module as supporting Enhanced Orthogonal Persistence and + // install it, so that upgrades require `wasm_memory_persistence = keep`. + // The bundled module is gzipped and must be decompressed before appending + // the custom section: the IC reads the uncompressed size of gzipped + // modules from the trailing bytes of the gzip stream, so bytes appended + // after the stream would corrupt the module. + let test_canister = create_canister(&env, canister_ids.station); + let mut module_bytes = Vec::new(); + GzDecoder::new(get_canister_wasm("test_canister").as_slice()) + .read_to_end(&mut module_bytes) + .unwrap(); + append_wasm_custom_section( + &mut module_bytes, + "icp:private enhanced-orthogonal-persistence", + &[], + ); + let module_hash = hash(&module_bytes); + env.install_canister( + test_canister, + module_bytes.clone(), + vec![], + Some(canister_ids.station), + ); + + permit_change_operation(&env, &canister_ids); + set_four_eyes_on_change(&env, &canister_ids); + + let config = DfxOrbitTestConfig { + canister_ids: vec![(String::from("test"), test_canister)], + ..Default::default() + }; + + let mut wasm = NamedTempFile::new().unwrap(); + wasm.write_all(&module_bytes).unwrap(); + + let install_args = |wasm_memory_persistence| RequestCanisterInstallArgs { + canister: String::from("test"), + mode: CanisterInstallModeArgs::Upgrade, + wasm_memory_persistence, + skip_pre_upgrade: false, + wasm: wasm.path().as_os_str().to_str().unwrap().to_string(), + argument: None, + arg_file: None, + asset_canister: None, + }; + let args_without_keep = install_args(None); + let args_with_keep = install_args(Some(WasmMemoryPersistenceArgs::Keep)); + + let (request_without_keep, request_with_keep) = dfx_orbit_test(&mut env, config, async { + let dfx_orbit = setup_dfx_orbit(canister_ids.station).await; + + let mut requests = Vec::new(); + for args in [args_without_keep, args_with_keep] { + let request = RequestArgs { + title: None, + summary: None, + action: RequestArgsActions::Canister(RequestCanisterArgs { + action: RequestCanisterActionArgs::Install(args), + }), + } + .into_request(&dfx_orbit) + .await + .unwrap(); + + let response = dfx_orbit.station.request(request).await.unwrap(); + requests.push(response.request); + } + + let with_keep = requests.pop().unwrap(); + let without_keep = requests.pop().unwrap(); + (without_keep, with_keep) + }); + + // Without `keep`, the IC refuses to drop the main memory of a canister + // running an Enhanced Orthogonal Persistence module, so the upgrade fails. + submit_request_approval( + &env, + other_user, + canister_ids.station, + request_without_keep.clone(), + RequestApprovalStatusDTO::Approved, + ); + wait_for_request(&env, other_user, canister_ids.station, request_without_keep) + .expect_err("the IC must reject an EOP upgrade without wasm_memory_persistence = keep"); + + // With `keep`, the same upgrade executes successfully. + submit_request_approval( + &env, + other_user, + canister_ids.station, + request_with_keep.clone(), + RequestApprovalStatusDTO::Approved, + ); + wait_for_request(&env, other_user, canister_ids.station, request_with_keep).unwrap(); + + // The canister still runs the (EOP-marked) module after the upgrade. + let status = canister_status(&env, Some(canister_ids.station), test_canister); + assert_eq!(status.module_hash, Some(module_hash)); +} diff --git a/tools/dfx-orbit/README.md b/tools/dfx-orbit/README.md index 6824fba9d..744e48c1c 100644 --- a/tools/dfx-orbit/README.md +++ b/tools/dfx-orbit/README.md @@ -141,6 +141,32 @@ Then a verifier can verify this request, using: dfx-orbit verify [REQUEST_ID] canister install --mode upgrade [CANISTER_NAME] --wasm [WASM_PATH] ``` +##### Preserving main memory across upgrades + +When upgrading, you can control what happens to the canister's Wasm main memory +with `--wasm-memory-persistence`: + +- `keep` — preserve the main memory. This is **required** for Motoko canisters + that use [Enhanced Orthogonal Persistence](https://docs.internetcomputer.org/motoko/orthogonal-persistence/enhanced#ic-main-memory-retention) — + the IC refuses to upgrade them otherwise, as a safety check against + accidentally dropping their main memory. +- `replace` — clear the main memory (the IC default). + +You can also pass `--skip-pre-upgrade` to skip the canister's `pre_upgrade` +hook, which is useful for recovery when the existing hook traps. + +``` +dfx-orbit request canister install --mode upgrade [CANISTER_NAME] --wasm [WASM_PATH] --wasm-memory-persistence keep +``` + +Both flags are only valid together with `--mode upgrade`. When verifying, pass +the same flags so the verifier checks that the request carries the expected +upgrade options: + +``` +dfx-orbit verify [REQUEST_ID] canister install --mode upgrade [CANISTER_NAME] --wasm [WASM_PATH] --wasm-memory-persistence keep +``` + ### Upload assets to a canister We will assume that Orbit is a controller of the asset canister. diff --git a/tools/dfx-orbit/src/canister.rs b/tools/dfx-orbit/src/canister.rs index 8f672f9b7..a5c3c0469 100644 --- a/tools/dfx-orbit/src/canister.rs +++ b/tools/dfx-orbit/src/canister.rs @@ -7,10 +7,9 @@ mod install; mod settings; mod util; -pub use self::{ - call::RequestCanisterCallArgs, install::CanisterInstallModeArgs, - install::RequestCanisterInstallArgs, settings::RequestCanisterUpdateSettingsArgs, -}; +pub use call::RequestCanisterCallArgs; +pub use install::{CanisterInstallModeArgs, RequestCanisterInstallArgs, WasmMemoryPersistenceArgs}; +pub use settings::RequestCanisterUpdateSettingsArgs; // TODO: Support Canister create + integration test // TODO: Canister get response functionality diff --git a/tools/dfx-orbit/src/canister/install.rs b/tools/dfx-orbit/src/canister/install.rs index d4978e210..c9c9ba517 100644 --- a/tools/dfx-orbit/src/canister/install.rs +++ b/tools/dfx-orbit/src/canister/install.rs @@ -6,8 +6,9 @@ use clap::{Parser, ValueEnum}; use orbit_essentials::types::WasmModuleExtraChunks; use sha2::{Digest, Sha256}; use station_api::{ - CanisterInstallMode, ChangeExternalCanisterOperationDTO, ChangeExternalCanisterOperationInput, - GetRequestResponse, RequestOperationDTO, RequestOperationInput, + CanisterInstallMode, CanisterUpgradeOptionsInput, ChangeExternalCanisterOperationDTO, + ChangeExternalCanisterOperationInput, GetRequestResponse, RequestOperationDTO, + RequestOperationInput, WasmMemoryPersistence, }; use std::{collections::HashMap, fmt::Write, path::PathBuf}; @@ -19,6 +20,17 @@ pub struct RequestCanisterInstallArgs { /// The installation mode. #[clap(long, value_enum, rename_all = "kebab-case")] pub mode: CanisterInstallModeArgs, + /// Controls whether the canister's Wasm main memory is kept or replaced during + /// an upgrade. Only valid with `--mode upgrade`. Motoko canisters that use + /// Enhanced Orthogonal Persistence require `keep` — the IC refuses to upgrade + /// them otherwise. When omitted, the IC default (`replace`) applies. + #[clap(long, value_enum, rename_all = "kebab-case")] + pub wasm_memory_persistence: Option, + /// Skip the canister's `pre_upgrade` hook during an upgrade. Only valid with + /// `--mode upgrade`. Useful for recovery when the existing `pre_upgrade` hook + /// traps. + #[clap(long)] + pub skip_pre_upgrade: bool, /// The path to the wasm file to install (can also be a wasm.gz). #[clap(short, long)] pub wasm: String, @@ -57,7 +69,7 @@ impl RequestCanisterInstallArgs { let canister_id = dfx_orbit.canister_id(&self.canister)?; let (module, arg) = self.load_module_and_args()?; - let mode = self.mode.into(); + let mode = self.install_mode()?; let (module, module_extra_chunks) = if let Some(ref asset_canister) = self.asset_canister { let asset_canister_id = dfx_orbit.canister_id(asset_canister)?; @@ -127,8 +139,13 @@ impl RequestCanisterInstallArgs { op.canister_id ); } - if CanisterInstallModeArgs::from(op.mode.clone()) != self.mode { - bail!("Canister install mode {:?} does not match", op.mode); + let expected_mode = self.install_mode()?; + if op.mode != expected_mode { + bail!( + "Canister install mode {:?} does not match expected {:?}", + op.mode, + expected_mode + ); } if op.module_checksum != module_checksum { log_hashes( @@ -160,6 +177,52 @@ impl RequestCanisterInstallArgs { Ok((module, args)) } + + /// Builds the Orbit `CanisterInstallMode`, folding the upgrade-only flags + /// (`--wasm-memory-persistence`, `--skip-pre-upgrade`) into the upgrade + /// options record. The flags are rejected for `install`/`reinstall` since + /// the IC only honours them on upgrades. + fn install_mode(&self) -> anyhow::Result { + match self.mode { + CanisterInstallModeArgs::Install => { + self.err_if_upgrade_flags_were_passed()?; + Ok(CanisterInstallMode::Install) + } + CanisterInstallModeArgs::Reinstall => { + self.err_if_upgrade_flags_were_passed()?; + Ok(CanisterInstallMode::Reinstall) + } + CanisterInstallModeArgs::Upgrade => { + // Collapse to `None` when neither flag is set so the request stays + // byte-for-byte identical to a plain `--mode upgrade`. + let options = if self.wasm_memory_persistence.is_some() || self.skip_pre_upgrade { + let skip_pre_upgrade = if self.skip_pre_upgrade { + Some(true) + } else { + None + }; + Some(CanisterUpgradeOptionsInput { + wasm_memory_persistence: self + .wasm_memory_persistence + .map(WasmMemoryPersistence::from), + skip_pre_upgrade, + }) + } else { + None + }; + Ok(CanisterInstallMode::Upgrade(options)) + } + } + } + + fn err_if_upgrade_flags_were_passed(&self) -> anyhow::Result<()> { + if self.wasm_memory_persistence.is_some() || self.skip_pre_upgrade { + bail!( + "--wasm-memory-persistence and --skip-pre-upgrade are only valid with --mode upgrade" + ); + } + Ok(()) + } } /// Canister installation mode equivalent to `dfx canister install --mode XXX` and `orbit_station_api::CanisterInstallMode`. @@ -193,6 +256,26 @@ impl From for CanisterInstallModeArgs { } } +/// Whether the canister's Wasm main memory is kept or replaced on upgrade. +/// Equivalent to `dfx canister install --wasm-memory-persistence XXX` and +/// `orbit_station_api::WasmMemoryPersistence`. +#[derive(Copy, Clone, Eq, PartialEq, Debug, ValueEnum)] +pub enum WasmMemoryPersistenceArgs { + /// Keep the canister's main memory (required for Motoko Enhanced Orthogonal Persistence). + Keep, + /// Replace (clear) the canister's main memory. + Replace, +} + +impl From for WasmMemoryPersistence { + fn from(mode: WasmMemoryPersistenceArgs) -> Self { + match mode { + WasmMemoryPersistenceArgs::Keep => Self::Keep, + WasmMemoryPersistenceArgs::Replace => Self::Replace, + } + } +} + impl DfxOrbit { pub(crate) fn display_change_canister_operation( &self, @@ -213,6 +296,16 @@ impl DfxOrbit { }; writeln!(output, "Mode: {mode}")?; + if let CanisterInstallMode::Upgrade(Some(options)) = &op.mode { + if let Some(persistence) = &options.wasm_memory_persistence { + let persistence = format!("{persistence:?}").to_lowercase(); + writeln!(output, "Wasm memory persistence: {persistence}")?; + } + if let Some(skip_pre_upgrade) = options.skip_pre_upgrade { + writeln!(output, "Skip pre-upgrade: {skip_pre_upgrade}")?; + } + } + writeln!(output, "Module checksum: {}", &op.module_checksum)?; if let Some(arg_checksum) = &op.arg_checksum { writeln!(output, "Argument checksum: {arg_checksum}")?; @@ -220,3 +313,120 @@ impl DfxOrbit { Ok(()) } } + +#[cfg(test)] +mod tests { + #![allow(clippy::unwrap_used)] + + use super::*; + + fn args(mode: CanisterInstallModeArgs) -> RequestCanisterInstallArgs { + RequestCanisterInstallArgs { + canister: "test".to_string(), + mode, + wasm_memory_persistence: None, + skip_pre_upgrade: false, + wasm: "test.wasm".to_string(), + argument: None, + arg_file: None, + asset_canister: None, + } + } + + #[test] + fn plain_modes_produce_no_upgrade_options() { + assert_eq!( + args(CanisterInstallModeArgs::Install) + .install_mode() + .unwrap(), + CanisterInstallMode::Install + ); + assert_eq!( + args(CanisterInstallModeArgs::Reinstall) + .install_mode() + .unwrap(), + CanisterInstallMode::Reinstall + ); + assert_eq!( + args(CanisterInstallModeArgs::Upgrade) + .install_mode() + .unwrap(), + CanisterInstallMode::Upgrade(None) + ); + } + + #[test] + fn upgrade_forwards_wasm_memory_persistence() { + let mut args = args(CanisterInstallModeArgs::Upgrade); + args.wasm_memory_persistence = Some(WasmMemoryPersistenceArgs::Keep); + + assert_eq!( + args.install_mode().unwrap(), + CanisterInstallMode::Upgrade(Some(CanisterUpgradeOptionsInput { + wasm_memory_persistence: Some(WasmMemoryPersistence::Keep), + skip_pre_upgrade: None, + })) + ); + } + + #[test] + fn upgrade_forwards_skip_pre_upgrade() { + let mut args = args(CanisterInstallModeArgs::Upgrade); + args.skip_pre_upgrade = true; + + assert_eq!( + args.install_mode().unwrap(), + CanisterInstallMode::Upgrade(Some(CanisterUpgradeOptionsInput { + wasm_memory_persistence: None, + skip_pre_upgrade: Some(true), + })) + ); + } + + #[test] + fn upgrade_forwards_both_flags() { + let mut args = args(CanisterInstallModeArgs::Upgrade); + args.wasm_memory_persistence = Some(WasmMemoryPersistenceArgs::Replace); + args.skip_pre_upgrade = true; + + assert_eq!( + args.install_mode().unwrap(), + CanisterInstallMode::Upgrade(Some(CanisterUpgradeOptionsInput { + wasm_memory_persistence: Some(WasmMemoryPersistence::Replace), + skip_pre_upgrade: Some(true), + })) + ); + } + + #[test] + fn upgrade_flags_rejected_for_install_and_reinstall() { + const EXPECTED_ERROR: &str = + "--wasm-memory-persistence and --skip-pre-upgrade are only valid with --mode upgrade"; + + for mode in [ + CanisterInstallModeArgs::Install, + CanisterInstallModeArgs::Reinstall, + ] { + let mut install_args_with_wasm_memory_persistence = args(mode); + install_args_with_wasm_memory_persistence.wasm_memory_persistence = + Some(WasmMemoryPersistenceArgs::Keep); + let error = install_args_with_wasm_memory_persistence + .install_mode() + .unwrap_err(); + assert!( + error.to_string().contains(EXPECTED_ERROR), + "unexpected error: {error}" + ); + + let mut install_args_with_skip_pre_upgrade = args(mode); + install_args_with_skip_pre_upgrade.skip_pre_upgrade = true; + let error = install_args_with_skip_pre_upgrade + .install_mode() + .unwrap_err(); + assert!( + error.to_string().contains(EXPECTED_ERROR), + "unexpected error: {error}" + ); + } + } +}