Remove stale mid-registrations

Cargo.lock

@@ -3961,6 +3961,7 @@ dependencies = [
  "clap 4.5.4",
  "fastrand 2.1.0",
  "futures",
+ "ipnetwork 0.16.0",
  "log",
  "nym-authenticator-requests",
  "nym-bin-common",
@@ -3981,6 +3982,7 @@ dependencies = [
  "serde_json",
  "thiserror",
  "tokio",
+ "tokio-stream",
  "tokio-util",
  "url",
 ]

common/wireguard-types/src/registration.rs

@@ -7,6 +7,7 @@ use base64::{engine::general_purpose, Engine};
 use dashmap::DashMap;
 use serde::{Deserialize, Serialize};
 use std::net::IpAddr;
+use std::time::SystemTime;
 use std::{fmt, ops::Deref, str::FromStr};
 
 #[cfg(feature = "verify")]
@@ -18,13 +19,13 @@ use sha2::Sha256;
 
 pub type GatewayClientRegistry = DashMap<PeerPublicKey, GatewayClient>;
 pub type PendingRegistrations = DashMap<PeerPublicKey, RegistrationData>;
-pub type PrivateIPs = DashMap<IpAddr, Free>;
+pub type PrivateIPs = DashMap<IpAddr, Taken>;
 
 #[cfg(feature = "verify")]
 pub type HmacSha256 = Hmac<Sha256>;
 
 pub type Nonce = u64;
-pub type Free = bool;
+pub type Taken = Option<SystemTime>;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(tag = "type", rename_all = "camelCase")]

service-providers/authenticator/Cargo.toml

@@ -16,12 +16,14 @@ bytes = { workspace = true }
 clap = { workspace = true }
 fastrand = { workspace = true }
 futures = { workspace = true }
+ipnetwork = { workspace = true }
 log = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "net"] }
+tokio-stream = { workspace = true }
 tokio-util = { workspace = true, features = ["codec"] }
 url = { workspace = true }
 

service-providers/authenticator/src/authenticator.rs

@@ -4,6 +4,7 @@
 use std::path::Path;
 
 use futures::channel::oneshot;
+use ipnetwork::IpNetwork;
 use nym_client_core::{HardcodedTopologyProvider, TopologyProvider};
 use nym_sdk::mixnet::Recipient;
 use nym_task::{TaskClient, TaskHandle};
