AkurAI-Drive / backend/src/api.rs

Latest change 6e3a0824f4ec4d2bea2524a20b4960c7ab9e2766 - backend: akurai-drive-api (OIDC RP + per-tenant S3 provisioning) by olibuijr

use axum::{
    Json,
    body::Bytes,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
use serde_json::json;

use crate::config;
use crate::db::with_db;
use crate::session;

fn unauthorized() -> Response {
    (
        StatusCode::UNAUTHORIZED,
        Json(json!({"authenticated": false})),
    )
        .into_response()
}

/// `GET /api/me` — report the current identity, or 401 if no valid session.
pub async fn me(headers: HeaderMap) -> Response {
    let cfg = config::get();
    match session::lookup_session(&headers) {
        Some(s) => {
            let is_admin = s.groups.iter().any(|g| g == &cfg.drive_admin_group);
            Json(json!({
                "authenticated": true,
                "sub": s.sub,
                "email": s.email,
                "groups": s.groups,
                "is_admin": is_admin,
                "tenant_id": s.tenant_id,
            }))
            .into_response()
        }
        None => unauthorized(),
    }
}

/// `GET /api/storage` — list the buckets owned by this session's `sub`.
pub async fn storage_list(headers: HeaderMap) -> Response {
    let s = match session::lookup_session(&headers) {
        Some(s) => s,
        None => return unauthorized(),
    };

    let buckets: Vec<serde_json::Value> = with_db(|conn| {
        let mut stmt = match conn.prepare(
            "SELECT id, org, plan, bucket, region, created
             FROM storage WHERE owner_sub = ?1 ORDER BY created DESC",
        ) {
            Ok(stmt) => stmt,
            Err(_) => return Vec::new(),
        };
        let rows = stmt.query_map(rusqlite::params![s.sub], |row| {
            Ok(json!({
                "id": row.get::<_, String>(0)?,
                "org": row.get::<_, String>(1)?,
                "plan": row.get::<_, String>(2)?,
                "bucket": row.get::<_, String>(3)?,
                "region": row.get::<_, String>(4)?,
                "created": row.get::<_, i64>(5)?,
            }))
        });
        match rows {
            Ok(iter) => iter.filter_map(Result::ok).collect(),
            Err(_) => Vec::new(),
        }
    });

    Json(json!({ "buckets": buckets })).into_response()
}

/// `POST /api/storage` — gated by the `X-Drive-Access-Code` header AND a valid
/// session, then provisions an S3 bucket via the `aws` CLI and records it.
pub async fn storage_create(headers: HeaderMap, body: Bytes) -> Response {
    let cfg = config::get();

    // 1. Access-code gate (checked before everything else).
    let provided_code = headers
        .get("x-drive-access-code")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");
    if provided_code != cfg.drive_access_code || cfg.drive_access_code.is_empty() {
        return (
            StatusCode::FORBIDDEN,
            Json(json!({"error": "invalid access code"})),
        )
            .into_response();
    }

    // 2. Session required.
    let s = match session::lookup_session(&headers) {
        Some(s) => s,
        None => return unauthorized(),
    };

    // 3. Parse the JSON body.
    let parsed: serde_json::Value = serde_json::from_slice(&body).unwrap_or_else(|_| json!({}));
    let org = parsed
        .get("org")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();
    let email = parsed
        .get("email")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();
    let plan = parsed
        .get("plan")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();

    if org.is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(json!({"error": "org is required"})),
        )
            .into_response();
    }

    // 4. Provision the bucket via the aws CLI (blocking) on the blocking pool.
    let bucket = crate::util::bucket_name(&cfg.drive_bucket_prefix, &org);
    let region = cfg.aws_region.clone();

    let provision = {
        let bucket = bucket.clone();
        let region = region.clone();
        tokio::task::spawn_blocking(move || crate::util::create_bucket(&bucket, &region)).await
    };

    match provision {
        Ok(Ok(())) => {}
        Ok(Err(stderr)) => {
            return (StatusCode::BAD_GATEWAY, Json(json!({"error": stderr}))).into_response();
        }
        Err(e) => {
            tracing::error!("provision task panicked: {e}");
            return (
                StatusCode::BAD_GATEWAY,
                Json(json!({"error": "provisioning failed"})),
            )
                .into_response();
        }
    }

    // 5. Record it.
    let id = crate::util::token();
    let now = chrono::Utc::now().timestamp();
    with_db(|conn| {
        conn.execute(
            "INSERT INTO storage (id, owner_sub, org, email, plan, bucket, region, created)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
            rusqlite::params![id, s.sub, org, email, plan, bucket, region, now],
        )
        .ok();
    });

    (
        StatusCode::CREATED,
        Json(json!({ "bucket": bucket, "region": region, "plan": plan })),
    )
        .into_response()
}