feat(cast): subscribe to logs using websockets (foundry-rs#5743)

* feat(cast): subscribe to logs

* undo generic signal

* fix tokio signal feature

---------

Co-authored-by: Enrique Ortiz <hi@enriqueortiz.dev>

Author: 2 people

crates/cast/Cargo.toml

@@ -60,7 +60,7 @@ regex = { version = "1", default-features = false }
 rpassword = "7"
 semver = "1"
 tempfile = "3"
-tokio = { version = "1", features = ["macros"] }
+tokio = { version = "1", features = ["macros", "signal"] }
 tracing = "0.1"
 yansi = "0.5"
 

crates/cast/bin/cmd/logs.rs

@@ -1,16 +1,17 @@
+use std::{io, str::FromStr};
+
 use cast::Cast;
 use clap::Parser;
-use ethers::{
+use ethers::{providers::Middleware, types::NameOrAddress};
+use ethers_core::{
     abi::{Address, Event, RawTopicFilter, Topic, TopicFilter},
-    providers::Middleware,
-    types::{BlockId, BlockNumber, Filter, FilterBlockOption, NameOrAddress, ValueOrArray, H256},
+    types::{BlockId, BlockNumber, Filter, FilterBlockOption, ValueOrArray, H256},
 };
 use eyre::Result;
 use foundry_cli::{opts::EthereumOpts, utils};
 use foundry_common::abi::{get_event, parse_tokens};
 use foundry_config::Config;
 use itertools::Itertools;
-use std::str::FromStr;
 
 /// CLI arguments for `cast logs`.
 #[derive(Debug, Parser)]
@@ -44,7 +45,12 @@ pub struct LogsArgs {
     #[clap(value_name = "TOPICS_OR_ARGS")]
     topics_or_args: Vec<String>,
 
-    /// Print the logs as JSON.
+    /// If the RPC type and endpoints supports `eth_subscribe` stream logs instead of printing and
+    /// exiting. Will continue until interrupted or TO_BLOCK is reached.
+    #[clap(long)]
+    subscribe: bool,
+
+    /// Print the logs as JSON.s
     #[clap(long, short, help_heading = "Display options")]
     json: bool,
 
@@ -55,12 +61,21 @@ pub struct LogsArgs {
 impl LogsArgs {
     pub async fn run(self) -> Result<()> {
         let LogsArgs {
-            from_block, to_block, address, topics_or_args, sig_or_topic, json, eth, ..
+            from_block,
+            to_block,
+            address,
+            sig_or_topic,
+            topics_or_args,
+            subscribe,
+            json,
+            eth,
         } = self;
 
         let config = Config::from(&eth);
         let provider = utils::get_provider(&config)?;
 
+        let cast = Cast::new(&provider);
+
         let address = match address {
             Some(address) => {
                 let address = match address {
@@ -72,48 +87,29 @@ impl LogsArgs {
             None => None,
         };
 
-        let from_block = convert_block_number(&provider, from_block).await?;
-        let to_block = convert_block_number(&provider, to_block).await?;
-
-        let cast = Cast::new(&provider);
+        let from_block = cast.convert_block_number(from_block).await?;
+        let to_block = cast.convert_block_number(to_block).await?;
 
         let filter = build_filter(from_block, to_block, address, sig_or_topic, topics_or_args)?;
 
-        let logs = cast.filter_logs(filter, json).await?;
+        if !subscribe {
+            let logs = cast.filter_logs(filter, json).await?;
 
-        println!("{}", logs);
+            println!("{}", logs);
 
-        Ok(())
-    }
-}
+            return Ok(())
+        }
 
-/// Converts a block identifier into a block number.
-///
-/// If the block identifier is a block number, then this function returns the block number. If the
-/// block identifier is a block hash, then this function returns the block number of that block
-/// hash. If the block identifier is `None`, then this function returns `None`.
-async fn convert_block_number<M: Middleware>(
-    provider: M,
-    block: Option<BlockId>,
-) -> Result<Option<BlockNumber>, eyre::Error>
-where
-    M::Error: 'static,
-{
-    match block {
-        Some(block) => match block {
-            BlockId::Number(block_number) => Ok(Some(block_number)),
-            BlockId::Hash(hash) => {
-                let block = provider.get_block(hash).await?;
-                Ok(block.map(|block| block.number.unwrap()).map(BlockNumber::from))
-            }
-        },
-        None => Ok(None),
+        let mut stdout = io::stdout();
+        cast.subscribe(filter, &mut stdout, json).await?;
+
+        Ok(())
     }
 }
 
-// First tries to parse the `sig_or_topic` as an event signature. If successful, `topics_or_args` is
-// parsed as indexed inputs and converted to topics. Otherwise, `sig_or_topic` is prepended to
-// `topics_or_args` and used as raw topics.
+/// Builds a Filter by first trying to parse the `sig_or_topic` as an event signature. If
+/// successful, `topics_or_args` is parsed as indexed inputs and converted to topics. Otherwise,
+/// `sig_or_topic` is prepended to `topics_or_args` and used as raw topics.
 fn build_filter(
     from_block: Option<BlockNumber>,
     to_block: Option<BlockNumber>,
@@ -154,7 +150,7 @@ fn build_filter(
     Ok(filter)
 }
 
-// Creates a TopicFilter for the given event signature and arguments.
+/// Creates a TopicFilter from the given event signature and arguments.
 fn build_filter_event_sig(event: Event, args: Vec<String>) -> Result<TopicFilter, eyre::Error> {
     let args = args.iter().map(|arg| arg.as_str()).collect::<Vec<_>>();
 
@@ -195,7 +191,7 @@ fn build_filter_event_sig(event: Event, args: Vec<String>) -> Result<TopicFilter
     Ok(event.filter(raw)?)
 }
 
-// Creates a TopicFilter from raw topic hashes.
+/// Creates a TopicFilter from raw topic hashes.
 fn build_filter_topics(topics: Vec<String>) -> Result<TopicFilter, eyre::Error> {
     let mut topics = topics
         .into_iter()
@@ -214,8 +210,11 @@ fn build_filter_topics(topics: Vec<String>) -> Result<TopicFilter, eyre::Error>
 
 #[cfg(test)]
 mod tests {
+    use std::str::FromStr;
+
     use super::*;
     use ethers::types::H160;
+    use ethers_core::types::H256;
 
     const ADDRESS: &str = "0x4D1A2e2bB4F88F0250f26Ffff098B0b30B26BF38";
     const TRANSFER_SIG: &str = "Transfer(address indexed,address indexed,uint256)";

crates/cast/src/lib.rs

@@ -13,22 +13,25 @@ use ethers_core::{
     },
 };
 use ethers_etherscan::{errors::EtherscanError, Client};
-use ethers_providers::{Middleware, PendingTransaction};
+use ethers_providers::{Middleware, PendingTransaction, PubsubClient};
 use evm_disassembler::{disassemble_bytes, disassemble_str, format_operations};
 use eyre::{Context, Result};
 use foundry_common::{abi::encode_args, fmt::*, TransactionReceiptWithRevertReason};
 pub use foundry_evm::*;
+use futures::{future::Either, FutureExt, StreamExt};
 use rayon::prelude::*;
 pub use rusoto_core::{
     credential::ChainProvider as AwsChainProvider, region::Region as AwsRegion,
     request::HttpClient as AwsHttpClient, Client as AwsClient,
 };
 pub use rusoto_kms::KmsClient;
 use std::{
+    io,
     path::PathBuf,
     str::FromStr,
     sync::atomic::{AtomicBool, Ordering},
 };
+use tokio::signal::ctrl_c;
 pub use tx::TxBuilder;
 use tx::{TxBuilderOutput, TxBuilderPeekOutput};
 
@@ -816,6 +819,147 @@ where
         };
         Ok(res)
     }
+
+    /// Converts a block identifier into a block number.
+    ///
+    /// If the block identifier is a block number, then this function returns the block number. If
+    /// the block identifier is a block hash, then this function returns the block number of
+    /// that block hash. If the block identifier is `None`, then this function returns `None`.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use cast::Cast;
+    /// use ethers_providers::{Provider, Http};
+    /// use ethers_core::types::{BlockId, BlockNumber};
+    /// use std::convert::TryFrom;
+    ///
+    /// # async fn foo() -> eyre::Result<()> {
+    /// let provider = Provider::<Http>::try_from("http://localhost:8545")?;
+    /// let cast = Cast::new(provider);
+    ///
+    /// let block_number = cast.convert_block_number(Some(BlockId::Number(BlockNumber::from(5)))).await?;
+    /// assert_eq!(block_number, Some(BlockNumber::from(5)));
+    ///
+    /// let block_number = cast.convert_block_number(Some(BlockId::Hash("0x1234".parse().unwrap()))).await?;
+    /// assert_eq!(block_number, Some(BlockNumber::from(1234)));
+    ///
+    /// let block_number = cast.convert_block_number(None).await?;
+    /// assert_eq!(block_number, None);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn convert_block_number(
+        &self,
+        block: Option<BlockId>,
+    ) -> Result<Option<BlockNumber>, eyre::Error> {
+        match block {
+            Some(block) => match block {
+                BlockId::Number(block_number) => Ok(Some(block_number)),
+                BlockId::Hash(hash) => {
+                    let block = self.provider.get_block(hash).await?;
+                    Ok(block.map(|block| block.number.unwrap()).map(BlockNumber::from))
+                }
+            },
+            None => Ok(None),
+        }
+    }
+
+    /// Sets up a subscription to the given filter and writes the logs to the given output.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use cast::Cast;
+    /// use ethers_core::abi::Address;
+    /// use ethers_providers::{Provider, Ws};
+    /// use ethers_core::types::Filter;
+    /// use std::{str::FromStr, convert::TryFrom};
+    /// use std::io;
+    ///
+    /// # async fn foo() -> eyre::Result<()> {
+    /// let provider = Provider::new(Ws::connect("wss://localhost:8545").await?);
+    /// let cast = Cast::new(provider);
+    ///
+    /// let filter = Filter::new().address(Address::from_str("0x00000000006c3852cbEf3e08E8dF289169EdE581")?);
+    /// let mut output = io::stdout();
+    /// cast.subscribe(filter, &mut output, false).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn subscribe(
+        &self,
+        filter: Filter,
+        output: &mut dyn io::Write,
+        to_json: bool,
+    ) -> Result<()>
+    where
+        <M as Middleware>::Provider: PubsubClient,
+    {
+        // Initialize the subscription stream for logs
+        let mut subscription = self.provider.subscribe_logs(&filter).await?;
+
+        // Check if a to_block is specified, if so, subscribe to blocks
+        let mut block_subscription = if filter.get_to_block().is_some() {
+            Some(self.provider.subscribe_blocks().await?)
+        } else {
+            None
+        };
+
+        let to_block_number = filter.get_to_block();
+
+        // If output should be JSON, start with an opening bracket
+        if to_json {
+            write!(output, "[")?;
+        }
+
+        let mut first = true;
+
+        loop {
+            tokio::select! {
+                // If block subscription is present, listen to it to avoid blocking indefinitely past the desired to_block
+                block = if let Some(bs) = &mut block_subscription {
+                    Either::Left(bs.next().fuse())
+                } else {
+                    Either::Right(futures::future::pending())
+                } => {
+                    if let (Some(block), Some(to_block)) = (block, to_block_number) {
+                        if block.number.map_or(false, |bn| bn > to_block) {
+                            break;
+                        }
+                    }
+                },
+                // Process incoming log
+                log = subscription.next() => {
+                    if to_json {
+                        if !first {
+                            write!(output, ",")?;
+                        }
+                        first = false;
+                        let log_str = serde_json::to_string(&log).unwrap();
+                        write!(output, "{}", log_str)?;
+                    } else {
+                        let log_str = log.pretty()
+                            .replacen('\n', "- ", 1)  // Remove empty first line
+                            .replace('\n', "\n  ");  // Indent
+                        writeln!(output, "{}", log_str)?;
+                    }
+                },
+                // Break on cancel signal, to allow for closing JSON bracket
+                _ = ctrl_c() => {
+                    break;
+                },
+                else => break,
+            }
+        }
+
+        // If output was JSON, end with a closing bracket
+        if to_json {
+            write!(output, "]")?;
+        }
+
+        Ok(())
+    }
 }
 
 pub struct InterfaceSource {