this repo has no description
lewis.moe
1use super::helpers::{create_access_token_with_delegation, verify_pkce};
2use super::types::{TokenGrant, TokenResponse, ValidatedTokenRequest};
3use crate::config::AuthConfig;
4use crate::delegation;
5use crate::oauth::{
6 AuthFlowState, ClientAuth, ClientMetadataCache, DPoPVerifier, OAuthError, RefreshToken,
7 TokenData, TokenId,
8 db::{self, RefreshTokenLookup},
9 scopes::expand_include_scopes,
10 verify_client_auth,
11};
12use crate::state::AppState;
13use axum::Json;
14use axum::http::HeaderMap;
15use chrono::{Duration, Utc};
16
17const ACCESS_TOKEN_EXPIRY_SECONDS: i64 = 300;
18const REFRESH_TOKEN_EXPIRY_DAYS_CONFIDENTIAL: i64 = 60;
19const REFRESH_TOKEN_EXPIRY_DAYS_PUBLIC: i64 = 14;
20
21pub async fn handle_authorization_code_grant(
22 state: AppState,
23 _headers: HeaderMap,
24 request: ValidatedTokenRequest,
25 dpop_proof: Option<String>,
26) -> Result<(HeaderMap, Json<TokenResponse>), OAuthError> {
27 tracing::info!(
28 has_dpop = dpop_proof.is_some(),
29 client_id = ?request.client_auth.client_id,
30 "Authorization code grant requested"
31 );
32 let (code, code_verifier, redirect_uri) = match request.grant {
33 TokenGrant::AuthorizationCode {
34 code,
35 code_verifier,
36 redirect_uri,
37 } => (code, code_verifier, redirect_uri),
38 _ => {
39 return Err(OAuthError::InvalidRequest(
40 "Expected authorization_code grant".to_string(),
41 ));
42 }
43 };
44 let auth_request = db::consume_authorization_request_by_code(&state.db, &code)
45 .await?
46 .ok_or_else(|| OAuthError::InvalidGrant("Invalid or expired code".to_string()))?;
47
48 let flow_state = AuthFlowState::from_request_data(&auth_request);
49 if flow_state.is_expired() {
50 return Err(OAuthError::InvalidGrant(
51 "Authorization code has expired".to_string(),
52 ));
53 }
54 if !flow_state.can_exchange() {
55 return Err(OAuthError::InvalidGrant(
56 "Authorization not completed".to_string(),
57 ));
58 }
59
60 if let Some(request_client_id) = &request.client_auth.client_id
61 && request_client_id != &auth_request.client_id
62 {
63 return Err(OAuthError::InvalidGrant("client_id mismatch".to_string()));
64 }
65 let did = flow_state.did().unwrap().to_string();
66 let client_metadata_cache = ClientMetadataCache::new(3600);
67 let client_metadata = client_metadata_cache.get(&auth_request.client_id).await?;
68 let client_auth = if let (Some(assertion), Some(assertion_type)) = (
69 &request.client_auth.client_assertion,
70 &request.client_auth.client_assertion_type,
71 ) {
72 if assertion_type != "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" {
73 return Err(OAuthError::InvalidClient(
74 "Unsupported client_assertion_type".to_string(),
75 ));
76 }
77 ClientAuth::PrivateKeyJwt {
78 client_assertion: assertion.clone(),
79 }
80 } else if let Some(secret) = &request.client_auth.client_secret {
81 ClientAuth::SecretPost {
82 client_secret: secret.clone(),
83 }
84 } else {
85 ClientAuth::None
86 };
87 verify_client_auth(&client_metadata_cache, &client_metadata, &client_auth).await?;
88 verify_pkce(&auth_request.parameters.code_challenge, &code_verifier)?;
89 if let Some(req_redirect_uri) = &redirect_uri
90 && req_redirect_uri != &auth_request.parameters.redirect_uri
91 {
92 return Err(OAuthError::InvalidGrant(
93 "redirect_uri mismatch".to_string(),
94 ));
95 }
96 let dpop_jkt = if let Some(proof) = &dpop_proof {
97 let config = AuthConfig::get();
98 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
99 let pds_hostname =
100 std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
101 let token_endpoint = format!("https://{}/oauth/token", pds_hostname);
102 let result = verifier.verify_proof(proof, "POST", &token_endpoint, None)?;
103 if !db::check_and_record_dpop_jti(&state.db, &result.jti).await? {
104 return Err(OAuthError::InvalidDpopProof(
105 "DPoP proof has already been used".to_string(),
106 ));
107 }
108 if let Some(expected_jkt) = &auth_request.parameters.dpop_jkt
109 && result.jkt.as_str() != expected_jkt
110 {
111 return Err(OAuthError::InvalidDpopProof(
112 "DPoP key binding mismatch".to_string(),
113 ));
114 }
115 Some(result.jkt.as_str().to_string())
116 } else if auth_request.parameters.dpop_jkt.is_some() || client_metadata.requires_dpop() {
117 return Err(OAuthError::UseDpopNonce(
118 DPoPVerifier::new(AuthConfig::get().dpop_secret().as_bytes()).generate_nonce(),
119 ));
120 } else {
121 None
122 };
123 let token_id = TokenId::generate();
124 let refresh_token = RefreshToken::generate();
125 let now = Utc::now();
126
127 let (raw_scope, controller_did) = if let Some(ref controller) = auth_request.controller_did {
128 let grant = delegation::get_delegation(&state.db, &did, controller)
129 .await
130 .ok()
131 .flatten();
132 let granted_scopes = grant.map(|g| g.granted_scopes).unwrap_or_default();
133 let requested = auth_request
134 .parameters
135 .scope
136 .as_deref()
137 .unwrap_or("atproto");
138 let intersected = delegation::intersect_scopes(requested, &granted_scopes);
139 (Some(intersected), Some(controller.clone()))
140 } else {
141 (auth_request.parameters.scope.clone(), None)
142 };
143
144 let final_scope = if let Some(ref scope) = raw_scope {
145 if scope.contains("include:") {
146 Some(expand_include_scopes(scope).await)
147 } else {
148 raw_scope
149 }
150 } else {
151 raw_scope
152 };
153
154 let access_token = create_access_token_with_delegation(
155 &token_id.0,
156 &did,
157 dpop_jkt.as_deref(),
158 final_scope.as_deref(),
159 controller_did.as_deref(),
160 )?;
161 let stored_client_auth = auth_request.client_auth.unwrap_or(ClientAuth::None);
162 let refresh_expiry_days = if matches!(stored_client_auth, ClientAuth::None) {
163 REFRESH_TOKEN_EXPIRY_DAYS_PUBLIC
164 } else {
165 REFRESH_TOKEN_EXPIRY_DAYS_CONFIDENTIAL
166 };
167 let mut stored_parameters = auth_request.parameters.clone();
168 stored_parameters.dpop_jkt = dpop_jkt.clone();
169 let token_data = TokenData {
170 did: did.clone(),
171 token_id: token_id.0.clone(),
172 created_at: now,
173 updated_at: now,
174 expires_at: now + Duration::days(refresh_expiry_days),
175 client_id: auth_request.client_id.clone(),
176 client_auth: stored_client_auth,
177 device_id: auth_request.device_id,
178 parameters: stored_parameters,
179 details: None,
180 code: None,
181 current_refresh_token: Some(refresh_token.0.clone()),
182 scope: final_scope.clone(),
183 controller_did: controller_did.clone(),
184 };
185 db::create_token(&state.db, &token_data).await?;
186 tracing::info!(
187 did = %did,
188 token_id = %token_id.0,
189 client_id = %auth_request.client_id,
190 "Authorization code grant completed, token created"
191 );
192 tokio::spawn({
193 let pool = state.db.clone();
194 let did_clone = did.clone();
195 async move {
196 if let Err(e) = db::enforce_token_limit_for_user(&pool, &did_clone).await {
197 tracing::warn!("Failed to enforce token limit for user: {:?}", e);
198 }
199 }
200 });
201 let mut response_headers = HeaderMap::new();
202 let config = AuthConfig::get();
203 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
204 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
205 Ok((
206 response_headers,
207 Json(TokenResponse {
208 access_token,
209 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
210 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
211 refresh_token: Some(refresh_token.0),
212 scope: final_scope,
213 sub: Some(did),
214 }),
215 ))
216}
217
218pub async fn handle_refresh_token_grant(
219 state: AppState,
220 _headers: HeaderMap,
221 request: ValidatedTokenRequest,
222 dpop_proof: Option<String>,
223) -> Result<(HeaderMap, Json<TokenResponse>), OAuthError> {
224 let refresh_token_str = match request.grant {
225 TokenGrant::RefreshToken { refresh_token } => refresh_token,
226 _ => {
227 return Err(OAuthError::InvalidRequest(
228 "Expected refresh_token grant".to_string(),
229 ));
230 }
231 };
232 let token_prefix = &refresh_token_str[..std::cmp::min(16, refresh_token_str.len())];
233 tracing::info!(
234 refresh_token_prefix = %token_prefix,
235 has_dpop = dpop_proof.is_some(),
236 "Refresh token grant requested"
237 );
238
239 let lookup = db::lookup_refresh_token(&state.db, &refresh_token_str).await?;
240 let token_state = lookup.state();
241 tracing::debug!(state = %token_state, "Refresh token state");
242
243 let (db_id, token_data) = match lookup {
244 RefreshTokenLookup::Valid { db_id, token_data } => (db_id, token_data),
245 RefreshTokenLookup::InGracePeriod {
246 db_id: _,
247 token_data,
248 rotated_at,
249 } => {
250 tracing::info!(
251 refresh_token_prefix = %token_prefix,
252 rotated_at = %rotated_at,
253 "Refresh token reuse within grace period, returning existing tokens"
254 );
255 let dpop_jkt = token_data.parameters.dpop_jkt.as_deref();
256 let access_token = create_access_token_with_delegation(
257 &token_data.token_id,
258 &token_data.did,
259 dpop_jkt,
260 token_data.scope.as_deref(),
261 token_data.controller_did.as_deref(),
262 )?;
263 let mut response_headers = HeaderMap::new();
264 let config = AuthConfig::get();
265 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
266 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
267 return Ok((
268 response_headers,
269 Json(TokenResponse {
270 access_token,
271 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
272 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
273 refresh_token: token_data.current_refresh_token,
274 scope: token_data.scope,
275 sub: Some(token_data.did),
276 }),
277 ));
278 }
279 RefreshTokenLookup::Used { original_token_id } => {
280 tracing::warn!(
281 refresh_token_prefix = %token_prefix,
282 "Refresh token reuse detected, revoking token family"
283 );
284 db::delete_token_family(&state.db, original_token_id).await?;
285 return Err(OAuthError::InvalidGrant(
286 "Refresh token reuse detected, token family revoked".to_string(),
287 ));
288 }
289 RefreshTokenLookup::Expired { db_id } => {
290 tracing::warn!(refresh_token_prefix = %token_prefix, "Refresh token has expired");
291 db::delete_token_family(&state.db, db_id).await?;
292 return Err(OAuthError::InvalidGrant(
293 "Refresh token has expired".to_string(),
294 ));
295 }
296 RefreshTokenLookup::NotFound => {
297 tracing::warn!(refresh_token_prefix = %token_prefix, "Refresh token not found");
298 return Err(OAuthError::InvalidGrant(
299 "Invalid refresh token".to_string(),
300 ));
301 }
302 };
303 let dpop_jkt = if let Some(proof) = &dpop_proof {
304 let config = AuthConfig::get();
305 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
306 let pds_hostname =
307 std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
308 let token_endpoint = format!("https://{}/oauth/token", pds_hostname);
309 let result = verifier.verify_proof(proof, "POST", &token_endpoint, None)?;
310 if !db::check_and_record_dpop_jti(&state.db, &result.jti).await? {
311 return Err(OAuthError::InvalidDpopProof(
312 "DPoP proof has already been used".to_string(),
313 ));
314 }
315 if let Some(expected_jkt) = &token_data.parameters.dpop_jkt
316 && result.jkt.as_str() != expected_jkt
317 {
318 return Err(OAuthError::InvalidDpopProof(
319 "DPoP key binding mismatch".to_string(),
320 ));
321 }
322 Some(result.jkt.as_str().to_string())
323 } else if token_data.parameters.dpop_jkt.is_some() {
324 return Err(OAuthError::InvalidRequest(
325 "DPoP proof required".to_string(),
326 ));
327 } else {
328 None
329 };
330 let new_refresh_token = RefreshToken::generate();
331 let refresh_expiry_days = if matches!(token_data.client_auth, ClientAuth::None) {
332 REFRESH_TOKEN_EXPIRY_DAYS_PUBLIC
333 } else {
334 REFRESH_TOKEN_EXPIRY_DAYS_CONFIDENTIAL
335 };
336 let new_expires_at = Utc::now() + Duration::days(refresh_expiry_days);
337 db::rotate_token(&state.db, db_id, &new_refresh_token.0, new_expires_at).await?;
338 tracing::info!(
339 did = %token_data.did,
340 new_expires_at = %new_expires_at,
341 "Refresh token rotated successfully"
342 );
343 let access_token = create_access_token_with_delegation(
344 &token_data.token_id,
345 &token_data.did,
346 dpop_jkt.as_deref(),
347 token_data.scope.as_deref(),
348 token_data.controller_did.as_deref(),
349 )?;
350 let mut response_headers = HeaderMap::new();
351 let config = AuthConfig::get();
352 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
353 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
354 Ok((
355 response_headers,
356 Json(TokenResponse {
357 access_token,
358 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
359 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
360 refresh_token: Some(new_refresh_token.0),
361 scope: token_data.scope,
362 sub: Some(token_data.did),
363 }),
364 ))
365}