diff --git a/client/network/src/protocol/notifications/behaviour.rs b/client/network/src/protocol/notifications/behaviour.rs index 1e0f4f1b639ff..e1ccdf6920406 100644 --- a/client/network/src/protocol/notifications/behaviour.rs +++ b/client/network/src/protocol/notifications/behaviour.rs @@ -415,7 +415,7 @@ impl Notifications { /// Disconnects the given peer if we are connected to it. pub fn disconnect_peer(&mut self, peer_id: &PeerId, set_id: sc_peerset::SetId) { trace!(target: "sub-libp2p", "External API => Disconnect({}, {:?})", peer_id, set_id); - self.disconnect_peer_inner(peer_id, set_id); + self.disconnect_peer_inner(peer_id, set_id, None); } /// Inner implementation of `disconnect_peer`. @@ -423,6 +423,7 @@ impl Notifications { &mut self, peer_id: &PeerId, set_id: sc_peerset::SetId, + ban: Option, ) { let mut entry = if let Entry::Occupied(entry) = self.peers.entry((*peer_id, set_id)) { entry @@ -441,7 +442,12 @@ impl Notifications { PeerState::DisabledPendingEnable { connections, timer_deadline, timer: _ } => { trace!(target: "sub-libp2p", "PSM <= Dropped({}, {:?})", peer_id, set_id); self.peerset.dropped(set_id, *peer_id, DropReason::Unknown); - *entry.into_mut() = PeerState::Disabled { connections, backoff_until: Some(timer_deadline) } + let backoff_until = Some(if let Some(ban) = ban { + cmp::max(timer_deadline, Instant::now() + ban) + } else { + timer_deadline + }); + *entry.into_mut() = PeerState::Disabled { connections, backoff_until } }, // Enabled => Disabled. @@ -489,7 +495,8 @@ impl Notifications { .iter() .any(|(_, s)| matches!(s, ConnectionState::Opening))); - *entry.into_mut() = PeerState::Disabled { connections, backoff_until: None } + let backoff_until = ban.map(|dur| Instant::now() + dur); + *entry.into_mut() = PeerState::Disabled { connections, backoff_until } }, // Incoming => Disabled. @@ -524,6 +531,13 @@ impl Notifications { *connec_state = ConnectionState::Closing; } + let backoff_until = match (backoff_until, ban) { + (Some(a), Some(b)) => Some(cmp::max(a, Instant::now() + b)), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(Instant::now() + b), + (None, None) => None, + }; + debug_assert!(!connections .iter() .any(|(_, s)| matches!(s, ConnectionState::OpenDesiredByRemote)));