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