1use axum::{
10 extract::{Path, Query, State},
11 http::StatusCode,
12 response::Json,
13};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::sync::Arc;
17use tokio::sync::RwLock;
18
19use crate::auth::token_lifecycle::{
20 extract_token_id, ClockSkewState, KeyRotationState, RevokedToken, TokenLifecycleManager,
21 TokenRevocationStore,
22};
23
24#[derive(Clone)]
26pub struct TokenLifecycleState {
27 pub manager: Arc<TokenLifecycleManager>,
29}
30
31#[derive(Debug, Deserialize)]
33pub struct RevokeTokenRequest {
34 pub token: Option<String>,
36 pub token_id: Option<String>,
38 pub user_id: Option<String>,
40 pub reason: String,
42}
43
44#[derive(Debug, Deserialize)]
46pub struct RevokeUserTokensRequest {
47 pub user_id: String,
49 pub reason: String,
51}
52
53#[derive(Debug, Deserialize)]
55pub struct RotateKeyRequest {
56 pub new_key_id: String,
58 pub grace_period_seconds: Option<i64>,
60}
61
62#[derive(Debug, Deserialize)]
64pub struct ClockSkewRequest {
65 pub skew_seconds: i64,
67 pub duration_seconds: Option<u64>,
69}
70
71#[derive(Debug, Deserialize)]
73pub struct ForceRefreshFailureRequest {
74 pub user_id: String,
76 pub failure_type: String,
78}
79
80#[derive(Debug, Deserialize)]
82pub struct RevokeMidSessionRequest {
83 pub user_id: String,
85 pub delay_seconds: u64,
87}
88
89pub async fn revoke_token(
91 State(state): State<TokenLifecycleState>,
92 Json(request): Json<RevokeTokenRequest>,
93) -> Result<Json<serde_json::Value>, StatusCode> {
94 let token_id = if let Some(token) = request.token {
95 extract_token_id(&token)
96 } else if let Some(tid) = request.token_id {
97 tid
98 } else {
99 return Err(StatusCode::BAD_REQUEST);
100 };
101
102 state
103 .manager
104 .revocation
105 .revoke_token(
106 token_id.clone(),
107 request.user_id,
108 request.reason,
109 None, )
111 .await;
112
113 Ok(Json(serde_json::json!({
114 "success": true,
115 "token_id": token_id,
116 "message": "Token revoked successfully"
117 })))
118}
119
120pub async fn revoke_user_tokens(
122 State(state): State<TokenLifecycleState>,
123 Json(request): Json<RevokeUserTokensRequest>,
124) -> Result<Json<serde_json::Value>, StatusCode> {
125 state
126 .manager
127 .revocation
128 .revoke_user_tokens(request.user_id.clone(), request.reason)
129 .await;
130
131 Ok(Json(serde_json::json!({
132 "success": true,
133 "user_id": request.user_id,
134 "message": "All user tokens revoked successfully"
135 })))
136}
137
138pub async fn get_token_status(
140 State(state): State<TokenLifecycleState>,
141 Query(params): Query<HashMap<String, String>>,
142) -> Result<Json<serde_json::Value>, StatusCode> {
143 let token_id = if let Some(token) = params.get("token") {
144 extract_token_id(token)
145 } else if let Some(tid) = params.get("token_id") {
146 tid.clone()
147 } else {
148 return Err(StatusCode::BAD_REQUEST);
149 };
150
151 if let Some(revoked) = state.manager.revocation.get_revocation_status(&token_id).await {
152 Ok(Json(serde_json::json!({
153 "revoked": true,
154 "revoked_at": revoked.revoked_at,
155 "reason": revoked.reason
156 })))
157 } else {
158 Ok(Json(serde_json::json!({
159 "revoked": false
160 })))
161 }
162}
163
164pub async fn rotate_keys(
166 State(state): State<TokenLifecycleState>,
167 Json(request): Json<RotateKeyRequest>,
168) -> Result<Json<serde_json::Value>, StatusCode> {
169 state
174 .manager
175 .key_rotation
176 .rotate_key(request.new_key_id.clone())
177 .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
178
179 Ok(Json(serde_json::json!({
180 "success": true,
181 "new_key_id": request.new_key_id,
182 "message": "Key rotation initiated. Update OIDC configuration to use the new key."
183 })))
184}
185
186pub async fn get_active_keys(
188 State(state): State<TokenLifecycleState>,
189) -> Json<serde_json::Value> {
190 let keys = state.manager.key_rotation.get_active_keys().await;
191 Json(serde_json::json!({
192 "keys": keys.iter().map(|k| serde_json::json!({
193 "kid": k.kid,
194 "created_at": k.created_at,
195 "inactive_at": k.inactive_at,
196 "is_primary": k.is_primary
197 })).collect::<Vec<_>>()
198 }))
199}
200
201pub async fn set_clock_skew(
203 State(state): State<TokenLifecycleState>,
204 Json(request): Json<ClockSkewRequest>,
205) -> Result<Json<serde_json::Value>, StatusCode> {
206 state.manager.clock_skew.set_skew(request.skew_seconds).await;
207
208 if let Some(duration) = request.duration_seconds {
210 let state_clone = state.clone();
211 let skew_value = request.skew_seconds;
212 tokio::spawn(async move {
213 tokio::time::sleep(tokio::time::Duration::from_secs(duration)).await;
214 state_clone.manager.clock_skew.set_skew(0).await;
215 });
216 }
217
218 Ok(Json(serde_json::json!({
219 "success": true,
220 "skew_seconds": request.skew_seconds,
221 "message": "Clock skew set successfully"
222 })))
223}
224
225pub async fn get_clock_skew(
227 State(state): State<TokenLifecycleState>,
228) -> Json<serde_json::Value> {
229 let skew = state.manager.clock_skew.get_skew().await;
230 let adjusted_time = state.manager.clock_skew.get_adjusted_time().await;
231 let server_time = chrono::Utc::now().timestamp();
232
233 Json(serde_json::json!({
234 "skew_seconds": skew,
235 "server_time": server_time,
236 "adjusted_time": adjusted_time
237 }))
238}
239
240pub async fn force_refresh_failure(
242 State(state): State<TokenLifecycleState>,
243 Json(request): Json<ForceRefreshFailureRequest>,
244) -> Result<Json<serde_json::Value>, StatusCode> {
245 let reason = format!("test_scenario:{}", request.failure_type);
247 state
248 .manager
249 .revocation
250 .revoke_user_tokens(request.user_id.clone(), reason)
251 .await;
252
253 Ok(Json(serde_json::json!({
254 "success": true,
255 "user_id": request.user_id,
256 "failure_type": request.failure_type,
257 "message": "Refresh token failure simulated"
258 })))
259}
260
261pub async fn revoke_mid_session(
263 State(state): State<TokenLifecycleState>,
264 Json(request): Json<RevokeMidSessionRequest>,
265) -> Result<Json<serde_json::Value>, StatusCode> {
266 let state_clone = state.clone();
267 let user_id = request.user_id.clone();
268 let delay = request.delay_seconds;
269
270 tokio::spawn(async move {
271 tokio::time::sleep(tokio::time::Duration::from_secs(delay)).await;
272 state_clone
273 .manager
274 .revocation
275 .revoke_user_tokens(user_id, "mid_session_revocation".to_string())
276 .await;
277 });
278
279 Ok(Json(serde_json::json!({
280 "success": true,
281 "user_id": request.user_id,
282 "delay_seconds": request.delay_seconds,
283 "message": format!("Token will be revoked in {} seconds", request.delay_seconds)
284 })))
285}
286
287pub fn token_lifecycle_router(state: TokenLifecycleState) -> axum::Router {
289 use axum::routing::{get, post};
290
291 axum::Router::new()
292 .route("/tokens/revoke", post(revoke_token))
293 .route("/tokens/revoke/user", post(revoke_user_tokens))
294 .route("/tokens/status", get(get_token_status))
295 .route("/keys/rotate", post(rotate_keys))
296 .route("/keys/active", get(get_active_keys))
297 .route("/clock/skew", post(set_clock_skew))
298 .route("/clock/skew", get(get_clock_skew))
299 .route("/test/force-refresh-failure", post(force_refresh_failure))
300 .route("/test/revoke-mid-session", post(revoke_mid_session))
301 .with_state(state)
302}
303