Skip to main content

max / makenotwork

8.8 KB · 266 lines History Blame Raw
1 //! Authentication queries: login tokens, password resets, lockouts.
2
3 use chrono::{DateTime, Utc};
4 use sqlx::PgPool;
5
6 use super::models::DbLoginToken;
7 use super::{LoginTokenId, UserId};
8 use crate::error::Result;
9
10 /// Result of an atomic failed-login increment.
11 pub(crate) struct FailedLoginResult {
12 /// New failed_login_attempts count after increment.
13 pub attempts: i32,
14 /// Whether the account was just locked by this increment.
15 pub just_locked: bool,
16 }
17
18 /// Atomically increment failed login attempts and lock the account if the
19 /// threshold is reached. This prevents race conditions where concurrent
20 /// requests could each pass the lockout check before either increments.
21 ///
22 /// The UPDATE uses a single SQL statement with conditional locked_until
23 /// assignment, so PostgreSQL's row-level locking serializes concurrent callers.
24 #[tracing::instrument(skip_all)]
25 pub(crate) async fn increment_failed_login(
26 pool: &PgPool,
27 user_id: UserId,
28 max_attempts: i32,
29 lockout_minutes: i64,
30 ) -> Result<FailedLoginResult> {
31 // Capture the pre-update row under a row lock so `just_locked` can be
32 // derived from the *same* predicate the CASE uses to set the lock. Reading
33 // the old values in a CTE avoids the trap where a bare `RETURNING` column is
34 // the post-update value: the previous code computed `just_locked` as
35 // `(failed_login_attempts = $2)` against the new count, which only matched
36 // the exact-threshold attempt and read false on a *re-lock* after an expired
37 // window (counter already >= threshold), so the lockout notification was
38 // silently skipped on every re-lock. Now `just_locked` is exactly "the lock
39 // was (re)set on this call".
40 let row = sqlx::query!(
41 r#"
42 WITH prev AS (
43 SELECT failed_login_attempts AS old_attempts, locked_until AS old_lock
44 FROM users WHERE id = $1 FOR UPDATE
45 )
46 UPDATE users u
47 SET failed_login_attempts = u.failed_login_attempts + 1,
48 last_failed_login_at = NOW(),
49 locked_until = CASE
50 -- Set/refresh the lock only when reaching the threshold AND not
51 -- already inside an active lock window. Without the second clause,
52 -- every failed attempt during a lock pushed `locked_until` forward
53 -- another window, letting an attacker keep a victim perpetually
54 -- locked with one wrong password every <LOCKOUT_MINUTES. Now an
55 -- active lock simply runs out; a fresh failure after it expires
56 -- re-locks (the counter isn't reset until a successful login).
57 WHEN u.failed_login_attempts + 1 >= $2
58 AND (u.locked_until IS NULL OR u.locked_until <= NOW())
59 THEN NOW() + ($3 || ' minutes')::interval
60 ELSE u.locked_until
61 END
62 FROM prev
63 WHERE u.id = $1
64 RETURNING
65 u.failed_login_attempts,
66 (prev.old_attempts + 1 >= $2
67 AND (prev.old_lock IS NULL OR prev.old_lock <= NOW())) AS "just_locked!"
68 "#,
69 user_id as UserId,
70 max_attempts,
71 lockout_minutes.to_string(),
72 )
73 .fetch_one(pool)
74 .await?;
75
76 Ok(FailedLoginResult {
77 attempts: row.failed_login_attempts,
78 just_locked: row.just_locked,
79 })
80 }
81
82 /// Reset failed login attempts (on successful login)
83 #[tracing::instrument(skip_all)]
84 pub(crate) async fn reset_failed_login(pool: &PgPool, user_id: UserId) -> Result<()> {
85 sqlx::query!(
86 "UPDATE users SET failed_login_attempts = 0, locked_until = NULL WHERE id = $1",
87 user_id as UserId,
88 )
89 .execute(pool)
90 .await?;
91
92 Ok(())
93 }
94
95 /// Create a one-time login token
96 #[tracing::instrument(skip_all)]
97 pub(crate) async fn create_login_token(
98 pool: &PgPool,
99 user_id: UserId,
100 token_hash: &str,
101 expires_at: DateTime<Utc>,
102 ) -> Result<DbLoginToken> {
103 // runtime-checked: binds a chrono DateTime<Utc> param; a bind param's type can't be overridden in the macro when sqlx time+chrono features are unified.
104 let token = sqlx::query_as::<_, DbLoginToken>(
105 r"
106 INSERT INTO login_tokens (user_id, token_hash, expires_at)
107 VALUES ($1, $2, $3)
108 RETURNING *
109 ",
110 )
111 .bind(user_id)
112 .bind(token_hash)
113 .bind(expires_at)
114 .fetch_one(pool)
115 .await?;
116
117 Ok(token)
118 }
119
120 /// Atomically consume a login token: mark it used and return it in one step.
121 ///
122 /// Returns `Some(token)` if the token was valid and successfully consumed,
123 /// or `None` if the token was already used, expired, or does not exist.
124 /// Because this is a single UPDATE with `used_at IS NULL` in the WHERE clause,
125 /// concurrent requests for the same token will never both succeed.
126 #[tracing::instrument(skip_all)]
127 pub(crate) async fn consume_login_token(
128 pool: &PgPool,
129 token_hash: &str,
130 ) -> Result<Option<DbLoginToken>> {
131 let token = sqlx::query_as!(
132 DbLoginToken,
133 r#"
134 UPDATE login_tokens
135 SET used_at = NOW()
136 WHERE token_hash = $1
137 AND used_at IS NULL
138 AND expires_at > NOW()
139 RETURNING
140 id AS "id: LoginTokenId",
141 user_id AS "user_id: UserId",
142 token_hash,
143 expires_at AS "expires_at: chrono::DateTime<chrono::Utc>",
144 used_at AS "used_at: chrono::DateTime<chrono::Utc>",
145 created_at AS "created_at: chrono::DateTime<chrono::Utc>"
146 "#,
147 token_hash,
148 )
149 .fetch_optional(pool)
150 .await?;
151
152 Ok(token)
153 }
154
155 /// Create a single-use password reset token (mirrors [`create_login_token`]).
156 #[tracing::instrument(skip_all)]
157 pub(crate) async fn create_password_reset_token(
158 pool: &PgPool,
159 user_id: UserId,
160 token_hash: &str,
161 expires_at: DateTime<Utc>,
162 ) -> Result<()> {
163 // runtime-checked: binds a chrono DateTime<Utc> param; a bind param's type can't be overridden in the macro when sqlx time+chrono features are unified.
164 sqlx::query(
165 r"
166 INSERT INTO password_reset_tokens (user_id, token_hash, expires_at)
167 VALUES ($1, $2, $3)
168 ",
169 )
170 .bind(user_id)
171 .bind(token_hash)
172 .bind(expires_at)
173 .execute(pool)
174 .await?;
175
176 Ok(())
177 }
178
179 /// Check whether a password reset token is currently valid (unused, unexpired),
180 /// returning the user it belongs to, without consuming it. Used to decide
181 /// whether to render the reset form. The token is only spent on submit via
182 /// [`consume_password_reset_token`].
183 #[tracing::instrument(skip_all)]
184 pub(crate) async fn peek_password_reset_token(
185 pool: &PgPool,
186 token_hash: &str,
187 ) -> Result<Option<UserId>> {
188 let user_id = sqlx::query_scalar!(
189 r#"
190 SELECT user_id AS "user_id: UserId" FROM password_reset_tokens
191 WHERE token_hash = $1
192 AND used_at IS NULL
193 AND expires_at > NOW()
194 "#,
195 token_hash,
196 )
197 .fetch_optional(pool)
198 .await?;
199
200 Ok(user_id)
201 }
202
203 /// Atomically consume a password reset token, returning the user it belongs to.
204 ///
205 /// Returns `None` if the token was already used, expired, or does not exist.
206 /// The single UPDATE with `used_at IS NULL` in the WHERE clause guarantees a
207 /// replay (or a concurrent double-submit) can never succeed twice.
208 #[tracing::instrument(skip_all)]
209 pub(crate) async fn consume_password_reset_token(
210 pool: &PgPool,
211 token_hash: &str,
212 ) -> Result<Option<UserId>> {
213 let user_id = sqlx::query_scalar!(
214 r#"
215 UPDATE password_reset_tokens
216 SET used_at = NOW()
217 WHERE token_hash = $1
218 AND used_at IS NULL
219 AND expires_at > NOW()
220 RETURNING user_id AS "user_id: UserId"
221 "#,
222 token_hash,
223 )
224 .fetch_optional(pool)
225 .await?;
226
227 Ok(user_id)
228 }
229
230 /// Delete consumed or expired password reset tokens (housekeeping so the table
231 /// doesn't accumulate dead rows). Keeps recently-used rows briefly for audit.
232 #[tracing::instrument(skip_all)]
233 pub(crate) async fn prune_password_reset_tokens(pool: &PgPool) -> Result<u64> {
234 let result = sqlx::query!(
235 r#"
236 DELETE FROM password_reset_tokens
237 WHERE expires_at < NOW() - interval '7 days'
238 OR used_at < NOW() - interval '7 days'
239 "#,
240 )
241 .execute(pool)
242 .await?;
243
244 Ok(result.rows_affected())
245 }
246
247 /// Invalidate every outstanding reset token for a user. Called after a
248 /// successful reset so any other links mailed to the same account (e.g. a
249 /// double request) are dead, matching the old hash-binding's "completing a
250 /// reset kills all outstanding links" behavior.
251 #[tracing::instrument(skip_all)]
252 pub(crate) async fn invalidate_password_reset_tokens(pool: &PgPool, user_id: UserId) -> Result<()> {
253 sqlx::query!(
254 r#"
255 UPDATE password_reset_tokens
256 SET used_at = NOW()
257 WHERE user_id = $1 AND used_at IS NULL
258 "#,
259 user_id as UserId,
260 )
261 .execute(pool)
262 .await?;
263
264 Ok(())
265 }
266