From d57992040e7d7b0f445ed40996e5dafecd985584 Mon Sep 17 00:00:00 2001
From: pgherveou <pgherveou@gmail.com>
Date: Mon, 25 Sep 2023 13:37:20 +0200
Subject: [PATCH] Add integrity check

---
 substrate/frame/contracts/src/migration.rs | 26 ++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/substrate/frame/contracts/src/migration.rs b/substrate/frame/contracts/src/migration.rs
index 3e3d6f37884c7..54d7cac3d227e 100644
--- a/substrate/frame/contracts/src/migration.rs
+++ b/substrate/frame/contracts/src/migration.rs
@@ -328,6 +328,32 @@ impl<T: Config, const TEST_ALL_STEPS: bool> OnRuntimeUpgrade for Migration<T, TE
 		);
 		Ok(Default::default())
 	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
+		if !TEST_ALL_STEPS {
+			return Ok(())
+		}
+
+		log::info!(target: LOG_TARGET, "=== POST UPGRADE CHECKS ===");
+
+		// Ensure that that the hashing algorithm is correct for all storage maps.
+		if let Some(hash) = crate::CodeInfoOf::<T>::iter_keys().next() {
+			crate::CodeInfoOf::<T>::get(hash).expect("CodeInfo exists for hash; qed");
+		}
+		if let Some(hash) = crate::PristineCode::<T>::iter_keys().next() {
+			crate::PristineCode::<T>::get(hash).expect("PristineCode exists for hash; qed");
+		}
+		if let Some(account_id) = crate::ContractInfoOf::<T>::iter_keys().next() {
+			crate::ContractInfoOf::<T>::get(account_id)
+				.expect("ContractInfo exists for account_id; qed");
+		}
+		if let Some(nonce) = crate::DeletionQueue::<T>::iter_keys().next() {
+			crate::DeletionQueue::<T>::get(nonce).expect("DeletionQueue exists for nonce; qed");
+		}
+
+		Ok(())
+	}
 }
 
 /// The result of running the migration.