auth

src/agent.rs

@@ -21,11 +21,12 @@ impl Agent {
         initial_message: String,
     ) -> Self {
         let intro = format!(
-                    "You are an autonomous agent. You have access to tools that can help 
-                    you achieve your goals. Use them wisely. Do not ask for clarification and use the
-                     answer tool once you to give your final answer. current date is {}",
-                    Utc::now().to_rfc3339()
-                );
+            "You are an autonomous agent. You have access to tools that can help 
+                    you achieve your goals. Use them wisely. The user is unable to respond to you
+                    so do not ask for clarification and use the
+                    answer tool once you to give your final answer. current date is {}",
+            Utc::now().to_rfc3339()
+        );
         println!("initial_message: {}", intro);
         let messages = vec![
             Message {

src/auth.rs

@@ -0,0 +1,149 @@
+use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Claims {
+    pub sub: String,
+    pub exp: usize,
+    pub iat: usize,
+    pub iss: String,
+    pub aud: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct Jwk {
+    kty: String,
+    kid: String,
+    n: String,
+    e: String,
+    alg: Option<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct Jwks {
+    keys: Vec<Jwk>,
+}
+
+pub struct JwksVerifier {
+    issuer: String,
+    jwks_uri: String,
+    keys: Arc<RwLock<Vec<Jwk>>>,
+    client: Client,
+}
+
+impl JwksVerifier {
+    pub async fn new(issuer: String) -> Result<Self, Box<dyn std::error::Error>> {
+        let client = Client::new();
+        // Authentik OIDC discovery
+        let discovery_url = format!(
+            "{}/.well-known/openid-configuration",
+            issuer.trim_end_matches('/')
+        );
+        let config: serde_json::Value = client.get(&discovery_url).send().await?.json().await?;
+
+        let jwks_uri = config["jwks_uri"]
+            .as_str()
+            .ok_or("Missing jwks_uri in discovery")?
+            .to_string();
+
+        let verifier = Self {
+            issuer,
+            jwks_uri,
+            keys: Arc::new(RwLock::new(Vec::new())),
+            client,
+        };
+
+        verifier.refresh_keys().await?;
+        Ok(verifier)
+    }
+
+    pub async fn refresh_keys(&self) -> Result<(), Box<dyn std::error::Error>> {
+        let jwks: Jwks = self.client.get(&self.jwks_uri).send().await?.json().await?;
+        let mut keys = self.keys.write().await;
+        *keys = jwks.keys;
+        Ok(())
+    }
+
+    pub async fn verify(&self, token: &str) -> Result<Claims, Box<dyn std::error::Error>> {
+        let header = decode_header(token)?;
+        let kid = header.kid.ok_or("Missing kid in token header")?;
+
+        let keys = self.keys.read().await;
+        let jwk = keys
+            .iter()
+            .find(|k| k.kid == kid)
+            .ok_or("Key not found in JWKS")?;
+
+        let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
+
+        let mut validation = Validation::new(Algorithm::RS256);
+        validation.set_issuer(&[self.issuer.clone()]);
+        // Aud validation might need careful config, usually it's the client_id
+        validation.validate_aud = false;
+
+        let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
+        Ok(token_data.claims)
+    }
+}
+
+pub struct Authenticator {
+    client_id: String,
+    client_secret: String,
+    token_url: String,
+    client: Client,
+}
+
+impl Authenticator {
+    pub async fn new(
+        issuer: String,
+        client_id: String,
+        client_secret: String,
+    ) -> Result<Self, Box<dyn std::error::Error>> {
+        let client = Client::new();
+        let discovery_url = format!(
+            "{}/.well-known/openid-configuration",
+            issuer.trim_end_matches('/')
+        );
+        let config: serde_json::Value = client.get(&discovery_url).send().await?.json().await?;
+
+        let token_url = config["token_endpoint"]
+            .as_str()
+            .ok_or("Missing token_endpoint in discovery")?
+            .to_string();
+
+        Ok(Self {
+            client_id,
+            client_secret,
+            token_url,
+            client,
+        })
+    }
+
+    pub async fn exchange_code(
+        &self,
+        code: String,
+        redirect_uri: String,
+    ) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
+        let params = [
+            ("grant_type", "authorization_code"),
+            ("code", &code),
+            ("redirect_uri", &redirect_uri),
+            ("client_id", &self.client_id),
+            ("client_secret", &self.client_secret),
+        ];
+
+        let res = self
+            .client
+            .post(&self.token_url)
+            .form(&params)
+            .send()
+            .await?
+            .json()
+            .await?;
+
+        Ok(res)
+    }
+}

src/main.rs

@@ -1,5 +1,6 @@
 mod agent;
 mod api;
+mod auth;
 mod entities;
 mod scheduler;
 mod server;

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()