refactor(katana): move the rpc address log

bin/katana/src/cli/node.rs

@@ -234,8 +234,7 @@ impl NodeArgs {
         if !self.silent {
             #[allow(deprecated)]
             let genesis = &node.backend.chain_spec.genesis;
-            let server_address = node.rpc_config.socket_addr();
-            print_intro(&self, genesis, &server_address);
+            print_intro(&self, genesis);
         }
 
         // Launch the node
@@ -369,7 +368,7 @@ impl NodeArgs {
     }
 }
 
-fn print_intro(args: &NodeArgs, genesis: &Genesis, address: &SocketAddr) {
+fn print_intro(args: &NodeArgs, genesis: &Genesis) {
     let mut accounts = genesis.accounts().peekable();
     let account_class_hash = accounts.peek().map(|e| e.1.class_hash());
     let seed = &args.starknet.seed;
@@ -381,7 +380,6 @@ fn print_intro(args: &NodeArgs, genesis: &Genesis, address: &SocketAddr) {
             serde_json::json!({
                 "accounts": accounts.map(|a| serde_json::json!(a)).collect::<Vec<_>>(),
                 "seed": format!("{}", seed),
-                "address": format!("{address}"),
             })
         )
     } else {
@@ -412,13 +410,6 @@ ACCOUNTS SEED
 {seed}
     "
         );
-
-        let addr = format!(
-            "🚀 JSON-RPC server started: {}",
-            Style::new().red().apply_to(format!("http://{address}"))
-        );
-
-        println!("\n{addr}\n\n",);
     }
 }
 

crates/katana/node-bindings/Cargo.toml

@@ -9,6 +9,7 @@ version.workspace = true
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+regex.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 starknet.workspace = true

crates/katana/node-bindings/src/json.rs

@@ -1,54 +1,71 @@
+#![allow(dead_code)]
+
 //! Utilities for parsing the logs in JSON format. This is when katana is run with `--json-log`.
 //!
 //! When JSON log is enabled, the startup details are all printed in a single log message.
-//! Example startup log in JSON format:
-//!
-//! ```json
-//! {"timestamp":"2024-07-06T03:35:00.410846Z","level":"INFO","fields":{"message":"{\"accounts\":[[\
-//! "318027405971194400117186968443431282813445578359155272415686954645506762954\",{\"balance\":\"
-//! 0x21e19e0c9bab2400000\",\"class_hash\":\"
-//! 0x5400e90f7e0ae78bd02c77cd75527280470e2fe19c54970dd79dc37a9d3645c\",\"private_key\":\"
-//! 0x2bbf4f9fd0bbb2e60b0316c1fe0b76cf7a4d0198bd493ced9b8df2a3a24d68a\",\"public_key\":\"
-//! 0x640466ebd2ce505209d3e5c4494b4276ed8f1cde764d757eb48831961f7cdea\"}]],\"address\":\"0.0.0.0:
-//! 5050\",\"seed\":\"0\"}"},"target":"katana::cli"}
-//! ```
-#![allow(dead_code)]
+
+use std::net::SocketAddr;
 
 use serde::Deserialize;
 
