refactoring

src/server/auth.rs

@@ -0,0 +1,178 @@
+use axum::{
+    Json, RequestPartsExt,
+    extract::{FromRef, FromRequestParts, Query, State},
+    http::{StatusCode, request::Parts},
+    response::IntoResponse,
+};
+use axum_extra::{
+    TypedHeader,
+    extract::cookie::{Cookie, CookieJar, SameSite},
+    headers::{Authorization, authorization::Bearer},
+};
+use std::sync::Arc;
+
+use super::AppState;
+use crate::domain::auth::Claims;
+
+#[allow(dead_code)]
+pub struct AuthenticatedUser(pub Claims);
+
+#[axum::async_trait]
+impl<S> FromRequestParts<S> for AuthenticatedUser
+where
+    Arc<AppState>: axum::extract::FromRef<S>,
+    S: Send + Sync,
+{
+    type Rejection = (StatusCode, String);
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        let app_state = Arc::<AppState>::from_ref(state);
+
+        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 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))
+    }
+}
+
+#[derive(serde::Deserialize)]
+pub struct AuthCallbackQuery {
+    pub code: String,
+    pub redirect_uri: String,
+}
+
+#[derive(serde::Deserialize)]
+pub struct RefreshRequest {
+    pub refresh_token: Option<String>,
+}
+
+pub async fn auth_refresh(
+    State(state): State<Arc<AppState>>,
+    jar: CookieJar,
+    Json(payload): Json<RefreshRequest>,
+) -> 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(refresh_token)
+        .await
+        .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?;
+
+    let jar = update_auth_cookies(jar, &data, &state.config);
+    Ok((jar, Json(data)))
+}
+
+pub async fn auth_callback(
+    State(state): State<Arc<AppState>>,
+    jar: CookieJar,
+    Query(query): Query<AuthCallbackQuery>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let data = state
+        .authenticator
+        .exchange_code(query.code, query.redirect_uri)
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Token exchange failed: {}", e),
+            )
+        })?;
+
+    let jar = update_auth_cookies(jar, &data, &state.config);
+    Ok((jar, Json(data)))
+}
+
+pub async fn auth_logout(State(state): State<Arc<AppState>>, jar: CookieJar) -> impl IntoResponse {
+    let jar = clear_auth_cookies(jar, &state.config);
+    (jar, StatusCode::NO_CONTENT)
+}
+
+pub async fn auth_session(user: AuthenticatedUser) -> Json<serde_json::Value> {
+    Json(serde_json::json!({
+        "authenticated": true,
+        "user": user.0
+    }))
+}
+
+fn secure(config: &crate::config::Config) -> bool {
+    config.cookie_secure
+}
+
+pub fn update_auth_cookies(
+    jar: CookieJar,
+    data: &serde_json::Value,
+    config: &crate::config::Config,
+) -> 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(config))
+            .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(config))
+            .build();
+        jar = jar.add(cookie);
+    }
+
+    jar
+}
+
+pub fn clear_auth_cookies(jar: CookieJar, config: &crate::config::Config) -> 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(config))
+            .max_age(cookie::time::Duration::seconds(0))
+            .build();
+        jar = jar.add(cookie);
+    }
+
+    jar
+}

src/server/mod.rs

