ferro-oci-server 1.0.0

// SPDX-License-Identifier: Apache-2.0
//! Axum router factory that wires the `/v2/**` OCI endpoints.
//!
//! Spec: OCI Distribution Spec v1.1 §3 "API".
//!
//! The router is stateful — it takes an [`AppState`] carrying the blob
//! store and the registry-metadata plane. Callers build the state once
//! at boot and then call [`router`] to obtain an `axum::Router` they
//! can mount under `/`.

use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};

use axum::Router;
use axum::extract::DefaultBodyLimit;
use axum::routing::{delete, get, post};

use crate::handlers::{base, catalog};
use crate::registry::RegistryMeta;
use ferro_blob_store::SharedBlobStore;

/// Maximum request body size (bytes) accepted on the `/v2/**` surface.
///
/// Axum's [`DefaultBodyLimit`] defaults to 2 MiB, which silently rejects
/// the manifest and blob-chunk pushes real clients send: the OCI
/// Distribution Spec expects registries to accept manifests of at least
/// 4 MiB, and individual blob chunks are routinely larger. We raise the
/// limit to 512 MiB — comfortably above any manifest and large enough
/// for a generous single blob chunk, while still bounding the bytes a
/// single request can buffer. Total per-session upload growth is
/// separately bounded by [`crate::upload::MAX_UPLOAD_SESSION_BYTES`].
///
/// The OCI surface routes every method through one `/v2/{*rest}`
/// wildcard, so a single body limit covers manifests and blobs alike;
/// the value is chosen to satisfy the larger (blob) case.
pub const MAX_BODY_BYTES: usize = 512 * 1024 * 1024;

/// Shared HTTP handler state.
pub struct AppState {
    /// Blob-bytes plane.
    pub blob_store: SharedBlobStore,
    /// Metadata plane (manifests, tags, upload sessions, referrers).
    pub registry: Arc<dyn RegistryMeta>,
    /// Best-effort count of distinct blobs this server has written.
    ///
    /// Maintained incrementally by the blob handlers (`+1` on a
    /// successful blob `put`, `-1` on a successful `delete`) so the
    /// `/metrics` scrape can report `ferrooci_storage_blobs` in O(1)
    /// instead of an O(number-of-blobs) `BlobStore::list()` filesystem
    /// scan on every request — an open `/metrics` endpoint would
    /// otherwise become a cheap amplification vector.
    ///
    /// It is a best-effort gauge: it counts blob writes/deletes observed
    /// through *this* process and does not deduplicate a re-`put` of an
    /// already-present digest, nor does it reflect blobs written by other
    /// processes against a shared filesystem store. The honest claim is
    /// "blobs written via this server instance".
    blob_count: Arc<AtomicI64>,
    /// Maximum total bytes a single in-flight upload session may buffer
    /// before the server refuses further chunks (memory-exhaustion `DoS`
    /// bound). Defaults to [`crate::upload::MAX_UPLOAD_SESSION_BYTES`];
    /// overridable via [`AppState::with_max_upload_session_bytes`] (used
    /// by tests to exercise the cap without allocating gigabytes).
    max_upload_session_bytes: u64,
    /// Serialises the `contains` → `put` → gauge-increment critical
    /// section for blob writes (R3-3).
    ///
    /// The `storage_blobs` gauge must reflect *distinct* blobs, so it is
    /// only incremented when a digest is genuinely new. The naive
    /// "check `contains`, then `put`, then `inc` if it was absent" sequence
    /// races: two concurrent uploads of the *same* digest can both observe
    /// it absent and both increment, drifting the gauge above the true
    /// count. Holding this mutex across the whole check+put+inc region
    /// (without touching `ferro-blob-store`) means only the first inserter
    /// of a given digest increments. It is a single accounting mutex rather
    /// than a per-digest map: blob writes are not the hot path, and a flat
    /// mutex keeps the invariant obviously correct.
    blob_accounting: tokio::sync::Mutex<()>,
}

