diff --git a/Cargo.lock b/Cargo.lock index 1a9f8e6..44e16d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,6 +217,7 @@ dependencies = [ "axum", "axum-core", "bytes", + "cookie", "fastrand", "futures-util", "headers", @@ -325,6 +326,7 @@ dependencies = [ "axum-extra", "base64", "chrono", + "cookie", "dashmap", "jsonwebtoken", "migration", @@ -486,6 +488,17 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "cookie" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + [[package]] name = "core-foundation" version = "0.9.4" diff --git a/Cargo.toml b/Cargo.toml index e995879..27633df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,4 +22,5 @@ tokio-cron-scheduler = "0.15.1" dashmap = "6.1.0" jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] } base64 = "0.22.1" -axum-extra = { version = "0.9", features = ["typed-header"] } +axum-extra = { version = "0.9", features = ["typed-header", "cookie"] } +cookie = "0.18" diff --git a/frontend/src/main.js b/frontend/src/main.js index 3e5100f..a1a4c66 100644 --- a/frontend/src/main.js +++ b/frontend/src/main.js @@ -16,8 +16,7 @@ const state = { selectedRunId: null, currentView: 'dashboard', // 'dashboard' or 'task' isEditing: false, - token: localStorage.getItem('auth_token'), - refreshToken: localStorage.getItem('refresh_token') + isAuthenticated: false }; // DOM elements @@ -58,29 +57,15 @@ const presetBtns = document.querySelectorAll('.btn-preset'); // Wrapper for fetch to include Authorization header async function fetchWithAuth(url, options = {}) { - if (!state.token) { - showLogin(); - throw new Error('Not authenticated'); - } + let response = await fetch(url, { ...options, credentials: 'include' }); - const headers = { - ...options.headers, - 'Authorization': `Bearer ${state.token}` - }; - - let response = await fetch(url, { ...options, headers }); - - if (response.status === 401 && state.refreshToken) { + if (response.status === 401) { // Try to refresh token try { const success = await attemptTokenRefresh(); if (success) { - // Retry original request with new token - const newHeaders = { - ...options.headers, - 'Authorization': `Bearer ${state.token}` - }; - response = await fetch(url, { ...options, headers: newHeaders }); + // Retry original request + response = await fetch(url, { ...options, credentials: 'include' }); } } catch (error) { console.error('Token refresh failed:', error); @@ -100,20 +85,12 @@ async function attemptTokenRefresh() { const response = await fetch(`${API_URL}/auth/refresh`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ refresh_token: state.refreshToken }) + credentials: 'include', + body: JSON.stringify({ refresh_token: '' }) }); if (response.ok) { - const data = await response.json(); - if (data.access_token) { - state.token = data.access_token; - localStorage.setItem('auth_token', data.access_token); - if (data.refresh_token) { - state.refreshToken = data.refresh_token; - localStorage.setItem('refresh_token', data.refresh_token); - } - return true; - } + return true; } } catch (error) { console.error('Error during token refresh:', error); @@ -514,10 +491,12 @@ async function showLogin() { } async function logout() { - state.token = null; - state.refreshToken = null; - localStorage.removeItem('auth_token'); - localStorage.removeItem('refresh_token'); + try { + await fetch(`${API_URL}/auth/logout`, { method: 'POST', credentials: 'include' }); + } catch (error) { + console.error('Logout failed:', error); + } + state.isAuthenticated = false; showLogin(); } @@ -531,16 +510,13 @@ async function handleCallback() { loginOverlay.classList.add('hidden'); try { - const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`); + const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`, { + credentials: 'include' + }); const data = await response.json(); - if (data.access_token) { - state.token = data.access_token; - localStorage.setItem('auth_token', data.access_token); - if (data.refresh_token) { - state.refreshToken = data.refresh_token; - localStorage.setItem('refresh_token', data.refresh_token); - } + if (response.ok) { + state.isAuthenticated = true; callbackOverlay.classList.add('hidden'); appEl.classList.remove('hidden'); initializeApp(); @@ -564,11 +540,6 @@ logoutBtn.addEventListener('click', () => { }); async function initializeApp() { - if (!state.token) { - showLogin(); - return; - } - appEl.classList.remove('hidden'); loginOverlay.classList.add('hidden'); diff --git a/src/server.rs b/src/server.rs index 2033484..db9fac3 100644 --- a/src/server.rs +++ b/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> { .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 { let app_state = Arc::::from_ref(state); - let TypedHeader(Authorization(bearer)) = parts - .extract::>>() - .await - .map_err(|_| { - ( - StatusCode::UNAUTHORIZED, - "Missing or invalid Authorization header".to_string(), - ) - })?; + let token = if let Ok(TypedHeader(Authorization(bearer))) = + parts.extract::>>().await + { + Some(bearer.token().to_string()) + } else { + let jar = parts.extract::().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, } async fn auth_refresh( State(state): State>, + jar: CookieJar, Json(payload): Json, -) -> Result, (StatusCode, String)> { - state +) -> Result { + 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>, + jar: CookieJar, Query(query): Query, -) -> Result, (StatusCode, String)> { - state +) -> Result { + 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(