Enables ENR auto-update based on new listen address (#610)

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.

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()
     }

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> {