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