AkurAI-Drive / backend/src/api.rs

Latest change 671df2d6131cdc6d029891e742e00bf5deff56fb - fix(api): /api/me accepts a Bearer device token (not just the session cookie) by olibuijr

use axum::{
    Json,
    body::Bytes,
    extract::Query,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
use serde::Deserialize;
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()
}

fn bad_request(msg: &str) -> Response {
    (StatusCode::BAD_REQUEST, Json(json!({ "error": msg }))).into_response()
}

/// 404 for any bucket the caller doesn't own — we don't reveal whether it exists.
fn bucket_not_found() -> Response {
    (
        StatusCode::NOT_FOUND,
        Json(json!({ "error": "bucket not found" })),
    )
        .into_response()
}

/// The authenticated caller, resolved from EITHER a `drive_session` cookie OR an
/// `Authorization: Bearer <device-token>` header.
struct CurrentUser {
    sub: String,
    #[allow(dead_code)]
    email: String,
}

fn current_user(headers: &HeaderMap) -> Option<CurrentUser> {
    if let Some(s) = session::lookup_session(headers) {
        return Some(CurrentUser {
            sub: s.sub,
            email: s.email,
        });
    }
    crate::device::lookup_bearer(headers).map(|d| CurrentUser {
        sub: d.sub,
        email: d.email,
    })
}

/// Ownership gate: returns the bucket's region iff `sub` owns it in `storage`,
/// else None (caller responds 404). EVERY file op runs through this.
fn owned_region(sub: &str, bucket: &str) -> Option<String> {
    with_db(|conn| {
        conn.query_row(
            "SELECT region FROM storage WHERE owner_sub = ?1 AND bucket = ?2 LIMIT 1",
            rusqlite::params![sub, bucket],
            |row| row.get::<_, String>(0),
        )
        .ok()
    })
}

/// `GET /api/me` — identity from EITHER a `drive_session` cookie (full: groups +
/// is_admin + tenant) OR a Bearer device token (sub + email; groups aren't stored
/// on the token, so they're empty for the native client).
pub async fn me(headers: HeaderMap) -> Response {
    let cfg = config::get();
    if let Some(s) = session::lookup_session(&headers) {
        let is_admin = s.groups.iter().any(|g| g == &cfg.drive_admin_group);
        return Json(json!({
            "authenticated": true,
            "sub": s.sub,
            "email": s.email,
            "groups": s.groups,
            "is_admin": is_admin,
            "tenant_id": s.tenant_id,
        }))
        .into_response();
    }
    if let Some(d) = crate::device::lookup_bearer(&headers) {
        return Json(json!({
            "authenticated": true,
            "sub": d.sub,
            "email": d.email,
            "groups": [],
            "is_admin": false,
        }))
        .into_response();
    }
    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()
}

// ── Native sync-client file API ──────────────────────────────────────────────
// All endpoints below accept EITHER a `drive_session` cookie OR a Bearer device
// token (via `current_user`), and validate bucket ownership on every operation.

/// `GET /api/buckets` — the caller's buckets from `storage`.
pub async fn buckets(headers: HeaderMap) -> Response {
    let u = match current_user(&headers) {
        Some(u) => u,
        None => return unauthorized(),
    };

    let buckets: Vec<serde_json::Value> = with_db(|conn| {
        let mut stmt = match conn.prepare(
            "SELECT bucket, region, plan, org, 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![u.sub], |row| {
            Ok(json!({
                "bucket": row.get::<_, String>(0)?,
                "region": row.get::<_, String>(1)?,
                "plan": row.get::<_, String>(2)?,
                "org": row.get::<_, String>(3)?,
                "created": row.get::<_, i64>(4)?,
            }))
        });
        match rows {
            Ok(iter) => iter.filter_map(Result::ok).collect(),
            Err(_) => Vec::new(),
        }
    });

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

#[derive(Deserialize)]
pub struct FilesQuery {
    bucket: Option<String>,
    prefix: Option<String>,
}

