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