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

alibi_sqlx/store/
sessions.rs
