feat: add associated types as tx hooks args (Sovereign-Labs#1225)

* feat: add associated types as tx hooks args

This commit introduces associated types to the `TxHooks` trait,
enhancing its ability to define custom interactions between the trait
and its implementing modules. The associated types allow for greater
flexibility in the implementation of `TxHooks`, enabling diverse use
cases beyond the current `Accounts` module.

The current `Accounts` implementation of `TxHooks` serves the purpose of
extracting a public key from a transaction and mapping it to a rollup
address, creating an execution context. This functionality is achieved
through the `pre_dispatch_tx_hook` signature which returns a
`C::Address`. However, this tight coupling with the `accounts` module
restricts future expansions of `TxHooks`.

With this modification, hooks gain increased versatility and modules are
free to implement them as per their requirements. Additionally, runtimes
can consume these implementations, utilizing them as essential building
blocks, while the banks module is empowered to enforce a gas cap on the
pre-dispatch hook.

* stf blueprint runtime tx hooks should return address

* add sequencer as argument of context

* fix missing test

* gen random sequencer for test

examples/demo-rollup/src/celestia_rollup.rs

@@ -4,7 +4,7 @@ use demo_stf::runtime::Runtime;
 use sov_celestia_adapter::verifier::{CelestiaSpec, CelestiaVerifier, RollupParams};
 use sov_celestia_adapter::{CelestiaConfig, CelestiaService};
 use sov_modules_api::default_context::{DefaultContext, ZkDefaultContext};
-use sov_modules_api::Spec;
+use sov_modules_api::{Address, Spec};
 use sov_modules_rollup_blueprint::{RollupBlueprint, WalletBlueprint};
 use sov_modules_stf_blueprint::kernels::basic::BasicKernel;
 use sov_modules_stf_blueprint::StfBlueprint;
@@ -57,12 +57,15 @@ impl RollupBlueprint for CelestiaDemoRollup {
         ledger_db: &sov_db::ledger_db::LedgerDB,
         da_service: &Self::DaService,
     ) -> Result<jsonrpsee::RpcModule<()>, anyhow::Error> {
+        // TODO set the sequencer address
+        let sequencer = Address::new([0; 32]);
+
         #[allow(unused_mut)]
         let mut rpc_methods = sov_modules_rollup_blueprint::register_rpc::<
             Self::NativeRuntime,
             Self::NativeContext,
             Self::DaService,
-        >(storage, ledger_db, da_service)?;
+        >(storage, ledger_db, da_service, sequencer)?;
 
         #[cfg(feature = "experimental")]
         crate::eth::register_ethereum::<Self::DaService>(

examples/demo-rollup/src/mock_rollup.rs

@@ -4,7 +4,7 @@ use demo_stf::runtime::Runtime;
 use sov_db::ledger_db::LedgerDB;
 use sov_mock_da::{MockDaConfig, MockDaService, MockDaSpec};
 use sov_modules_api::default_context::{DefaultContext, ZkDefaultContext};
-use sov_modules_api::Spec;
+use sov_modules_api::{Address, Spec};
 use sov_modules_rollup_blueprint::RollupBlueprint;
 use sov_modules_stf_blueprint::kernels::basic::BasicKernel;
 use sov_modules_stf_blueprint::StfBlueprint;
@@ -55,12 +55,15 @@ impl RollupBlueprint for MockDemoRollup {
         ledger_db: &LedgerDB,
         da_service: &Self::DaService,
     ) -> Result<jsonrpsee::RpcModule<()>, anyhow::Error> {
+        // TODO set the sequencer address
+        let sequencer = Address::new([0; 32]);
+
         #[allow(unused_mut)]
         let mut rpc_methods = sov_modules_rollup_blueprint::register_rpc::<
             Self::NativeRuntime,
             Self::NativeContext,
             Self::DaService,
-        >(storage, ledger_db, da_service)?;
+        >(storage, ledger_db, da_service, sequencer)?;
 
         #[cfg(feature = "experimental")]
         crate::eth::register_ethereum::<Self::DaService>(

examples/demo-rollup/stf/src/hooks_impl.rs

@@ -1,7 +1,8 @@
+use sov_accounts::AccountsTxHook;
 use sov_modules_api::hooks::{ApplyBlobHooks, FinalizeHook, SlotHooks, TxHooks};
 use sov_modules_api::transaction::Transaction;
 use sov_modules_api::{AccessoryWorkingSet, Context, Spec, WorkingSet};
-use sov_modules_stf_blueprint::SequencerOutcome;
+use sov_modules_stf_blueprint::{RuntimeTxHook, SequencerOutcome};
 #[cfg(feature = "experimental")]
 use sov_rollup_interface::da::BlockHeaderTrait;
 use sov_rollup_interface::da::{BlobReaderTrait, DaSpec};
@@ -13,24 +14,30 @@ use crate::runtime::Runtime;
 
 impl<C: Context, Da: DaSpec> TxHooks for Runtime<C, Da> {
     type Context = C;
+    type PreArg = RuntimeTxHook<C>;
+    type PreResult = C;
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         working_set: &mut WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
-        // Before executing a transaction, retrieve the sender's address from the accounts module
-        // and check the nonce
-        self.accounts.pre_dispatch_tx_hook(tx, working_set)
+        arg: RuntimeTxHook<C>,
+    ) -> anyhow::Result<C> {
+        let RuntimeTxHook { height, sequencer } = arg;
+        let AccountsTxHook { sender, sequencer } =
+            self.accounts
+                .pre_dispatch_tx_hook(tx, working_set, sequencer)?;
+        Ok(C::new(sender, sequencer, height))
     }
 
     fn post_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
+        ctx: &C,
         working_set: &mut WorkingSet<C>,
     ) -> anyhow::Result<()> {
-        // After executing each transaction, update the nonce
-        self.accounts.post_dispatch_tx_hook(tx, working_set)
+        self.accounts.post_dispatch_tx_hook(tx, ctx, working_set)?;
+        Ok(())
     }
 }
 

examples/simple-nft-module/tests/nft_test.rs

@@ -17,6 +17,7 @@ fn genesis_and_mint() {
     let admin = generate_address("admin");
     let owner1 = generate_address("owner2");
     let owner2 = generate_address("owner2");
+    let sequencer = generate_address("sequencer");
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
         owners: vec![(0, owner1)],
@@ -38,7 +39,7 @@ fn genesis_and_mint() {
 
     // Mint, anybody can mint
     let mint_message = CallMessage::Mint { id: 1 };
-    let owner2_context = C::new(owner2, 1);
+    let owner2_context = C::new(owner2, sequencer, 1);
     nft.call(mint_message.clone(), &owner2_context, &mut working_set)
         .expect("Minting failed");
 
@@ -61,9 +62,10 @@ fn genesis_and_mint() {
 fn transfer() {
     // Preparation
     let admin = generate_address("admin");
-    let admin_context = C::new(admin, 1);
+    let sequencer = generate_address("sequencer");
+    let admin_context = C::new(admin, sequencer, 1);
     let owner1 = generate_address("owner2");
-    let owner1_context = C::new(owner1, 1);
+    let owner1_context = C::new(owner1, sequencer, 1);
     let owner2 = generate_address("owner2");
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
@@ -117,9 +119,10 @@ fn transfer() {
 fn burn() {
     // Preparation
     let admin = generate_address("admin");
-    let admin_context = C::new(admin, 1);
+    let sequencer = generate_address("sequencer");
+    let admin_context = C::new(admin, sequencer, 1);
     let owner1 = generate_address("owner2");
-    let owner1_context = C::new(owner1, 1);
+    let owner1_context = C::new(owner1, sequencer, 1);
     let config: NonFungibleTokenConfig<C> = NonFungibleTokenConfig {
         admin,
         owners: vec![(0, owner1)],

full-node/sov-sequencer/src/batch_builder.rs

@@ -50,6 +50,7 @@ pub struct FiFoStrictBatchBuilder<C: Context, R: DispatchCall<Context = C>> {
     runtime: R,
     max_batch_size_bytes: usize,
     current_storage: C::Storage,
+    sequencer: C::Address,
 }
 
 impl<C, R> FiFoStrictBatchBuilder<C, R>
@@ -63,13 +64,15 @@ where
         mempool_max_txs_count: usize,
         runtime: R,
         current_storage: C::Storage,
+        sequencer: C::Address,
     ) -> Self {
         Self {
             mempool: VecDeque::new(),
             mempool_max_txs_count,
             max_batch_size_bytes,
             runtime,
             current_storage,
+            sequencer,
         }
     }
 }
@@ -139,7 +142,7 @@ where
                 // TODO: Bug(!), because potential discrepancy. Should be resolved by https://github.com/Sovereign-Labs/sovereign-sdk/issues/434
                 let sender_address: C::Address = pooled.tx.pub_key().to_address();
                 // FIXME! This should use the correct height
-                let ctx = C::new(sender_address, 0);
+                let ctx = C::new(sender_address, self.sequencer.clone(), 0);
 
                 if let Err(error) = self.runtime.dispatch_call(msg, &mut working_set, &ctx) {
                     warn!(%error, tx = hex::encode(&pooled.raw), "Error during transaction dispatch");
@@ -182,7 +185,9 @@ mod tests {
     use sov_modules_api::default_signature::DefaultPublicKey;
     use sov_modules_api::macros::DefaultRuntime;
     use sov_modules_api::transaction::Transaction;
-    use sov_modules_api::{Context, DispatchCall, EncodeCall, Genesis, MessageCodec, PrivateKey};
+    use sov_modules_api::{
+        Address, Context, DispatchCall, EncodeCall, Genesis, MessageCodec, PrivateKey,
+    };
     use sov_rollup_interface::services::batch_builder::BatchBuilder;
     use sov_state::{DefaultStorageSpec, ProverStorage, Storage};
     use sov_value_setter::{CallMessage, ValueSetter, ValueSetterConfig};
@@ -246,11 +251,13 @@ mod tests {
     ) {
         let storage = ProverStorage::<DefaultStorageSpec>::with_path(tmpdir.path()).unwrap();
 
+        let sequencer = Address::from([0; 32]);
         let batch_builder = FiFoStrictBatchBuilder::new(
             batch_size_bytes,
             MAX_TX_POOL_SIZE,
             TestRuntime::<C>::default(),
             storage.clone(),
+            sequencer,
         );
         (batch_builder, storage)
     }

fuzz/fuzz_targets/accounts_call.rs

@@ -15,8 +15,8 @@ use sov_modules_api::{Context, Module, PrivateKey, Spec, WorkingSet};
 type C = DefaultContext;
 
 // Check well-formed calls
-fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
-    let (iterations, seed, keys) = input;
+fuzz_target!(|input: (u16, [u8; 32], [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
+    let (iterations, seed, sequencer, keys) = input;
     if iterations < 1024 {
         // pointless to setup & run a small iterations count
         return Corpus::Reject;
@@ -42,6 +42,7 @@ fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
     let storage = <C as Spec>::Storage::with_path(tmpdir.path()).unwrap();
     let working_set = &mut WorkingSet::new(storage);
 
+    let sequencer = <C as Spec>::Address::from(sequencer);
     let config: AccountConfig<C> = keys.iter().map(|k| k.pub_key()).collect();
     let accounts: Accounts<C> = Accounts::default();
     accounts.genesis(&config, working_set).unwrap();
@@ -54,7 +55,7 @@ fuzz_target!(|input: (u16, [u8; 32], Vec<DefaultPrivateKey>)| -> Corpus {
     for i in 0..iterations {
         // we use slices for better select performance
         let sender = addresses.choose(rng).unwrap();
-        let context = C::new(*sender, i as u64);
+        let context = C::new(*sender, sequencer, i as u64);
 
         // clear previous state
         let previous = state.get(sender).unwrap().as_hex();

fuzz/fuzz_targets/bank_call.rs

@@ -9,12 +9,12 @@ use sov_state::ProverStorage;
 
 type C = DefaultContext;
 
-fuzz_target!(|input: (&[u8], [u8; 32])| {
-    let (data, sender) = input;
+fuzz_target!(|input: (&[u8], [u8; 32], [u8; 32])| {
+    let (data, sender, sequencer) = input;
     if let Ok(msgs) = serde_json::from_slice::<Vec<CallMessage<C>>>(data) {
         let tmpdir = tempfile::tempdir().unwrap();
         let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
-        let ctx = C::new(sender.into(), 1);
+        let ctx = C::new(sender.into(), sequencer.into(), 1);
         let bank = Bank::default();
         for msg in msgs {
             bank.call(msg, &ctx, &mut working_set).ok();

module-system/module-implementations/examples/sov-accessory-state/tests/test.rs

@@ -18,7 +18,8 @@ fn test_accessory_value_setter() {
     let mut working_set_for_check: WorkingSet<DefaultContext> = WorkingSet::new(storage.clone());
 
     let admin = Address::from([1; 32]);
-    let context = DefaultContext::new(admin, 1);
+    let sequencer = Address::from([2; 32]);
+    let context = DefaultContext::new(admin, sequencer, 1);
 
     let module = AccessorySetter::<DefaultContext>::default();
 

module-system/module-implementations/examples/sov-value-setter/src/tests.rs

@@ -10,11 +10,12 @@ fn test_value_setter() {
     let tmpdir = tempfile::tempdir().unwrap();
     let mut working_set = WorkingSet::new(ProverStorage::with_path(tmpdir.path()).unwrap());
     let admin = Address::from([1; 32]);
+    let sequencer = Address::from([2; 32]);
     // Test Native-Context
     #[cfg(feature = "native")]
     {
         let config = ValueSetterConfig { admin };
-        let context = DefaultContext::new(admin, 1);
+        let context = DefaultContext::new(admin, sequencer, 1);
         test_value_setter_helper(context, &config, &mut working_set);
     }
 
@@ -23,7 +24,7 @@ fn test_value_setter() {
     // Test Zk-Context
     {
         let config = ValueSetterConfig { admin };
-        let zk_context = ZkDefaultContext::new(admin, 1);
+        let zk_context = ZkDefaultContext::new(admin, sequencer, 1);
         let mut zk_working_set = WorkingSet::with_witness(ZkStorage::new(), witness);
         test_value_setter_helper(zk_context, &config, &mut zk_working_set);
     }
@@ -63,6 +64,7 @@ fn test_value_setter_helper<C: Context>(
 #[test]
 fn test_err_on_sender_is_not_admin() {
     let sender = Address::from([1; 32]);
+    let sequencer = Address::from([2; 32]);
 
     let tmpdir = tempfile::tempdir().unwrap();
     let backing_store = ProverStorage::with_path(tmpdir.path()).unwrap();
@@ -75,7 +77,7 @@ fn test_err_on_sender_is_not_admin() {
         let config = ValueSetterConfig {
             admin: sender_not_admin,
         };
-        let context = DefaultContext::new(sender, 1);
+        let context = DefaultContext::new(sender, sequencer, 1);
         test_err_on_sender_is_not_admin_helper(context, &config, &mut native_working_set);
     }
     let (_, witness) = native_working_set.checkpoint().freeze();
@@ -86,7 +88,7 @@ fn test_err_on_sender_is_not_admin() {
             admin: sender_not_admin,
         };
         let zk_backing_store = ZkStorage::new();
-        let zk_context = ZkDefaultContext::new(sender, 1);
+        let zk_context = ZkDefaultContext::new(sender, sequencer, 1);
         let zk_working_set = &mut WorkingSet::with_witness(zk_backing_store, witness);
         test_err_on_sender_is_not_admin_helper(zk_context, &config, zk_working_set);
     }

module-system/module-implementations/examples/sov-vec-setter/tests/tests.rs

@@ -5,24 +5,25 @@ use sov_vec_setter::{CallMessage, VecSetter, VecSetterConfig};
 
 // rustfmt doesn't like long lines, but it's easier to read in this case.
 #[rustfmt::skip]
-fn test_cases() -> Vec<(Address, CallMessage, Option<Vec<u32>>)> {
+fn test_cases() -> Vec<(Address, Address, CallMessage, Option<Vec<u32>>)> {
     let admin = Address::from([1; 32]);
     let not_admin = Address::from([2; 32]);
+    let sequencer = Address::from([3; 32]);
 
     // (sender, call, expected vec contents or None if call should fail)
     vec![
-        (admin, CallMessage::PushValue(1), Some(vec![1])),
-        (admin, CallMessage::PushValue(2), Some(vec![1, 2])),
-        (admin, CallMessage::PopValue, Some(vec![1])),
-        (not_admin, CallMessage::PopValue, None),
-        (admin, CallMessage::PopValue, Some(vec![])),
-        (not_admin, CallMessage::SetValue { index: 0, value: 10 }, None),
-        (admin, CallMessage::SetValue { index: 0, value: 10 }, None),
-        (admin, CallMessage::PushValue(8), Some(vec![8])),
-        (admin, CallMessage::SetValue { index: 0, value: 10 }, Some(vec![10])),
-        (admin, CallMessage::PushValue(0), Some(vec![10, 0])),
-        (admin, CallMessage::SetAllValues(vec![11, 12]), Some(vec![11, 12])),
-        (not_admin, CallMessage::SetAllValues(vec![]), None),
+        (admin, sequencer, CallMessage::PushValue(1), Some(vec![1])),
+        (admin, sequencer, CallMessage::PushValue(2), Some(vec![1, 2])),
+        (admin, sequencer, CallMessage::PopValue, Some(vec![1])),
+        (not_admin, sequencer, CallMessage::PopValue, None),
+        (admin, sequencer, CallMessage::PopValue, Some(vec![])),
+        (not_admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, None),
+        (admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, None),
+        (admin, sequencer, CallMessage::PushValue(8), Some(vec![8])),
+        (admin, sequencer, CallMessage::SetValue { index: 0, value: 10 }, Some(vec![10])),
+        (admin, sequencer, CallMessage::PushValue(0), Some(vec![10, 0])),
+        (admin, sequencer, CallMessage::SetAllValues(vec![11, 12]), Some(vec![11, 12])),
+        (not_admin, sequencer, CallMessage::SetAllValues(vec![]), None),
     ]
 }
 
@@ -39,8 +40,8 @@ fn test_vec_setter_calls() {
     let vec_setter = VecSetter::default();
     vec_setter.genesis(&config, &mut working_set).unwrap();
 
-    for (sender, call, expected_contents) in test_cases().iter().cloned() {
-        let context = DefaultContext::new(sender, 1);
+    for (sender, sequencer, call, expected_contents) in test_cases().iter().cloned() {
+        let context = DefaultContext::new(sender, sequencer, 1);
 
         let call_result = vec_setter.call(call, &context, &mut working_set);
 

module-system/module-implementations/integration-tests/src/chain_state/helpers.rs

@@ -7,7 +7,7 @@ use sov_modules_api::{
     AccessoryWorkingSet, BlobReaderTrait, Context, DaSpec, DispatchCall, Genesis, MessageCodec,
     PublicKey, Spec,
 };
-use sov_modules_stf_blueprint::{Runtime, SequencerOutcome};
+use sov_modules_stf_blueprint::{Runtime, RuntimeTxHook, SequencerOutcome};
 use sov_state::Storage;
 use sov_value_setter::{ValueSetter, ValueSetterConfig};
 
@@ -20,18 +20,25 @@ pub(crate) struct TestRuntime<C: Context, Da: DaSpec> {
 
 impl<C: Context, Da: DaSpec> TxHooks for TestRuntime<C, Da> {
     type Context = C;
+    type PreArg = RuntimeTxHook<C>;
+    type PreResult = C;
 
     fn pre_dispatch_tx_hook(
         &self,
         tx: &Transaction<Self::Context>,
         _working_set: &mut sov_modules_api::WorkingSet<C>,
-    ) -> anyhow::Result<<Self::Context as Spec>::Address> {
-        Ok(tx.pub_key().to_address())
+        arg: RuntimeTxHook<C>,
+    ) -> anyhow::Result<C> {
+        let RuntimeTxHook { height, sequencer } = arg;
+        let sender = tx.pub_key().to_address();
+        let sequencer = sequencer.to_address();
+        Ok(C::new(sender, sequencer, height))
     }
 
     fn post_dispatch_tx_hook(
         &self,
         _tx: &Transaction<Self::Context>,
+        _ctx: &C,
         _working_set: &mut sov_modules_api::WorkingSet<C>,
     ) -> anyhow::Result<()> {
         Ok(())