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