maarcadetweet/crates/pds-server/tests/blob_integration.rs

use serde_json::{json, Value};
use std::time::Duration;

const PDS_URL: &str = "http://127.0.0.1:2583";

async fn client() -> reqwest::Client {
    reqwest::Client::builder()
        .timeout(Duration::from_secs(5))
        .build()
        .unwrap()
}

async fn wait_for_pds() -> bool {
    let c = client().await;
    for _ in 0..20 {
        if let Ok(r) = c.get(format!("{}/healthz", PDS_URL)).send().await {
            if r.status().is_success() {
                return true;
            }
        }
        tokio::time::sleep(Duration::from_millis(250)).await;
    }
    false
}

async fn db_pool() -> Option<sqlx::PgPool> {
    let url = std::env::var("DATABASE_URL_PDS")
        .unwrap_or_else(|_| "postgres://pds:pds@127.0.0.1:5434/pds".to_string());
    sqlx::postgres::PgPoolOptions::new()
        .max_connections(2)
        .acquire_timeout(Duration::from_secs(2))
        .connect(&url)
        .await
        .ok()
}

async fn fresh_user(prefix: &str) -> (reqwest::Client, String, String) {
    let c = client().await;
    let handle = format!(
        "{}_{}.maarcadetweet.local",
        prefix,
        uuid::Uuid::new_v4().simple()
    );
    let acc: Value = c
        .post(format!("{}/xrpc/com.atproto.server.createAccount", PDS_URL))
        .json(&json!({"handle": handle, "password": "hunter2hunter2"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let did = acc["did"].as_str().unwrap().to_string();
    let jwt = acc["access_jwt"].as_str().unwrap().to_string();
    (c, did, jwt)
}

/// Seed a single post so the repo has a head-commit / repo_blocks
/// row. The exact content doesn't matter — we just need *any* block
/// under the user's DID so that `repo_blocks` is non-empty and the
/// `getBlob` handler's "user has a repo" gate passes.
async fn seed_any_record(c: &reqwest::Client, did: &str, jwt: &str) {
    let r: Value = c
        .post(format!("{}/xrpc/com.atproto.repo.createRecord", PDS_URL))
        .bearer_auth(jwt)
        .json(&json!({
            "repo": did,
            "collection": "app.twi.post",
            "record": {
                "text": "blob seed",
                "createdAt": "2026-07-05T12:00:00Z",
            },
        }))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert!(r["uri"].is_string(), "seed createRecord: {:?}", r);
}

/// Compute a CIDv1 + sha256 CID for the given blob bytes — the same
/// shape PDS clients use to reference uploaded blobs.
fn blob_cid(bytes: &[u8]) -> cid::Cid {
    at_crypto::cid::cid_for_raw(0x55, at_crypto::cid::sha256(bytes)).unwrap()
}

#[tokio::test]
async fn get_blob_returns_value_block() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let Some(pool) = db_pool().await else {
        eprintln!("no PDS database reachable, skipping");
        return;
    };
    let (c, did, jwt) = fresh_user("blob").await;
    seed_any_record(&c, &did, &jwt).await;

    let payload: Vec<u8> = b"hello, blob! \xe2\x98\x83 \xf0\x9f\x9a\x80".to_vec();
    let cid = blob_cid(&payload);

    let cid_bytes: Vec<u8> = cid.to_bytes();
    sqlx::query(
        r#"INSERT INTO repo_blocks (did, cid, block, size)
           VALUES ($1, $2, $3, $4)
           ON CONFLICT (did, cid) DO NOTHING"#,
    )
    .bind(&did)
    .bind(cid_bytes.as_slice())
    .bind(payload.as_slice())
    .bind(payload.len() as i32)
    .execute(&pool)
    .await
    .unwrap();

    let url = format!(
        "{}/xrpc/com.atproto.sync.getBlob?did={}&cid={}",
        PDS_URL, did, cid
    );
    let resp = c.get(&url).send().await.unwrap();
    assert_eq!(resp.status().as_u16(), 200, "expected 200 for {url}");
    let bytes = resp.bytes().await.unwrap();
    assert_eq!(bytes.as_ref(), payload.as_slice());
}

#[tokio::test]
async fn get_blob_returns_404_for_unknown_cid() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let Some(_pool) = db_pool().await else {
        eprintln!("no PDS database reachable, skipping");
        return;
    };
    let (c, did, jwt) = fresh_user("blobmiss").await;
    seed_any_record(&c, &did, &jwt).await;

    let cid = blob_cid(b"definitely-not-uploaded");

    let resp = c
        .get(format!(
            "{}/xrpc/com.atproto.sync.getBlob?did={}&cid={}",
            PDS_URL, did, cid
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 400);
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["error"], json!("BlobNotFound"));
}

#[tokio::test]
async fn get_blob_rejects_invalid_cid() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, _jwt) = fresh_user("blobcid").await;
    let resp = c
        .get(format!(
            "{}/xrpc/com.atproto.sync.getBlob?did={}&cid=not-a-cid",
            PDS_URL, did
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 400);
    let body: Value = resp.json().await.unwrap();
    assert_eq!(body["error"], json!("InvalidRequest"));
}

