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