fix: reduce memory usage of SBOM browser

tools/harbor-sbom-browser/Cargo.toml

@@ -7,17 +7,16 @@ edition = "2021"
 
 [dependencies]
 axum = "0.8.9"
-axum-extra = { version = "0.12.6", features = ["erased-json"] }
 base64 = "0.23.1"
 futures = "0.3.30"
 lazy_static = "1.4.0"
 regex = "1.10.3"
 reqwest = { version = "0.13.5", features = ["json", "rustls"], default-features = false }
 serde = { version = "1.0.196", features = ["derive"] }
-serde_json = "1.0.113"
+serde_json = { version = "1.0.113", features = ["raw_value"] }
 snafu = "0.9.2"
 strum = { version = "0.28.0", features = ["derive"] }
-tokio = { version = "1.36.0", features = ["macros", "process", "rt-multi-thread"] }
+tokio = { version = "1.36.0", features = ["macros", "process", "rt-multi-thread", "sync", "time"] }
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
 urlencoding = "2.1.3"

tools/harbor-sbom-browser/src/handlers/mod.rs

@@ -1,2 +1,3 @@
 pub mod artifact_tree;
+pub mod robots;
 pub mod sbom;

tools/harbor-sbom-browser/src/handlers/robots.rs

@@ -0,0 +1,7 @@
+/// Crawlers request every SBOM download link (one `cosign` run each) in bursts, which is not
+/// what the SBOM browser is for. The links already carry `rel='nofollow'`, but that is only a hint.
+const ROBOTS_TXT: &str = "User-agent: *\nDisallow: /sbom/\n";
+
+pub async fn robots_txt() -> &'static str {
+    ROBOTS_TXT
+}

tools/harbor-sbom-browser/src/handlers/sbom.rs

@@ -2,13 +2,12 @@ use axum::{
     extract::Path,
     http::{header, HeaderMap},
 };
-use axum_extra::response::ErasedJson;
 
 use crate::utils::{verify_attestation, DownloadSbomError};
 
 pub async fn download(
     Path((repository, digest)): Path<(String, String)>,
-) -> Result<(HeaderMap, ErasedJson), DownloadSbomError> {
+) -> Result<(HeaderMap, String), DownloadSbomError> {
     let attestation = verify_attestation(&repository, &digest).await?;
     let mut headers = HeaderMap::new();
     headers.insert(header::CONTENT_TYPE, "application/json".parse().unwrap());
@@ -18,5 +17,8 @@ pub async fn download(
             .parse()
             .unwrap(),
     );
-    Ok((headers, ErasedJson::pretty(attestation.predicate)))
+    // The SBOM is returned as it is in the attestation instead of being pretty-printed, which would
+    // need the whole document parsed into memory.
+    let sbom: Box<str> = attestation.predicate.into();
+    Ok((headers, sbom.into_string()))
 }

tools/harbor-sbom-browser/src/main.rs

