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