From c5aa51dca79c62d887d74798057e679ab95a9741 Mon Sep 17 00:00:00 2001
From: Gorka Irazoqui <gorka.irazoki@gmail.com>
Date: Mon, 23 Nov 2020 11:20:37 +0100
Subject: [PATCH] fix(chain manager): drop beacon-coincident outbound peers if
 different superblock consensus

---
 node/.Cargo.toml.swp                         | Bin 0 -> 1024 bytes
 node/src/actors/chain_manager/mod.rs         |   5 +++-
 node/src/actors/messages.rs                  |   7 ++++++
 node/src/actors/sessions_manager/handlers.rs |  23 ++++++++++++++++---
 node/src/actors/sessions_manager/mod.rs      |  16 +++++++++----
 5 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 node/.Cargo.toml.swp

diff --git a/node/.Cargo.toml.swp b/node/.Cargo.toml.swp
new file mode 100644
index 0000000000000000000000000000000000000000..cb506d0ded5de69a091d3e6b8763158900541383
GIT binary patch
literal 1024
zcmYc?$V<%2S1{7E)H7y40?#EG7}7I~607pFGm(U`GGXd;Lozb+vI7!RbVE!m3^0}I
zXXNLm>cdp)yQCIpm*f}dmuHsbrIzRxl@^!i=jErQ>N_VErRVFF<mcvKnlUOe8Umvs
LKz0a(qU!_zSbief

literal 0
HcmV?d00001

diff --git a/node/src/actors/chain_manager/mod.rs b/node/src/actors/chain_manager/mod.rs
index d4de81af4..f441b0fd6 100644
--- a/node/src/actors/chain_manager/mod.rs
+++ b/node/src/actors/chain_manager/mod.rs
@@ -70,7 +70,7 @@ use crate::{
         inventory_manager::InventoryManager,
         json_rpc::JsonRpcServer,
         messages::{
-            AddItem, AddItems, AddTransaction, Anycast, BlockNotify, Broadcast,
+            AddItem, AddItems, AddTransaction, Anycast, BlockNotify, Broadcast, DropOutboundPeers,
             GetBlocksEpochRange, GetItemBlock, NodeStatusNotify, SendInventoryItem,
             SendInventoryRequest, SendLastBeacon, SendSuperBlockVote, StoreInventoryItem,
             SuperBlockNotify,
@@ -1368,6 +1368,9 @@ impl ChainManager {
                         "Superblock consensus {} different from current superblock",
                         target_superblock_hash
                     );
+                    let sessions_manager_addr = SessionsManager::from_registry();
+
+                    sessions_manager_addr.do_send(DropOutboundPeers {});
                     act.initialize_from_storage(ctx);
                     act.update_state_machine(StateMachine::WaitingConsensus);
 
diff --git a/node/src/actors/messages.rs b/node/src/actors/messages.rs
index 5c6dd113e..0108649ba 100644
--- a/node/src/actors/messages.rs
+++ b/node/src/actors/messages.rs
@@ -941,6 +941,13 @@ impl Message for LogMessage {
     type Result = SessionsUnitResult;
 }
 
+/// Drop all outbound peers
+#[derive(Clone, Debug)]
+pub struct DropOutboundPeers {}
+impl Message for DropOutboundPeers {
+    type Result = ();
+}
+
 /// Set the LastBeacon
 #[derive(Clone, Debug)]
 pub struct SetLastBeacon {
diff --git a/node/src/actors/sessions_manager/handlers.rs b/node/src/actors/sessions_manager/handlers.rs
index 23437f7f2..876b1972c 100644
--- a/node/src/actors/sessions_manager/handlers.rs
+++ b/node/src/actors/sessions_manager/handlers.rs
@@ -1,6 +1,7 @@
 use std::{
     fmt::{Debug, Display},
     marker::Send,
+    net::SocketAddr,
 };
 
 use actix::{
@@ -16,9 +17,10 @@ use crate::actors::{
     chain_manager::ChainManager,
     codec::P2PCodec,
     messages::{
-        AddConsolidatedPeer, AddPeers, Anycast, Broadcast, Consolidate, Create, EpochNotification,
-        GetConsolidatedPeers, LogMessage, NumSessions, NumSessionsResult, PeerBeacon, Register,
-        RemoveAddressesFromTried, SessionsUnitResult, SetLastBeacon, TryMineBlock, Unregister,
+        AddConsolidatedPeer, AddPeers, Anycast, Broadcast, Consolidate, Create, DropOutboundPeers,
+        EpochNotification, GetConsolidatedPeers, LogMessage, NumSessions, NumSessionsResult,
+        PeerBeacon, Register, RemoveAddressesFromTried, SessionsUnitResult, SetLastBeacon,
+        TryMineBlock, Unregister,
     },
     peers_manager::PeersManager,
     session::Session,
@@ -468,3 +470,18 @@ impl Handler<SetLastBeacon> for SessionsManager {
         self.last_beacon = Some(msg.beacon);
     }
 }
+
+impl Handler<DropOutboundPeers> for SessionsManager {
+    type Result = <DropOutboundPeers as Message>::Result;
+
+    fn handle(&mut self, _msg: DropOutboundPeers, _ctx: &mut Context<Self>) -> Self::Result {
+        let outbound_peers: Vec<SocketAddr> = self
+            .sessions
+            .outbound_consolidated
+            .collection
+            .keys()
+            .cloned()
+            .collect();
+        self.drop_outbound_peers(outbound_peers.as_ref());
+    }
+}
diff --git a/node/src/actors/sessions_manager/mod.rs b/node/src/actors/sessions_manager/mod.rs
index 65cf8f31f..9b2cae2b1 100644
--- a/node/src/actors/sessions_manager/mod.rs
+++ b/node/src/actors/sessions_manager/mod.rs
@@ -260,12 +260,8 @@ impl SessionsManager {
                             );
                         }
                         // Unregister peers out of consensus
+                        act.drop_outbound_peers(peers_to_unregister.as_ref());
                         for peer in peers_to_unregister {
-                            if let Some(a) =
-                                act.sessions.outbound_consolidated.collection.get(&peer)
-                            {
-                                a.reference.do_send(CloseSession);
-                            }
                             peers_to_keep.remove(&peer);
                         }
                         // Mark remaining peers as safu
@@ -291,6 +287,16 @@ impl SessionsManager {
                 .cloned(),
         );
     }
+
+    /// Drop outbound peers
+    fn drop_outbound_peers(&mut self, peers_to_unregister: &[SocketAddr]) {
+        // Unregister peers out of consensus
+        for peer in peers_to_unregister {
+            if let Some(a) = self.sessions.outbound_consolidated.collection.get(&peer) {
+                a.reference.do_send(CloseSession);
+            }
+        }
+    }
 }
 
 /// Required traits for being able to retrieve SessionsManager address from registry