@@ -102,8 +103,13 @@ impl Authenticator {
 
         let self_address = *mixnet_client.nym_address();
 
+        let private_ip_network = IpNetwork::new(
+            self.config.authenticator.private_ip,
+            self.config.authenticator.private_network_prefix,
+        )?;
         let mixnet_listener = crate::mixnet_listener::MixnetListener::new(
             self.config,
+            private_ip_network,
             self.wireguard_gateway_data,
             mixnet_client,
             task_handle,

service-providers/authenticator/src/error.rs

@@ -1,6 +1,7 @@
 // Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
+use ipnetwork::IpNetworkError;
 use nym_client_core::error::ClientCoreError;
 use nym_id::NymIdError;
 
@@ -49,6 +50,9 @@ pub enum AuthenticatorError {
     #[error("I/O error: {0}")]
     IoError(#[from] std::io::Error),
 
+    #[error("{0}")]
+    IpNetworkError(#[from] IpNetworkError),
+
     #[error("mac does not verify")]
     MacVerificationFailure,
 
@@ -60,6 +64,9 @@ pub enum AuthenticatorError {
 
     #[error("registration is not in progress for the given key")]
     RegistrationNotInProgress,
+
+    #[error("internal data corruption: {0}")]
+    InternalDataCorruption(String),
 }
 
 pub type Result<T> = std::result::Result<T, AuthenticatorError>;

service-providers/authenticator/src/mixnet_listener.rs

@@ -1,10 +1,14 @@
 // Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use std::sync::Arc;
+use std::{
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
 
 use crate::error::AuthenticatorError;
 use futures::StreamExt;
+use ipnetwork::IpNetwork;
 use nym_authenticator_requests::v1::{
     self,
     request::{AuthenticatorRequest, AuthenticatorRequestData},
@@ -19,10 +23,12 @@ use nym_wireguard_types::{
     GatewayClient, InitMessage, PeerPublicKey,
 };
 use rand::{prelude::IteratorRandom, thread_rng};
+use tokio_stream::wrappers::IntervalStream;
 
 use crate::{config::Config, error::*};
 
 type AuthenticatorHandleResult = Result<AuthenticatorResponse>;
+const DEFAULT_REGISTRATION_TIMEOUT_CHECK: Duration = Duration::from_secs(60); // 1 minute
 
 pub(crate) struct MixnetListener {
     // The configuration for the mixnet listener
@@ -40,22 +46,30 @@ pub(crate) struct MixnetListener {
     pub(crate) wireguard_gateway_data: WireguardGatewayData,
 
     pub(crate) free_private_network_ips: Arc<PrivateIPs>,
+
+    pub(crate) timeout_check_interval: IntervalStream,
 }
 
 impl MixnetListener {
     pub fn new(
         config: Config,
+        private_ip_network: IpNetwork,
         wireguard_gateway_data: WireguardGatewayData,
         mixnet_client: nym_sdk::mixnet::MixnetClient,
         task_handle: TaskHandle,
     ) -> Self {
+        let timeout_check_interval =
+            IntervalStream::new(tokio::time::interval(DEFAULT_REGISTRATION_TIMEOUT_CHECK));
         MixnetListener {
             config,
             mixnet_client,
             task_handle,
             registration_in_progres: Default::default(),
             wireguard_gateway_data,
-            free_private_network_ips: Default::default(),
+            free_private_network_ips: Arc::new(
+                private_ip_network.iter().map(|ip| (ip, None)).collect(),
+            ),
+            timeout_check_interval,
         }
     }
 
@@ -75,6 +89,38 @@ impl MixnetListener {
         Ok(())
     }
 
+    fn remove_stale_registrations(&self) -> Result<()> {
+        for reg in self.registration_in_progres.iter().map(|reg| reg.clone()) {
+            let mut ip = self
+                .free_private_network_ips
+                .get_mut(&reg.gateway_data.private_ip)
+                .ok_or(AuthenticatorError::InternalDataCorruption(format!(
+                    "IP {} should be present",
+                    reg.gateway_data.private_ip
+                )))?;
+
+            let timestamp = ip.ok_or(AuthenticatorError::InternalDataCorruption(format!(
+                "timestamp should be set for IP {}",
+                ip.key()
+            )))?;
+            let duration = SystemTime::now().duration_since(timestamp).map_err(|_| {
+                AuthenticatorError::InternalDataCorruption(format!(
+                    "set timestamp shouldn't have been set in the future"
+                ))
+            })?;
+            if duration > DEFAULT_REGISTRATION_TIMEOUT_CHECK {
+                *ip = None;
+                self.registration_in_progres
+                    .remove(&reg.gateway_data.pub_key());
+                log::debug!(
+                    "Removed stale registration of {}",
+                    reg.gateway_data.pub_key()
+                );
+            }
+        }
+        Ok(())
+    }
+
     fn on_initial_request(
         &mut self,
         init_message: InitMessage,
@@ -99,7 +145,7 @@ impl MixnetListener {
                 .ok_or(AuthenticatorError::InternalError(String::from(
                     "could not find private IP",
                 )))?;
-            *private_ip_ref = true;
+            *private_ip_ref = None;
             Some(gateway_client.clone())
         } else {
             None
@@ -110,11 +156,11 @@ impl MixnetListener {
         let mut private_ip_ref = self
             .free_private_network_ips
             .iter_mut()
-            .filter(|r| **r)
+            .filter(|r| r.is_none())
             .choose(&mut thread_rng())
             .ok_or(AuthenticatorError::NoFreeIp)?;
         // mark it as used, even though it's not final
-        *private_ip_ref = false;
+        *private_ip_ref = Some(SystemTime::now());
         let gateway_data = GatewayClient::new(
             self.wireguard_gateway_data.keypair().private_key(),
             remote_public.inner(),
@@ -231,6 +277,11 @@ impl MixnetListener {
                 _ = task_client.recv() => {
                     log::debug!("Authenticator [main loop]: received shutdown");
                 },
+                _ = self.timeout_check_interval.next() => {
+                    if let Err(e) = self.remove_stale_registrations() {
+                        log::error!("Could not clear stale registrations. The registration process might get jammed soon - {:?}", e);
+                    }
+                }
                 msg = self.mixnet_client.next() => {
                     if let Some(msg) = msg {
                         match self.on_reconstructed_message(msg).await {