greentic_session/
inmemory.rs

1use crate::error::SessionResult;
2use crate::error::{invalid_argument, not_found};
3use crate::store::SessionStore;
4use greentic_types::{EnvId, SessionData, SessionKey, TeamId, TenantCtx, TenantId, UserId};
5use parking_lot::RwLock;
6use std::collections::HashMap;
7use uuid::Uuid;
8
9/// Simple in-memory implementation backed by hash maps.
10pub struct InMemorySessionStore {
11    sessions: RwLock<HashMap<SessionKey, SessionData>>,
12    user_index: RwLock<HashMap<UserLookupKey, SessionKey>>,
13}
14
15impl Default for InMemorySessionStore {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl InMemorySessionStore {
22    /// Constructs an empty store.
23    pub fn new() -> Self {
24        Self {
25            sessions: RwLock::new(HashMap::new()),
26            user_index: RwLock::new(HashMap::new()),
27        }
28    }
29
30    fn next_key() -> SessionKey {
31        SessionKey::new(Uuid::new_v4().to_string())
32    }
33
34    fn ensure_alignment(ctx: &TenantCtx, data: &SessionData) -> SessionResult<()> {
35        if ctx.env != data.tenant_ctx.env || ctx.tenant_id != data.tenant_ctx.tenant_id {
36            return Err(invalid_argument(
37                "session data tenant context does not match provided TenantCtx",
38            ));
39        }
40        Ok(())
41    }
42
43    fn lookup_from_ctx(ctx: &TenantCtx) -> Option<UserLookupKey> {
44        let user = ctx.user_id.clone().or_else(|| ctx.user.clone())?;
45        Some(UserLookupKey::from_ctx(ctx, &user))
46    }
47
48    fn lookup_from_data(data: &SessionData) -> Option<UserLookupKey> {
49        let user = data
50            .tenant_ctx
51            .user_id
52            .clone()
53            .or_else(|| data.tenant_ctx.user.clone())?;
54        Some(UserLookupKey::from_ctx(&data.tenant_ctx, &user))
55    }
56
57    fn record_user_mapping(
58        &self,
59        ctx_hint: Option<&TenantCtx>,
60        data: &SessionData,
61        key: &SessionKey,
62    ) {
63        let lookup =
64            Self::lookup_from_data(data).or_else(|| ctx_hint.and_then(Self::lookup_from_ctx));
65        if let Some(entry) = lookup {
66            self.user_index.write().insert(entry, key.clone());
67        }
68    }
69
70    fn purge_user_mapping(&self, data: &SessionData, key: &SessionKey) {
71        if let Some(entry) = Self::lookup_from_data(data) {
72            let mut guard = self.user_index.write();
73            if guard
74                .get(&entry)
75                .map(|existing| existing == key)
76                .unwrap_or(false)
77            {
78                guard.remove(&entry);
79            }
80        }
81    }
82}
83
84impl SessionStore for InMemorySessionStore {
85    fn create_session(&self, ctx: &TenantCtx, data: SessionData) -> SessionResult<SessionKey> {
86        Self::ensure_alignment(ctx, &data)?;
87        let key = Self::next_key();
88        self.sessions.write().insert(key.clone(), data.clone());
89        self.record_user_mapping(Some(ctx), &data, &key);
90        Ok(key)
91    }
92
93    fn get_session(&self, key: &SessionKey) -> SessionResult<Option<SessionData>> {
94        Ok(self.sessions.read().get(key).cloned())
95    }
96
97    fn update_session(&self, key: &SessionKey, data: SessionData) -> SessionResult<()> {
98        let previous = self.sessions.write().insert(key.clone(), data.clone());
99        let Some(old) = previous else {
100            return Err(not_found(key));
101        };
102        self.purge_user_mapping(&old, key);
103        self.record_user_mapping(None, &data, key);
104        Ok(())
105    }
106
107    fn remove_session(&self, key: &SessionKey) -> SessionResult<()> {
108        if let Some(old) = self.sessions.write().remove(key) {
109            self.purge_user_mapping(&old, key);
110            Ok(())
111        } else {
112            Err(not_found(key))
113        }
114    }
115
116    fn find_by_user(
117        &self,
118        ctx: &TenantCtx,
119        user: &UserId,
120    ) -> SessionResult<Option<(SessionKey, SessionData)>> {
121        let lookup = UserLookupKey::from_ctx(ctx, user);
122        if let Some(stored_key) = self.user_index.read().get(&lookup).cloned() {
123            if let Some(data) = self.sessions.read().get(&stored_key).cloned() {
124                return Ok(Some((stored_key, data)));
125            }
126            self.user_index.write().remove(&lookup);
127        }
128        Ok(None)
129    }
130}
131
132#[derive(Clone, Eq, PartialEq, Hash)]
133struct UserLookupKey {
134    env: EnvId,
135    tenant: TenantId,
136    team: Option<TeamId>,
137    user: UserId,
138}
139
140impl UserLookupKey {
141    fn from_ctx(ctx: &TenantCtx, user: &UserId) -> Self {
142        Self {
143            env: ctx.env.clone(),
144            tenant: ctx.tenant_id.clone(),
145            team: ctx.team_id.clone().or_else(|| ctx.team.clone()),
146            user: user.clone(),
147        }
148    }
149}