AkurAI-Drive / backend/src/util.rs

Latest change f891f45338fe5e5d14370df77cc8001dceb06426 - Drive MCP: recover /mcp tools and make it an IdP OAuth protected resource by AkurAI Agent

use akurai_drive_contract::RemoteFile;
use uuid::Uuid;

/// Random opaque token (32 lowercase hex chars) for session ids and OAuth state.
pub fn token() -> String {
    Uuid::new_v4().simple().to_string()
}

/// 6 lowercase hex chars, used as the bucket-name suffix for uniqueness.
fn short_hex() -> String {
    Uuid::new_v4().simple().to_string()[..6].to_string()
}

/// Slugify an org name into `[a-z0-9-]`, collapsing runs and trimming dashes.
fn slug(org: &str) -> String {
    let mut out = String::new();
    let mut last_dash = true; // leading dashes are dropped
    for c in org.to_lowercase().chars() {
        if c.is_ascii_alphanumeric() {
            out.push(c);
            last_dash = false;
        } else if !last_dash {
            out.push('-');
            last_dash = true;
        }
    }
    out.trim_matches('-').to_string()
}

/// Build an S3-legal bucket name: `<prefix>-<slug(org)>-<6 hex>`, lowercased,
/// `[a-z0-9-]` only, <= 63 chars. The slug is truncated to keep the whole name
/// within the limit.
pub fn bucket_name(prefix: &str, org: &str) -> String {
    let hex = short_hex();
    let prefix = prefix.to_lowercase();
    // Reserved length: "<prefix>-" + "-<hex>".
    let reserved = prefix.len() + 1 + 1 + hex.len();
    let max_slug = 63usize.saturating_sub(reserved);

    let mut org_slug = slug(org);
    if org_slug.len() > max_slug {
        org_slug.truncate(max_slug);
        org_slug = org_slug.trim_end_matches('-').to_string();
    }

    let mut name = if org_slug.is_empty() {
        format!("{prefix}-{hex}")
    } else {
        format!("{prefix}-{org_slug}-{hex}")
    };

    // Final safety net: keep only legal chars and the length bound.
    name = name
        .chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '-')
        .collect();
    if name.len() > 63 {
        name.truncate(63);
    }
    name.trim_matches('-').to_string()
}

/// Provision an S3 bucket via the `aws` CLI. Credentials come from the process
/// environment (systemd EnvironmentFile). For us-east-1, no LocationConstraint is
/// allowed; any other region requires `--create-bucket-configuration`.
/// On success returns Ok(()); on non-zero exit returns the first stderr line.
pub fn create_bucket(bucket: &str, region: &str) -> Result<(), String> {
    let mut cmd = std::process::Command::new("aws");
    cmd.arg("s3api")
        .arg("create-bucket")
        .arg("--bucket")
        .arg(bucket)
        .arg("--region")
        .arg(region);

    if region != "us-east-1" {
        cmd.arg("--create-bucket-configuration")
            .arg(format!("LocationConstraint={region}"));
    }

    let output = cmd
        .output()
        .map_err(|e| format!("failed to run aws: {e}"))?;

    if output.status.success() {
        return Ok(());
    }
    Err(first_stderr_line(
        &output.stderr,
        "aws create-bucket failed",
    ))
}

/// List objects in `bucket` via `aws s3api list-objects-v2` (creds from the process
/// env). Returns the shared API contract's `RemoteFile` values. On non-zero exit
/// returns the first stderr line.
pub fn list_objects(
    bucket: &str,
    region: &str,
    prefix: Option<&str>,
) -> Result<Vec<RemoteFile>, String> {
    let mut cmd = std::process::Command::new("aws");
    cmd.arg("s3api")
        .arg("list-objects-v2")
        .arg("--bucket")
        .arg(bucket)
        .arg("--region")
        .arg(region)
        .arg("--output")
        .arg("json");
    if let Some(p) = prefix {
        cmd.arg("--prefix").arg(p);
    }

    let output = cmd
        .output()
        .map_err(|e| format!("failed to run aws: {e}"))?;
    if !output.status.success() {
        return Err(first_stderr_line(
            &output.stderr,
            "aws list-objects-v2 failed",
        ));
    }

    let parsed: serde_json::Value = serde_json::from_slice(&output.stdout)
        .map_err(|e| format!("aws json parse failed: {e}"))?;

    let files = parsed
        .get("Contents")
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .map(|o| RemoteFile {
                    key: o
                        .get("Key")
                        .and_then(|v| v.as_str())
                        .unwrap_or("")
                        .to_string(),
                    size: o.get("Size").and_then(|v| v.as_i64()).unwrap_or(0) as u64,
                    etag: o
                        .get("ETag")
                        .and_then(|v| v.as_str())
                        .unwrap_or("")
                        .trim_matches('"')
                        .to_string(),
                    last_modified: o
                        .get("LastModified")
                        .and_then(|v| v.as_str())
                        .unwrap_or("")
                        .to_string(),
                })
                .collect()
        })
        .unwrap_or_default();

    Ok(files)
}

