anvil/
pubsub.rs

1use crate::{
2    StorageInfo,
3    eth::{
4        backend::notifications::{ChainNotification, ChainNotifications},
5        error::to_rpc_result,
6    },
7};
8use alloy_consensus::{BlockHeader, TxReceipt};
9use alloy_network::{AnyRpcTransaction, Network};
10use alloy_primitives::{B256, TxHash};
11use alloy_rpc_types::{FilteredParams, Log, Transaction, pubsub::SubscriptionResult};
12use anvil_core::eth::{block::Block, subscription::SubscriptionId};
13use anvil_rpc::{request::Version, response::ResponseResult};
14use foundry_primitives::FoundryTxReceipt;
15use futures::{Stream, StreamExt, channel::mpsc::Receiver, ready};
16use serde::Serialize;
17use std::{
18    collections::VecDeque,
19    pin::Pin,
20    task::{Context, Poll},
21};
22use tokio::sync::mpsc::UnboundedReceiver;
23
24/// Listens for new blocks and matching logs emitted in that block
25pub struct LogsSubscription<N: Network> {
26    pub blocks: ChainNotifications,
27    pub storage: StorageInfo<N>,
28    pub filter: FilteredParams,
29    pub queued: VecDeque<Log>,
30    pub id: SubscriptionId,
31}
32
33impl<N: Network> std::fmt::Debug for LogsSubscription<N> {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        f.debug_struct("LogsSubscription")
36            .field("filter", &self.filter)
37            .field("id", &self.id)
38            .finish_non_exhaustive()
39    }
40}
41
42impl<N: Network> LogsSubscription<N>
43where
44    N::ReceiptEnvelope: TxReceipt<Log = alloy_primitives::Log>,
45{
46    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Option<EthSubscriptionResponse>> {
47        loop {
48            if let Some(log) = self.queued.pop_front() {
49                let params = EthSubscriptionParams {
50                    subscription: self.id.clone(),
51                    result: to_rpc_result(log),
52                };
53                return Poll::Ready(Some(EthSubscriptionResponse::new(params)));
54            }
55
56            if let Some(notification) = ready!(self.blocks.poll_next_unpin(cx)) {
57                let logs = match notification {
58                    ChainNotification::Block(block) => {
59                        let b = self.storage.block(block.hash);
60                        let receipts = self.storage.receipts(block.hash);
61                        if let (Some(receipts), Some(block)) = (receipts, b) {
62                            filter_logs(block, receipts, &self.filter)
63                        } else {
64                            continue;
65                        }
66                    }
67                    ChainNotification::RemovedLogs(logs) => {
68                        filter_removed_logs(&logs, &self.filter)
69                    }
70                };
71                if logs.is_empty() {
72                    // this ensures we poll the receiver until it is pending, in which case the
73                    // underlying `UnboundedReceiver` will register the new waker, see
74                    // [`futures::channel::mpsc::UnboundedReceiver::poll_next()`]
75                    continue;
76                }
77                self.queued.extend(logs)
78            } else {
79                return Poll::Ready(None);
80            }
81
82            if self.queued.is_empty() {
83                return Poll::Pending;
84            }
85        }
86    }
87}
88
89#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
90pub struct EthSubscriptionResponse {
91    jsonrpc: Version,
92    method: &'static str,
93    params: EthSubscriptionParams,
94}
95
96impl EthSubscriptionResponse {
97    pub const fn new(params: EthSubscriptionParams) -> Self {
98        Self { jsonrpc: Version::V2, method: "eth_subscription", params }
99    }
100}
101
102/// Represents the `params` field of an `eth_subscription` event
103#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
104pub struct EthSubscriptionParams {
105    subscription: SubscriptionId,
106    #[serde(flatten)]
107    result: ResponseResult,
108}
109
110/// Represents an ethereum Websocket subscription
111pub enum EthSubscription<N: Network> {
112    Logs(Box<LogsSubscription<N>>),
113    Header(ChainNotifications, StorageInfo<N>, SubscriptionId),
114    PendingTransactions(Receiver<TxHash>, SubscriptionId),
115    FullPendingTransactions(UnboundedReceiver<AnyRpcTransaction>, SubscriptionId),
116    Syncing(Option<SubscriptionId>),
117    TransactionReceipts(UnboundedReceiver<Vec<FoundryTxReceipt>>, SubscriptionId),
118}
119
120impl<N: Network> std::fmt::Debug for EthSubscription<N> {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        match self {
123            Self::Logs(_) => f.debug_tuple("Logs").finish(),
124            Self::Header(..) => f.debug_tuple("Header").finish(),
125            Self::PendingTransactions(..) => f.debug_tuple("PendingTransactions").finish(),
126            Self::FullPendingTransactions(..) => f.debug_tuple("FullPendingTransactions").finish(),
127            Self::Syncing(_) => f.debug_tuple("Syncing").finish(),
128            Self::TransactionReceipts(..) => f.debug_tuple("TransactionReceipts").finish(),
129        }
130    }
131}
132
133impl<N: Network> EthSubscription<N>
134where
135    N::ReceiptEnvelope: TxReceipt<Log = alloy_primitives::Log>,
136{
137    fn poll_response(&mut self, cx: &mut Context<'_>) -> Poll<Option<EthSubscriptionResponse>> {
138        match self {
139            Self::Logs(listener) => listener.poll(cx),
140            Self::Header(blocks, storage, id) => {
141                // this loop ensures we poll the receiver until it is pending, in which case the
142                // underlying `UnboundedReceiver` will register the new waker, see
143                // [`futures::channel::mpsc::UnboundedReceiver::poll_next()`]
144                loop {
145                    if let Some(notification) = ready!(blocks.poll_next_unpin(cx)) {
146                        if let Some(block) = notification.as_new_block()
147                            && let Some(block) = storage.eth_block(block.hash)
148                        {
149                            let params = EthSubscriptionParams {
150                                subscription: id.clone(),
151                                result: to_rpc_result(block),
152                            };
153                            return Poll::Ready(Some(EthSubscriptionResponse::new(params)));
154                        }
155                    } else {
156                        return Poll::Ready(None);
157                    }
158                }
159            }
160            Self::PendingTransactions(tx, id) => {
161                let res = ready!(tx.poll_next_unpin(cx))
162                    .map(SubscriptionResult::<Transaction>::TransactionHash)
163                    .map(to_rpc_result)
164                    .map(|result| {
165                        let params = EthSubscriptionParams { subscription: id.clone(), result };
166                        EthSubscriptionResponse::new(params)
167                    });
168                Poll::Ready(res)
169            }
170            Self::FullPendingTransactions(tx, id) => {
171                let res = ready!(tx.poll_recv(cx)).map(to_rpc_result).map(|result| {
172                    let params = EthSubscriptionParams { subscription: id.clone(), result };
173                    EthSubscriptionResponse::new(params)
174                });
175                Poll::Ready(res)
176            }
177            Self::Syncing(id) => {
178                if let Some(id) = id.take() {
179                    let params =
180                        EthSubscriptionParams { subscription: id, result: to_rpc_result(false) };
181                    return Poll::Ready(Some(EthSubscriptionResponse::new(params)));
182                }
183                Poll::Pending
184            }
185            Self::TransactionReceipts(receipts, id) => {
186                let res = ready!(receipts.poll_recv(cx)).map(to_rpc_result).map(|result| {
187                    let params = EthSubscriptionParams { subscription: id.clone(), result };
188                    EthSubscriptionResponse::new(params)
189                });
190                Poll::Ready(res)
191            }
192        }
193    }
194}
195
196impl<N: Network> Stream for EthSubscription<N>
197where
198    N::ReceiptEnvelope: TxReceipt<Log = alloy_primitives::Log>,
199{
200    type Item = serde_json::Value;
201
202    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
203        let pin = self.get_mut();
204        match ready!(pin.poll_response(cx)) {
205            None => Poll::Ready(None),
206            Some(res) => Poll::Ready(Some(serde_json::to_value(res).expect("can't fail;"))),
207        }
208    }
209}
210
211/// Returns all the logs that match the given filter
212pub fn filter_logs<R>(block: Block, receipts: Vec<R>, filter: &FilteredParams) -> Vec<Log>
213where
214    R: TxReceipt<Log = alloy_primitives::Log>,
215{
216    /// Determines whether to add this log
217    fn add_log(
218        block_hash: B256,
219        l: &alloy_primitives::Log,
220        block: &Block,
221        params: &FilteredParams,
222    ) -> bool {
223        if params.filter.is_some() {
224            let block_number = block.header.number();
225            if !params.filter_block_range(block_number)
226                || !params.filter_block_hash(block_hash)
227                || !params.filter_address(&l.address)
228                || !params.filter_topics(l.topics())
229            {
230                return false;
231            }
232        }
233        true
234    }
235
236    let block_hash = block.header.hash_slow();
237    let mut logs = vec![];
238    let mut log_index: u32 = 0;
239    for (receipt_index, receipt) in receipts.into_iter().enumerate() {
240        let transaction_hash = block.body.transactions[receipt_index].hash();
241        for log in receipt.logs() {
242            if add_log(block_hash, log, &block, filter) {
243                logs.push(Log {
244                    inner: log.clone(),
245                    block_hash: Some(block_hash),
246                    block_number: Some(block.header.number()),
247                    transaction_hash: Some(transaction_hash),
248                    transaction_index: Some(receipt_index as u64),
249                    log_index: Some(log_index as u64),
250                    removed: false,
251                    block_timestamp: Some(block.header.timestamp()),
252                });
253            }
254            log_index += 1;
255        }
256    }
257    logs
258}
259
260/// Returns all logs of reorged out blocks that match the given filter.
261///
262/// The logs are expected to already be marked as `removed` and carry the metadata of the block
263/// they were originally included in, so the filter is applied against that metadata.
264pub fn filter_removed_logs(logs: &[Log], filter: &FilteredParams) -> Vec<Log> {
265    logs.iter()
266        .filter(|log| {
267            filter.filter.is_none()
268                || (filter.filter_block_range(log.block_number.unwrap_or_default())
269                    && filter.filter_block_hash(log.block_hash.unwrap_or_default())
270                    && filter.filter_address(&log.inner.address)
271                    && filter.filter_topics(log.inner.topics()))
272        })
273        .cloned()
274        .collect()
275}
anvil/pubsub.rs

anvil/
pubsub.rs