impl AppState {
    /// Construct shared handler state from a blob store and a
    /// registry-metadata plane, wrapped in an [`Arc`] ready to hand to
    /// [`router`].
    ///
    /// ```no_run
    /// use std::sync::Arc;
    /// use ferro_blob_store::InMemoryBlobStore;
    /// use ferro_oci_server::{AppState, InMemoryRegistryMeta, router};
    ///
    /// let state = AppState::new(
    ///     Arc::new(InMemoryBlobStore::new()),
    ///     Arc::new(InMemoryRegistryMeta::new()),
    /// );
    /// let app = router(state);
    /// ```
    #[must_use]
    pub fn new(blob_store: SharedBlobStore, registry: Arc<dyn RegistryMeta>) -> Arc<Self> {
        Arc::new(Self {
            blob_store,
            registry,
            blob_count: Arc::new(AtomicI64::new(0)),
            max_upload_session_bytes: crate::upload::MAX_UPLOAD_SESSION_BYTES,
            blob_accounting: tokio::sync::Mutex::new(()),
        })
    }

    /// Build state with a custom upload-session size cap.
    ///
    /// Behaves like [`AppState::new`] but overrides the per-session byte
    /// cap (the memory-exhaustion `DoS` bound). Production callers use the
    /// default via [`AppState::new`]; this exists so tests can drive the
    /// cap-exceeded path with a handful of bytes instead of gigabytes.
    #[must_use]
    pub fn with_max_upload_session_bytes(
        blob_store: SharedBlobStore,
        registry: Arc<dyn RegistryMeta>,
        max_upload_session_bytes: u64,
    ) -> Arc<Self> {
        Arc::new(Self {
            blob_store,
            registry,
            blob_count: Arc::new(AtomicI64::new(0)),
            max_upload_session_bytes,
            blob_accounting: tokio::sync::Mutex::new(()),
        })
    }

    /// The per-session upload byte cap currently in force.
    #[must_use]
    pub const fn max_upload_session_bytes(&self) -> u64 {
        self.max_upload_session_bytes
    }

    /// Shared handle to the incremental blob counter.
    ///
    /// Cloned into the metrics layer so `/metrics` can report
    /// `ferrooci_storage_blobs` without a filesystem scan. See the
    /// [`blob_count`](AppState::blob_count) field docs for the exact
    /// semantics.
    #[must_use]
    pub fn blob_count_handle(&self) -> Arc<AtomicI64> {
        Arc::clone(&self.blob_count)
    }

    /// Record that one blob was written via this server.
    pub fn inc_blob_count(&self) {
        self.blob_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Put `body` under `digest`, incrementing the distinct-blobs gauge
    /// exactly once for a genuinely new digest (R3-3).
    ///
    /// The `contains` → `put` → `inc` sequence is held under the
    /// [`blob_accounting`](AppState::blob_accounting) mutex so two
    /// concurrent uploads of the *same* digest cannot both observe it
    /// absent and both increment the gauge (which would drift it above the
    /// true count). Only the first inserter increments. The blob store
    /// itself is content-addressed and idempotent on a duplicate `put`, so
    /// serialising here changes only the gauge accounting, not blob bytes.
    ///
    /// # Errors
    ///
    /// Propagates any [`ferro_blob_store::BlobStoreError`] from the
    /// underlying `contains` / `put`.
    pub async fn store_blob_counted(
        &self,
        digest: &ferro_blob_store::Digest,
        body: axum::body::Bytes,
    ) -> ferro_blob_store::Result<()> {
        let _accounting = self.blob_accounting.lock().await;
        let already_present = self.blob_store.contains(digest).await.unwrap_or(false);
        self.blob_store.put(digest, body).await?;
        if !already_present {
            self.inc_blob_count();
        }
        Ok(())
    }

    /// Record that one blob was deleted via this server, saturating at 0.
    pub fn dec_blob_count(&self) {
        // Saturate at zero: a delete of a blob this process never counted
        // (e.g. a pre-existing blob in a shared FS store) must not drive
        // the gauge negative.
        let _ = self.blob_count.fetch_update(
            Ordering::Relaxed,
            Ordering::Relaxed,
            |n| Some(n.saturating_sub(1).max(0)),
        );
    }

    /// Current best-effort blob count.
    #[must_use]
    pub fn blob_count(&self) -> i64 {
        self.blob_count.load(Ordering::Relaxed)
    }
}

#[cfg(test)]
mod app_state_tests {
    use std::sync::Arc;

