From 1a0b4198ba0ac46deaf5e4467028e0219db1a34d Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Tue, 10 Oct 2023 22:53:15 -0400
Subject: [PATCH] Correct the check for if we still need to set keys

The prior check had an edge case where once keys were pruned, it'd believe the
keys needed to be set ad infinitum.
---
 coordinator/src/main.rs                      | 31 +++++++++++++++-----
 coordinator/src/substrate/mod.rs             |  2 ++
 substrate/client/src/serai/validator_sets.rs | 10 ++++++-
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 0778fb16..e0891cf6 100644
--- a/coordinator/src/main.rs
+++ b/coordinator/src/main.rs
@@ -253,13 +253,30 @@ pub(crate) async fn scan_tributaries<
                         // creation
                         // TODO2: Differentiate connection errors from invariants
                         Err(e) => {
-                          // Check if this failed because the keys were already set by someone else
-                          // TODO: hash_with_keys is latest, yet we'll remove old keys from storage
-                          let hash_with_keys = serai.get_latest_block_hash().await.unwrap();
-                          if matches!(serai.get_keys(spec.set(), hash_with_keys).await, Ok(Some(_)))
-                          {
-                            log::info!("another coordinator set key pair for {:?}", set);
-                            break;
+                          if let Ok(latest) = serai.get_latest_block_hash().await {
+                            // Check if this failed because the keys were already set by someone
+                            // else
+                            if matches!(serai.get_keys(spec.set(), latest).await, Ok(Some(_))) {
+                              log::info!("another coordinator set key pair for {:?}", set);
+                              break;
+                            }
+
+                            // The above block may return false if the keys have been pruned from
+                            // the state
+                            // Check if this session is no longer the latest session, meaning it at
+                            // some point did set keys, and we're just operating off very
+                            // historical data
+                            if let Ok(Some(current_session)) =
+                              serai.get_session(spec.set().network, latest).await
+                            {
+                              if current_session.0 > spec.set().session.0 {
+                                log::warn!(
+                                  "trying to set keys for a set which isn't the latest {:?}",
+                                  set
+                                );
+                                break;
+                              }
+                            }
                           }
 
                           log::error!(
diff --git a/coordinator/src/substrate/mod.rs b/coordinator/src/substrate/mod.rs
index 60ca8d32..529273f8 100644
--- a/coordinator/src/substrate/mod.rs
+++ b/coordinator/src/substrate/mod.rs
@@ -243,6 +243,8 @@ async fn handle_block<D: Db, CNT: Clone + Fn(&mut D, TributarySpec), Pro: Proces
       panic!("NewSet event wasn't NewSet: {new_set:?}");
     };
 
+    // If this is Serai, do nothing
+    // We only coordinate/process external networks
     if set.network == NetworkId::Serai {
       continue;
     }
diff --git a/substrate/client/src/serai/validator_sets.rs b/substrate/client/src/serai/validator_sets.rs
index a00ada58..1c34ff68 100644
--- a/substrate/client/src/serai/validator_sets.rs
+++ b/substrate/client/src/serai/validator_sets.rs
@@ -2,7 +2,7 @@ use sp_core::sr25519::{Public, Signature};
 
 use serai_runtime::{validator_sets, ValidatorSets, Runtime};
 pub use validator_sets::primitives;
-use primitives::{ValidatorSet, KeyPair};
+use primitives::{Session, ValidatorSet, KeyPair};
 
 use subxt::utils::Encoded;
 
@@ -31,6 +31,14 @@ impl Serai {
       .await
   }
 
+  pub async fn get_session(
+    &self,
+    network: NetworkId,
+    at_hash: [u8; 32],
+  ) -> Result<Option<Session>, SeraiError> {
+    self.storage(PALLET, "CurrentSession", Some(vec![scale_value(network)]), at_hash).await
+  }
+
   pub async fn get_validator_set_participants(
     &self,
     network: NetworkId,