didbot_serve/

blobs.rs

//! The three blob methods, and the streaming that keeps them honest.
//!
//! `com.atproto.repo.uploadBlob`, `com.atproto.sync.getBlob` and
//! `com.atproto.sync.listBlobs`. What is here rather than in
//! [`crate::routes`] is everything to do with a body that does not fit in a
//! `Bytes`: an upload is consumed a frame at a time and handed straight to
//! the store, and a download is a stream over an open file. Neither direction
//! ever holds a whole blob.
//!
//! # Which repository an upload is for
//!
//! `com.atproto.repo.uploadBlob` takes no parameters. In atproto the account
//! comes from the request's authentication, and now that this route has one
//! — see [`Credential::AccountToken`](crate::auth::Credential::AccountToken) —
//! that is exactly where it comes from: the DID
//! [`auth::require_agent_token_or_dpop`] resolves *is* the repository, with no
//! separate field to disagree with it. This used to arrive in a
//! `didbot-repo` header, back when there was no session to take it from;
//! `plan/auth-types.md` names that header as the thing that "comes out the
//! moment a session exists", and this is that moment.

use std::pin::Pin;
use std::task::{Context, Poll};

use axum::body::{Body, Bytes};
use axum::extract::rejection::QueryRejection;
use axum::extract::{Query, State};
use axum::http::{header, HeaderMap, HeaderValue, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::Json;
use didbot_pds::{BlobError, EstopRefusal, Fetch, ProvisionError, DEFAULT_MIME_TYPE};
use futures_core::Stream;
use serde_json::json;

use crate::auth;
use crate::error::ApiError;
use crate::routes::AppState;
use crate::wire::{GetBlobQuery, ListBlobsQuery};

/// How much of a blob is read from the store per chunk on the way out.
///
/// Sixty-four kibibytes, matching the verification pass in the store: the
/// point of both numbers is that neither the syscall nor the buffer is the
/// cost, and a blob at the size limit never has more than this much of itself
/// in memory at once.
const CHUNK: usize = 64 * 1024;

/// How many `com.atproto.sync.getBlob` calls may hold a blocking-pool thread
/// at once, across every caller.
///
/// One read takes a thread twice: once to hash the file against the CID it is
/// filed under, and once for the task that reads the body out, which
/// [`BLOB_READAHEAD`]'s back-pressure parks for as long as the client is slow
/// to take it. That pool is the same one session creation, provisioning,
/// deletion and firehose publishing draw from, and it holds fewer threads
/// than [`crate::MAX_CONCURRENT_REQUESTS`] can put in front of it. One permit
/// covers both threads of one read, from the hash until the body ends.
///
/// A caller that finds every permit taken waits for one, and
/// [`crate::routes::REQUEST_TIMEOUT`]'s layer answers `503` if that wait
/// outlasts the request.
static BLOB_READ_CONCURRENCY: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(32);

/// `POST /xrpc/com.atproto.repo.uploadBlob`
///
/// The body is streamed into the store frame by frame and never assembled.
/// The blob's content identifier is the hash of the bytes, so it does not
/// exist until the last frame has arrived; until then the upload is
/// somewhere no reader can reach, and it becomes a blob in one step at the
/// end. A request that stops half way — a client that disconnects, a body
/// over the size limit — leaves nothing at all, because the store's upload
/// handle cleans up after itself when it is dropped without being committed.
///
/// An app's OAuth token uploads under a `blob:` scope, which names the types
/// it may upload. The declared type is checked against it before the body is
/// read, and the type the bytes are filed under is checked again before the
/// blob is kept. The account's own agent token is not scoped. How large a
/// blob may be is nobody's grant: the store's own limit bounds the body as it
/// arrives, and a policy judges the declared length before the body, again
/// at checkpoints while it arrives, and the size that came in once it is in.
/// A body whose length differs from its `Content-Length` is refused.
///
/// The blob reference in the response names the type the bytes carry a
/// signature for wherever that contradicts the `Content-Type` the request
/// declared, which is one of the two answers atproto's blob specification
/// allows. A client writes that reference into its record, so the record and
/// what `getBlob` answers say the same thing.
pub async fn upload_blob(
    State(state): State<AppState>,
    headers: HeaderMap,
    body: Body,
) -> Response {
    // The e-stop's write gate, unconditional regardless of credential and
    // ahead of the credential check — the stop refuses this upload the same
    // way it refuses every other write, and `docs/write-pipeline.md`'s stage 1
    // answers before anything is spent finding out who is asking.
    if let Err(halted) = state.estop.check_use(EstopRefusal::Operation) {
        return ApiError::from(halted).into_response();
    }
    let writer = match auth::require_agent_token_or_dpop(
        state.registry.as_ref(),
        &state.oauth,
        &headers,
        "POST",
        "/xrpc/com.atproto.repo.uploadBlob",
    ) {
        Ok(writer) => writer,
        Err(err) => return err.into_response(),
    };
    let did = writer.did().to_owned();
    // The lifecycle's half of the same gate every `com.atproto.repo.*` write
    // passes, in the same order; see `routes::check_write_not_halted`. A
    // blob is new data like any other, so an unclaimed server refuses it.
    if let Err(err) = crate::routes::require_record_writes(&state) {
        return err.into_response();
    }

    // `*/*` is what the lexicon declares, so anything is accepted. This is
    // the caller's claim about the body, and the claim is what policy judges;
    // the bytes settle it on commit, and the reference that comes back may
    // name a different type than this one.
    let mime_type = headers
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .filter(|value| !value.is_empty())
        .unwrap_or(DEFAULT_MIME_TYPE)
        .to_owned();
    let declared = headers
        .get(header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<u64>().ok());

    let client_id = match crate::routes::authorize_upload(&state, &writer, &mime_type) {
        Ok(authorized) => authorized.client_id,
        Err(err) => return err.into_response(),
    };
    let mut upload = match state
        .registry
        .begin_blob_from(&did, &mime_type, declared, client_id)
    {
        Ok(upload) => upload,
        Err(err) => return blob_error(&err).into_response(),
    };
    // Charged once policy has judged the upload, so an upload it refuses
    // spends nothing. One operation and no bytes: `[blobs]
    // account_quota_bytes` already bounds what this account's blobs hold,
    // and the log gets one entry.
    if let Err(err) = crate::routes::charge_write(&state, &did, 1, 0) {
        return err.into_response();
    }

    // `futures_core::Stream` is polled by hand because that is the only
    // stream trait this crate depends on; see the `futures-core` note in
    // Cargo.toml. The alternative is `body.collect()`, which is the defect
    // this whole route exists not to have.
    let stream = body.into_data_stream();
    let mut stream = std::pin::pin!(stream);
    // The leading bytes, which the type the blob is filed under is read from.
    let mut head = Vec::new();
    let mut received: u64 = 0;
    while let Some(frame) = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await {
        let chunk = match frame {
            Ok(chunk) => chunk,
            Err(err) => {
                // The upload is dropped here, which discards it. A client
                // that vanished mid-body has left nothing behind.
                tracing::info!(did, %err, "a blob upload's body ended early; discarding it");
                return ApiError::bad_request(format!("the request body ended early: {err}"))
                    .into_response();
            }
        };
        received += chunk.len() as u64;
        if let Some(declared) = declared.filter(|declared| received > *declared) {
            return length_differs(&did, declared, received);
        }
        let room = didbot_pds::blobs::SNIFF_BYTES.saturating_sub(head.len());
        head.extend_from_slice(&chunk[..room.min(chunk.len())]);
        if let Err(err) = upload.write(&chunk) {
            return blob_error(&ProvisionError::Blob(err)).into_response();
        }
    }
    if let Some(declared) = declared.filter(|declared| received != *declared) {
        return length_differs(&did, declared, received);
    }

    // The scope names types, and the bytes can overrule the declared one. A
    // refusal here drops the upload, which discards it.
    let filed = didbot_pds::blobs::filed_type(&mime_type, didbot_pds::blobs::sniff(&head));
    if filed != mime_type {
        if let Err(err) = crate::routes::authorize_upload(&state, &writer, filed) {
            return err.into_response();
        }
    }

    match upload.commit() {
        Ok(reference) => {
            tracing::info!(
                did,
                client = client_id.map(didbot_pds::policy::ClientId::as_str),
                cid = %reference.cid,
                size = reference.size,
                mime_type = %reference.mime_type,
                "uploaded a blob"
            );
            Json(json!({ "blob": reference.to_json() })).into_response()
        }
        Err(err) => blob_error(&ProvisionError::Blob(err)).into_response(),
    }
}

/// The answer to an upload body that is not the length its `Content-Length`
/// declared. atproto's blob specification: "If the actual blob upload size
/// differs from the `Content-Length` header, the server should reject the
/// upload." A body that runs past its length is refused at the frame that
/// runs past it, and no later frame is read.
fn length_differs(did: &str, declared: u64, received: u64) -> Response {
    tracing::info!(
        did,
        declared,
        received,
        "a blob upload's body is not the length it declared; discarding it"
    );
    let message = if received > declared {
        format!("the request body is longer than its Content-Length of {declared} bytes")
    } else {
        format!(
            "the request body is {received} bytes, shorter than its Content-Length of {declared}"
        )
    };
    ApiError::bad_request(message).into_response()
}

/// The content security policy every `com.atproto.sync.getBlob` answer
/// carries.
///
/// A blob is bytes an account uploaded, served from the same origin as
/// `/dashboard`. An HTML or SVG blob opened in a browser would otherwise run
/// its script as this origin, and a `fetch` from there carries the
/// dashboard's session cookie. `sandbox` gives the page an origin of its
/// own, and `default-src 'none'` lets it load nothing. An `<img>` on another
/// page is not bound by it, so an avatar still shows.
const BLOB_CONTENT_SECURITY_POLICY: &str = "default-src 'none'; sandbox";

/// `GET /xrpc/com.atproto.sync.getBlob`
///
/// The bytes as they were uploaded, having first been checked against the
/// name they are filed under. A store that finds the wrong bytes answers
/// `BlobCorrupt` and serves none of them: the CID is a hash, so a reader has
/// every reason to trust what comes back, and that is exactly why a server
/// that can tell must not serve it.
///
/// A store backed by object storage may answer with a redirect instead, and
/// this handler passes one through. See [`Fetch`].
///
/// Every answer carries [`BLOB_CONTENT_SECURITY_POLICY`].
pub async fn get_blob(
    State(state): State<AppState>,
    query: Result<Query<GetBlobQuery>, QueryRejection>,
) -> Response {
    let mut response = serve_blob(state, query).await;
    response.headers_mut().insert(
        header::CONTENT_SECURITY_POLICY,
        HeaderValue::from_static(BLOB_CONTENT_SECURITY_POLICY),
    );
    response
}

/// [`get_blob`]'s answer, before its content security policy.
async fn serve_blob(
    state: AppState,
    query: Result<Query<GetBlobQuery>, QueryRejection>,
) -> Response {
    let Query(query) = match query {
        Ok(query) => query,
        Err(rejection) => return ApiError::bad_request(rejection.body_text()).into_response(),
    };
    // The permit is acquired before `spawn_blocking`, not inside it: what it
    // bounds is blocking-pool threads (see `BLOB_READ_CONCURRENCY`), and
    // claiming one of those first would just hold it idle while this request
    // waits its turn. It is released at the end of this function on every
    // branch that does not stream, and moved into `BlobStream` on the one
    // that does.
    let permit = BLOB_READ_CONCURRENCY
        .acquire()
        .await
        .expect("this semaphore is never closed");
    // Off the async worker: `FileBlobStore::fetch` hashes the whole file
    // against the CID it is filed under before it hands back a single byte,
    // so this call is as long as the blob is.
    let registry = state.registry.clone();
    let (did, cid) = (query.did.clone(), query.cid.clone());
    let fetched = tokio::task::spawn_blocking(move || registry.fetch_blob(&did, &cid)).await;
    let fetched = match fetched {
        Ok(fetched) => fetched,
        Err(join_err) => {
            return crate::routes::blocking_task_failed("com.atproto.sync.getBlob", join_err)
                .into_response()
        }
    };
    match fetched {
        Ok(Fetch::Stream {
            mime_type,
            size,
            body,
        }) => (
            [
                (header::CONTENT_TYPE, mime_type),
                (header::CONTENT_LENGTH, size.to_string()),
                // A blob never changes under its own name, so a cache may
                // keep it as long as it likes. This is the one answer this
                // server gives where that is true by construction.
                (
                    header::CACHE_CONTROL,
                    "public, max-age=31536000, immutable".to_owned(),
                ),
            ],
            Body::from_stream(BlobStream::new(body, permit)),
        )
            .into_response(),
        Ok(Fetch::Redirect { url, expires_in }) => (
            StatusCode::FOUND,
            [
                (header::LOCATION, url),
                (
                    header::CACHE_CONTROL,
                    format!("private, max-age={}", expires_in.as_secs()),
                ),
            ],
        )
            .into_response(),
        Err(ProvisionError::UnknownAccount { did }) => {
            crate::routes::repo_not_served(&state, &did).into_response()
        }
        Err(err) => blob_error(&err).into_response(),
    }
}

/// `GET /xrpc/com.atproto.sync.listBlobs`
///
/// Every blob CID one account holds, in CID order, paged on the last CID of
/// the previous page. `since` narrows that to the blobs referenced by
/// records that changed after the named revision — see
/// [`Registry::list_blobs`](didbot_pds::Registry::list_blobs) for how, and for
/// what happens when this deployment cannot place the revision.
///
/// It takes no credential, and a `since` walks every record changed after
/// it, so it spends the caller's sync budget and runs on the blocking pool
/// like the other `com.atproto.sync` reads.
pub async fn list_blobs(
    State(state): State<AppState>,
    peer: crate::rate_limit::Peer,
    headers: HeaderMap,
    query: Result<Query<ListBlobsQuery>, QueryRejection>,
) -> Response {
    if let Err(refusal) = crate::routes::require_sync_budget(&state, &headers, peer) {
        return refusal.into_response();
    }
    let Query(query) = match query {
        Ok(query) => query,
        Err(rejection) => return ApiError::bad_request(rejection.body_text()).into_response(),
    };
    let listed = crate::routes::sync_read(&state, "com.atproto.sync.listBlobs", move |registry| {
        registry.list_blobs(
            &query.did,
            query.effective_limit(),
            query.cursor.as_deref(),
            query.since.as_deref(),
        )
    })
    .await;
    match listed {
        Ok((cids, Some(cursor))) => Json(json!({ "cids": cids, "cursor": cursor })).into_response(),
        Ok((cids, None)) => Json(json!({ "cids": cids })).into_response(),
        Err(response) => response,
    }
}

/// Maps a blob failure onto the error name a client switches on.
///
/// `BlobNotFound` and `RepoNotFound` are the lexicons' own names for
/// `getBlob`. The others are this server's, because the lexicons declare
/// nothing for them: a size limit and the quotas are a deployment's policy
/// rather than the protocol's, and corruption is not a thing the protocol
/// expects a server to be able to notice.
pub fn blob_error(err: &ProvisionError) -> ApiError {
    let ProvisionError::Blob(blob) = err else {
        return match err {
            ProvisionError::UnknownAccount { did } => ApiError::new(
                StatusCode::NOT_FOUND,
                "RepoNotFound",
                format!("no repository here for `{did}`"),
            ),
            other => ApiError::from(other),
        };
    };
    match blob {
        BlobError::NotFound { .. } => {
            ApiError::new(StatusCode::NOT_FOUND, "BlobNotFound", blob.to_string())
        }
        BlobError::MalformedCid { .. } => ApiError::bad_request(blob.to_string()),
        // 413 rather than 400, because the request is well formed and the
        // only thing wrong with it is its size — which is what that status
        // says and what a client backing off needs to be able to tell.
        BlobError::TooLarge { .. } => ApiError::new(
            StatusCode::PAYLOAD_TOO_LARGE,
            "BlobTooLarge",
            blob.to_string(),
        ),
        // 403: nothing about the request would make it succeed later without
        // the account first giving something back.
        BlobError::QuotaExceeded { .. } => ApiError::new(
            StatusCode::FORBIDDEN,
            "AccountQuotaExceeded",
            blob.to_string(),
        ),
        // 507, and the name a full log answers with: the server is out of
        // the room it set aside, and the same upload works once blobs are
        // collected or deleted somewhere on it.
        BlobError::TotalQuotaExceeded { .. } => ApiError::new(
            StatusCode::INSUFFICIENT_STORAGE,
            "StorageFull",
            blob.to_string(),
        ),
        BlobError::Corrupt { .. } => ApiError::new(
            StatusCode::INTERNAL_SERVER_ERROR,
            "BlobCorrupt",
            blob.to_string(),
        ),
        BlobError::Backend(_) => ApiError::new(
            StatusCode::INTERNAL_SERVER_ERROR,
            "InternalServerError",
            blob.to_string(),
        ),
        // The same answers a record write refused by policy gets.
        BlobError::PolicyRejected { reason } => ApiError::from(&ProvisionError::PolicyRejected {
            reason: reason.clone(),
        }),
        BlobError::PolicyFrozen { reason } => ApiError::from(&ProvisionError::PolicyFrozen {
            reason: reason.clone(),
        }),
    }
}

/// How many chunks the reading task may run ahead of the socket.
///
/// The back-pressure: the reader fills this and then waits, so a slow client
/// costs one task and [`BLOB_READAHEAD`] × [`CHUNK`] of buffer rather than
/// the whole blob in memory.
const BLOB_READAHEAD: usize = 4;

/// A blob on its way out, a chunk at a time.
///
/// The store hands back a plain [`std::io::Read`], and this turns it into the
/// stream `axum::body::Body` wants. The read is blocking, which is the same
/// posture the rest of this server takes towards its disk — the write-ahead
/// log is written from inside async handlers too — and it is what keeps
/// the store's interface one trait rather than one per runtime.
/// [`BLOB_READ_CONCURRENCY`] is what bounds how many of those reads run at
/// once.
struct BlobStream {
    chunks: tokio::sync::mpsc::Receiver<Result<Bytes, std::io::Error>>,
}

impl BlobStream {
    /// Reads `reader` on the blocking pool, handing chunks over a bounded
    /// channel.
    ///
    /// `std::io::Read` on a file blocks the thread that calls it, and before
    /// this that thread was an async worker: a deployment with two of them
    /// (the `t3.small` `infra/pds` deploys onto) had every route, the
    /// firehose included, waiting behind two blob reads. The read happens on
    /// a blocking thread now, and the channel is what keeps it from running
    /// ahead of the socket.
    ///
    /// `permit` is the caller's [`BLOB_READ_CONCURRENCY`] permit, held for as
    /// long as the reading task holds its thread and given back the moment
    /// that task ends — whether the blob ran out, the read failed, or the
    /// client went away.
    fn new(
        mut reader: Box<dyn std::io::Read + Send>,
        permit: tokio::sync::SemaphorePermit<'static>,
    ) -> Self {
        let (chunks, rx) = tokio::sync::mpsc::channel(BLOB_READAHEAD);
        tokio::task::spawn_blocking(move || {
            let _permit = permit;
            let mut buffer = vec![0u8; CHUNK];
            loop {
                let chunk = match reader.read(&mut buffer) {
                    Ok(0) => return,
                    Ok(read) => Ok(Bytes::copy_from_slice(&buffer[..read])),
                    // The stream ends here rather than repeating the failure:
                    // a body that has already begun cannot become an error
                    // response, so the honest thing left is a short one that
                    // the declared `Content-Length` makes visible as
                    // truncated.
                    Err(err) => {
                        let _ = chunks.blocking_send(Err(err));
                        return;
                    }
                };
                // A closed channel is a client that went away mid-body, and
                // the read stops with it rather than pulling the rest of the
                // blob off the disk for nobody.
                if chunks.blocking_send(chunk).is_err() {
                    return;
                }
            }
        });
        Self { chunks: rx }
    }
}

impl Stream for BlobStream {
    type Item = Result<Bytes, std::io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.get_mut().chunks.poll_recv(cx)
    }
}

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

    /// A reader whose first `read` parks until it is released, standing in
    /// for a disk that is slow — or, for `FileBlobStore`, for the whole-file
    /// re-hash it does before it hands back a byte.
    struct ParkedReader {
        go: Option<std::sync::mpsc::Receiver<()>>,
    }

    /// A one-permit semaphore of this test's own, standing in for
    /// [`BLOB_READ_CONCURRENCY`] so that what one test counts cannot be
    /// another test's blob read.
    fn own_semaphore() -> &'static tokio::sync::Semaphore {
        Box::leak(Box::new(tokio::sync::Semaphore::new(1)))
    }

    impl std::io::Read for ParkedReader {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            match self.go.take() {
                Some(go) => {
                    let _ = go.recv();
                    buf[..3].copy_from_slice(b"hi!");
                    Ok(3)
                }
                None => Ok(0),
            }
        }
    }

    /// **What a blob read used to cost.** `std::io::Read` blocks the thread
    /// that calls it, and `poll_next` was calling it on an async worker — so
    /// on the two-worker host `infra/pds` deploys onto, two slow blob reads
    /// were every route waiting, the firehose included.
    ///
    /// One worker here, so the property is exact rather than probabilistic:
    /// while the read is parked, another task still runs.
    #[tokio::test(flavor = "current_thread")]
    async fn a_blob_body_is_read_off_the_async_worker() {
        let (release, go) = std::sync::mpsc::channel();
        let semaphore = own_semaphore();
        let permit = semaphore.try_acquire().expect("a free permit");
        let mut stream = Box::pin(BlobStream::new(
            Box::new(ParkedReader { go: Some(go) }),
            permit,
        ));

        // The read is parked inside `ParkedReader::read`. The worker is free
        // only if that call is not on it.
        let ran = Arc::new(std::sync::atomic::AtomicBool::new(false));
        let flag = ran.clone();
        tokio::spawn(async move {
            flag.store(true, std::sync::atomic::Ordering::SeqCst);
        });
        let polled = std::future::poll_fn(|cx| {
            let _ = stream.as_mut().poll_next(cx);
            Poll::Ready(())
        });
        tokio::time::timeout(std::time::Duration::from_secs(5), polled)
            .await
            .expect("polling a parked blob body must not block the worker");
        tokio::task::yield_now().await;
        assert!(
            ran.load(std::sync::atomic::Ordering::SeqCst),
            "the parked read held the only async worker"
        );

        let _ = release.send(());
        let chunk = tokio::time::timeout(
            std::time::Duration::from_secs(5),
            std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)),
        )
        .await
        .expect("the released read answers")
        .expect("a chunk")
        .expect("read without error");
        assert_eq!(chunk.as_ref(), b"hi!");
    }

    /// **A download that ends early still gives its thread back.** The permit
    /// bounds blocking-pool threads and is held for the whole body, so the
    /// case to get right is the body that never finishes: a client that
    /// disconnects mid-download leaves the reading task parked in a read it
    /// only abandons on its next send.
    #[tokio::test]
    async fn a_blob_body_gives_its_permit_back_when_the_client_goes_away() {
        let (release, go) = std::sync::mpsc::channel();
        let semaphore = own_semaphore();
        let permit = semaphore.try_acquire().expect("a free permit");
        let stream = BlobStream::new(Box::new(ParkedReader { go: Some(go) }), permit);
        assert_eq!(
            semaphore.available_permits(),
            0,
            "a blob read in flight holds its permit"
        );

        drop(stream);
        let _ = release.send(());

        let returned = tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while semaphore.available_permits() == 0 {
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        })
        .await;
        assert!(
            returned.is_ok(),
            "the permit outlived the blob read it bounds"
        );
    }
}