    use ferro_blob_store::InMemoryBlobStore;

    use super::AppState;
    use crate::registry::InMemoryRegistryMeta;

    fn state() -> Arc<AppState> {
        AppState::new(
            Arc::new(InMemoryBlobStore::new()),
            Arc::new(InMemoryRegistryMeta::new()),
        )
    }

    #[test]
    fn blob_count_reflects_increments_and_saturating_decrements() {
        // `blob_count()` must return the live count, not a constant.
        // Kills `-> 0`, `-> 1`, and `-> -1` return-constant mutants by
        // asserting three distinct non-constant values (0, 2, 1).
        let st = state();
        assert_eq!(st.blob_count(), 0, "fresh state starts at 0");

        st.inc_blob_count();
        st.inc_blob_count();
        assert_eq!(st.blob_count(), 2, "two increments give 2");

        st.dec_blob_count();
        assert_eq!(st.blob_count(), 1, "one decrement gives 1");
    }

    #[test]
    fn blob_count_saturates_at_zero_never_negative() {
        // A decrement below zero saturates (never -1). This also rules
        // out the `-> -1` constant mutant from a different angle.
        let st = state();
        st.dec_blob_count();
        assert_eq!(st.blob_count(), 0, "decrement on empty stays at 0");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn r3_3_concurrent_same_digest_put_increments_gauge_once() {
        // R3-3: two concurrent uploads of the SAME digest must increment
        // the distinct-blobs gauge exactly once. `store_blob_counted`
        // serialises the contains→put→inc region under an accounting mutex
        // so only the first inserter counts; without it both could observe
        // the digest absent and both increment (gauge == 2 = drift).
        use bytes::Bytes;
        use ferro_blob_store::Digest;

        let st = state();
        let body = Bytes::from_static(b"the-same-blob-bytes");
        let digest = Digest::sha256_of(&body);

        // Spawn many concurrent puts of the identical digest to widen the
        // race window the old non-atomic code would have lost on.
        let mut handles = Vec::new();
        for _ in 0..16 {
            let st = Arc::clone(&st);
            let digest = digest.clone();
            let body = body.clone();
            handles.push(tokio::spawn(async move {
                st.store_blob_counted(&digest, body).await.expect("put");
            }));
        }
        for h in handles {
            h.await.expect("join");
        }

        assert_eq!(
            st.blob_count(),
            1,
            "concurrent puts of one digest must count it exactly once"
        );
    }

    #[tokio::test]
    async fn r3_3_distinct_digests_each_counted_once() {
        // Sanity: distinct digests each increment exactly once (the lock
        // serialises but does not over-deduplicate genuinely new blobs).
        use bytes::Bytes;
        use ferro_blob_store::Digest;

        let st = state();
        for i in 0..5u8 {
            let body = Bytes::from(vec![i; 4]);
            let digest = Digest::sha256_of(&body);
            st.store_blob_counted(&digest, body).await.expect("put");
            // A duplicate put of the same digest must not double-count.
            let body2 = Bytes::from(vec![i; 4]);
            let digest2 = Digest::sha256_of(&body2);
            st.store_blob_counted(&digest2, body2).await.expect("put2");
        }
        assert_eq!(st.blob_count(), 5, "five distinct blobs counted once each");
    }
}

/// Build the Axum router for every `/v2/**` OCI endpoint.
///
/// This router covers only the OCI Distribution Spec surface. To add
/// Kubernetes liveness / readiness probes (`/live`, `/healthz`,
/// `/ready`), merge [`probe_routes`] into the returned router with
/// [`Router::merge`].
pub fn router(state: Arc<AppState>) -> Router {
    Router::new()
        // Version / auth challenge.
        .route("/v2/", get(base::version_check))
        .route("/v2", get(base::version_check))
        // Catalog and tag listing.
        .route("/v2/_catalog", get(catalog::list_catalog))
        .route("/v2/{*rest}", get(dispatch::dispatch_get))
        .route("/v2/{*rest}", axum::routing::head(dispatch::dispatch_head))
        .route("/v2/{*rest}", delete(dispatch::dispatch_delete))
        // Blob uploads -- POST / PATCH / PUT.
        .route(
            "/v2/{*rest}",
            post(dispatch::dispatch_post)
                .patch(dispatch::dispatch_patch_inner)
                .put(dispatch::dispatch_put_inner),
        )
        // Raise the body limit above Axum's 2 MiB default so manifest
        // (≥4 MiB per spec) and blob-chunk pushes are accepted. See
        // [`MAX_BODY_BYTES`].
        .layer(DefaultBodyLimit::max(MAX_BODY_BYTES))
        .with_state(state)
}

/// Build a stateless router exposing Kubernetes-style health probes.
///
/// Routes:
///
/// - `GET /live` — liveness; returns `200 OK` with body `OK` as long as
///   the process is running and able to serve requests.
/// - `GET /healthz` — health; returns `200 OK` with JSON
///   `{"status":"ok"}`.
/// - `GET /ready` — readiness; returns `200 OK` with body `OK`. The
///   in-memory metadata plane and blob store are ready as soon as the
///   process is up, so this mirrors liveness today; a persistent backend
///   would gate this on a successful storage ping.
///
/// Merge into the OCI router at boot:
///
/// ```no_run
/// use std::sync::Arc;
/// use ferro_blob_store::InMemoryBlobStore;
/// use ferro_oci_server::{AppState, InMemoryRegistryMeta, probe_routes, router};
///
/// let state = AppState::new(
///     Arc::new(InMemoryBlobStore::new()),
///     Arc::new(InMemoryRegistryMeta::new()),
/// );
/// let app = router(state).merge(probe_routes());
/// ```
pub fn probe_routes() -> Router {
    use axum::Json;
    use axum::http::StatusCode;
    use serde_json::json;

    Router::new()
        .route("/live", get(|| async { (StatusCode::OK, "OK") }))
        .route(
            "/healthz",
            get(|| async { (StatusCode::OK, Json(json!({ "status": "ok" }))) }),
        )
        .route("/ready", get(|| async { (StatusCode::OK, "OK") }))
}

/// Small axum-aware dispatch layer.
///
/// Distribution routes use the `{name}` path parameter which can contain
/// slashes. Axum's `{*rest}` wildcard allows us to greedily capture the
/// path tail and then inspect the suffix ourselves. We then dispatch to
/// the real handler based on the suffix shape — `blobs/{digest}`,
/// `blobs/uploads/{uuid?}`, `manifests/{reference}`, `tags/list`, or
/// `referrers/{digest}`.
pub mod dispatch {
    use std::sync::Arc;

