From 2f2e13b6c2c8741ca9d825eaaef736ede484bc85 Mon Sep 17 00:00:00 2001
From: Amiti Uttarwar <amiti@uttarwar.org>
Date: Mon, 20 Jul 2020 14:24:48 -0700
Subject: [PATCH] [net/refactor] Simplify multiple-connection checks

Extract logic that check multiple connection types into interface functions &
structure as switch statements. This makes it very clear what touch points are
for accessing `m_conn_type` & using the switch statements enables the compiler
to warn if a new connection type is introduced but not handled for these cases.
---
 src/net.cpp            |  4 ++--
 src/net.h              | 15 +++++++++++++++
 src/net_processing.cpp |  4 ++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/net.cpp b/src/net.cpp
index 35daddbbe25..539ba9f40c9 100644
--- a/src/net.cpp
+++ b/src/net.cpp
@@ -1648,7 +1648,7 @@ void CConnman::ThreadDNSAddressSeed()
                     {
                         LOCK(cs_vNodes);
                         for (const CNode* pnode : vNodes) {
-                            nRelevant += pnode->fSuccessfullyConnected && !pnode->IsFeelerConn() && !pnode->IsAddrFetchConn() && !pnode->IsManualConn() && !pnode->IsInboundConn();
+                            if (pnode->fSuccessfullyConnected && pnode->IsOutboundOrBlockRelayConn()) ++nRelevant;
                         }
                     }
                     if (nRelevant >= 2) {
@@ -1758,7 +1758,7 @@ int CConnman::GetExtraOutboundCount()
     {
         LOCK(cs_vNodes);
         for (const CNode* pnode : vNodes) {
-            if (!pnode->IsInboundConn() && !pnode->IsManualConn() && !pnode->IsFeelerConn() && !pnode->fDisconnect && !pnode->IsAddrFetchConn() && pnode->fSuccessfullyConnected) {
+            if (pnode->fSuccessfullyConnected && !pnode->fDisconnect && pnode->IsOutboundOrBlockRelayConn()) {
                 ++nOutbound;
             }
         }
diff --git a/src/net.h b/src/net.h
index 4834050122c..444a1ceed25 100644
--- a/src/net.h
+++ b/src/net.h
@@ -789,6 +789,21 @@ public:
     std::atomic_bool fPauseRecv{false};
     std::atomic_bool fPauseSend{false};
 
+    bool IsOutboundOrBlockRelayConn() const {
+        switch(m_conn_type) {
+            case ConnectionType::OUTBOUND:
+            case ConnectionType::BLOCK_RELAY:
+                return true;
+            case ConnectionType::INBOUND:
+            case ConnectionType::MANUAL:
+            case ConnectionType::ADDR_FETCH:
+            case ConnectionType::FEELER:
+                return false;
+        }
+
+        assert(false);
+    }
+
     bool IsFullOutboundConn() const {
         return m_conn_type == ConnectionType::OUTBOUND;
     }
diff --git a/src/net_processing.cpp b/src/net_processing.cpp
index da2fa802634..f90ae658ec1 100644
--- a/src/net_processing.cpp
+++ b/src/net_processing.cpp
@@ -829,7 +829,7 @@ void UpdateLastBlockAnnounceTime(NodeId node, int64_t time_in_seconds)
 
 static bool IsOutboundDisconnectionCandidate(const CNode& node)
 {
-    return !(node.IsInboundConn() || node.IsManualConn() || node.IsFeelerConn() || node.IsAddrFetchConn());
+    return node.IsOutboundOrBlockRelayConn();
 }
 
 void PeerLogicValidation::InitializeNode(CNode *pnode) {
@@ -2324,7 +2324,7 @@ void ProcessMessage(
         {
             connman.SetServices(pfrom.addr, nServices);
         }
-        if (!pfrom.IsInboundConn() && !pfrom.IsFeelerConn() && !pfrom.IsManualConn() && !HasAllDesirableServiceFlags(nServices))
+        if (pfrom.ExpectServicesFromConn() && !HasAllDesirableServiceFlags(nServices))
         {
             LogPrint(BCLog::NET, "peer=%d does not offer the expected services (%08x offered, %08x expected); disconnecting\n", pfrom.GetId(), nServices, GetDesirableServiceFlags(nServices));
             pfrom.fDisconnect = true;