Skip to main content

alibi_sqlx/store/
sessions.rs

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