this repo has no description
lewis.moe
1use chrono::{DateTime, Duration, Utc};
2use rand::Rng;
3use sqlx::PgPool;
4use uuid::Uuid;
5
6use super::super::OAuthError;
7
8pub struct TwoFactorChallenge {
9 pub id: Uuid,
10 pub did: String,
11 pub request_uri: String,
12 pub code: String,
13 pub attempts: i32,
14 pub created_at: DateTime<Utc>,
15 pub expires_at: DateTime<Utc>,
16}
17
18pub fn generate_2fa_code() -> String {
19 let mut rng = rand::thread_rng();
20 let code: u32 = rng.gen_range(0..1_000_000);
21 format!("{:06}", code)
22}
23
24pub async fn create_2fa_challenge(
25 pool: &PgPool,
26 did: &str,
27 request_uri: &str,
28) -> Result<TwoFactorChallenge, OAuthError> {
29 let code = generate_2fa_code();
30 let expires_at = Utc::now() + Duration::minutes(10);
31
32 let row = sqlx::query!(
33 r#"
34 INSERT INTO oauth_2fa_challenge (did, request_uri, code, expires_at)
35 VALUES ($1, $2, $3, $4)
36 RETURNING id, did, request_uri, code, attempts, created_at, expires_at
37 "#,
38 did,
39 request_uri,
40 code,
41 expires_at,
42 )
43 .fetch_one(pool)
44 .await?;
45
46 Ok(TwoFactorChallenge {
47 id: row.id,
48 did: row.did,
49 request_uri: row.request_uri,
50 code: row.code,
51 attempts: row.attempts,
52 created_at: row.created_at,
53 expires_at: row.expires_at,
54 })
55}
56
57pub async fn get_2fa_challenge(
58 pool: &PgPool,
59 request_uri: &str,
60) -> Result<Option<TwoFactorChallenge>, OAuthError> {
61 let row = sqlx::query!(
62 r#"
63 SELECT id, did, request_uri, code, attempts, created_at, expires_at
64 FROM oauth_2fa_challenge
65 WHERE request_uri = $1
66 "#,
67 request_uri
68 )
69 .fetch_optional(pool)
70 .await?;
71
72 Ok(row.map(|r| TwoFactorChallenge {
73 id: r.id,
74 did: r.did,
75 request_uri: r.request_uri,
76 code: r.code,
77 attempts: r.attempts,
78 created_at: r.created_at,
79 expires_at: r.expires_at,
80 }))
81}
82
83pub async fn increment_2fa_attempts(pool: &PgPool, id: Uuid) -> Result<i32, OAuthError> {
84 let row = sqlx::query!(
85 r#"
86 UPDATE oauth_2fa_challenge
87 SET attempts = attempts + 1
88 WHERE id = $1
89 RETURNING attempts
90 "#,
91 id
92 )
93 .fetch_one(pool)
94 .await?;
95
96 Ok(row.attempts)
97}
98
99pub async fn delete_2fa_challenge(pool: &PgPool, id: Uuid) -> Result<(), OAuthError> {
100 sqlx::query!(
101 r#"
102 DELETE FROM oauth_2fa_challenge WHERE id = $1
103 "#,
104 id
105 )
106 .execute(pool)
107 .await?;
108
109 Ok(())
110}
111
112pub async fn delete_2fa_challenge_by_request_uri(
113 pool: &PgPool,
114 request_uri: &str,
115) -> Result<(), OAuthError> {
116 sqlx::query!(
117 r#"
118 DELETE FROM oauth_2fa_challenge WHERE request_uri = $1
119 "#,
120 request_uri
121 )
122 .execute(pool)
123 .await?;
124
125 Ok(())
126}
127
128pub async fn cleanup_expired_2fa_challenges(pool: &PgPool) -> Result<u64, OAuthError> {
129 let result = sqlx::query!(
130 r#"
131 DELETE FROM oauth_2fa_challenge WHERE expires_at < NOW()
132 "#
133 )
134 .execute(pool)
135 .await?;
136
137 Ok(result.rows_affected())
138}
139
140pub async fn check_user_2fa_enabled(pool: &PgPool, did: &str) -> Result<bool, OAuthError> {
141 let row = sqlx::query!(
142 r#"
143 SELECT two_factor_enabled
144 FROM users
145 WHERE did = $1
146 "#,
147 did
148 )
149 .fetch_optional(pool)
150 .await?;
151
152 Ok(row.map(|r| r.two_factor_enabled).unwrap_or(false))
153}