/// Delete a single object via `aws s3api delete-object` (creds from the process env).
/// On success returns Ok(()); on non-zero exit returns the first stderr line.
pub fn delete_object(bucket: &str, region: &str, key: &str) -> Result<(), String> {
    let output = std::process::Command::new("aws")
        .arg("s3api")
        .arg("delete-object")
        .arg("--bucket")
        .arg(bucket)
        .arg("--region")
        .arg(region)
        .arg("--key")
        .arg(key)
        .output()
        .map_err(|e| format!("failed to run aws: {e}"))?;

    if output.status.success() {
        return Ok(());
    }
    Err(first_stderr_line(
        &output.stderr,
        "aws delete-object failed",
    ))
}

/// Run `aws s3api <op> --bucket <bucket> --region <region> <args..>` (creds from
/// the process env). Returns stdout on success, else the first stderr line.
fn s3api(op: &str, bucket: &str, region: &str, args: &[&str]) -> Result<Vec<u8>, String> {
    let output = std::process::Command::new("aws")
        .args(["s3api", op, "--bucket", bucket, "--region", region])
        .args(args)
        .output()
        .map_err(|e| format!("failed to run aws: {e}"))?;
    if output.status.success() {
        return Ok(output.stdout);
    }
    Err(first_stderr_line(
        &output.stderr,
        &format!("aws {op} failed"),
    ))
}

/// Create an empty object (used as a virtual folder marker, key ending in '/').
pub fn put_empty_object(bucket: &str, region: &str, key: &str) -> Result<(), String> {
    s3api(
        "put-object",
        bucket,
        region,
        &["--key", key, "--body", "/dev/null"],
    )
    .map(|_| ())
}

/// Server-side copy of `src` to `dst` within one bucket.
pub fn copy_object(bucket: &str, region: &str, src: &str, dst: &str) -> Result<(), String> {
    let source = format!("{bucket}/{src}");
    s3api(
        "copy-object",
        bucket,
        region,
        &["--copy-source", &source, "--key", dst],
    )
    .map(|_| ())
}

/// Object metadata via `head-object`; `Ok(None)` when the key does not exist.
pub fn head_object(
    bucket: &str,
    region: &str,
    key: &str,
) -> Result<Option<serde_json::Value>, String> {
    match s3api(
        "head-object",
        bucket,
        region,
        &["--key", key, "--checksum-mode", "ENABLED"],
    ) {
        Ok(out) => serde_json::from_slice(&out)
            .map(Some)
            .map_err(|e| format!("aws json parse failed: {e}")),
        Err(e) if e.contains("404") || e.contains("Not Found") || e.contains("NoSuchKey") => {
            Ok(None)
        }
        Err(e) => Err(e),
    }
}

/// First non-empty, trimmed line of a command's stderr, or `fallback`.
fn first_stderr_line(stderr: &[u8], fallback: &str) -> String {
    String::from_utf8_lossy(stderr)
        .lines()
        .map(str::trim)
        .find(|line| !line.is_empty())
        .unwrap_or(fallback)
        .to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bucket_name_is_legal() {
        let n = bucket_name("akurai-drive", "Acme ehf.");
        assert!(n.starts_with("akurai-drive-acme-ehf-"));
        assert!(n.len() <= 63);
        assert!(
            n.chars()
                .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
        );
    }

    #[test]
    fn bucket_name_truncates_long_org() {
        let n = bucket_name("akurai-drive", &"x".repeat(200));
        assert!(n.len() <= 63);
    }

    #[test]
    fn bucket_name_handles_empty_slug() {
        let n = bucket_name("akurai-drive", "***");
        assert!(n.starts_with("akurai-drive-"));
        assert!(n.len() <= 63);
    }
}