Rework shielded transactions

node/Cargo.toml

@@ -26,6 +26,7 @@ clap = { workspace = true, features = ["derive"] }
 futures = { workspace = true, features = ["thread-pool"] }
 serde = { workspace = true, features = ["derive"] }
 hex.workspace = true
+tokio = { workspace = true, features = ["time"] }
 
 # Storage import
 memmap2.workspace = true
@@ -77,6 +78,7 @@ sp-keystore.workspace = true
 polkadot-sdk = { workspace = true, features = [
 	"cumulus-primitives-proof-size-hostfunction",
 ] }
+stp-io.workspace = true
 
 # These dependencies are used for the subtensor's RPCs
 jsonrpsee = { workspace = true, features = ["server"] }
@@ -94,6 +96,7 @@ pallet-transaction-payment-rpc-runtime-api.workspace = true
 # These dependencies are used for runtime benchmarking
 frame-benchmarking.workspace = true
 frame-benchmarking-cli.workspace = true
+pallet-subtensor.workspace = true
 
 # Needed for Frontier
 fc-mapping-sync.workspace = true
@@ -116,18 +119,8 @@ 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"
+stp-shield.workspace = true
+stc-shield.workspace = true
 
 # Local Dependencies
 node-subtensor-runtime = { workspace = true, features = ["std"] }
@@ -187,8 +180,8 @@ try-runtime = [
 	"pallet-commitments/try-runtime",
 	"pallet-drand/try-runtime",
 	"polkadot-sdk/try-runtime",
-	"pallet-shield/try-runtime",
 	"pallet-subtensor/try-runtime",
+	"pallet-shield/try-runtime",
 ]
 
 metadata-hash = ["node-subtensor-runtime/metadata-hash"]

node/src/benchmarking.rs

@@ -124,7 +124,7 @@ pub fn create_benchmark_extrinsic(
         .checked_next_power_of_two()
         .map(|c| c / 2)
         .unwrap_or(2) as u64;
-    let extra: runtime::TransactionExtensions =
+    let extra: runtime::TxExtension = (
         (
             frame_system::CheckNonZeroSender::<runtime::Runtime>::new(),
             frame_system::CheckSpecVersion::<runtime::Runtime>::new(),
@@ -136,32 +136,31 @@ pub fn create_benchmark_extrinsic(
             )),
             check_nonce::CheckNonce::<runtime::Runtime>::from(nonce),
             frame_system::CheckWeight::<runtime::Runtime>::new(),
-            transaction_payment_wrapper::ChargeTransactionPaymentWrapper::new(
-                pallet_transaction_payment::ChargeTransactionPayment::<runtime::Runtime>::from(0),
-            ),
+        ),
+        (
+            transaction_payment_wrapper::ChargeTransactionPaymentWrapper::new(0),
             sudo_wrapper::SudoTransactionExtension::<runtime::Runtime>::new(),
-            pallet_subtensor::transaction_extension::SubtensorTransactionExtension::<
-                runtime::Runtime,
-            >::new(),
+            pallet_shield::CheckShieldedTxValidity::<runtime::Runtime>::new(),
+            pallet_subtensor::SubtensorTransactionExtension::<runtime::Runtime>::new(),
             pallet_drand::drand_priority::DrandPriority::<runtime::Runtime>::new(),
-            frame_metadata_hash_extension::CheckMetadataHash::<runtime::Runtime>::new(true),
-        );
+        ),
+        frame_metadata_hash_extension::CheckMetadataHash::<runtime::Runtime>::new(true),
+    );
 
     let raw_payload = runtime::SignedPayload::from_raw(
         call.clone(),
         extra.clone(),
         (
-            (),
-            runtime::VERSION.spec_version,
-            runtime::VERSION.transaction_version,
-            genesis_hash,
-            best_hash,
-            (),
-            (),
-            (),
-            (),
-            (),
-            (),
+            (
+                (),
+                runtime::VERSION.spec_version,
+                runtime::VERSION.transaction_version,
+                genesis_hash,
+                best_hash,
+                (),
+                (),
+            ),
+            ((), (), (), (), ()),
             None,
         ),
     );

node/src/client.rs

@@ -16,5 +16,6 @@ pub type HostFunctions = (
     frame_benchmarking::benchmarking::HostFunctions,
     sp_crypto_ec_utils::bls12_381::host_calls::HostFunctions,
     ProofSize,
+    stp_io::SubtensorHostFunctions,
 );
 pub type RuntimeExecutor = WasmExecutor<HostFunctions>;

node/src/consensus/aura_consensus.rs

