fix: use cookies

src/server.rs

@@ -2,10 +2,12 @@ use axum::{
     Json, RequestPartsExt, Router,
     extract::{FromRef, FromRequestParts, Path, Query, State},
     http::{HeaderValue, StatusCode, request::Parts},
+    response::IntoResponse,
     routing::{get, post},
 };
 use axum_extra::{
     TypedHeader,
+    extract::cookie::{Cookie, CookieJar, SameSite},
     headers::{Authorization, authorization::Bearer},
 };
 use chrono::Utc;
@@ -14,7 +16,7 @@ use sea_orm::{
 };
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
-use tower_http::cors::{AllowOrigin, Any, CorsLayer};
+use tower_http::cors::{AllowOrigin, CorsLayer};
 use uuid::Uuid;
 
 use crate::agent::Agent;
@@ -128,6 +130,7 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
         .route("/api/runs/recent", get(get_recent_runs))
         .route("/api/auth/callback", get(auth_callback))
         .route("/api/auth/refresh", post(auth_refresh))
+        .route("/api/auth/logout", post(auth_logout))
         .layer(cors)
         .with_state(state);
 
@@ -152,18 +155,30 @@ fn build_cors_layer() -> CorsLayer {
             .collect();
 
         if values.is_empty() {
-            AllowOrigin::any()
+            AllowOrigin::mirror_request()
         } else {
             AllowOrigin::list(values)
         }
     } else {
-        AllowOrigin::any()
+        AllowOrigin::mirror_request()
     };
 
     CorsLayer::new()
         .allow_origin(allow_origin)
-        .allow_methods(Any)
-        .allow_headers(Any)
+        .allow_methods([
+            axum::http::Method::GET,
+            axum::http::Method::POST,
+            axum::http::Method::PUT,
+            axum::http::Method::PATCH,
+            axum::http::Method::DELETE,
+            axum::http::Method::OPTIONS,
+        ])
+        .allow_headers([
+            axum::http::header::CONTENT_TYPE,
+            axum::http::header::AUTHORIZATION,
+            axum::http::header::ACCEPT,
+        ])
+        .allow_credentials(true)
 }
 
 async fn list_tasks(
@@ -277,7 +292,11 @@ async fn execute_agent_run(
 
     let (logs, answer, status) = match agent.run().await {
         Ok((logs, answer)) => (logs, answer, "completed".to_string()),
-        Err(e) => (format!("Execution failed: {}", e), None, "failed".to_string()),
+        Err(e) => (
+            format!("Execution failed: {}", e),
+            None,
+            "failed".to_string(),
+        ),
     };
 
     // Update with final logs and status
@@ -345,26 +364,27 @@ where
     async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
         let app_state = Arc::<AppState>::from_ref(state);
 
-        let TypedHeader(Authorization(bearer)) = parts
-            .extract::<TypedHeader<Authorization<Bearer>>>()
-            .await
-            .map_err(|_| {
-                (
-                    StatusCode::UNAUTHORIZED,
-                    "Missing or invalid Authorization header".to_string(),
-                )
-            })?;
+        let token = if let Ok(TypedHeader(Authorization(bearer))) =
+            parts.extract::<TypedHeader<Authorization<Bearer>>>().await
+        {
+            Some(bearer.token().to_string())
+        } else {
+            let jar = parts.extract::<CookieJar>().await.unwrap();
+            jar.get("access_token")
+                .map(|cookie| cookie.value().to_string())
+        };
 
-        let claims = app_state
-            .verifier
-            .verify(bearer.token())
-            .await
-            .map_err(|e| {
-                (
-                    StatusCode::UNAUTHORIZED,
-                    format!("Token verification failed: {}", e),
-                )
-            })?;
+        let token = token.ok_or((
+            StatusCode::UNAUTHORIZED,
+            "Missing or invalid access token".to_string(),
+        ))?;
+
+        let claims = app_state.verifier.verify(&token).await.map_err(|e| {
+            (
+                StatusCode::UNAUTHORIZED,
+                format!("Token verification failed: {}", e),
+            )
+        })?;
 
         Ok(AuthenticatedUser(claims))
     }
@@ -378,36 +398,110 @@ pub struct AuthCallbackQuery {
 
 #[derive(Deserialize)]
 struct RefreshRequest {
-    refresh_token: String,
+    refresh_token: Option<String>,
 }
 
 async fn auth_refresh(
     State(state): State<Arc<AppState>>,
+    jar: CookieJar,
     Json(payload): Json<RefreshRequest>,
-) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
-    state
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let refresh_token = payload
+        .refresh_token
+        .filter(|token| !token.is_empty())
+        .or_else(|| {
+            jar.get("refresh_token")
+                .map(|cookie| cookie.value().to_string())
+        })
+        .ok_or((
+            StatusCode::UNAUTHORIZED,
+            "Missing refresh token".to_string(),
+        ))?;
+
+    let data = state
         .authenticator
-        .refresh_token(payload.refresh_token)
+        .refresh_token(refresh_token)
         .await
-        .map(Json)
-        .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))
+        .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?;
+
+    let jar = update_auth_cookies(jar, &data);
+    Ok((jar, Json(data)))
 }
 
 async fn auth_callback(
     State(state): State<Arc<AppState>>,
+    jar: CookieJar,
     Query(query): Query<AuthCallbackQuery>,
-) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
-    state
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let data = state
         .authenticator
         .exchange_code(query.code, query.redirect_uri)
         .await
-        .map(Json)
         .map_err(|e| {
             (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 format!("Token exchange failed: {}", e),
             )
-        })
+        })?;
+
+    let jar = update_auth_cookies(jar, &data);
+    Ok((jar, Json(data)))
+}
+
+async fn auth_logout(jar: CookieJar) -> impl IntoResponse {
+    let jar = clear_auth_cookies(jar);
+    (jar, StatusCode::NO_CONTENT)
+}
+
+fn secure() -> bool {
+    std::env::var("COOKIE_SECURE")
+        .map(|value| value == "true")
+        .unwrap_or(false)
+}
+
+fn update_auth_cookies(jar: CookieJar, data: &serde_json::Value) -> CookieJar {
+    let access_token = data.get("access_token");
+    let refresh_token = data.get("refresh_token");
+
+    let mut jar = jar;
+
+    if let Some(token) = access_token.and_then(|t| t.as_str()) {
+        let cookie = Cookie::build(("access_token", token.to_owned()))
+            .path("/")
+            .http_only(true)
+            .same_site(SameSite::Lax)
+            .secure(secure())
+            .build();
+        jar = jar.add(cookie);
+    }
+
+    if let Some(token) = refresh_token.and_then(|t| t.as_str()) {
+        let cookie = Cookie::build(("refresh_token", token.to_owned()))
+            .path("/")
+            .http_only(true)
+            .same_site(SameSite::Lax)
+            .secure(secure())
+            .build();
+        jar = jar.add(cookie);
+    }
+
+    jar
+}
+
+fn clear_auth_cookies(jar: CookieJar) -> CookieJar {
+    let mut jar = jar;
+    for name in ["access_token", "refresh_token"] {
+        let cookie = Cookie::build((name, ""))
+            .path("/")
+            .http_only(true)
+            .same_site(SameSite::Lax)
+            .secure(secure())
+            .max_age(cookie::time::Duration::seconds(0))
+            .build();
+        jar = jar.add(cookie);
+    }
+
+    jar
 }
 
 async fn get_task(