#[tokio::test]
async fn get_blob_shortcut_returns_value_block() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let Some(pool) = db_pool().await else {
        eprintln!("no PDS database reachable, skipping");
        return;
    };
    let (c, did, jwt) = fresh_user("blobshort").await;
    seed_any_record(&c, &did, &jwt).await;

    // Use a binary payload that won't match any known signature or
    // pass the ASCII-text heuristic, so the shortcut endpoint falls
    // back to `application/octet-stream`. (Phase 7: previously this
    // test used a plain-text payload; with magic-byte sniffing that
    // would be classified as `text/plain; charset=utf-8` instead of
    // the octet-stream default.)
    let payload: Vec<u8> = vec![0x00, 0x01, 0x02, 0xff, 0xfe, 0x80, 0x90];
    let cid = blob_cid(&payload);
    let cid_bytes: Vec<u8> = cid.to_bytes();

    sqlx::query(
        r#"INSERT INTO repo_blocks (did, cid, block, size)
           VALUES ($1, $2, $3, $4)
           ON CONFLICT (did, cid) DO NOTHING"#,
    )
    .bind(&did)
    .bind(cid_bytes.as_slice())
    .bind(payload.as_slice())
    .bind(payload.len() as i32)
    .execute(&pool)
    .await
    .unwrap();

    let resp = c
        .get(format!("{}/blob/{}", PDS_URL, cid))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 200);
    let ct = resp
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string();
    let bytes = resp.bytes().await.unwrap();
    assert_eq!(bytes.as_ref(), payload.as_slice());
    assert_eq!(ct, "application/octet-stream");
}

// -- Phase 7: uploadBlob tests ---------------------------------------------

/// A minimal PNG signature followed by enough bytes that the magic
/// detector recognises it. We don't need a fully-valid PNG for the
/// uploadBlob tests — we just need the first 8 bytes to match the
/// signature and the response to carry the correct `mimeType`.
fn png_bytes() -> Vec<u8> {
    let mut v = vec![0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a];
    v.extend_from_slice(&[0u8; 64]);
    v
}

