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