From d36241b4a12719800d04673d8f700ccf0d9253a8 Mon Sep 17 00:00:00 2001
From: Age Manning <Age@AgeManning.com>
Date: Tue, 27 Feb 2024 17:32:48 +1100
Subject: [PATCH] Track multiaddr in connection status (#5308)

* Record the multiaddr for connected peers
---
 .../src/peer_manager/peerdb.rs                |  4 +-
 .../src/peer_manager/peerdb/peer_info.rs      | 39 ++++++++++++-------
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
index a6bf3ffec..8aef0bad4 100644
--- a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
+++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
@@ -756,8 +756,8 @@ impl<TSpec: EthSpec> PeerDB<TSpec> {
 
                 // Update the connection state
                 match direction {
-                    ConnectionDirection::Incoming => info.connect_ingoing(Some(seen_address)),
-                    ConnectionDirection::Outgoing => info.connect_outgoing(Some(seen_address)),
+                    ConnectionDirection::Incoming => info.connect_ingoing(seen_address),
+                    ConnectionDirection::Outgoing => info.connect_outgoing(seen_address),
                 }
             }
 
diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs
index 44c54511d..c5e13c515 100644
--- a/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs
+++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs
@@ -307,13 +307,13 @@ impl<T: EthSpec> PeerInfo<T> {
 
     /// Checks if the peer is outbound-only
     pub fn is_outbound_only(&self) -> bool {
-        matches!(self.connection_status, Connected {n_in, n_out} if n_in == 0 && n_out > 0)
+        matches!(self.connection_status, Connected {n_in, n_out, ..} if n_in == 0 && n_out > 0)
     }
 
     /// Returns the number of connections with this peer.
     pub fn connections(&self) -> (u8, u8) {
         match self.connection_status {
-            Connected { n_in, n_out } => (n_in, n_out),
+            Connected { n_in, n_out, .. } => (n_in, n_out),
             _ => (0, 0),
         }
     }
@@ -421,7 +421,9 @@ impl<T: EthSpec> PeerInfo<T> {
 
     /// Modifies the status to Connected and increases the number of ingoing
     /// connections by one
-    pub(super) fn connect_ingoing(&mut self, seen_multiaddr: Option<Multiaddr>) {
+    pub(super) fn connect_ingoing(&mut self, multiaddr: Multiaddr) {
+        self.seen_multiaddrs.insert(multiaddr.clone());
+
         match &mut self.connection_status {
             Connected { n_in, .. } => *n_in += 1,
             Disconnected { .. }
@@ -429,19 +431,20 @@ impl<T: EthSpec> PeerInfo<T> {
             | Dialing { .. }
             | Disconnecting { .. }
             | Unknown => {
-                self.connection_status = Connected { n_in: 1, n_out: 0 };
+                self.connection_status = Connected {
+                    n_in: 1,
+                    n_out: 0,
+                    multiaddr,
+                };
                 self.connection_direction = Some(ConnectionDirection::Incoming);
             }
         }
-
-        if let Some(multiaddr) = seen_multiaddr {
-            self.seen_multiaddrs.insert(multiaddr);
-        }
     }
 
     /// Modifies the status to Connected and increases the number of outgoing
     /// connections by one
-    pub(super) fn connect_outgoing(&mut self, seen_multiaddr: Option<Multiaddr>) {
+    pub(super) fn connect_outgoing(&mut self, multiaddr: Multiaddr) {
+        self.seen_multiaddrs.insert(multiaddr.clone());
         match &mut self.connection_status {
             Connected { n_out, .. } => *n_out += 1,
             Disconnected { .. }
@@ -449,13 +452,14 @@ impl<T: EthSpec> PeerInfo<T> {
             | Dialing { .. }
             | Disconnecting { .. }
             | Unknown => {
-                self.connection_status = Connected { n_in: 0, n_out: 1 };
+                self.connection_status = Connected {
+                    n_in: 0,
+                    n_out: 1,
+                    multiaddr,
+                };
                 self.connection_direction = Some(ConnectionDirection::Outgoing);
             }
         }
-        if let Some(multiaddr) = seen_multiaddr {
-            self.seen_multiaddrs.insert(multiaddr);
-        }
     }
 
     #[cfg(test)]
@@ -487,6 +491,8 @@ pub enum ConnectionDirection {
 pub enum PeerConnectionStatus {
     /// The peer is connected.
     Connected {
+        /// The multiaddr that we are connected via.
+        multiaddr: Multiaddr,
         /// number of ingoing connections.
         n_in: u8,
         /// number of outgoing connections.
@@ -522,7 +528,12 @@ impl Serialize for PeerConnectionStatus {
     fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         let mut s = serializer.serialize_struct("connection_status", 6)?;
         match self {
-            Connected { n_in, n_out } => {
+            Connected {
+                n_in,
+                n_out,
+                multiaddr,
+            } => {
+                s.serialize_field("multiaddr", multiaddr)?;
                 s.serialize_field("status", "connected")?;
                 s.serialize_field("connections_in", n_in)?;
                 s.serialize_field("connections_out", n_out)?;