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 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 if fields.contains_key(*name) {
83 **destination = None;
84 }
85 }
86 if persist {
87 fields.apply_adapter_transforms_async().await?;
88 }
89 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 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}