diff --git a/src/example/paid_service/src/lib.rs b/src/example/paid_service/src/lib.rs index d69f986..740cbf5 100644 --- a/src/example/paid_service/src/lib.rs +++ b/src/example/paid_service/src/lib.rs @@ -1,8 +1,7 @@ mod state; use example_paid_service_api::InitArgs; -use ic_cdk::init; -use ic_cdk::{export_candid, update}; +use ic_cdk::{export_candid, init, post_upgrade, pre_upgrade, update}; use ic_papi_api::cycles::cycles_ledger_canister_id; use ic_papi_api::{PaymentError, PaymentType}; use ic_papi_guard::guards::PaymentGuardTrait; @@ -11,7 +10,7 @@ use ic_papi_guard::guards::{ caller_pays_icrc2_cycles::CallerPaysIcrc2CyclesPaymentGuard, caller_pays_icrc2_tokens::CallerPaysIcrc2TokensPaymentGuard, }; -use state::{set_init_args, PAYMENT_GUARD}; +use state::{get_init_args, set_init_args, PAYMENT_GUARD}; #[init] fn init(init_args: Option) { @@ -20,6 +19,40 @@ fn init(init_args: Option) { } } +/// Persists the init args to stable memory before a canister upgrade. +/// +/// The init args are held in a non-stable `thread_local` (see `state::INIT_ARGS`), which the IC +/// wipes on upgrade. Without this hook there is no `post_upgrade` counterpart to restore them, so +/// after any upgrade `payment_ledger()` — and therefore `cost_1b` — would trap with +/// "No init args provided". +#[pre_upgrade] +fn pre_upgrade() { + ic_cdk::storage::stable_save((get_init_args(),)) + .expect("Failed to save init args to stable memory"); +} + +/// Restores the init args after a canister upgrade. +/// +/// Resolution order: +/// 1. Args passed explicitly at upgrade time. This lets an operator upgrade from a version that +/// never persisted its args (e.g. one without `pre_upgrade`) and supply them in the same step. +/// 2. Args persisted to stable memory by `pre_upgrade`. +/// +/// Restoring from stable memory is tolerant of a missing or malformed payload: upgrading from a +/// version that did not run `pre_upgrade` leaves stable memory without a valid +/// `(Option,)`, so we fall back to `None` rather than trapping and aborting the upgrade. +#[post_upgrade] +fn post_upgrade(init_args: Option) { + let init_args = init_args.or_else(|| { + ic_cdk::storage::stable_restore::<(Option,)>() + .map(|(init_args,)| init_args) + .unwrap_or_default() + }); + if let Some(init_args) = init_args { + set_init_args(init_args); + } +} + #[update()] fn free() -> String { "Yes, I am free!".to_string() diff --git a/src/example/paid_service/src/state.rs b/src/example/paid_service/src/state.rs index ce688c2..41d2676 100644 --- a/src/example/paid_service/src/state.rs +++ b/src/example/paid_service/src/state.rs @@ -38,3 +38,8 @@ pub fn payment_ledger() -> Principal { pub fn set_init_args(init_args: InitArgs) { INIT_ARGS.set(Some(init_args)); } + +/// Returns the current init args, if any have been set. +pub fn get_init_args() -> Option { + INIT_ARGS.with(|init_args| init_args.borrow().clone()) +} diff --git a/src/example/paid_service/tests/it/main.rs b/src/example/paid_service/tests/it/main.rs index 30b991d..a3cf618 100644 --- a/src/example/paid_service/tests/it/main.rs +++ b/src/example/paid_service/tests/it/main.rs @@ -3,4 +3,5 @@ mod caller_pays_icrc2_cycles; mod caller_pays_icrc2_tokens; mod patron_pays_icrc2_cycles; mod patron_pays_icrc2_tokens; +mod upgrade; mod util; diff --git a/src/example/paid_service/tests/it/upgrade.rs b/src/example/paid_service/tests/it/upgrade.rs new file mode 100644 index 0000000..2b40290 --- /dev/null +++ b/src/example/paid_service/tests/it/upgrade.rs @@ -0,0 +1,54 @@ +//! Regression tests: the paid service must remain usable after a canister upgrade. +//! +//! The init args (which include the payment ledger) are held in a non-stable `thread_local`. +//! Without a `post_upgrade` hook to restore them, `cost_1b` traps with "No init args provided" +//! after any upgrade. See the `pre_upgrade`/`post_upgrade` hooks in `src/lib.rs`. +use crate::util::pic_canister::PicCanisterTrait; +use crate::util::test_environment::{PaidMethods, TestSetup, LEDGER_FEE}; +use example_paid_service_api::InitArgs; +use ic_papi_api::caller::CallerPaysIcrc2Tokens; +use ic_papi_api::cycles::cycles_ledger_canister_id; +use ic_papi_api::{PaymentError, PaymentType}; + +/// Drives a `cost_1b` call and asserts it succeeds, i.e. the payment ledger config is available. +fn assert_cost_1b_succeeds(setup: &TestSetup) { + let method = PaidMethods::Cost1b; + setup.user_approves_payment_for_paid_service(method.cost() + LEDGER_FEE); + let response: Result = setup.call_paid_service( + setup.user, + method, + PaymentType::CallerPaysIcrc2Tokens(CallerPaysIcrc2Tokens { + ledger: cycles_ledger_canister_id(), + }), + ); + assert_eq!( + response, + Ok("Yes, you paid 1 billion cycles!".to_string()), + "cost_1b should succeed after a canister upgrade", + ); +} + +/// Verifies that `cost_1b` still works after an upgrade that relies on stable memory. +/// +/// The upgrade happens before any call has lazily initialised the payment guard, so the guard is +/// first built (reading the init args) only after the upgrade. This is exactly the path that used +/// to trap when the init args were not persisted across upgrades. +#[test] +fn cost_1b_works_after_upgrade_restoring_from_stable_memory() { + let setup = TestSetup::default(); + // Upgrade with no explicit args: the init args must be restored from stable memory. + setup.upgrade_paid_service(None); + assert_cost_1b_succeeds(&setup); +} + +/// Verifies that init args supplied explicitly at upgrade time are used. +/// +/// This is the path an operator takes when upgrading from a version that never persisted its args +/// (e.g. one without `pre_upgrade`): the ledger config is provided in the upgrade itself. +#[test] +fn cost_1b_works_after_upgrade_with_explicit_args() { + let setup = TestSetup::default(); + let ledger = setup.ledger.canister_id(); + setup.upgrade_paid_service(Some(InitArgs { ledger })); + assert_cost_1b_succeeds(&setup); +} diff --git a/src/example/paid_service/tests/it/util/test_environment.rs b/src/example/paid_service/tests/it/util/test_environment.rs index 868aade..c369ec9 100644 --- a/src/example/paid_service/tests/it/util/test_environment.rs +++ b/src/example/paid_service/tests/it/util/test_environment.rs @@ -208,6 +208,21 @@ impl TestSetup { .expect("Failed to call the ledger to approve") .expect("Failed to approve the paid service to spend the user's ICRC-2 tokens"); } + /// Upgrades the paid service canister in place, exercising the pre/post-upgrade hooks. + /// + /// `init_args` are forwarded to the `post_upgrade` hook. Pass `None` to rely on the state + /// persisted to stable memory by `pre_upgrade`; pass `Some(..)` to supply the args explicitly. + pub fn upgrade_paid_service(&self, init_args: Option) { + self.pic + .upgrade_canister( + self.paid_service.canister_id(), + std::fs::read(PicCanister::cargo_wasm_path("example_paid_service")) + .expect("Could not read the paid service wasm"), + encode_one(init_args).expect("Failed to encode the upgrade args"), + None, + ) + .expect("Failed to upgrade the paid service canister"); + } /// Calls a paid service. #[allow(clippy::result_large_err)] pub fn call_paid_service(