1use std::{collections::HashMap, fmt::Debug, sync::Arc};
17
18use alloy::primitives::Address;
19use nautilus_core::{consts::NAUTILUS_USER_AGENT, string::secret::REDACTED};
20use nautilus_live::SocketControl;
21#[cfg(feature = "hypersync")]
22use nautilus_model::defi::DexType;
23use nautilus_model::defi::{
24 Block, Chain,
25 rpc::{RpcLog, RpcNodeWssResponse},
26};
27use nautilus_network::{
28 RECONNECTED,
29 http::USER_AGENT,
30 websocket::{TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler},
31};
32use tokio_tungstenite::tungstenite::Message;
33
34#[cfg(feature = "hypersync")]
35use crate::exchanges::get_dex_extended;
36use crate::rpc::{
37 error::BlockchainRpcClientError,
38 types::{BlockchainMessage, RpcEventType},
39 utils::{
40 extract_rpc_subscription_id, is_subscription_confirmation_response, is_subscription_event,
41 is_unsubscribe_confirmation_response,
42 },
43};
44
45pub struct CoreBlockchainRpcClient {
54 chain: Chain,
56 wss_rpc_url: String,
58 request_id: u64,
60 pending_subscription_request: HashMap<u64, RpcEventType>,
62 subscription_event_types: HashMap<String, RpcEventType>,
65 wss_client: Option<Arc<WebSocketClient>>,
67 wss_consumer_rx: Option<tokio::sync::mpsc::UnboundedReceiver<Message>>,
69 subscriptions: Arc<tokio::sync::RwLock<HashMap<RpcEventType, RpcSubscription>>>,
71 transport_backend: TransportBackend,
73 proxy_url: Option<String>,
75 socket_control: Option<SocketControl>,
76}
77
78impl Debug for CoreBlockchainRpcClient {
79 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80 f.debug_struct(stringify!(CoreBlockchainRpcClient))
81 .field("chain", &self.chain)
82 .field("wss_rpc_url", &REDACTED)
83 .field("request_id", &self.request_id)
84 .field(
85 "pending_subscription_request",
86 &self.pending_subscription_request,
87 )
88 .field("subscription_event_types", &self.subscription_event_types)
89 .field(
90 "wss_client",
91 &self.wss_client.as_ref().map(|_| "<WebSocketClient>"),
92 )
93 .field(
94 "wss_consumer_rx",
95 &self.wss_consumer_rx.as_ref().map(|_| "<Receiver>"),
96 )
97 .field("confirmed_subscriptions", &"<RwLock<HashMap>>")
98 .finish()
99 }
100}
101
102#[derive(Debug, Clone)]
103struct RpcSubscription {
104 name: String,
105 filter: Option<serde_json::Value>,
106}
107
108impl RpcSubscription {
109 fn new(name: &str) -> Self {
110 Self {
111 name: name.to_string(),
112 filter: None,
113 }
114 }
115
116 fn pool_logs(addresses: &[Address], event_signature: String) -> Self {
117 let mut addresses: Vec<String> = addresses
118 .iter()
119 .map(|address| format!("{address:?}"))
120 .collect();
121 addresses.sort();
122
123 Self {
124 name: "logs".to_string(),
125 filter: Some(serde_json::json!({
126 "address": addresses,
127 "topics": [event_signature],
128 })),
129 }
130 }
131
132 fn params(&self) -> Vec<serde_json::Value> {
133 let mut params = vec![serde_json::json!(self.name)];
134 if let Some(filter) = &self.filter {
135 params.push(filter.clone());
136 }
137 params
138 }
139}
140
141impl CoreBlockchainRpcClient {
142 #[must_use]
143 pub fn new(chain: Chain, wss_rpc_url: String, proxy_url: Option<String>) -> Self {
144 Self {
145 chain,
146 wss_rpc_url,
147 request_id: 1,
148 wss_client: None,
149 pending_subscription_request: HashMap::new(),
150 subscription_event_types: HashMap::new(),
151 wss_consumer_rx: None,
152 subscriptions: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
153 transport_backend: TransportBackend::default(),
154 proxy_url,
155 socket_control: None,
156 }
157 }
158
159 #[must_use]
161 pub fn with_transport_backend(mut self, backend: TransportBackend) -> Self {
162 self.transport_backend = backend;
163 self
164 }
165
166 pub fn set_transport_backend(&mut self, backend: TransportBackend) {
168 self.transport_backend = backend;
169 }
170
171 pub fn set_socket_control(&mut self, control: SocketControl) {
173 self.socket_control = Some(control);
174 }
175
176 pub async fn connect(&mut self) -> anyhow::Result<()> {
185 let (handler, rx) = channel_message_handler();
186
187 let heartbeat_interval = 30;
189
190 let config = WebSocketConfig {
191 url: self.wss_rpc_url.clone(),
192 headers: vec![(USER_AGENT.to_string(), NAUTILUS_USER_AGENT.to_string())],
193 heartbeat_interval_secs: Some(heartbeat_interval),
194 heartbeat_payload: None,
195 connect_timeout_ms: Some(10_000),
196 reconnect_delay_initial_ms: Some(1_000),
197 reconnect_delay_max_ms: Some(30_000),
198 reconnect_backoff_factor: Some(2.0),
199 reconnect_jitter_ms: Some(1_000),
200 reconnect_max_attempts: None,
201 heartbeat_timeout_secs: None,
202 idle_timeout_ms: None,
203 backend: self.transport_backend,
204 proxy_url: self.proxy_url.clone(),
205 };
206
207 let client = WebSocketClient::builder()
208 .config(config)
209 .message_handler(handler)
210 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
211 .connect()
212 .await?;
213 let reconnect_handle = client.reconnect_handle();
214 if let Some(control) = &self.socket_control {
215 control.register(move || reconnect_handle.request_reconnect());
216 }
217
218 self.wss_client = Some(Arc::new(client));
219 self.wss_consumer_rx = Some(rx);
220
221 Ok(())
222 }
223
224 async fn subscribe_events(
226 &mut self,
227 event_type: RpcEventType,
228 subscription: RpcSubscription,
229 ) -> Result<(), BlockchainRpcClientError> {
230 if self.subscriptions.read().await.contains_key(&event_type) {
231 return Ok(());
232 }
233
234 self.send_subscription_request(event_type, subscription)
235 .await
236 }
237
238 async fn replace_subscription(
239 &mut self,
240 event_type: RpcEventType,
241 subscription: RpcSubscription,
242 ) -> Result<(), BlockchainRpcClientError> {
243 self.unsubscribe_event_type(event_type).await?;
244 self.send_subscription_request(event_type, subscription)
245 .await
246 }
247
248 async fn send_subscription_request(
249 &mut self,
250 event_type: RpcEventType,
251 subscription: RpcSubscription,
252 ) -> Result<(), BlockchainRpcClientError> {
253 if let Some(client) = &self.wss_client {
254 log::debug!(
255 "Subscribing to '{}' on chain '{}'",
256 subscription.name,
257 self.chain.name
258 );
259 let msg = serde_json::json!({
260 "method": "eth_subscribe",
261 "id": self.request_id,
262 "jsonrpc": "2.0",
263 "params": subscription.params()
264 });
265 self.pending_subscription_request
266 .insert(self.request_id, event_type);
267 self.request_id += 1;
268
269 if let Err(e) = client.send_text(msg.to_string(), None).await {
270 log::error!("Error sending subscribe message: {e:?}");
271 }
272
273 let mut confirmed = self.subscriptions.write().await;
275 confirmed.insert(event_type, subscription);
276
277 Ok(())
278 } else {
279 Err(BlockchainRpcClientError::ClientError(String::from(
280 "Client not connected",
281 )))
282 }
283 }
284
285 async fn resubscribe_all(&mut self) -> Result<(), BlockchainRpcClientError> {
287 let subscriptions = self.subscriptions.read().await;
288
289 if subscriptions.is_empty() {
290 log::debug!(
291 "No subscriptions to re-establish for chain '{}'",
292 self.chain.name
293 );
294 return Ok(());
295 }
296
297 log::info!(
298 "Re-establishing {} subscription(s) for chain '{}'",
299 subscriptions.len(),
300 self.chain.name
301 );
302
303 let subs_to_restore: Vec<(RpcEventType, RpcSubscription)> = subscriptions
304 .iter()
305 .map(|(event_type, subscription)| (*event_type, subscription.clone()))
306 .collect();
307
308 drop(subscriptions);
309
310 for (event_type, subscription) in subs_to_restore {
311 self.send_subscription_request(event_type, subscription)
312 .await?;
313 }
314
315 Ok(())
316 }
317
318 async fn unsubscribe_events(
320 &self,
321 subscription_id: String,
322 ) -> Result<(), BlockchainRpcClientError> {
323 if let Some(client) = &self.wss_client {
324 log::debug!(
325 "Unsubscribing from '{}' on chain {}",
326 subscription_id,
327 self.chain.name
328 );
329 let msg = serde_json::json!({
330 "method": "eth_unsubscribe",
331 "id": 1,
332 "jsonrpc": "2.0",
333 "params": [subscription_id]
334 });
335
336 if let Err(e) = client.send_text(msg.to_string(), None).await {
337 log::error!("Error sending unsubscribe message: {e:?}");
338 }
339 Ok(())
340 } else {
341 Err(BlockchainRpcClientError::ClientError(String::from(
342 "Client not connected",
343 )))
344 }
345 }
346
347 async fn unsubscribe_event_type(
348 &mut self,
349 event_type: RpcEventType,
350 ) -> Result<(), BlockchainRpcClientError> {
351 let subscription_ids_to_remove: Vec<String> = self
352 .subscription_event_types
353 .iter()
354 .filter(|(_, active_event_type)| **active_event_type == event_type)
355 .map(|(id, _)| id.clone())
356 .collect();
357
358 for id in subscription_ids_to_remove {
359 self.unsubscribe_events(id.clone()).await?;
360 self.subscription_event_types.remove(&id);
361 }
362
363 self.pending_subscription_request
364 .retain(|_, pending_event_type| *pending_event_type != event_type);
365 self.subscriptions.write().await.remove(&event_type);
366
367 Ok(())
368 }
369
370 pub async fn wait_on_rpc_channel(&mut self) -> Option<Message> {
372 self.wss_consumer_rx.as_mut()?.recv().await
373 }
374
375 pub async fn next_rpc_message(
383 &mut self,
384 ) -> Result<BlockchainMessage, BlockchainRpcClientError> {
385 while let Some(msg) = self.wait_on_rpc_channel().await {
386 match msg {
387 Message::Text(text) => {
388 if text == RECONNECTED {
389 log::info!("Detected reconnection for chain '{}'", self.chain.name);
390
391 if let Err(e) = self.resubscribe_all().await {
392 log::error!("Failed to re-establish subscriptions: {e:?}");
393 }
394 continue;
395 }
396
397 match serde_json::from_str::<serde_json::Value>(&text) {
398 Ok(json) => {
399 if is_unsubscribe_confirmation_response(&json) {
400 log::debug!(
401 "Received unsubscribe confirmation on chain '{}'",
402 self.chain.name
403 );
404 continue;
405 } else if is_subscription_confirmation_response(&json) {
406 let subscription_request_id = json
407 .get("id")
408 .and_then(serde_json::Value::as_u64)
409 .ok_or_else(|| {
410 BlockchainRpcClientError::InternalRpcClientError(
411 "Missing subscription request id".to_string(),
412 )
413 })?;
414 let result = json
415 .get("result")
416 .and_then(serde_json::Value::as_str)
417 .ok_or_else(|| {
418 BlockchainRpcClientError::InternalRpcClientError(
419 "Missing subscription id".to_string(),
420 )
421 })?;
422 let Some(event_type) = self
423 .pending_subscription_request
424 .remove(&subscription_request_id)
425 else {
426 log::debug!(
427 "Unsubscribing from stale subscription confirmation '{}' on chain '{}'",
428 result,
429 self.chain.name
430 );
431 self.unsubscribe_events(result.to_string()).await?;
432 continue;
433 };
434
435 if self.subscriptions.read().await.contains_key(&event_type) {
436 self.subscription_event_types
437 .insert(result.to_string(), event_type);
438 } else {
439 self.unsubscribe_events(result.to_string()).await?;
440 }
441 continue;
442 } else if is_subscription_event(&json) {
443 let subscription_id = match extract_rpc_subscription_id(&json) {
444 Some(id) => id,
445 None => {
446 return Err(BlockchainRpcClientError::InternalRpcClientError(
447 "Error parsing subscription id from valid rpc response"
448 .to_string(),
449 ));
450 }
451 };
452
453 if let Some(event_type) =
454 self.subscription_event_types.get(subscription_id).copied()
455 {
456 match event_type {
457 RpcEventType::NewBlock => {
458 return match serde_json::from_value::<
459 RpcNodeWssResponse<Block>,
460 >(
461 json
462 ) {
463 Ok(block_response) => {
464 let block = block_response.params.result;
465 Ok(BlockchainMessage::Block(block))
466 }
467 Err(e) => Err(
468 BlockchainRpcClientError::MessageParsingError(
469 format!(
470 "Error parsing rpc response to block with error {e}"
471 ),
472 ),
473 ),
474 };
475 }
476 RpcEventType::PoolSwap(_)
477 | RpcEventType::PoolMint(_)
478 | RpcEventType::PoolBurn(_)
479 | RpcEventType::PoolCollect(_)
480 | RpcEventType::PoolFlash(_)
481 | RpcEventType::PoolFeeProtocolUpdate(_)
482 | RpcEventType::PoolFeeProtocolCollect(_) => {
483 let log = Self::parse_rpc_log_response(json)?;
484
485 if let Some(message) = self
486 .blockchain_message_from_pool_log(
487 event_type, &log,
488 )?
489 {
490 return Ok(message);
491 }
492 continue;
493 }
494 }
495 }
496 return Err(BlockchainRpcClientError::InternalRpcClientError(
497 format!(
498 "Event type not found for defined subscription id {subscription_id}"
499 ),
500 ));
501 }
502 return Err(BlockchainRpcClientError::UnsupportedRpcResponseType(
503 json.to_string(),
504 ));
505 }
506 Err(e) => {
507 return Err(BlockchainRpcClientError::MessageParsingError(
508 e.to_string(),
509 ));
510 }
511 }
512 }
513 Message::Pong(_) => {}
514 _ => {
515 return Err(BlockchainRpcClientError::UnsupportedRpcResponseType(
516 msg.to_string(),
517 ));
518 }
519 }
520 }
521
522 Err(BlockchainRpcClientError::NoMessageReceived)
523 }
524
525 fn parse_rpc_log_response(json: serde_json::Value) -> Result<RpcLog, BlockchainRpcClientError> {
526 serde_json::from_value::<RpcNodeWssResponse<RpcLog>>(json)
527 .map(|response| response.params.result)
528 .map_err(|e| {
529 BlockchainRpcClientError::MessageParsingError(format!(
530 "Error parsing rpc response to log with error {e}"
531 ))
532 })
533 }
534
535 #[cfg(feature = "hypersync")]
536 fn blockchain_message_from_pool_log(
537 &self,
538 event_type: RpcEventType,
539 log: &RpcLog,
540 ) -> Result<Option<BlockchainMessage>, BlockchainRpcClientError> {
541 if log.removed {
542 log::debug!(
543 "Skipping removed pool log on chain '{}' for event {:?}",
544 self.chain.name,
545 event_type
546 );
547 return Ok(None);
548 }
549
550 let dex = Self::pool_event_dex(event_type)?;
551 let dex_extended = get_dex_extended(self.chain.name, &dex).ok_or_else(|| {
552 BlockchainRpcClientError::InternalRpcClientError(format!(
553 "DEX {dex} is not registered for chain {}",
554 self.chain.name
555 ))
556 })?;
557
558 match event_type {
559 RpcEventType::PoolSwap(_) => dex_extended
560 .parse_swap_event_rpc(log)
561 .map(BlockchainMessage::SwapEvent),
562 RpcEventType::PoolMint(_) => dex_extended
563 .parse_mint_event_rpc(log)
564 .map(BlockchainMessage::MintEvent),
565 RpcEventType::PoolBurn(_) => dex_extended
566 .parse_burn_event_rpc(log)
567 .map(BlockchainMessage::BurnEvent),
568 RpcEventType::PoolCollect(_) => dex_extended
569 .parse_collect_event_rpc(log)
570 .map(BlockchainMessage::CollectEvent),
571 RpcEventType::PoolFlash(_) => dex_extended
572 .parse_flash_event_rpc(log)
573 .map(BlockchainMessage::FlashEvent),
574 RpcEventType::PoolFeeProtocolUpdate(_) => dex_extended
575 .parse_fee_protocol_update_event_rpc(log)
576 .map(BlockchainMessage::FeeProtocolUpdateEvent),
577 RpcEventType::PoolFeeProtocolCollect(_) => dex_extended
578 .parse_fee_protocol_collect_event_rpc(log)
579 .map(BlockchainMessage::FeeProtocolCollectEvent),
580 RpcEventType::NewBlock => Err(anyhow::anyhow!(
581 "NewBlock event type cannot parse pool logs"
582 )),
583 }
584 .map(Some)
585 .map_err(|e| BlockchainRpcClientError::MessageParsingError(e.to_string()))
586 }
587
588 #[cfg(not(feature = "hypersync"))]
589 fn blockchain_message_from_pool_log(
590 &self,
591 event_type: RpcEventType,
592 log: &RpcLog,
593 ) -> Result<Option<BlockchainMessage>, BlockchainRpcClientError> {
594 if log.removed {
595 log::debug!(
596 "Skipping removed pool log on chain '{}' for event {:?}",
597 self.chain.name,
598 event_type
599 );
600 return Ok(None);
601 }
602
603 Err(BlockchainRpcClientError::UnsupportedRpcResponseType(
604 format!("RPC pool log parsing for {event_type:?} requires the hypersync feature"),
605 ))
606 }
607
608 #[cfg(feature = "hypersync")]
609 fn pool_event_dex(event_type: RpcEventType) -> Result<DexType, BlockchainRpcClientError> {
610 match event_type {
611 RpcEventType::PoolSwap(dex)
612 | RpcEventType::PoolMint(dex)
613 | RpcEventType::PoolBurn(dex)
614 | RpcEventType::PoolCollect(dex)
615 | RpcEventType::PoolFlash(dex)
616 | RpcEventType::PoolFeeProtocolUpdate(dex)
617 | RpcEventType::PoolFeeProtocolCollect(dex) => Ok(dex),
618 RpcEventType::NewBlock => Err(BlockchainRpcClientError::InternalRpcClientError(
619 "NewBlock event type has no DEX".to_string(),
620 )),
621 }
622 }
623
624 pub async fn subscribe_blocks(&mut self) -> Result<(), BlockchainRpcClientError> {
630 self.subscribe_events(RpcEventType::NewBlock, RpcSubscription::new("newHeads"))
631 .await
632 }
633
634 pub async fn subscribe_pool_events(
640 &mut self,
641 event_type: RpcEventType,
642 addresses: &[Address],
643 event_signature: String,
644 ) -> Result<(), BlockchainRpcClientError> {
645 if matches!(event_type, RpcEventType::NewBlock) {
646 return Err(BlockchainRpcClientError::InvalidParameters(
647 "NewBlock is not a pool event subscription".to_string(),
648 ));
649 }
650
651 if addresses.is_empty() {
652 return self.unsubscribe_event_type(event_type).await;
653 }
654
655 self.replace_subscription(
656 event_type,
657 RpcSubscription::pool_logs(addresses, event_signature),
658 )
659 .await
660 }
661
662 pub async fn unsubscribe_blocks(&mut self) -> Result<(), BlockchainRpcClientError> {
668 self.unsubscribe_event_type(RpcEventType::NewBlock).await
669 }
670}
671
672#[cfg(test)]
673mod tests {
674 use alloy::primitives::address;
675 use nautilus_model::defi::{Chain, DexType};
676 use rstest::rstest;
677
678 use super::*;
679
680 #[rstest]
681 fn debug_redacts_websocket_rpc_url() {
682 const USERINFO_SECRET: &str = "core-wss-userinfo-secret";
683 const PATH_SECRET: &str = "core-wss-path-secret";
684 const QUERY_SECRET: &str = "core-wss-query-secret";
685 let wss_rpc_url = format!(
686 "wss://rpc-user:{USERINFO_SECRET}@rpc.example.com/{PATH_SECRET}?api_key={QUERY_SECRET}"
687 );
688 let client = CoreBlockchainRpcClient::new(
689 Chain::from_chain_id(1)
690 .expect("Ethereum chain should exist")
691 .clone(),
692 wss_rpc_url.clone(),
693 None,
694 );
695
696 let debug = format!("{client:?}");
697
698 assert!(debug.contains("wss_rpc_url: \"<redacted>\""));
699 assert!(!debug.contains(USERINFO_SECRET));
700 assert!(!debug.contains(PATH_SECRET));
701 assert!(!debug.contains(QUERY_SECRET));
702 assert!(!debug.contains(&wss_rpc_url));
703 }
704
705 #[rstest]
706 fn pool_logs_subscription_params_use_logs_filter_with_sorted_addresses() {
707 let event_signature =
708 "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67".to_string();
709 let subscription = RpcSubscription::pool_logs(
710 &[
711 address!("2222222222222222222222222222222222222222"),
712 address!("1111111111111111111111111111111111111111"),
713 ],
714 event_signature.clone(),
715 );
716
717 let params = subscription.params();
718 let filter = ¶ms[1];
719
720 assert_eq!(params[0], serde_json::json!("logs"));
721 assert_eq!(
722 filter["address"],
723 serde_json::json!([
724 "0x1111111111111111111111111111111111111111",
725 "0x2222222222222222222222222222222222222222",
726 ])
727 );
728 assert_eq!(filter["topics"], serde_json::json!([event_signature]));
729 }
730
731 #[cfg(feature = "hypersync")]
732 #[rstest]
733 fn pool_event_dex_rejects_block_event_type() {
734 assert!(CoreBlockchainRpcClient::pool_event_dex(RpcEventType::NewBlock).is_err());
735 assert_eq!(
736 CoreBlockchainRpcClient::pool_event_dex(RpcEventType::PoolSwap(DexType::UniswapV3))
737 .unwrap(),
738 DexType::UniswapV3
739 );
740 }
741
742 #[rstest]
743 fn removed_pool_log_returns_no_message() {
744 let client = CoreBlockchainRpcClient::new(
745 Chain::from_chain_id(1)
746 .expect("Ethereum chain should exist")
747 .clone(),
748 "ws://127.0.0.1:9".to_string(),
749 None,
750 );
751 let log = RpcLog {
752 removed: true,
753 log_index: Some("0x0".to_string()),
754 transaction_index: Some("0x0".to_string()),
755 transaction_hash: Some("0x1".to_string()),
756 block_hash: Some("0x1".to_string()),
757 block_number: Some("0x1".to_string()),
758 address: "0x1111111111111111111111111111111111111111".to_string(),
759 data: "0x".to_string(),
760 topics: vec![],
761 };
762
763 let message = client
764 .blockchain_message_from_pool_log(RpcEventType::PoolSwap(DexType::UniswapV3), &log)
765 .expect("removed logs should not fail conversion");
766
767 assert!(message.is_none());
768 }
769
770 #[tokio::test]
771 async fn next_rpc_message_skips_unsubscribe_confirmation() {
772 let mut client = CoreBlockchainRpcClient::new(
773 Chain::from_chain_id(1)
774 .expect("Ethereum chain should exist")
775 .clone(),
776 "ws://127.0.0.1:9".to_string(),
777 None,
778 );
779 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
780 client.wss_consumer_rx = Some(rx);
781 tx.send(Message::Text(
782 serde_json::json!({"jsonrpc": "2.0", "id": 1, "result": true})
783 .to_string()
784 .into(),
785 ))
786 .expect("unsubscribe ack should enqueue");
787 drop(tx);
788
789 let error = client
790 .next_rpc_message()
791 .await
792 .expect_err("unsubscribe ack should be skipped");
793
794 assert!(matches!(error, BlockchainRpcClientError::NoMessageReceived));
795 }
796}