goat.navy
forked from
tranquil.farm/tranquil-pds
Our Personal Data Server from scratch!
11 kB
346 lines
1use super::*;
2
3#[derive(Debug, Deserialize)]
4pub struct Authorize2faQuery {
5 pub request_uri: String,
6 pub channel: Option<String>,
7}
8
9#[derive(Debug, Deserialize)]
10pub struct Authorize2faSubmit {
11 pub request_uri: String,
12 pub code: String,
13 #[serde(default)]
14 pub trust_device: bool,
15}
16
17const MAX_2FA_ATTEMPTS: i32 = 5;
18
19pub async fn authorize_2fa_get(
20 State(state): State<AppState>,
21 Query(query): Query<Authorize2faQuery>,
22) -> Response {
23 let twofa_request_id = RequestId::from(query.request_uri.clone());
24 let challenge = match state.repos.oauth.get_2fa_challenge(&twofa_request_id).await {
25 Ok(Some(c)) => c,
26 Ok(None) => {
27 return redirect_to_frontend_error(
28 "invalid_request",
29 "No 2FA challenge found. Please start over.",
30 );
31 }
32 Err(_) => {
33 return redirect_to_frontend_error(
34 "server_error",
35 "An error occurred. Please try again.",
36 );
37 }
38 };
39 if challenge.expires_at < Utc::now() {
40 let _ = state.repos.oauth.delete_2fa_challenge(challenge.id).await;
41 return redirect_to_frontend_error(
42 "invalid_request",
43 "2FA code has expired. Please start over.",
44 );
45 }
46 let _request_data = match state
47 .repos
48 .oauth
49 .get_authorization_request(&twofa_request_id)
50 .await
51 {
52 Ok(Some(d)) => d,
53 Ok(None) => {
54 return redirect_to_frontend_error(
55 "invalid_request",
56 "Authorization request not found. Please start over.",
57 );
58 }
59 Err(_) => {
60 return redirect_to_frontend_error(
61 "server_error",
62 "An error occurred. Please try again.",
63 );
64 }
65 };
66 let channel = query.channel.as_deref().unwrap_or("email");
67 redirect_see_other(&format!(
68 "/app/oauth/2fa?request_uri={}&channel={}",
69 url_encode(&query.request_uri),
70 url_encode(channel)
71 ))
72}
73
74pub async fn authorize_2fa_post(
75 State(state): State<AppState>,
76 _rate_limit: OAuthRateLimited<OAuthAuthorizeLimit>,
77 headers: HeaderMap,
78 client_ip: ClientIp,
79 Json(form): Json<Authorize2faSubmit>,
80) -> Response {
81 let json_error = |status: StatusCode, error: &str, description: &str| -> Response {
82 (
83 status,
84 Json(serde_json::json!({
85 "error": error,
86 "error_description": description
87 })),
88 )
89 .into_response()
90 };
91 let twofa_post_request_id = RequestId::from(form.request_uri.clone());
92 let request_data = match state
93 .repos
94 .oauth
95 .get_authorization_request(&twofa_post_request_id)
96 .await
97 {
98 Ok(Some(d)) => d,
99 Ok(None) => {
100 return json_error(
101 StatusCode::BAD_REQUEST,
102 "invalid_request",
103 "Authorization request not found.",
104 );
105 }
106 Err(_) => {
107 return json_error(
108 StatusCode::INTERNAL_SERVER_ERROR,
109 "server_error",
110 "An error occurred.",
111 );
112 }
113 };
114 if request_data.expires_at < Utc::now() {
115 let _ = state
116 .repos
117 .oauth
118 .delete_authorization_request(&twofa_post_request_id)
119 .await;
120 return json_error(
121 StatusCode::BAD_REQUEST,
122 "invalid_request",
123 "Authorization request has expired.",
124 );
125 }
126 let challenge = state
127 .repos
128 .oauth
129 .get_2fa_challenge(&twofa_post_request_id)
130 .await
131 .ok()
132 .flatten();
133 if let Some(challenge) = challenge {
134 if challenge.expires_at < Utc::now() {
135 let _ = state.repos.oauth.delete_2fa_challenge(challenge.id).await;
136 return json_error(
137 StatusCode::BAD_REQUEST,
138 "invalid_request",
139 "2FA code has expired. Please start over.",
140 );
141 }
142 if challenge.attempts >= MAX_2FA_ATTEMPTS {
143 let _ = state.repos.oauth.delete_2fa_challenge(challenge.id).await;
144 return json_error(
145 StatusCode::FORBIDDEN,
146 "access_denied",
147 "Too many failed attempts. Please start over.",
148 );
149 }
150 let code_valid: bool = form
151 .code
152 .trim()
153 .as_bytes()
154 .ct_eq(challenge.code.as_bytes())
155 .into();
156 if !code_valid {
157 let _ = state.repos.oauth.increment_2fa_attempts(challenge.id).await;
158 return json_error(
159 StatusCode::FORBIDDEN,
160 "invalid_code",
161 "Invalid verification code. Please try again.",
162 );
163 }
164 let _ = state.repos.oauth.delete_2fa_challenge(challenge.id).await;
165 let code = AuthorizationCode::generate();
166 let device_id = extract_device_cookie(&headers);
167 let twofa_totp_device_id = device_id.clone();
168 if state
169 .repos
170 .oauth
171 .update_authorization_request(
172 &twofa_post_request_id,
173 &challenge.did,
174 twofa_totp_device_id.as_ref(),
175 &code,
176 )
177 .await
178 .is_err()
179 {
180 return json_error(
181 StatusCode::INTERNAL_SERVER_ERROR,
182 "server_error",
183 "An error occurred. Please try again.",
184 );
185 }
186 let redirect_url = build_intermediate_redirect_url(
187 &request_data.parameters.redirect_uri,
188 &code.0,
189 request_data.parameters.state.as_deref(),
190 request_data.parameters.response_mode.map(|m| m.as_str()),
191 );
192 return Json(serde_json::json!({
193 "redirect_uri": redirect_url
194 }))
195 .into_response();
196 }
197 let did_str = match &request_data.did {
198 Some(d) => d.clone(),
199 None => {
200 return json_error(
201 StatusCode::BAD_REQUEST,
202 "invalid_request",
203 "No 2FA challenge found. Please start over.",
204 );
205 }
206 };
207 let did: tranquil_types::Did = match did_str.parse() {
208 Ok(d) => d,
209 Err(_) => {
210 return json_error(
211 StatusCode::BAD_REQUEST,
212 "invalid_request",
213 "Invalid DID format.",
214 );
215 }
216 };
217 if !tranquil_api::server::has_totp_enabled(&state, &did).await {
218 return json_error(
219 StatusCode::BAD_REQUEST,
220 "invalid_request",
221 "No 2FA challenge found. Please start over.",
222 );
223 }
224 let _rate_proof = match check_user_rate_limit::<TotpVerifyLimit>(&state, &did).await {
225 Ok(proof) => proof,
226 Err(_) => {
227 return json_error(
228 StatusCode::TOO_MANY_REQUESTS,
229 "RateLimitExceeded",
230 "Too many verification attempts. Please try again in a few minutes.",
231 );
232 }
233 };
234 let totp_valid =
235 tranquil_api::server::verify_totp_or_backup_for_user(&state, &did, &form.code).await;
236 if !totp_valid {
237 return json_error(
238 StatusCode::FORBIDDEN,
239 "invalid_code",
240 "Invalid verification code. Please try again.",
241 );
242 }
243 let mut device_id = extract_device_cookie(&headers);
244 let mut new_cookie: Option<String> = None;
245 if form.trust_device {
246 let trust_device_id = match &device_id {
247 Some(existing_id) => existing_id.clone(),
248 None => {
249 let new_device_id = DeviceId::generate();
250 let device_data = DeviceData {
251 session_id: SessionId::generate(),
252 user_agent: extract_user_agent(&headers),
253 ip_address: client_ip.into_string(),
254 last_seen_at: Utc::now(),
255 };
256 if state
257 .repos
258 .oauth
259 .create_device(&new_device_id, &device_data)
260 .await
261 .is_ok()
262 {
263 new_cookie = Some(make_device_cookie(&new_device_id));
264 device_id = Some(new_device_id.clone());
265 }
266 new_device_id
267 }
268 };
269 let _ = state
270 .repos
271 .oauth
272 .upsert_account_device(&did, &trust_device_id)
273 .await;
274 let _ =
275 tranquil_api::server::trust_device(state.repos.oauth.as_ref(), &trust_device_id, &did)
276 .await;
277 }
278 let requested_scope_str = request_data
279 .parameters
280 .scope
281 .as_deref()
282 .unwrap_or("atproto");
283 let requested_scopes: Vec<String> = requested_scope_str
284 .split_whitespace()
285 .map(|s| s.to_string())
286 .collect();
287 let needs_consent = should_show_consent(
288 state.repos.oauth.as_ref(),
289 &did,
290 &request_data.parameters.client_id,
291 &requested_scopes,
292 )
293 .await
294 .unwrap_or(true);
295 if needs_consent {
296 let consent_url = format!(
297 "/app/oauth/consent?request_uri={}",
298 url_encode(&form.request_uri)
299 );
300 if let Some(cookie) = new_cookie {
301 return (
302 StatusCode::OK,
303 [(SET_COOKIE, cookie)],
304 Json(serde_json::json!({"redirect_uri": consent_url})),
305 )
306 .into_response();
307 }
308 return Json(serde_json::json!({"redirect_uri": consent_url})).into_response();
309 }
310 let code = AuthorizationCode::generate();
311 let twofa_final_device_id = device_id.clone();
312 if state
313 .repos
314 .oauth
315 .update_authorization_request(
316 &twofa_post_request_id,
317 &did,
318 twofa_final_device_id.as_ref(),
319 &code,
320 )
321 .await
322 .is_err()
323 {
324 return json_error(
325 StatusCode::INTERNAL_SERVER_ERROR,
326 "server_error",
327 "An error occurred. Please try again.",
328 );
329 }
330 let redirect_url = build_intermediate_redirect_url(
331 &request_data.parameters.redirect_uri,
332 &code.0,
333 request_data.parameters.state.as_deref(),
334 request_data.parameters.response_mode.map(|m| m.as_str()),
335 );
336 if let Some(cookie) = new_cookie {
337 (
338 StatusCode::OK,
339 [(SET_COOKIE, cookie)],
340 Json(serde_json::json!({"redirect_uri": redirect_url})),
341 )
342 .into_response()
343 } else {
344 Json(serde_json::json!({"redirect_uri": redirect_url})).into_response()
345 }
346}