tiny_web/sys/
session.rs

1use std::{
2    collections::{btree_map::Entry, BTreeMap},
3    sync::{Arc, LazyLock},
4};
5
6use chrono::Local;
7use ring::rand::{SecureRandom, SystemRandom};
8use sha3::{Digest, Sha3_512};
9use tiny_web_macro::fnv1a_64;
10use tokio::sync::Mutex;
11
12use crate::{fnv1a_64, StrOrI64};
13use tiny_web_macro::fnv1a_64 as m_fnv1a_64;
14
15use super::{data::Data, dbs::adapter::DB, request::Request};
16
17/// Temporary cache for install mode only
18static INSTALL_CACHE: LazyLock<Mutex<BTreeMap<i64, Data>>> = LazyLock::new(|| Mutex::new(BTreeMap::new()));
19
20/// Types of flash messages
21#[repr(i16)]
22#[derive(Debug)]
23pub enum Flash {
24    Info = 1,
25    Success = 2,
26    Warning = 3,
27    Error = 4,
28}
29
30impl From<i16> for Flash {
31    fn from(value: i16) -> Self {
32        match value {
33            1 => Flash::Info,
34            2 => Flash::Success,
35            3 => Flash::Warning,
36            4 => Flash::Error,
37            #[cfg(not(debug_assertions))]
38            _ => Flash::Error,
39            #[cfg(debug_assertions)]
40            _ => panic!("Invalid value for Status"),
41        }
42    }
43}
44
45impl From<Flash> for i16 {
46    fn from(value: Flash) -> Self {
47        value as i16
48    }
49}
50
51/// User session
52///
53///  # Values
54///
55/// * `id: i64` - session_id from database.
56/// * `lang_id: i64` - lang_id from database.
57/// * `user_id: i64` - user_id from database.
58/// * `role_id: i64` - role_id from database.
59/// * `pub key: String` - Cookie key.
60/// * `data: HashMap<String, Data>` - User data from database.
61/// * `change: bool` - User data is changed.
62#[derive(Debug)]
63pub struct Session {
64    /// session_id from database
65    id: i64,
66    /// Default lang_id for user
67    lang_id: i64,
68    /// user_id from database
69    pub user_id: i64,
70    /// role_id from database
71    pub role_id: i64,
72    /// Cookie key (session value)
73    pub key: String,
74    /// Session key
75    pub session_key: Arc<String>,
76    /// User data from database
77    data: BTreeMap<i64, Data>,
78    /// User data is changed
79    change: bool,
80}
81
82impl Session {
83    /// Create new session
84    pub(crate) fn new(lang_id: i64, salt: &str, ip: &str, agent: &str, host: &str, session_key: Arc<String>) -> Session {
85        Session {
86            id: 0,
87            lang_id,
88            user_id: 0,
89            role_id: 0,
90            key: Session::generate_session(salt, ip, agent, host),
91            session_key,
92            data: BTreeMap::new(),
93            change: false,
94        }
95    }
96
97    /// Create new session by cookie (session) key
98    pub(crate) fn with_key(lang_id: i64, key: String, session_key: Arc<String>) -> Session {
99        Session {
100            id: 0,
101            lang_id,
102            user_id: 0,
103            role_id: 0,
104            key,
105            session_key,
106            data: BTreeMap::new(),
107            change: false,
108        }
109    }
110
111    /// Load session from database
112    pub(crate) async fn load_session(key: String, db: Arc<DB>, lang_id: i64, session_key: Arc<String>) -> Session {
113        if db.in_use() {
114            let res = match db.query_prepare(fnv1a_64!("lib_get_session"), &[&key], false).await {
115                Some(r) => r,
116                None => return Session::with_key(lang_id, key, session_key),
117            };
118            if res.is_empty() {
119                return Session::with_key(lang_id, key, session_key);
120            }
121            let row = if let Data::Vec(row) = unsafe { res.get_unchecked(0) } {
122                row
123            } else {
124                return Session::with_key(lang_id, key, session_key);
125            };
126            if row.len() != 5 {
127                return Session::with_key(lang_id, key, session_key);
128            }
129            let session_id = if let Data::I64(val) = unsafe { row.get_unchecked(0) } {
130                *val
131            } else {
132                return Session::with_key(lang_id, key, session_key);
133            };
134            let user_id = if let Data::I64(val) = unsafe { row.get_unchecked(1) } {
135                *val
136            } else {
137                return Session::with_key(lang_id, key, session_key);
138            };
139            let role_id = if let Data::I64(val) = unsafe { row.get_unchecked(2) } {
140                *val
141            } else {
142                return Session::with_key(lang_id, key, session_key);
143            };
144            let data = if let Data::Raw(val) = unsafe { row.get_unchecked(3) } {
145                val.to_owned()
146            } else {
147                return Session::with_key(lang_id, key, session_key);
148            };
149            let lang_id = if let Data::I64(val) = unsafe { row.get_unchecked(4) } {
150                *val
151            } else {
152                return Session::with_key(lang_id, key, session_key);
153            };
154
155            let res = if data.is_empty() {
156                BTreeMap::new()
157            } else {
158                bincode::deserialize::<BTreeMap<i64, Data>>(&data[..]).unwrap_or_else(|_| BTreeMap::new())
159            };
160            Session {
161                id: session_id,
162                lang_id,
163                user_id,
164                role_id,
165                key,
166                session_key,
167                data: res,
168                change: false,
169            }
170        } else {
171            let cache = INSTALL_CACHE.lock().await;
172            match cache.get(&fnv1a_64(key.as_bytes())) {
173                Some(map) => {
174                    let data: BTreeMap<i64, Data> = map.clone().into();
175                    let lang_id = if let Some(Data::I64(lang)) = data.get(&m_fnv1a_64!("lang_id")) {
176                        *lang
177                    } else {
178                        return Session::with_key(lang_id, key, session_key);
179                    };
180                    let data = if let Some(Data::Map(data)) = data.get(&m_fnv1a_64!("data")) {
181                        data.clone()
182                    } else {
183                        return Session::with_key(lang_id, key, session_key);
184                    };
185                    Session {
186                        id: 0,
187                        lang_id,
188                        user_id: 0,
189                        role_id: 0,
190                        key,
191                        session_key,
192                        data,
193                        change: false,
194                    }
195                }
196                None => Session::with_key(lang_id, key, session_key),
197            }
198        }
199    }
200
201    /// Save session into database
202    pub(crate) async fn save_session(db: Arc<DB>, session: &Session, request: &Request) {
203        if session.change {
204            if db.in_use() {
205                let data = bincode::serialize(&session.data).unwrap_or_else(|_| Vec::new());
206                if session.id > 0 {
207                    db.execute_prepare(
208                        fnv1a_64!("lib_set_session"),
209                        &[&session.user_id, &session.lang_id, &data, &request.ip, &request.agent, &session.id],
210                    )
211                    .await;
212                } else {
213                    db.execute_prepare(
214                        fnv1a_64!("lib_add_session"),
215                        &[&session.user_id, &session.lang_id, &session.key, &data, &request.ip, &request.agent],
216                    )
217                    .await;
218                };
219            } else {
220                let mut cache = INSTALL_CACHE.lock().await;
221                let mut data = BTreeMap::new();
222                data.insert(m_fnv1a_64!("lang_id"), session.lang_id.into());
223                data.insert(m_fnv1a_64!("data"), session.data.clone().into());
224
225                cache.insert(fnv1a_64(session.key.as_bytes()), data.into());
226            }
227        }
228    }
229
230    /// Generete new session key
231    fn generate_session(salt: &str, ip: &str, agent: &str, host: &str) -> String {
232        // Generate a new cookie
233        let time = Local::now().format("%Y.%m.%d %H:%M:%S%.9f %:z").to_string();
234        let cook = format!("{}{}{}{}{}", salt, ip, agent, host, time);
235        let mut hasher = Sha3_512::new();
236        hasher.update(cook.as_bytes());
237        format!("{:#x}", hasher.finalize())
238    }
239
240    pub fn generate_salt(&self) -> String {
241        let time = Local::now().format("%Y.%m.%d %H:%M:%S%.9f %:z").to_string();
242        let cook = format!("{}{}", self.key, time);
243        let mut hasher = Sha3_512::new();
244        hasher.update(cook.as_bytes());
245        let rand = format!("!@#$^&*()_-+=,<.>/?|{:#x}", hasher.finalize());
246        self.shuffle_string(&rand)
247    }
248
249    fn shuffle_string(&self, s: &str) -> String {
250        let mut chars: Vec<char> = s.chars().collect();
251        let len = chars.len();
252        let rng = SystemRandom::new();
253
254        for i in (1..len).rev() {
255            let mut buf = [0u8; 8];
256            let _ = rng.fill(&mut buf);
257            let rand_index = (u64::from_ne_bytes(buf) % (i as u64 + 1)) as usize;
258            chars.swap(i, rand_index);
259        }
260        for c in chars.iter_mut() {
261            let mut buf = [0u8; 1];
262            let _ = rng.fill(&mut buf);
263            if buf[0] % 2 == 0 {
264                *c = c.to_ascii_uppercase();
265            }
266        }
267        let mut str: String = chars.into_iter().collect();
268        str.truncate(32);
269        str
270    }
271
272    /// Set lang_id
273    pub fn set_lang_id(&mut self, lang_id: i64) {
274        if self.lang_id != lang_id {
275            self.lang_id = lang_id;
276            self.change = true;
277        }
278    }
279
280    /// Get lang_id
281    pub fn get_lang_id(&self) -> i64 {
282        self.lang_id
283    }
284
285    /// Set session data
286    pub fn set<T>(&mut self, key: impl StrOrI64, value: T)
287    where
288        T: Into<Data>,
289    {
290        self.change = true;
291        self.data.insert(key.to_i64(), value.into());
292    }
293
294    /// Get session data for reading
295    pub fn get(&mut self, key: impl StrOrI64) -> Option<&Data> {
296        self.change = true;
297        self.data.get(&key.to_i64())
298    }
299
300    /// Getting session data by deleting it
301    pub fn take(&mut self, key: impl StrOrI64) -> Option<Data> {
302        self.change = true;
303        self.data.remove(&key.to_i64())
304    }
305
306    /// Remove session data
307    pub fn remove(&mut self, key: impl StrOrI64) {
308        self.change = true;
309        self.data.remove(&key.to_i64());
310    }
311
312    /// Clear session data
313    pub fn clear(&mut self) {
314        self.change = true;
315        self.data.clear();
316    }
317
318    /// Set flash message to session data
319    pub(crate) fn set_flash(&mut self, kind: Flash, value: String) {
320        self.change = true;
321        match self.data.entry(fnv1a_64!("flash-message")) {
322            Entry::Vacant(entry) => {
323                entry.insert(vec![Data::I16(kind.into()), Data::String(value)].into());
324            }
325            Entry::Occupied(mut entry) => {
326                let e = entry.get_mut();
327                if let Data::Vec(vec) = e {
328                    vec.push(Data::I16(kind.into()));
329                    vec.push(Data::String(value));
330                } else {
331                    *e = vec![Data::I16(kind.into()), Data::String(value)].into();
332                }
333            }
334        }
335    }
336
337    /// Take flash message from session data
338    pub(crate) fn take_flash(&mut self) -> Option<Vec<(Flash, String)>> {
339        self.change = true;
340        if let Data::Vec(vec) = self.data.remove(&fnv1a_64!("flash-message"))? {
341            if vec.len() % 2 != 0 {
342                return None;
343            }
344            let mut result = Vec::with_capacity(vec.len() / 2);
345            let mut iter = vec.into_iter();
346            while let (Some(first), Some(second)) = (iter.next(), iter.next()) {
347                match (first, second) {
348                    (Data::I16(num), Data::String(text)) => result.push((num.into(), text)),
349                    _ => return None,
350                }
351            }
352
353            Some(result)
354        } else {
355            None
356        }
357    }
358}
tiny_web/sys/session.rs

tiny_web/sys/
session.rs
