diff --git a/Cargo.lock b/Cargo.lock index d242b5922f..a70fd85b18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6020,6 +6020,15 @@ dependencies = [ "serde", ] +[[package]] +name = "hybrid-array" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2d35805454dc9f8662a98d6d61886ffe26bd465f5960e0e55345c70d5c0d2a9" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "0.14.32" @@ -6766,6 +6775,16 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "kem" +version = "0.3.0-pre.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b8645470337db67b01a7f966decf7d0bafedbae74147d33e641c67a91df239f" +dependencies = [ + "rand_core 0.6.4", + "zeroize", +] + [[package]] name = "keystream" version = "1.0.0" @@ -7787,6 +7806,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "ml-kem" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97befee0c869cb56f3118f49d0f9bb68c9e3f380dec23c1100aedc4ec3ba239a" +dependencies = [ + "hybrid-array", + "kem", + "rand_core 0.6.4", + "sha3", +] + [[package]] name = "mmr-gadget" version = "46.0.0" @@ -8159,7 +8190,10 @@ checksum = "43794a0ace135be66a25d3ae77d41b91615fb68ae937f904090203e81f755b65" name = "node-subtensor" version = "4.0.0-dev" dependencies = [ + "anyhow", "async-trait", + "blake2 0.10.6", + "chacha20poly1305", "clap", "fc-api", "fc-aura", @@ -8179,19 +8213,26 @@ dependencies = [ "frame-system-rpc-runtime-api", "futures", "hex", + "hkdf", "jsonrpsee", "log", "memmap2 0.9.8", + "ml-kem", "node-subtensor-runtime", "num-traits", "pallet-commitments", "pallet-drand", + "pallet-shield", + "pallet-subtensor", "pallet-subtensor-swap-rpc", "pallet-subtensor-swap-runtime-api", "pallet-transaction-payment", "pallet-transaction-payment-rpc", "pallet-transaction-payment-rpc-runtime-api", + "parity-scale-codec", "polkadot-sdk", + "rand 0.8.5", + "rand_core 0.9.3", "sc-basic-authorship", "sc-chain-spec", "sc-chain-spec-derive", @@ -8219,6 +8260,7 @@ dependencies = [ "sc-transaction-pool-api", "serde", "serde_json", + "sha2 0.10.9", "sp-api", "sp-block-builder", "sp-blockchain", @@ -8243,7 +8285,10 @@ dependencies = [ "substrate-prometheus-endpoint", "subtensor-custom-rpc", "subtensor-custom-rpc-runtime-api", + "subtensor-macros", "subtensor-runtime-common", + "tokio", + "x25519-dalek", ] [[package]] @@ -8301,6 +8346,7 @@ dependencies = [ "pallet-safe-mode", "pallet-scheduler", "pallet-session", + "pallet-shield", "pallet-staking", "pallet-staking-reward-curve", "pallet-staking-reward-fn", @@ -10521,6 +10567,26 @@ dependencies = [ "sp-session", ] +[[package]] +name = "pallet-shield" +version = "0.0.1" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "pallet-aura", + "pallet-timestamp", + "parity-scale-codec", + "scale-info", + "sp-consensus-aura", + "sp-core", + "sp-io", + "sp-runtime", + "sp-std", + "sp-weights", + "subtensor-macros", +] + [[package]] name = "pallet-skip-feeless-payment" version = "16.0.0" diff --git a/Cargo.toml b/Cargo.toml index 6139004914..b70d09681b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -285,6 +285,9 @@ sha2 = { version = "0.10.8", default-features = false } rand_chacha = { version = "0.3.1", default-features = false } tle = { git = "https://github.com/ideal-lab5/timelock", rev = "5416406cfd32799e31e1795393d4916894de4468", default-features = false } +pallet-shield = { path = "pallets/shield", default-features = false } +ml-kem = { version = "0.2.0", default-features = true } + # Primitives [profile.release] diff --git a/node/Cargo.toml b/node/Cargo.toml index 12a4b71189..1d2351c265 100644 --- a/node/Cargo.toml +++ b/node/Cargo.toml @@ -114,6 +114,21 @@ fc-aura.workspace = true fp-consensus.workspace = true num-traits = { workspace = true, features = ["std"] } +# Mev Shield +pallet-shield.workspace = true +tokio = { version = "1.38", features = ["time"] } +x25519-dalek = "2" +hkdf = "0.12" +chacha20poly1305 = { version = "0.10", features = ["std"] } +codec.workspace = true +rand.workspace = true +sha2.workspace = true +anyhow.workspace = true +pallet-subtensor.workspace = true +ml-kem.workspace = true +rand_core = "0.9.3" +blake2 = "0.10.6" + # Local Dependencies node-subtensor-runtime = { workspace = true, features = ["std"] } subtensor-runtime-common = { workspace = true, features = ["std"] } @@ -121,6 +136,7 @@ subtensor-custom-rpc = { workspace = true, features = ["std"] } subtensor-custom-rpc-runtime-api = { workspace = true, features = ["std"] } pallet-subtensor-swap-rpc = { workspace = true, features = ["std"] } pallet-subtensor-swap-runtime-api = { workspace = true, features = ["std"] } +subtensor-macros.workspace = true [build-dependencies] substrate-build-script-utils.workspace = true @@ -138,6 +154,7 @@ default = ["rocksdb", "sql", "txpool"] fast-runtime = [ "node-subtensor-runtime/fast-runtime", "subtensor-runtime-common/fast-runtime", + "pallet-subtensor/fast-runtime", ] sql = ["fc-db/sql", "fc-mapping-sync/sql"] txpool = ["fc-rpc/txpool", "fc-rpc-core/txpool"] @@ -154,7 +171,10 @@ runtime-benchmarks = [ "pallet-drand/runtime-benchmarks", "pallet-transaction-payment/runtime-benchmarks", "polkadot-sdk/runtime-benchmarks", + "pallet-subtensor/runtime-benchmarks", + "pallet-shield/runtime-benchmarks", ] + pow-faucet = [] # Enable features that allow the runtime to be tried and debugged. Name might be subject to change @@ -167,6 +187,8 @@ try-runtime = [ "pallet-commitments/try-runtime", "pallet-drand/try-runtime", "polkadot-sdk/try-runtime", + "pallet-shield/try-runtime", + "pallet-subtensor/try-runtime", ] metadata-hash = ["node-subtensor-runtime/metadata-hash"] diff --git a/node/src/lib.rs b/node/src/lib.rs index c447a07309..ab4a409e1b 100644 --- a/node/src/lib.rs +++ b/node/src/lib.rs @@ -4,5 +4,6 @@ pub mod client; pub mod conditional_evm_block_import; pub mod consensus; pub mod ethereum; +pub mod mev_shield; pub mod rpc; pub mod service; diff --git a/node/src/main.rs b/node/src/main.rs index 64f25acc67..7adffa0ae9 100644 --- a/node/src/main.rs +++ b/node/src/main.rs @@ -10,6 +10,7 @@ mod command; mod conditional_evm_block_import; mod consensus; mod ethereum; +mod mev_shield; mod rpc; mod service; diff --git a/node/src/mev_shield/author.rs b/node/src/mev_shield/author.rs new file mode 100644 index 0000000000..20ac53702a --- /dev/null +++ b/node/src/mev_shield/author.rs @@ -0,0 +1,447 @@ +use chacha20poly1305::{ + KeyInit, XChaCha20Poly1305, XNonce, + aead::{Aead, Payload}, +}; +use ml_kem::{EncodedSizeUser, KemCore, MlKem768}; +use node_subtensor_runtime as runtime; +use rand::rngs::OsRng; +use sp_core::blake2_256; +use sp_runtime::KeyTypeId; +use std::sync::{Arc, Mutex}; +use subtensor_macros::freeze_struct; +use tokio::time::sleep; + +/// Parameters controlling time windows inside the slot. +#[freeze_struct("5c7ce101b36950de")] +#[derive(Clone)] +pub struct TimeParams { + pub slot_ms: u64, + pub announce_at_ms: u64, + pub decrypt_window_ms: u64, +} + +/// Holds the current/next ML‑KEM keypairs and their 32‑byte fingerprints. +#[freeze_struct("5e3c8209248282c3")] +#[derive(Clone)] +pub struct ShieldKeys { + pub current_sk: Vec, // ML‑KEM secret key bytes (encoded form) + pub current_pk: Vec, // ML‑KEM public key bytes (encoded form) + pub current_fp: [u8; 32], // blake2_256(pk) + pub next_sk: Vec, + pub next_pk: Vec, + pub next_fp: [u8; 32], +} + +impl ShieldKeys { + pub fn new() -> Self { + let (sk, pk) = MlKem768::generate(&mut OsRng); + + let sk_bytes = sk.as_bytes(); + let pk_bytes = pk.as_bytes(); + let sk_slice: &[u8] = sk_bytes.as_ref(); + let pk_slice: &[u8] = pk_bytes.as_ref(); + + let current_sk = sk_slice.to_vec(); + let current_pk = pk_slice.to_vec(); + let current_fp = blake2_256(pk_slice); + + let (nsk, npk) = MlKem768::generate(&mut OsRng); + let nsk_bytes = nsk.as_bytes(); + let npk_bytes = npk.as_bytes(); + let nsk_slice: &[u8] = nsk_bytes.as_ref(); + let npk_slice: &[u8] = npk_bytes.as_ref(); + let next_sk = nsk_slice.to_vec(); + let next_pk = npk_slice.to_vec(); + let next_fp = blake2_256(npk_slice); + + Self { + current_sk, + current_pk, + current_fp, + next_sk, + next_pk, + next_fp, + } + } + + pub fn roll_for_next_slot(&mut self) { + // Move next -> current + self.current_sk = core::mem::take(&mut self.next_sk); + self.current_pk = core::mem::take(&mut self.next_pk); + self.current_fp = self.next_fp; + + // Generate fresh next + let (nsk, npk) = MlKem768::generate(&mut OsRng); + let nsk_bytes = nsk.as_bytes(); + let npk_bytes = npk.as_bytes(); + let nsk_slice: &[u8] = nsk_bytes.as_ref(); + let npk_slice: &[u8] = npk_bytes.as_ref(); + self.next_sk = nsk_slice.to_vec(); + self.next_pk = npk_slice.to_vec(); + self.next_fp = blake2_256(npk_slice); + } +} + +impl Default for ShieldKeys { + fn default() -> Self { + Self::new() + } +} + +/// Shared context state. +#[freeze_struct("62af7d26cf7c1271")] +#[derive(Clone)] +pub struct ShieldContext { + pub keys: Arc>, + pub timing: TimeParams, +} + +/// Derive AEAD key directly from the 32‑byte ML‑KEM shared secret. +pub fn derive_aead_key(ss: &[u8]) -> [u8; 32] { + let mut key = [0u8; 32]; + let n = ss.len().min(32); + + if let (Some(dst), Some(src)) = (key.get_mut(..n), ss.get(..n)) { + dst.copy_from_slice(src); + } + key +} + +/// Plain XChaCha20-Poly1305 decrypt helper +pub fn aead_decrypt( + key: [u8; 32], + nonce24: [u8; 24], + ciphertext: &[u8], + aad: &[u8], +) -> Option> { + let aead = XChaCha20Poly1305::new((&key).into()); + aead.decrypt( + XNonce::from_slice(&nonce24), + Payload { + msg: ciphertext, + aad, + }, + ) + .ok() +} + +const AURA_KEY_TYPE: KeyTypeId = KeyTypeId(*b"aura"); + +/// Start background tasks: +/// - per-slot ML‑KEM key rotation +/// - at ~announce_at_ms announce the next key bytes on chain, +pub fn spawn_author_tasks( + task_spawner: &sc_service::SpawnTaskHandle, + client: Arc, + pool: Arc, + keystore: sp_keystore::KeystorePtr, + timing: TimeParams, +) -> ShieldContext +where + B: sp_runtime::traits::Block, + C: sc_client_api::HeaderBackend + sc_client_api::BlockchainEvents + Send + Sync + 'static, + Pool: sc_transaction_pool_api::TransactionPool + Send + Sync + 'static, + B::Extrinsic: From, +{ + let ctx = ShieldContext { + keys: Arc::new(Mutex::new(ShieldKeys::new())), + timing: timing.clone(), + }; + + let aura_keys: Vec = keystore.sr25519_public_keys(AURA_KEY_TYPE); + + let local_aura_pub = match aura_keys.first().copied() { + Some(k) => k, + None => { + log::warn!( + target: "mev-shield", + "spawn_author_tasks: no local Aura sr25519 key in keystore; \ + this node will NOT announce MEV-Shield keys" + ); + return ctx; + } + }; + + let ctx_clone = ctx.clone(); + let client_clone = client.clone(); + let pool_clone = pool.clone(); + let keystore_clone = keystore.clone(); + + // Slot tick / key-announce loop. + task_spawner.spawn( + "mev-shield-keys-and-announce", + None, + async move { + use futures::StreamExt; + use sp_consensus::BlockOrigin; + + let slot_ms = timing.slot_ms; + + // Clamp announce_at_ms so it never exceeds slot_ms. + let mut announce_at_ms = timing.announce_at_ms; + if announce_at_ms > slot_ms { + log::warn!( + target: "mev-shield", + "spawn_author_tasks: announce_at_ms ({announce_at_ms}) > slot_ms ({slot_ms}); clamping to slot_ms", + ); + announce_at_ms = slot_ms; + } + let tail_ms = slot_ms.saturating_sub(announce_at_ms); + + log::debug!( + target: "mev-shield", + "author timing: slot_ms={slot_ms} announce_at_ms={announce_at_ms} (effective) tail_ms={tail_ms}", + ); + + let mut import_stream = client_clone.import_notification_stream(); + let mut local_nonce: u32 = 0; + + while let Some(notif) = import_stream.next().await { + // ✅ Only act on blocks that this node authored. + if notif.origin != BlockOrigin::Own { + continue; + } + + // This block is the start of a slot for which we are the author. + let (curr_pk_len, next_pk_len) = match ctx_clone.keys.lock() { + Ok(k) => (k.current_pk.len(), k.next_pk.len()), + Err(e) => { + log::debug!( + target: "mev-shield", + "spawn_author_tasks: failed to lock ShieldKeys (poisoned?): {e:?}", + ); + continue; + } + }; + + log::debug!( + target: "mev-shield", + "Slot start (local author): (pk sizes: curr={curr_pk_len}B, next={next_pk_len}B)", + ); + + // Wait until the announce window in this slot. + if announce_at_ms > 0 { + sleep(std::time::Duration::from_millis(announce_at_ms)).await; + } + + // Read the next key we intend to use for the following block. + let next_pk = match ctx_clone.keys.lock() { + Ok(k) => k.next_pk.clone(), + Err(e) => { + log::debug!( + target: "mev-shield", + "spawn_author_tasks: failed to lock ShieldKeys for next_pk: {e:?}", + ); + continue; + } + }; + + // Submit announce_next_key once, signed with the local Aura authority that authors this block + match submit_announce_extrinsic::( + client_clone.clone(), + pool_clone.clone(), + keystore_clone.clone(), + local_aura_pub, + next_pk.clone(), + local_nonce, + ) + .await + { + Ok(()) => { + local_nonce = local_nonce.saturating_add(1); + } + Err(e) => { + let msg = format!("{e:?}"); + // If the nonce is stale, bump once and retry. + if msg.contains("InvalidTransaction::Stale") || msg.contains("Stale") { + if submit_announce_extrinsic::( + client_clone.clone(), + pool_clone.clone(), + keystore_clone.clone(), + local_aura_pub, + next_pk, + local_nonce.saturating_add(1), + ) + .await + .is_ok() + { + local_nonce = local_nonce.saturating_add(2); + } else { + log::debug!( + target: "mev-shield", + "announce_next_key retry failed after stale nonce: {e:?}" + ); + } + } else { + log::debug!( + target: "mev-shield", + "announce_next_key submit error: {e:?}" + ); + } + } + } + + // Sleep the remainder of the slot (if any). + if tail_ms > 0 { + sleep(std::time::Duration::from_millis(tail_ms)).await; + } + + // Roll keys for the next block. + match ctx_clone.keys.lock() { + Ok(mut k) => { + k.roll_for_next_slot(); + log::debug!( + target: "mev-shield", + "Rolled ML-KEM key at slot boundary", + ); + } + Err(e) => { + log::debug!( + target: "mev-shield", + "spawn_author_tasks: failed to lock ShieldKeys for roll_for_next_slot: {e:?}", + ); + } + } + } + }, + ); + + ctx +} + +/// Build & submit the signed `announce_next_key` extrinsic OFF-CHAIN +pub async fn submit_announce_extrinsic( + client: Arc, + pool: Arc, + keystore: sp_keystore::KeystorePtr, + aura_pub: sp_core::sr25519::Public, + next_public_key: Vec, + nonce: u32, +) -> anyhow::Result<()> +where + B: sp_runtime::traits::Block, + C: sc_client_api::HeaderBackend + Send + Sync + 'static, + Pool: sc_transaction_pool_api::TransactionPool + Send + Sync + 'static, + B::Extrinsic: From, + B::Hash: AsRef<[u8]>, +{ + use node_subtensor_runtime as runtime; + use runtime::{RuntimeCall, SignedPayload, UncheckedExtrinsic}; + + use sc_transaction_pool_api::TransactionSource; + use sp_core::H256; + use sp_runtime::codec::Encode; + use sp_runtime::{ + AccountId32, BoundedVec, MultiSignature, + generic::Era, + traits::{ConstU32, TransactionExtension}, + }; + + // Helper: map a Block hash to H256 + fn to_h256>(h: H) -> H256 { + let bytes = h.as_ref(); + let mut out = [0u8; 32]; + + if bytes.is_empty() { + return H256(out); + } + + let n = bytes.len().min(32); + let src_start = bytes.len().saturating_sub(n); + let dst_start = 32usize.saturating_sub(n); + + let src_slice = bytes.get(src_start..).and_then(|s| s.get(..n)); + + if let (Some(dst), Some(src)) = (out.get_mut(dst_start..32), src_slice) { + dst.copy_from_slice(src); + H256(out) + } else { + // Extremely unlikely; fall back to zeroed H256 if indices are somehow invalid. + H256([0u8; 32]) + } + } + + type MaxPk = ConstU32<2048>; + let public_key: BoundedVec = BoundedVec::try_from(next_public_key) + .map_err(|_| anyhow::anyhow!("public key too long (>2048 bytes)"))?; + + // 1) The runtime call carrying public key bytes. + let call = RuntimeCall::MevShield(pallet_shield::Call::announce_next_key { public_key }); + + type Extra = runtime::TransactionExtensions; + let extra: Extra = + ( + frame_system::CheckNonZeroSender::::new(), + frame_system::CheckSpecVersion::::new(), + frame_system::CheckTxVersion::::new(), + frame_system::CheckGenesis::::new(), + frame_system::CheckEra::::from(Era::Immortal), + node_subtensor_runtime::check_nonce::CheckNonce::::from(nonce).into(), + frame_system::CheckWeight::::new(), + node_subtensor_runtime::transaction_payment_wrapper::ChargeTransactionPaymentWrapper::< + runtime::Runtime, + >::new(pallet_transaction_payment::ChargeTransactionPayment::< + runtime::Runtime, + >::from(0u64)), + pallet_subtensor::transaction_extension::SubtensorTransactionExtension::< + runtime::Runtime, + >::new(), + pallet_drand::drand_priority::DrandPriority::::new(), + frame_metadata_hash_extension::CheckMetadataHash::::new(false), + ); + + type Implicit = >::Implicit; + + let info = client.info(); + let genesis_h256: H256 = to_h256(info.genesis_hash); + + let implicit: Implicit = ( + (), // CheckNonZeroSender + runtime::VERSION.spec_version, // CheckSpecVersion + runtime::VERSION.transaction_version, // CheckTxVersion + genesis_h256, // CheckGenesis + genesis_h256, // CheckEra (Immortal) + (), // CheckNonce (additional part) + (), // CheckWeight + (), // ChargeTransactionPaymentWrapper (additional part) + (), // SubtensorTransactionExtension (additional part) + (), // DrandPriority + None, // CheckMetadataHash (disabled) + ); + + // Build the exact signable payload. + let payload: SignedPayload = SignedPayload::from_raw(call.clone(), extra.clone(), implicit); + + let raw_payload = payload.encode(); + + // Sign with the local Aura key. + let sig_opt = keystore + .sr25519_sign(AURA_KEY_TYPE, &aura_pub, &raw_payload) + .map_err(|e| anyhow::anyhow!("keystore sr25519_sign error: {e:?}"))?; + let sig = sig_opt + .ok_or_else(|| anyhow::anyhow!("keystore sr25519_sign returned None for Aura key"))?; + + let signature: MultiSignature = sig.into(); + + let who: AccountId32 = aura_pub.into(); + let address = sp_runtime::MultiAddress::Id(who); + + let uxt: UncheckedExtrinsic = UncheckedExtrinsic::new_signed(call, address, signature, extra); + + let xt_bytes = uxt.encode(); + let xt_hash = sp_core::hashing::blake2_256(&xt_bytes); + let xt_hash_hex = hex::encode(xt_hash); + + let opaque: sp_runtime::OpaqueExtrinsic = uxt.into(); + let xt: ::Extrinsic = opaque.into(); + + pool.submit_one(info.best_hash, TransactionSource::Local, xt) + .await?; + + log::debug!( + target: "mev-shield", + "announce_next_key submitted: xt=0x{xt_hash_hex}, nonce={nonce}", + ); + + Ok(()) +} diff --git a/node/src/mev_shield/mod.rs b/node/src/mev_shield/mod.rs new file mode 100644 index 0000000000..91817097bf --- /dev/null +++ b/node/src/mev_shield/mod.rs @@ -0,0 +1,2 @@ +pub mod author; +pub mod proposer; diff --git a/node/src/mev_shield/proposer.rs b/node/src/mev_shield/proposer.rs new file mode 100644 index 0000000000..7bacbf90e3 --- /dev/null +++ b/node/src/mev_shield/proposer.rs @@ -0,0 +1,816 @@ +use super::author::ShieldContext; +use futures::StreamExt; +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_runtime::{AccountId32, MultiSignature, OpaqueExtrinsic}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::time::sleep; + +/// Buffer of wrappers keyed by the block number in which they were included. +#[derive(Default, Clone)] +struct WrapperBuffer { + by_id: HashMap< + H256, + ( + Vec, // ciphertext blob + u64, // originating block number + AccountId32, // wrapper author + ), + >, +} + +impl WrapperBuffer { + fn upsert(&mut self, id: H256, block_number: u64, author: AccountId32, ciphertext: Vec) { + self.by_id.insert(id, (ciphertext, block_number, author)); + } + + /// Drain only wrappers whose `block_number` matches the given `block`. + /// - Wrappers with `block_number > block` are kept for future decrypt windows. + /// - Wrappers with `block_number < block` are considered stale and dropped. + fn drain_for_block( + &mut self, + block: u64, + ) -> Vec<(H256, u64, sp_runtime::AccountId32, Vec)> { + let mut ready = Vec::new(); + let mut kept_future: usize = 0; + let mut dropped_past: usize = 0; + + self.by_id.retain(|id, (ct, block_number, who)| { + if *block_number == block { + // Ready to process now; remove from buffer. + ready.push((*id, *block_number, who.clone(), ct.clone())); + false + } else if *block_number > block { + // Not yet reveal time; keep for future blocks. + kept_future = kept_future.saturating_add(1); + true + } else { + // block_number < block => stale / missed reveal window; drop. + dropped_past = dropped_past.saturating_add(1); + log::debug!( + target: "mev-shield", + "revealer: dropping stale wrapper id=0x{} block_number={} < curr_block={}", + hex::encode(id.as_bytes()), + *block_number, + block + ); + false + } + }); + + log::debug!( + target: "mev-shield", + "revealer: drain_for_block(block={}): ready={}, kept_future={}, dropped_past={}", + block, + ready.len(), + kept_future, + dropped_past + ); + + ready + } +} + +/// Start a background worker that: +/// • watches imported blocks and captures `MevShield::submit_encrypted` +/// • buffers those wrappers per originating block, +/// • during the last `decrypt_window_ms` of the slot: decrypt & submit unsigned `execute_revealed` +pub fn spawn_revealer( + task_spawner: &SpawnTaskHandle, + client: Arc, + pool: Arc, + ctx: ShieldContext, +) where + B: sp_runtime::traits::Block, + C: sc_client_api::HeaderBackend + + sc_client_api::BlockchainEvents + + sc_client_api::BlockBackend + + Send + + Sync + + 'static, + 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; + + let buffer: Arc> = Arc::new(Mutex::new(WrapperBuffer::default())); + + // ── 1) buffer wrappers ─────────────────────────────────────── + { + let client = Arc::clone(&client); + let buffer = Arc::clone(&buffer); + + task_spawner.spawn( + "mev-shield-buffer-wrappers", + None, + async move { + log::debug!(target: "mev-shield", "buffer-wrappers task started"); + let mut import_stream = client.import_notification_stream(); + + while let Some(notif) = import_stream.next().await { + let at_hash = notif.hash; + let block_number_u64: u64 = (*notif.header.number()).saturated_into(); + + log::debug!( + target: "mev-shield", + "imported block hash={:?} number={} origin={:?}", + at_hash, + block_number_u64, + notif.origin + ); + + match client.block_body(at_hash) { + Ok(Some(body)) => { + log::debug!( + target: "mev-shield", + " block has {} extrinsics", + body.len() + ); + + for (idx, opaque_xt) in body.into_iter().enumerate() { + let encoded = opaque_xt.encode(); + log::debug!( + target: "mev-shield", + " [xt #{idx}] opaque len={} bytes", + encoded.len() + ); + + let uxt: RUnchecked = match RUnchecked::decode(&mut &encoded[..]) { + Ok(u) => u, + Err(e) => { + log::debug!( + target: "mev-shield", + " [xt #{idx}] failed to decode UncheckedExtrinsic: {e:?}", + ); + continue; + } + }; + + log::debug!( + target: "mev-shield", + " [xt #{idx}] decoded call: {:?}", + &uxt.0.function + ); + + let author_opt: Option = + match &uxt.0.preamble { + sp_runtime::generic::Preamble::Signed( + addr, + _sig, + _ext, + ) => match addr.clone() { + Address::Id(acc) => Some(acc), + Address::Address32(bytes) => { + Some(sp_runtime::AccountId32::new(bytes)) + } + _ => None, + }, + _ => None, + }; + + let Some(author) = author_opt else { + log::debug!( + target: "mev-shield", + " [xt #{idx}] not a Signed(AccountId32) extrinsic; skipping" + ); + continue; + }; + + if let node_subtensor_runtime::RuntimeCall::MevShield( + pallet_shield::Call::submit_encrypted { + commitment, + ciphertext, + }, + ) = &uxt.0.function + { + let payload = + (author.clone(), *commitment, ciphertext).encode(); + let id = H256(sp_core::hashing::blake2_256(&payload)); + + log::debug!( + target: "mev-shield", + " [xt #{idx}] buffered submit_encrypted: id=0x{}, block_number={}, author={}, ct_len={}, commitment={:?}", + hex::encode(id.as_bytes()), + block_number_u64, + author, + ciphertext.len(), + commitment + ); + + if let Ok(mut buf) = buffer.lock() { + buf.upsert( + id, + block_number_u64, + author, + ciphertext.to_vec(), + ); + } else { + log::debug!( + target: "mev-shield", + " [xt #{idx}] failed to lock WrapperBuffer; dropping wrapper" + ); + } + } + } + } + Ok(None) => log::debug!( + target: "mev-shield", + " block_body returned None for hash={at_hash:?}", + ), + Err(e) => log::debug!( + target: "mev-shield", + " block_body error for hash={at_hash:?}: {e:?}", + ), + } + } + }, + ); + } + + // ── 2) decrypt window revealer ────────────────────────────── + { + let client = Arc::clone(&client); + let pool = Arc::clone(&pool); + let buffer = Arc::clone(&buffer); + let ctx = ctx.clone(); + + task_spawner.spawn( + "mev-shield-last-window-revealer", + None, + async move { + log::debug!(target: "mev-shield", "last-window-revealer task started"); + + // Respect the configured slot_ms, but clamp the decrypt window so it never + // exceeds the slot length (important for fast runtimes). + let slot_ms = ctx.timing.slot_ms; + let mut decrypt_window_ms = ctx.timing.decrypt_window_ms; + + if decrypt_window_ms > slot_ms { + log::warn!( + target: "mev-shield", + "spawn_revealer: decrypt_window_ms ({decrypt_window_ms}) > slot_ms ({slot_ms}); clamping to slot_ms", + ); + decrypt_window_ms = slot_ms; + } + + let tail_ms = slot_ms.saturating_sub(decrypt_window_ms); + + log::debug!( + target: "mev-shield", + "revealer timing: slot_ms={slot_ms} decrypt_window_ms={decrypt_window_ms} (effective) tail_ms={tail_ms}", + ); + + loop { + log::debug!( + target: "mev-shield", + "revealer: sleeping {tail_ms} ms before decrypt window (slot_ms={slot_ms}, decrypt_window_ms={decrypt_window_ms})", + ); + + if tail_ms > 0 { + sleep(Duration::from_millis(tail_ms)).await; + } + + // Snapshot the current ML‑KEM secret (but *not* any epoch). + let snapshot_opt = match ctx.keys.lock() { + Ok(k) => { + let sk_hash = sp_core::hashing::blake2_256(&k.current_sk); + Some(( + k.current_sk.clone(), + k.current_pk.len(), + k.next_pk.len(), + sk_hash, + )) + } + Err(e) => { + log::debug!( + target: "mev-shield", + "revealer: failed to lock ShieldKeys (poisoned?): {e:?}", + ); + None + } + }; + + let (curr_sk_bytes, curr_pk_len, next_pk_len, sk_hash) = + match snapshot_opt { + Some(v) => v, + None => { + // Skip this decrypt window entirely, without holding any guard. + if decrypt_window_ms > 0 { + sleep(Duration::from_millis(decrypt_window_ms)).await; + } + continue; + } + }; + + // Use best block number as the “epoch” for which we reveal. + 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={}", + curr_block, + curr_sk_bytes.len(), + hex::encode(sk_hash), + curr_pk_len, + next_pk_len + ); + + // Only process wrappers whose originating block matches the current block. + let drained: Vec<(H256, u64, sp_runtime::AccountId32, Vec)> = + match buffer.lock() { + Ok(mut buf) => buf.drain_for_block(curr_block), + Err(e) => { + log::debug!( + target: "mev-shield", + "revealer: failed to lock WrapperBuffer for drain_for_block: {e:?}", + ); + Vec::new() + } + }; + + log::debug!( + target: "mev-shield", + "revealer: drained {} buffered wrappers for current block={}", + drained.len(), + curr_block + ); + + 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={}", + hex::encode(id.as_bytes()), + block_number, + curr_block, + author, + blob.len() + ); + + // 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, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: blob too short or invalid length prefix", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let kem_end = match 2usize.checked_add(kem_len) { + Some(v) => v, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: kem_len overflow", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let nonce_end = match kem_end.checked_add(24usize) { + Some(v) => v, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: nonce range overflow", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let kem_ct_bytes = match blob.get(2..kem_end) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: blob too short for kem_ct (kem_len={}, total={})", + hex::encode(id.as_bytes()), + kem_len, + blob.len() + ); + continue; + } + }; + + let nonce_bytes = match blob.get(kem_end..nonce_end) { + Some(s) if s.len() == 24 => s, + _ => { + log::debug!( + target: "mev-shield", + " id=0x{}: blob too short for 24-byte nonce (kem_len={}, total={})", + hex::encode(id.as_bytes()), + kem_len, + blob.len() + ); + continue; + } + }; + + let aead_body = match blob.get(nonce_end..) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: blob has no AEAD body", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + 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{}", + hex::encode(id.as_bytes()), + kem_len, + hex::encode(kem_ct_hash), + hex::encode(nonce_bytes), + aead_body.len(), + hex::encode(aead_body_hash), + ); + + // Rebuild DecapsulationKey and decapsulate. + let enc_sk = + match Encoded::>::try_from( + &curr_sk_bytes[..], + ) { + Ok(e) => e, + Err(e) => { + log::debug!( + target: "mev-shield", + " id=0x{}: DecapsulationKey::try_from(sk_bytes) failed (len={}, err={:?})", + hex::encode(id.as_bytes()), + curr_sk_bytes.len(), + e + ); + continue; + } + }; + let sk = DecapsulationKey::::from_bytes(&enc_sk); + + let ct = match Ciphertext::::try_from(kem_ct_bytes) { + Ok(c) => c, + Err(e) => { + log::debug!( + target: "mev-shield", + " id=0x{}: Ciphertext::try_from failed: {:?}", + hex::encode(id.as_bytes()), + e + ); + continue; + } + }; + + let ss = match sk.decapsulate(&ct) { + Ok(s) => s, + Err(_) => { + log::debug!( + target: "mev-shield", + " id=0x{}: ML-KEM decapsulate() failed", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let ss_bytes: &[u8] = ss.as_ref(); + if ss_bytes.len() != 32 { + log::debug!( + target: "mev-shield", + " id=0x{}: shared secret len={} != 32; skipping", + hex::encode(id.as_bytes()), + ss_bytes.len() + ); + continue; + } + let mut ss32 = [0u8; 32]; + 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); + + log::debug!( + target: "mev-shield", + " id=0x{}: decapsulated shared_secret_len=32 shared_secret_hash=0x{}", + hex::encode(id.as_bytes()), + hex::encode(ss_hash) + ); + log::debug!( + target: "mev-shield", + " id=0x{}: derived AEAD key hash=0x{} (direct-from-ss)", + hex::encode(id.as_bytes()), + hex::encode(key_hash) + ); + + let mut nonce24 = [0u8; 24]; + nonce24.copy_from_slice(nonce_bytes); + + log::debug!( + target: "mev-shield", + " id=0x{}: attempting AEAD decrypt nonce=0x{} ct_len={}", + hex::encode(id.as_bytes()), + hex::encode(nonce24), + aead_body.len() + ); + + let plaintext = match crate::mev_shield::author::aead_decrypt( + aead_key, + nonce24, + aead_body, + &[], + ) { + Some(pt) => pt, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: AEAD decrypt FAILED with direct-from-ss key; ct_hash=0x{}", + hex::encode(id.as_bytes()), + hex::encode(aead_body_hash), + ); + continue; + } + }; + + log::debug!( + target: "mev-shield", + " id=0x{}: AEAD decrypt OK, plaintext_len={}", + hex::encode(id.as_bytes()), + plaintext.len() + ); + + type RuntimeNonce = + ::Nonce; + + // Safely parse plaintext layout without panics. + // Layout: signer (32) || nonce (4) || call (..) + // || sig_kind (1) || sig (64) + let min_plain_len: usize = 32usize + .saturating_add(4) + .saturating_add(1) + .saturating_add(1) + .saturating_add(64); + if plaintext.len() < min_plain_len { + log::debug!( + target: "mev-shield", + " id=0x{}: plaintext too short ({}) for expected layout", + hex::encode(id.as_bytes()), + plaintext.len() + ); + continue; + } + + let signer_raw = match plaintext.get(0..32) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: missing signer bytes", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let nonce_le = match plaintext.get(32..36) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: missing nonce bytes", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let sig_off = match plaintext.len().checked_sub(65) { + Some(off) if off >= 36 => off, + _ => { + log::debug!( + target: "mev-shield", + " id=0x{}: invalid plaintext length for signature split", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let call_bytes = match plaintext.get(36..sig_off) { + Some(s) => s, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: missing call bytes", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let sig_kind = match plaintext.get(sig_off) { + Some(b) => *b, + None => { + log::debug!( + target: "mev-shield", + " id=0x{}: missing signature kind byte", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + 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 => { + log::debug!( + target: "mev-shield", + " id=0x{}: missing signature bytes", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + + let signer_array: [u8; 32] = match signer_raw.try_into() { + Ok(a) => a, + Err(_) => { + log::debug!( + target: "mev-shield", + " id=0x{}: signer_raw not 32 bytes", + hex::encode(id.as_bytes()) + ); + continue; + } + }; + 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 inner_call: node_subtensor_runtime::RuntimeCall = + match Decode::decode(&mut &call_bytes[..]) { + Ok(c) => c, + Err(e) => { + log::debug!( + target: "mev-shield", + " id=0x{}: failed to decode RuntimeCall (len={}): {:?}", + hex::encode(id.as_bytes()), + call_bytes.len(), + e + ); + continue; + } + }; + + 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)) + } 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() + ); + continue; + }; + + log::debug!( + target: "mev-shield", + " id=0x{}: decrypted wrapper: signer={}, nonce={}, call={:?}", + hex::encode(id.as_bytes()), + signer, + raw_nonce_u32, + inner_call + ); + + let reveal = node_subtensor_runtime::RuntimeCall::MevShield( + pallet_shield::Call::execute_revealed { + id, + signer: signer.clone(), + nonce: account_nonce, + call: Box::new(inner_call), + signature, + }, + ); + + to_submit.push((id, reveal)); + } + + // Submit locally. + let at = client.info().best_hash; + log::debug!( + target: "mev-shield", + "revealer: submitting {} execute_revealed calls at best_hash={:?}", + to_submit.len(), + at + ); + + for (id, call) in to_submit.into_iter() { + let uxt: node_subtensor_runtime::UncheckedExtrinsic = + node_subtensor_runtime::UncheckedExtrinsic::new_bare(call); + let xt_bytes = uxt.encode(); + + log::debug!( + target: "mev-shield", + " id=0x{}: encoded UncheckedExtrinsic len={}", + hex::encode(id.as_bytes()), + xt_bytes.len() + ); + + match OpaqueExtrinsic::from_bytes(&xt_bytes) { + Ok(opaque) => { + match pool + .submit_one(at, TransactionSource::Local, opaque) + .await + { + Ok(_) => { + let xt_hash = + sp_core::hashing::blake2_256(&xt_bytes); + log::debug!( + target: "mev-shield", + " id=0x{}: submit_one(execute_revealed) OK, xt_hash=0x{}", + hex::encode(id.as_bytes()), + hex::encode(xt_hash) + ); + } + Err(e) => { + log::debug!( + target: "mev-shield", + " id=0x{}: submit_one(execute_revealed) FAILED: {:?}", + hex::encode(id.as_bytes()), + e + ); + } + } + } + Err(e) => { + log::debug!( + target: "mev-shield", + " id=0x{}: OpaqueExtrinsic::from_bytes failed: {:?}", + hex::encode(id.as_bytes()), + e + ); + } + } + } + + // Let the decrypt window elapse. + if decrypt_window_ms > 0 { + sleep(Duration::from_millis(decrypt_window_ms)).await; + } + } + }, + ); + } +} diff --git a/node/src/service.rs b/node/src/service.rs index 2ef1904f08..d32aceea9c 100644 --- a/node/src/service.rs +++ b/node/src/service.rs @@ -34,6 +34,7 @@ use crate::ethereum::{ BackendType, EthConfiguration, FrontierBackend, FrontierPartialComponents, StorageOverride, StorageOverrideHandler, db_config_dir, new_frontier_partial, spawn_frontier_tasks, }; +use crate::mev_shield::{author, proposer}; const LOG_TARGET: &str = "node-service"; @@ -534,6 +535,48 @@ where ) .await; + // ==== MEV-SHIELD HOOKS ==== + let mut mev_timing: Option = None; + + if role.is_authority() { + let slot_duration = consensus_mechanism.slot_duration(&client)?; + let slot_duration_ms: u64 = u64::try_from(slot_duration.as_millis()).unwrap_or(u64::MAX); + + // For 12s blocks: announce ≈ 7s, decrypt window ≈ 3s. + // For 250ms blocks: announce ≈ 145ms, decrypt window ≈ 62ms, etc. + let announce_at_ms_raw = slot_duration_ms.saturating_mul(7).saturating_div(12); + + let decrypt_window_ms = slot_duration_ms.saturating_mul(3).saturating_div(12); + + // Ensure announce_at_ms + decrypt_window_ms never exceeds slot_ms. + let max_announce = slot_duration_ms.saturating_sub(decrypt_window_ms); + let announce_at_ms = announce_at_ms_raw.min(max_announce); + + let timing = author::TimeParams { + slot_ms: slot_duration_ms, + announce_at_ms, + decrypt_window_ms, + }; + mev_timing = Some(timing.clone()); + + // Start author-side tasks with dynamic timing. + let mev_ctx = author::spawn_author_tasks::( + &task_manager.spawn_handle(), + client.clone(), + transaction_pool.clone(), + keystore_container.keystore(), + timing.clone(), + ); + + // Start last-portion-of-slot revealer (decrypt -> execute_revealed). + proposer::spawn_revealer::( + &task_manager.spawn_handle(), + client.clone(), + transaction_pool.clone(), + mev_ctx.clone(), + ); + } + if role.is_authority() { // manual-seal authorship if let Some(sealing) = sealing { @@ -548,7 +591,6 @@ where telemetry.as_ref(), commands_stream, )?; - log::info!("Manual Seal Ready"); return Ok(task_manager); } @@ -562,6 +604,26 @@ where ); let slot_duration = consensus_mechanism.slot_duration(&client)?; + + let start_fraction: f32 = { + let (slot_ms, decrypt_ms) = mev_timing + .as_ref() + .map(|t| (t.slot_ms, t.decrypt_window_ms)) + .unwrap_or((slot_duration.as_millis(), 3_000)); + + let guard_ms: u64 = 200; // small cushion so reveals hit the pool first + let after_decrypt_ms = slot_ms.saturating_sub(decrypt_ms).saturating_add(guard_ms); + + let f_raw = if slot_ms > 0 { + (after_decrypt_ms as f32) / (slot_ms as f32) + } else { + // Extremely defensive fallback; should never happen in practice. + 0.75 + }; + + f_raw.clamp(0.50, 0.98) + }; + let create_inherent_data_providers = move |_, ()| async move { CM::create_inherent_data_providers(slot_duration) }; @@ -579,7 +641,7 @@ where force_authoring, backoff_authoring_blocks, keystore: keystore_container.keystore(), - block_proposal_slot_portion: SlotProportion::new(2f32 / 3f32), + block_proposal_slot_portion: SlotProportion::new(start_fraction), max_block_proposal_slot_portion: None, telemetry: telemetry.as_ref().map(|x| x.handle()), }, diff --git a/pallets/shield/Cargo.toml b/pallets/shield/Cargo.toml new file mode 100644 index 0000000000..c0038f2b92 --- /dev/null +++ b/pallets/shield/Cargo.toml @@ -0,0 +1,73 @@ +[package] +name = "pallet-shield" +description = "FRAME pallet for opt-in, per-block ephemeral-key encrypted transactions, MEV-shielded." +authors = ["Subtensor Contributors "] +version = "0.0.1" +license = "Unlicense" +edition.workspace = true +homepage = "https://github.com/opentensor/subtensor" +repository = "https://github.com/opentensor/subtensor" +publish = false + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true, features = ["derive", "max-encoded-len"] } +scale-info = { workspace = true, features = ["derive"] } + +subtensor-macros.workspace = true + +# FRAME core +frame-support.workspace = true +frame-system.workspace = true +frame-benchmarking = { workspace = true, optional = true } + +# Substrate primitives +sp-core.workspace = true +sp-runtime.workspace = true +sp-io.workspace = true +sp-std.workspace = true +sp-weights.workspace = true + +# Pallets used in Config +pallet-timestamp.workspace = true +pallet-aura.workspace = true +sp-consensus-aura.workspace = true + +[dev-dependencies] + +[features] +default = [] + +std = [ + "codec/std", + "scale-info/std", + "frame-support/std", + "frame-system/std", + "frame-benchmarking?/std", + "sp-core/std", + "sp-runtime/std", + "sp-io/std", + "sp-std/std", + "sp-weights/std", + "pallet-timestamp/std", + "pallet-aura/std", + "sp-consensus-aura/std", +] + +runtime-benchmarks = [ + "frame-benchmarking/runtime-benchmarks", + "frame-support/runtime-benchmarks", + "frame-system/runtime-benchmarks", + "sp-runtime/runtime-benchmarks", + "pallet-timestamp/runtime-benchmarks", +] + +try-runtime = [ + "frame-support/try-runtime", + "frame-system/try-runtime", + "sp-runtime/try-runtime", + "pallet-aura/try-runtime", + "pallet-timestamp/try-runtime", +] diff --git a/pallets/shield/src/benchmarking.rs b/pallets/shield/src/benchmarking.rs new file mode 100644 index 0000000000..8e1d370e7f --- /dev/null +++ b/pallets/shield/src/benchmarking.rs @@ -0,0 +1,202 @@ +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 sp_std::{boxed::Box, vec, vec::Vec}; + +/// Helper to build bounded bytes (public key) of a given length. +fn bounded_pk(len: usize) -> BoundedVec> { + let v = vec![7u8; len]; + BoundedVec::>::try_from(v).expect("within bound; qed") +} + +/// Helper to build bounded bytes (ciphertext) of a given length. +fn bounded_ct(len: usize) -> BoundedVec> { + let v = vec![0u8; len]; + 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: +/// `::AuthorityId: From`. +fn seed_aura_authority_from_sr25519(pubkey: &sr25519::Public) +where + T: pallet::Config + pallet_aura::Config, + ::AuthorityId: From, +{ + let auth_id: ::AuthorityId = (*pubkey).into(); + pallet_aura::Authorities::::mutate(|auths| { + let _ = auths.try_push(auth_id); + }); +} + +#[benchmarks( + where + // Needed to build a concrete inner call and convert into T::RuntimeCall. + ::RuntimeCall: From>, + // Needed so we can seed Authorities from a dev sr25519 pubkey. + ::AuthorityId: From, +)] +mod benches { + use super::*; + + /// Benchmark `announce_next_key`. + #[benchmark] + fn announce_next_key() { + // Generate a deterministic dev key in the host keystore (for benchmarks). + // Any 4-byte KeyTypeId works for generation; it does not affect AccountId derivation. + const KT: KeyTypeId = KeyTypeId(*b"benc"); + let alice_pub: sr25519::Public = sr25519_generate(KT, Some("//Alice".as_bytes().to_vec())); + let alice_acc: AccountId32 = alice_pub.into(); + + // Make this account an Aura authority for the generic runtime. + seed_aura_authority_from_sr25519::(&alice_pub); + + // Valid Kyber768 public key length per pallet check. + const KYBER768_PK_LEN: usize = 1184; + let public_key: BoundedVec> = bounded_pk::<2048>(KYBER768_PK_LEN); + + // Measure: dispatch the extrinsic. + #[extrinsic_call] + announce_next_key(RawOrigin::Signed(alice_acc.clone()), public_key.clone()); + + // Assert: NextKey should be set exactly. + let stored = NextKey::::get().expect("must be set by announce_next_key"); + assert_eq!(stored, public_key); + } + + /// Benchmark `submit_encrypted`. + #[benchmark] + fn submit_encrypted() { + // Any whitelisted caller is fine (no authority requirement). + let who: T::AccountId = whitelisted_caller(); + + // Dummy commitment and ciphertext (bounded to 8192). + let commitment: T::Hash = ::Hashing::hash(b"bench-commitment"); + const CT_DEFAULT_LEN: usize = 256; + let ciphertext: BoundedVec> = super::bounded_ct::<8192>(CT_DEFAULT_LEN); + + // Pre-compute expected id to assert postconditions. + let id: T::Hash = + ::Hashing::hash_of(&(who.clone(), commitment, &ciphertext)); + + // Measure: dispatch the extrinsic. + #[extrinsic_call] + submit_encrypted( + RawOrigin::Signed(who.clone()), + commitment, + ciphertext.clone(), + ); + + // Assert: stored under expected id. + let got = Submissions::::get(id).expect("submission must exist"); + assert_eq!(got.author, who); + assert_eq!( + got.commitment, + ::Hashing::hash(b"bench-commitment") + ); + assert_eq!(got.ciphertext.as_slice(), ciphertext.as_slice()); + } + + /// Benchmark `execute_revealed`. + #[benchmark] + fn execute_revealed() { + // 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). + 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(); + + // 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. + const CT_DEFAULT_LEN: usize = 64; + let ciphertext: BoundedVec> = super::bounded_ct::<8192>(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. + let sub = Submission::, ::Hash> { + author: signer.clone(), + commitment, + ciphertext: ciphertext.clone(), + submitted_in: frame_system::Pallet::::block_number(), + }; + Submissions::::insert(id, sub); + + // 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. + #[extrinsic_call] + execute_revealed( + RawOrigin::None, + id, + signer.clone(), + nonce, + Box::new(inner_call.clone()), + signature.clone(), + ); + + // Assert: submission consumed, signer nonce bumped to 1. + 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 new file mode 100644 index 0000000000..1e1605b633 --- /dev/null +++ b/pallets/shield/src/lib.rs @@ -0,0 +1,347 @@ +// pallets/mev-shield/src/lib.rs +#![cfg_attr(not(feature = "std"), no_std)] + +pub use pallet::*; +#[cfg(feature = "runtime-benchmarks")] +mod benchmarking; + +#[cfg(test)] +pub mod mock; + +#[cfg(test)] +mod tests; + +#[frame_support::pallet] +pub mod pallet { + use super::*; + use codec::Encode; + use frame_support::{ + dispatch::{GetDispatchInfo, PostDispatchInfo}, + pallet_prelude::*, + traits::ConstU32, + weights::Weight, + }; + use frame_system::pallet_prelude::*; + use sp_consensus_aura::sr25519::AuthorityId as AuraAuthorityId; + use sp_core::ByteArray; + use sp_runtime::transaction_validity::{ + InvalidTransaction, TransactionSource, ValidTransaction, + }; + use sp_runtime::{ + AccountId32, DispatchErrorWithPostInfo, MultiSignature, RuntimeDebug, + traits::{BadOrigin, Dispatchable, Hash, SaturatedConversion, Verify, Zero}, + }; + use sp_std::{marker::PhantomData, prelude::*}; + use subtensor_macros::freeze_struct; + + /// Origin helper: ensure the signer is an Aura authority (no session/authorship). + pub struct EnsureAuraAuthority(PhantomData); + + pub trait AuthorityOriginExt { + type AccountId; + + fn ensure_validator(origin: Origin) -> Result; + } + + impl AuthorityOriginExt> for EnsureAuraAuthority + where + T: frame_system::Config + + pallet_aura::Config, + { + type AccountId = AccountId32; + + fn ensure_validator(origin: OriginFor) -> Result { + let who: AccountId32 = frame_system::ensure_signed(origin)?; + + let aura_id = + ::from_slice(who.as_ref()).map_err(|_| BadOrigin)?; + + let is_validator = pallet_aura::Authorities::::get() + .into_iter() + .any(|id| id == aura_id); + + if is_validator { + Ok(who) + } else { + Err(BadOrigin) + } + } + } + + // ----------------- Types ----------------- + + /// AEAD‑independent commitment over the revealed payload. + #[freeze_struct("66e393c88124f360")] + #[derive(Encode, Decode, Clone, PartialEq, Eq, RuntimeDebug, TypeInfo, MaxEncodedLen)] + pub struct Submission { + pub author: AccountId, + pub commitment: Hash, + pub ciphertext: BoundedVec>, + pub submitted_in: BlockNumber, + } + + // ----------------- Config ----------------- + + #[pallet::config] + pub trait Config: + frame_system::Config>> + + pallet_timestamp::Config + + pallet_aura::Config + { + type RuntimeCall: Parameter + + sp_runtime::traits::Dispatchable< + RuntimeOrigin = Self::RuntimeOrigin, + PostInfo = PostDispatchInfo, + > + GetDispatchInfo; + + type AuthorityOrigin: AuthorityOriginExt; + } + + #[pallet::pallet] + pub struct Pallet(_); + + // ----------------- Storage ----------------- + + #[pallet::storage] + pub type CurrentKey = StorageValue<_, BoundedVec>, OptionQuery>; + + #[pallet::storage] + pub type NextKey = StorageValue<_, BoundedVec>, OptionQuery>; + + #[pallet::storage] + pub type Submissions = StorageMap< + _, + Blake2_128Concat, + T::Hash, + Submission, T::Hash>, + OptionQuery, + >; + + // ----------------- Events & Errors ----------------- + + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// Encrypted wrapper accepted. + EncryptedSubmitted { id: T::Hash, who: T::AccountId }, + /// Decrypted call executed. + DecryptedExecuted { id: T::Hash, signer: T::AccountId }, + /// Decrypted execution rejected. + DecryptedRejected { + id: T::Hash, + reason: DispatchErrorWithPostInfo, + }, + } + + #[pallet::error] + pub enum Error { + SubmissionAlreadyExists, + MissingSubmission, + CommitmentMismatch, + SignatureInvalid, + NonceMismatch, + BadPublicKeyLen, + } + + // ----------------- Hooks ----------------- + + #[pallet::hooks] + impl Hooks> for Pallet { + fn on_initialize(_n: BlockNumberFor) -> Weight { + if let Some(next) = >::take() { + >::put(&next); + } + T::DbWeight::get().reads_writes(1, 2) + } + } + + // ----------------- Calls ----------------- + + #[pallet::call] + impl Pallet { + #[pallet::call_index(0)] + #[pallet::weight( + Weight::from_parts(9_979_000, 0) + .saturating_add(T::DbWeight::get().reads(1_u64)) + .saturating_add(T::DbWeight::get().writes(1_u64)) + )] + pub fn announce_next_key( + origin: OriginFor, + public_key: BoundedVec>, + ) -> DispatchResult { + // Only a current Aura validator may call this (signed account ∈ Aura authorities) + T::AuthorityOrigin::ensure_validator(origin)?; + + const MAX_KYBER768_PK_LENGTH: usize = 1184; + ensure!( + public_key.len() == MAX_KYBER768_PK_LENGTH, + Error::::BadPublicKeyLen + ); + + NextKey::::put(public_key.clone()); + + Ok(()) + } + + /// Users submit an encrypted wrapper. + /// + /// `commitment` is `blake2_256(raw_payload)`, where: + /// raw_payload = signer || nonce || SCALE(call) + /// + /// `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) + .saturating_add(T::DbWeight::get().reads(1_u64)) + .saturating_add(T::DbWeight::get().writes(1_u64)), + DispatchClass::Normal, + Pays::Yes, + ))] + pub fn submit_encrypted( + origin: OriginFor, + commitment: T::Hash, + ciphertext: BoundedVec>, + ) -> DispatchResult { + let who = ensure_signed(origin)?; + + let id: T::Hash = T::Hashing::hash_of(&(who.clone(), commitment, &ciphertext)); + let sub = Submission::, T::Hash> { + author: who.clone(), + commitment, + ciphertext, + submitted_in: >::block_number(), + }; + ensure!( + !Submissions::::contains_key(id), + Error::::SubmissionAlreadyExists + ); + Submissions::::insert(id, sub); + Self::deposit_event(Event::EncryptedSubmitted { id, who }); + Ok(()) + } + + /// Executed by the block author. + #[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)))] + #[allow(clippy::useless_conversion)] + pub fn execute_revealed( + origin: OriginFor, + id: T::Hash, + signer: T::AccountId, + nonce: T::Nonce, + call: Box<::RuntimeCall>, + signature: MultiSignature, + ) -> DispatchResultWithPostInfo { + ensure_none(origin)?; + + 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()); + + // 1) 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. + 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); + + // 4) Dispatch inner call from signer. + let info = call.get_dispatch_info(); + let required = info.call_weight.saturating_add(info.extension_weight); + + let origin_signed = frame_system::RawOrigin::Signed(signer.clone()).into(); + let res = (*call).dispatch(origin_signed); + + match res { + Ok(post) => { + let actual = post.actual_weight.unwrap_or(required); + Self::deposit_event(Event::DecryptedExecuted { id, signer }); + Ok(PostDispatchInfo { + actual_weight: Some(actual), + pays_fee: Pays::No, + }) + } + Err(e) => { + Self::deposit_event(Event::DecryptedRejected { id, reason: e }); + Ok(PostDispatchInfo { + actual_weight: Some(required), + pays_fee: Pays::No, + }) + } + } + } + } + + impl Pallet { + /// Build the raw payload bytes used for both: + /// - `commitment = blake2_256(raw_payload)` + /// - signature message (after domain separation). + /// + /// Layout: + /// signer (32B) || nonce (u32 LE) || SCALE(call) + fn build_raw_payload_bytes( + signer: &T::AccountId, + nonce: T::Nonce, + 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(call.encode()); + + out + } + } + + #[pallet::validate_unsigned] + impl ValidateUnsigned for Pallet { + type Call = Call; + + fn validate_unsigned(source: TransactionSource, call: &Self::Call) -> TransactionValidity { + match call { + Call::execute_revealed { id, .. } => { + match source { + // Only allow locally-submitted / already-in-block txs. + TransactionSource::Local | TransactionSource::InBlock => { + ValidTransaction::with_tag_prefix("mev-shield-exec") + .priority(u64::MAX) + .longevity(64) // long because propagate(false) + .and_provides(id) // dedupe by wrapper id + .propagate(false) // CRITICAL: no gossip, stays on author node + .build() + } + _ => InvalidTransaction::Call.into(), + } + } + + _ => InvalidTransaction::Call.into(), + } + } + } +} diff --git a/pallets/shield/src/mock.rs b/pallets/shield/src/mock.rs new file mode 100644 index 0000000000..0732670406 --- /dev/null +++ b/pallets/shield/src/mock.rs @@ -0,0 +1,145 @@ +use crate as pallet_mev_shield; + +use frame_support::{construct_runtime, derive_impl, parameter_types, traits::Everything}; +use frame_system as system; + +use sp_consensus_aura::sr25519::AuthorityId as AuraId; +use sp_core::{ConstU32, H256}; +use sp_runtime::traits::BadOrigin; +use sp_runtime::{ + AccountId32, BuildStorage, + traits::{BlakeTwo256, IdentityLookup}, +}; + +// ----------------------------------------------------------------------------- +// Mock runtime +// ----------------------------------------------------------------------------- + +pub type UncheckedExtrinsic = system::mocking::MockUncheckedExtrinsic; +pub type Block = system::mocking::MockBlock; + +construct_runtime!( + pub enum Test { + System: frame_system = 0, + Timestamp: pallet_timestamp = 1, + Aura: pallet_aura = 2, + MevShield: pallet_mev_shield = 3, + } +); + +// A concrete nonce type used in tests. +pub type TestNonce = u64; + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl system::Config for Test { + // Basic system config + type BaseCallFilter = Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + + type RuntimeOrigin = RuntimeOrigin; + type RuntimeCall = RuntimeCall; + type RuntimeEvent = RuntimeEvent; + + type Nonce = TestNonce; + type Hash = H256; + type Hashing = BlakeTwo256; + + type AccountId = AccountId32; + type Lookup = IdentityLookup; + type Block = Block; + + type BlockHashCount = (); + type Version = (); + type PalletInfo = PalletInfo; + + type AccountData = (); + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = (); + type OnSetCode = (); + + // Max number of consumer refs per account. + type MaxConsumers = ConstU32<16>; +} + +parameter_types! { + pub const MinimumPeriod: u64 = 1; +} + +impl pallet_timestamp::Config for Test { + type Moment = u64; + type OnTimestampSet = (); + type MinimumPeriod = MinimumPeriod; + type WeightInfo = (); +} + +// Aura mock configuration +parameter_types! { + pub const MaxAuthorities: u32 = 32; + pub const AllowMultipleBlocksPerSlot: bool = false; + pub const SlotDuration: u64 = 6000; +} + +impl pallet_aura::Config for Test { + type AuthorityId = AuraId; + // For tests we don't need dynamic disabling; just use unit type. + type DisabledValidators = (); + type MaxAuthorities = MaxAuthorities; + type AllowMultipleBlocksPerSlot = AllowMultipleBlocksPerSlot; + type SlotDuration = SlotDuration; +} + +// ----------------------------------------------------------------------------- +// Authority origin for tests – root-only +// ----------------------------------------------------------------------------- + +/// For tests, treat Root as the “validator set” and return a dummy AccountId. +pub struct TestAuthorityOrigin; + +impl pallet_mev_shield::AuthorityOriginExt for TestAuthorityOrigin { + type AccountId = AccountId32; + + fn ensure_validator(origin: RuntimeOrigin) -> Result { + // Must be a signed origin. + let who: AccountId32 = frame_system::ensure_signed(origin).map_err(|_| BadOrigin)?; + + // Interpret the AccountId bytes as an AuraId, just like the real pallet. + let aura_id = + ::from_slice(who.as_ref()).map_err(|_| BadOrigin)?; + + // Check membership in the Aura validator set. + let is_validator = pallet_aura::Authorities::::get() + .into_iter() + .any(|id| id == aura_id); + + if is_validator { + Ok(who) + } else { + Err(BadOrigin) + } + } +} + +// ----------------------------------------------------------------------------- +// MevShield Config +// ----------------------------------------------------------------------------- + +impl pallet_mev_shield::Config for Test { + type RuntimeCall = RuntimeCall; + type AuthorityOrigin = TestAuthorityOrigin; +} + +// ----------------------------------------------------------------------------- +// new_test_ext +// ----------------------------------------------------------------------------- + +pub fn new_test_ext() -> sp_io::TestExternalities { + // Use the construct_runtime!-generated genesis config. + RuntimeGenesisConfig::default() + .build_storage() + .expect("RuntimeGenesisConfig builds valid default genesis storage") + .into() +} diff --git a/pallets/shield/src/tests.rs b/pallets/shield/src/tests.rs new file mode 100644 index 0000000000..e3bc630014 --- /dev/null +++ b/pallets/shield/src/tests.rs @@ -0,0 +1,328 @@ +use crate as pallet_mev_shield; +use crate::mock::*; + +use codec::Encode; +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 pallet_mev_shield::{ + Call as MevShieldCall, CurrentKey, Event as MevShieldEvent, 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, +}; + +// ----------------------------------------------------------------------------- +// Helpers +// ----------------------------------------------------------------------------- + +/// Deterministic sr25519 pair for tests (acts as "Alice"). +fn test_sr25519_pair() -> sr25519::Pair { + sr25519::Pair::from_seed(&[1u8; 32]) +} + +/// Reproduce the pallet's raw payload layout: +/// signer (32B) || nonce (u32 LE) || SCALE(call) +fn build_raw_payload_bytes_for_test( + signer: &AccountId32, + nonce: TestNonce, + 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(call.encode()); + out +} + +#[test] +fn authority_can_announce_next_key_and_on_initialize_rolls_it() { + new_test_ext().execute_with(|| { + const KYBER_PK_LEN: usize = 1184; + let pk_bytes = vec![7u8; KYBER_PK_LEN]; + let bounded_pk: BoundedVec> = + BoundedVec::truncate_from(pk_bytes.clone()); + + // Seed Aura authorities with a single validator and derive the matching account. + let validator_pair = test_sr25519_pair(); + let validator_account: AccountId32 = validator_pair.public().into(); + let validator_aura_id: ::AuthorityId = + validator_pair.public().into(); + + // Authorities storage expects a BoundedVec. + let authorities: BoundedVec< + ::AuthorityId, + ::MaxAuthorities, + > = BoundedVec::truncate_from(vec![validator_aura_id.clone()]); + pallet_aura::Authorities::::put(authorities); + + assert!(CurrentKey::::get().is_none()); + assert!(NextKey::::get().is_none()); + + // Signed by an Aura validator -> passes TestAuthorityOrigin::ensure_validator. + assert_ok!(MevShield::announce_next_key( + RuntimeOrigin::signed(validator_account.clone()), + bounded_pk.clone(), + )); + + // NextKey storage updated + let next = NextKey::::get().expect("NextKey should be set"); + assert_eq!(next, pk_bytes); + + // Roll on new block + MevShield::on_initialize(2); + + let curr = CurrentKey::::get().expect("CurrentKey should be set"); + assert_eq!(curr, pk_bytes); + + assert!(NextKey::::get().is_none()); + }); +} + +#[test] +fn announce_next_key_rejects_non_validator_origins() { + new_test_ext().execute_with(|| { + const KYBER_PK_LEN: usize = 1184; + + // Validator account: bytes match the Aura authority we put into storage. + let validator_pair = test_sr25519_pair(); + let validator_account: AccountId32 = validator_pair.public().into(); + let validator_aura_id: ::AuthorityId = + validator_pair.public().into(); + + // Non‑validator is some other key (not in Aura::Authorities). + let non_validator_pair = sr25519::Pair::from_seed(&[2u8; 32]); + let non_validator: AccountId32 = non_validator_pair.public().into(); + + // Only the validator is in the Aura validator set. + let authorities: BoundedVec< + ::AuthorityId, + ::MaxAuthorities, + > = BoundedVec::truncate_from(vec![validator_aura_id.clone()]); + pallet_aura::Authorities::::put(authorities); + + let pk_bytes = vec![9u8; KYBER_PK_LEN]; + let bounded_pk: BoundedVec> = + BoundedVec::truncate_from(pk_bytes.clone()); + + // 1) Signed non‑validator origin must fail with BadOrigin. + assert_noop!( + MevShield::announce_next_key( + RuntimeOrigin::signed(non_validator.clone()), + bounded_pk.clone(), + ), + sp_runtime::DispatchError::BadOrigin + ); + + // 2) Unsigned origin must also fail with BadOrigin. + assert_noop!( + MevShield::announce_next_key(RuntimeOrigin::none(), bounded_pk.clone(),), + sp_runtime::DispatchError::BadOrigin + ); + + // 3) Signed validator origin succeeds (sanity check). + assert_ok!(MevShield::announce_next_key( + RuntimeOrigin::signed(validator_account.clone()), + bounded_pk.clone(), + )); + + let next = NextKey::::get().expect("NextKey must be set by validator"); + assert_eq!(next, pk_bytes); + }); +} + +#[test] +fn submit_encrypted_stores_submission_and_emits_event() { + new_test_ext().execute_with(|| { + let pair = test_sr25519_pair(); + let who: AccountId32 = pair.public().into(); + + System::set_block_number(10); + + let commitment = + ::Hashing::hash(b"test-mevshield-commitment"); + let ciphertext_bytes = vec![1u8, 2, 3, 4]; + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(ciphertext_bytes.clone()); + + assert_ok!(MevShield::submit_encrypted( + RuntimeOrigin::signed(who.clone()), + commitment, + ciphertext.clone(), + )); + + let id = ::Hashing::hash_of(&( + who.clone(), + commitment, + &ciphertext, + )); + + let stored = Submissions::::get(id).expect("submission stored"); + assert_eq!(stored.author, who); + assert_eq!(stored.commitment, commitment); + assert_eq!(stored.ciphertext.to_vec(), ciphertext_bytes); + assert_eq!(stored.submitted_in, 10); + + let events = System::events(); + let last = events.last().expect("at least one event").event.clone(); + + assert!( + matches!( + last, + RuntimeEvent::MevShield( + MevShieldEvent::::EncryptedSubmitted { id: ev_id, who: ev_who } + ) + if ev_id == id && ev_who == who + ), + "expected EncryptedSubmitted event with correct id & who", + ); + }); +} + +#[test] +fn execute_revealed_happy_path_verifies_and_executes_inner_call() { + new_test_ext().execute_with(|| { + let pair = test_sr25519_pair(); + let signer: AccountId32 = pair.public().into(); + + // Inner call – System.remark; must dispatch successfully. + let inner_call = RuntimeCall::System(frame_system::Call::::remark { + remark: b"hello-mevshield".to_vec(), + }); + + let nonce: TestNonce = Zero::zero(); + assert_eq!(System::account_nonce(&signer), nonce); + + let payload_bytes = build_raw_payload_bytes_for_test(&signer, nonce, &inner_call); + + let commitment = ::Hashing::hash(payload_bytes.as_ref()); + + let ciphertext_bytes = vec![9u8, 9, 9, 9]; + let ciphertext: BoundedVec> = + BoundedVec::truncate_from(ciphertext_bytes.clone()); + + System::set_block_number(1); + + // Wrapper author == signer for simplest path + assert_ok!(MevShield::submit_encrypted( + RuntimeOrigin::signed(signer.clone()), + commitment, + ciphertext.clone(), + )); + + let id = ::Hashing::hash_of(&( + signer.clone(), + commitment, + &ciphertext, + )); + + // Build message "mev-shield:v1" || genesis_hash || payload + 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(); + + let result = MevShield::execute_revealed( + RuntimeOrigin::none(), + id, + signer.clone(), + nonce, + Box::new(inner_call.clone()), + signature, + ); + + assert_ok!(result); + + // 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 + .last() + .expect("an event should be emitted") + .event + .clone(); + + assert!( + matches!( + last, + RuntimeEvent::MevShield( + MevShieldEvent::::DecryptedExecuted { id: ev_id, signer: ev_signer } + ) + if ev_id == id && ev_signer == signer + ), + "expected DecryptedExecuted event" + ); + }); +} + +#[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 signature: MultiSignature = sr25519::Signature::from_raw([0u8; 64]).into(); + + let call = MevShieldCall::::execute_revealed { + id, + signer, + nonce, + call: Box::new(inner_call), + signature, + }; + + let validity = MevShield::validate_unsigned(TransactionSource::Local, &call); + assert_ok!(validity); + }); +} + +#[test] +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 signature: MultiSignature = sr25519::Signature::from_raw([1u8; 64]).into(); + + let call = MevShieldCall::::execute_revealed { + id, + signer, + nonce, + call: Box::new(inner_call), + signature, + }; + + let validity = MevShield::validate_unsigned(TransactionSource::InBlock, &call); + assert_ok!(validity); + }); +} diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index d534dbb0c6..912866463e 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -2362,7 +2362,9 @@ mod dispatches { /// #[pallet::call_index(122)] #[pallet::weight(( - Weight::from_parts(19_420_000, 0).saturating_add(T::DbWeight::get().writes(4_u64)), + Weight::from_parts(19_420_000, 0) + .saturating_add(T::DbWeight::get().reads(2_u64)) + .saturating_add(T::DbWeight::get().writes(4_u64)), DispatchClass::Normal, Pays::Yes ))] @@ -2381,7 +2383,9 @@ mod dispatches { /// --- Sets root claim number (sudo extrinsic). Zero disables auto-claim. #[pallet::call_index(123)] #[pallet::weight(( - Weight::from_parts(4_000_000, 0).saturating_add(T::DbWeight::get().writes(1_u64)), + Weight::from_parts(4_000_000, 0) + .saturating_add(T::DbWeight::get().reads(0_u64)) + .saturating_add(T::DbWeight::get().writes(1_u64)), DispatchClass::Operational, Pays::Yes ))] @@ -2401,7 +2405,9 @@ mod dispatches { /// --- Sets root claim threshold for subnet (sudo or owner origin). #[pallet::call_index(124)] #[pallet::weight(( - Weight::from_parts(5_711_000, 0).saturating_add(T::DbWeight::get().writes(1_u64)), + Weight::from_parts(5_711_000, 0) + .saturating_add(T::DbWeight::get().reads(0_u64)) + .saturating_add(T::DbWeight::get().writes(1_u64)), DispatchClass::Operational, Pays::Yes ))] diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index 9760ac1b53..b3aced2160 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -149,6 +149,9 @@ ark-serialize = { workspace = true, features = ["derive"] } # Crowdloan pallet-crowdloan.workspace = true +# Mev Shield +pallet-shield.workspace = true + ethereum.workspace = true [dev-dependencies] @@ -271,6 +274,7 @@ std = [ "pallet-contracts/std", "subtensor-chain-extensions/std", "ethereum/std", + "pallet-shield/std", ] runtime-benchmarks = [ "frame-benchmarking/runtime-benchmarks", @@ -315,6 +319,7 @@ runtime-benchmarks = [ # Smart Tx fees pallet "subtensor-transaction-fee/runtime-benchmarks", + "pallet-shield/runtime-benchmarks", ] try-runtime = [ "frame-try-runtime/try-runtime", @@ -362,5 +367,6 @@ try-runtime = [ "pallet-drand/try-runtime", "runtime-common/try-runtime", "pallet-aura/try-runtime", + "pallet-shield/try-runtime", ] metadata-hash = ["substrate-wasm-builder/metadata-hash"] diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 2171708685..adff037ac8 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -29,6 +29,7 @@ use frame_system::{EnsureRoot, EnsureRootWithSuccess, EnsureSigned}; use pallet_commitments::{CanCommit, OnMetadataCommitment}; use pallet_grandpa::{AuthorityId as GrandpaId, fg_primitives}; use pallet_registry::CanRegisterIdentity; +pub use pallet_shield; use pallet_subtensor::rpc_info::{ delegate_info::DelegateInfo, dynamic_info::DynamicInfo, @@ -119,6 +120,22 @@ impl frame_system::offchain::SigningTypes for Runtime { type Signature = Signature; } +impl pallet_shield::Config for Runtime { + type RuntimeCall = RuntimeCall; + type AuthorityOrigin = pallet_shield::EnsureAuraAuthority; +} + +parameter_types! { + /// Milliseconds per slot; use the chain’s configured slot duration. + pub const ShieldSlotMs: u64 = SLOT_DURATION; + /// Emit the *next* ephemeral public key event at 7s. + pub const ShieldAnnounceAtMs: u64 = 7_000; + /// Old key remains accepted until 9s (2s grace). + pub const ShieldGraceMs: u64 = 2_000; + /// Last 3s of the slot reserved for decrypt+execute. + pub const ShieldDecryptWindowMs: u64 = 3_000; +} + impl frame_system::offchain::CreateTransactionBase for Runtime where RuntimeCall: From, @@ -220,7 +237,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { // `spec_version`, and `authoring_version` are the same between Wasm and native. // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use // the compatible custom types. - spec_version: 348, + spec_version: 349, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, @@ -1569,6 +1586,7 @@ construct_runtime!( Crowdloan: pallet_crowdloan = 27, Swap: pallet_subtensor_swap = 28, Contracts: pallet_contracts = 29, + MevShield: pallet_shield = 30, } ); @@ -1650,6 +1668,7 @@ mod benches { [pallet_drand, Drand] [pallet_crowdloan, Crowdloan] [pallet_subtensor_swap, Swap] + [pallet_shield, MevShield] ); } diff --git a/scripts/benchmark_action.sh b/scripts/benchmark_action.sh index a475211163..29bd7744d3 100755 --- a/scripts/benchmark_action.sh +++ b/scripts/benchmark_action.sh @@ -1,13 +1,14 @@ #!/usr/bin/env bash set -euo pipefail -PALLET_LIST=(subtensor admin_utils commitments drand) +PALLET_LIST=(subtensor admin_utils commitments drand shield) declare -A DISPATCH_PATHS=( [subtensor]="../pallets/subtensor/src/macros/dispatches.rs" [admin_utils]="../pallets/admin-utils/src/lib.rs" [commitments]="../pallets/commitments/src/lib.rs" [drand]="../pallets/drand/src/lib.rs" + [shield]="../pallets/shield/src/lib.rs" [swap]="../pallets/swap/src/pallet/mod.rs" ) diff --git a/scripts/benchmark_all.sh b/scripts/benchmark_all.sh index 22d23483f3..99bdb60338 100755 --- a/scripts/benchmark_all.sh +++ b/scripts/benchmark_all.sh @@ -6,6 +6,7 @@ pallets=( "pallet_commitments" "pallet_drand" "pallet_admin_utils" + "pallet_shield" ) RUNTIME_WASM=./target/production/wbuild/node-subtensor-runtime/node_subtensor_runtime.compact.compressed.wasm