/// `com.atproto.uploadBlob` — happy path. POST a small binary blob
/// with a `Content-Type` header, then read it back via
/// `com.atproto.sync.getBlob` and verify the bytes match and the
/// server-side CID matches what `sha256 + CIDv1-raw` would compute
/// locally.
#[tokio::test]
async fn upload_blob_persists_and_round_trips() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("upl").await;
    seed_any_record(&c, &did, &jwt).await;

    let payload = png_bytes();

    let resp = c
        .post(format!("{}/xrpc/com.atproto.uploadBlob", PDS_URL))
        .bearer_auth(&jwt)
        .header("Content-Type", "image/png")
        .body(payload.clone())
        .send()
        .await
        .unwrap();
    assert_eq!(
        resp.status().as_u16(),
        200,
        "uploadBlob should succeed"
    );
    let body: Value = resp.json().await.unwrap();
    let returned_cid = body["blob"]["ref"]["$link"]
        .as_str()
        .expect("blob.ref.$link should be a string")
        .to_string();
    let returned_size = body["blob"]["size"].as_u64().unwrap();
    let returned_mime = body["blob"]["mimeType"].as_str().unwrap();
    assert_eq!(returned_mime, "image/png");
    assert_eq!(returned_size, payload.len() as u64);

    // The returned CID must match what we compute locally from the
    // payload bytes (CIDv1-raw + SHA-256).
    let expected_cid = blob_cid(&payload).to_string();
    assert_eq!(returned_cid, expected_cid);

    // Read the blob back via the spec endpoint and confirm bytes
    // match.
    let get_resp = c
        .get(format!(
            "{}/xrpc/com.atproto.sync.getBlob?did={}&cid={}",
            PDS_URL, did, returned_cid
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(get_resp.status().as_u16(), 200);
    let bytes = get_resp.bytes().await.unwrap();
    assert_eq!(bytes.as_ref(), payload.as_slice());
}

/// `com.atproto.uploadBlob` rejects payloads larger than the 1 MiB
/// limit with `413 Payload Too Large`.
#[tokio::test]
async fn upload_blob_rejects_oversized() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("upbig").await;
    seed_any_record(&c, &did, &jwt).await;

    // 2 MiB payload. The PDS's `DefaultBodyLimit::max(1 MiB)` layer
    // rejects the request before our handler sees it, so we expect
    // 413 from axum's body extractor.
    let payload = vec![0u8; 2 * 1024 * 1024];

    let resp = c
        .post(format!("{}/xrpc/com.atproto.uploadBlob", PDS_URL))
        .bearer_auth(&jwt)
        .header("Content-Type", "image/png")
        .body(payload)
        .send()
        .await
        .unwrap();
    assert_eq!(
        resp.status().as_u16(),
        413,
        "oversized upload must return 413"
    );
}

/// `com.atproto.uploadBlob` rejects requests with no `Authorization`
/// header. Returns `401 Unauthenticated`.
#[tokio::test]
async fn upload_blob_unauthenticated_rejected() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let payload = png_bytes();
    let resp = client()
        .await
        .post(format!("{}/xrpc/com.atproto.uploadBlob", PDS_URL))
        .header("Content-Type", "image/png")
        .body(payload)
        .send()
        .await
        .unwrap();
    assert_eq!(
        resp.status().as_u16(),
        401,
        "no bearer token must return 401"
    );
}

/// `com.atproto.sync.getBlob` returns the resolved `Content-Type` for
/// a previously-uploaded blob. We POST a PNG, GET it back, and
/// verify the response advertises `image/png` rather than the old
/// `application/octet-stream` default.
#[tokio::test]
async fn get_blob_returns_detected_mime_type() {
    if !wait_for_pds().await {
        eprintln!("pds not running, skipping");
        return;
    }
    let (c, did, jwt) = fresh_user("upmime").await;
    seed_any_record(&c, &did, &jwt).await;

    let payload = png_bytes();
    let up: Value = c
        .post(format!("{}/xrpc/com.atproto.uploadBlob", PDS_URL))
        .bearer_auth(&jwt)
        .header("Content-Type", "image/png")
        .body(payload.clone())
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let cid = up["blob"]["ref"]["$link"].as_str().unwrap().to_string();

    let resp = c
        .get(format!(
            "{}/xrpc/com.atproto.sync.getBlob?did={}&cid={}",
            PDS_URL, did, cid
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 200);
    let ct = resp
        .headers()
        .get("content-type")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string();
    let bytes = resp.bytes().await.unwrap();
    assert_eq!(bytes.as_ref(), payload.as_slice());
    assert_eq!(
        ct, "image/png",
        "getBlob should serve the stored mime type, not the octet-stream default"
    );
}