    use axum::body::Bytes;
    use axum::extract::{Path, Query, State};
    use axum::http::{HeaderMap, Method, StatusCode};
    use axum::response::{IntoResponse, Response};

    use super::AppState;
    use crate::error::{OciError, OciErrorCode};
    use crate::handlers::{blob, blob_upload, manifest as manifest_h, referrers, tags};

    /// Split `rest` into `(name, suffix)` where `suffix` is one of
    /// `blobs/...`, `manifests/...`, `tags/list`, `referrers/...`.
    fn split_rest(rest: &str) -> Option<(&str, &str)> {
        // Walk the path and find the last segment boundary where the
        // suffix starts with a known keyword. This accommodates
        // multi-level names like `my-org/library/alpine`.
        let keywords = ["blobs/", "manifests/", "tags/list", "referrers/"];
        for kw in keywords {
            if let Some(idx) = rest.rfind(kw) {
                // Ensure the match is preceded by a `/` or is at idx 0
                // after the name prefix.
                if idx == 0 {
                    return None;
                }
                if &rest[idx - 1..idx] != "/" {
                    continue;
                }
                let name = &rest[..idx - 1];
                let suffix = &rest[idx..];
                return Some((name, suffix));
            }
        }
        None
    }

    /// Decode the rest into (name, suffix) or return `NAME_INVALID`.
    fn decode(rest: &str) -> Result<(String, String), OciError> {
        let (name, suffix) = split_rest(rest).ok_or_else(|| {
            OciError::new(OciErrorCode::NameUnknown, format!("cannot route `{rest}`"))
        })?;
        Ok((name.to_owned(), suffix.to_owned()))
    }

