Merge branch 'ulan/follow-1187-rc--2023-08-09_23-01' into 'rc--2023-0…

…8-09_23-01'

[hotfix] FOLLOW-1187: Fix the `apply_balance_changes()` function

This fixes a bug in the function that may cause incorrect bookkeeping
of consumed cycles due to saturating arithmetic operations. 

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

rs/execution_environment/src/execution/response/tests.rs

@@ -2505,3 +2505,56 @@ fn subnet_available_memory_does_not_change_on_cleanup_resume_failure() {
         test.subnet_available_memory().get_execution_memory()
     );
 }
+
+#[test]
+fn cycles_balance_changes_applied_correctly() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let a_id = test
+        .universal_canister_with_cycles(Cycles::new(10_000_000_000_000))
+        .unwrap();
+    let b_id = test
+        .universal_canister_with_cycles(Cycles::new(81_000_000_000))
+        .unwrap();
+
+    test.ingress(
+        b_id,
+        "update",
+        wasm()
+            .call_with_cycles(
+                a_id.get(),
+                "update",
+                call_args().other_side(wasm().accept_cycles(Cycles::new(u128::MAX))),
+                Cycles::new(60_000_000_000),
+            )
+            .build(),
+    )
+    .unwrap();
+
+    let mut b = wasm().accept_cycles(Cycles::new(u128::MAX));
+
+    for _ in 0..400 {
+        b = b.call_simple(a_id, "update", call_args());
+    }
+
+    let b = b.push_int(42).reply_int().build();
+
+    let a = wasm()
+        .call_with_cycles(
+            b_id.get(),
+            "update",
+            call_args().other_side(b.clone()),
+            Cycles::new(5_000_000_000_000),
+        )
+        .build();
+    let a_balance_old = test.canister_state(a_id).system_state.balance();
+    let b_balance_old = test.canister_state(b_id).system_state.balance();
+    let res = test.ingress(a_id, "update", a).unwrap();
+    match res {
+        WasmResult::Reply(_) => {}
+        WasmResult::Reject(msg) => unreachable!("rejected : {}", msg),
+    }
+    let a_balance_new = test.canister_state(a_id).system_state.balance();
+    let b_balance_new = test.canister_state(b_id).system_state.balance();
+
+    assert!(a_balance_old + b_balance_old > a_balance_new + b_balance_new);
+}

rs/system_api/src/sandbox_safe_system_state.rs

@@ -422,42 +422,44 @@ impl SystemStateChanges {
         Ok(())
     }
 
-    /// Applies the balance change to the given state.
-    pub fn apply_balance_changes(&self, state: &mut SystemState) {
-        let balance_before = state.balance();
-        let mut removed_consumed_cycles = Cycles::new(0);
-        for (use_case, amount) in self.consumed_cycles_by_use_case.iter() {
-            state.remove_cycles(*amount, *use_case);
-            removed_consumed_cycles += *amount;
+    /// Returns `self.cycles_balance_change` without cycles that are accounted
+    /// for in `self.consumed_cycles_by_use_case`.
+    fn cycles_balance_change_without_consumed_cycles_by_use_case(&self) -> CyclesBalanceChange {
+        // `self.cycles_balance_change` consists of:
+        // - CyclesBalanceChange::added(cycles_accepted_from_the_call_context)
+        // - CyclesBalanceChange::remove(cycles_sent_via_outgoing_calls)
+        // - CyclesBalanceChange::remove(cycles_consumed_by_various_fees)
+        // This loop removes the last part from `self.cycles_balance_change`.
+        let mut result = self.cycles_balance_change;
+        for (_use_case, amount) in self.consumed_cycles_by_use_case.iter() {
+            result = result + CyclesBalanceChange::added(*amount)
         }
+        result
+    }
 
-        // The final balance of the canister should reflect `cycles_balance_change`.
-        // Since we removed `removed_consumed_cycles` above, we need to add it back
-        // to the `cycles_balance_change` to make sure the final balance of the canister is correct.
-        match self.cycles_balance_change {
+    /// Applies the balance change to the given state.
+    pub fn apply_balance_changes(&self, state: &mut SystemState) {
+        let initial_balance = state.balance();
+        let non_consumed_cycles_change =
+            self.cycles_balance_change_without_consumed_cycles_by_use_case();
+        match non_consumed_cycles_change {
             CyclesBalanceChange::Added(added) => {
-                // When 'cycles_balance_change' is positive we should add 'removed_consumed_cycles'.
-                state.add_cycles(added + removed_consumed_cycles, CyclesUseCase::NonConsumed);
-                debug_assert_eq!(balance_before + added, state.balance());
+                state.add_cycles(added, CyclesUseCase::NonConsumed)
             }
             CyclesBalanceChange::Removed(removed) => {
-                // When 'cycles_balance_change' is negative we are adding 'removed_consumed_cycles'
-                // but additionaly we should take care about the sign of the sum, which will
-                // determine whether we are adding or removing cycles from the balance.
-                if removed_consumed_cycles > removed {
-                    state.add_cycles(
-                        removed_consumed_cycles - removed,
-                        CyclesUseCase::NonConsumed,
-                    );
-                } else {
-                    state.remove_cycles(
-                        removed - removed_consumed_cycles,
-                        CyclesUseCase::NonConsumed,
-                    );
-                }
-                debug_assert_eq!(balance_before - removed, state.balance());
+                state.remove_cycles(removed, CyclesUseCase::NonConsumed)
             }
         }
+        for (use_case, amount) in self.consumed_cycles_by_use_case.iter() {
+            state.remove_cycles(*amount, *use_case);
+        }
+        // All changes applied above should be equivalent to simply applying
+        // `self.cycles_balance_change` to the initial balance.
+        let expected_balance = match self.cycles_balance_change {
+            CyclesBalanceChange::Added(added) => initial_balance + added,
+            CyclesBalanceChange::Removed(removed) => initial_balance - removed,
+        };
+        assert_eq!(state.balance(), expected_balance);
     }
 
     fn add_consumed_cycles(&mut self, consumed_cycles: &[(CyclesUseCase, Cycles)]) {