remove unnecessary MisBehaviorResult

network/src/tentacle/protocols/discovery/addr.rs

@@ -1,7 +1,7 @@
 use std::collections::hash_map::RandomState;
 
 use bloom_filters::{BloomFilter, DefaultBuildHashKernels, StableBloomFilter};
-use p2p::{context::SessionContext, multiaddr::Multiaddr, ProtocolId, SessionId};
+use p2p::{multiaddr::Multiaddr, ProtocolId, SessionId};
 
 use crate::Flags;
 
@@ -21,29 +21,13 @@ pub enum Misbehavior {
     InvalidData,
 }
 
-/// Misbehavior report result
-pub enum MisbehaveResult {
-    /// Disconnect this peer
-    Disconnect,
-}
-
-impl MisbehaveResult {
-    pub fn is_disconnect(&self) -> bool {
-        match self {
-            MisbehaveResult::Disconnect => true,
-            // _ => false,
-        }
-    }
-}
-
 // FIXME: Should be peer store?
 pub trait AddressManager {
     fn register(&self, id: SessionId, pid: ProtocolId, version: &str);
     fn unregister(&self, id: SessionId, pid: ProtocolId);
     fn is_valid_addr(&self, addr: &Multiaddr) -> bool;
     fn add_new_addr(&mut self, session_id: SessionId, addr: (Multiaddr, Flags));
     fn add_new_addrs(&mut self, session_id: SessionId, addrs: Vec<(Multiaddr, Flags)>);
-    fn misbehave(&mut self, session: &SessionContext, kind: &Misbehavior) -> MisbehaveResult;
     fn get_random(&mut self, n: usize, target: Flags) -> Vec<(Multiaddr, Flags)>;
     fn required_flags(&self) -> Flags;
     fn node_flags(&self, id: SessionId) -> Option<Flags>;

network/src/tentacle/protocols/discovery/mod.rs

@@ -16,7 +16,7 @@ use p2p::{
 use rand::seq::SliceRandom;
 
 pub use self::{
-    addr::{AddrKnown, AddressManager, MisbehaveResult, Misbehavior},
+    addr::{AddrKnown, AddressManager, Misbehavior},
     protocol::{DiscoveryMessage, Node, Nodes},
     state::SessionState,
 };
@@ -95,10 +95,6 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
         let session = context.session;
         trace!("[received message]: length={}", data.len());
 
-        let mgr = &mut self.addr_mgr;
-        let mut check =
-            |behavior: Misbehavior| -> bool { mgr.misbehave(session, &behavior).is_disconnect() };
-
         match decode(&data) {
             Some(item) => {
                 match item {
@@ -109,10 +105,13 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
                         required_flags,
                     } => {
                         if let Some(state) = self.sessions.get_mut(&session.id) {
-                            if state.received_get_nodes && check(Misbehavior::DuplicateGetNodes) {
-                                if context.disconnect(session.id).await.is_err() {
-                                    debug!("disconnect {:?} send fail", session.id)
-                                }
+                            if state.received_get_nodes {
+                                disconnect_due_to_misbehavior(
+                                    context,
+                                    session,
+                                    &Misbehavior::DuplicateGetNodes,
+                                )
+                                .await;
                                 return;
                             }
 
@@ -171,22 +170,20 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
                     }
                     DiscoveryMessage::Nodes(nodes) => {
                         if let Some(misbehavior) = verify_nodes_message(&nodes) {
-                            if check(misbehavior) {
-                                if context.disconnect(session.id).await.is_err() {
-                                    debug!("disconnect {:?} send fail", session.id)
-                                }
-                                return;
-                            }
+                            disconnect_due_to_misbehavior(context, session, &misbehavior).await;
+                            return;
                         }
 
                         if let Some(state) = self.sessions.get_mut(&session.id) {
                             if !nodes.announce && state.received_nodes {
                                 warn!("already received Nodes(announce=false) message");
-                                if check(Misbehavior::DuplicateFirstNodes)
-                                    && context.disconnect(session.id).await.is_err()
-                                {
-                                    debug!("disconnect {:?} send fail", session.id)
-                                }
+                                disconnect_due_to_misbehavior(
+                                    context,
+                                    session,
+                                    &Misbehavior::DuplicateFirstNodes,
+                                )
+                                .await;
+                                return;
                             } else {
                                 let addrs = nodes
                                     .items
@@ -211,14 +208,8 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
                 }
             }
             None => {
-                if self
-                    .addr_mgr
-                    .misbehave(session, &Misbehavior::InvalidData)
-                    .is_disconnect()
-                    && context.disconnect(session.id).await.is_err()
-                {
-                    debug!("disconnect {:?} send fail", session.id)
-                }
+                disconnect_due_to_misbehavior(context, session, &Misbehavior::InvalidData).await;
+                return;
             }
         }
     }
@@ -356,16 +347,6 @@ impl AddressManager for DiscoveryAddressManager {
         }
     }
 
-    fn misbehave(&mut self, session: &SessionContext, behavior: &Misbehavior) -> MisbehaveResult {
-        error!(
-            "DiscoveryProtocol detects abnormal behavior, session: {:?}, behavior: {:?}",
-            session, behavior
-        );
-
-        // FIXME:
-        MisbehaveResult::Disconnect
-    }
-
     fn get_random(&mut self, n: usize, flags: Flags) -> Vec<(Multiaddr, Flags)> {
         let fetch_random_addrs = self
             .network_state
@@ -395,3 +376,17 @@ impl AddressManager for DiscoveryAddressManager {
         })
     }
 }
+
+async fn disconnect_due_to_misbehavior(
+    context: ProtocolContextMutRef<'_>,
+    session: &SessionContext,
+    behavior: &Misbehavior,
+) {
+    error!(
+        "DiscoveryProtocol detects abnormal behavior, session: {:?}, behavior: {:?}",
+        session, behavior
+    );
+    if context.disconnect(session.id).await.is_err() {
+        debug!("disconnect {:?} send fail", session.id)
+    }
+}