Suggestions to simplify Session management

src/config.rs

@@ -1,9 +1,13 @@
 //! A set of configuration parameters to tune the discovery protocol.
 use crate::{
-    handler::MIN_SESSIONS_UNREACHABLE_ENR, kbucket::MAX_NODES_PER_BUCKET, Enr, Executor,
-    ListenConfig, PermitBanList, RateLimiter, RateLimiterBuilder,
+    kbucket::MAX_NODES_PER_BUCKET, Enr, Executor, ListenConfig, PermitBanList, RateLimiter,
+    RateLimiterBuilder,
 };
 
+/// The minimum number of unreachable Sessions a node must allow. This enables the network to
+/// boostrap.
+const MIN_SESSIONS_UNREACHABLE_ENR: usize = 10;
+
 use std::{ops::RangeInclusive, time::Duration};
 
 /// Configuration parameters that define the performance of the discovery network.
@@ -100,7 +104,7 @@ pub struct Discv5Config {
     /// The max limit for peers with unreachable ENRs. Benevolent examples of such peers are peers
     /// that are discovering their externally reachable socket, nodes must assist at least one
     /// such peer in discovering their reachable socket via ip voting, and peers behind symmetric
-    /// NAT. Default is no limit. Minimum is 1.
+    /// NAT. Default is no limit. Minimum is 10.
     pub unreachable_enr_limit: Option<usize>,
 
     /// The unused port range to try and bind to when testing if this node is behind NAT based on

src/handler/nat_hole_punch/mod.rs

@@ -11,7 +11,7 @@ mod error;
 mod utils;
 
 pub use error::Error;
-pub use utils::NatHolePunchUtils;
+pub use utils::NatUtils;
 
 #[async_trait::async_trait]
 pub trait HolePunchNat {

src/handler/nat_hole_punch/utils.rs

@@ -1,18 +1,15 @@
 use std::{
     net::{IpAddr, SocketAddr, UdpSocket},
     ops::RangeInclusive,
-    pin::Pin,
-    task::{Context, Poll},
     time::Duration,
 };
 
-use derive_more::{Deref, DerefMut};
+use delay_map::HashSetDelay;
 use enr::NodeId;
-use futures::{channel::mpsc, Stream, StreamExt};
 use lru::LruCache;
 use rand::Rng;
 
-use crate::{lru_time_cache::LruTimeCache, node_info::NodeAddress, Enr, IpMode};
+use crate::{node_info::NodeAddress, Enr, IpMode};
 
 /// The expected shortest lifetime in most NAT configurations of a punched hole in seconds.
 pub const DEFAULT_HOLE_PUNCH_LIFETIME: u64 = 20;
@@ -22,7 +19,7 @@ pub const PORT_BIND_TRIES: usize = 4;
 pub const USER_AND_DYNAMIC_PORTS: RangeInclusive<u16> = 1025..=u16::MAX;
 
 /// Aggregates types necessary to implement nat hole punching for [`crate::handler::Handler`].
-pub struct NatHolePunchUtils {
+pub struct NatUtils {
     /// Ip mode as set in config.
     pub ip_mode: IpMode,
     /// This node has been observed to be behind a NAT.
@@ -36,29 +33,33 @@ pub struct NatHolePunchUtils {
     pub new_peer_latest_relay_cache: LruCache<NodeId, NodeAddress>,
     /// Keeps track if this node needs to send a packet to a peer in order to keep a hole punched
     /// for it in its NAT.
-    hole_punch_tracker: NatHolePunchTracker,
+    pub hole_punch_tracker: HashSetDelay<SocketAddr>,
     /// Ports to trie to bind to check if this node is behind NAT.
     pub unused_port_range: Option<RangeInclusive<u16>>,
     /// If the filter is enabled this sets the default timeout for bans enacted by the filter.
     pub ban_duration: Option<Duration>,
+    /// The number of unreachable ENRs we store at most in our session cache.
+    pub unreachable_enr_limit: Option<usize>,
 }
 
-impl NatHolePunchUtils {
-    pub fn new<'a>(
-        listen_sockets: impl Iterator<Item = &'a SocketAddr>,
+impl NatUtils {
+    pub fn new(
+        listen_sockets: &[SocketAddr],
         local_enr: &Enr,
         ip_mode: IpMode,
         unused_port_range: Option<RangeInclusive<u16>>,
         ban_duration: Option<Duration>,
         session_cache_capacity: usize,
+        unreachable_enr_limit: Option<usize>,
     ) -> Self {
-        let mut nat_hole_puncher = NatHolePunchUtils {
+        let mut nat_hole_puncher = NatUtils {
             ip_mode,
             is_behind_nat: None,
             new_peer_latest_relay_cache: LruCache::new(session_cache_capacity),
-            hole_punch_tracker: NatHolePunchTracker::new(session_cache_capacity),
+            hole_punch_tracker: HashSetDelay::new(Duration::from_secs(DEFAULT_HOLE_PUNCH_LIFETIME)),
             unused_port_range,
             ban_duration,
+            unreachable_enr_limit,
         };
         // Optimistically only test one advertised socket, ipv4 has precedence. If it is
         // reachable, assumption is made that also the other ip version socket is reachable.
@@ -86,7 +87,7 @@ impl NatHolePunchUtils {
         if self.is_behind_nat == Some(false) {
             return;
         }
-        self.hole_punch_tracker.insert(peer_socket, ());
+        self.hole_punch_tracker.insert(peer_socket);
     }
 
     pub fn untrack(&mut self, peer_socket: &SocketAddr) {
@@ -95,13 +96,16 @@ impl NatHolePunchUtils {
 
     /// Called when a new observed address is reported at start up or after a
     /// [`crate::Discv5Event::SocketUpdated`].
-    pub fn set_is_behind_nat<'a>(
+    pub fn set_is_behind_nat(
         &mut self,
-        mut listen_sockets: impl Iterator<Item = &'a SocketAddr>,
+        listen_sockets: &[SocketAddr],
         observed_ip: Option<IpAddr>,
         observed_port: Option<u16>,
     ) {
-        if !listen_sockets.any(|listen_socket| Some(listen_socket.port()) == observed_port) {
+        if !listen_sockets
+            .iter()
+            .any(|listen_socket| Some(listen_socket.port()) == observed_port)
+        {
             self.is_behind_nat = Some(true);
             return;
         }
@@ -121,40 +125,10 @@ impl NatHolePunchUtils {
             }
         });
     }
-}
-
-impl Stream for NatHolePunchUtils {
-    type Item = SocketAddr;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        // Until ip voting is done and an observed public address is finalised, all nodes act as
-        // if they are behind a NAT.
-        if self.is_behind_nat == Some(false) || self.hole_punch_tracker.len() == 0 {
-            return Poll::Pending;
-        }
-        self.hole_punch_tracker.expired_entries.poll_next_unpin(cx)
-    }
-}
-
-#[derive(Deref, DerefMut)]
-struct NatHolePunchTracker {
-    #[deref]
-    #[deref_mut]
-    cache: LruTimeCache<SocketAddr, ()>,
-    expired_entries: mpsc::Receiver<SocketAddr>,
-}
-
-impl NatHolePunchTracker {
-    fn new(session_cache_capacity: usize) -> Self {
-        let (tx, rx) = futures::channel::mpsc::channel::<SocketAddr>(session_cache_capacity);
-        Self {
-            cache: LruTimeCache::new_with_expiry_feedback(
-                Duration::from_secs(DEFAULT_HOLE_PUNCH_LIFETIME),
-                Some(session_cache_capacity),
-                tx,
-            ),
-            expired_entries: rx,
-        }
+    /// Determines if an ENR is reachable or not based on its assigned keys.
+    pub fn is_enr_reachable(enr: &Enr) -> bool {
+        enr.udp4_socket().is_some() || enr.udp6_socket().is_some()
     }
 }
 

src/handler/sessions/session.rs → src/handler/session.rs

@@ -10,17 +10,12 @@ use crate::{
     Discv5Error, Enr,
 };
 
-// If the message nonce length is ever set below 4 bytes this will explode. The packet
-// size constants shouldn't be modified.
-const _: () = assert!(MESSAGE_NONCE_LENGTH > 4);
-
 use enr::{CombinedKey, NodeId};
 use parking_lot::RwLock;
 use std::sync::Arc;
-use tracing::warn;
 use zeroize::Zeroize;
 
-#[derive(Zeroize, PartialEq, Clone, Copy)]
+#[derive(Zeroize, PartialEq)]
 pub(crate) struct Keys {
     /// The encryption key.
     encryption_key: [u8; 16],
@@ -39,7 +34,6 @@ impl From<([u8; 16], [u8; 16])> for Keys {
 
 /// A Session containing the encryption/decryption keys. These are kept individually for a given
 /// node.
-#[derive(Clone)]
 pub(crate) struct Session {
     /// The current keys used to encrypt/decrypt messages.
     keys: Keys,
@@ -174,60 +168,23 @@ impl Session {
     pub(crate) fn establish_from_challenge(
         local_key: Arc<RwLock<CombinedKey>>,
         local_id: &NodeId,
-        challenge: Challenge,
+        remote_id: &NodeId,
+        challenge_data: ChallengeData,
         id_nonce_sig: &[u8],
         ephem_pubkey: &[u8],
-        enr_record: Option<Enr>,
-        node_address: &NodeAddress,
-        session_limiter: impl FnOnce(&NodeAddress, &Enr) -> Option<Result<(), Discv5Error>>,
+        session_enr: Enr,
     ) -> Result<(Session, Enr), Discv5Error> {
-        // check and verify a potential ENR update
-
-        let Challenge { data, remote_enr } = challenge;
-
-        // Avoid cloning an ENR
-        let session_enr = match (enr_record, remote_enr) {
-            (Some(new_enr), Some(known_enr)) => {
-                if new_enr.seq() > known_enr.seq() {
-                    MostRecentEnr::Handshake {
-                        enr: new_enr,
-                        challenge_enr: Some(known_enr),
-                    }
-                } else {
-                    MostRecentEnr::Challenge { enr: known_enr }
-                }
-            }
-            (Some(new_enr), None) => MostRecentEnr::Handshake {
-                enr: new_enr,
-                challenge_enr: None,
-            },
-            (None, Some(known_enr)) => MostRecentEnr::Challenge { enr: known_enr },
-            (None, None) => {
-                warn!(
-                        "Peer did not respond with their ENR. Session could not be established. Node: {}",
-                        node_address.node_id
-                    );
-                return Err(Discv5Error::SessionNotEstablished);
-            }
-        };
-
-        // Avoid unnecessary key derivation computation, first verify session candidate against
-        // current sessions state.
-        session_limiter(node_address, session_enr.enr()).transpose()?;
-
-        let remote_public_key = session_enr.enr().public_key();
-
         // verify the auth header nonce
         if !crypto::verify_authentication_nonce(
-            &remote_public_key,
+            &session_enr.public_key(),
             ephem_pubkey,
-            &data,
+            &challenge_data,
             local_id,
             id_nonce_sig,
         ) {
             let challenge = Challenge {
-                data,
-                remote_enr: session_enr.into_challenge_enr(),
+                data: challenge_data,
+                remote_enr: Some(session_enr),
             };
             return Err(Discv5Error::InvalidChallengeSignature(challenge));
         }
@@ -239,8 +196,8 @@ impl Session {
         let (decryption_key, encryption_key) = crypto::derive_keys_from_pubkey(
             &local_key.read(),
             local_id,
-            &node_address.node_id,
-            &data,
+            remote_id,
+            &challenge_data,
             ephem_pubkey,
         )?;
 
@@ -249,7 +206,6 @@ impl Session {
             decryption_key,
         };
 
-        let session_enr = session_enr.into_most_recent_enr();
         Ok((Session::new(keys), session_enr))
     }
 
@@ -307,41 +263,6 @@ impl Session {
     }
 }
 
-enum MostRecentEnr {
-    Challenge {
-        enr: Enr,
-    },
-    Handshake {
-        enr: Enr,
-        challenge_enr: Option<Enr>,
-    },
-}
-
-impl MostRecentEnr {
-    fn enr(&self) -> &Enr {
-        match self {
-            Self::Challenge { enr } => enr,
-            Self::Handshake { enr, .. } => enr,
-        }
-    }
-
-    fn into_most_recent_enr(self) -> Enr {
-        match self {
-            MostRecentEnr::Challenge { enr } => enr,
-            MostRecentEnr::Handshake { enr, .. } => enr,
-        }
-    }
-    fn into_challenge_enr(self) -> Option<Enr> {
-        match self {
-            MostRecentEnr::Challenge { enr } => Some(enr),
-            MostRecentEnr::Handshake {
-                enr: _,
-                challenge_enr,
-            } => challenge_enr,
-        }
-    }
-}
-
 #[cfg(test)]
 pub(crate) fn build_dummy_session() -> Session {
     Session::new(Keys {