refactor(katana): chain id type (#1351)

dojoengine · Jan 11, 2024 · 70360e2 · 70360e2
1 parent 7157b18
commit 70360e2
Show file tree

Hide file tree

Showing 21 changed files with 330 additions and 62 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/dojo-test-utils/Cargo.toml b/crates/dojo-test-utils/Cargo.toml
@@ -18,6 +18,7 @@ dojo-lang = { path = "../dojo-lang" }
 dojo-world = { path = "../dojo-world", features = [ "manifest", "migration" ] }
 jsonrpsee = { version = "0.16.2", features = [ "server" ] }
 katana-core = { path = "../katana/core" }
+katana-primitives = { path = "../katana/primitives" }
 katana-rpc = { path = "../katana/rpc" }
 scarb-ui.workspace = true
 scarb.workspace = true

diff --git a/crates/dojo-test-utils/src/sequencer.rs b/crates/dojo-test-utils/src/sequencer.rs
@@ -4,6 +4,7 @@ use jsonrpsee::core::Error;
 pub use katana_core::backend::config::{Environment, StarknetConfig};
 use katana_core::sequencer::KatanaSequencer;
 pub use katana_core::sequencer::SequencerConfig;
+use katana_primitives::chain::ChainId;
 use katana_rpc::api::ApiKind;
 use katana_rpc::config::ServerConfig;
 use katana_rpc::{spawn, NodeHandle};
@@ -79,7 +80,7 @@ impl TestSequencer {
 pub fn get_default_test_starknet_config() -> StarknetConfig {
     StarknetConfig {
         disable_fee: true,
-        env: Environment { chain_id: "SN_GOERLI".into(), ..Default::default() },
+        env: Environment { chain_id: ChainId::GOERLI, ..Default::default() },
         ..Default::default()
     }
 }
diff --git a/crates/katana/Cargo.toml b/crates/katana/Cargo.toml
@@ -11,6 +11,7 @@ clap.workspace = true
 clap_complete.workspace = true
 console.workspace = true
 katana-core = { path = "core" }
+katana-primitives = { path = "primitives" }
 katana-rpc = { path = "rpc" }
 metrics = { path = "../metrics" }
 metrics-process.workspace = true

diff --git a/crates/katana/core/src/backend/config.rs b/crates/katana/core/src/backend/config.rs
@@ -1,6 +1,6 @@
 use blockifier::block_context::{BlockContext, FeeTokenAddresses, GasPrices};
+use katana_primitives::chain::ChainId;
 use starknet_api::block::{BlockNumber, BlockTimestamp};
-use starknet_api::core::ChainId;
 use url::Url;
 
 use crate::constants::{
@@ -24,7 +24,7 @@ impl StarknetConfig {
     pub fn block_context(&self) -> BlockContext {
         BlockContext {
             block_number: BlockNumber::default(),
-            chain_id: ChainId(self.env.chain_id.clone()),
+            chain_id: self.env.chain_id.into(),
             block_timestamp: BlockTimestamp::default(),
             sequencer_address: (*SEQUENCER_ADDRESS).into(),
             // As the fee has two currencies, we also have to adjust their addresses.
@@ -67,7 +67,7 @@ impl Default for StarknetConfig {
 
 #[derive(Debug, Clone)]
 pub struct Environment {
-    pub chain_id: String,
+    pub chain_id: ChainId,
     pub gas_price: u128,
     pub invoke_max_steps: u32,
     pub validate_max_steps: u32,
@@ -77,7 +77,7 @@ impl Default for Environment {
     fn default() -> Self {
         Self {
             gas_price: DEFAULT_GAS_PRICE,
-            chain_id: "KATANA".to_string(),
+            chain_id: ChainId::parse("KATANA").unwrap(),
             invoke_max_steps: DEFAULT_INVOKE_MAX_STEPS,
             validate_max_steps: DEFAULT_VALIDATE_MAX_STEPS,
         }

diff --git a/crates/katana/core/src/backend/mod.rs b/crates/katana/core/src/backend/mod.rs
@@ -4,6 +4,7 @@ use blockifier::block_context::BlockContext;
 use katana_primitives::block::{
     Block, FinalityStatus, GasPrices, Header, PartialHeader, SealedBlockWithStatus,
 };
+use katana_primitives::chain::ChainId;
 use katana_primitives::contract::ContractAddress;
 use katana_primitives::receipt::Receipt;
 use katana_primitives::state::StateUpdatesWithDeclaredClasses;
@@ -20,7 +21,6 @@ use starknet::core::utils::parse_cairo_short_string;
 use starknet::providers::jsonrpc::HttpTransport;
 use starknet::providers::{JsonRpcClient, Provider};
 use starknet_api::block::{BlockNumber, BlockTimestamp};
-use starknet_api::core::ChainId;
 use tracing::{info, trace};
 
 pub mod config;
@@ -40,6 +40,8 @@ pub struct Backend {
     pub config: RwLock<StarknetConfig>,
     /// stores all block related data in memory
     pub blockchain: Blockchain,
+    /// The chain id.
+    pub chain_id: ChainId,
     /// The chain environment values.
     pub env: Arc<RwLock<Env>>,
     pub block_context_generator: RwLock<BlockContextGenerator>,
@@ -57,7 +59,9 @@ impl Backend {
             .with_balance(*DEFAULT_PREFUNDED_ACCOUNT_BALANCE)
             .generate();
 
-        let blockchain: Blockchain = if let Some(forked_url) = &config.fork_rpc_url {
+        let (blockchain, chain_id): (Blockchain, ChainId) = if let Some(forked_url) =
+            &config.fork_rpc_url
+        {
             let provider = Arc::new(JsonRpcClient::new(HttpTransport::new(forked_url.clone())));
             let forked_chain_id = provider.chain_id().await.unwrap();
 
@@ -79,7 +83,8 @@ impl Backend {
             block_context.block_number = BlockNumber(block.block_number);
             block_context.block_timestamp = BlockTimestamp(block.timestamp);
             block_context.sequencer_address = ContractAddress(block.sequencer_address).into();
-            block_context.chain_id = ChainId(parse_cairo_short_string(&forked_chain_id).unwrap());
+            block_context.chain_id =
+                starknet_api::core::ChainId(parse_cairo_short_string(&forked_chain_id).unwrap());
 
             trace!(
                 target: "backend",
@@ -89,7 +94,7 @@ impl Backend {
                 forked_url
             );
 
-            Blockchain::new_from_forked(
+            let blockchain = Blockchain::new_from_forked(
                 ForkedProvider::new(provider, forked_block_num.into()),
                 block.block_hash,
                 block.parent_hash,
@@ -101,10 +106,14 @@ impl Backend {
                     _ => panic!("unable to fork for non-accepted block"),
                 },
             )
-            .expect("able to create forked blockchain")
+            .expect("able to create forked blockchain");
+
+            (blockchain, forked_chain_id.into())
         } else {
-            Blockchain::new_with_genesis(InMemoryProvider::new(), &block_context)
-                .expect("able to create blockchain from genesis block")
+            let blockchain = Blockchain::new_with_genesis(InMemoryProvider::new(), &block_context)
+                .expect("able to create blockchain from genesis block");
+
+            (blockchain, config.env.chain_id)
         };
 
         let env = Env { block: block_context };
@@ -115,6 +124,7 @@ impl Backend {
         }
 
         Self {
+            chain_id,
             accounts,
             blockchain,
             config: RwLock::new(config),

diff --git a/crates/katana/core/src/sequencer.rs b/crates/katana/core/src/sequencer.rs
@@ -10,6 +10,7 @@ use katana_executor::blockifier::state::StateRefDb;
 use katana_executor::blockifier::utils::EntryPointCall;
 use katana_executor::blockifier::PendingState;
 use katana_primitives::block::{BlockHash, BlockHashOrNumber, BlockIdOrTag, BlockNumber};
+use katana_primitives::chain::ChainId;
 use katana_primitives::contract::{
     ClassHash, CompiledContractClass, ContractAddress, Nonce, StorageKey, StorageValue,
 };
@@ -26,7 +27,6 @@ use katana_provider::traits::transaction::{
     ReceiptProvider, TransactionProvider, TransactionsProviderExt,
 };
 use starknet::core::types::{BlockTag, EmittedEvent, EventsPage, FeeEstimate};
-use starknet_api::core::ChainId;
 
 use crate::backend::config::StarknetConfig;
 use crate::backend::contract::StarknetContract;
@@ -215,7 +215,7 @@ impl KatanaSequencer {
     }
 
     pub fn chain_id(&self) -> ChainId {
-        self.backend.env.read().block.chain_id.clone()
+        self.backend.chain_id
     }
 
     pub fn block_number(&self) -> BlockNumber {

diff --git a/crates/katana/core/src/service/messaging/ethereum.rs b/crates/katana/core/src/service/messaging/ethereum.rs
@@ -8,6 +8,7 @@ use ethers::prelude::*;
 use ethers::providers::{Http, Provider};
 use ethers::types::{Address, BlockNumber, Log};
 use k256::ecdsa::SigningKey;
+use katana_primitives::chain::ChainId;
 use katana_primitives::receipt::MessageToL1;
 use katana_primitives::transaction::L1HandlerTx;
 use katana_primitives::utils::transaction::compute_l1_message_hash;
@@ -127,7 +128,7 @@ impl Messenger for EthereumMessaging {
         &self,
         from_block: u64,
         max_blocks: u64,
-        chain_id: FieldElement,
+        chain_id: ChainId,
     ) -> MessengerResult<(u64, Vec<Self::MessageTransaction>)> {
         let chain_latest_block: u64 = self
             .provider
@@ -206,7 +207,7 @@ impl Messenger for EthereumMessaging {
     }
 }
 
-fn l1_handler_tx_from_log(log: Log, chain_id: FieldElement) -> MessengerResult<L1HandlerTx> {
+fn l1_handler_tx_from_log(log: Log, chain_id: ChainId) -> MessengerResult<L1HandlerTx> {
     let parsed_log = <LogMessageToL2 as EthLogDecode>::decode_log(&log.into()).map_err(|e| {
         error!(target: LOG_TARGET, "Log parsing failed {e}");
         Error::GatherError
@@ -259,6 +260,7 @@ fn felt_from_address(v: Address) -> FieldElement {
 #[cfg(test)]
 mod tests {
 
+    use katana_primitives::chain::{ChainId, NamedChainId};
     use starknet::macros::{felt, selector};
 
     use super::*;
@@ -299,7 +301,7 @@ mod tests {
         };
 
         // SN_GOERLI.
-        let chain_id = starknet::macros::felt!("0x534e5f474f45524c49");
+        let chain_id = ChainId::Named(NamedChainId::Goerli);
         let to_address = FieldElement::from_hex_be(to_address).unwrap();
         let from_address = FieldElement::from_hex_be(from_address).unwrap();
 

diff --git a/crates/katana/core/src/service/messaging/mod.rs b/crates/katana/core/src/service/messaging/mod.rs
@@ -39,12 +39,12 @@ mod starknet;
 
 use std::path::Path;
 
-use ::starknet::core::types::FieldElement;
 use ::starknet::providers::ProviderError as StarknetProviderError;
 use anyhow::Result;
 use async_trait::async_trait;
 use ethereum::EthereumMessaging;
 use ethers::providers::ProviderError as EthereumProviderError;
+use katana_primitives::chain::ChainId;
 use katana_primitives::receipt::MessageToL1;
 use serde::Deserialize;
 use tracing::{error, info};
@@ -145,7 +145,7 @@ pub trait Messenger {
         &self,
         from_block: u64,
         max_blocks: u64,
-        chain_id: FieldElement,
+        chain_id: ChainId,
     ) -> MessengerResult<(u64, Vec<Self::MessageTransaction>)>;
 
     /// Computes the hash of the given messages and sends them to the settlement chain.

diff --git a/crates/katana/core/src/service/messaging/service.rs b/crates/katana/core/src/service/messaging/service.rs
@@ -3,7 +3,6 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Duration;
 
-use ::starknet::core::types::FieldElement;
 use futures::{Future, FutureExt, Stream};
 use katana_primitives::block::BlockHashOrNumber;
 use katana_primitives::receipt::MessageToL1;
@@ -76,17 +75,14 @@ impl MessagingService {
         backend: Arc<Backend>,
         from_block: u64,
     ) -> MessengerResult<(u64, usize)> {
-        let chain_id = FieldElement::from_hex_be(&backend.env.read().block.chain_id.as_hex())
-            .expect("failed to parse katana chain id");
-
         // 200 avoids any possible rejection from RPC with possibly lot's of messages.
         // TODO: May this be configurable?
         let max_block = 200;
 
         match messenger.as_ref() {
             MessengerMode::Ethereum(inner) => {
                 let (block_num, txs) =
-                    inner.gather_messages(from_block, max_block, chain_id).await?;
+                    inner.gather_messages(from_block, max_block, backend.chain_id).await?;
                 let txs_count = txs.len();
 
                 txs.into_iter().for_each(|tx| {
@@ -101,7 +97,7 @@ impl MessagingService {
             #[cfg(feature = "starknet-messaging")]
             MessengerMode::Starknet(inner) => {
                 let (block_num, txs) =
-                    inner.gather_messages(from_block, max_block, chain_id).await?;
+                    inner.gather_messages(from_block, max_block, backend.chain_id).await?;
                 let txs_count = txs.len();
 
                 txs.into_iter().for_each(|tx| {

diff --git a/crates/katana/core/src/service/messaging/starknet.rs b/crates/katana/core/src/service/messaging/starknet.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use async_trait::async_trait;
+use katana_primitives::chain::ChainId;
 use katana_primitives::receipt::MessageToL1;
 use katana_primitives::transaction::L1HandlerTx;
 use katana_primitives::utils::transaction::compute_l1_message_hash;
@@ -163,7 +164,7 @@ impl Messenger for StarknetMessaging {
         &self,
         from_block: u64,
         max_blocks: u64,
-        chain_id: FieldElement,
+        chain_id: ChainId,
     ) -> MessengerResult<(u64, Vec<Self::MessageTransaction>)> {
         let chain_latest_block: u64 = match self.provider.block_number().await {
             Ok(n) => n,
@@ -306,7 +307,7 @@ fn parse_messages(messages: &[MessageToL1]) -> MessengerResult<(Vec<FieldElement
     Ok((hashes, calls))
 }
 
-fn l1_handler_tx_from_event(event: &EmittedEvent, chain_id: FieldElement) -> Result<L1HandlerTx> {
+fn l1_handler_tx_from_event(event: &EmittedEvent, chain_id: ChainId) -> Result<L1HandlerTx> {
     if event.keys[0] != selector!("MessageSentToAppchain") {
         debug!(
             target: LOG_TARGET,
@@ -429,7 +430,7 @@ mod tests {
         let from_address = selector!("from_address");
         let to_address = selector!("to_address");
         let selector = selector!("selector");
-        let chain_id = selector!("KATANA");
+        let chain_id = ChainId::parse("KATANA").unwrap();
         let nonce = FieldElement::ONE;
         let calldata = vec![from_address, FieldElement::THREE];
 
@@ -438,7 +439,7 @@ mod tests {
             to_address,
             selector,
             &calldata,
-            chain_id,
+            chain_id.into(),
             nonce,
         );
 
@@ -512,7 +513,7 @@ mod tests {
             transaction_hash,
         };
 
-        let _tx = l1_handler_tx_from_event(&event, FieldElement::ZERO).unwrap();
+        let _tx = l1_handler_tx_from_event(&event, ChainId::default()).unwrap();
     }
 
     #[test]
@@ -536,6 +537,6 @@ mod tests {
             transaction_hash,
         };
 
-        let _tx = l1_handler_tx_from_event(&event, FieldElement::ZERO).unwrap();
+        let _tx = l1_handler_tx_from_event(&event, ChainId::default()).unwrap();
     }
 }
diff --git a/crates/katana/primitives/Cargo.toml b/crates/katana/primitives/Cargo.toml
@@ -14,6 +14,8 @@ serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
 starknet.workspace = true
+strum.workspace = true
+strum_macros.workspace = true
 thiserror.workspace = true
 
 blockifier.workspace = true