    /// GET dispatcher.
    pub async fn dispatch_get(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        Query(params): Query<std::collections::BTreeMap<String, String>>,
        headers: HeaderMap,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(
            state,
            name,
            suffix,
            Method::GET,
            headers,
            params,
            Bytes::new(),
        )
        .await
    }

    /// HEAD dispatcher.
    pub async fn dispatch_head(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        headers: HeaderMap,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(
            state,
            name,
            suffix,
            Method::HEAD,
            headers,
            std::collections::BTreeMap::default(),
            Bytes::new(),
        )
        .await
    }

    /// DELETE dispatcher.
    pub async fn dispatch_delete(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        headers: HeaderMap,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(
            state,
            name,
            suffix,
            Method::DELETE,
            headers,
            std::collections::BTreeMap::default(),
            Bytes::new(),
        )
        .await
    }

    /// POST dispatcher (blob upload init).
    pub async fn dispatch_post(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        Query(params): Query<std::collections::BTreeMap<String, String>>,
        headers: HeaderMap,
        body: Bytes,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(state, name, suffix, Method::POST, headers, params, body).await
    }

    /// PATCH dispatcher.
    pub async fn dispatch_patch_inner(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        headers: HeaderMap,
        body: Bytes,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(
            state,
            name,
            suffix,
            Method::PATCH,
            headers,
            std::collections::BTreeMap::default(),
            body,
        )
        .await
    }

    /// PUT dispatcher.
    pub async fn dispatch_put_inner(
        State(state): State<Arc<AppState>>,
        Path(rest): Path<String>,
        Query(params): Query<std::collections::BTreeMap<String, String>>,
        headers: HeaderMap,
        body: Bytes,
    ) -> Response {
        let (name, suffix) = match decode(&rest) {
            Ok(v) => v,
            Err(e) => return e.into_response(),
        };
        dispatch_inner(state, name, suffix, Method::PUT, headers, params, body).await
    }

