From 862860ecc89e077331b72eb66cd7305b931fb9e8 Mon Sep 17 00:00:00 2001
From: Guillaume Thiolliere <gui.thiolliere@gmail.com>
Date: Sat, 10 Aug 2024 00:06:59 +0900
Subject: [PATCH] StorageWeightReclaim: Fix issue when underestimating refund.
 (#5273)

The code do reduce or increase the weight by comparing
`benchmarked_weight` and `consumed_weight`.

But `benchmarked_weight` is the pre dispatch weight. not the post
dispatch weight that is actually written into the block weight by
`CheckWeight`.

So in case the consumed weight was: `pre dispatch weight > consumed
weight > post dispatch weight` then the reclaim code was reducing the
block weight instead of increasing it.

Might explain this issue even better
https://github.com/paritytech/polkadot-sdk/issues/5229

@skunert
@s0me0ne-unkn0wn
---
 .../storage-weight-reclaim/src/lib.rs         | 48 +++++++++++++++++--
 prdoc/pr_5273.prdoc                           | 10 ++++
 2 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 prdoc/pr_5273.prdoc

diff --git a/cumulus/primitives/storage-weight-reclaim/src/lib.rs b/cumulus/primitives/storage-weight-reclaim/src/lib.rs
index f48dd927ee9..5984fa77a2c 100644
--- a/cumulus/primitives/storage-weight-reclaim/src/lib.rs
+++ b/cumulus/primitives/storage-weight-reclaim/src/lib.rs
@@ -165,15 +165,14 @@ where
 			);
 			return Ok(())
 		};
-		let benchmarked_weight = info.weight.proof_size();
-		let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);
-
 		// Unspent weight according to the `actual_weight` from `PostDispatchInfo`
 		// This unspent weight will be refunded by the `CheckWeight` extension, so we need to
 		// account for that.
 		let unspent = post_info.calc_unspent(info).proof_size();
-		let storage_size_diff =
-			benchmarked_weight.saturating_sub(unspent).abs_diff(consumed_weight as u64);
+		let benchmarked_weight = info.weight.proof_size().saturating_sub(unspent);
+		let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);
+
+		let storage_size_diff = benchmarked_weight.abs_diff(consumed_weight as u64);
 
 		// This value will be reclaimed by [`frame_system::CheckWeight`], so we need to calculate
 		// that in.
@@ -294,6 +293,45 @@ mod tests {
 		})
 	}
 
+	#[test]
+	fn underestimating_refund() {
+		// We fixed a bug where `pre dispatch info weight > consumed weight > post info weight`
+		// resulted in error.
+
+		// The real cost will be 100 bytes of storage size
+		let mut test_ext = setup_test_externalities(&[0, 100]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 500
+			let info = DispatchInfo { weight: Weight::from_parts(0, 101), ..Default::default() };
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(0, 99)),
+				pays_fee: Default::default(),
+			};
+
+			assert_ok!(CheckWeight::<Test>::do_pre_dispatch(&info, LEN));
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(0));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			// We expect an accrue of 1
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(get_storage_weight().total().proof_size(), 1250);
+		})
+	}
+
 	#[test]
 	fn does_nothing_without_extension() {
 		let mut test_ext = new_test_ext();
diff --git a/prdoc/pr_5273.prdoc b/prdoc/pr_5273.prdoc
new file mode 100644
index 00000000000..981172c6c13
--- /dev/null
+++ b/prdoc/pr_5273.prdoc
@@ -0,0 +1,10 @@
+title: Fix storage weight reclaim bug.
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      A bug in storage weight reclaim signed extension is fixed. The bug was causing an underestimate of the proof size when the post dispatch info was underestimating the proof size and the pre dispatch info was overestimating the proof size at the same time.
+
+crates:
+  - name: cumulus-primitives-storage-weight-reclaim
+    bump: patch
-- 
GitLab