@@ -0,0 +1,135 @@
+pub mod auth;
+pub mod tasks;
+
+use axum::{
+    Router,
+    http::HeaderValue,
+    routing::{get, post},
+};
+use migration::{Migrator, MigratorTrait};
+use sea_orm::{Database, DatabaseConnection, EntityTrait};
+use std::sync::Arc;
+use tower_http::cors::{AllowOrigin, CorsLayer};
+
+use crate::entities::task::Entity as Task;
+use crate::scheduler::Scheduler;
+
+#[derive(Clone)]
+pub struct AppState {
+    pub db: DatabaseConnection,
+    pub scheduler: Arc<Scheduler>,
+    pub config: Arc<crate::config::Config>,
+    pub verifier: Arc<crate::domain::auth::JwksVerifier>,
+    pub authenticator: Arc<crate::domain::auth::Authenticator>,
+}
+
+pub async fn start(config: crate::config::Config) -> Result<(), Box<dyn std::error::Error>> {
+    let db = Database::connect(&config.database_url).await?;
+    Migrator::up(&db, None).await?;
+
+    let config = Arc::new(config);
+
+    let scheduler = Arc::new(Scheduler::new(db.clone(), config.clone()).await?);
+
+    // Load existing scheduled tasks
+    let existing_tasks = Task::find().all(&db).await?;
+    for task in existing_tasks {
+        if let Some(cron) = task.cron {
+            let _ = scheduler.add_task_job(task.id, &cron).await;
+        }
+    }
+
+    let verifier = Arc::new(
+        crate::domain::auth::JwksVerifier::new(
+            config.authentik_issuer.clone(),
+            config.authentik_client_id.clone(),
+        )
+        .await?,
+    );
+    let authenticator = Arc::new(
+        crate::domain::auth::Authenticator::new(
+            config.authentik_issuer.clone(),
+            config.authentik_client_id.clone(),
+            config.authentik_client_secret.clone(),
+        )
+        .await?,
+    );
+
+    let state = Arc::new(AppState {
+        db,
+        scheduler,
+        config: config.clone(),
+        verifier,
+        authenticator,
+    });
+
+    let cors = build_cors_layer(&config);
+
+    let app = Router::new()
+        .route("/api/tasks", post(tasks::create_task).get(tasks::list_tasks))
+        .route("/api/tasks/:id", get(tasks::get_task).put(tasks::update_task))
+        .route("/api/tasks/:id/runs", post(tasks::rerun_task))
+        .route("/api/runs/recent", get(tasks::get_recent_runs))
+        .route("/api/auth/session", get(auth::auth_session))
+        .route("/api/auth/callback", get(auth::auth_callback))
+        .route("/api/auth/refresh", post(auth::auth_refresh))
+        .route("/api/auth/logout", post(auth::auth_logout))
+        .layer(cors)
+        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
+            axum::http::header::CONTENT_SECURITY_POLICY,
+            HeaderValue::from_static("default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data:;"),
+        ))
+        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
+            axum::http::header::X_CONTENT_TYPE_OPTIONS,
+            HeaderValue::from_static("nosniff"),
+        ))
+        .layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
+            axum::http::header::REFERRER_POLICY,
+            HeaderValue::from_static("strict-origin-when-cross-origin"),
+        ))
+        .layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit
+        .with_state(state);
+
+    let addr = format!("0.0.0.0:{}", config.port);
+    let listener = tokio::net::TcpListener::bind(&addr).await?;
+    println!("Server running on http://localhost:{}", config.port);
+    axum::serve(listener, app).await?;
+
+    Ok(())
+}
+
+fn build_cors_layer(config: &crate::config::Config) -> CorsLayer {
+    let allow_origin = if let Some(origins) = &config.cors_allowed_origins {
+        let values: Vec<HeaderValue> = origins
+            .split(',')
+            .map(|origin| origin.trim())
+            .filter(|origin| !origin.is_empty())
+            .filter_map(|origin| HeaderValue::from_str(origin).ok())
+            .collect();
+
+        if values.is_empty() {
+            AllowOrigin::mirror_request()
+        } else {
+            AllowOrigin::list(values)
+        }
+    } else {
+        AllowOrigin::mirror_request()
+    };
+
+    CorsLayer::new()
+        .allow_origin(allow_origin)
+        .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)
+}

src/server/tasks.rs

