diff --git a/node/src/mev_shield/proposer.rs b/node/src/mev_shield/proposer.rs index 7bacbf90e3..20fb6c790b 100644 --- a/node/src/mev_shield/proposer.rs +++ b/node/src/mev_shield/proposer.rs @@ -4,8 +4,8 @@ use ml_kem::kem::{Decapsulate, DecapsulationKey}; use ml_kem::{Ciphertext, Encoded, EncodedSizeUser, MlKem768, MlKem768Params}; use sc_service::SpawnTaskHandle; use sc_transaction_pool_api::{TransactionPool, TransactionSource}; -use sp_core::H256; -use sp_runtime::traits::Header; +use sp_core::{H256, sr25519}; +use sp_runtime::traits::{Header, SaturatedConversion}; use sp_runtime::{AccountId32, MultiSignature, OpaqueExtrinsic}; use std::{ collections::HashMap, @@ -14,6 +14,8 @@ use std::{ }; use tokio::time::sleep; +const KEY_FP_LEN: usize = 32; + /// Buffer of wrappers keyed by the block number in which they were included. #[derive(Default, Clone)] struct WrapperBuffer { @@ -57,7 +59,7 @@ impl WrapperBuffer { dropped_past = dropped_past.saturating_add(1); log::debug!( target: "mev-shield", - "revealer: dropping stale wrapper id=0x{} block_number={} < curr_block={}", + "revealer: dropping stale wrapper id=0x{} block_number={} < block={}", hex::encode(id.as_bytes()), *block_number, block @@ -99,7 +101,6 @@ pub fn spawn_revealer( Pool: TransactionPool + Send + Sync + 'static, { use codec::{Decode, Encode}; - use sp_runtime::traits::SaturatedConversion; type Address = sp_runtime::MultiAddress; type RUnchecked = node_subtensor_runtime::UncheckedExtrinsic; @@ -281,7 +282,7 @@ pub fn spawn_revealer( sleep(Duration::from_millis(tail_ms)).await; } - // Snapshot the current ML‑KEM secret (but *not* any epoch). + // Snapshot the current ML‑KEM secret. let snapshot_opt = match ctx.keys.lock() { Ok(k) => { let sk_hash = sp_core::hashing::blake2_256(&k.current_sk); @@ -313,12 +314,12 @@ pub fn spawn_revealer( } }; - // Use best block number as the “epoch” for which we reveal. + // Use best block number as the block whose submissions we reveal now. let curr_block: u64 = client.info().best_number.saturated_into(); log::debug!( target: "mev-shield", - "revealer: decrypt window start. block={} sk_len={} sk_hash=0x{} curr_pk_len={} next_pk_len={}", + "revealer: decrypt window start. reveal_block={} sk_len={} sk_hash=0x{} curr_pk_len={} next_pk_len={}", curr_block, curr_sk_bytes.len(), hex::encode(sk_hash), @@ -326,7 +327,7 @@ pub fn spawn_revealer( next_pk_len ); - // Only process wrappers whose originating block matches the current block. + // Only process wrappers whose originating block matches the reveal_block. let drained: Vec<(H256, u64, sp_runtime::AccountId32, Vec)> = match buffer.lock() { Ok(mut buf) => buf.drain_for_block(curr_block), @@ -341,17 +342,18 @@ pub fn spawn_revealer( log::debug!( target: "mev-shield", - "revealer: drained {} buffered wrappers for current block={}", + "revealer: drained {} buffered wrappers for reveal_block={}", drained.len(), curr_block ); - let mut to_submit: Vec<(H256, node_subtensor_runtime::RuntimeCall)> = Vec::new(); + let mut to_submit: Vec<(H256, node_subtensor_runtime::RuntimeCall)> = + Vec::new(); for (id, block_number, author, blob) in drained.into_iter() { log::debug!( target: "mev-shield", - "revealer: candidate id=0x{} block_number={} (curr_block={}) author={} blob_len={}", + "revealer: candidate id=0x{} submitted_in={} (reveal_block={}) author={} blob_len={}", hex::encode(id.as_bytes()), block_number, curr_block, @@ -360,80 +362,128 @@ pub fn spawn_revealer( ); // Safely parse blob: [u16 kem_len][kem_ct][nonce24][aead_ct] - let kem_len: usize = match blob - .get(0..2) - .and_then(|two| <[u8; 2]>::try_from(two).ok()) - { - Some(arr) => u16::from_le_bytes(arr) as usize, + if blob.len() < 2 { + log::debug!( + target: "mev-shield", + " id=0x{}: blob too short to contain kem_len", + hex::encode(id.as_bytes()) + ); + continue; + } + + let mut cursor: usize = 0; + + // 1) kem_len (u16 LE) + let kem_len_end = match cursor.checked_add(2usize) { + Some(e) => e, None => { log::debug!( target: "mev-shield", - " id=0x{}: blob too short or invalid length prefix", + " id=0x{}: kem_len range overflow", hex::encode(id.as_bytes()) ); continue; } }; - let kem_end = match 2usize.checked_add(kem_len) { - Some(v) => v, + let kem_len_slice = match blob.get(cursor..kem_len_end) { + Some(s) => s, None => { log::debug!( target: "mev-shield", - " id=0x{}: kem_len overflow", + " id=0x{}: blob too short for kem_len bytes (cursor={} end={})", + hex::encode(id.as_bytes()), + cursor, + kem_len_end + ); + continue; + } + }; + + let kem_len_bytes: [u8; 2] = match kem_len_slice.try_into() { + Ok(arr) => arr, + Err(_) => { + log::debug!( + target: "mev-shield", + " id=0x{}: kem_len slice not 2 bytes", hex::encode(id.as_bytes()) ); continue; } }; - let nonce_end = match kem_end.checked_add(24usize) { - Some(v) => v, + let kem_len = u16::from_le_bytes(kem_len_bytes) as usize; + cursor = kem_len_end; + + // 2) KEM ciphertext + let kem_ct_end = match cursor.checked_add(kem_len) { + Some(e) => e, None => { log::debug!( target: "mev-shield", - " id=0x{}: nonce range overflow", - hex::encode(id.as_bytes()) + " id=0x{}: kem_ct range overflow (cursor={} kem_len={})", + hex::encode(id.as_bytes()), + cursor, + kem_len ); continue; } }; - let kem_ct_bytes = match blob.get(2..kem_end) { + let kem_ct_bytes = match blob.get(cursor..kem_ct_end) { Some(s) => s, None => { log::debug!( target: "mev-shield", - " id=0x{}: blob too short for kem_ct (kem_len={}, total={})", + " id=0x{}: blob too short for kem_ct (cursor={} end={})", hex::encode(id.as_bytes()), - kem_len, - blob.len() + cursor, + kem_ct_end ); continue; } }; + cursor = kem_ct_end; - let nonce_bytes = match blob.get(kem_end..nonce_end) { - Some(s) if s.len() == 24 => s, - _ => { + // 3) Nonce (24 bytes) + const NONCE_LEN: usize = 24; + let nonce_end = match cursor.checked_add(NONCE_LEN) { + Some(e) => e, + None => { log::debug!( target: "mev-shield", - " id=0x{}: blob too short for 24-byte nonce (kem_len={}, total={})", + " id=0x{}: nonce range overflow (cursor={})", hex::encode(id.as_bytes()), - kem_len, - blob.len() + cursor ); continue; } }; - let aead_body = match blob.get(nonce_end..) { + let nonce_bytes = match blob.get(cursor..nonce_end) { Some(s) => s, None => { log::debug!( target: "mev-shield", - " id=0x{}: blob has no AEAD body", - hex::encode(id.as_bytes()) + " id=0x{}: blob too short for nonce24 (cursor={} end={})", + hex::encode(id.as_bytes()), + cursor, + nonce_end + ); + continue; + } + }; + cursor = nonce_end; + + // 4) AEAD body (rest) + let aead_body = match blob.get(cursor..) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: blob too short for aead_body (cursor={})", + hex::encode(id.as_bytes()), + cursor ); continue; } @@ -441,6 +491,7 @@ pub fn spawn_revealer( let kem_ct_hash = sp_core::hashing::blake2_256(kem_ct_bytes); let aead_body_hash = sp_core::hashing::blake2_256(aead_body); + log::debug!( target: "mev-shield", " id=0x{}: kem_len={} kem_ct_hash=0x{} nonce=0x{} aead_body_len={} aead_body_hash=0x{}", @@ -510,8 +561,9 @@ pub fn spawn_revealer( ss32.copy_from_slice(ss_bytes); let ss_hash = sp_core::hashing::blake2_256(&ss32); - let aead_key = crate::mev_shield::author::derive_aead_key(&ss32); - let key_hash = sp_core::hashing::blake2_256(&aead_key); + let aead_key = + crate::mev_shield::author::derive_aead_key(&ss32); + let key_hash_dbg = sp_core::hashing::blake2_256(&aead_key); log::debug!( target: "mev-shield", @@ -523,7 +575,7 @@ pub fn spawn_revealer( target: "mev-shield", " id=0x{}: derived AEAD key hash=0x{} (direct-from-ss)", hex::encode(id.as_bytes()), - hex::encode(key_hash) + hex::encode(key_hash_dbg) ); let mut nonce24 = [0u8; 24]; @@ -547,7 +599,7 @@ pub fn spawn_revealer( None => { log::debug!( target: "mev-shield", - " id=0x{}: AEAD decrypt FAILED with direct-from-ss key; ct_hash=0x{}", + " id=0x{}: AEAD decrypt FAILED; ct_hash=0x{}", hex::encode(id.as_bytes()), hex::encode(aead_body_hash), ); @@ -562,23 +614,26 @@ pub fn spawn_revealer( plaintext.len() ); - type RuntimeNonce = - ::Nonce; - // Safely parse plaintext layout without panics. - // Layout: signer (32) || nonce (4) || call (..) - // || sig_kind (1) || sig (64) + // + // Layout: + // signer (32) + // key_hash (32) == Hashing::hash(NextKey_bytes) at submit time + // call (..) + // sig_kind (1) + // sig (64) let min_plain_len: usize = 32usize - .saturating_add(4) - .saturating_add(1) - .saturating_add(1) - .saturating_add(64); + .saturating_add(KEY_FP_LEN) + .saturating_add(1usize) + .saturating_add(64usize); + if plaintext.len() < min_plain_len { log::debug!( target: "mev-shield", - " id=0x{}: plaintext too short ({}) for expected layout", + " id=0x{}: plaintext too short ({}) for expected layout (min={})", hex::encode(id.as_bytes()), - plaintext.len() + plaintext.len(), + min_plain_len ); continue; } @@ -595,33 +650,40 @@ pub fn spawn_revealer( } }; - let nonce_le = match plaintext.get(32..36) { - Some(s) => s, - None => { + let key_hash_raw = match plaintext.get(32..32usize.saturating_add(KEY_FP_LEN)) + { + Some(s) if s.len() == KEY_FP_LEN => s, + _ => { log::debug!( target: "mev-shield", - " id=0x{}: missing nonce bytes", + " id=0x{}: missing or malformed key_hash bytes", hex::encode(id.as_bytes()) ); continue; } }; - let sig_off = match plaintext.len().checked_sub(65) { - Some(off) if off >= 36 => off, + // sig_off = len - 65 (sig_kind + 64-byte sig) + let sig_min_offset: usize = + 32usize.saturating_add(KEY_FP_LEN); + + let sig_off = match plaintext.len().checked_sub(65usize) { + Some(off) if off >= sig_min_offset => off, _ => { log::debug!( target: "mev-shield", - " id=0x{}: invalid plaintext length for signature split", - hex::encode(id.as_bytes()) + " id=0x{}: invalid plaintext length for signature split (len={})", + hex::encode(id.as_bytes()), + plaintext.len() ); continue; } }; - let call_bytes = match plaintext.get(36..sig_off) { - Some(s) => s, - None => { + let call_start: usize = sig_min_offset; + let call_bytes = match plaintext.get(call_start..sig_off) { + Some(s) if !s.is_empty() => s, + _ => { log::debug!( target: "mev-shield", " id=0x{}: missing call bytes", @@ -643,24 +705,13 @@ pub fn spawn_revealer( } }; - let sig_start = match sig_off.checked_add(1) { - Some(v) => v, - None => { - log::debug!( - target: "mev-shield", - " id=0x{}: sig_start overflow", - hex::encode(id.as_bytes()) - ); - continue; - } - }; - - let sig_raw = match plaintext.get(sig_start..) { - Some(s) => s, - None => { + let sig_bytes_start = sig_off.saturating_add(1usize); + let sig_bytes = match plaintext.get(sig_bytes_start..) { + Some(s) if s.len() == 64 => s, + _ => { log::debug!( target: "mev-shield", - " id=0x{}: missing signature bytes", + " id=0x{}: signature bytes not 64 bytes", hex::encode(id.as_bytes()) ); continue; @@ -680,19 +731,9 @@ pub fn spawn_revealer( }; let signer = sp_runtime::AccountId32::new(signer_array); - let nonce_array: [u8; 4] = match nonce_le.try_into() { - Ok(a) => a, - Err(_) => { - log::debug!( - target: "mev-shield", - " id=0x{}: nonce bytes not 4 bytes", - hex::encode(id.as_bytes()) - ); - continue; - } - }; - let raw_nonce_u32 = u32::from_le_bytes(nonce_array); - let account_nonce: RuntimeNonce = raw_nonce_u32.saturated_into(); + let mut fp_array = [0u8; KEY_FP_LEN]; + fp_array.copy_from_slice(key_hash_raw); + let key_hash_h256 = H256(fp_array); let inner_call: node_subtensor_runtime::RuntimeCall = match Decode::decode(&mut &call_bytes[..]) { @@ -710,27 +751,27 @@ pub fn spawn_revealer( }; let signature: MultiSignature = - if sig_kind == 0x01 && sig_raw.len() == 64 { - let mut raw = [0u8; 64]; - raw.copy_from_slice(sig_raw); - MultiSignature::from(sp_core::sr25519::Signature::from_raw(raw)) + if sig_kind == 0x01 { + let mut raw_sig = [0u8; 64]; + raw_sig.copy_from_slice(sig_bytes); + MultiSignature::from(sr25519::Signature::from_raw(raw_sig)) } else { log::debug!( target: "mev-shield", " id=0x{}: unsupported signature format kind=0x{:02x}, len={}", hex::encode(id.as_bytes()), sig_kind, - sig_raw.len() + sig_bytes.len() ); continue; }; log::debug!( target: "mev-shield", - " id=0x{}: decrypted wrapper: signer={}, nonce={}, call={:?}", + " id=0x{}: decrypted wrapper: signer={}, key_hash=0x{}, call={:?}", hex::encode(id.as_bytes()), signer, - raw_nonce_u32, + hex::encode(key_hash_h256.as_bytes()), inner_call ); @@ -738,7 +779,7 @@ pub fn spawn_revealer( pallet_shield::Call::execute_revealed { id, signer: signer.clone(), - nonce: account_nonce, + key_hash: key_hash_h256.into(), call: Box::new(inner_call), signature, }, diff --git a/pallets/shield/src/benchmarking.rs b/pallets/shield/src/benchmarking.rs index 8e1d370e7f..5a82c6310d 100644 --- a/pallets/shield/src/benchmarking.rs +++ b/pallets/shield/src/benchmarking.rs @@ -1,21 +1,11 @@ use super::*; -use codec::Encode; use frame_benchmarking::v2::*; -use frame_system::RawOrigin; - use frame_support::{BoundedVec, pallet_prelude::ConstU32}; -use frame_system::pallet_prelude::BlockNumberFor; - -use sp_core::crypto::KeyTypeId; -use sp_core::sr25519; -use sp_io::crypto::{sr25519_generate, sr25519_sign}; - -use sp_runtime::{ - AccountId32, MultiSignature, - traits::{Hash as HashT, SaturatedConversion, Zero}, -}; - +use frame_system::{RawOrigin, pallet_prelude::BlockNumberFor}; +use sp_core::{crypto::KeyTypeId, sr25519}; +use sp_io::crypto::sr25519_generate; +use sp_runtime::{AccountId32, MultiSignature, traits::Hash as HashT}; use sp_std::{boxed::Box, vec, vec::Vec}; /// Helper to build bounded bytes (public key) of a given length. @@ -30,25 +20,6 @@ fn bounded_ct(len: usize) -> BoundedVec> { BoundedVec::>::try_from(v).expect("within bound; qed") } -/// Build the raw payload bytes used by `commitment` & signature verification in the pallet. -/// Layout: signer (32B) || nonce (u32 LE) || SCALE(call) -fn build_payload_bytes( - signer: &T::AccountId, - nonce: ::Nonce, - call: &::RuntimeCall, -) -> Vec { - let mut out = Vec::new(); - out.extend_from_slice(signer.as_ref()); - - // canonicalize nonce to u32 LE - let n_u32: u32 = nonce.saturated_into(); - out.extend_from_slice(&n_u32.to_le_bytes()); - - // append SCALE-encoded call - out.extend(call.encode()); - out -} - /// Seed Aura authorities so `EnsureAuraAuthority` passes for a given sr25519 pubkey. /// /// We avoid requiring `ByteArray` on `AuthorityId` by relying on: @@ -135,68 +106,87 @@ mod benches { /// Benchmark `execute_revealed`. #[benchmark] fn execute_revealed() { - // Generate a dev sr25519 key in the host keystore and derive the account. + use codec::Encode; + use frame_support::BoundedVec; + use sp_core::{crypto::KeyTypeId, sr25519}; + use sp_io::crypto::{sr25519_generate, sr25519_sign}; + use sp_runtime::traits::Zero; + + // 1) Generate a dev sr25519 key in the host keystore and derive the account. const KT: KeyTypeId = KeyTypeId(*b"benc"); let signer_pub: sr25519::Public = sr25519_generate(KT, Some("//Alice".as_bytes().to_vec())); let signer: AccountId32 = signer_pub.into(); - // Inner call that will be executed as the signer (cheap & always available). + // 2) Inner call that will be executed as the signer (cheap & always available). let inner_call: ::RuntimeCall = frame_system::Call::::remark { remark: vec![1, 2, 3], } .into(); - // Nonce must match current system nonce (fresh account => 0). - let nonce: ::Nonce = 0u32.into(); + // 3) Simulate the MEV‑Shield key epoch at the current block. + // + // In the real system, KeyHashByBlock[submitted_in] is filled by on_initialize + // as hash(CurrentKey). For the benchmark we just use a dummy value and + // insert it directly. + let submitted_in: BlockNumberFor = frame_system::Pallet::::block_number(); + let dummy_epoch_bytes: &[u8] = b"benchmark-epoch-key"; + let key_hash: ::Hash = + ::Hashing::hash(dummy_epoch_bytes); + KeyHashByBlock::::insert(submitted_in, key_hash); + + // 4) Build payload and commitment exactly how the pallet expects: + // payload = signer (32B) || key_hash (T::Hash bytes) || SCALE(call) + let mut payload_bytes = Vec::new(); + payload_bytes.extend_from_slice(signer.as_ref()); + payload_bytes.extend_from_slice(key_hash.as_ref()); + payload_bytes.extend(inner_call.encode()); - // Build payload and commitment exactly how the pallet expects. - let payload_bytes = super::build_payload_bytes::(&signer, nonce, &inner_call); let commitment: ::Hash = ::Hashing::hash(payload_bytes.as_slice()); - // Ciphertext is stored in the submission but not used by `execute_revealed`; keep small. + // 5) Ciphertext is stored in the submission but not used by `execute_revealed`; + // keep it small and arbitrary. const CT_DEFAULT_LEN: usize = 64; - let ciphertext: BoundedVec> = super::bounded_ct::<8192>(CT_DEFAULT_LEN); + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(vec![0u8; CT_DEFAULT_LEN]); // The submission `id` must match pallet's hashing scheme in submit_encrypted. let id: ::Hash = ::Hashing::hash_of( &(signer.clone(), commitment, &ciphertext), ); - // Seed the Submissions map with the expected entry. + // 6) Seed the Submissions map with the expected entry. let sub = Submission::, ::Hash> { author: signer.clone(), commitment, ciphertext: ciphertext.clone(), - submitted_in: frame_system::Pallet::::block_number(), + submitted_in, }; Submissions::::insert(id, sub); - // Domain-separated signing as in pallet: "mev-shield:v1" || genesis_hash || payload + // 7) Domain-separated signing as in pallet: + // "mev-shield:v1" || genesis_hash || payload let zero: BlockNumberFor = Zero::zero(); let genesis = frame_system::Pallet::::block_hash(zero); let mut msg = b"mev-shield:v1".to_vec(); msg.extend_from_slice(genesis.as_ref()); msg.extend_from_slice(&payload_bytes); - // Sign using the host keystore and wrap into MultiSignature. let sig = sr25519_sign(KT, &signer_pub, &msg).expect("signing should succeed in benches"); let signature: MultiSignature = sig.into(); - // Measure: dispatch the unsigned extrinsic (RawOrigin::None) with a valid wrapper. + // 8) Measure: dispatch the unsigned extrinsic (RawOrigin::None) with a valid wrapper. #[extrinsic_call] execute_revealed( RawOrigin::None, id, signer.clone(), - nonce, + key_hash, Box::new(inner_call.clone()), signature.clone(), ); - // Assert: submission consumed, signer nonce bumped to 1. + // 9) Assert: submission consumed. assert!(Submissions::::get(id).is_none()); - let new_nonce = frame_system::Pallet::::account_nonce(&signer); - assert_eq!(new_nonce, 1u32.into()); } } diff --git a/pallets/shield/src/lib.rs b/pallets/shield/src/lib.rs index 3fb44d9617..87e06949fe 100644 --- a/pallets/shield/src/lib.rs +++ b/pallets/shield/src/lib.rs @@ -28,8 +28,8 @@ pub mod pallet { InvalidTransaction, TransactionSource, ValidTransaction, }; use sp_runtime::{ - AccountId32, DispatchErrorWithPostInfo, MultiSignature, RuntimeDebug, - traits::{BadOrigin, Dispatchable, Hash, SaturatedConversion, Verify, Zero}, + AccountId32, DispatchErrorWithPostInfo, MultiSignature, RuntimeDebug, Saturating, + traits::{BadOrigin, Dispatchable, Hash, Verify}, }; use sp_std::{marker::PhantomData, prelude::*}; use subtensor_macros::freeze_struct; @@ -102,12 +102,15 @@ pub mod pallet { // ----------------- Storage ----------------- + /// Current ML‑KEM‑768 public key bytes (encoded form). #[pallet::storage] pub type CurrentKey = StorageValue<_, BoundedVec>, OptionQuery>; + /// Next ML‑KEM‑768 public key bytes, announced by the block author. #[pallet::storage] pub type NextKey = StorageValue<_, BoundedVec>, OptionQuery>; + /// Buffered encrypted submissions, indexed by wrapper id. #[pallet::storage] pub type Submissions = StorageMap< _, @@ -117,6 +120,14 @@ pub mod pallet { OptionQuery, >; + /// Hash(CurrentKey) per block, used to bind `key_hash` to the epoch at submit time. + #[pallet::storage] + pub type KeyHashByBlock = + StorageMap<_, Blake2_128Concat, BlockNumberFor, T::Hash, OptionQuery>; + + /// How many recent blocks of key-epoch hashes we retain. + const KEY_EPOCH_HISTORY: u32 = 100; + // ----------------- Events & Errors ----------------- #[pallet::event] @@ -135,23 +146,89 @@ pub mod pallet { #[pallet::error] pub enum Error { + /// A submission with the same id already exists in `Submissions`. SubmissionAlreadyExists, + /// The referenced submission id does not exist in `Submissions`. MissingSubmission, + /// The recomputed commitment does not match the stored commitment. CommitmentMismatch, + /// The provided signature over the payload is invalid. SignatureInvalid, - NonceMismatch, + /// The announced ML‑KEM public key length is invalid. BadPublicKeyLen, + /// The MEV‑Shield key epoch for this submission has expired and is no longer accepted. + KeyExpired, + /// The provided `key_hash` does not match the expected epoch key hash. + KeyHashMismatch, } // ----------------- Hooks ----------------- #[pallet::hooks] impl Hooks> for Pallet { - fn on_initialize(_n: BlockNumberFor) -> Weight { - if let Some(next) = >::take() { - >::put(&next); + fn on_initialize(n: BlockNumberFor) -> Weight { + let db_weight = T::DbWeight::get(); + let mut reads: u64 = 0; + let mut writes: u64 = 0; + + // 1) Roll NextKey -> CurrentKey if a next key is present. + reads = reads.saturating_add(1); + writes = writes.saturating_add(1); + let mut current_opt: Option>> = + if let Some(next) = NextKey::::take() { + CurrentKey::::put(&next); + writes = writes.saturating_add(1); + Some(next) + } else { + None + }; + + // 2) If we didn't roll, read the existing CurrentKey exactly once. + if current_opt.is_none() { + reads = reads.saturating_add(1); + current_opt = CurrentKey::::get(); + } + + // 3) Maintain KeyHashByBlock entry for this block: + match current_opt { + Some(current) => { + let epoch_hash: T::Hash = T::Hashing::hash(current.as_ref()); + KeyHashByBlock::::insert(n, epoch_hash); + writes = writes.saturating_add(1); + } + None => { + KeyHashByBlock::::remove(n); + writes = writes.saturating_add(1); + } } - T::DbWeight::get().reads_writes(1, 2) + + // 4) Prune old epoch hashes with a sliding window of size KEY_EPOCH_HISTORY. + let depth: BlockNumberFor = KEY_EPOCH_HISTORY.into(); + if n >= depth { + let prune_bn = n.saturating_sub(depth); + KeyHashByBlock::::remove(prune_bn); + writes = writes.saturating_add(1); + } + + // 5) TTL-based pruning of stale submissions. + let ttl: BlockNumberFor = KEY_EPOCH_HISTORY.into(); + let threshold: BlockNumberFor = n.saturating_sub(ttl); + + let mut to_remove: Vec = Vec::new(); + + for (id, sub) in Submissions::::iter() { + reads = reads.saturating_add(1); + if sub.submitted_in < threshold { + to_remove.push(id); + } + } + + for id in to_remove { + Submissions::::remove(id); + writes = writes.saturating_add(1); + } + + db_weight.reads_writes(reads, writes) } } @@ -159,6 +236,8 @@ pub mod pallet { #[pallet::call] impl Pallet { + /// Announce the ML‑KEM public key that will become `CurrentKey` in + /// the following block. #[pallet::call_index(0)] #[pallet::weight(( Weight::from_parts(9_979_000, 0) @@ -187,15 +266,29 @@ pub mod pallet { /// Users submit an encrypted wrapper. /// - /// `commitment` is `blake2_256(raw_payload)`, where: - /// raw_payload = signer || nonce || SCALE(call) + /// Client‑side: + /// + /// 1. Read `NextKey` (ML‑KEM public key bytes) from storage. + /// 2. Compute `key_hash = Hashing::hash(NextKey_bytes)`. + /// 3. Build: + /// + /// raw_payload = signer (32B AccountId) + /// || key_hash (32B Hash) + /// || SCALE(call) + /// + /// 4. `commitment = Hashing::hash(raw_payload)`. + /// 5. Signature message: + /// + /// "mev-shield:v1" || genesis_hash || raw_payload + /// + /// 6. Encrypt: + /// + /// plaintext = raw_payload || sig_kind || signature(64B) + /// + /// with ML‑KEM‑768 + XChaCha20‑Poly1305, producing + /// + /// ciphertext = [u16 kem_len] || kem_ct || nonce24 || aead_ct /// - /// `ciphertext` is constructed as: - /// [u16 kem_len] || kem_ct || nonce24 || aead_ct - /// where: - /// - `kem_ct` is the ML‑KEM ciphertext (encapsulated shared secret) - /// - `aead_ct` is XChaCha20‑Poly1305 over: - /// signer || nonce || SCALE(call) || sig_kind || signature #[pallet::call_index(1)] #[pallet::weight(( Weight::from_parts(13_980_000, 0) @@ -227,12 +320,21 @@ pub mod pallet { Ok(()) } - /// Executed by the block author. + /// Executed by the block author after decrypting a batch of wrappers. + /// + /// The author passes in: + /// + /// * `id` – wrapper id (hash of (author, commitment, ciphertext)) + /// * `signer` – account that should be treated as the origin of `call` + /// * `key_hash` – 32‑byte hash the client embedded (and signed) in the payload + /// * `call` – inner RuntimeCall to execute on behalf of `signer` + /// * `signature` – MultiSignature over the domain‑separated payload + /// #[pallet::call_index(2)] #[pallet::weight(( Weight::from_parts(77_280_000, 0) .saturating_add(T::DbWeight::get().reads(4_u64)) - .saturating_add(T::DbWeight::get().writes(2_u64)), + .saturating_add(T::DbWeight::get().writes(1_u64)), DispatchClass::Operational, Pays::No ))] @@ -241,38 +343,43 @@ pub mod pallet { origin: OriginFor, id: T::Hash, signer: T::AccountId, - nonce: T::Nonce, + key_hash: T::Hash, call: Box<::RuntimeCall>, signature: MultiSignature, ) -> DispatchResultWithPostInfo { + // Unsigned: only the author node may inject this via ValidateUnsigned. ensure_none(origin)?; + // 1) Load and consume the submission. let Some(sub) = Submissions::::take(id) else { return Err(Error::::MissingSubmission.into()); }; - let payload_bytes = Self::build_raw_payload_bytes(&signer, nonce, call.as_ref()); + // 2) Bind to the MEV‑Shield key epoch at submit time. + let expected_key_hash = + KeyHashByBlock::::get(sub.submitted_in).ok_or(Error::::KeyExpired)?; - // 1) Commitment check against on-chain stored commitment. + ensure!(key_hash == expected_key_hash, Error::::KeyHashMismatch); + + // 3) Rebuild the same payload bytes the client used for both + // commitment and signature. + let payload_bytes = Self::build_raw_payload_bytes(&signer, &key_hash, call.as_ref()); + + // 4) Commitment check against on-chain stored commitment. let recomputed: T::Hash = T::Hashing::hash(&payload_bytes); ensure!(sub.commitment == recomputed, Error::::CommitmentMismatch); - // 2) Signature check over the same payload. + // 5) Signature check over the same payload, with domain separation + // and genesis hash to make signatures chain‑bound. let genesis = frame_system::Pallet::::block_hash(BlockNumberFor::::zero()); let mut msg = b"mev-shield:v1".to_vec(); msg.extend_from_slice(genesis.as_ref()); msg.extend_from_slice(&payload_bytes); - ensure!( - signature.verify(msg.as_slice(), &signer), - Error::::SignatureInvalid - ); - // 3) Nonce check & bump. - let acc = frame_system::Pallet::::account_nonce(&signer); - ensure!(acc == nonce, Error::::NonceMismatch); - frame_system::Pallet::::inc_account_nonce(&signer); + let sig_ok = signature.verify(msg.as_slice(), &signer); + ensure!(sig_ok, Error::::SignatureInvalid); - // 4) Dispatch inner call from signer. + // 6) Dispatch inner call from signer. let info = call.get_dispatch_info(); let required = info.call_weight.saturating_add(info.extension_weight); @@ -301,26 +408,22 @@ pub mod pallet { impl Pallet { /// Build the raw payload bytes used for both: - /// - `commitment = blake2_256(raw_payload)` - /// - signature message (after domain separation). + /// + /// * `commitment = T::Hashing::hash(raw_payload)` + /// * signature message (after domain separation) /// /// Layout: - /// signer (32B) || nonce (u32 LE) || SCALE(call) + /// + /// signer (32B) || key_hash (T::Hash bytes) || SCALE(call) fn build_raw_payload_bytes( signer: &T::AccountId, - nonce: T::Nonce, + key_hash: &T::Hash, call: &::RuntimeCall, ) -> Vec { let mut out = Vec::new(); out.extend_from_slice(signer.as_ref()); - - // We canonicalise nonce to u32 LE for the payload. - let n_u32: u32 = nonce.saturated_into(); - out.extend_from_slice(&n_u32.to_le_bytes()); - - // Append SCALE-encoded call. + out.extend_from_slice(key_hash.as_ref()); out.extend(call.encode()); - out } } diff --git a/pallets/shield/src/tests.rs b/pallets/shield/src/tests.rs index e3bc630014..2a09bd43a4 100644 --- a/pallets/shield/src/tests.rs +++ b/pallets/shield/src/tests.rs @@ -6,17 +6,19 @@ use frame_support::pallet_prelude::ValidateUnsigned; use frame_support::traits::ConstU32 as FrameConstU32; use frame_support::traits::Hooks; use frame_support::{BoundedVec, assert_noop, assert_ok}; +use frame_system::pallet_prelude::BlockNumberFor; use pallet_mev_shield::{ - Call as MevShieldCall, CurrentKey, Event as MevShieldEvent, NextKey, Submissions, + Call as MevShieldCall, CurrentKey, Event as MevShieldEvent, KeyHashByBlock, NextKey, + Submissions, }; use sp_core::Pair; use sp_core::sr25519; -use sp_runtime::traits::Hash; -use sp_runtime::{ - AccountId32, MultiSignature, - traits::{SaturatedConversion, Zero}, - transaction_validity::TransactionSource, -}; +use sp_runtime::traits::{Hash, SaturatedConversion}; +use sp_runtime::{AccountId32, MultiSignature, transaction_validity::TransactionSource}; + +// Type aliases for convenience in tests. +type TestHash = ::Hash; +type TestBlockNumber = BlockNumberFor; // ----------------------------------------------------------------------------- // Helpers @@ -28,24 +30,25 @@ fn test_sr25519_pair() -> sr25519::Pair { } /// Reproduce the pallet's raw payload layout: -/// signer (32B) || nonce (u32 LE) || SCALE(call) +/// signer (32B) || key_hash (Hash bytes) || SCALE(call) fn build_raw_payload_bytes_for_test( signer: &AccountId32, - nonce: TestNonce, + key_hash: &TestHash, call: &RuntimeCall, ) -> Vec { let mut out = Vec::new(); out.extend_from_slice(signer.as_ref()); - - let n_u32: u32 = nonce.saturated_into(); - out.extend_from_slice(&n_u32.to_le_bytes()); - + out.extend_from_slice(key_hash.as_ref()); out.extend(call.encode()); out } +// ----------------------------------------------------------------------------- +// Tests +// ----------------------------------------------------------------------------- + #[test] -fn authority_can_announce_next_key_and_on_initialize_rolls_it() { +fn authority_can_announce_next_key_and_on_initialize_rolls_it_and_records_epoch_hash() { new_test_ext().execute_with(|| { const KYBER_PK_LEN: usize = 1184; let pk_bytes = vec![7u8; KYBER_PK_LEN]; @@ -78,13 +81,22 @@ fn authority_can_announce_next_key_and_on_initialize_rolls_it() { let next = NextKey::::get().expect("NextKey should be set"); assert_eq!(next, pk_bytes); - // Roll on new block - MevShield::on_initialize(2); + // Simulate beginning of block #2. + let block_two: TestBlockNumber = 2u64.saturated_into(); + MevShield::on_initialize(block_two); + // CurrentKey should now equal the previously announced NextKey. let curr = CurrentKey::::get().expect("CurrentKey should be set"); assert_eq!(curr, pk_bytes); + // And NextKey cleared. assert!(NextKey::::get().is_none()); + + // Key hash for this block should be recorded and equal hash(CurrentKey_bytes). + let expected_hash: TestHash = ::Hashing::hash(curr.as_ref()); + let recorded = + KeyHashByBlock::::get(block_two).expect("epoch key hash must be recorded"); + assert_eq!(recorded, expected_hash); }); } @@ -199,18 +211,22 @@ fn execute_revealed_happy_path_verifies_and_executes_inner_call() { remark: b"hello-mevshield".to_vec(), }); - let nonce: TestNonce = Zero::zero(); - assert_eq!(System::account_nonce(&signer), nonce); + // Choose a deterministic epoch key hash and wire it up for block #1. + let key_hash: TestHash = ::Hashing::hash(b"epoch-key"); + let payload_bytes = build_raw_payload_bytes_for_test(&signer, &key_hash, &inner_call); - let payload_bytes = build_raw_payload_bytes_for_test(&signer, nonce, &inner_call); - - let commitment = ::Hashing::hash(payload_bytes.as_ref()); + let commitment: TestHash = + ::Hashing::hash(payload_bytes.as_ref()); let ciphertext_bytes = vec![9u8, 9, 9, 9]; let ciphertext: BoundedVec> = BoundedVec::truncate_from(ciphertext_bytes.clone()); + // All submissions in this test happen at block #1. System::set_block_number(1); + let submitted_in = System::block_number(); + // Record epoch hash for that block, as on_initialize would do. + KeyHashByBlock::::insert(submitted_in, key_hash); // Wrapper author == signer for simplest path assert_ok!(MevShield::submit_encrypted( @@ -219,7 +235,7 @@ fn execute_revealed_happy_path_verifies_and_executes_inner_call() { ciphertext.clone(), )); - let id = ::Hashing::hash_of(&( + let id: TestHash = ::Hashing::hash_of(&( signer.clone(), commitment, &ciphertext, @@ -238,7 +254,7 @@ fn execute_revealed_happy_path_verifies_and_executes_inner_call() { RuntimeOrigin::none(), id, signer.clone(), - nonce, + key_hash, Box::new(inner_call.clone()), signature, ); @@ -248,10 +264,6 @@ fn execute_revealed_happy_path_verifies_and_executes_inner_call() { // Submission consumed assert!(Submissions::::get(id).is_none()); - // Nonce bumped once - let expected_nonce: TestNonce = (1u32).saturated_into(); - assert_eq!(System::account_nonce(&signer), expected_nonce); - // Last event is DecryptedExecuted let events = System::events(); let last = events @@ -273,24 +285,273 @@ fn execute_revealed_happy_path_verifies_and_executes_inner_call() { }); } +#[test] +fn execute_revealed_fails_on_key_hash_mismatch() { + new_test_ext().execute_with(|| { + let pair = test_sr25519_pair(); + let signer: AccountId32 = pair.public().into(); + + let inner_call = RuntimeCall::System(frame_system::Call::::remark { + remark: b"bad-key-hash".to_vec(), + }); + + System::set_block_number(5); + let submitted_in = System::block_number(); + + // Epoch hash recorded for this block: + let correct_key_hash: TestHash = + ::Hashing::hash(b"correct-epoch"); + KeyHashByBlock::::insert(submitted_in, correct_key_hash); + + // But we build payload & commitment with a *different* key_hash. + let wrong_key_hash: TestHash = + ::Hashing::hash(b"wrong-epoch"); + + let payload_bytes = build_raw_payload_bytes_for_test(&signer, &wrong_key_hash, &inner_call); + let commitment: TestHash = + ::Hashing::hash(payload_bytes.as_ref()); + + let ciphertext_bytes = vec![0u8; 4]; + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(ciphertext_bytes); + + assert_ok!(MevShield::submit_encrypted( + RuntimeOrigin::signed(signer.clone()), + commitment, + ciphertext.clone(), + )); + + let id: TestHash = ::Hashing::hash_of(&( + signer.clone(), + commitment, + &ciphertext, + )); + + let genesis = System::block_hash(0); + let mut msg = b"mev-shield:v1".to_vec(); + msg.extend_from_slice(genesis.as_ref()); + msg.extend_from_slice(&payload_bytes); + + let sig_sr25519 = pair.sign(&msg); + let signature: MultiSignature = sig_sr25519.into(); + + // execute_revealed should fail with KeyHashMismatch. + let res = MevShield::execute_revealed( + RuntimeOrigin::none(), + id, + signer.clone(), + wrong_key_hash, + Box::new(inner_call.clone()), + signature, + ); + assert_noop!(res, pallet_mev_shield::Error::::KeyHashMismatch); + }); +} + +#[test] +fn execute_revealed_rejects_replay_for_same_wrapper_id() { + new_test_ext().execute_with(|| { + let pair = test_sr25519_pair(); + let signer: AccountId32 = pair.public().into(); + + let inner_call = RuntimeCall::System(frame_system::Call::::remark { + remark: b"replay-test".to_vec(), + }); + + System::set_block_number(10); + let submitted_in = System::block_number(); + + let key_hash: TestHash = ::Hashing::hash(b"replay-epoch"); + KeyHashByBlock::::insert(submitted_in, key_hash); + + let payload_bytes = build_raw_payload_bytes_for_test(&signer, &key_hash, &inner_call); + let commitment: TestHash = + ::Hashing::hash(payload_bytes.as_ref()); + + let ciphertext_bytes = vec![7u8; 16]; + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(ciphertext_bytes.clone()); + + assert_ok!(MevShield::submit_encrypted( + RuntimeOrigin::signed(signer.clone()), + commitment, + ciphertext.clone(), + )); + + let id: TestHash = ::Hashing::hash_of(&( + signer.clone(), + commitment, + &ciphertext, + )); + + let genesis = System::block_hash(0); + let mut msg = b"mev-shield:v1".to_vec(); + msg.extend_from_slice(genesis.as_ref()); + msg.extend_from_slice(&payload_bytes); + + let sig_sr25519 = pair.sign(&msg); + let signature: MultiSignature = sig_sr25519.into(); + + // First execution succeeds. + assert_ok!(MevShield::execute_revealed( + RuntimeOrigin::none(), + id, + signer.clone(), + key_hash, + Box::new(inner_call.clone()), + signature.clone(), + )); + + // Second execution with the same id must fail with MissingSubmission. + let res = MevShield::execute_revealed( + RuntimeOrigin::none(), + id, + signer.clone(), + key_hash, + Box::new(inner_call.clone()), + signature, + ); + assert_noop!(res, pallet_mev_shield::Error::::MissingSubmission); + }); +} + +#[test] +fn key_hash_by_block_prunes_old_entries() { + new_test_ext().execute_with(|| { + // This must match the constant configured in the pallet. + const KEEP: u64 = 100; + const TOTAL: u64 = KEEP + 5; + + // For each block n, set a CurrentKey and call on_initialize(n), + // which will record KeyHashByBlock[n] and prune old entries. + for n in 1..=TOTAL { + let key_bytes = vec![n as u8; 32]; + let bounded: BoundedVec> = + BoundedVec::truncate_from(key_bytes.clone()); + + CurrentKey::::put(bounded.clone()); + + let bn: TestBlockNumber = n.saturated_into(); + MevShield::on_initialize(bn); + } + + // The oldest block that should still be kept after TOTAL blocks. + let oldest_kept: u64 = if TOTAL > KEEP { TOTAL - KEEP + 1 } else { 1 }; + + // Blocks strictly before oldest_kept must be pruned. + for old in 0..oldest_kept { + let bn: TestBlockNumber = old.saturated_into(); + assert!( + KeyHashByBlock::::get(bn).is_none(), + "block {bn:?} should have been pruned" + ); + } + + // Blocks from oldest_kept..=TOTAL must still have entries. + for recent in oldest_kept..=TOTAL { + let bn: TestBlockNumber = recent.saturated_into(); + assert!( + KeyHashByBlock::::get(bn).is_some(), + "block {bn:?} should be retained" + ); + } + + // Additionally, assert we never exceed the configured cap. + let mut count: u64 = 0; + for bn in 0..=TOTAL { + let bn_t: TestBlockNumber = bn.saturated_into(); + if KeyHashByBlock::::get(bn_t).is_some() { + count += 1; + } + } + let expected = KEEP.min(TOTAL); + assert_eq!( + count, expected, + "expected at most {expected} entries in KeyHashByBlock after pruning, got {count}" + ); + }); +} + +#[test] +fn submissions_pruned_after_ttl_window() { + new_test_ext().execute_with(|| { + // This must match KEY_EPOCH_HISTORY in the pallet. + const KEEP: u64 = 100; + const TOTAL: u64 = KEEP + 5; + + let pair = test_sr25519_pair(); + let who: AccountId32 = pair.public().into(); + + // Helper: create a submission at a specific block with a tagged commitment. + let make_submission = |block: u64, tag: &[u8]| -> TestHash { + System::set_block_number(block); + let commitment: TestHash = ::Hashing::hash(tag); + let ciphertext_bytes = vec![block as u8; 4]; + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(ciphertext_bytes); + + assert_ok!(MevShield::submit_encrypted( + RuntimeOrigin::signed(who.clone()), + commitment, + ciphertext.clone(), + )); + + ::Hashing::hash_of(&( + who.clone(), + commitment, + &ciphertext, + )) + }; + + // With n = TOTAL and depth = KEEP, prune_before = n - KEEP = 5. + let stale_block1: u64 = 1; // < 5, should be pruned + let stale_block2: u64 = 4; // < 5, should be pruned + let keep_block1: u64 = 5; // == prune_before, should be kept + let keep_block2: u64 = TOTAL; // latest, should be kept + + let id_stale1 = make_submission(stale_block1, b"stale-1"); + let id_stale2 = make_submission(stale_block2, b"stale-2"); + let id_keep1 = make_submission(keep_block1, b"keep-1"); + let id_keep2 = make_submission(keep_block2, b"keep-2"); + + // Sanity: all are present before pruning. + assert!(Submissions::::get(id_stale1).is_some()); + assert!(Submissions::::get(id_stale2).is_some()); + assert!(Submissions::::get(id_keep1).is_some()); + assert!(Submissions::::get(id_keep2).is_some()); + + // Run on_initialize at block TOTAL, triggering TTL pruning over Submissions. + let n_final: TestBlockNumber = TOTAL.saturated_into(); + MevShield::on_initialize(n_final); + + // Submissions with submitted_in < prune_before (5) should be gone. + assert!(Submissions::::get(id_stale1).is_none()); + assert!(Submissions::::get(id_stale2).is_none()); + + // Submissions at or after prune_before should remain. + assert!(Submissions::::get(id_keep1).is_some()); + assert!(Submissions::::get(id_keep2).is_some()); + }); +} + #[test] fn validate_unsigned_accepts_local_source_for_execute_revealed() { new_test_ext().execute_with(|| { let pair = test_sr25519_pair(); let signer: AccountId32 = pair.public().into(); - let nonce: TestNonce = Zero::zero(); let inner_call = RuntimeCall::System(frame_system::Call::::remark { remark: b"noop-local".to_vec(), }); - let id = ::Hashing::hash(b"mevshield-id-local"); + let id: TestHash = ::Hashing::hash(b"mevshield-id-local"); + let key_hash: TestHash = ::Hashing::hash(b"epoch-for-local"); let signature: MultiSignature = sr25519::Signature::from_raw([0u8; 64]).into(); let call = MevShieldCall::::execute_revealed { id, signer, - nonce, + key_hash, call: Box::new(inner_call), signature, }; @@ -305,19 +566,20 @@ fn validate_unsigned_accepts_inblock_source_for_execute_revealed() { new_test_ext().execute_with(|| { let pair = test_sr25519_pair(); let signer: AccountId32 = pair.public().into(); - let nonce: TestNonce = Zero::zero(); let inner_call = RuntimeCall::System(frame_system::Call::::remark { remark: b"noop-inblock".to_vec(), }); - let id = ::Hashing::hash(b"mevshield-id-inblock"); + let id: TestHash = ::Hashing::hash(b"mevshield-id-inblock"); + let key_hash: TestHash = + ::Hashing::hash(b"epoch-for-inblock"); let signature: MultiSignature = sr25519::Signature::from_raw([1u8; 64]).into(); let call = MevShieldCall::::execute_revealed { id, signer, - nonce, + key_hash, call: Box::new(inner_call), signature, }; diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index edf561810e..ef36b17921 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -1084,7 +1084,7 @@ mod dispatches { #[pallet::call_index(71)] #[pallet::weight((Weight::from_parts(161_700_000, 0) .saturating_add(T::DbWeight::get().reads(16_u64)) - .saturating_add(T::DbWeight::get().writes(9)), DispatchClass::Operational, Pays::Yes))] + .saturating_add(T::DbWeight::get().writes(11_u64)), DispatchClass::Operational, Pays::Yes))] pub fn swap_coldkey( origin: OriginFor, old_coldkey: T::AccountId,