/// `GET /api/files?bucket=<b>&prefix=<p>` — list objects (ownership-gated).
pub async fn files_list(headers: HeaderMap, Query(q): Query<FilesQuery>) -> Response {
    let u = match current_user(&headers) {
        Some(u) => u,
        None => return unauthorized(),
    };
    let bucket = match q.bucket {
        Some(b) if !b.trim().is_empty() => b.trim().to_string(),
        _ => return bad_request("bucket is required"),
    };
    let region = match owned_region(&u.sub, &bucket) {
        Some(r) => r,
        None => return bucket_not_found(),
    };
    let prefix = q.prefix.filter(|p| !p.is_empty());

    let result = {
        let bucket = bucket.clone();
        tokio::task::spawn_blocking(move || {
            crate::util::list_objects(&bucket, &region, prefix.as_deref())
        })
        .await
    };

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

/// `POST /api/files/presign` `{bucket,key,op:"put"|"get",content_type?}` — return a
/// SigV4 presigned URL (15-min expiry). `content_type` is accepted but not signed:
/// signed headers are `host` only, so the client may set Content-Type freely on PUT.
pub async fn files_presign(headers: HeaderMap, body: Bytes) -> Response {
    let cfg = config::get();
    let u = match current_user(&headers) {
        Some(u) => u,
        None => return unauthorized(),
    };

    let parsed: serde_json::Value = serde_json::from_slice(&body).unwrap_or_else(|_| json!({}));
    let bucket = parsed
        .get("bucket")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();
    let key = parsed
        .get("key")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();
    let op = parsed
        .get("op")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_lowercase();

    if bucket.is_empty() || key.is_empty() {
        return bad_request("bucket and key are required");
    }
    let method = match op.as_str() {
        "put" => "PUT",
        "get" => "GET",
        _ => return bad_request("op must be 'put' or 'get'"),
    };

    let region = match owned_region(&u.sub, &bucket) {
        Some(r) => r,
        None => return bucket_not_found(),
    };

    if cfg.aws_access_key_id.is_empty() || cfg.aws_secret_access_key.is_empty() {
        tracing::error!("presign requested but AWS credentials are not configured");
        return (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({ "error": "signing credentials not configured" })),
        )
            .into_response();
    }

    let url = crate::sigv4::presign(&crate::sigv4::Sigv4Params {
        access_key: &cfg.aws_access_key_id,
        secret_key: &cfg.aws_secret_access_key,
        region: &region,
        bucket: &bucket,
        key: &key,
        method,
        expires: 900,
        now: chrono::Utc::now(),
    });

    Json(json!({ "url": url, "method": method, "expires_in": 900 })).into_response()
}

/// `DELETE /api/files` `{bucket,key}` — delete one object (ownership-gated).
pub async fn files_delete(headers: HeaderMap, body: Bytes) -> Response {
    let u = match current_user(&headers) {
        Some(u) => u,
        None => return unauthorized(),
    };

    let parsed: serde_json::Value = serde_json::from_slice(&body).unwrap_or_else(|_| json!({}));
    let bucket = parsed
        .get("bucket")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();
    let key = parsed
        .get("key")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .trim()
        .to_string();

    if bucket.is_empty() || key.is_empty() {
        return bad_request("bucket and key are required");
    }
    let region = match owned_region(&u.sub, &bucket) {
        Some(r) => r,
        None => return bucket_not_found(),
    };

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

    match result {
        Ok(Ok(())) => Json(json!({ "ok": true })).into_response(),
        Ok(Err(stderr)) => {
            (StatusCode::BAD_GATEWAY, Json(json!({ "error": stderr }))).into_response()
        }
        Err(e) => {
            tracing::error!("delete task panicked: {e}");
            (
                StatusCode::BAD_GATEWAY,
                Json(json!({ "error": "delete failed" })),
            )
                .into_response()
        }
    }
}