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    /// 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<bool> {
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(false);
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(false);
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(true)
427    }
428    async fn end_user_sessions_preserving(&self, user_id: &str) -> AuthResult<bool> {
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(false);
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(true)
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_by_id(&self, id: &str) -> AuthResult<Option<S::Session>> {
493        let mut sql = model::select_model::<S::Session>(self.exec());
494        sql.push(" WHERE ");
495        sql.compare_model::<S::Session>(
496            <S::Session as SqlxModel>::TABLE,
497            S::Session::id_column(),
498            " = ",
499            S::Session::parse_id(id)?,
500        );
501        self.exec().fetch_optional(sql).await
502    }
503    async fn get_session(&self, token: &str) -> AuthResult<Option<S::Session>> {
504        self.exec()
505            .fetch_optional(Self::active_session_by_token(self.exec(), token))
506            .await
507    }
508
509    async fn get_sessions_by_tokens(&self, tokens: &[String]) -> AuthResult<Vec<S::Session>> {
510        let table = <S::Session as SqlxModel>::TABLE;
511        let mut sql = model::select_model::<S::Session>(self.exec());
512        sql.push(" WHERE ");
513        if tokens.is_empty() {
514            sql.push("1 = 2");
515        } else {
516            sql.column(table, S::Session::token_column());
517            sql.push(" IN ");
518            sql.bind_list(
519                tokens.iter().cloned().map(|token| {
520                    S::Session::column_value(S::Session::token_column(), token.into())
521                }),
522            );
523        }
524        sql.push(" LIMIT ");
525        sql.bind(self.find_many_limit());
526        self.exec().fetch_all(sql).await
527    }
528
529    async fn get_user_sessions(&self, user_id: &str) -> AuthResult<Vec<S::Session>> {
530        let user_id = <S::Session as SqlxSessionModel>::parse_user_id(user_id)?;
531        let table = <S::Session as SqlxModel>::TABLE;
532        let mut sql = model::select_model::<S::Session>(self.exec());
533        sql.push(" WHERE ");
534        sql.compare_model::<S::Session>(table, S::Session::user_id_column(), " = ", user_id);
535        sql.push(" ORDER BY ");
536        sql.column(table, S::Session::created_at_column());
537        sql.push(" ASC");
538        self.exec().fetch_all(sql).await
539    }
540
541    async fn update_session_fields(
542        &self,
543        token: &str,
544        fields: FieldValues,
545    ) -> AuthResult<Option<S::Session>> {
546        self.update_session_with_fields(token, None, fields).await
547    }
548
549    async fn update_session_expiry(
550        &self,
551        token: &str,
552        expires_at: DateTime<Utc>,
553    ) -> AuthResult<()> {
554        self.refresh_session(token, expires_at)
555            .await?
556            .map(|_| ())
557            .ok_or(AuthError::SessionNotFound)
558    }
559
560    async fn refresh_session(
561        &self,
562        token: &str,
563        expires_at: DateTime<Utc>,
564    ) -> AuthResult<Option<S::Session>> {
565        self.update_session_with_fields(token, Some(expires_at), FieldValues::default())
566            .await
567    }
568
569    async fn refresh_session_with_fields(
570        &self,
571        token: &str,
572        expires_at: DateTime<Utc>,
573        fields: FieldValues,
574    ) -> AuthResult<Option<S::Session>> {
575        self.update_session_with_fields(token, Some(expires_at), fields)
576            .await
577    }
578
579    async fn delete_session(&self, token: &str) -> AuthResult<()> {
580        let session = self.get_session(token).await?;
581        let hook_context = self.hook_context(None);
582        if let Some(session) = &session {
583            for hook in self.hooks() {
584                if hook
585                    .before_delete_session(session, &hook_context)
586                    .await?
587                    .is_cancelled()
588                {
589                    return Ok(());
590                }
591            }
592        }
593        let table = <S::Session as SqlxModel>::TABLE;
594        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
595        sql.ident(table);
596        sql.push(" WHERE ");
597        sql.compare_model::<S::Session>(table, S::Session::token_column(), " = ", token);
598        _ = self.exec().execute(sql).await?;
599        if let Some(session) = &session {
600            for hook in self.hooks() {
601                hook.after_delete_session(session, &hook_context).await?;
602            }
603        }
604        Ok(())
605    }
606
607    async fn delete_user_sessions(&self, user_id: &str) -> AuthResult<()> {
608        let user_id = <S::Session as SqlxSessionModel>::parse_user_id(user_id)?;
609        let table = <S::Session as SqlxModel>::TABLE;
610        let sessions = if self.hooks().is_empty() {
611            Vec::new()
612        } else {
613            let mut query = model::select_model::<S::Session>(self.exec());
614            query.push(" WHERE ");
615            query.compare_model::<S::Session>(
616                table,
617                S::Session::user_id_column(),
618                " = ",
619                user_id.clone(),
620            );
621            self.exec().fetch_all::<S::Session>(query).await?
622        };
623        let hook_context = self.hook_context(None);
624        for session in &sessions {
625            for hook in self.hooks() {
626                if hook
627                    .before_delete_session(session, &hook_context)
628                    .await?
629                    .is_cancelled()
630                {
631                    return Ok(());
632                }
633            }
634        }
635        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
636        sql.ident(table);
637        sql.push(" WHERE ");
638        sql.compare_model::<S::Session>(table, S::Session::user_id_column(), " = ", user_id);
639        _ = self.exec().execute(sql).await?;
640        for session in &sessions {
641            for hook in self.hooks() {
642                hook.after_delete_session(session, &hook_context).await?;
643            }
644        }
645        Ok(())
646    }
647
648    async fn delete_expired_sessions(&self) -> AuthResult<usize> {
649        let table = <S::Session as SqlxModel>::TABLE;
650        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
651        sql.ident(table);
652        sql.push(" WHERE ");
653        sql.compare_model::<S::Session>(
654            table,
655            S::Session::expires_at_column(),
656            " < ",
657            S::Session::timestamp_value(S::Session::expires_at_column(), Utc::now()),
658        );
659        let deleted = self.exec().execute(sql).await?;
660        usize::try_from(deleted)
661            .map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
662    }
663
664    async fn update_session_active_organization(
665        &self,
666        token: &str,
667        organization_id: Option<&str>,
668    ) -> AuthResult<S::Session> {
669        let Some(model) = self.get_session(token).await? else {
670            return Err(AuthError::SessionNotFound);
671        };
672
673        let mut active = model.into_active();
674        S::Session::set_active_organization_id(&mut active, organization_id.map(str::to_owned));
675        S::Session::set_updated_at(&mut active, Utc::now());
676        model::update::<S::Session>(self.exec(), &active)
677            .await?
678            .ok_or_else(record_not_updated)
679    }
680
681    async fn update_session_active_team(
682        &self,
683        token: &str,
684        team_id: Option<&str>,
685    ) -> AuthResult<S::Session> {
686        let model = self
687            .get_session(token)
688            .await?
689            .ok_or(AuthError::SessionNotFound)?;
690        let mut active = model.into_active();
691        S::Session::set_active_team_id(&mut active, team_id.map(str::to_owned))?;
692        S::Session::set_updated_at(&mut active, Utc::now());
693        model::update::<S::Session>(self.exec(), &active)
694            .await?
695            .ok_or_else(record_not_updated)
696    }
697}