1use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use mas_data_model::{Clock, UlidExt as _, UserRegistrationToken};
10use mas_storage::{
11 Page, Pagination,
12 pagination::Node,
13 user::{UserRegistrationTokenFilter, UserRegistrationTokenRepository},
14};
15use rand::Rng;
16use sea_query::{Condition, Expr, ExprTrait, PostgresQueryBuilder, Query, enum_def};
17use sea_query_sqlx::SqlxBinder;
18use sqlx::PgConnection;
19use ulid::Ulid;
20use uuid::Uuid;
21
22use crate::{
23 DatabaseInconsistencyError,
24 errors::DatabaseError,
25 filter::{Filter, StatementExt},
26 iden::UserRegistrationTokens,
27 pagination::QueryBuilderExt,
28 tracing::ExecuteExt,
29};
30
31pub struct PgUserRegistrationTokenRepository<'c> {
34 conn: &'c mut PgConnection,
35}
36
37impl<'c> PgUserRegistrationTokenRepository<'c> {
38 pub fn new(conn: &'c mut PgConnection) -> Self {
41 Self { conn }
42 }
43}
44
45#[derive(Debug, Clone, sqlx::FromRow)]
46#[enum_def]
47struct UserRegistrationTokenLookup {
48 user_registration_token_id: Uuid,
49 token: String,
50 username: Option<String>,
51 email: Option<String>,
52 passwordless: bool,
53 room_id: Option<String>,
54 usage_limit: Option<i32>,
55 times_used: i32,
56 created_at: DateTime<Utc>,
57 last_used_at: Option<DateTime<Utc>>,
58 expires_at: Option<DateTime<Utc>>,
59 revoked_at: Option<DateTime<Utc>>,
60}
61
62impl Node<Ulid> for UserRegistrationTokenLookup {
63 fn cursor(&self) -> Ulid {
64 self.user_registration_token_id.into()
65 }
66}
67
68impl Filter for UserRegistrationTokenFilter {
69 fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
70 sea_query::Condition::all()
71 .add_option(self.has_been_used().map(|has_been_used| {
72 if has_been_used {
73 Expr::col((
74 UserRegistrationTokens::Table,
75 UserRegistrationTokens::TimesUsed,
76 ))
77 .gt(0)
78 } else {
79 Expr::col((
80 UserRegistrationTokens::Table,
81 UserRegistrationTokens::TimesUsed,
82 ))
83 .eq(0)
84 }
85 }))
86 .add_option(self.is_revoked().map(|is_revoked| {
87 if is_revoked {
88 Expr::col((
89 UserRegistrationTokens::Table,
90 UserRegistrationTokens::RevokedAt,
91 ))
92 .is_not_null()
93 } else {
94 Expr::col((
95 UserRegistrationTokens::Table,
96 UserRegistrationTokens::RevokedAt,
97 ))
98 .is_null()
99 }
100 }))
101 .add_option(self.is_expired().map(|is_expired| {
102 if is_expired {
103 Condition::all()
104 .add(
105 Expr::col((
106 UserRegistrationTokens::Table,
107 UserRegistrationTokens::ExpiresAt,
108 ))
109 .is_not_null(),
110 )
111 .add(
112 Expr::col((
113 UserRegistrationTokens::Table,
114 UserRegistrationTokens::ExpiresAt,
115 ))
116 .lt(Expr::val(self.now())),
117 )
118 } else {
119 Condition::any()
120 .add(
121 Expr::col((
122 UserRegistrationTokens::Table,
123 UserRegistrationTokens::ExpiresAt,
124 ))
125 .is_null(),
126 )
127 .add(
128 Expr::col((
129 UserRegistrationTokens::Table,
130 UserRegistrationTokens::ExpiresAt,
131 ))
132 .gte(Expr::val(self.now())),
133 )
134 }
135 }))
136 .add_option(self.is_valid().map(|is_valid| {
137 let valid = Condition::all()
138 .add(
140 Condition::any()
141 .add(
142 Expr::col((
143 UserRegistrationTokens::Table,
144 UserRegistrationTokens::UsageLimit,
145 ))
146 .is_null(),
147 )
148 .add(
149 Expr::col((
150 UserRegistrationTokens::Table,
151 UserRegistrationTokens::TimesUsed,
152 ))
153 .lt(Expr::col((
154 UserRegistrationTokens::Table,
155 UserRegistrationTokens::UsageLimit,
156 ))),
157 ),
158 )
159 .add(
161 Expr::col((
162 UserRegistrationTokens::Table,
163 UserRegistrationTokens::RevokedAt,
164 ))
165 .is_null(),
166 )
167 .add(
169 Condition::any()
170 .add(
171 Expr::col((
172 UserRegistrationTokens::Table,
173 UserRegistrationTokens::ExpiresAt,
174 ))
175 .is_null(),
176 )
177 .add(
178 Expr::col((
179 UserRegistrationTokens::Table,
180 UserRegistrationTokens::ExpiresAt,
181 ))
182 .gte(Expr::val(self.now())),
183 ),
184 );
185
186 if is_valid { valid } else { valid.not() }
187 }))
188 }
189}
190
191impl TryFrom<UserRegistrationTokenLookup> for UserRegistrationToken {
192 type Error = DatabaseInconsistencyError;
193
194 fn try_from(res: UserRegistrationTokenLookup) -> Result<Self, Self::Error> {
195 let id = Ulid::from(res.user_registration_token_id);
196
197 let usage_limit = res
198 .usage_limit
199 .map(u32::try_from)
200 .transpose()
201 .map_err(|e| {
202 DatabaseInconsistencyError::on("user_registration_tokens")
203 .column("usage_limit")
204 .row(id)
205 .source(e)
206 })?;
207
208 let times_used = res.times_used.try_into().map_err(|e| {
209 DatabaseInconsistencyError::on("user_registration_tokens")
210 .column("times_used")
211 .row(id)
212 .source(e)
213 })?;
214
215 Ok(UserRegistrationToken {
216 id,
217 token: res.token,
218 username: res.username,
219 email: res.email,
220 passwordless: res.passwordless,
221 room_id: res.room_id,
222 usage_limit,
223 times_used,
224 created_at: res.created_at,
225 last_used_at: res.last_used_at,
226 expires_at: res.expires_at,
227 revoked_at: res.revoked_at,
228 })
229 }
230}
231
232#[async_trait]
233impl UserRegistrationTokenRepository for PgUserRegistrationTokenRepository<'_> {
234 type Error = DatabaseError;
235
236 #[tracing::instrument(
237 name = "db.user_registration_token.list",
238 skip_all,
239 fields(
240 db.query.text,
241 ),
242 err,
243 )]
244 async fn list(
245 &mut self,
246 filter: UserRegistrationTokenFilter,
247 pagination: Pagination,
248 ) -> Result<Page<UserRegistrationToken>, Self::Error> {
249 let (sql, arguments) = Query::select()
250 .expr_as(
251 Expr::col((
252 UserRegistrationTokens::Table,
253 UserRegistrationTokens::UserRegistrationTokenId,
254 )),
255 UserRegistrationTokenLookupIden::UserRegistrationTokenId,
256 )
257 .expr_as(
258 Expr::col((UserRegistrationTokens::Table, UserRegistrationTokens::Token)),
259 UserRegistrationTokenLookupIden::Token,
260 )
261 .expr_as(
262 Expr::col((
263 UserRegistrationTokens::Table,
264 UserRegistrationTokens::Username,
265 )),
266 UserRegistrationTokenLookupIden::Username,
267 )
268 .expr_as(
269 Expr::col((UserRegistrationTokens::Table, UserRegistrationTokens::Email)),
270 UserRegistrationTokenLookupIden::Email,
271 )
272 .expr_as(
273 Expr::col((
274 UserRegistrationTokens::Table,
275 UserRegistrationTokens::Passwordless,
276 )),
277 UserRegistrationTokenLookupIden::Passwordless,
278 )
279 .expr_as(
280 Expr::col((
281 UserRegistrationTokens::Table,
282 UserRegistrationTokens::RoomId,
283 )),
284 UserRegistrationTokenLookupIden::RoomId,
285 )
286 .expr_as(
287 Expr::col((
288 UserRegistrationTokens::Table,
289 UserRegistrationTokens::UsageLimit,
290 )),
291 UserRegistrationTokenLookupIden::UsageLimit,
292 )
293 .expr_as(
294 Expr::col((
295 UserRegistrationTokens::Table,
296 UserRegistrationTokens::TimesUsed,
297 )),
298 UserRegistrationTokenLookupIden::TimesUsed,
299 )
300 .expr_as(
301 Expr::col((
302 UserRegistrationTokens::Table,
303 UserRegistrationTokens::CreatedAt,
304 )),
305 UserRegistrationTokenLookupIden::CreatedAt,
306 )
307 .expr_as(
308 Expr::col((
309 UserRegistrationTokens::Table,
310 UserRegistrationTokens::LastUsedAt,
311 )),
312 UserRegistrationTokenLookupIden::LastUsedAt,
313 )
314 .expr_as(
315 Expr::col((
316 UserRegistrationTokens::Table,
317 UserRegistrationTokens::ExpiresAt,
318 )),
319 UserRegistrationTokenLookupIden::ExpiresAt,
320 )
321 .expr_as(
322 Expr::col((
323 UserRegistrationTokens::Table,
324 UserRegistrationTokens::RevokedAt,
325 )),
326 UserRegistrationTokenLookupIden::RevokedAt,
327 )
328 .from(UserRegistrationTokens::Table)
329 .apply_filter(filter)
330 .generate_pagination(
331 (
332 UserRegistrationTokens::Table,
333 UserRegistrationTokens::UserRegistrationTokenId,
334 ),
335 pagination,
336 )
337 .build_sqlx(PostgresQueryBuilder);
338
339 let edges: Vec<UserRegistrationTokenLookup> = sqlx::query_as_with(&sql, arguments)
340 .traced()
341 .fetch_all(&mut *self.conn)
342 .await?;
343
344 let page = pagination
345 .process(edges)
346 .try_map(UserRegistrationToken::try_from)?;
347
348 Ok(page)
349 }
350
351 #[tracing::instrument(
352 name = "db.user_registration_token.count",
353 skip_all,
354 fields(
355 db.query.text,
356 user_registration_token.filter = ?filter,
357 ),
358 err,
359 )]
360 async fn count(&mut self, filter: UserRegistrationTokenFilter) -> Result<usize, Self::Error> {
361 let (sql, values) = Query::select()
362 .expr(
363 Expr::col((
364 UserRegistrationTokens::Table,
365 UserRegistrationTokens::UserRegistrationTokenId,
366 ))
367 .count(),
368 )
369 .from(UserRegistrationTokens::Table)
370 .apply_filter(filter)
371 .build_sqlx(PostgresQueryBuilder);
372
373 let count: i64 = sqlx::query_scalar_with(&sql, values)
374 .traced()
375 .fetch_one(&mut *self.conn)
376 .await?;
377
378 count
379 .try_into()
380 .map_err(DatabaseError::to_invalid_operation)
381 }
382
383 #[tracing::instrument(
384 name = "db.user_registration_token.lookup",
385 skip_all,
386 fields(
387 db.query.text,
388 user_registration_token.id = %id,
389 ),
390 err,
391 )]
392 async fn lookup(&mut self, id: Ulid) -> Result<Option<UserRegistrationToken>, Self::Error> {
393 let res = sqlx::query_as!(
394 UserRegistrationTokenLookup,
395 r#"
396 SELECT user_registration_token_id,
397 token,
398 username,
399 email,
400 passwordless,
401 room_id,
402 usage_limit,
403 times_used,
404 created_at,
405 last_used_at,
406 expires_at,
407 revoked_at
408 FROM user_registration_tokens
409 WHERE user_registration_token_id = $1
410 "#,
411 Uuid::from(id)
412 )
413 .traced()
414 .fetch_optional(&mut *self.conn)
415 .await?;
416
417 let Some(res) = res else {
418 return Ok(None);
419 };
420
421 Ok(Some(res.try_into()?))
422 }
423
424 #[tracing::instrument(
425 name = "db.user_registration_token.find_by_token",
426 skip_all,
427 fields(
428 db.query.text,
429 token = %token,
430 ),
431 err,
432 )]
433 async fn find_by_token(
434 &mut self,
435 token: &str,
436 ) -> Result<Option<UserRegistrationToken>, Self::Error> {
437 let res = sqlx::query_as!(
438 UserRegistrationTokenLookup,
439 r#"
440 SELECT user_registration_token_id,
441 token,
442 username,
443 email,
444 passwordless,
445 room_id,
446 usage_limit,
447 times_used,
448 created_at,
449 last_used_at,
450 expires_at,
451 revoked_at
452 FROM user_registration_tokens
453 WHERE token = $1
454 "#,
455 token
456 )
457 .traced()
458 .fetch_optional(&mut *self.conn)
459 .await?;
460
461 let Some(res) = res else {
462 return Ok(None);
463 };
464
465 Ok(Some(res.try_into()?))
466 }
467
468 #[tracing::instrument(
469 name = "db.user_registration_token.add",
470 skip_all,
471 fields(
472 db.query.text,
473 user_registration_token.token = %token,
474 ),
475 err,
476 )]
477 async fn add(
478 &mut self,
479 rng: &mut (dyn Rng + Send),
480 clock: &dyn mas_data_model::Clock,
481 token: String,
482 usage_limit: Option<u32>,
483 expires_at: Option<DateTime<Utc>>,
484 username: Option<String>,
485 email: Option<String>,
486 passwordless: bool,
487 room_id: Option<String>,
488 ) -> Result<UserRegistrationToken, Self::Error> {
489 let created_at = clock.now();
490 let id = Ulid::from_datetime_with_rng(created_at, rng);
491
492 let usage_limit_i32 = usage_limit
493 .map(i32::try_from)
494 .transpose()
495 .map_err(DatabaseError::to_invalid_operation)?;
496
497 sqlx::query!(
498 r#"
499 INSERT INTO user_registration_tokens
500 (user_registration_token_id, token, usage_limit, created_at, expires_at, username, email, passwordless, room_id)
501 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
502 "#,
503 Uuid::from(id),
504 &token,
505 usage_limit_i32,
506 created_at,
507 expires_at,
508 username.as_deref(),
509 email.as_deref(),
510 passwordless,
511 room_id.as_deref(),
512 )
513 .traced()
514 .execute(&mut *self.conn)
515 .await?;
516
517 Ok(UserRegistrationToken {
518 id,
519 token,
520 username,
521 email,
522 passwordless,
523 room_id,
524 usage_limit,
525 times_used: 0,
526 created_at,
527 last_used_at: None,
528 expires_at,
529 revoked_at: None,
530 })
531 }
532
533 #[tracing::instrument(
534 name = "db.user_registration_token.use_token",
535 skip_all,
536 fields(
537 db.query.text,
538 user_registration_token.id = %token.id,
539 ),
540 err,
541 )]
542 async fn use_token(
543 &mut self,
544 clock: &dyn Clock,
545 token: UserRegistrationToken,
546 ) -> Result<UserRegistrationToken, Self::Error> {
547 let now = clock.now();
548 let new_times_used = sqlx::query_scalar!(
549 r#"
550 UPDATE user_registration_tokens
551 SET times_used = times_used + 1,
552 last_used_at = $2
553 WHERE user_registration_token_id = $1 AND revoked_at IS NULL
554 RETURNING times_used
555 "#,
556 Uuid::from(token.id),
557 now,
558 )
559 .traced()
560 .fetch_one(&mut *self.conn)
561 .await?;
562
563 let new_times_used = new_times_used
564 .try_into()
565 .map_err(DatabaseError::to_invalid_operation)?;
566
567 Ok(UserRegistrationToken {
568 times_used: new_times_used,
569 last_used_at: Some(now),
570 ..token
571 })
572 }
573
574 #[tracing::instrument(
575 name = "db.user_registration_token.revoke",
576 skip_all,
577 fields(
578 db.query.text,
579 user_registration_token.id = %token.id,
580 ),
581 err,
582 )]
583 async fn revoke(
584 &mut self,
585 clock: &dyn Clock,
586 mut token: UserRegistrationToken,
587 ) -> Result<UserRegistrationToken, Self::Error> {
588 let revoked_at = clock.now();
589 let res = sqlx::query!(
590 r#"
591 UPDATE user_registration_tokens
592 SET revoked_at = $2
593 WHERE user_registration_token_id = $1
594 "#,
595 Uuid::from(token.id),
596 revoked_at,
597 )
598 .traced()
599 .execute(&mut *self.conn)
600 .await?;
601
602 DatabaseError::ensure_affected_rows(&res, 1)?;
603
604 token.revoked_at = Some(revoked_at);
605
606 Ok(token)
607 }
608
609 #[tracing::instrument(
610 name = "db.user_registration_token.unrevoke",
611 skip_all,
612 fields(
613 db.query.text,
614 user_registration_token.id = %token.id,
615 ),
616 err,
617 )]
618 async fn unrevoke(
619 &mut self,
620 mut token: UserRegistrationToken,
621 ) -> Result<UserRegistrationToken, Self::Error> {
622 let res = sqlx::query!(
623 r#"
624 UPDATE user_registration_tokens
625 SET revoked_at = NULL
626 WHERE user_registration_token_id = $1
627 "#,
628 Uuid::from(token.id),
629 )
630 .traced()
631 .execute(&mut *self.conn)
632 .await?;
633
634 DatabaseError::ensure_affected_rows(&res, 1)?;
635
636 token.revoked_at = None;
637
638 Ok(token)
639 }
640
641 #[tracing::instrument(
642 name = "db.user_registration_token.set_expiry",
643 skip_all,
644 fields(
645 db.query.text,
646 user_registration_token.id = %token.id,
647 ),
648 err,
649 )]
650 async fn set_expiry(
651 &mut self,
652 mut token: UserRegistrationToken,
653 expires_at: Option<DateTime<Utc>>,
654 ) -> Result<UserRegistrationToken, Self::Error> {
655 let res = sqlx::query!(
656 r#"
657 UPDATE user_registration_tokens
658 SET expires_at = $2
659 WHERE user_registration_token_id = $1
660 "#,
661 Uuid::from(token.id),
662 expires_at,
663 )
664 .traced()
665 .execute(&mut *self.conn)
666 .await?;
667
668 DatabaseError::ensure_affected_rows(&res, 1)?;
669
670 token.expires_at = expires_at;
671
672 Ok(token)
673 }
674
675 #[tracing::instrument(
676 name = "db.user_registration_token.set_usage_limit",
677 skip_all,
678 fields(
679 db.query.text,
680 user_registration_token.id = %token.id,
681 ),
682 err,
683 )]
684 async fn set_usage_limit(
685 &mut self,
686 mut token: UserRegistrationToken,
687 usage_limit: Option<u32>,
688 ) -> Result<UserRegistrationToken, Self::Error> {
689 let usage_limit_i32 = usage_limit
690 .map(i32::try_from)
691 .transpose()
692 .map_err(DatabaseError::to_invalid_operation)?;
693
694 let res = sqlx::query!(
695 r#"
696 UPDATE user_registration_tokens
697 SET usage_limit = $2
698 WHERE user_registration_token_id = $1
699 "#,
700 Uuid::from(token.id),
701 usage_limit_i32,
702 )
703 .traced()
704 .execute(&mut *self.conn)
705 .await?;
706
707 DatabaseError::ensure_affected_rows(&res, 1)?;
708
709 token.usage_limit = usage_limit;
710
711 Ok(token)
712 }
713
714 #[tracing::instrument(
715 name = "db.user_registration_token.set_username",
716 skip_all,
717 fields(
718 db.query.text,
719 user_registration_token.id = %token.id,
720 ),
721 err,
722 )]
723 async fn set_username(
724 &mut self,
725 mut token: UserRegistrationToken,
726 username: Option<String>,
727 ) -> Result<UserRegistrationToken, Self::Error> {
728 let res = sqlx::query!(
729 r#"
730 UPDATE user_registration_tokens
731 SET username = $2
732 WHERE user_registration_token_id = $1
733 "#,
734 Uuid::from(token.id),
735 username.as_deref(),
736 )
737 .traced()
738 .execute(&mut *self.conn)
739 .await?;
740
741 DatabaseError::ensure_affected_rows(&res, 1)?;
742
743 token.username = username;
744
745 Ok(token)
746 }
747
748 #[tracing::instrument(
749 name = "db.user_registration_token.set_email",
750 skip_all,
751 fields(
752 db.query.text,
753 user_registration_token.id = %token.id,
754 ),
755 err,
756 )]
757 async fn set_email(
758 &mut self,
759 mut token: UserRegistrationToken,
760 email: Option<String>,
761 ) -> Result<UserRegistrationToken, Self::Error> {
762 let res = sqlx::query!(
763 r#"
764 UPDATE user_registration_tokens
765 SET email = $2
766 WHERE user_registration_token_id = $1
767 "#,
768 Uuid::from(token.id),
769 email.as_deref(),
770 )
771 .traced()
772 .execute(&mut *self.conn)
773 .await?;
774
775 DatabaseError::ensure_affected_rows(&res, 1)?;
776
777 token.email = email;
778
779 Ok(token)
780 }
781}
782
783#[cfg(test)]
784mod tests {
785 use chrono::Duration;
786 use mas_data_model::{Clock as _, clock::MockClock};
787 use mas_storage::{Pagination, user::UserRegistrationTokenFilter};
788 use rand::SeedableRng;
789 use rand_chacha::ChaChaRng;
790 use sqlx::PgPool;
791
792 use crate::PgRepository;
793
794 #[sqlx::test(migrator = "crate::MIGRATOR")]
795 async fn test_unrevoke(pool: PgPool) {
796 let mut rng = ChaChaRng::seed_from_u64(42);
797 let clock = MockClock::default();
798
799 let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
800
801 let token = repo
803 .user_registration_token()
804 .add(
805 &mut rng,
806 &clock,
807 "test_token".to_owned(),
808 None,
809 None,
810 None,
811 None,
812 false,
813 None,
814 )
815 .await
816 .unwrap();
817
818 let revoked_token = repo
820 .user_registration_token()
821 .revoke(&clock, token)
822 .await
823 .unwrap();
824
825 assert!(revoked_token.revoked_at.is_some());
827
828 let unrevoked_token = repo
830 .user_registration_token()
831 .unrevoke(revoked_token)
832 .await
833 .unwrap();
834
835 assert!(unrevoked_token.revoked_at.is_none());
837
838 let non_revoked_filter = UserRegistrationTokenFilter::new(clock.now()).with_revoked(false);
840 let page = repo
841 .user_registration_token()
842 .list(non_revoked_filter, Pagination::first(10))
843 .await
844 .unwrap();
845
846 assert!(page.edges.iter().any(|t| t.node.id == unrevoked_token.id));
847 }
848
849 #[sqlx::test(migrator = "crate::MIGRATOR")]
850 async fn test_set_expiry(pool: PgPool) {
851 let mut rng = ChaChaRng::seed_from_u64(42);
852 let clock = MockClock::default();
853
854 let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
855
856 let token = repo
858 .user_registration_token()
859 .add(
860 &mut rng,
861 &clock,
862 "test_token_expiry".to_owned(),
863 None,
864 None,
865 None,
866 None,
867 false,
868 None,
869 )
870 .await
871 .unwrap();
872
873 assert!(token.expires_at.is_none());
875
876 let future_time = clock.now() + Duration::days(30);
878 let updated_token = repo
879 .user_registration_token()
880 .set_expiry(token, Some(future_time))
881 .await
882 .unwrap();
883
884 assert_eq!(updated_token.expires_at, Some(future_time));
886
887 let final_token = repo
889 .user_registration_token()
890 .set_expiry(updated_token, None)
891 .await
892 .unwrap();
893
894 assert!(final_token.expires_at.is_none());
896 }
897
898 #[sqlx::test(migrator = "crate::MIGRATOR")]
899 async fn test_set_usage_limit(pool: PgPool) {
900 let mut rng = ChaChaRng::seed_from_u64(42);
901 let clock = MockClock::default();
902
903 let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
904
905 let token = repo
907 .user_registration_token()
908 .add(
909 &mut rng,
910 &clock,
911 "test_token_limit".to_owned(),
912 None,
913 None,
914 None,
915 None,
916 false,
917 None,
918 )
919 .await
920 .unwrap();
921
922 assert!(token.usage_limit.is_none());
924
925 let updated_token = repo
927 .user_registration_token()
928 .set_usage_limit(token, Some(5))
929 .await
930 .unwrap();
931
932 assert_eq!(updated_token.usage_limit, Some(5));
934
935 let changed_token = repo
937 .user_registration_token()
938 .set_usage_limit(updated_token, Some(10))
939 .await
940 .unwrap();
941
942 assert_eq!(changed_token.usage_limit, Some(10));
944
945 let final_token = repo
947 .user_registration_token()
948 .set_usage_limit(changed_token, None)
949 .await
950 .unwrap();
951
952 assert!(final_token.usage_limit.is_none());
954 }
955
956 #[sqlx::test(migrator = "crate::MIGRATOR")]
957 async fn test_room_id(pool: PgPool) {
958 let mut rng = ChaChaRng::seed_from_u64(42);
959 let clock = MockClock::default();
960
961 let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
962
963 for room_id in [None, Some("!abc:example.com".to_owned())] {
964 let token = repo
965 .user_registration_token()
966 .add(
967 &mut rng,
968 &clock,
969 format!("token_{room_id:?}"),
970 None,
971 None,
972 None,
973 None,
974 false,
975 room_id.clone(),
976 )
977 .await
978 .unwrap();
979 assert_eq!(token.room_id, room_id);
980
981 let looked_up = repo
982 .user_registration_token()
983 .lookup(token.id)
984 .await
985 .unwrap()
986 .unwrap();
987 assert_eq!(looked_up, token);
988
989 let found = repo
990 .user_registration_token()
991 .find_by_token(&token.token)
992 .await
993 .unwrap()
994 .unwrap();
995 assert_eq!(found, token);
996 }
997
998 let page = repo
999 .user_registration_token()
1000 .list(
1001 UserRegistrationTokenFilter::new(clock.now()),
1002 Pagination::first(10),
1003 )
1004 .await
1005 .unwrap();
1006 let mut room_ids: Vec<_> = page.edges.into_iter().map(|e| e.node.room_id).collect();
1007 room_ids.sort();
1008 assert_eq!(room_ids, [None, Some("!abc:example.com".to_owned())]);
1009 }
1010
1011 #[sqlx::test(migrator = "crate::MIGRATOR")]
1012 async fn test_list_and_count(pool: PgPool) {
1013 let mut rng = ChaChaRng::seed_from_u64(42);
1014 let clock = MockClock::default();
1015
1016 let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
1017
1018 let _token1 = repo
1021 .user_registration_token()
1022 .add(
1023 &mut rng,
1024 &clock,
1025 "token1".to_owned(),
1026 None,
1027 None,
1028 None,
1029 None,
1030 false,
1031 None,
1032 )
1033 .await
1034 .unwrap();
1035
1036 let token2 = repo
1038 .user_registration_token()
1039 .add(
1040 &mut rng,
1041 &clock,
1042 "token2".to_owned(),
1043 None,
1044 None,
1045 None,
1046 None,
1047 false,
1048 None,
1049 )
1050 .await
1051 .unwrap();
1052 let token2 = repo
1053 .user_registration_token()
1054 .use_token(&clock, token2)
1055 .await
1056 .unwrap();
1057
1058 let past_time = clock.now() - Duration::days(1);
1060 let token3 = repo
1061 .user_registration_token()
1062 .add(
1063 &mut rng,
1064 &clock,
1065 "token3".to_owned(),
1066 None,
1067 Some(past_time),
1068 None,
1069 None,
1070 false,
1071 None,
1072 )
1073 .await
1074 .unwrap();
1075
1076 let token4 = repo
1078 .user_registration_token()
1079 .add(
1080 &mut rng,
1081 &clock,
1082 "token4".to_owned(),
1083 None,
1084 None,
1085 None,
1086 None,
1087 false,
1088 None,
1089 )
1090 .await
1091 .unwrap();
1092 let token4 = repo
1093 .user_registration_token()
1094 .revoke(&clock, token4)
1095 .await
1096 .unwrap();
1097
1098 let empty_filter = UserRegistrationTokenFilter::new(clock.now());
1100 let page = repo
1101 .user_registration_token()
1102 .list(empty_filter, Pagination::first(10))
1103 .await
1104 .unwrap();
1105 assert_eq!(page.edges.len(), 4);
1106
1107 let count = repo
1109 .user_registration_token()
1110 .count(empty_filter)
1111 .await
1112 .unwrap();
1113 assert_eq!(count, 4);
1114
1115 let used_filter = UserRegistrationTokenFilter::new(clock.now()).with_been_used(true);
1117 let page = repo
1118 .user_registration_token()
1119 .list(used_filter, Pagination::first(10))
1120 .await
1121 .unwrap();
1122 assert_eq!(page.edges.len(), 1);
1123 assert_eq!(page.edges[0].node.id, token2.id);
1124
1125 let unused_filter = UserRegistrationTokenFilter::new(clock.now()).with_been_used(false);
1127 let page = repo
1128 .user_registration_token()
1129 .list(unused_filter, Pagination::first(10))
1130 .await
1131 .unwrap();
1132 assert_eq!(page.edges.len(), 3);
1133
1134 let expired_filter = UserRegistrationTokenFilter::new(clock.now()).with_expired(true);
1136 let page = repo
1137 .user_registration_token()
1138 .list(expired_filter, Pagination::first(10))
1139 .await
1140 .unwrap();
1141 assert_eq!(page.edges.len(), 1);
1142 assert_eq!(page.edges[0].node.id, token3.id);
1143
1144 let not_expired_filter = UserRegistrationTokenFilter::new(clock.now()).with_expired(false);
1145 let page = repo
1146 .user_registration_token()
1147 .list(not_expired_filter, Pagination::first(10))
1148 .await
1149 .unwrap();
1150 assert_eq!(page.edges.len(), 3);
1151
1152 let revoked_filter = UserRegistrationTokenFilter::new(clock.now()).with_revoked(true);
1154 let page = repo
1155 .user_registration_token()
1156 .list(revoked_filter, Pagination::first(10))
1157 .await
1158 .unwrap();
1159 assert_eq!(page.edges.len(), 1);
1160 assert_eq!(page.edges[0].node.id, token4.id);
1161
1162 let not_revoked_filter = UserRegistrationTokenFilter::new(clock.now()).with_revoked(false);
1163 let page = repo
1164 .user_registration_token()
1165 .list(not_revoked_filter, Pagination::first(10))
1166 .await
1167 .unwrap();
1168 assert_eq!(page.edges.len(), 3);
1169
1170 let valid_filter = UserRegistrationTokenFilter::new(clock.now()).with_valid(true);
1172 let page = repo
1173 .user_registration_token()
1174 .list(valid_filter, Pagination::first(10))
1175 .await
1176 .unwrap();
1177 assert_eq!(page.edges.len(), 2);
1178
1179 let invalid_filter = UserRegistrationTokenFilter::new(clock.now()).with_valid(false);
1180 let page = repo
1181 .user_registration_token()
1182 .list(invalid_filter, Pagination::first(10))
1183 .await
1184 .unwrap();
1185 assert_eq!(page.edges.len(), 2);
1186
1187 let combined_filter = UserRegistrationTokenFilter::new(clock.now())
1189 .with_been_used(false)
1190 .with_revoked(true);
1191 let page = repo
1192 .user_registration_token()
1193 .list(combined_filter, Pagination::first(10))
1194 .await
1195 .unwrap();
1196 assert_eq!(page.edges.len(), 1);
1197 assert_eq!(page.edges[0].node.id, token4.id);
1198
1199 let page = repo
1201 .user_registration_token()
1202 .list(empty_filter, Pagination::first(2))
1203 .await
1204 .unwrap();
1205 assert_eq!(page.edges.len(), 2);
1206 }
1207}