1use bytes::Bytes;
4use sova_store::{BoxFuture, KvStore};
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7use std::time::Duration;
8
9pub const SESSION_USER_KEY: &str = "_sova_uid";
11
12pub trait SessionStore: Send + Sync + 'static {
14 fn load<'a>(&'a self, id: &'a str) -> BoxFuture<'a, Option<HashMap<String, String>>>;
15 fn save<'a>(
16 &'a self,
17 id: &'a str,
18 data: &'a HashMap<String, String>,
19 ttl: Duration,
20 ) -> BoxFuture<'a, ()>;
21 fn destroy<'a>(&'a self, id: &'a str) -> BoxFuture<'a, ()>;
22 fn destroy_user<'a>(
24 &'a self,
25 user_id: &'a str,
26 keep_sid: Option<&'a str>,
27 ) -> BoxFuture<'a, u64>;
28}
29
30#[derive(Clone)]
32pub struct SessionStoreHandle(pub Arc<dyn SessionStore>);
33
34impl SessionStoreHandle {
35 pub fn inner(&self) -> &Arc<dyn SessionStore> {
36 &self.0
37 }
38}
39
40pub(crate) fn encode(data: &HashMap<String, String>) -> Bytes {
41 let mut out = String::new();
42 for (k, v) in data {
43 out.push_str(k);
44 out.push('\0');
45 out.push_str(v);
46 out.push('\n');
47 }
48 Bytes::from(out)
49}
50
51pub(crate) fn decode(bytes: &[u8]) -> HashMap<String, String> {
52 let mut map = HashMap::new();
53 for line in std::str::from_utf8(bytes).unwrap_or("").lines() {
54 if let Some((k, v)) = line.split_once('\0') {
55 map.insert(k.to_string(), v.to_string());
56 }
57 }
58 map
59}
60
61fn index_key(user_id: &str) -> String {
62 format!("uids:{user_id}")
63}
64
65fn decode_set(raw: &[u8]) -> HashSet<String> {
66 serde_json::from_slice(raw).unwrap_or_default()
67}
68
69fn encode_set(set: &HashSet<String>) -> Bytes {
70 Bytes::from(serde_json::to_vec(set).unwrap_or_else(|_| b"[]".to_vec()))
71}
72
73pub struct KvSessionStore {
77 kv: Arc<dyn KvStore>,
78}
79
80impl KvSessionStore {
81 pub fn new(kv: Arc<dyn KvStore>) -> Self {
82 Self { kv }
83 }
84
85 async fn track(&self, user_id: &str, sid: &str, ttl: Duration) {
86 if user_id.is_empty() || sid.is_empty() {
87 return;
88 }
89 let key = index_key(user_id);
90 let mut set = self
91 .kv
92 .get(&key)
93 .await
94 .map(|b| decode_set(&b))
95 .unwrap_or_default();
96 set.insert(sid.to_string());
97 self.kv.set(&key, encode_set(&set), Some(ttl)).await;
98 }
99
100 async fn untrack(&self, user_id: &str, sid: &str, ttl: Duration) {
101 if user_id.is_empty() || sid.is_empty() {
102 return;
103 }
104 let key = index_key(user_id);
105 let Some(raw) = self.kv.get(&key).await else {
106 return;
107 };
108 let mut set = decode_set(&raw);
109 if set.remove(sid) {
110 if set.is_empty() {
111 self.kv.remove(&key).await;
112 } else {
113 self.kv.set(&key, encode_set(&set), Some(ttl)).await;
114 }
115 }
116 }
117}
118
119impl SessionStore for KvSessionStore {
120 fn load<'a>(&'a self, id: &'a str) -> BoxFuture<'a, Option<HashMap<String, String>>> {
121 Box::pin(async move { self.kv.get(id).await.map(|b| decode(&b)) })
122 }
123
124 fn save<'a>(
125 &'a self,
126 id: &'a str,
127 data: &'a HashMap<String, String>,
128 ttl: Duration,
129 ) -> BoxFuture<'a, ()> {
130 Box::pin(async move {
131 if let Some(prev) = self.kv.get(id).await.map(|b| decode(&b)) {
132 let old_uid = prev.get(SESSION_USER_KEY).map(String::as_str);
133 let new_uid = data.get(SESSION_USER_KEY).map(String::as_str);
134 if old_uid != new_uid {
135 if let Some(u) = old_uid {
136 self.untrack(u, id, ttl).await;
137 }
138 }
139 }
140 self.kv.set(id, encode(data), Some(ttl)).await;
141 if let Some(uid) = data.get(SESSION_USER_KEY).filter(|s| !s.is_empty()) {
142 self.track(uid, id, ttl).await;
143 }
144 })
145 }
146
147 fn destroy<'a>(&'a self, id: &'a str) -> BoxFuture<'a, ()> {
148 Box::pin(async move {
149 if let Some(data) = self.kv.get(id).await.map(|b| decode(&b)) {
150 if let Some(uid) = data.get(SESSION_USER_KEY).filter(|s| !s.is_empty()) {
151 self.untrack(uid, id, Duration::from_secs(86400)).await;
153 }
154 }
155 self.kv.remove(id).await;
156 })
157 }
158
159 fn destroy_user<'a>(
160 &'a self,
161 user_id: &'a str,
162 keep_sid: Option<&'a str>,
163 ) -> BoxFuture<'a, u64> {
164 Box::pin(async move {
165 if user_id.is_empty() {
166 return 0;
167 }
168 let key = index_key(user_id);
169 let Some(raw) = self.kv.get(&key).await else {
170 return 0;
171 };
172 let set = decode_set(&raw);
173 let mut kept = HashSet::new();
174 let mut n = 0u64;
175 for sid in set {
176 if keep_sid.is_some_and(|k| k == sid) {
177 kept.insert(sid);
178 continue;
179 }
180 self.kv.remove(&sid).await;
181 n += 1;
182 }
183 if kept.is_empty() {
184 self.kv.remove(&key).await;
185 } else {
186 self.kv
187 .set(&key, encode_set(&kept), Some(Duration::from_secs(86400)))
188 .await;
189 }
190 n
191 })
192 }
193}