-#[derive(Deserialize)]
-pub struct JsonLogMessage {
+#[derive(Deserialize, Debug)]
+pub struct JsonLog<T = serde_json::Value> {
     pub timestamp: String,
     pub level: String,
-    pub fields: JsonLogFields,
+    pub fields: Fields<T>,
     pub target: String,
 }
 
-#[derive(Deserialize)]
-pub struct JsonLogFields {
-    #[serde(deserialize_with = "deserialize_katana_info")]
-    pub message: KatanaInfo,
-}
-
-fn deserialize_katana_info<'de, D>(deserializer: D) -> Result<KatanaInfo, D::Error>
-where
-    D: serde::Deserializer<'de>,
-{
-    let s = String::deserialize(deserializer)?;
-    serde_json::from_str(&s).map_err(serde::de::Error::custom)
+#[derive(Deserialize, Debug)]
+pub struct Fields<T = serde_json::Value> {
+    pub message: String,
+    #[serde(flatten)]
+    pub other: T,
 }
 
-#[derive(Deserialize)]
+/// Katana startup log message. The object is included as a string in the `message` field. Hence we
+/// have to parse it separately unlike the [`RpcAddr`] where we can directly deserialize using the
+/// Fields generic parameter.
+///
+/// Example:
+///
+/// ```json
+/// {
+///   "timestamp": "2024-10-10T14:55:04.452924Z",
+///   "level": "INFO",
+///   "fields": {
+///     "message": "{\"accounts\":[[\"0x6162896d1d7ab204c7ccac6dd5f8e9e7c25ecd5ae4fcb4ad32e57786bb46e03\",{\"balance\":\"0x21e19e0c9bab2400000\",\"class_hash\":\"0x5400e90f7e0ae78bd02c77cd75527280470e2fe19c54970dd79dc37a9d3645c\",\"private_key\":\"0x1800000000300000180000000000030000000000003006001800006600\",\"public_key\":\"0x2b191c2f3ecf685a91af7cf72a43e7b90e2e41220175de5c4f7498981b10053\"}]],\"seed\":\"0\"}"
+///   },
+///   "target": "katana::cli"
+/// }
+/// ```
+#[derive(Deserialize, Debug)]
 pub struct KatanaInfo {
     pub seed: String,
-    pub address: String,
     pub accounts: Vec<(String, AccountInfo)>,
 }
 
-#[derive(Deserialize)]
+impl TryFrom<String> for KatanaInfo {
+    type Error = serde_json::Error;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        serde_json::from_str(&value)
+    }
+}
+
+#[derive(Deserialize, Debug)]
 pub struct AccountInfo {
     pub balance: String,
     pub class_hash: String,
     pub private_key: String,
     pub public_key: String,
 }
+
+/// {
+///     "message": "RPC server started.",
+///     "addr": "127.0.0.1:5050"
+/// }
+#[derive(Deserialize, Debug)]
+pub struct RpcAddr {
+    pub addr: SocketAddr,
+}

crates/katana/node-bindings/src/lib.rs

@@ -6,21 +6,23 @@
 
 mod json;
 
+use std::borrow::Cow;
 use std::io::{BufRead, BufReader};
 use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::process::{Child, Command};
 use std::str::FromStr;
 use std::time::{Duration, Instant};
 
+use json::RpcAddr;
 use starknet::core::types::{Felt, FromStrError};
 use starknet::macros::short_string;
 use starknet::signers::SigningKey;
 use thiserror::Error;
 use tracing::trace;
 use url::Url;
 
-use crate::json::{JsonLogMessage, KatanaInfo};
+use crate::json::{JsonLog, KatanaInfo};
 
 /// How long we will wait for katana to indicate that it is ready.
 const KATANA_STARTUP_TIMEOUT_MILLIS: u64 = 10_000;
@@ -125,13 +127,22 @@ pub enum Error {
     MissingAccountPrivateKey,
 
     /// A line indicating the instance address was found but the actual value was not.
-    #[error("missing account private key")]
+    #[error("missing rpc server address")]
     MissingSocketAddr,
 
     #[error("encountered unexpected format: {0}")]
     UnexpectedFormat(String),
+
+    #[error("failed to match regex: {0}")]
+    Regex(#[from] regex::Error),
+
+    #[error("expected logs to be in JSON format: {0}")]
+    ExpectedJsonFormat(#[from] serde_json::Error),
 }
 
+/// The string indicator from which the RPC server address can be extracted from.
+const RPC_ADDR_LOG_SUBSTR: &str = "RPC server started.";
+
 /// Builder for launching `katana`.
 ///
 /// # Panics
@@ -411,7 +422,6 @@ impl Katana {
         if let Some(db_dir) = self.db_dir {
             cmd.arg("--db-dir").arg(db_dir);
         }
-
         if let Some(rpc_url) = self.rpc_url {
             cmd.arg("--rpc-url").arg(rpc_url);
         }
@@ -501,33 +511,36 @@ impl Katana {
             trace!(line);
 
             if self.json_log {
-                if let Ok(log) = serde_json::from_str::<JsonLogMessage>(&line) {
-                    let KatanaInfo { address, accounts: account_infos, .. } = log.fields.message;
-
-                    let addr = SocketAddr::from_str(&address)?;
-                    port = addr.port();
-
-                    for (address, info) in account_infos {
-                        let address = Felt::from_str(&address)?;
-                        let private_key = Felt::from_str(&info.private_key)?;
-                        let key = SigningKey::from_secret_scalar(private_key);
-                        accounts.push(Account { address, private_key: Some(key) });
-                    }
-
+                dbg!(&line);
+
+                // Because we using a concrete type for rpc addr log, we need to parse this first.
+                // Otherwise if we were to inverse the if statements, the else block
+                // would never be executed as all logs can be parsed as `JsonLog`.
+                if let Ok(log) = dbg!(serde_json::from_str::<JsonLog<RpcAddr>>(&line)) {
+                    debug_assert!(log.fields.message.contains(RPC_ADDR_LOG_SUBSTR));
+                    port = log.fields.other.addr.port();
+                    // We can safely break here as we don't need any information after the rpc
+                    // address
                     break;
                 }
+                // Parse all logs as generic logs
+                else if let Ok(info) = serde_json::from_str::<JsonLog>(&line) {
+                    // Check if this log is a katana startup info log
+                    if let Ok(info) = KatanaInfo::try_from(info.fields.message) {
+                        for (address, info) in info.accounts {
+                            let address = Felt::from_str(&address)?;
+                            let private_key = Felt::from_str(&info.private_key)?;
+                            let key = SigningKey::from_secret_scalar(private_key);
+                            accounts.push(Account { address, private_key: Some(key) });
+                        }
+
+                        continue;
+                    }
+                }
             } else {
-                const URL_PREFIX: &str = "🚀 JSON-RPC server started:";
-                if line.starts_with(URL_PREFIX) {
-                    // <🚀 JSON-RPC server started: http://0.0.0.0:5050>
-                    let line = line.strip_prefix(URL_PREFIX).ok_or(Error::MissingSocketAddr)?;
-                    let addr = line.trim();
-
-                    // parse the actual port
-                    let addr = addr.strip_prefix("http://").unwrap_or(addr);
-                    let addr = SocketAddr::from_str(addr)?;
+                if line.contains(RPC_ADDR_LOG_SUBSTR) {
+                    let addr = parse_rpc_addr_log(&line)?;
                     port = addr.port();
-
                     // The address is the last thing to be displayed so we can safely break here.
                     break;
                 }
@@ -577,21 +590,47 @@ impl Katana {
     }
 }
 
+/// Removes ANSI escape codes from a string.
+///
+/// This is useful for removing the color codes from the katana output.
+fn clean_ansi_escape_codes(input: &str) -> Result<Cow<'_, str>, Error> {
+    let re = regex::Regex::new(r"\x1b\[[0-9;]*[a-zA-Z]")?;
+    Ok(re.replace_all(input, ""))
+}
+
+// Example RPC address log format (ansi color codes removed):
+// 2024-10-10T14:20:53.563106Z  INFO rpc: RPC server started. addr=127.0.0.1:60373
+fn parse_rpc_addr_log(log: &str) -> Result<SocketAddr, Error> {
+    // remove any ANSI escape codes from the log.
+    let cleaned = clean_ansi_escape_codes(log)?;
+
+    // This will separate the log into two parts as separated by `addr=` str and we take
+    // only the second part which is the address.
+    let addr_part = cleaned.split("addr=").nth(1).ok_or(Error::MissingSocketAddr)?;
+    let addr = addr_part.trim();
+    Ok(SocketAddr::from_str(addr)?)
+}
+
 #[cfg(test)]
 mod tests {
     use starknet::providers::jsonrpc::HttpTransport;
     use starknet::providers::{JsonRpcClient, Provider};
 
     use super::*;
 
-    #[test]
-    fn can_launch_katana() {
+    #[tokio::test]
+    async fn can_launch_katana() {
+        // this will launch katana with random ports
         let katana = Katana::new().spawn();
         // assert some default values
         assert_eq!(katana.accounts().len(), 10);
         assert_eq!(katana.chain_id(), short_string!("KATANA"));
         // assert that all accounts have private key
         assert!(katana.accounts().iter().all(|a| a.private_key.is_some()));
+
+        let provider = JsonRpcClient::new(HttpTransport::new(katana.endpoint_url()));
+        let result = provider.chain_id().await;
+        assert!(result.is_ok());
     }
 
     #[test]
@@ -622,11 +661,15 @@ mod tests {
         let _ = Katana::new().block_time(500).spawn();
     }
 
-    #[test]
-    fn can_launch_katana_with_specific_port() {
+    #[tokio::test]
+    async fn can_launch_katana_with_specific_port() {
         let specific_port = 49999;
         let katana = Katana::new().port(specific_port).spawn();
         assert_eq!(katana.port(), specific_port);
+
+        let provider = JsonRpcClient::new(HttpTransport::new(katana.endpoint_url()));
+        let result = provider.chain_id().await;
+        assert!(result.is_ok());
     }
 
     #[tokio::test]
@@ -652,4 +695,15 @@ mod tests {
         assert!(db_path.exists());
         assert!(db_path.is_dir());
     }
+
+    #[test]
+    fn test_parse_rpc_addr_log() {
+        // actual rpc log from katana
+        let log = "\u{1b}[2m2024-10-10T14:48:55.397891Z\u{1b}[0m \u{1b}[32m INFO\u{1b}[0m \
+                   \u{1b}[2mrpc\u{1b}[0m\u{1b}[2m:\u{1b}[0m RPC server started. \
+                   \u{1b}[3maddr\u{1b}[0m\u{1b}[2m=\u{1b}[0m127.0.0.1:60817\n";
+        let addr = parse_rpc_addr_log(log).unwrap();
+        assert_eq!(addr.ip().to_string(), "127.0.0.1");
+        assert_eq!(addr.port(), 60817);
+    }
 }

crates/katana/node/src/lib.rs

@@ -355,6 +355,8 @@ pub async fn spawn<EF: ExecutorFactory>(
     let addr = server.local_addr()?;
     let handle = server.start(methods)?;
 
+    info!(target: "rpc", %addr, "RPC server started.");
+
     Ok(RpcServer { handle, addr })
 }