Skip to main content

alibi_sqlx/store/
sessions.rs

1use super::SqlxStore;
2use crate::error::record_not_updated;
3use crate::model::{self, SqlxModel};
4use crate::pool::{Exec, SqlxTransaction};
5use crate::schema::{AuthSchema, SqlxSessionModel};
6use crate::sql::Sql;
7use alibi_core::error::{AuthError, AuthResult};
8use alibi_core::store::SessionStore;
9use alibi_core::store::adapter::cancelled_by_hook;
10use alibi_core::types::CreateSession;
11use async_trait::async_trait;
12use chrono::{DateTime, Utc};
13
14impl<S> SqlxStore<S>
15where
16    S: AuthSchema,
17    S::Session: SqlxSessionModel,
18{
19    fn normalize_session_client_field(value: Option<String>) -> Option<String> {
20        value.map_or_else(|| Some(String::new()), Some)
21    }
22
23    /// `SELECT ... WHERE token = ? AND active = true LIMIT 1`.
24    fn active_session_by_token(exec: Exec<'_>, token: &str) -> Sql {
25        let table = <S::Session as SqlxModel>::TABLE;
26        let mut sql = model::select_model::<S::Session>(exec);
27        sql.push(" WHERE ");
28        sql.compare_model::<S::Session>(table, S::Session::token_column(), " = ", token);
29        sql.push(" AND ");
30        sql.compare_model::<S::Session>(table, S::Session::active_column(), " = ", true);
31        sql.push(" LIMIT 1");
32        sql
33    }
34
35    async fn create_session_with_connection(
36        &self,
37        exec: Exec<'_>,
38        tx: Option<&SqlxTransaction>,
39        mut create_session: CreateSession,
40        persist: bool,
41        complete: bool,
42    ) -> AuthResult<S::Session> {
43        if create_session.token.is_none() {
44            create_session.token = Some(alibi_core::utils::sessions::generate_session_token());
45        }
46        let hook_context = self.hook_context(tx);
47        for hook in self.hooks() {
48            if hook
49                .before_create_session(&mut create_session, &hook_context)
50                .await
51                .map_err(alibi_core::store::adapter::callback_error)?
52                .is_cancelled()
53            {
54                return Err(AuthError::SessionCreationCancelled);
55            }
56        }
57        let now = Utc::now();
58        create_session.ip_address = Self::normalize_session_client_field(create_session.ip_address);
59        create_session.user_agent = Self::normalize_session_client_field(create_session.user_agent);
60        let token = create_session
61            .token
62            .take()
63            .unwrap_or_else(alibi_core::utils::sessions::generate_session_token);
64        let mut fields = std::mem::take(&mut create_session.additional_fields);
65        let mut typed_fields = [
66            (
67                "activeOrganizationId",
68                &mut create_session.active_organization_id,
69            ),
70            ("activeTeamId", &mut create_session.active_team_id),
71            ("impersonatedBy", &mut create_session.impersonated_by),
72        ];
73        for (name, destination) in &mut typed_fields {
74            if let Some(value) = destination.as_ref() {
75                fields.preserve_creation_value(
76                    name,
77                    alibi_core::utils::json::JsValue::String(value.clone()),
78                );
79            }
80            // Configured values now belong to the adapter input. A transform
81            // that omits one must also omit its original typed creation value.
82            if fields.contains_key(*name) {
83                **destination = None;
84            }
85        }
86        if persist {
87            fields.apply_adapter_transforms_async().await?;
88        }
89        // These fields are persisted by new_active, including on handwritten
90        // schemas with renamed columns and no generic additional-field bindings.
91        for (name, destination) in typed_fields {
92            if let Some(value) = fields.shift_remove(name) {
93                *destination = crate::additional_fields::prepare_string_value(
94                    exec,
95                    crate::additional_fields::raw_value(&value)?,
96                )
97                .await?;
98            }
99        }
100        let mut active = S::Session::new_active(None, token, create_session, now);
101        if !fields.is_empty() {
102            for (column, value) in S::Session::additional_field_bindings(&fields, exec.engine())? {
103                let value = crate::additional_fields::prepare_value(
104                    exec,
105                    <S::Session as SqlxModel>::column_kind(column),
106                    value,
107                )
108                .await?;
109                S::Session::set_additional_field(&mut active, column, value, exec.engine())?;
110            }
111        }
112        let session = if persist {
113            model::insert::<S::Session>(exec, &active).await?
114        } else {
115            S::Session::materialize_secondary(active)?
116        };
117        if tx.is_none() && complete {
118            for hook in self.hooks() {
119                hook.after_create_session(&session, &hook_context)
120                    .await
121                    .map_err(alibi_core::store::adapter::callback_error)?;
122            }
123        }
124        Ok(session)
125    }
126
127    pub(crate) async fn prepare_secondary_update_with_connection(
128        &self,
129        exec: Exec<'_>,
130        tx: Option<&SqlxTransaction>,
131        session: S::Session,
132        expires_at: Option<DateTime<Utc>>,
133        mut fields: alibi_core::field_policy::FieldValues,
134    ) -> AuthResult<Option<(S::Session, alibi_core::field_policy::FieldValues)>> {
135        use alibi_core::AuthSession;
136        let hook_context = self.hook_context(tx);
137        for hook in self.hooks() {
138            if hook
139                .before_update_session(session.token(), &mut fields, &hook_context)
140                .await?
141                .is_cancelled()
142            {
143                return Ok(None);
144            }
145        }
146        // The cache stores hook output before SQL adapter input transformations.
147        let mut active = session.into_active();
148        let backend = exec.engine();
149        for (column, value) in S::Session::additional_field_bindings(&fields, backend)? {
150            let value = crate::additional_fields::prepare_value(
151                exec,
152                <S::Session as SqlxModel>::column_kind(column),
153                value,
154            )
155            .await?;
156            S::Session::set_additional_field(&mut active, column, value, backend)?;
157        }
158        if let Some(expiry) = expires_at {
159            S::Session::set_expires_at(&mut active, expiry);
160        }
161        S::Session::set_updated_at(&mut active, Utc::now());
162        let session = S::Session::materialize_secondary(active)?;
163        Ok(Some((session, fields)))
164    }
165
166    pub(crate) async fn complete_secondary_update_with_connection(
167        &self,
168        exec: Exec<'_>,
169        tx: Option<&SqlxTransaction>,
170        session: S::Session,
171        expires_at: Option<DateTime<Utc>>,
172        mut fields: alibi_core::field_policy::FieldValues,
173        persist: bool,
174    ) -> AuthResult<Option<S::Session>> {
175        use alibi_core::AuthSession;
176        let context = self.hook_context(tx);
177        let session = if persist {
178            fields.apply_adapter_transforms_async().await?;
179            let Some(current) = exec
180                .fetch_optional::<S::Session>(Self::active_session_by_token(exec, session.token()))
181                .await?
182            else {
183                for hook in self.hooks().iter().filter(|_| tx.is_none()) {
184                    hook.after_update_session_missing(session.token(), &context)
185                        .await?;
186                }
187                return Ok(None);
188            };
189            let mut active = current.into_active();
190            let backend = exec.engine();
191            for (column, value) in S::Session::additional_field_bindings(&fields, backend)? {
192                let value = crate::additional_fields::prepare_value(
193                    exec,
194                    <S::Session as SqlxModel>::column_kind(column),
195                    value,
196                )
197                .await?;
198                S::Session::set_additional_field(&mut active, column, value, backend)?;
199            }
200            if let Some(expiry) = expires_at {
201                S::Session::set_expires_at(&mut active, expiry);
202            }
203            S::Session::set_updated_at(&mut active, session.updated_at());
204            if let Some(model) = model::update::<S::Session>(exec, &active).await? {
205                model
206            } else {
207                for hook in self.hooks().iter().filter(|_| tx.is_none()) {
208                    hook.after_update_session_missing(session.token(), &context)
209                        .await?;
210                }
211                return Ok(None);
212            }
213        } else {
214            session
215        };
216        for hook in self.hooks().iter().filter(|_| tx.is_none()) {
217            hook.after_update_session(&session, &context).await?;
218        }
219        Ok(Some(session))
220    }
221
222    pub(crate) async fn prepare_secondary_session_in_tx(
223        &self,
224        tx: &SqlxTransaction,
225        input: CreateSession,
226        persist: bool,
227    ) -> AuthResult<S::Session> {
228        self.create_session_with_connection(Exec::Tx(tx), Some(tx), input, persist, false)
229            .await
230    }
231
232    pub(crate) async fn create_session_in_tx(
233        &self,
234        tx: &SqlxTransaction,
235        create_session: CreateSession,
236    ) -> AuthResult<S::Session> {
237        self.create_session_with_connection(Exec::Tx(tx), Some(tx), create_session, true, true)
238            .await
239    }
240}
241
242pub(super) enum SessionScope<'a> {
243    Team(Option<&'a str>),
244    Organization(Option<&'a str>),
245}
246
247impl<S> SqlxStore<S>
248where
249    S: AuthSchema,
250    S::Session: SqlxSessionModel,
251{
252    async fn update_session_with_fields(
253        &self,
254        token: &str,
255        expires_at: Option<DateTime<Utc>>,
256        mut fields: alibi_core::field_policy::FieldValues,
257    ) -> AuthResult<Option<S::Session>> {
258        let hook_context = self.hook_context(None);
259        for hook in self.hooks() {
260            if hook
261                .before_update_session(token, &mut fields, &hook_context)
262                .await?
263                .is_cancelled()
264            {
265                return Ok(None);
266            }
267        }
268        fields.apply_adapter_transforms_async().await?;
269        let table = <S::Session as SqlxModel>::TABLE;
270        let mut query = model::select_model::<S::Session>(self.exec());
271        query.push(" WHERE ");
272        query.compare_model::<S::Session>(table, S::Session::token_column(), " = ", token);
273        if expires_at.is_some() {
274            query.push(" AND ");
275            query.compare_model::<S::Session>(table, S::Session::active_column(), " = ", true);
276        }
277        query.push(" LIMIT 1");
278        let Some(model) = self.exec().fetch_optional::<S::Session>(query).await? else {
279            for hook in self.hooks() {
280                hook.after_update_session_missing(token, &hook_context)
281                    .await?;
282            }
283            return Ok(None);
284        };
285        let mut active = model.into_active();
286        let backend = self.exec().engine();
287        if !fields.is_empty() {
288            for (column, value) in S::Session::additional_field_bindings(&fields, backend)? {
289                let value = crate::additional_fields::prepare_value(
290                    self.exec(),
291                    <S::Session as SqlxModel>::column_kind(column),
292                    value,
293                )
294                .await?;
295                S::Session::set_additional_field(&mut active, column, value, backend)?;
296            }
297        }
298        if let Some(expires_at) = expires_at {
299            S::Session::set_expires_at(&mut active, expires_at);
300        }
301        S::Session::set_updated_at(&mut active, Utc::now());
302        let Some(session) = model::update::<S::Session>(self.exec(), &active).await? else {
303            for hook in self.hooks() {
304                hook.after_update_session_missing(token, &hook_context)
305                    .await?;
306            }
307            return Ok(None);
308        };
309        for hook in self.hooks() {
310            hook.after_update_session(&session, &hook_context).await?;
311        }
312        Ok(Some(session))
313    }
314
315    pub(super) async fn update_session_scope_with_connection(
316        &self,
317        exec: Exec<'_>,
318        token: &str,
319        scope: SessionScope<'_>,
320    ) -> AuthResult<S::Session> {
321        let model = exec
322            .fetch_optional::<S::Session>(Self::active_session_by_token(exec, token))
323            .await?
324            .ok_or(AuthError::SessionNotFound)?;
325        let mut active = model.into_active();
326        match scope {
327            SessionScope::Team(team) => {
328                S::Session::set_active_team_id(&mut active, team.map(str::to_owned))?;
329            }
330            SessionScope::Organization(organization) => {
331                S::Session::set_active_organization_id(
332                    &mut active,
333                    organization.map(str::to_owned),
334                );
335            }
336        }
337        S::Session::set_updated_at(&mut active, Utc::now());
338        model::update::<S::Session>(exec, &active)
339            .await?
340            .ok_or_else(record_not_updated)
341    }
342}
343
344#[async_trait]
345impl<S> SessionStore<S> for SqlxStore<S>
346where
347    S: AuthSchema + Send + Sync,
348    S::Session: SqlxSessionModel,
349{
350    async fn prepare_secondary_session_creation(
351        &self,
352        input: CreateSession,
353        persist: bool,
354    ) -> AuthResult<S::Session> {
355        self.create_session_with_connection(self.exec(), None, input, persist, false)
356            .await
357    }
358    async fn complete_secondary_session_creation(&self, session: &S::Session) -> AuthResult<()> {
359        let context = self.hook_context(None);
360        for hook in self.hooks() {
361            hook.after_create_session(session, &context)
362                .await
363                .map_err(alibi_core::store::adapter::callback_error)?;
364        }
365        Ok(())
366    }
367
368    async fn prepare_secondary_session_update(
369        &self,
370        session: S::Session,
371        expires_at: Option<DateTime<Utc>>,
372        fields: alibi_core::field_policy::FieldValues,
373    ) -> AuthResult<Option<(S::Session, alibi_core::field_policy::FieldValues)>> {
374        self.prepare_secondary_update_with_connection(
375            self.exec(),
376            None,
377            session,
378            expires_at,
379            fields,
380        )
381        .await
382    }
383    async fn complete_secondary_session_update(
384        &self,
385        session: S::Session,
386        expires_at: Option<DateTime<Utc>>,
387        fields: alibi_core::field_policy::FieldValues,
388        persist: bool,
389    ) -> AuthResult<Option<S::Session>> {
390        self.complete_secondary_update_with_connection(
391            self.exec(),
392            None,
393            session,
394            expires_at,
395            fields,
396            persist,
397        )
398        .await
399    }
400
401    async fn end_session_preserving(&self, token: &str) -> AuthResult<()> {
402        use alibi_core::AuthSession;
403        let now = Utc::now();
404        let Some(session) = self
405            .get_session(token)
406            .await?
407            .filter(|session| session.expires_at() > now)
408        else {
409            return Ok(());
410        };
411        let context = self.hook_context(None);
412        for hook in self.hooks() {
413            if hook
414                .before_delete_session(&session, &context)
415                .await?
416                .is_cancelled()
417            {
418                return Err(cancelled_by_hook("session deletion"));
419            }
420        }
421        let table = <S::Session as SqlxModel>::TABLE;
422        let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
423        sql.ident(table);
424        sql.push(" SET ");
425        sql.assign(
426            S::Session::expires_at_column(),
427            S::Session::timestamp_value(S::Session::expires_at_column(), now),
428        );
429        sql.push(" WHERE ");
430        sql.compare_model::<S::Session>(table, S::Session::token_column(), " = ", token);
431        sql.push(" AND ");
432        sql.compare_model::<S::Session>(
433            table,
434            S::Session::expires_at_column(),
435            " > ",
436            S::Session::timestamp_value(S::Session::expires_at_column(), now),
437        );
438        _ = self.exec().execute(sql).await?;
439        for hook in self.hooks() {
440            hook.after_delete_session(&session, &context).await?;
441        }
442        Ok(())
443    }
444    async fn end_user_sessions_preserving(&self, user_id: &str) -> AuthResult<()> {
445        let now = Utc::now();
446        let user_id = S::Session::parse_user_id(user_id)?;
447        let table = <S::Session as SqlxModel>::TABLE;
448        let mut live = model::select_model::<S::Session>(self.exec());
449        live.push(" WHERE ");
450        live.compare_model::<S::Session>(
451            table,
452            S::Session::user_id_column(),
453            " = ",
454            user_id.clone(),
455        );
456        live.push(" AND ");
457        live.compare_model::<S::Session>(
458            table,
459            S::Session::expires_at_column(),
460            " > ",
461            S::Session::timestamp_value(S::Session::expires_at_column(), now),
462        );
463        live.push(" LIMIT ");
464        live.bind(self.find_many_limit());
465        let sessions: Vec<S::Session> = self.exec().fetch_all(live).await?;
466        let context = self.hook_context(None);
467        for session in &sessions {
468            for hook in self.hooks() {
469                if hook
470                    .before_delete_session(session, &context)
471                    .await?
472                    .is_cancelled()
473                {
474                    return Err(cancelled_by_hook("session deletion"));
475                }
476            }
477        }
478        let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
479        sql.ident(table);
480        sql.push(" SET ");
481        sql.assign(
482            S::Session::expires_at_column(),
483            S::Session::timestamp_value(S::Session::expires_at_column(), now),
484        );
485        sql.push(" WHERE ");
486        sql.compare_model::<S::Session>(table, S::Session::user_id_column(), " = ", user_id);
487        sql.push(" AND ");
488        sql.compare_model::<S::Session>(
489            table,
490            S::Session::expires_at_column(),
491            " > ",
492            S::Session::timestamp_value(S::Session::expires_at_column(), now),
493        );
494        _ = self.exec().execute(sql).await?;
495        for session in &sessions {
496            for hook in self.hooks() {
497                hook.after_delete_session(session, &context).await?;
498            }
499        }
500        Ok(())
501    }
502
503    async fn create_session(&self, create_session: CreateSession) -> AuthResult<S::Session> {
504        self.create_session_with_connection(self.exec(), None, create_session, true, true)
505            .await
506    }
507
508    async fn get_session(&self, token: &str) -> AuthResult<Option<S::Session>> {
509        self.exec()
510            .fetch_optional(Self::active_session_by_token(self.exec(), token))
511            .await
512    }
513
514    async fn get_sessions_by_tokens(&self, tokens: &[String]) -> AuthResult<Vec<S::Session>> {
515        let table = <S::Session as SqlxModel>::TABLE;
516        let mut sql = model::select_model::<S::Session>(self.exec());
517        sql.push(" WHERE ");
518        if tokens.is_empty() {
519            sql.push("1 = 2");
520        } else {
521            sql.column(table, S::Session::token_column());
522            sql.push(" IN ");
523            sql.bind_list(
524                tokens.iter().cloned().map(|token| {
525                    S::Session::column_value(S::Session::token_column(), token.into())
526                }),
527            );
528        }
529        sql.push(" AND ");
530        sql.compare_model::<S::Session>(table, S::Session::active_column(), " = ", true);
531        sql.push(" LIMIT ");
532        sql.bind(self.find_many_limit());
533        self.exec().fetch_all(sql).await
534    }
535
536    async fn get_user_sessions(&self, user_id: &str) -> AuthResult<Vec<S::Session>> {
537        let user_id = <S::Session as SqlxSessionModel>::parse_user_id(user_id)?;
538        let table = <S::Session as SqlxModel>::TABLE;
539        let mut sql = model::select_model::<S::Session>(self.exec());
540        sql.push(" WHERE ");
541        sql.compare_model::<S::Session>(table, S::Session::user_id_column(), " = ", user_id);
542        sql.push(" AND ");
543        sql.compare_model::<S::Session>(table, S::Session::active_column(), " = ", true);
544        sql.push(" ORDER BY ");
545        sql.column(table, S::Session::created_at_column());
546        sql.push(" ASC");
547        self.exec().fetch_all(sql).await
548    }
549
550    async fn update_session_fields(
551        &self,
552        token: &str,
553        fields: alibi_core::field_policy::FieldValues,
554    ) -> AuthResult<Option<S::Session>> {
555        self.update_session_with_fields(token, None, fields).await
556    }
557
558    async fn update_session_expiry(
559        &self,
560        token: &str,
561        expires_at: DateTime<Utc>,
562    ) -> AuthResult<()> {
563        self.refresh_session(token, expires_at)
564            .await?
565            .map(|_| ())
566            .ok_or(AuthError::SessionNotFound)
567    }
568
569    async fn refresh_session(
570        &self,
571        token: &str,
572        expires_at: DateTime<Utc>,
573    ) -> AuthResult<Option<S::Session>> {
574        self.update_session_with_fields(token, Some(expires_at), Default::default())
575            .await
576    }
577
578    async fn refresh_session_with_fields(
579        &self,
580        token: &str,
581        expires_at: DateTime<Utc>,
582        fields: alibi_core::field_policy::FieldValues,
583    ) -> AuthResult<Option<S::Session>> {
584        self.update_session_with_fields(token, Some(expires_at), fields)
585            .await
586    }
587
588    async fn delete_session(&self, token: &str) -> AuthResult<()> {
589        let session = self.get_session(token).await?;
590        let hook_context = self.hook_context(None);
591        if let Some(session) = &session {
592            for hook in self.hooks() {
593                if hook
594                    .before_delete_session(session, &hook_context)
595                    .await?
596                    .is_cancelled()
597                {
598                    return Err(cancelled_by_hook("session deletion"));
599                }
600            }
601        }
602        let table = <S::Session as SqlxModel>::TABLE;
603        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
604        sql.ident(table);
605        sql.push(" WHERE ");
606        sql.compare_model::<S::Session>(table, S::Session::token_column(), " = ", token);
607        _ = self.exec().execute(sql).await?;
608        if let Some(session) = &session {
609            for hook in self.hooks() {
610                hook.after_delete_session(session, &hook_context).await?;
611            }
612        }
613        Ok(())
614    }
615
616    async fn delete_user_sessions(&self, user_id: &str) -> AuthResult<()> {
617        let user_id = <S::Session as SqlxSessionModel>::parse_user_id(user_id)?;
618        let table = <S::Session as SqlxModel>::TABLE;
619        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
620        sql.ident(table);
621        sql.push(" WHERE ");
622        sql.compare_model::<S::Session>(table, S::Session::user_id_column(), " = ", user_id);
623        self.exec().execute(sql).await.map(drop)
624    }
625
626    async fn delete_expired_sessions(&self) -> AuthResult<usize> {
627        let table = <S::Session as SqlxModel>::TABLE;
628        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
629        sql.ident(table);
630        sql.push(" WHERE ");
631        sql.compare_model::<S::Session>(
632            table,
633            S::Session::expires_at_column(),
634            " < ",
635            S::Session::timestamp_value(S::Session::expires_at_column(), Utc::now()),
636        );
637        sql.push(" OR ");
638        sql.compare_model::<S::Session>(table, S::Session::active_column(), " = ", false);
639        let deleted = self.exec().execute(sql).await?;
640        usize::try_from(deleted)
641            .map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
642    }
643
644    async fn update_session_active_organization(
645        &self,
646        token: &str,
647        organization_id: Option<&str>,
648    ) -> AuthResult<S::Session> {
649        let Some(model) = self.get_session(token).await? else {
650            return Err(AuthError::SessionNotFound);
651        };
652
653        let mut active = model.into_active();
654        S::Session::set_active_organization_id(&mut active, organization_id.map(str::to_owned));
655        S::Session::set_updated_at(&mut active, Utc::now());
656        model::update::<S::Session>(self.exec(), &active)
657            .await?
658            .ok_or_else(record_not_updated)
659    }
660
661    async fn update_session_active_team(
662        &self,
663        token: &str,
664        team_id: Option<&str>,
665    ) -> AuthResult<S::Session> {
666        let model = self
667            .get_session(token)
668            .await?
669            .ok_or(AuthError::SessionNotFound)?;
670        let mut active = model.into_active();
671        S::Session::set_active_team_id(&mut active, team_id.map(str::to_owned))?;
672        S::Session::set_updated_at(&mut active, Utc::now());
673        model::update::<S::Session>(self.exec(), &active)
674            .await?
675            .ok_or_else(record_not_updated)
676    }
677}