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