Merge pull request #2075 from opentensor/fix/no-epoch-on-inconsistent-data

Epoch safeguards

Author: sam0x17

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -256,7 +256,9 @@ impl<T: Config> Pallet<T> {
                 log::warn!("Failed to reveal commits for subnet {netuid} due to error: {e:?}");
             };
             // Pass on subnets that have not reached their tempo.
-            if Self::should_run_epoch(netuid, current_block) {
+            if Self::should_run_epoch(netuid, current_block)
+                && Self::is_epoch_input_state_consistent(netuid)
+            {
                 // Restart counters.
                 BlocksSinceLastStep::<T>::insert(netuid, 0);
                 LastMechansimStepBlock::<T>::insert(netuid, current_block);

pallets/subtensor/src/epoch/run_epoch.rs

@@ -1,6 +1,6 @@
 use super::*;
 use crate::epoch::math::*;
-use alloc::collections::BTreeMap;
+use alloc::collections::{BTreeMap, BTreeSet};
 use frame_support::IterableStorageDoubleMap;
 use safe_math::*;
 use sp_std::collections::btree_map::IntoIter;
@@ -1625,4 +1625,20 @@ impl<T: Config> Pallet<T> {
 
         Ok(())
     }
+
+    /// This function ensures major assumptions made by epoch function:
+    ///   1. Keys map has no duplicate hotkeys
+    ///
+    pub fn is_epoch_input_state_consistent(netuid: NetUid) -> bool {
+        // Check if Keys map has duplicate hotkeys or uids
+        let mut hotkey_set: BTreeSet<T::AccountId> = BTreeSet::new();
+        // `iter_prefix` over a double map yields (uid, value) for the given first key.
+        for (_uid, hotkey) in Keys::<T>::iter_prefix(netuid) {
+            if !hotkey_set.insert(hotkey) {
+                log::error!("Duplicate hotkeys detected for netuid {netuid}");
+                return false;
+            }
+        }
+        true
+    }
 }

pallets/subtensor/src/tests/epoch.rs

@@ -3884,3 +3884,55 @@ fn test_last_update_size_mismatch() {
         assert_eq!(SubtensorModule::get_dividends_for_uid(netuid, uid), 0);
     });
 }
+
+#[test]
+fn empty_ok() {
+    new_test_ext(1).execute_with(|| {
+        let netuid: NetUid = 155.into();
+        assert!(Pallet::<Test>::is_epoch_input_state_consistent(netuid));
+    });
+}
+
+#[test]
+fn unique_hotkeys_and_uids_ok() {
+    new_test_ext(1).execute_with(|| {
+        let netuid: NetUid = 155.into();
+
+        // (netuid, uid) -> hotkey (AccountId = U256)
+        Keys::<Test>::insert(netuid, 0u16, U256::from(1u64));
+        Keys::<Test>::insert(netuid, 1u16, U256::from(2u64));
+        Keys::<Test>::insert(netuid, 2u16, U256::from(3u64));
+
+        assert!(Pallet::<Test>::is_epoch_input_state_consistent(netuid));
+    });
+}
+
+#[test]
+fn duplicate_hotkey_within_same_netuid_fails() {
+    new_test_ext(1).execute_with(|| {
+        let netuid: NetUid = 155.into();
+
+        // Same hotkey mapped from two different UIDs in the SAME netuid
+        let hk = U256::from(42u64);
+        Keys::<Test>::insert(netuid, 0u16, hk);
+        Keys::<Test>::insert(netuid, 1u16, U256::from(42u64)); // duplicate hotkey
+
+        assert!(!Pallet::<Test>::is_epoch_input_state_consistent(netuid));
+    });
+}
+
+#[test]
+fn same_hotkey_across_different_netuids_is_ok() {
+    new_test_ext(1).execute_with(|| {
+        let net_a: NetUid = 10.into();
+        let net_b: NetUid = 11.into();
+
+        // Same hotkey appears once in each netuid — each net checks independently.
+        let hk = U256::from(777u64);
+        Keys::<Test>::insert(net_a, 0u16, hk);
+        Keys::<Test>::insert(net_b, 0u16, hk);
+
+        assert!(Pallet::<Test>::is_epoch_input_state_consistent(net_a));
+        assert!(Pallet::<Test>::is_epoch_input_state_consistent(net_b));
+    });
+}