    #[allow(clippy::too_many_arguments)]
    async fn dispatch_inner(
        state: Arc<AppState>,
        name: String,
        suffix: String,
        method: Method,
        headers: HeaderMap,
        params: std::collections::BTreeMap<String, String>,
        body: Bytes,
    ) -> Response {
        // Tag listing.
        if suffix == "tags/list" {
            return if method == Method::GET {
                tags::list_tags(&state, &name, &params)
                    .await
                    .into_response()
            } else {
                OciError::new(OciErrorCode::Unsupported, "unsupported method")
                    .with_status(StatusCode::METHOD_NOT_ALLOWED)
                    .into_response()
            };
        }
        // Referrers.
        if let Some(rest) = suffix.strip_prefix("referrers/") {
            return if method == Method::GET {
                referrers::get_referrers(&state, &name, rest, &params)
                    .await
                    .into_response()
            } else {
                OciError::new(OciErrorCode::Unsupported, "unsupported method")
                    .with_status(StatusCode::METHOD_NOT_ALLOWED)
                    .into_response()
            };
        }
        // Manifests.
        if let Some(rest) = suffix.strip_prefix("manifests/") {
            return match method {
                Method::GET => manifest_h::get_manifest(&state, &name, rest, &headers)
                    .await
                    .into_response(),
                Method::HEAD => manifest_h::head_manifest(&state, &name, rest)
                    .await
                    .into_response(),
                Method::PUT => manifest_h::put_manifest(&state, &name, rest, &headers, body)
                    .await
                    .into_response(),
                Method::DELETE => manifest_h::delete_manifest(&state, &name, rest)
                    .await
                    .into_response(),
                _ => OciError::new(OciErrorCode::Unsupported, "unsupported method")
                    .with_status(StatusCode::METHOD_NOT_ALLOWED)
                    .into_response(),
            };
        }
        // Blob uploads.
        if let Some(rest) = suffix.strip_prefix("blobs/uploads/") {
            let uuid = rest.trim_end_matches('/');
            return match method {
                Method::POST => {
                    // `rest` is "" for the "/blobs/uploads/" endpoint.
                    blob_upload::init_upload(&state, &name, &headers, &params, body)
                        .await
                        .into_response()
                }
                Method::PATCH => blob_upload::patch_upload(&state, &name, uuid, &headers, body)
                    .await
                    .into_response(),
                Method::PUT => blob_upload::finish_upload(&state, &name, uuid, &params, body)
                    .await
                    .into_response(),
                Method::GET => blob_upload::get_upload_status(&state, &name, uuid)
                    .await
                    .into_response(),
                Method::DELETE => blob_upload::cancel_upload(&state, &name, uuid)
                    .await
                    .into_response(),
                _ => OciError::new(OciErrorCode::Unsupported, "unsupported method")
                    .with_status(StatusCode::METHOD_NOT_ALLOWED)
                    .into_response(),
            };
        }
        // Blobs (by digest).
        if let Some(rest) = suffix.strip_prefix("blobs/") {
            return match method {
                Method::GET => blob::get_blob(&state, &name, rest).await.into_response(),
                Method::HEAD => blob::head_blob(&state, &name, rest).await.into_response(),
                Method::DELETE => blob::delete_blob(&state, &name, rest).await.into_response(),
                _ => OciError::new(OciErrorCode::Unsupported, "unsupported method")
                    .with_status(StatusCode::METHOD_NOT_ALLOWED)
                    .into_response(),
            };
        }
        OciError::new(
            OciErrorCode::NameUnknown,
            format!("cannot route `{name}/{suffix}`"),
        )
        .into_response()
    }

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

        #[test]
        fn split_simple_manifest_path() {
            let (name, suffix) = split_rest("alpine/manifests/latest").expect("split");
            assert_eq!(name, "alpine");
            assert_eq!(suffix, "manifests/latest");
        }

        #[test]
        fn split_nested_blob_path() {
            let (name, suffix) = split_rest("my-org/lib/alpine/blobs/uploads/abc").expect("split");
            assert_eq!(name, "my-org/lib/alpine");
            assert_eq!(suffix, "blobs/uploads/abc");
        }

        #[test]
        fn split_tags_list() {
            let (name, suffix) = split_rest("lib/alpine/tags/list").expect("split");
            assert_eq!(name, "lib/alpine");
            assert_eq!(suffix, "tags/list");
        }

        #[test]
        fn split_referrers() {
            let (name, suffix) = split_rest("lib/alpine/referrers/sha256:abcd").expect("split");
            assert_eq!(name, "lib/alpine");
            assert_eq!(suffix, "referrers/sha256:abcd");
        }

        #[test]
        fn split_none_for_bare_name() {
            assert!(split_rest("alpine").is_none());
        }
    }
}