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 let token_id = TokenId::generate();
120 let refresh_token = RefreshToken::generate();
121 let now = Utc::now();
122
123 let (raw_scope, controller_did) = if let Some(ref controller) = auth_request.controller_did {
124 let grant = delegation::get_delegation(&state.db, &did, controller)
125 .await
126 .ok()
127 .flatten();
128 let granted_scopes = grant.map(|g| g.granted_scopes).unwrap_or_default();
129 let requested = auth_request
130 .parameters
131 .scope
132 .as_deref()
133 .unwrap_or("atproto");
134 let intersected = delegation::intersect_scopes(requested, &granted_scopes);
135 (Some(intersected), Some(controller.clone()))
136 } else {
137 (auth_request.parameters.scope.clone(), None)
138 };
139
140 let final_scope = if let Some(ref scope) = raw_scope {
141 if scope.contains("include:") {
142 Some(expand_include_scopes(scope).await)
143 } else {
144 raw_scope
145 }
146 } else {
147 raw_scope
148 };
149
150 let access_token = create_access_token_with_delegation(
151 &token_id.0,
152 &did,
153 dpop_jkt.as_deref(),
154 final_scope.as_deref(),
155 controller_did.as_deref(),
156 )?;
157 let stored_client_auth = auth_request.client_auth.unwrap_or(ClientAuth::None);
158 let refresh_expiry_days = if matches!(stored_client_auth, ClientAuth::None) {
159 REFRESH_TOKEN_EXPIRY_DAYS_PUBLIC
160 } else {
161 REFRESH_TOKEN_EXPIRY_DAYS_CONFIDENTIAL
162 };
163 let mut stored_parameters = auth_request.parameters.clone();
164 stored_parameters.dpop_jkt = dpop_jkt.clone();
165 let token_data = TokenData {
166 did: did.clone(),
167 token_id: token_id.0.clone(),
168 created_at: now,
169 updated_at: now,
170 expires_at: now + Duration::days(refresh_expiry_days),
171 client_id: auth_request.client_id.clone(),
172 client_auth: stored_client_auth,
173 device_id: auth_request.device_id,
174 parameters: stored_parameters,
175 details: None,
176 code: None,
177 current_refresh_token: Some(refresh_token.0.clone()),
178 scope: final_scope.clone(),
179 controller_did: controller_did.clone(),
180 };
181 db::create_token(&state.db, &token_data).await?;
182 tokio::spawn({
183 let pool = state.db.clone();
184 let did_clone = did.clone();
185 async move {
186 if let Err(e) = db::enforce_token_limit_for_user(&pool, &did_clone).await {
187 tracing::warn!("Failed to enforce token limit for user: {:?}", e);
188 }
189 }
190 });
191 let mut response_headers = HeaderMap::new();
192 let config = AuthConfig::get();
193 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
194 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
195 Ok((
196 response_headers,
197 Json(TokenResponse {
198 access_token,
199 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
200 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
201 refresh_token: Some(refresh_token.0),
202 scope: final_scope,
203 sub: Some(did),
204 }),
205 ))
206}
207
208pub async fn handle_refresh_token_grant(
209 state: AppState,
210 _headers: HeaderMap,
211 request: ValidatedTokenRequest,
212 dpop_proof: Option<String>,
213) -> Result<(HeaderMap, Json<TokenResponse>), OAuthError> {
214 let refresh_token_str = match request.grant {
215 TokenGrant::RefreshToken { refresh_token } => refresh_token,
216 _ => {
217 return Err(OAuthError::InvalidRequest(
218 "Expected refresh_token grant".to_string(),
219 ));
220 }
221 };
222 let token_prefix = &refresh_token_str[..std::cmp::min(16, refresh_token_str.len())];
223 tracing::info!(
224 refresh_token_prefix = %token_prefix,
225 has_dpop = dpop_proof.is_some(),
226 "Refresh token grant requested"
227 );
228
229 let lookup = db::lookup_refresh_token(&state.db, &refresh_token_str).await?;
230 let token_state = lookup.state();
231 tracing::debug!(state = %token_state, "Refresh token state");
232
233 let (db_id, token_data) = match lookup {
234 RefreshTokenLookup::Valid { db_id, token_data } => (db_id, token_data),
235 RefreshTokenLookup::InGracePeriod {
236 db_id: _,
237 token_data,
238 rotated_at,
239 } => {
240 tracing::info!(
241 refresh_token_prefix = %token_prefix,
242 rotated_at = %rotated_at,
243 "Refresh token reuse within grace period, returning existing tokens"
244 );
245 let dpop_jkt = token_data.parameters.dpop_jkt.as_deref();
246 let access_token = create_access_token_with_delegation(
247 &token_data.token_id,
248 &token_data.did,
249 dpop_jkt,
250 token_data.scope.as_deref(),
251 token_data.controller_did.as_deref(),
252 )?;
253 let mut response_headers = HeaderMap::new();
254 let config = AuthConfig::get();
255 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
256 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
257 return Ok((
258 response_headers,
259 Json(TokenResponse {
260 access_token,
261 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
262 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
263 refresh_token: token_data.current_refresh_token,
264 scope: token_data.scope,
265 sub: Some(token_data.did),
266 }),
267 ));
268 }
269 RefreshTokenLookup::Used { original_token_id } => {
270 tracing::warn!(
271 refresh_token_prefix = %token_prefix,
272 "Refresh token reuse detected, revoking token family"
273 );
274 db::delete_token_family(&state.db, original_token_id).await?;
275 return Err(OAuthError::InvalidGrant(
276 "Refresh token reuse detected, token family revoked".to_string(),
277 ));
278 }
279 RefreshTokenLookup::Expired { db_id } => {
280 tracing::warn!(refresh_token_prefix = %token_prefix, "Refresh token has expired");
281 db::delete_token_family(&state.db, db_id).await?;
282 return Err(OAuthError::InvalidGrant(
283 "Refresh token has expired".to_string(),
284 ));
285 }
286 RefreshTokenLookup::NotFound => {
287 tracing::warn!(refresh_token_prefix = %token_prefix, "Refresh token not found");
288 return Err(OAuthError::InvalidGrant(
289 "Invalid refresh token".to_string(),
290 ));
291 }
292 };
293 let dpop_jkt = if let Some(proof) = &dpop_proof {
294 let config = AuthConfig::get();
295 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
296 let pds_hostname =
297 std::env::var("PDS_HOSTNAME").unwrap_or_else(|_| "localhost".to_string());
298 let token_endpoint = format!("https://{}/oauth/token", pds_hostname);
299 let result = verifier.verify_proof(proof, "POST", &token_endpoint, None)?;
300 if !db::check_and_record_dpop_jti(&state.db, &result.jti).await? {
301 return Err(OAuthError::InvalidDpopProof(
302 "DPoP proof has already been used".to_string(),
303 ));
304 }
305 if let Some(expected_jkt) = &token_data.parameters.dpop_jkt
306 && result.jkt.as_str() != expected_jkt
307 {
308 return Err(OAuthError::InvalidDpopProof(
309 "DPoP key binding mismatch".to_string(),
310 ));
311 }
312 Some(result.jkt.as_str().to_string())
313 } else if token_data.parameters.dpop_jkt.is_some() {
314 return Err(OAuthError::InvalidRequest(
315 "DPoP proof required".to_string(),
316 ));
317 } else {
318 None
319 };
320 let new_token_id = TokenId::generate();
321 let new_refresh_token = RefreshToken::generate();
322 let refresh_expiry_days = if matches!(token_data.client_auth, ClientAuth::None) {
323 REFRESH_TOKEN_EXPIRY_DAYS_PUBLIC
324 } else {
325 REFRESH_TOKEN_EXPIRY_DAYS_CONFIDENTIAL
326 };
327 let new_expires_at = Utc::now() + Duration::days(refresh_expiry_days);
328 db::rotate_token(
329 &state.db,
330 db_id,
331 &new_token_id.0,
332 &new_refresh_token.0,
333 new_expires_at,
334 )
335 .await?;
336 tracing::info!(
337 did = %token_data.did,
338 new_expires_at = %new_expires_at,
339 "Refresh token rotated successfully"
340 );
341 let access_token = create_access_token_with_delegation(
342 &new_token_id.0,
343 &token_data.did,
344 dpop_jkt.as_deref(),
345 token_data.scope.as_deref(),
346 token_data.controller_did.as_deref(),
347 )?;
348 let mut response_headers = HeaderMap::new();
349 let config = AuthConfig::get();
350 let verifier = DPoPVerifier::new(config.dpop_secret().as_bytes());
351 response_headers.insert("DPoP-Nonce", verifier.generate_nonce().parse().unwrap());
352 Ok((
353 response_headers,
354 Json(TokenResponse {
355 access_token,
356 token_type: if dpop_jkt.is_some() { "DPoP" } else { "Bearer" }.to_string(),
357 expires_in: ACCESS_TOKEN_EXPIRY_SECONDS as u64,
358 refresh_token: Some(new_refresh_token.0),
359 scope: token_data.scope,
360 sub: Some(token_data.did),
361 }),
362 ))
363}