greentic_session/
inmemory.rs

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