snap_control/pg_wap/
session_manager.rs

1// Copyright 2026 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Session manager for the PathGuard WAP API, responsible for managing client authentication and
15//! TCP sessions.
16
17use std::{
18    collections::HashMap,
19    net::IpAddr,
20    sync::{Arc, Mutex},
21    time::Instant,
22};
23
24use crate::pg_wap::auth::AuthInfo;
25
26/// Manages client authentication and TCP sessions.
27#[derive(Clone)]
28pub struct WapSessionManager {
29    inner: Arc<Mutex<WapSessionManagerInner>>,
30}
31
32impl WapSessionManager {
33    /// Authenticate an IP address to be able to open sessions.
34    ///
35    /// The authenticated session is appended to the set of authenticated sessions for the given IP
36    /// address.
37    pub fn add_session_authentication(&self, auth_info: AuthInfo) {
38        let mut inner = self.inner.lock().unwrap();
39        inner.add_session_authentication(auth_info);
40    }
41
42    /// Find the set of authenticated sessions for the given client IP address.
43    pub fn authenticated_sessions_for_ip(&self, client_addr: IpAddr) -> Vec<AuthInfo> {
44        let mut inner = self.inner.lock().unwrap();
45        inner.authenticated_sessions_for_ip(Instant::now(), client_addr)
46    }
47}
48
49/// Internal session manager state.
50pub struct WapSessionManagerInner {
51    sessions: HashMap<IpAddr, Vec<AuthInfo>>,
52}
53
54impl WapSessionManager {
55    /// Create a new [WapSessionManager].
56    pub fn new() -> Self {
57        let inner = Arc::new(Mutex::new(WapSessionManagerInner {
58            sessions: HashMap::new(),
59        }));
60
61        Self { inner }
62    }
63}
64
65impl Default for WapSessionManager {
66    fn default() -> Self {
67        Self::new()
68    }
69}
70
71impl WapSessionManagerInner {
72    /// Authenticate a IP to be able to open sessions.
73    ///
74    /// The authenticated session is appended to the set of authenticated sessions for the given IP
75    /// address.
76    pub fn add_session_authentication(&mut self, auth_info: AuthInfo) {
77        // XXX(uniquefine): Currently the sessions map can grow unlimited e.g. due to a malicious
78        // client. We should add a limit on the maximum number of entries in the
79        // map/sessions in an auth_info.
80        self.sessions
81            .entry(auth_info.ip)
82            .or_default()
83            .push(auth_info);
84    }
85
86    /// Find the set of authenticated sessions for the given client IP address.
87    pub fn authenticated_sessions_for_ip(
88        &mut self,
89        now: Instant,
90        client_addr: IpAddr,
91    ) -> Vec<AuthInfo> {
92        let Some(auth_infos) = self.sessions.get_mut(&client_addr) else {
93            return Vec::new();
94        };
95
96        // Remove expired auths and return the remaining ones.
97        // XXX(uniquefine): We should consider adding an active cleanup of expired auths.
98        auth_infos.retain(|auth| auth.valid_until > now);
99        if auth_infos.is_empty() {
100            self.sessions.remove(&client_addr);
101            Vec::new()
102        } else {
103            auth_infos.clone()
104        }
105    }
106}
snap_control/pg_wap/session_manager.rs

snap_control/pg_wap/
session_manager.rs
