fix: implement token expiration mechanism for attested contracts (#1976)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: nacho.d.g <iduartgomez@users.noreply.github.com>
Co-authored-by: nacho.d.g <iduartgomez@gmail.com>

Author: 4 people

apps/freenet-ping/app/tests/common/mod.rs

@@ -112,6 +112,8 @@ pub async fn base_node_test_config_with_rng(
         ws_api: WebsocketApiArgs {
             address: Some(Ipv4Addr::LOCALHOST.into()),
             ws_api_port: Some(ws_api_port),
+            token_ttl_seconds: None,
+            token_cleanup_interval_seconds: None,
         },
         network_api: NetworkArgs {
             public_address: Some(Ipv4Addr::LOCALHOST.into()),

crates/core/src/client_events/websocket.rs

@@ -1,9 +1,11 @@
 use std::{
     collections::{HashMap, VecDeque},
-    sync::{Arc, OnceLock, RwLock},
+    sync::{Arc, OnceLock},
     time::Duration,
 };
 
+use dashmap::DashMap;
+
 use axum::{
     extract::{
         ws::{Message, WebSocket},
@@ -43,7 +45,7 @@ impl std::ops::Deref for WebSocketRequest {
     }
 }
 
-pub(crate) struct WebSocketProxy {
+pub struct WebSocketProxy {
     proxy_server_request: mpsc::Receiver<ClientConnection>,
     response_channels: HashMap<ClientId, mpsc::UnboundedSender<HostCallbackResult>>,
 }
@@ -53,10 +55,7 @@ const PARALLELISM: usize = 10; // TODO: get this from config, or whatever optima
 impl WebSocketProxy {
     pub fn create_router(server_routing: Router) -> (Self, Router) {
         // Create a default empty attested contracts map
-        let attested_contracts = Arc::new(RwLock::new(HashMap::<
-            AuthToken,
-            (ContractInstanceId, ClientId),
-        >::new()));
+        let attested_contracts = Arc::new(DashMap::new());
         Self::create_router_with_attested_contracts(server_routing, attested_contracts)
     }
 
@@ -290,27 +289,27 @@ async fn websocket_commands(
     Extension(attested_contracts): Extension<AttestedContractMap>,
 ) -> Response {
     let on_upgrade = move |ws: WebSocket| async move {
-        // Get the data we need and immediately drop the lock
+        // Get the data we need from the DashMap
         let auth_and_instance = if let Some(token) = auth_token.as_ref() {
-            let attested_contracts_read = attested_contracts.read().unwrap();
-
             // Only collect and log map contents when trace is enabled
             if tracing::enabled!(tracing::Level::TRACE) {
-                let map_contents: Vec<_> = attested_contracts_read.keys().cloned().collect();
+                let map_contents: Vec<_> =
+                    attested_contracts.iter().map(|e| e.key().clone()).collect();
                 tracing::trace!(?token, "attested_contracts map keys: {:?}", map_contents);
             }
 
-            if let Some((cid, _)) = attested_contracts_read.get(token) {
-                tracing::trace!(?token, ?cid, "Found token in attested_contracts map");
-                Some((token.clone(), *cid))
+            if let Some(entry) = attested_contracts.get(token) {
+                let attested = entry.value();
+                tracing::trace!(?token, contract_id = ?attested.contract_id, "Found token in attested_contracts map");
+                Some((token.clone(), attested.contract_id))
             } else {
                 tracing::warn!(?token, "Auth token not found in attested_contracts map");
                 None
             }
         } else {
             tracing::trace!("No auth token provided in WebSocket request");
             None
-        }; // RwLockReadGuard is dropped here
+        };
 
         // Only evaluate auth_and_instance for trace when trace is enabled
         if tracing::enabled!(tracing::Level::TRACE) {

crates/core/src/config/mod.rs

@@ -101,6 +101,8 @@ impl Default for ConfigArgs {
             ws_api: WebsocketApiArgs {
                 address: Some(default_listening_address()),
                 ws_api_port: Some(default_http_gateway_port()),
+                token_ttl_seconds: None,
+                token_cleanup_interval_seconds: None,
             },
             secrets: Default::default(),
             log_level: Some(tracing::log::LevelFilter::Info),
@@ -230,6 +232,12 @@ impl ConfigArgs {
             self.mode.get_or_insert(cfg.mode);
             self.ws_api.address.get_or_insert(cfg.ws_api.address);
             self.ws_api.ws_api_port.get_or_insert(cfg.ws_api.port);
+            self.ws_api
+                .token_ttl_seconds
+                .get_or_insert(cfg.ws_api.token_ttl_seconds);
+            self.ws_api
+                .token_cleanup_interval_seconds
+                .get_or_insert(cfg.ws_api.token_cleanup_interval_seconds);
             self.log_level.get_or_insert(cfg.log_level);
             self.config_paths.merge(cfg.config_paths.as_ref().clone());
         }
@@ -361,6 +369,14 @@ impl ConfigArgs {
                     .ws_api
                     .ws_api_port
                     .unwrap_or(default_http_gateway_port()),
+                token_ttl_seconds: self
+                    .ws_api
+                    .token_ttl_seconds
+                    .unwrap_or(default_token_ttl_seconds()),
+                token_cleanup_interval_seconds: self
+                    .ws_api
+                    .token_cleanup_interval_seconds
+                    .unwrap_or(default_token_cleanup_interval_seconds()),
             },
             secrets,
             log_level: self.log_level.unwrap_or(tracing::log::LevelFilter::Info),
@@ -616,6 +632,19 @@ pub struct WebsocketApiArgs {
     #[arg(long, env = "WS_API_PORT")]
     #[serde(rename = "ws-api-port", skip_serializing_if = "Option::is_none")]
     pub ws_api_port: Option<u16>,
+
+    /// Token time-to-live in seconds (default is 86400 = 24 hours)
+    #[arg(long, env = "TOKEN_TTL_SECONDS")]
+    #[serde(rename = "token-ttl-seconds", skip_serializing_if = "Option::is_none")]
+    pub token_ttl_seconds: Option<u64>,
+
+    /// Token cleanup interval in seconds (default is 300 = 5 minutes)
+    #[arg(long, env = "TOKEN_CLEANUP_INTERVAL_SECONDS")]
+    #[serde(
+        rename = "token-cleanup-interval-seconds",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub token_cleanup_interval_seconds: Option<u64>,
 }
 
 #[derive(Debug, Copy, Clone, Serialize, Deserialize)]
@@ -627,13 +656,36 @@ pub struct WebsocketApiConfig {
     /// Port to expose api on
     #[serde(default = "default_http_gateway_port", rename = "ws-api-port")]
     pub port: u16,
+
+    /// Token time-to-live in seconds
+    #[serde(default = "default_token_ttl_seconds", rename = "token-ttl-seconds")]
+    pub token_ttl_seconds: u64,
+
+    /// Token cleanup interval in seconds
+    #[serde(
+        default = "default_token_cleanup_interval_seconds",
+        rename = "token-cleanup-interval-seconds"
+    )]
+    pub token_cleanup_interval_seconds: u64,
+}
+
+#[inline]
+const fn default_token_ttl_seconds() -> u64 {
+    86400 // 24 hours
+}
+
+#[inline]
+const fn default_token_cleanup_interval_seconds() -> u64 {
+    300 // 5 minutes
 }
 
 impl From<SocketAddr> for WebsocketApiConfig {
     fn from(addr: SocketAddr) -> Self {
         Self {
             address: addr.ip(),
             port: addr.port(),
+            token_ttl_seconds: default_token_ttl_seconds(),
+            token_cleanup_interval_seconds: default_token_cleanup_interval_seconds(),
         }
     }
 }
@@ -644,6 +696,8 @@ impl Default for WebsocketApiConfig {
         Self {
             address: default_listening_address(),
             port: default_http_gateway_port(),
+            token_ttl_seconds: default_token_ttl_seconds(),
+            token_cleanup_interval_seconds: default_token_cleanup_interval_seconds(),
         }
     }
 }

crates/core/src/lib.rs

@@ -60,7 +60,7 @@ pub mod dev_tool {
     use super::*;
     pub use crate::config::Config;
     pub use client_events::{
-        test::MemoryEventsGen, test::NetworkEventGenerator, ClientEventsProxy, ClientId,
+        test::MemoryEventsGen, test::NetworkEventGenerator, AuthToken, ClientEventsProxy, ClientId,
         OpenRequest,
     };
     pub use contract::{storages::Storage, Executor, OperationMode};

crates/core/src/node/mod.rs

@@ -1334,9 +1334,8 @@ pub async fn run_local_node(
             ClientRequest::DelegateOp(op) => {
                 let attested_contract = token.and_then(|token| {
                     gw.attested_contracts
-                        .read()
-                        .ok()
-                        .and_then(|guard| guard.get(&token).map(|(t, _)| *t))
+                        .get(&token)
+                        .map(|entry| entry.value().contract_id)
                 });
                 let op_name = match op {
                     DelegateRequest::RegisterDelegate { .. } => "RegisterDelegate",
@@ -1356,17 +1355,6 @@ pub async fn run_local_node(
                 if let Some(cause) = cause {
                     tracing::info!("disconnecting cause: {cause}");
                 }
-                // FIXME: We're not removing tokens on disconnect to allow WebSocket connections
-                // to use them for authentication. We should implement a proper token expiration
-                // mechanism instead of keeping them forever or removing them immediately.
-                // if let Ok(mut guard) = gw.attested_contracts.write() {
-                //     if let Some(rm_token) = guard
-                //         .iter()
-                //         .find_map(|(k, (_, eid))| (eid == &id).then(|| k.clone()))
-                //     {
-                //         guard.remove(&rm_token);
-                //     }
-                // }
                 continue;
             }
             _ => Err(ExecutorError::other(anyhow::anyhow!("not supported"))),

crates/core/src/server/http_gateway.rs

@@ -1,6 +1,9 @@
 use std::collections::HashMap;
 use std::net::{IpAddr, SocketAddr};
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
+use std::time::Instant;
+
+use dashmap::DashMap;
 
 use axum::extract::Path;
 use axum::response::IntoResponse;
@@ -31,19 +34,42 @@ impl std::ops::Deref for HttpGatewayRequest {
     }
 }
 
-pub type AttestedContractMap = Arc<RwLock<HashMap<AuthToken, (ContractInstanceId, ClientId)>>>;
+/// Represents an attested contract entry with metadata for token expiration.
+#[derive(Clone, Debug)]
+pub struct AttestedContract {
+    /// The contract instance ID
+    pub contract_id: ContractInstanceId,
+    /// The client ID associated with this token
+    pub client_id: ClientId,
+    /// Timestamp of when the token was last accessed (for expiration tracking)
+    pub last_accessed: Instant,
+}
+
+impl AttestedContract {
+    /// Create a new attested contract entry
+    pub fn new(contract_id: ContractInstanceId, client_id: ClientId) -> Self {
+        Self {
+            contract_id,
+            client_id,
+            last_accessed: Instant::now(),
+        }
+    }
+}
+
+/// Maps authentication tokens to attested contract metadata.
+pub type AttestedContractMap = Arc<DashMap<AuthToken, AttestedContract>>;
 
 /// A gateway to access and interact with contracts through an HTTP interface.
-pub(crate) struct HttpGateway {
-    pub attested_contracts: AttestedContractMap,
+pub struct HttpGateway {
+    pub(crate) attested_contracts: AttestedContractMap,
     proxy_server_request: mpsc::Receiver<ClientConnection>,
     response_channels: HashMap<ClientId, mpsc::UnboundedSender<HostCallbackResult>>,
 }
 
 impl HttpGateway {
     /// Returns the uninitialized axum router to compose with other routing handling or websockets.
     pub fn as_router(socket: &SocketAddr) -> (Self, Router) {
-        let attested_contracts = Arc::new(RwLock::new(HashMap::new()));
+        let attested_contracts = Arc::new(DashMap::new());
         Self::as_router_with_attested_contracts(socket, attested_contracts)
     }
 
@@ -54,6 +80,12 @@ impl HttpGateway {
     ) -> (Self, Router) {
         Self::create_router_v1_with_attested_contracts(socket, attested_contracts)
     }
+
+    /// Returns a reference to the attested contracts map (for integration testing).
+    /// This allows tests to verify token expiration behavior.
+    pub fn attested_contracts(&self) -> &AttestedContractMap {
+        &self.attested_contracts
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -81,10 +113,9 @@ impl ClientEventsProxy for HttpGateway {
                             .send(HostCallbackResult::NewId { id: cli_id })
                             .map_err(|_e| ErrorKind::NodeUnavailable)?;
                         if let Some((assigned_token, contract)) = assigned_token {
+                            let attested = AttestedContract::new(contract, cli_id);
                             self.attested_contracts
-                                .write()
-                                .map_err(|_| ErrorKind::FailedOperation)?
-                                .insert(assigned_token.clone(), (contract, cli_id));
+                                .insert(assigned_token.clone(), attested);
                             tracing::debug!(
                                 ?assigned_token,
                                 ?contract,