@@ -15,6 +15,7 @@ async fn main() {
 
     let app = Router::new()
         .route("/", get(artifact_tree::render_as_html))
+        .route("/robots.txt", get(robots::robots_txt))
         .route("/sbom/{repository}/{digest}", get(sbom::download))
         .with_state(cached_rendered_artifact_tree);
     let listener = tokio::net::TcpListener::bind("0.0.0.0:9000").await.unwrap();

tools/harbor-sbom-browser/src/structs.rs

@@ -2,7 +2,7 @@ use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
 use serde::Deserialize;
-use serde_json::Value;
+use serde_json::value::RawValue;
 
 /// How long a cached object stays valid. This is kept short, because artifacts are deleted from the
 /// registry regularly (e.g. dev builds), and links to deleted artifacts do not work any more.
@@ -25,7 +25,9 @@ pub struct Artifact {
 
 #[derive(Deserialize, Debug)]
 pub struct InTotoAttestation {
-    pub predicate: Value,
+    /// Kept as raw JSON text. SBOMs can be tens of MB, and parsing them into a `serde_json::Value`
+    /// tree only to serialize them again costs several times their size in memory.
+    pub predicate: Box<RawValue>,
 }
 #[derive(Deserialize, Debug)]
 pub struct Dsse {

tools/harbor-sbom-browser/src/utils.rs

@@ -1,5 +1,7 @@
+use std::time::Duration;
+
 use crate::structs::{Dsse, InTotoAttestation};
-use axum::http::StatusCode;
+use axum::http::{header, StatusCode};
 use axum::response::{IntoResponse, Response};
 use base64::prelude::BASE64_STANDARD;
 use base64::Engine;
@@ -9,11 +11,25 @@ use snafu::ResultExt;
 use snafu::Snafu;
 use strum::{EnumDiscriminants, IntoStaticStr};
 use tokio::process::Command;
+use tokio::sync::Semaphore;
 use tracing::{error, warn};
 
+/// How many `cosign verify-attestation` processes may run at the same time. Each one needs up to
+/// about 170 MiB for the larger SBOMs, and crawlers request dozens of SBOMs at once, which
+/// otherwise gets the container OOMKilled.
+const MAX_CONCURRENT_COSIGN_RUNS: usize = 4;
+
+/// How long a request waits for a free cosign slot before it is rejected. Waiting requests cost
+/// next to no memory, but without a bound a crawler burst queues for minutes.
+const COSIGN_PERMIT_TIMEOUT: Duration = Duration::from_secs(10);
+
+/// Sent with the 503 response when all cosign slots are taken.
+const RETRY_AFTER_SECONDS: &str = "30";
+
 lazy_static! {
     static ref SHA256_REGEX: Regex = Regex::new(r"^[a-f0-9]{64}$").unwrap();
     static ref ALPHANUMERIC_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9\-]+$").unwrap();
+    static ref COSIGN_PERMITS: Semaphore = Semaphore::new(MAX_CONCURRENT_COSIGN_RUNS);
 }
 
 #[derive(Snafu, Debug, EnumDiscriminants)]
@@ -43,13 +59,16 @@ pub enum DownloadSbomError {
     ParseInTotoAttestation { source: serde_json::Error },
     #[snafu(display("failed to execute cosign"))]
     CosignExecution { source: std::io::Error },
+    #[snafu(display("too many SBOM downloads at the moment, try again later"))]
+    CosignBusy,
 }
 
 impl DownloadSbomError {
     fn status_code(&self) -> StatusCode {
         match self {
             Self::InvalidSbomParameters => StatusCode::BAD_REQUEST,
             Self::SbomNotFound { .. } => StatusCode::NOT_FOUND,
+            Self::CosignBusy => StatusCode::SERVICE_UNAVAILABLE,
             _ => StatusCode::INTERNAL_SERVER_ERROR,
         }
     }
@@ -66,6 +85,14 @@ impl IntoResponse for DownloadSbomError {
         } else {
             warn!("error: {:?}", self);
         }
+        if matches!(self, Self::CosignBusy) {
+            return (
+                status_code,
+                [(header::RETRY_AFTER, RETRY_AFTER_SECONDS)],
+                self.to_string(),
+            )
+                .into_response();
+        }
         (status_code, self.to_string()).into_response()
     }
 }
@@ -77,6 +104,11 @@ pub async fn verify_attestation(
     if !SHA256_REGEX.is_match(digest) || !ALPHANUMERIC_REGEX.is_match(repository) {
         return Err(DownloadSbomError::InvalidSbomParameters);
     }
+    // Held until the attestation is parsed, because parsing also holds copies of the whole SBOM.
+    let _permit = tokio::time::timeout(COSIGN_PERMIT_TIMEOUT, COSIGN_PERMITS.acquire())
+        .await
+        .map_err(|_| DownloadSbomError::CosignBusy)?
+        .expect("the cosign semaphore is never closed");
     let cmd_output = Command::new("cosign")
         .arg("verify-attestation")
         .arg("--type")
@@ -112,13 +144,18 @@ pub async fn verify_attestation(
         });
     }
 
-    let output = String::from_utf8_lossy(&cmd_output.stdout);
-    let dsse = serde_json::from_str::<Dsse>(&output).context(ParseDsseSnafu)?;
+    parse_attestation(cmd_output.stdout)
+}
+
+fn parse_attestation(cosign_stdout: Vec<u8>) -> Result<InTotoAttestation, DownloadSbomError> {
+    let dsse = serde_json::from_slice::<Dsse>(&cosign_stdout).context(ParseDsseSnafu)?;
+    // Drop the cosign stdout to free memory before decoding the payload, which is a copy of the stdout.
+    drop(cosign_stdout);
     let attestation_bytes = BASE64_STANDARD
         .decode(dsse.payload)
         .context(DecodeDssePayloadSnafu)?;
     let attestation_string =
         std::str::from_utf8(&attestation_bytes).context(ParseDssePayloadAsStringSnafu)?;
     serde_json::from_str::<InTotoAttestation>(attestation_string)
         .context(ParseInTotoAttestationSnafu)
-}
+}