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