@@ -0,0 +1,179 @@
+use axum::{
+    Json,
+    extract::{Path, State},
+    http::StatusCode,
+};
+use sea_orm::{EntityTrait, QueryOrder, QuerySelect};
+use std::sync::Arc;
+use uuid::Uuid;
+
+use super::AppState;
+use super::auth::AuthenticatedUser;
+use crate::domain::tasks::{
+    self, CreateTaskRequest, RecentRunResponse, TaskResponse, UpdateTaskRequest,
+};
+
+pub async fn list_tasks(
+    _user: AuthenticatedUser,
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<Vec<TaskResponse>>, (StatusCode, String)> {
+    let tasks = crate::entities::task::Entity::find()
+        .find_with_related(crate::entities::task_run::Entity)
+        .order_by_desc(crate::entities::task::Column::CreatedAt)
+        .all(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let response = tasks
+        .into_iter()
+        .map(|(t, mut runs)| {
+            runs.sort_by(|a, b| b.created_at.cmp(&a.created_at));
+            TaskResponse {
+                id: t.id,
+                goal: t.goal,
+                cron: t.cron,
+                created_at: t.created_at,
+                runs: runs
+                    .into_iter()
+                    .map(|r| tasks::TaskRunResponse {
+                        id: r.id,
+                        status: r.status,
+                        logs: r.logs,
+                        answer: r.answer,
+                        created_at: r.created_at,
+                    })
+                    .collect(),
+            }
+        })
+        .collect();
+
+    Ok(Json(response))
+}
+
+pub async fn create_task(
+    _user: AuthenticatedUser,
+    State(state): State<Arc<AppState>>,
+    Json(payload): Json<CreateTaskRequest>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let task_id = Uuid::new_v4();
+
+    let new_task = crate::entities::task::ActiveModel {
+        id: sea_orm::Set(task_id),
+        goal: sea_orm::Set(payload.goal.clone()),
+        cron: sea_orm::Set(payload.cron.clone()),
+        created_at: sea_orm::Set(chrono::Utc::now().into()),
+    };
+
+    use sea_orm::ActiveModelTrait;
+    new_task
+        .insert(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    if let Some(cron) = &payload.cron {
+        let _ = state.scheduler.add_task_job(task_id, cron).await;
+    } else {
+        let _ = state.scheduler.remove_task_job(task_id).await;
+    }
+
+    tasks::get_task_inner(task_id, &state.db)
+        .await
+        .map(Json)
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
+}
+
+pub async fn rerun_task(
+    _user: AuthenticatedUser,
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<Uuid>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let task = crate::entities::task::Entity::find_by_id(id)
+        .one(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?;
+
+    tasks::execute_agent_run(
+        &state.db,
+        &state.scheduler,
+        &state.config,
+        task.id,
+        task.goal,
+    )
+    .await
+    .map(Json)
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
+}
+
+pub async fn update_task(
+    _user: AuthenticatedUser,
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<Uuid>,
+    Json(payload): Json<UpdateTaskRequest>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let mut task: crate::entities::task::ActiveModel =
+        crate::entities::task::Entity::find_by_id(id)
+            .one(&state.db)
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+            .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?
+            .into();
+
+    task.goal = sea_orm::Set(payload.goal.clone());
+    task.cron = sea_orm::Set(payload.cron.clone());
+
+    use sea_orm::ActiveModelTrait;
+    task.update(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    if let Some(cron) = &payload.cron {
+        let _ = state.scheduler.add_task_job(id, cron).await;
+    } else {
+        let _ = state.scheduler.remove_task_job(id).await;
+    }
+
+    tasks::get_task_inner(id, &state.db)
+        .await
+        .map(Json)
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
+}
+
+pub async fn get_task(
+    _user: AuthenticatedUser,
+    Path(id): Path<Uuid>,
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    tasks::get_task_inner(id, &state.db)
+        .await
+        .map(Json)
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
+}
+
+pub async fn get_recent_runs(
+    _user: AuthenticatedUser,
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<Vec<RecentRunResponse>>, (StatusCode, String)> {
+    let results = crate::entities::task_run::Entity::find()
+        .find_also_related(crate::entities::task::Entity)
+        .order_by_desc(crate::entities::task_run::Column::CreatedAt)
+        .limit(50)
+        .all(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let response = results
+        .into_iter()
+        .filter_map(|(run, task_opt)| {
+            task_opt.map(|task| RecentRunResponse {
+                id: run.id,
+                task_id: run.task_id,
+                goal: task.goal,
+                status: run.status,
+                created_at: run.created_at,
+            })
+        })
+        .collect();
+
+    Ok(Json(response))
+}