1use crate::{
2 connection::ConnWrapper,
3 credentials::EthstatsCredentials,
4 error::EthStatsError,
5 events::{
6 AuthMsg, BlockMsg, BlockStats, HistoryMsg, LatencyMsg, NodeInfo, NodeStats, PendingMsg,
7 PendingStats, PingMsg, StatsMsg, TxStats, UncleStats,
8 },
9};
10use alloy_consensus::{BlockHeader, Sealable};
11use alloy_primitives::U256;
12use reth_chain_state::{CanonStateNotification, CanonStateSubscriptions};
13use reth_network_api::{NetworkInfo, Peers};
14use reth_primitives_traits::{Block, BlockBody};
15use reth_storage_api::{BlockReader, BlockReaderIdExt, NodePrimitivesProvider};
16use reth_transaction_pool::TransactionPool;
17
18use chrono::Local;
19use serde_json::Value;
20use std::{
21 str::FromStr,
22 sync::Arc,
23 time::{Duration, Instant},
24};
25use tokio::{
26 sync::{mpsc, Mutex, RwLock},
27 time::{interval, sleep, timeout},
28};
29use tokio_stream::StreamExt;
30use tokio_tungstenite::connect_async;
31use tracing::{debug, info};
32use url::Url;
33
34const HISTORY_UPDATE_RANGE: u64 = 50;
36const RECONNECT_INTERVAL: Duration = Duration::from_secs(5);
38const PING_TIMEOUT: Duration = Duration::from_secs(5);
40const REPORT_INTERVAL: Duration = Duration::from_secs(15);
42const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
44const READ_TIMEOUT: Duration = Duration::from_secs(30);
46
47#[derive(Debug)]
54pub struct EthStatsService<Network, Provider, Pool> {
55 credentials: EthstatsCredentials,
57 conn: Arc<RwLock<Option<ConnWrapper>>>,
59 last_ping: Arc<Mutex<Option<Instant>>>,
61 network: Network,
63 provider: Provider,
65 pool: Pool,
67}
68
69impl<Network, Provider, Pool> EthStatsService<Network, Provider, Pool>
70where
71 Network: NetworkInfo + Peers,
72 Provider: BlockReaderIdExt + CanonStateSubscriptions,
73 Pool: TransactionPool,
74{
75 pub async fn new(
83 url: &str,
84 network: Network,
85 provider: Provider,
86 pool: Pool,
87 ) -> Result<Self, EthStatsError> {
88 let credentials = EthstatsCredentials::from_str(url)?;
89 let service = Self {
90 credentials,
91 conn: Arc::new(RwLock::new(None)),
92 last_ping: Arc::new(Mutex::new(None)),
93 network,
94 provider,
95 pool,
96 };
97 service.connect().await?;
98
99 Ok(service)
100 }
101
102 async fn connect(&self) -> Result<(), EthStatsError> {
107 debug!(
108 target: "ethstats",
109 "Attempting to connect to EthStats server at {}", self.credentials.host
110 );
111 let full_url = format!("ws://{}/api", self.credentials.host);
112 let url = Url::parse(&full_url).map_err(EthStatsError::Url)?;
113
114 match timeout(CONNECT_TIMEOUT, connect_async(url.as_str())).await {
115 Ok(Ok((ws_stream, _))) => {
116 debug!(
117 target: "ethstats",
118 "Successfully connected to EthStats server at {}", self.credentials.host
119 );
120 let conn: ConnWrapper = ConnWrapper::new(ws_stream);
121 *self.conn.write().await = Some(conn.clone());
122 self.login().await?;
123 Ok(())
124 }
125 Ok(Err(e)) => Err(EthStatsError::WebSocket(e)),
126 Err(_) => {
127 debug!(target: "ethstats", "Connection to EthStats server timed out");
128 Err(EthStatsError::Timeout)
129 }
130 }
131 }
132
133 async fn login(&self) -> Result<(), EthStatsError> {
138 debug!(
139 target: "ethstats",
140 "Attempting to login to EthStats server as node_id {}", self.credentials.node_id
141 );
142 let conn = self.conn.read().await;
143 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
144
145 let network_status = self
146 .network
147 .network_status()
148 .await
149 .map_err(|e| EthStatsError::AuthError(e.to_string()))?;
150 let id = &self.credentials.node_id;
151 let secret = &self.credentials.secret;
152 let protocol = network_status
153 .capabilities
154 .iter()
155 .map(|cap| format!("{}/{}", cap.name, cap.version))
156 .collect::<Vec<_>>()
157 .join(", ");
158 let port = self.network.local_addr().port() as u64;
159
160 let auth = AuthMsg {
161 id: id.clone(),
162 secret: secret.clone(),
163 info: NodeInfo {
164 name: id.clone(),
165 node: network_status.client_version.clone(),
166 port,
167 network: self.network.chain_id().to_string(),
168 protocol,
169 api: "No".to_string(),
170 os: std::env::consts::OS.into(),
171 os_ver: std::env::consts::ARCH.into(),
172 client: "0.1.1".to_string(),
173 history: true,
174 },
175 };
176
177 let message = auth.generate_login_message();
178 conn.write_json(&message).await?;
179
180 let response =
181 timeout(READ_TIMEOUT, conn.read_json()).await.map_err(|_| EthStatsError::Timeout)??;
182
183 if let Some(ack) = response.get("emit") &&
184 ack.get(0) == Some(&Value::String("ready".to_string()))
185 {
186 info!(
187 target: "ethstats",
188 "Login successful to EthStats server as node_id {}", self.credentials.node_id
189 );
190 return Ok(());
191 }
192
193 debug!(target: "ethstats", "Login failed: Unauthorized or unexpected login response");
194 Err(EthStatsError::AuthError("Unauthorized or unexpected login response".into()))
195 }
196
197 async fn report_stats(&self) -> Result<(), EthStatsError> {
202 let conn = self.conn.read().await;
203 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
204
205 let stats_msg = StatsMsg {
206 id: self.credentials.node_id.clone(),
207 stats: NodeStats {
208 active: true,
209 syncing: self.network.is_syncing(),
210 peers: self.network.num_connected_peers() as u64,
211 gas_price: 0, uptime: 100,
213 },
214 };
215
216 let message = stats_msg.generate_stats_message();
217 conn.write_json(&message).await?;
218
219 Ok(())
220 }
221
222 async fn send_ping(&self) -> Result<(), EthStatsError> {
227 let conn = self.conn.read().await;
228 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
229
230 let ping_time = Instant::now();
231 *self.last_ping.lock().await = Some(ping_time);
232
233 let client_time = Local::now().format("%Y-%m-%d %H:%M:%S%.f %:z %Z").to_string();
234 let ping_msg = PingMsg { id: self.credentials.node_id.clone(), client_time };
235
236 let message = ping_msg.generate_ping_message();
237 conn.write_json(&message).await?;
238
239 let active_ping = self.last_ping.clone();
241 let conn_ref = self.conn.clone();
242 tokio::spawn(async move {
243 sleep(PING_TIMEOUT).await;
244 let mut active = active_ping.lock().await;
245 if active.is_some() {
246 debug!(target: "ethstats", "Ping timeout");
247 *active = None;
248 if let Some(conn) = conn_ref.write().await.take() {
250 let _ = conn.close().await;
251 }
252 }
253 });
254
255 Ok(())
256 }
257
258 async fn report_latency(&self) -> Result<(), EthStatsError> {
263 let conn = self.conn.read().await;
264 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
265
266 let mut active = self.last_ping.lock().await;
267 if let Some(start) = active.take() {
268 let latency = start.elapsed().as_millis() as u64 / 2;
269
270 debug!(target: "ethstats", "Reporting latency: {}ms", latency);
271
272 let latency_msg = LatencyMsg { id: self.credentials.node_id.clone(), latency };
273
274 let message = latency_msg.generate_latency_message();
275 conn.write_json(&message).await?
276 }
277
278 Ok(())
279 }
280
281 async fn report_pending(&self) -> Result<(), EthStatsError> {
286 let conn = self.conn.read().await;
287 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
288 let pending = self.pool.pool_size().pending as u64;
289
290 debug!(target: "ethstats", "Reporting pending txs: {}", pending);
291
292 let pending_msg =
293 PendingMsg { id: self.credentials.node_id.clone(), stats: PendingStats { pending } };
294
295 let message = pending_msg.generate_pending_message();
296 conn.write_json(&message).await?;
297
298 Ok(())
299 }
300
301 async fn report_block(
310 &self,
311 head: Option<CanonStateNotification<<Provider as NodePrimitivesProvider>::Primitives>>,
312 ) -> Result<(), EthStatsError> {
313 let conn = self.conn.read().await;
314 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
315
316 let block_number = if let Some(head) = head {
317 head.tip().header().number()
318 } else {
319 self.provider
320 .best_block_number()
321 .map_err(|e| EthStatsError::DataFetchError(e.to_string()))?
322 };
323
324 match self.provider.block_by_id(block_number.into()) {
325 Ok(Some(block)) => {
326 let block_msg = BlockMsg {
327 id: self.credentials.node_id.clone(),
328 block: self.block_to_stats(&block)?,
329 };
330
331 debug!(target: "ethstats", "Reporting block: {}", block_number);
332
333 let message = block_msg.generate_block_message();
334 conn.write_json(&message).await?;
335 }
336 Ok(None) => {
337 debug!(target: "ethstats", "Block {} not found", block_number);
339 return Err(EthStatsError::BlockNotFound(block_number));
340 }
341 Err(e) => {
342 debug!(target: "ethstats", "Error fetching block {}: {}", block_number, e);
343 return Err(EthStatsError::DataFetchError(e.to_string()));
344 }
345 };
346
347 Ok(())
348 }
349
350 fn block_to_stats(
358 &self,
359 block: &<Provider as BlockReader>::Block,
360 ) -> Result<BlockStats, EthStatsError> {
361 let body = block.body();
362 let header = block.header();
363
364 let txs = body.transaction_hashes_iter().copied().map(|hash| TxStats { hash }).collect();
365
366 Ok(BlockStats {
367 number: U256::from(header.number()),
368 hash: header.hash_slow(),
369 parent_hash: header.parent_hash(),
370 timestamp: U256::from(header.timestamp()),
371 miner: header.beneficiary(),
372 gas_used: header.gas_used(),
373 gas_limit: header.gas_limit(),
374 diff: header.difficulty().to_string(),
375 total_diff: "0".into(),
376 txs,
377 tx_root: header.transactions_root(),
378 root: header.state_root(),
379 uncles: UncleStats(vec![]),
380 })
381 }
382
383 async fn report_history(&self, list: Option<&Vec<u64>>) -> Result<(), EthStatsError> {
392 let conn = self.conn.read().await;
393 let conn = conn.as_ref().ok_or(EthStatsError::NotConnected)?;
394
395 let indexes = if let Some(list) = list {
396 list
397 } else {
398 let best_block_number = self
399 .provider
400 .best_block_number()
401 .map_err(|e| EthStatsError::DataFetchError(e.to_string()))?;
402
403 let start = best_block_number.saturating_sub(HISTORY_UPDATE_RANGE);
404
405 &(start..=best_block_number).collect()
406 };
407
408 let mut blocks = Vec::with_capacity(indexes.len());
409 for &block_number in indexes {
410 match self.provider.block_by_id(block_number.into()) {
411 Ok(Some(block)) => {
412 blocks.push(block);
413 }
414 Ok(None) => {
415 debug!(target: "ethstats", "Block {} not found", block_number);
417 break;
418 }
419 Err(e) => {
420 debug!(target: "ethstats", "Error fetching block {}: {}", block_number, e);
421 break;
422 }
423 }
424 }
425
426 let history: Vec<BlockStats> =
427 blocks.iter().map(|block| self.block_to_stats(block)).collect::<Result<_, _>>()?;
428
429 if history.is_empty() {
430 debug!(target: "ethstats", "No history to send to stats server");
431 } else {
432 debug!(
433 target: "ethstats",
434 "Sending historical blocks to ethstats, first: {}, last: {}",
435 history.first().unwrap().number,
436 history.last().unwrap().number
437 );
438 }
439
440 let history_msg = HistoryMsg { id: self.credentials.node_id.clone(), history };
441
442 let message = history_msg.generate_history_message();
443 conn.write_json(&message).await?;
444
445 Ok(())
446 }
447
448 async fn report(&self) -> Result<(), EthStatsError> {
453 self.send_ping().await?;
454 self.report_block(None).await?;
455 self.report_pending().await?;
456 self.report_stats().await?;
457
458 Ok(())
459 }
460
461 async fn handle_message(&self, msg: Value) -> Result<(), EthStatsError> {
487 let emit = match msg.get("emit") {
488 Some(emit) => emit,
489 None => {
490 debug!(target: "ethstats", "Stats server sent non-broadcast, msg {}", msg);
491 return Err(EthStatsError::InvalidRequest);
492 }
493 };
494
495 let command = match emit.get(0) {
496 Some(Value::String(command)) => command.as_str(),
497 _ => {
498 debug!(target: "ethstats", "Invalid stats server message type, msg {}", msg);
499 return Err(EthStatsError::InvalidRequest);
500 }
501 };
502
503 match command {
504 "node-pong" => {
505 self.report_latency().await?;
506 }
507 "history" => {
508 let block_numbers = emit
509 .get(1)
510 .and_then(|v| v.as_object())
511 .and_then(|obj| obj.get("list"))
512 .and_then(|v| v.as_array());
513
514 if block_numbers.is_none() {
515 self.report_history(None).await?;
516
517 return Ok(());
518 }
519
520 let block_numbers = block_numbers
521 .unwrap()
522 .iter()
523 .map(|val| {
524 val.as_u64().ok_or_else(|| {
525 debug!(
526 target: "ethstats",
527 "Invalid stats history block number, msg {}", msg
528 );
529 EthStatsError::InvalidRequest
530 })
531 })
532 .collect::<Result<_, _>>()?;
533
534 self.report_history(Some(&block_numbers)).await?;
535 }
536 other => debug!(target: "ethstats", "Unhandled command: {}", other),
537 }
538
539 Ok(())
540 }
541
542 pub async fn run(self) {
554 let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
556 let (message_tx, mut message_rx) = mpsc::channel(32);
557 let (head_tx, mut head_rx) = mpsc::channel(10);
558
559 let read_handle = {
561 let conn = self.conn.clone();
562 let message_tx = message_tx.clone();
563 let shutdown_tx = shutdown_tx.clone();
564
565 tokio::spawn(async move {
566 loop {
567 let conn = conn.read().await;
568 if let Some(conn) = conn.as_ref() {
569 match conn.read_json().await {
570 Ok(msg) => {
571 if message_tx.send(msg).await.is_err() {
572 break;
573 }
574 }
575 Err(e) => {
576 debug!(target: "ethstats", "Read error: {}", e);
577 break;
578 }
579 }
580 } else {
581 sleep(RECONNECT_INTERVAL).await;
582 }
583 }
584
585 let _ = shutdown_tx.send(()).await;
586 })
587 };
588
589 let canonical_stream_handle = {
590 let mut canonical_stream = self.provider.canonical_state_stream();
591 let head_tx = head_tx.clone();
592 let shutdown_tx = shutdown_tx.clone();
593
594 tokio::spawn(async move {
595 loop {
596 let head = canonical_stream.next().await;
597 if let Some(head) = head &&
598 head_tx.send(head).await.is_err()
599 {
600 break;
601 }
602 }
603
604 let _ = shutdown_tx.send(()).await;
605 })
606 };
607
608 let mut pending_tx_receiver = self.pool.pending_transactions_listener();
609
610 let mut report_interval = interval(REPORT_INTERVAL);
612 let mut reconnect_interval = interval(RECONNECT_INTERVAL);
613
614 loop {
616 tokio::select! {
617 _ = shutdown_rx.recv() => {
619 info!(target: "ethstats", "Shutting down ethstats service");
620 break;
621 }
622
623 Some(msg) = message_rx.recv() => {
625 if let Err(e) = self.handle_message(msg).await {
626 debug!(target: "ethstats", "Error handling message: {}", e);
627 self.disconnect().await;
628 }
629 }
630
631 Some(head) = head_rx.recv() => {
633 if let Err(e) = self.report_block(Some(head)).await {
634 debug!(target: "ethstats", "Failed to report block: {}", e);
635 self.disconnect().await;
636 }
637
638 if let Err(e) = self.report_pending().await {
639 debug!(target: "ethstats", "Failed to report pending: {}", e);
640 self.disconnect().await;
641 }
642 }
643
644 _= pending_tx_receiver.recv() => {
646 if let Err(e) = self.report_pending().await {
647 debug!(target: "ethstats", "Failed to report pending: {}", e);
648 self.disconnect().await;
649 }
650 }
651
652 _ = report_interval.tick() => {
654 if let Err(e) = self.report().await {
655 debug!(target: "ethstats", "Failed to report: {}", e);
656 self.disconnect().await;
657 }
658 }
659
660 _ = reconnect_interval.tick(), if self.conn.read().await.is_none() => {
662 match self.connect().await {
663 Ok(_) => info!(target: "ethstats", "Reconnected successfully"),
664 Err(e) => debug!(target: "ethstats", "Reconnect failed: {}", e),
665 }
666 }
667 }
668 }
669
670 self.disconnect().await;
672
673 read_handle.abort();
675 canonical_stream_handle.abort();
676 }
677
678 async fn disconnect(&self) {
683 if let Some(conn) = self.conn.write().await.take() &&
684 let Err(e) = conn.close().await
685 {
686 debug!(target: "ethstats", "Error closing connection: {}", e);
687 }
688 }
689
690 #[cfg(test)]
692 pub async fn is_connected(&self) -> bool {
693 self.conn.read().await.is_some()
694 }
695}
696
697#[cfg(test)]
698mod tests {
699 use super::*;
700 use futures_util::{SinkExt, StreamExt};
701 use reth_network_api::noop::NoopNetwork;
702 use reth_storage_api::noop::NoopProvider;
703 use reth_transaction_pool::noop::NoopTransactionPool;
704 use serde_json::json;
705 use tokio::net::TcpListener;
706 use tokio_tungstenite::tungstenite::protocol::{frame::Utf8Bytes, Message};
707
708 const TEST_HOST: &str = "127.0.0.1";
709 const TEST_PORT: u16 = 0; async fn setup_mock_server() -> (String, tokio::task::JoinHandle<()>) {
712 let listener = TcpListener::bind((TEST_HOST, TEST_PORT)).await.unwrap();
713 let addr = listener.local_addr().unwrap();
714
715 let handle = tokio::spawn(async move {
716 let (stream, _) = listener.accept().await.unwrap();
717 let mut ws_stream = tokio_tungstenite::accept_async(stream).await.unwrap();
718
719 if let Some(Ok(Message::Text(text))) = ws_stream.next().await {
721 let value: serde_json::Value = serde_json::from_str(&text).unwrap();
722 if value["emit"][0] == "hello" {
723 let response = json!({
724 "emit": ["ready", []]
725 });
726 ws_stream
727 .send(Message::Text(Utf8Bytes::from(response.to_string())))
728 .await
729 .unwrap();
730 }
731 }
732
733 while let Some(Ok(msg)) = ws_stream.next().await {
735 if let Message::Text(text) = msg &&
736 text.contains("node-ping")
737 {
738 let pong = json!({
739 "emit": ["node-pong", {"id": "test-node"}]
740 });
741 ws_stream.send(Message::Text(Utf8Bytes::from(pong.to_string()))).await.unwrap();
742 }
743 }
744 });
745
746 (addr.to_string(), handle)
747 }
748
749 #[tokio::test]
750 async fn test_connection_and_login() {
751 let (server_url, server_handle) = setup_mock_server().await;
752 let ethstats_url = format!("test-node:test-secret@{server_url}");
753
754 let network = NoopNetwork::default();
755 let provider = NoopProvider::default();
756 let pool = NoopTransactionPool::default();
757
758 let service = EthStatsService::new(ðstats_url, network, provider, pool)
759 .await
760 .expect("Service should connect");
761
762 assert!(service.is_connected().await, "Service should be connected");
764
765 server_handle.abort();
767 }
768
769 #[tokio::test]
770 async fn test_history_command_handling() {
771 let (server_url, server_handle) = setup_mock_server().await;
772 let ethstats_url = format!("test-node:test-secret@{server_url}");
773
774 let network = NoopNetwork::default();
775 let provider = NoopProvider::default();
776 let pool = NoopTransactionPool::default();
777
778 let service = EthStatsService::new(ðstats_url, network, provider, pool)
779 .await
780 .expect("Service should connect");
781
782 let history_cmd = json!({
784 "emit": ["history", {"list": [1, 2, 3]}]
785 });
786
787 service.handle_message(history_cmd).await.expect("History command should be handled");
788
789 server_handle.abort();
791 }
792
793 #[tokio::test]
794 async fn test_invalid_url_handling() {
795 let network = NoopNetwork::default();
796 let provider = NoopProvider::default();
797 let pool = NoopTransactionPool::default();
798
799 let result = EthStatsService::new(
801 "test-node@localhost",
802 network.clone(),
803 provider.clone(),
804 pool.clone(),
805 )
806 .await;
807 assert!(
808 matches!(result, Err(EthStatsError::InvalidUrl(_))),
809 "Should detect invalid URL format"
810 );
811
812 let result = EthStatsService::new("invalid-url", network, provider, pool).await;
814 assert!(
815 matches!(result, Err(EthStatsError::InvalidUrl(_))),
816 "Should detect invalid URL format"
817 );
818 }
819}