Rework shielded transactions

.github/workflows/e2e.yml

@@ -1,4 +1,15 @@
 name: E2E Tests
 
+concurrency:
+  group: e2e-${{ github.ref }}
+  cancel-in-progress: true
+
 on:
-  workflow_dispatch:
+  pull_request:
+
+  workflow_dispatch:
+    inputs:
+      verbose:
+        description: "Output more information when triggered manually"
+        required: false
+        default: ""

common/src/lib.rs

@@ -15,8 +15,10 @@ use sp_runtime::{
 use subtensor_macros::freeze_struct;
 
 pub use currency::*;
+pub use transaction_error::*;
 
 mod currency;
+mod transaction_error;
 
 /// Balance of an account.
 pub type Balance = u64;

common/src/transaction_error.rs

@@ -0,0 +1,70 @@
+use sp_runtime::transaction_validity::{InvalidTransaction, TransactionValidityError};
+
+#[derive(Debug, PartialEq)]
+pub enum CustomTransactionError {
+    ColdkeyInSwapSchedule,
+    StakeAmountTooLow,
+    BalanceTooLow,
+    SubnetNotExists,
+    HotkeyAccountDoesntExist,
+    NotEnoughStakeToWithdraw,
+    RateLimitExceeded,
+    InsufficientLiquidity,
+    SlippageTooHigh,
+    TransferDisallowed,
+    HotKeyNotRegisteredInNetwork,
+    InvalidIpAddress,
+    ServingRateLimitExceeded,
+    InvalidPort,
+    BadRequest,
+    ZeroMaxAmount,
+    InvalidRevealRound,
+    CommitNotFound,
+    CommitBlockNotInRevealRange,
+    InputLengthsUnequal,
+    UidNotFound,
+    EvmKeyAssociateRateLimitExceeded,
+    ColdkeySwapDisputed,
+    InvalidRealAccount,
+    FailedShieldedTxParsing,
+    InvalidShieldedTxPubKeyHash,
+}
+
+impl From<CustomTransactionError> for u8 {
+    fn from(variant: CustomTransactionError) -> u8 {
+        match variant {
+            CustomTransactionError::ColdkeyInSwapSchedule => 0,
+            CustomTransactionError::StakeAmountTooLow => 1,
+            CustomTransactionError::BalanceTooLow => 2,
+            CustomTransactionError::SubnetNotExists => 3,
+            CustomTransactionError::HotkeyAccountDoesntExist => 4,
+            CustomTransactionError::NotEnoughStakeToWithdraw => 5,
+            CustomTransactionError::RateLimitExceeded => 6,
+            CustomTransactionError::InsufficientLiquidity => 7,
+            CustomTransactionError::SlippageTooHigh => 8,
+            CustomTransactionError::TransferDisallowed => 9,
+            CustomTransactionError::HotKeyNotRegisteredInNetwork => 10,
+            CustomTransactionError::InvalidIpAddress => 11,
+            CustomTransactionError::ServingRateLimitExceeded => 12,
+            CustomTransactionError::InvalidPort => 13,
+            CustomTransactionError::BadRequest => 255,
+            CustomTransactionError::ZeroMaxAmount => 14,
+            CustomTransactionError::InvalidRevealRound => 15,
+            CustomTransactionError::CommitNotFound => 16,
+            CustomTransactionError::CommitBlockNotInRevealRange => 17,
+            CustomTransactionError::InputLengthsUnequal => 18,
+            CustomTransactionError::UidNotFound => 19,
+            CustomTransactionError::EvmKeyAssociateRateLimitExceeded => 20,
+            CustomTransactionError::ColdkeySwapDisputed => 21,
+            CustomTransactionError::InvalidRealAccount => 22,
+            CustomTransactionError::FailedShieldedTxParsing => 23,
+            CustomTransactionError::InvalidShieldedTxPubKeyHash => 24,
+        }
+    }
+}
+
+impl From<CustomTransactionError> for TransactionValidityError {
+    fn from(variant: CustomTransactionError) -> Self {
+        TransactionValidityError::Invalid(InvalidTransaction::Custom(variant.into()))
+    }
+}

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,33 @@ 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_shield::CheckShieldedTxValidity::<runtime::Runtime>::new(),
             pallet_subtensor::transaction_extension::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

@@ -28,14 +28,17 @@ use sp_inherents::CreateInherentDataProviders;
 use sp_keystore::KeystorePtr;
 use sp_runtime::traits::Block as BlockT;
 use sp_runtime::traits::NumberFor;
+use stc_shield::InherentDataProvider as ShieldInherentDataProvider;
 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 +104,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 = ShieldInherentDataProvider::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 +131,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

@@ -31,7 +31,9 @@ use sp_consensus_slots::SlotDuration;
 use sp_inherents::CreateInherentDataProviders;
 use sp_keystore::KeystorePtr;
 use sp_runtime::traits::NumberFor;
+use stc_shield::InherentDataProvider as ShieldInherentDataProvider;
 use std::{error::Error, sync::Arc};
+use stp_shield::ShieldKeystorePtr;
 
 pub struct BabeConsensus {
     babe_link: Option<BabeLink<Block>>,
@@ -42,6 +44,7 @@ impl ConsensusMechanism for BabeConsensus {
     type InherentDataProviders = (
         sp_consensus_babe::inherents::InherentDataProvider,
         sp_timestamp::InherentDataProvider,
+        stc_shield::InherentDataProvider,
     );
 
     #[allow(clippy::expect_used)]
@@ -111,18 +114,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 +141,8 @@ impl ConsensusMechanism for BabeConsensus {
                 *timestamp,
                 slot_duration,
             );
-        Ok((slot, timestamp))
+        let shield = ShieldInherentDataProvider::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;