From 62d66f1c10acefca87425464102bb5239d67f809 Mon Sep 17 00:00:00 2001
From: Age Manning <Age@AgeManning.com>
Date: Wed, 20 Nov 2019 15:03:50 +1100
Subject: [PATCH] Enables ENR auto-update based on new listen address (#610)

---
 beacon_node/eth2-libp2p/Cargo.toml       |  4 +--
 beacon_node/eth2-libp2p/src/behaviour.rs |  5 +++
 beacon_node/eth2-libp2p/src/discovery.rs | 18 +++++++++++
 beacon_node/eth2-libp2p/src/service.rs   | 40 +++++++++++++++++++++++-
 4 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/beacon_node/eth2-libp2p/Cargo.toml b/beacon_node/eth2-libp2p/Cargo.toml
index 8982e1766..51d596b1e 100644
--- a/beacon_node/eth2-libp2p/Cargo.toml
+++ b/beacon_node/eth2-libp2p/Cargo.toml
@@ -8,8 +8,8 @@ edition = "2018"
 clap = "2.33.0"
 hex = "0.3"
 #SigP repository 
-libp2p =  { git = "https://github.com/SigP/rust-libp2p", rev = "cdd5251d29e21a01aa2ffed8cb577a37a0f9e2eb" }
-enr =  { git = "https://github.com/SigP/rust-libp2p/", rev = "cdd5251d29e21a01aa2ffed8cb577a37a0f9e2eb", features = ["serde"] }
+libp2p =  { git = "https://github.com/SigP/rust-libp2p", rev = "1295ff592a94d19f23f176712d6d04af4db6e698" }
+enr =  { git = "https://github.com/SigP/rust-libp2p/", rev = "1295ff592a94d19f23f176712d6d04af4db6e698", features = ["serde"] }
 types = { path =  "../../eth2/types" }
 serde = "1.0.102"
 serde_derive = "1.0.102"
diff --git a/beacon_node/eth2-libp2p/src/behaviour.rs b/beacon_node/eth2-libp2p/src/behaviour.rs
index d8301ad8b..6a97ac9e4 100644
--- a/beacon_node/eth2-libp2p/src/behaviour.rs
+++ b/beacon_node/eth2-libp2p/src/behaviour.rs
@@ -222,6 +222,11 @@ impl<TSubstream: AsyncRead + AsyncWrite> Behaviour<TSubstream> {
     pub fn connected_peers(&self) -> usize {
         self.discovery.connected_peers()
     }
+
+    /// Informs the discovery behaviour if a new IP/Port is set at the application layer
+    pub fn update_local_enr_socket(&mut self, socket: std::net::SocketAddr, is_tcp: bool) {
+        self.discovery.update_local_enr(socket, is_tcp);
+    }
 }
 
 /// The types of events than can be obtained from polling the behaviour.
diff --git a/beacon_node/eth2-libp2p/src/discovery.rs b/beacon_node/eth2-libp2p/src/discovery.rs
index 380914af5..8ad58300a 100644
--- a/beacon_node/eth2-libp2p/src/discovery.rs
+++ b/beacon_node/eth2-libp2p/src/discovery.rs
@@ -103,6 +103,24 @@ impl<TSubstream> Discovery<TSubstream> {
         })
     }
 
+    /// Allows the application layer to update the `ip` and `port` of the local ENR. The second
+    /// parameter defines whether the port is a TPC port. If false, this is interpreted as a UDP
+    /// port.
+    pub fn update_local_enr(&mut self, socket: std::net::SocketAddr, is_tcp: bool) {
+        // discv5 checks to see if an update is necessary before performing it, so we do not
+        // need to check here
+        if self.discovery.update_local_enr_socket(socket, is_tcp) {
+            let enr = self.discovery.local_enr();
+            info!(
+            self.log,
+            "ENR Updated";
+            "enr" => enr.to_base64(), 
+            "seq" => enr.seq(),
+            "address" => format!("{:?}", socket));
+        }
+    }
+
+    /// Return the nodes local ENR.
     pub fn local_enr(&self) -> &Enr {
         self.discovery.local_enr()
     }
diff --git a/beacon_node/eth2-libp2p/src/service.rs b/beacon_node/eth2-libp2p/src/service.rs
index 2ffafb855..b7ede9dd8 100644
--- a/beacon_node/eth2-libp2p/src/service.rs
+++ b/beacon_node/eth2-libp2p/src/service.rs
@@ -30,6 +30,8 @@ pub struct Service {
     pub swarm: Swarm<Libp2pStream, Libp2pBehaviour>,
     /// This node's PeerId.
     pub local_peer_id: PeerId,
+    /// Indicates if the listening address have been verified and compared to the expected ENR.
+    pub verified_listen_address: bool,
     /// The libp2p logger handle.
     pub log: slog::Logger,
 }
@@ -151,6 +153,7 @@ impl Service {
         Ok(Service {
             local_peer_id,
             swarm,
+            verified_listen_address: false,
             log,
         })
     }
@@ -189,7 +192,18 @@ impl Stream for Service {
                     }
                 },
                 Ok(Async::Ready(None)) => unreachable!("Swarm stream shouldn't end"),
-                Ok(Async::NotReady) => break,
+                Ok(Async::NotReady) => {
+                    // check to see if the address is different to the config. If so, update our ENR
+                    if !self.verified_listen_address {
+                        let multiaddr = Swarm::listeners(&self.swarm).next();
+                        if let Some(multiaddr) = multiaddr {
+                            if let Some(socket_addr) = multiaddr_to_socket_addr(multiaddr) {
+                                self.swarm.update_local_enr_socket(socket_addr, true);
+                            }
+                        }
+                    }
+                    break;
+                }
                 _ => break,
             }
         }
@@ -197,6 +211,30 @@ impl Stream for Service {
     }
 }
 
+/// Converts a multiaddr to a `SocketAddr` if the multiaddr has the TCP/IP form. Libp2p currently
+/// only supports TCP, so the UDP case is currently ignored.
+fn multiaddr_to_socket_addr(multiaddr: &Multiaddr) -> Option<std::net::SocketAddr> {
+    let protocols = multiaddr.iter().collect::<Vec<_>>();
+    // assume the IP protocol
+    match protocols[0] {
+        Protocol::Ip4(address) => {
+            if let Protocol::Tcp(port) = protocols[1] {
+                Some(std::net::SocketAddr::new(address.into(), port))
+            } else {
+                None
+            }
+        }
+        Protocol::Ip6(address) => {
+            if let Protocol::Tcp(port) = protocols[1] {
+                Some(std::net::SocketAddr::new(address.into(), port))
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
+
 /// The implementation supports TCP/IP, WebSockets over TCP/IP, secio as the encryption layer, and
 /// mplex or yamux as the multiplexing layer.
 fn build_transport(local_private_key: Keypair) -> Boxed<(PeerId, StreamMuxerBox), Error> {