From 72b93ce39c9bd81f3cfb45bddff60bac30103d84 Mon Sep 17 00:00:00 2001 From: Antonio Ventilii Date: Wed, 15 Jul 2026 09:29:53 +0200 Subject: [PATCH 1/2] fix(paid_service): persist init args across canister upgrades MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The paid_service example stored its init args (including the payment ledger) only in a non-stable `thread_local` (`INIT_ARGS`), set solely in `#[init]`. With no `#[post_upgrade]` hook, the IC wiped this state on every upgrade, leaving it `None`. The `cost_1b` update method lazily builds `PAYMENT_GUARD`, whose initializer reads `payment_ledger()` and traps with "No init args provided" when the state is missing — a persistent failure of `cost_1b` after any upgrade until the canister is reinstalled. Add `#[pre_upgrade]`/`#[post_upgrade]` hooks that save the init args to stable memory and restore them, plus a regression test that upgrades the canister and asserts `cost_1b` still succeeds. --- src/example/paid_service/src/lib.rs | 27 +++++++++++-- src/example/paid_service/src/state.rs | 5 +++ src/example/paid_service/tests/it/main.rs | 1 + src/example/paid_service/tests/it/upgrade.rs | 40 +++++++++++++++++++ .../tests/it/util/test_environment.rs | 13 ++++++ 5 files changed, 83 insertions(+), 3 deletions(-) create mode 100644 src/example/paid_service/tests/it/upgrade.rs diff --git a/src/example/paid_service/src/lib.rs b/src/example/paid_service/src/lib.rs index d69f986..e0e2575 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,28 @@ 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 from stable memory after a canister upgrade. +#[post_upgrade] +fn post_upgrade() { + let (init_args,): (Option,) = + ic_cdk::storage::stable_restore().expect("Failed to restore init args from stable memory"); + 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..6542055 --- /dev/null +++ b/src/example/paid_service/tests/it/upgrade.rs @@ -0,0 +1,40 @@ +//! Regression test: 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::test_environment::{PaidMethods, TestSetup, LEDGER_FEE}; +use ic_papi_api::caller::CallerPaysIcrc2Tokens; +use ic_papi_api::cycles::cycles_ledger_canister_id; +use ic_papi_api::{PaymentError, PaymentType}; + +/// Verifies that `cost_1b` still works after the canister has been upgraded. +/// +/// 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() { + let setup = TestSetup::default(); + let method = PaidMethods::Cost1b; + + // Upgrade the canister. Without persistence, this would drop the init args. + setup.upgrade_paid_service(); + + // Pre-approve payment. + setup.user_approves_payment_for_paid_service(method.cost() + LEDGER_FEE); + + // The call must still succeed: the payment ledger config survived the upgrade. + 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", + ); +} 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..f9eefd5 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,19 @@ 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. + pub fn upgrade_paid_service(&self) { + 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"), + // The `post_upgrade` hook takes no arguments; it restores state from stable memory. + vec![], + None, + ) + .expect("Failed to upgrade the paid service canister"); + } /// Calls a paid service. #[allow(clippy::result_large_err)] pub fn call_paid_service( From e5ae1ae4656bf6117412971e86c38e7fc745e514 Mon Sep 17 00:00:00 2001 From: Antonio Ventilii Date: Wed, 15 Jul 2026 10:50:00 +0200 Subject: [PATCH 2/2] fix(paid-service): tolerate upgrades from un-persisted versions Address review feedback: the previous post_upgrade called stable_restore().expect(...) unconditionally, so upgrading from a version without pre_upgrade (e.g. current main) would find no valid candid payload in stable memory and trap, aborting the upgrade. Make post_upgrade: - accept an optional InitArgs argument (matching init's signature, so the candid interface is unchanged), used in preference to stable memory so operators can supply the ledger config while upgrading from an un-persisted version; and - tolerate a missing/malformed stable payload by falling back to None instead of trapping. Add a regression test for the explicit-args upgrade path. --- src/example/paid_service/src/lib.rs | 20 ++++++-- src/example/paid_service/tests/it/upgrade.rs | 46 ++++++++++++------- .../tests/it/util/test_environment.rs | 8 ++-- 3 files changed, 51 insertions(+), 23 deletions(-) diff --git a/src/example/paid_service/src/lib.rs b/src/example/paid_service/src/lib.rs index e0e2575..740cbf5 100644 --- a/src/example/paid_service/src/lib.rs +++ b/src/example/paid_service/src/lib.rs @@ -31,11 +31,23 @@ fn pre_upgrade() { .expect("Failed to save init args to stable memory"); } -/// Restores the init args from stable memory after a canister upgrade. +/// 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() { - let (init_args,): (Option,) = - ic_cdk::storage::stable_restore().expect("Failed to restore init args from stable memory"); +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); } diff --git a/src/example/paid_service/tests/it/upgrade.rs b/src/example/paid_service/tests/it/upgrade.rs index 6542055..2b40290 100644 --- a/src/example/paid_service/tests/it/upgrade.rs +++ b/src/example/paid_service/tests/it/upgrade.rs @@ -1,30 +1,19 @@ -//! Regression test: the paid service must remain usable after a canister upgrade. +//! 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}; -/// Verifies that `cost_1b` still works after the canister has been upgraded. -/// -/// 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() { - let setup = TestSetup::default(); +/// 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; - - // Upgrade the canister. Without persistence, this would drop the init args. - setup.upgrade_paid_service(); - - // Pre-approve payment. setup.user_approves_payment_for_paid_service(method.cost() + LEDGER_FEE); - - // The call must still succeed: the payment ledger config survived the upgrade. let response: Result = setup.call_paid_service( setup.user, method, @@ -38,3 +27,28 @@ fn cost_1b_works_after_upgrade() { "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 f9eefd5..c369ec9 100644 --- a/src/example/paid_service/tests/it/util/test_environment.rs +++ b/src/example/paid_service/tests/it/util/test_environment.rs @@ -209,14 +209,16 @@ impl TestSetup { .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. - pub fn upgrade_paid_service(&self) { + /// + /// `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"), - // The `post_upgrade` hook takes no arguments; it restores state from stable memory. - vec![], + encode_one(init_args).expect("Failed to encode the upgrade args"), None, ) .expect("Failed to upgrade the paid service canister");