fix(paid-service): persist init args across canister upgrades

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<InitArgs>) {
@@ -20,6 +19,40 @@ fn init(init_args: Option<InitArgs>) {
     }
 }
 
+/// 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<InitArgs>,)`, so we fall back to `None` rather than trapping and aborting the upgrade.
+#[post_upgrade]
+fn post_upgrade(init_args: Option<InitArgs>) {
+    let init_args = init_args.or_else(|| {
+        ic_cdk::storage::stable_restore::<(Option<InitArgs>,)>()
+            .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()

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<InitArgs> {
+    INIT_ARGS.with(|init_args| init_args.borrow().clone())
+}

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;

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<String, PaymentError> = 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);
+}

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<InitArgs>) {
+        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(