From ce2ca6230c6fa723cf5c5526f3db0b758514486e Mon Sep 17 00:00:00 2001
From: kylezs <zsembery.kyle@gmail.com>
Date: Wed, 23 Oct 2024 15:17:41 +0200
Subject: [PATCH] fix: handle status check and corrupt storage better

---
 state-chain/pallets/cf-elections/src/lib.rs   | 22 +++++++++++--------
 .../runtime/src/chainflip/solana_elections.rs | 17 +++++---------
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/state-chain/pallets/cf-elections/src/lib.rs b/state-chain/pallets/cf-elections/src/lib.rs
index c0af636382e..2ff650369ac 100644
--- a/state-chain/pallets/cf-elections/src/lib.rs
+++ b/state-chain/pallets/cf-elections/src/lib.rs
@@ -1677,12 +1677,7 @@ pub mod pallet {
 		}
 
 		pub fn ensure_initialized() -> Result<(), DispatchError> {
-			if Status::<T, I>::get().is_some() {
-				Ok(())
-			} else {
-				Self::deposit_event(Event::<T, I>::Uninitialized);
-				Err(Error::<T, I>::Uninitialized.into())
-			}
+			Self::with_status_check(|| Ok(()))
 		}
 
 		/// Provides access into the ElectoralSystem's current election
@@ -1695,11 +1690,20 @@ pub mod pallet {
 		>(
 			f: F,
 		) -> Result<R, DispatchError> {
-			if Status::<T, I>::get().is_some() {
+			Self::with_status_check(|| {
 				let mut election_identifiers =
 					ElectionProperties::<T, I>::iter_keys().collect::<Vec<_>>();
 				election_identifiers.sort();
-				Self::handle_corrupt_storage(f(election_identifiers)).map_err(Into::into)
+				f(election_identifiers)
+			})
+		}
+
+		// TODO: make handle_corrupt storage private
+		pub fn with_status_check<R, F: FnOnce() -> Result<R, CorruptStorageError>>(
+			f: F,
+		) -> Result<R, DispatchError> {
+			if Status::<T, I>::get().is_some() {
+				Self::handle_corrupt_storage(f()).map_err(Into::into)
 			} else {
 				Self::deposit_event(Event::<T, I>::Uninitialized);
 				Err(Error::<T, I>::Uninitialized.into())
@@ -2058,7 +2062,7 @@ pub mod pallet {
 			Ok(*election_identifier.unique_monotonic())
 		}
 
-		pub fn handle_corrupt_storage<Ok>(
+		fn handle_corrupt_storage<Ok>(
 			result: Result<Ok, CorruptStorageError>,
 		) -> Result<Ok, Error<T, I>> {
 			match result {
diff --git a/state-chain/runtime/src/chainflip/solana_elections.rs b/state-chain/runtime/src/chainflip/solana_elections.rs
index 50c9c489a98..e3a976293dc 100644
--- a/state-chain/runtime/src/chainflip/solana_elections.rs
+++ b/state-chain/runtime/src/chainflip/solana_elections.rs
@@ -418,18 +418,15 @@ impl SolanaNonceWatch for SolanaNonceTrackingTrigger {
 		nonce_account: SolAddress,
 		previous_nonce_value: SolHash,
 	) -> DispatchResult {
-		// TODO: Check if safe. We are not checking if initialised or not here - we were before in
-		// with_electoral_access
-		// TODO: Look at handling the corrupt storage elsewhere.
-		Ok(pallet_cf_elections::Pallet::<Runtime, SolanaInstance>::handle_corrupt_storage(
+		Ok(pallet_cf_elections::Pallet::<Runtime, SolanaInstance>::with_status_check(|| {
 			SolanaNonceTracking::watch_for_change::<
 				DerivedElectoralAccess<
 					_,
 					SolanaNonceTracking,
 					RunnerStorageAccess<Runtime, SolanaInstance>,
 				>,
-			>(nonce_account, previous_nonce_value),
-		)?)
+			>(nonce_account, previous_nonce_value)
+		})?)
 	}
 }
 
@@ -439,16 +436,14 @@ impl ElectionEgressWitnesser for SolanaEgressWitnessingTrigger {
 	type Chain = SolanaCrypto;
 
 	fn watch_for_egress_success(signature: SolSignature) -> DispatchResult {
-		// TODO: Check if safe. We are not checking if initialised or not here - we were before in
-		// with_electoral_access
-		Ok(pallet_cf_elections::Pallet::<Runtime, SolanaInstance>::handle_corrupt_storage(
+		Ok(pallet_cf_elections::Pallet::<Runtime, SolanaInstance>::with_status_check(|| {
 			SolanaEgressWitnessing::watch_for_egress::<
 				DerivedElectoralAccess<
 					_,
 					SolanaEgressWitnessing,
 					RunnerStorageAccess<Runtime, SolanaInstance>,
 				>,
-			>(signature),
-		)?)
+			>(signature)
+		})?)
 	}
 }