greentic_session/
inmemory.rs1use 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
9pub 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 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}