Skip to main content

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                        // Rule out nonmatching blocks before requiring data that may be pruned.
60                        if !self.filter.filter_block_range(block.header.number)
61                            || !self.filter.filter_block_hash(block.hash)
62                            || block.header.logs_bloom.is_zero()
63                            || self.filter.filter.as_ref().is_some_and(|filter| {
64                                !filter.matches_bloom(block.header.logs_bloom)
65                            })
66                        {
67                            continue;
68                        }
69                        let hash = block.hash;
70                        let (Some(block), Some(receipts)) =
71                            (self.storage.block(hash), self.storage.receipts(hash))
72                        else {
73                            tracing::warn!(target: "rpc::ws", ?hash, "Closing logs subscription because block data is unavailable");
74                            return Poll::Ready(None);
75                        };
76                        if block.body.transactions.len() != receipts.len()
77                            || (block.body.transactions.is_empty()
78                                && block.header.transactions_root()
79                                    != alloy_consensus::constants::EMPTY_ROOT_HASH)
80                        {
81                            tracing::warn!(target: "rpc::ws", ?hash, "Closing logs subscription because block data was pruned");
82                            return Poll::Ready(None);
83                        }
84                        filter_logs(block, receipts, &self.filter)
85                    }
86                    ChainNotification::RemovedLogs(logs) => {
87                        filter_removed_logs(&logs, &self.filter)
88                    }
89                };
90                if logs.is_empty() {
91                    // this ensures we poll the receiver until it is pending, in which case the
92                    // underlying `UnboundedReceiver` will register the new waker, see
93                    // [`futures::channel::mpsc::UnboundedReceiver::poll_next()`]
94                    continue;
95                }
96                self.queued.extend(logs)
97            } else {
98                return Poll::Ready(None);
99            }
100
101            if self.queued.is_empty() {
102                return Poll::Pending;
103            }
104        }
105    }
106}
107
108#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
109pub struct EthSubscriptionResponse {
110    jsonrpc: Version,
111    method: &'static str,
112    params: EthSubscriptionParams,
113}
114
115impl EthSubscriptionResponse {
116    pub const fn new(params: EthSubscriptionParams) -> Self {
117        Self { jsonrpc: Version::V2, method: "eth_subscription", params }
118    }
119}
120
121/// Represents the `params` field of an `eth_subscription` event
122#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
123pub struct EthSubscriptionParams {
124    subscription: SubscriptionId,
125    #[serde(flatten)]
126    result: ResponseResult,
127}
128
129/// Represents an ethereum Websocket subscription
130pub enum EthSubscription<N: Network> {
131    Logs(Box<LogsSubscription<N>>),
132    Header(ChainNotifications, StorageInfo<N>, SubscriptionId),
133    PendingTransactions(Receiver<TxHash>, SubscriptionId),
134    FullPendingTransactions(UnboundedReceiver<AnyRpcTransaction>, SubscriptionId),
135    Syncing(Option<SubscriptionId>),
136    TransactionReceipts(UnboundedReceiver<Vec<FoundryTxReceipt>>, SubscriptionId),
137}
138
139impl<N: Network> std::fmt::Debug for EthSubscription<N> {
140    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
141        match self {
142            Self::Logs(_) => f.debug_tuple("Logs").finish(),
143            Self::Header(..) => f.debug_tuple("Header").finish(),
144            Self::PendingTransactions(..) => f.debug_tuple("PendingTransactions").finish(),
145            Self::FullPendingTransactions(..) => f.debug_tuple("FullPendingTransactions").finish(),
146            Self::Syncing(_) => f.debug_tuple("Syncing").finish(),
147            Self::TransactionReceipts(..) => f.debug_tuple("TransactionReceipts").finish(),
148        }
149    }
150}
151
152impl<N: Network> EthSubscription<N>
153where
154    N::ReceiptEnvelope: TxReceipt<Log = alloy_primitives::Log>,
155{
156    fn poll_response(&mut self, cx: &mut Context<'_>) -> Poll<Option<EthSubscriptionResponse>> {
157        match self {
158            Self::Logs(listener) => listener.poll(cx),
159            Self::Header(blocks, storage, id) => {
160                // this loop ensures we poll the receiver until it is pending, in which case the
161                // underlying `UnboundedReceiver` will register the new waker, see
162                // [`futures::channel::mpsc::UnboundedReceiver::poll_next()`]
163                loop {
164                    if let Some(notification) = ready!(blocks.poll_next_unpin(cx)) {
165                        if let Some(block) = notification.as_new_block()
166                            && let Some(block) = storage.eth_block(block.hash)
167                        {
168                            let params = EthSubscriptionParams {
169                                subscription: id.clone(),
170                                result: to_rpc_result(block),
171                            };
172                            return Poll::Ready(Some(EthSubscriptionResponse::new(params)));
173                        }
174                    } else {
175                        return Poll::Ready(None);
176                    }
177                }
178            }
179            Self::PendingTransactions(tx, id) => {
180                let res = ready!(tx.poll_next_unpin(cx))
181                    .map(SubscriptionResult::<Transaction>::TransactionHash)
182                    .map(to_rpc_result)
183                    .map(|result| {
184                        let params = EthSubscriptionParams { subscription: id.clone(), result };
185                        EthSubscriptionResponse::new(params)
186                    });
187                Poll::Ready(res)
188            }
189            Self::FullPendingTransactions(tx, id) => {
190                let res = ready!(tx.poll_recv(cx)).map(to_rpc_result).map(|result| {
191                    let params = EthSubscriptionParams { subscription: id.clone(), result };
192                    EthSubscriptionResponse::new(params)
193                });
194                Poll::Ready(res)
195            }
196            Self::Syncing(id) => {
197                if let Some(id) = id.take() {
198                    let params =
199                        EthSubscriptionParams { subscription: id, result: to_rpc_result(false) };
200                    return Poll::Ready(Some(EthSubscriptionResponse::new(params)));
201                }
202                Poll::Pending
203            }
204            Self::TransactionReceipts(receipts, id) => {
205                let res = ready!(receipts.poll_recv(cx)).map(to_rpc_result).map(|result| {
206                    let params = EthSubscriptionParams { subscription: id.clone(), result };
207                    EthSubscriptionResponse::new(params)
208                });
209                Poll::Ready(res)
210            }
211        }
212    }
213}
214
215impl<N: Network> Stream for EthSubscription<N>
216where
217    N::ReceiptEnvelope: TxReceipt<Log = alloy_primitives::Log>,
218{
219    type Item = serde_json::Value;
220
221    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
222        let pin = self.get_mut();
223        match ready!(pin.poll_response(cx)) {
224            None => Poll::Ready(None),
225            Some(res) => Poll::Ready(Some(serde_json::to_value(res).expect("can't fail;"))),
226        }
227    }
228}
229
230/// Returns all the logs that match the given filter
231pub fn filter_logs<R>(block: Block, receipts: Vec<R>, filter: &FilteredParams) -> Vec<Log>
232where
233    R: TxReceipt<Log = alloy_primitives::Log>,
234{
235    /// Determines whether to add this log
236    fn add_log(
237        block_hash: B256,
238        l: &alloy_primitives::Log,
239        block: &Block,
240        params: &FilteredParams,
241    ) -> bool {
242        if params.filter.is_some() {
243            let block_number = block.header.number();
244            if !params.filter_block_range(block_number)
245                || !params.filter_block_hash(block_hash)
246                || !params.filter_address(&l.address)
247                || !params.filter_topics(l.topics())
248            {
249                return false;
250            }
251        }
252        true
253    }
254
255    let block_hash = block.header.hash_slow();
256    let mut logs = vec![];
257    let mut log_index: u32 = 0;
258    for (receipt_index, receipt) in receipts.into_iter().enumerate() {
259        let transaction_hash = block.body.transactions[receipt_index].hash();
260        for log in receipt.logs() {
261            if add_log(block_hash, log, &block, filter) {
262                logs.push(Log {
263                    inner: log.clone(),
264                    block_hash: Some(block_hash),
265                    block_number: Some(block.header.number()),
266                    transaction_hash: Some(transaction_hash),
267                    transaction_index: Some(receipt_index as u64),
268                    log_index: Some(log_index as u64),
269                    removed: false,
270                    block_timestamp: Some(block.header.timestamp()),
271                });
272            }
273            log_index += 1;
274        }
275    }
276    logs
277}
278
279/// Returns all logs of reorged out blocks that match the given filter.
280///
281/// The logs are expected to already be marked as `removed` and carry the metadata of the block
282/// they were originally included in, so the filter is applied against that metadata.
283pub fn filter_removed_logs(logs: &[Log], filter: &FilteredParams) -> Vec<Log> {
284    logs.iter()
285        .filter(|log| {
286            filter.filter.is_none()
287                || (filter.filter_block_range(log.block_number.unwrap_or_default())
288                    && filter.filter_block_hash(log.block_hash.unwrap_or_default())
289                    && filter.filter_address(&log.inner.address)
290                    && filter.filter_topics(log.inner.topics()))
291        })
292        .cloned()
293        .collect()
294}