From 95f4925aa9bb47d83cc1d71d5818c512deabbbd9 Mon Sep 17 00:00:00 2001
From: rustaceanrob <rob.netzke@gmail.com>
Date: Fri, 29 Aug 2025 14:14:37 +0100
Subject: [PATCH] Track stale blocks

Keep track of when blocks are announced by `inv` or `headers` to
determine if a block is being withheld or stale.
---
 src/lib.rs | 35 ++++++++++++++++++++++-------------
 src/net.rs | 30 ++++++++++++++++++------------
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 44ede77..090bdf1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -101,6 +101,15 @@ impl ConnectionMetrics {
         now.duration_since(self.start_time)
     }
 
+    /// Is the last block considered stale according to the timeout.
+    pub fn stale_block(&self, timeout: Duration) -> bool {
+        if let Ok(lock) = self.timed_messages.lock() {
+            let now = Instant::now();
+            return now.duration_since(lock.last_block) > timeout;
+        }
+        false
+    }
+
     /// Has the connection failed to respond to a ping after the given duration.
     pub fn ping_timed_out(&self, timeout: Duration) -> bool {
         if let Ok(lock) = self.outbound_ping_state.lock() {
@@ -188,10 +197,13 @@ pub enum TimedMessage {
 }
 
 #[derive(Debug, Clone)]
-struct TimedMessages(HashMap<TimedMessage, MessageRate>);
+struct TimedMessages {
+    tracked: HashMap<TimedMessage, MessageRate>,
+    last_block: Instant,
+}
 
 impl TimedMessages {
-    fn new() -> Self {
+    fn new(now: Instant) -> Self {
         let mut map = HashMap::with_capacity(4);
         for key in [
             TimedMessage::BlockHeaders,
@@ -202,12 +214,15 @@ impl TimedMessages {
         ] {
             map.insert(key, MessageRate::new());
         }
-        Self(map)
+        Self {
+            tracked: map,
+            last_block: now,
+        }
     }
 
     fn add_single(&mut self, message: TimedMessage, now: Instant) {
         let val = self
-            .0
+            .tracked
             .get_mut(&message)
             .expect("all timed messages are in the map");
         val.add_single_message(now);
@@ -215,25 +230,19 @@ impl TimedMessages {
 
     fn add_many(&mut self, message: TimedMessage, num_messages: usize, now: Instant) {
         let val = self
-            .0
+            .tracked
             .get_mut(&message)
             .expect("all timed messages are in the map");
         val.add_messages(num_messages, now);
     }
 
     fn message_rate(&self, message: TimedMessage) -> &MessageRate {
-        self.0
+        self.tracked
             .get(&message)
             .expect("all timed messages are in the map")
     }
 }
 
-impl Default for TimedMessages {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 #[derive(Debug, Clone, Copy)]
 enum OutboundPing {
     Waiting { nonce: u64, then: Instant },
@@ -294,7 +303,7 @@ mod tests {
     fn test_timed_messages() {
         let now = Instant::now();
         let later = now.checked_add(Duration::from_secs(10)).unwrap();
-        let mut timed_messages = TimedMessages::new();
+        let mut timed_messages = TimedMessages::new(now);
         timed_messages.add_many(TimedMessage::Addr, 1_000, now);
         let msg_per_secs = timed_messages
             .message_rate(TimedMessage::Addr)
diff --git a/src/net.rs b/src/net.rs
index ffde394..4c784be 100644
--- a/src/net.rs
+++ b/src/net.rs
@@ -109,7 +109,8 @@ impl ConnectionExt for ConnectionConfig {
                         for response in responses {
                             write_half.write_message(response, &mut tcp_stream)?;
                         }
-                        let timed_messages = Arc::new(Mutex::new(TimedMessages::new()));
+                        let timed_messages =
+                            Arc::new(Mutex::new(TimedMessages::new(Instant::now())));
                         let outbound_ping = Arc::new(Mutex::new(OutboundPing::LastReceived {
                             then: Instant::now(),
                         }));
@@ -307,6 +308,7 @@ impl ConnectionReader {
                 NetworkMessage::Headers(_) => {
                     if let Ok(mut lock) = self.timed_messages.lock() {
                         lock.add_single(TimedMessage::BlockHeaders, Instant::now());
+                        lock.last_block = Instant::now();
                     }
                 }
                 NetworkMessage::CFilter(_) => {
@@ -342,17 +344,21 @@ impl ConnectionReader {
                     let now = Instant::now();
                     if let Ok(mut lock) = self.timed_messages.lock() {
                         for inv in payload {
-                            match inv {
-                                Inventory::WTx(_) => {
-                                    lock.add_single(TimedMessage::TransactionAnnouncement, now);
-                                }
-                                Inventory::Transaction(_) => {
-                                    lock.add_single(TimedMessage::TransactionAnnouncement, now);
-                                }
-                                Inventory::WitnessTransaction(_) => {
-                                    lock.add_single(TimedMessage::TransactionAnnouncement, now);
-                                }
-                                _ => (),
+                            if matches!(
+                                inv,
+                                Inventory::WTx(_)
+                                    | Inventory::WitnessTransaction(_)
+                                    | Inventory::Transaction(_)
+                            ) {
+                                lock.add_single(TimedMessage::TransactionAnnouncement, now);
+                            }
+                            if matches!(
+                                inv,
+                                Inventory::Block(_)
+                                    | Inventory::WitnessBlock(_)
+                                    | Inventory::CompactBlock(_)
+                            ) {
+                                lock.last_block = now;
                             }
                         }
                     }