diff --git a/proxy_agent/src/shared_state/proxy_server_wrapper.rs b/proxy_agent/src/shared_state/proxy_server_wrapper.rs index e9ef02a2..45a620f3 100644 --- a/proxy_agent/src/shared_state/proxy_server_wrapper.rs +++ b/proxy_agent/src/shared_state/proxy_server_wrapper.rs @@ -21,8 +21,16 @@ use crate::common::logger; use crate::common::result::Result; use crate::proxy::User; use std::collections::HashMap; +use std::time::{Duration, Instant}; use tokio::sync::{mpsc, oneshot}; +const USER_CACHE_TTL: Duration = Duration::from_secs(5 * 60); + +struct CachedUser { + user: User, + cached_at: Instant, +} + enum ProxyServerAction { AddUser { user: User, @@ -46,20 +54,37 @@ pub struct ProxyServerSharedState(mpsc::Sender); impl ProxyServerSharedState { pub fn start_new() -> Self { + Self::start_new_with_cache_ttl(USER_CACHE_TTL) + } + + fn start_new_with_cache_ttl(cache_ttl: Duration) -> Self { let (tx, mut rx) = mpsc::channel(100); tokio::spawn(async move { - let mut users: HashMap = HashMap::new(); + let mut users: HashMap = HashMap::new(); while let Some(action) = rx.recv().await { match action { ProxyServerAction::AddUser { user, response } => { let id = user.logon_id; - users.insert(id, user); + users.insert( + id, + CachedUser { + user, + cached_at: Instant::now(), + }, + ); if response.send(()).is_err() { logger::write_warning(format!("Failed to send response to ProxyServerAction::AddUser with id '{id}'")); } } ProxyServerAction::GetUser { user_id, response } => { - let user = users.get(&user_id).cloned(); + // only return the user if it is still valid (not expired) + let user = users + .get(&user_id) + .filter(|cached_user| cached_user.cached_at.elapsed() < cache_ttl) + .map(|cached_user| cached_user.user.clone()); + if user.is_none() { + users.remove(&user_id); + } if response.send(user).is_err() { logger::write_warning(format!("Failed to send response to ProxyServerAction::GetUser with id '{user_id}'")); } @@ -131,7 +156,6 @@ impl ProxyServerSharedState { .map_err(|e| Error::RecvError("ProxyServerAction::GetUsersCount".to_string(), e)) } - // TODO:: need caller to refresh the users info regularly pub async fn clear_users(&self) -> Result<()> { let (tx, rx) = oneshot::channel(); self.0 @@ -144,3 +168,43 @@ impl ProxyServerSharedState { .map_err(|e| Error::RecvError("ProxyServerAction::ClearUsers".to_string(), e)) } } + +#[cfg(test)] +mod tests { + use super::ProxyServerSharedState; + use crate::proxy::User; + use std::time::Duration; + + #[tokio::test] + async fn expired_user_is_invalidated() { + let state = ProxyServerSharedState::start_new_with_cache_ttl(Duration::ZERO); + state + .add_user(User { + logon_id: 42, + user_name: "user".to_string(), + user_groups: vec!["old-group".to_string()], + }) + .await + .unwrap(); + + assert!(state.get_user(42).await.unwrap().is_none()); + assert_eq!(state.get_users_count().await.unwrap(), 0); + } + + #[tokio::test] + async fn unexpired_user_is_returned() { + let state = ProxyServerSharedState::start_new_with_cache_ttl(Duration::from_secs(60)); + state + .add_user(User { + logon_id: 42, + user_name: "user".to_string(), + user_groups: vec!["current-group".to_string()], + }) + .await + .unwrap(); + + let user = state.get_user(42).await.unwrap().unwrap(); + assert_eq!(user.user_groups, vec!["current-group"]); + assert_eq!(state.get_users_count().await.unwrap(), 1); + } +}