AkurAI-Drive / backend/src/api.rs

Latest change beff27fc3a062a220cde556e7bb28da49f6d04ea - Drive login: recover oidc.rs (IdP PKCE + introspection) from deployed binary by AkurAI Agent

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;
use akurai_drive_contract::{
    Bucket, BucketsResponse, FilesResponse, Identity, Presign, StorageBucket, StorageResponse,
};

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 an `akurai_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 and the canonical telemetry linkage for one Drive workspace.
pub(crate) struct StorageAccess {
    pub(crate) region: String,
    pub(crate) telemetry: crate::telemetry::StorageTelemetryContext,
}

/// Returns workspace access only when `sub` owns the bucket. Legacy rows with no
/// canonical tenant ID remain operational but intentionally do not emit telemetry.
pub(crate) fn storage_access(sub: &str, bucket: &str) -> Option<StorageAccess> {
    with_db(|conn| {
        conn.query_row(
            "SELECT region, COALESCE(organization_id, ''), COALESCE(workspace_id, ''), plan
             FROM storage WHERE owner_sub = ?1 AND bucket = ?2 LIMIT 1",
            rusqlite::params![sub, bucket],
            |row| {
                Ok(StorageAccess {
                    region: row.get(0)?,
                    telemetry: crate::telemetry::StorageTelemetryContext {
                        organization_id: row.get(1)?,
                        workspace_id: row.get(2)?,
                        plan: row.get(3)?,
                    },
                })
            },
        )
        .ok()
    })
}

/// `GET /api/me` — identity from EITHER an `akurai_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(Identity {
            authenticated: true,
            sub: s.sub,
            email: s.email,
            groups: s.groups,
            is_admin,
            tenant_id: s.tenant_id,
        })
        .into_response();
    }
    if let Some(d) = crate::device::lookup_bearer(&headers) {
        return Json(Identity {
            authenticated: true,
            sub: d.sub,
            email: d.email,
            ..Identity::default()
        })
        .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<StorageBucket> = 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(StorageBucket {
                id: row.get(0)?,
                org: row.get(1)?,
                plan: row.get(2)?,
                bucket: row.get(3)?,
                region: row.get(4)?,
                created: row.get(5)?,
            })
        });
        match rows {
            Ok(iter) => iter.filter_map(Result::ok).collect(),
            Err(_) => Vec::new(),
        }
    });

    Json(StorageResponse { 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. Only the IDP tenant is a canonical organization ID; the
    // submitted organization label remains presentation data and is never telemetry.
    let id = crate::util::token();
    let now = chrono::Utc::now().timestamp();
    let organization_id = s.tenant_id.clone().unwrap_or_default();
    with_db(|conn| {
        conn.execute(
            "INSERT INTO storage
             (id, owner_sub, org, email, plan, bucket, region, created, organization_id, workspace_id)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
            rusqlite::params![
                id,
                s.sub,
                org,
                email,
                plan,
                bucket,
                region,
                now,
                organization_id,
                id,
            ],
        )
        .ok();
    });

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

// ── Native sync-client file API ──────────────────────────────────────────────
// All endpoints below accept EITHER an `akurai_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(),
    };

    Json(BucketsResponse {
        buckets: owned_buckets(&u.sub),
    })
    .into_response()
}

/// The buckets `sub` owns, newest first.
pub(crate) fn owned_buckets(sub: &str) -> Vec<Bucket> {
    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![sub], |row| {
            Ok(Bucket {
                bucket: row.get(0)?,
                region: row.get(1)?,
                plan: row.get(2)?,
                org: row.get(3)?,
                created: row.get(4)?,
            })
        });
        match rows {
            Ok(iter) => iter.filter_map(Result::ok).collect(),
            Err(_) => Vec::new(),
        }
    })
}

#[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 access = match storage_access(&u.sub, &bucket) {
        Some(access) => access,
        None => return bucket_not_found(),
    };
    let region = access.region.clone();
    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)) => {
            if !files.is_empty() {
                let storage_used_bytes = files
                    .iter()
                    .fold(0_u64, |total, file| total.saturating_add(file.size));
                crate::telemetry::observe_upload(&access.telemetry, storage_used_bytes);
            }
            Json(FilesResponse { 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 access = match storage_access(&u.sub, &bucket) {
        Some(access) => access,
        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: &access.region,
        bucket: &bucket,
        key: &key,
        method,
        expires: 900,
        now: chrono::Utc::now(),
    });

    if method == "GET" {
        crate::telemetry::observe_retrieval(&access.telemetry);
    }

    Json(Presign {
        url,
        method: method.to_string(),
        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 access = match storage_access(&u.sub, &bucket) {
        Some(access) => access,
        None => return bucket_not_found(),
    };

    let result = {
        let (bucket, key) = (bucket.clone(), key.clone());
        let region = access.region.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()
        }
    }
}

/// `POST /api/client-logs` — forward browser logs without exposing the Monitor token.
pub async fn client_logs(Json(body): Json<serde_json::Value>) -> Response {
    let Some(raw_logs) = body.get("logs").and_then(serde_json::Value::as_array) else {
        return bad_request("expected { logs: [...] }");
    };

    let now = chrono::Utc::now().timestamp();
    let logs: Vec<serde_json::Value> = raw_logs
        .iter()
        .take(500)
        .map(|log| {
            let line = log
                .get("line")
                .map(|value| match value {
                    serde_json::Value::String(value) => value.clone(),
                    value => value.to_string(),
                })
                .unwrap_or_default();
            let line = if line.len() > 16_384 {
                line.chars().take(16_384).collect()
            } else {
                line
            };
            let level = match log.get("level").and_then(serde_json::Value::as_str) {
                Some("trace" | "debug" | "info" | "warn" | "error") => log["level"].clone(),
                _ => json!("info"),
            };

            json!({
                "source": "akurai-drive-api-client",
                "line": line,
                "level": level,
                "target": log.get("target").and_then(serde_json::Value::as_str),
                "origin": "client",
                "fields": log.get("fields").cloned().unwrap_or(serde_json::Value::Null),
                "host_id": 1,
                "application_id": 15,
                "ts": log.get("ts").and_then(serde_json::Value::as_i64).unwrap_or(now),
            })
        })
        .collect();
    let inserted = logs.len();

    let url = std::env::var("MONITOR_INGEST_URL").ok();
    let token = std::env::var("MONITOR_INGEST_TOKEN").ok();
    if let (Some(url), Some(token)) = (url, token) {
        let endpoint = format!("{}/api/ingest", url.trim_end_matches('/'));
        let payload = json!({ "logs": logs });
        if let Err(error) = tokio::task::spawn_blocking(move || {
            ureq::post(&endpoint)
                .set("Authorization", &format!("Bearer {token}"))
                .set("Content-Type", "application/json")
                .send_string(&payload.to_string())
                .map(|_| ())
                .map_err(|e| e.to_string())
        })
        .await
        {
            tracing::warn!("client log shipping task failed: {error}");
        }
    }

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