auth

src/server.rs

@@ -1,16 +1,20 @@
 use axum::{
-    Json, Router,
-    extract::{Path, State},
-    http::StatusCode,
+    Json, RequestPartsExt, Router,
+    extract::{FromRef, FromRequestParts, Path, Query, State},
+    http::{StatusCode, request::Parts},
     routing::{get, post},
 };
+use axum_extra::{
+    TypedHeader,
+    headers::{Authorization, authorization::Bearer},
+};
 use chrono::Utc;
 use sea_orm::{
     ActiveModelTrait, Database, DatabaseConnection, EntityTrait, QueryOrder, QuerySelect, Set,
 };
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
-use tower_http::cors::CorsLayer;
+use tower_http::cors::{Any, CorsLayer};
 use uuid::Uuid;
 
 use crate::agent::Agent;
@@ -25,6 +29,8 @@ pub struct AppState {
     pub scheduler: Arc<Scheduler>,
     pub zen_api_key: Option<String>,
     pub tavily_api_key: Option<String>,
+    pub verifier: Arc<crate::auth::JwksVerifier>,
+    pub authenticator: Arc<crate::auth::Authenticator>,
 }
 
 #[derive(Deserialize)]
@@ -84,20 +90,43 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
         }
     }
 
+    let authentik_issuer =
+        std::env::var("AUTHENTIK_ISSUER").map_err(|_| "AUTHENTIK_ISSUER not set")?;
+    let authentik_client_id =
+        std::env::var("AUTHENTIK_CLIENT_ID").map_err(|_| "AUTHENTIK_CLIENT_ID not set")?;
+    let authentik_client_secret =
+        std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?;
+
+    let verifier = Arc::new(crate::auth::JwksVerifier::new(authentik_issuer.clone()).await?);
+    let authenticator = Arc::new(
+        crate::auth::Authenticator::new(
+            authentik_issuer,
+            authentik_client_id,
+            authentik_client_secret,
+        )
+        .await?,
+    );
+
     let state = Arc::new(AppState {
         db,
         scheduler,
         zen_api_key,
         tavily_api_key,
+        verifier,
+        authenticator,
     });
 
-    let cors = CorsLayer::permissive();
+    let cors = CorsLayer::new()
+        .allow_origin(Any)
+        .allow_methods(Any)
+        .allow_headers(Any);
 
     let app = Router::new()
         .route("/tasks", post(create_task).get(list_tasks))
         .route("/tasks/:id", get(get_task).put(update_task))
         .route("/tasks/:id/runs", post(rerun_task))
         .route("/runs/recent", get(get_recent_runs))
+        .route("/auth/callback", get(auth_callback))
         .layer(cors)
         .with_state(state);
 
@@ -109,6 +138,7 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
 }
 
 async fn list_tasks(
+    _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
 ) -> Result<Json<Vec<TaskResponse>>, (StatusCode, String)> {
     let tasks = Task::find()
@@ -142,6 +172,7 @@ async fn list_tasks(
 }
 
 async fn create_task(
+    _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
     Json(payload): Json<CreateTaskRequest>,
 ) -> Result<Json<TaskResponse>, (StatusCode, String)> {
@@ -166,10 +197,11 @@ async fn create_task(
         let _ = state.scheduler.remove_task_job(task_id).await;
     }
 
-    execute_agent_run(state, task_id, payload.goal).await
+    get_task_inner(task_id, &state).await.map(Json)
 }
 
 async fn rerun_task(
+    _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
     Path(id): Path<Uuid>,
 ) -> Result<Json<TaskResponse>, (StatusCode, String)> {
@@ -234,10 +266,11 @@ async fn execute_agent_run(
         .await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    get_task(Path(task_id), State(state)).await
+    get_task_inner(task_id, &state).await.map(Json)
 }
 
 async fn update_task(
+    _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
     Path(id): Path<Uuid>,
     Json(payload): Json<UpdateTaskRequest>,
@@ -262,29 +295,95 @@ async fn update_task(
         let _ = state.scheduler.remove_task_job(id).await;
     }
 
-    get_task(Path(id), State(state)).await
+    get_task_inner(id, &state).await.map(Json)
+}
+
+pub struct AuthenticatedUser(pub crate::auth::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 TypedHeader(Authorization(bearer)) = parts
+            .extract::<TypedHeader<Authorization<Bearer>>>()
+            .await
+            .map_err(|_| {
+                (
+                    StatusCode::UNAUTHORIZED,
+                    "Missing or invalid Authorization header".to_string(),
+                )
+            })?;
+
+        let claims = app_state
+            .verifier
+            .verify(bearer.token())
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::UNAUTHORIZED,
+                    format!("Token verification failed: {}", e),
+                )
+            })?;
+
+        Ok(AuthenticatedUser(claims))
+    }
+}
+
+#[derive(Deserialize)]
+pub struct AuthCallbackQuery {
+    pub code: String,
+    pub redirect_uri: String,
+}
+
+async fn auth_callback(
+    State(state): State<Arc<AppState>>,
+    Query(query): Query<AuthCallbackQuery>,
+) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
+    state
+        .authenticator
+        .exchange_code(query.code, query.redirect_uri)
+        .await
+        .map(Json)
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Token exchange failed: {}", e),
+            )
+        })
 }
 
 async fn get_task(
+    _user: AuthenticatedUser,
     Path(id): Path<Uuid>,
     State(state): State<Arc<AppState>>,
 ) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    get_task_inner(id, &state).await.map(Json)
+}
+
+async fn get_task_inner(id: Uuid, state: &AppState) -> Result<TaskResponse, (StatusCode, String)> {
     let results = Task::find_by_id(id)
         .find_with_related(TaskRun)
         .all(&state.db)
         .await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    let (task, runs) = results
+    let (t, runs) = results
         .into_iter()
         .next()
         .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?;
 
-    Ok(Json(TaskResponse {
-        id: task.id,
-        goal: task.goal,
-        cron: task.cron,
-        created_at: task.created_at,
+    Ok(TaskResponse {
+        id: t.id,
+        goal: t.goal,
+        cron: t.cron,
+        created_at: t.created_at,
         runs: runs
             .into_iter()
             .map(|r| TaskRunResponse {
@@ -295,10 +394,11 @@ async fn get_task(
                 created_at: r.created_at,
             })
             .collect(),
-    }))
+    })
 }
 
 async fn get_recent_runs(
+    _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
 ) -> Result<Json<Vec<RecentRunResponse>>, (StatusCode, String)> {
     let results = TaskRun::find()