From 4a036d26c97acd3e984ffadf34ac1f4a5eee940c Mon Sep 17 00:00:00 2001
From: quake <quake.wang@gmail.com>
Date: Mon, 11 Nov 2024 10:54:34 +0900
Subject: [PATCH] fix: try to fix relay protocol hardfork issue

---
 src/protocols/relayer.rs | 25 +++++++++++++++----------
 src/subcmds.rs           |  2 --
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/protocols/relayer.rs b/src/protocols/relayer.rs
index a93fa7b..088f5f1 100644
--- a/src/protocols/relayer.rs
+++ b/src/protocols/relayer.rs
@@ -12,7 +12,6 @@ use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 
 use crate::protocols::{Peers, BAD_MESSAGE_BAN_TIME};
-use crate::storage::Storage;
 
 const CHECK_PENDING_TXS_TOKEN: u64 = 0;
 
@@ -23,7 +22,6 @@ pub(crate) struct RelayProtocol {
     // Pending transactions which are waiting for relay
     pending_txs: Arc<RwLock<PendingTxs>>,
     consensus: Consensus,
-    storage: Storage,
     v3: bool,
 }
 
@@ -88,7 +86,6 @@ impl RelayProtocol {
         pending_txs: Arc<RwLock<PendingTxs>>,
         connected_peers: Arc<Peers>,
         consensus: Consensus,
-        storage: Storage,
         v3: bool,
     ) -> Self {
         Self {
@@ -96,7 +93,6 @@ impl RelayProtocol {
             pending_txs,
             connected_peers,
             consensus,
-            storage,
             v3,
         }
     }
@@ -116,17 +112,26 @@ impl CKBProtocolHandler for RelayProtocol {
         peer: PeerIndex,
         version: &str,
     ) {
-        let epoch = self
+        let prove_state_epoch = self
             .connected_peers
             .get_state(&peer)
-            .map(|peer_state| {
+            .and_then(|peer_state| {
                 peer_state
                     .get_prove_state()
                     .map(|s| s.get_last_header().header().epoch())
-                    .unwrap_or_else(|| self.storage.get_last_state().1.raw().epoch().unpack())
-                    .number()
-            })
-            .unwrap_or_default();
+            });
+
+        let epoch = match prove_state_epoch {
+            Some(epoch) => epoch.number(),
+            None => {
+                debug!("RelayProtocol.connected peer={} failed to get epoch, ignore and close the protocol", peer);
+                let _ = nc
+                    .p2p_control()
+                    .expect("p2p_control should be exist")
+                    .close_protocol(peer, nc.protocol_id());
+                return;
+            }
+        };
 
         let ckb2023 = self
             .consensus
diff --git a/src/subcmds.rs b/src/subcmds.rs
index f30534e..b3b822d 100644
--- a/src/subcmds.rs
+++ b/src/subcmds.rs
@@ -72,14 +72,12 @@ impl RunConfig {
             pending_txs.clone(),
             Arc::clone(&peers),
             consensus.clone(),
-            storage.clone(),
             false,
         );
         let relay_protocol_v3 = RelayProtocol::new(
             pending_txs.clone(),
             Arc::clone(&peers),
             consensus.clone(),
-            storage.clone(),
             true,
         );
         let light_client: Box<dyn CKBProtocolHandler> = Box::new(LightClientProtocol::new(