@@ -29,13 +29,15 @@ use sp_keystore::KeystorePtr;
 use sp_runtime::traits::Block as BlockT;
 use sp_runtime::traits::NumberFor;
 use std::{error::Error, sync::Arc};
+use stp_shield::ShieldKeystorePtr;
 
 pub struct AuraConsensus;
 
 impl ConsensusMechanism for AuraConsensus {
     type InherentDataProviders = (
         sp_consensus_aura::inherents::InherentDataProvider,
         sp_timestamp::InherentDataProvider,
+        stc_shield::InherentDataProvider,
     );
 
     fn start_authoring<C, SC, I, PF, SO, L, CIDP, BS, Error>(
@@ -101,18 +103,21 @@ impl ConsensusMechanism for AuraConsensus {
 
     fn create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn Error + Send + Sync>> {
         let timestamp = sp_timestamp::InherentDataProvider::from_system_time();
         let slot =
             sp_consensus_aura::inherents::InherentDataProvider::from_timestamp_and_slot_duration(
                 *timestamp,
                 slot_duration,
             );
-        Ok((slot, timestamp))
+        let shield = stc_shield::InherentDataProvider::new(shield_keystore);
+        Ok((slot, timestamp, shield))
     }
 
     fn pending_create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn Error + Send + Sync>> {
         let current = sp_timestamp::InherentDataProvider::from_system_time();
         let next_slot = current
@@ -125,7 +130,8 @@ impl ConsensusMechanism for AuraConsensus {
                 *timestamp,
                 slot_duration,
             );
-        Ok((slot, timestamp))
+        let shield = stc_shield::InherentDataProvider::new(shield_keystore);
+        Ok((slot, timestamp, shield))
     }
 
     fn new() -> Self {

node/src/consensus/babe_consensus.rs

@@ -32,6 +32,7 @@ use sp_inherents::CreateInherentDataProviders;
 use sp_keystore::KeystorePtr;
 use sp_runtime::traits::NumberFor;
 use std::{error::Error, sync::Arc};
+use stp_shield::ShieldKeystorePtr;
 
 pub struct BabeConsensus {
     babe_link: Option<BabeLink<Block>>,
@@ -42,6 +43,7 @@ impl ConsensusMechanism for BabeConsensus {
     type InherentDataProviders = (
         sp_consensus_babe::inherents::InherentDataProvider,
         sp_timestamp::InherentDataProvider,
+        stc_shield::InherentDataProvider,
     );
 
     #[allow(clippy::expect_used)]
@@ -111,18 +113,21 @@ impl ConsensusMechanism for BabeConsensus {
 
     fn create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn Error + Send + Sync>> {
         let timestamp = sp_timestamp::InherentDataProvider::from_system_time();
         let slot =
             sp_consensus_babe::inherents::InherentDataProvider::from_timestamp_and_slot_duration(
                 *timestamp,
                 slot_duration,
             );
-        Ok((slot, timestamp))
+        let shield = stc_shield::InherentDataProvider::new(shield_keystore);
+        Ok((slot, timestamp, shield))
     }
 
     fn pending_create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn Error + Send + Sync>> {
         let current = sp_timestamp::InherentDataProvider::from_system_time();
         let next_slot = current
@@ -135,7 +140,8 @@ impl ConsensusMechanism for BabeConsensus {
                 *timestamp,
                 slot_duration,
             );
-        Ok((slot, timestamp))
+        let shield = stc_shield::InherentDataProvider::new(shield_keystore);
+        Ok((slot, timestamp, shield))
     }
 
     fn new() -> Self {

node/src/consensus/consensus_mechanism.rs

@@ -24,6 +24,7 @@ use sp_keystore::KeystorePtr;
 use sp_runtime::traits::NumberFor;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
+use stp_shield::ShieldKeystorePtr;
 
 use crate::client::FullClient;
 use crate::service::BIQ;
@@ -85,11 +86,13 @@ pub trait ConsensusMechanism {
     /// Creates IDPs for the consensus mechanism.
     fn create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn std::error::Error + Send + Sync>>;
 
     /// Creates IDPs for the consensus mechanism for pending blocks.
     fn pending_create_inherent_data_providers(
         slot_duration: SlotDuration,
+        shield_keystore: ShieldKeystorePtr,
     ) -> Result<Self::InherentDataProviders, Box<dyn std::error::Error + Send + Sync>>;
 
     /// Creates the frontier consensus data provider with this mechanism.

node/src/lib.rs

@@ -4,6 +4,5 @@ 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;

node/src/main.rs

@@ -10,7 +10,6 @@ mod command;
 mod conditional_evm_block_import;
 mod consensus;
 mod ethereum;
-mod mev_shield;
 mod rpc;
 mod service;