From 730e9d11d8106e66ace0afaaf53e014c2d09a8e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 11:54:33 +0000 Subject: [PATCH 1/8] Fix apply of created files from diff caches A diff archive has no delta for a file the patch creates, yet the disk stager counted a cached diff archive as covering the whole patch. With only diffs on disk, `apply --offline` passed the gate, patched the modified files, then failed on the created file's missing blob and left the package half-patched; online `apply` never fetched that blob at all. Coverage is now per file: a diff covers only files with a before-hash, and created files need their blob. Online, a cached diff archive no longer suppresses the download, and the top-up fetches just the created files' blobs. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01WzdTEhubve9yWfqBE7vAsB --- Cargo.lock | 1 + crates/socket-patch-cli/Cargo.toml | 2 + .../src/commands/fetch_stage.rs | 130 +++++++++- .../tests/diff_created_file_e2e.rs | 239 ++++++++++++++++++ 4 files changed, 359 insertions(+), 13 deletions(-) create mode 100644 crates/socket-patch-cli/tests/diff_created_file_e2e.rs diff --git a/Cargo.lock b/Cargo.lock index 7b687168..a829bbe7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1769,6 +1769,7 @@ dependencies = [ "hex", "libc", "portable-pty", + "qbsdiff", "regex", "reqwest", "semver", diff --git a/crates/socket-patch-cli/Cargo.toml b/crates/socket-patch-cli/Cargo.toml index 4858d339..64152732 100644 --- a/crates/socket-patch-cli/Cargo.toml +++ b/crates/socket-patch-cli/Cargo.toml @@ -74,6 +74,8 @@ sha1 = { workspace = true } # scan_vendor_e2e builds pristine registry tarballs for the auto-fetch tests. tar = { workspace = true } flate2 = { workspace = true } +# diff_created_file_e2e builds a real diff archive. +qbsdiff = { workspace = true } # update_fixture builds the Windows-shaped release archive for self-update e2e. zip = { workspace = true } hex = { workspace = true } diff --git a/crates/socket-patch-cli/src/commands/fetch_stage.rs b/crates/socket-patch-cli/src/commands/fetch_stage.rs index e60bddb9..c3e11faa 100644 --- a/crates/socket-patch-cli/src/commands/fetch_stage.rs +++ b/crates/socket-patch-cli/src/commands/fetch_stage.rs @@ -201,10 +201,11 @@ fn format_blob_fallback(diff_failed: usize, blobs: usize) -> String { } /// The manifest PURLs with no usable local source. A patch is "locally -/// applicable" iff at least one of: -/// - every `after_hash` blob it references is on disk, OR -/// - its diff archive is on disk, OR -/// - its package archive is on disk. +/// applicable" iff its package archive is on disk, or every file it +/// touches has its `after_hash` blob on disk or is covered by the patch's +/// diff archive. A diff covers only files that exist before the patch: a +/// created file (empty `before_hash`) has nothing to diff against, so it +/// always needs its blob. /// /// The patch pipeline picks whichever is present per file. Shared by the /// offline gate (probed against `.socket/`) and the post-download gate @@ -219,13 +220,13 @@ fn patches_without_source<'m>( .patches .iter() .filter_map(|(purl, record)| { - let all_blobs_present = record - .files - .values() - .all(|f| !missing_blobs.contains(&f.after_hash)); let diff_present = !missing_diff_archives.contains(&record.uuid); let pkg_present = !missing_package_archives.contains(&record.uuid); - if all_blobs_present || diff_present || pkg_present { + let files_covered = record.files.values().all(|f| { + !missing_blobs.contains(&f.after_hash) + || (diff_present && !f.before_hash.is_empty()) + }); + if pkg_present || files_covered { None } else { Some(purl.as_str()) @@ -234,6 +235,33 @@ fn patches_without_source<'m>( .collect() } +/// `manifest` cut down to the files a diff archive cannot patch (created +/// files, whose `before_hash` is empty): the blobs a diff-mode fetch still +/// needs even when every diff archive is present. +pub(crate) fn files_diffs_cannot_cover(manifest: &PatchManifest) -> PatchManifest { + let patches = manifest + .patches + .iter() + .filter_map(|(purl, record)| { + let files: HashMap<_, _> = record + .files + .iter() + .filter(|(_, f)| f.before_hash.is_empty()) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + (!files.is_empty()).then(|| { + let mut record = record.clone(); + record.files = files; + (purl.clone(), record) + }) + }) + .collect(); + PatchManifest { + patches, + setup: manifest.setup.clone(), + } +} + /// Mirror `src`'s files into `dst` by hardlink (copy fallback). Pre-seeds the /// overlay tempdir with everything already cached so only the gap downloads. async fn overlay_dir(src: &Path, dst: &Path) { @@ -311,12 +339,14 @@ pub(crate) async fn stage_patch_sources( // locally. We honor `--download-mode` for the primary fetch when there's // actually a gap to close. Skip the archive fetch entirely when all file // blobs are already present locally — the pipeline will succeed via the - // blob path, so an archive fetch would be wasted round-trips. + // blob path, so an archive fetch would be wasted round-trips. Cached + // diff archives can still leave a patch uncovered (a created file), and + // the blob top-up below closes that gap. let download_needed = !common.offline && match download_mode { DownloadMode::File => !missing_blobs.is_empty(), DownloadMode::Diff if missing_blobs.is_empty() => false, - DownloadMode::Diff => !missing_diff_archives.is_empty(), + DownloadMode::Diff => !missing_diff_archives.is_empty() || !no_source_purls.is_empty(), }; if !download_needed { @@ -374,15 +404,25 @@ pub(crate) async fn stage_patch_sources( // For non-file modes, automatically fetch any still-missing file blobs as // a fallback. Patches that lack the requested mode on the server will // still apply via the legacy blob path. + // + // With every diff archive already cached, only the files no diff can + // patch are fetched: that is the gap that triggered this download. let mut blob_fetch_failed = false; if download_mode != DownloadMode::File { - let still_missing_blobs = get_missing_blobs(manifest, &staged.blobs).await; + let created_only; + let blob_scope = if missing_diff_archives.is_empty() { + created_only = files_diffs_cannot_cover(manifest); + &created_only + } else { + manifest + }; + let still_missing_blobs = get_missing_blobs(blob_scope, &staged.blobs).await; if !still_missing_blobs.is_empty() { status.set(format_blob_fallback( fetch_result.failed, still_missing_blobs.len(), )); - let blob_result = fetch_missing_blobs(manifest, &staged.blobs, client, None).await; + let blob_result = fetch_missing_blobs(blob_scope, &staged.blobs, client, None).await; status.finish(); if !quiet { for line in format_fetch_summary(&blob_result, BLOB, true) { @@ -972,6 +1012,70 @@ mod tests { ); } + /// A diff archive cannot patch a file the patch creates (nothing to diff + /// against), so it covers such a patch only together with the created + /// file's blob: without it, offline staging is Unavailable up front + /// instead of passing the gate and failing mid-apply. + #[tokio::test] + async fn stage_offline_diff_archive_does_not_cover_a_created_file() { + let tmp = tempfile::tempdir().unwrap(); + let socket_dir = tmp.path().join(".socket"); + std::fs::create_dir_all(socket_dir.join("diffs")).unwrap(); + std::fs::write( + socket_dir.join("diffs").join(format!("{UUID}.tar.gz")), + b"x", + ) + .unwrap(); + let created = "c".repeat(64); + let mut manifest = manifest_with_one_patch(); + manifest + .patches + .get_mut("pkg:npm/left-pad@1.3.0") + .unwrap() + .files + .insert( + "new.js".to_string(), + PatchFileInfo { + before_hash: String::new(), + after_hash: created.clone(), + }, + ); + + let outcome = + stage_patch_sources(&offline_args(), &manifest, &socket_dir, &offline_client()) + .await + .expect("no hard failure"); + assert!(matches!(outcome, StageOutcome::Unavailable)); + + std::fs::create_dir_all(socket_dir.join("blobs")).unwrap(); + std::fs::write(socket_dir.join("blobs").join(&created), b"new").unwrap(); + let outcome = + stage_patch_sources(&offline_args(), &manifest, &socket_dir, &offline_client()) + .await + .expect("no hard failure"); + assert!( + matches!(outcome, StageOutcome::Ready(_)), + "diff for the modified file + blob for the created one covers the patch" + ); + } + + #[test] + fn files_diffs_cannot_cover_keeps_only_created_files() { + let mut manifest = manifest_with_one_patch(); + assert!(files_diffs_cannot_cover(&manifest).patches.is_empty()); + let record = manifest.patches.get_mut("pkg:npm/left-pad@1.3.0").unwrap(); + record.files.insert( + "new.js".to_string(), + PatchFileInfo { + before_hash: String::new(), + after_hash: "c".repeat(64), + }, + ); + let cut = files_diffs_cannot_cover(&manifest); + let files: Vec<&String> = cut.patches["pkg:npm/left-pad@1.3.0"].files.keys().collect(); + assert_eq!(files, ["new.js"]); + } + /// The vendor (in-memory) stager documents the opposite policy: a diff /// archive is NOT sufficient (auto-force can need the full after-blob), /// so the same fixture that satisfies the disk stager is Unavailable diff --git a/crates/socket-patch-cli/tests/diff_created_file_e2e.rs b/crates/socket-patch-cli/tests/diff_created_file_e2e.rs new file mode 100644 index 00000000..c2e1f104 --- /dev/null +++ b/crates/socket-patch-cli/tests/diff_created_file_e2e.rs @@ -0,0 +1,239 @@ +//! A diff archive only carries deltas for files that exist before the +//! patch: a file the patch CREATES (empty `beforeHash`) has nothing to +//! diff against, so the patch service leaves it out and the pipeline can +//! only apply it from its after-blob (or a package archive). These tests +//! pin that the disk stager and `repair` both treat a diff archive as +//! covering only the files it can actually patch: +//! +//! - `apply --offline` with just the diff archive on disk fails closed +//! with the "no local source" report and leaves every file untouched, +//! instead of passing the gate and failing mid-apply; +//! - online `apply` with the diff archive cached still fetches the +//! created file's blob; +//! - a default (diff-mode) `repair` also downloads the created file's +//! blob, so a later `apply --offline` succeeds. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +use flate2::write::GzEncoder; +use flate2::Compression; +use qbsdiff::Bsdiff; +use sha2::{Digest, Sha256}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const ORG_SLUG: &str = "test-org"; +const UUID: &str = "67676767-6767-4767-8767-676767676767"; +const PURL: &str = "pkg:npm/created-file-test@1.0.0"; +const BEFORE: &[u8] = b"module.exports = 'before, and long enough to diff';\n"; +const AFTER: &[u8] = b"module.exports = 'after!, and long enough to diff';\n"; +const CREATED: &[u8] = b"module.exports = 'a brand new file';\n"; + +fn binary() -> PathBuf { + env!("CARGO_BIN_EXE_socket-patch").into() +} + +fn git_sha256(content: &[u8]) -> String { + let header = format!("blob {}\0", content.len()); + let mut hasher = Sha256::new(); + hasher.update(header.as_bytes()); + hasher.update(content); + hex::encode(hasher.finalize()) +} + +fn run_cli(root: &Path, argv: &[&str], mock_uri: Option<&str>) -> (i32, String, String) { + let mut cmd = Command::new(binary()); + cmd.args(argv).current_dir(root); + for (key, _) in std::env::vars_os() { + if key.to_string_lossy().starts_with("SOCKET_") + && key.to_string_lossy() != "SOCKET_NO_CONFIG" + { + cmd.env_remove(&key); + } + } + cmd.env("SOCKET_TELEMETRY_DISABLED", "1"); + // Offline runs get a dead endpoint: any request they make fails. + let api_url = mock_uri.unwrap_or("http://127.0.0.1:1"); + cmd.env("SOCKET_API_URL", api_url) + .env("SOCKET_API_TOKEN", "fake-token-for-test") + .env("SOCKET_ORG_SLUG", ORG_SLUG); + let out = cmd.output().expect("run socket-patch"); + ( + out.status.code().unwrap_or(-1), + String::from_utf8_lossy(&out.stdout).into_owned(), + String::from_utf8_lossy(&out.stderr).into_owned(), + ) +} + +/// The diff archive the service serves for this patch: a bsdiff delta for +/// the modified file and nothing for the created one. +fn diff_archive() -> Vec { + let mut delta = Vec::new(); + Bsdiff::new(BEFORE, AFTER) + .compare(std::io::Cursor::new(&mut delta)) + .unwrap(); + let mut builder = tar::Builder::new(GzEncoder::new(Vec::new(), Compression::default())); + let mut header = tar::Header::new_gnu(); + header.set_size(delta.len() as u64); + header.set_mode(0o644); + header.set_cksum(); + builder + .append_data(&mut header, "index.js", delta.as_slice()) + .unwrap(); + builder.into_inner().unwrap().finish().unwrap() +} + +/// An npm project with the unpatched package installed and a manifest +/// whose patch modifies `index.js` and creates `new.js` (sorted after +/// `index.js`, so a mid-apply failure would leave `index.js` patched). Returns the +/// installed package dir. +fn seed_project(root: &Path) -> PathBuf { + std::fs::write( + root.join("package.json"), + r#"{"name":"created-file-root","version":"0.0.0"}"#, + ) + .unwrap(); + let pkg = root.join("node_modules").join("created-file-test"); + std::fs::create_dir_all(&pkg).unwrap(); + std::fs::write( + pkg.join("package.json"), + r#"{"name":"created-file-test","version":"1.0.0"}"#, + ) + .unwrap(); + std::fs::write(pkg.join("index.js"), BEFORE).unwrap(); + + let socket = root.join(".socket"); + std::fs::create_dir_all(&socket).unwrap(); + std::fs::write( + socket.join("manifest.json"), + serde_json::to_vec_pretty(&serde_json::json!({ + "patches": { + PURL: { + "uuid": UUID, + "exportedAt": "2026-01-01T00:00:00Z", + "files": { + "package/index.js": { + "beforeHash": git_sha256(BEFORE), + "afterHash": git_sha256(AFTER), + }, + "package/new.js": { + "beforeHash": "", + "afterHash": git_sha256(CREATED), + } + }, + "vulnerabilities": {}, + "description": "creates a file", + "license": "MIT", + "tier": "free", + } + } + })) + .unwrap(), + ) + .unwrap(); + pkg +} + +fn seed_cached_diff_archive(root: &Path) { + let diffs = root.join(".socket").join("diffs"); + std::fs::create_dir_all(&diffs).unwrap(); + std::fs::write(diffs.join(format!("{UUID}.tar.gz")), diff_archive()).unwrap(); +} + +async fn mount_blob(mock: &MockServer, content: &'static [u8]) { + Mock::given(method("GET")) + .and(path(format!( + "/v0/orgs/{ORG_SLUG}/patches/blob/{}", + git_sha256(content) + ))) + .respond_with(ResponseTemplate::new(200).set_body_bytes(content.to_vec())) + .mount(mock) + .await; +} + +fn assert_fully_patched(pkg: &Path) { + assert_eq!(std::fs::read(pkg.join("index.js")).unwrap(), AFTER); + assert_eq!(std::fs::read(pkg.join("new.js")).unwrap(), CREATED); +} + +#[test] +fn offline_apply_with_only_a_diff_archive_reports_the_created_file_gap() { + let tmp = tempfile::tempdir().unwrap(); + let pkg = seed_project(tmp.path()); + seed_cached_diff_archive(tmp.path()); + + let (code, stdout, stderr) = run_cli(tmp.path(), &["apply", "--offline"], None); + + assert_eq!(code, 1, "stdout={stdout}\nstderr={stderr}"); + assert_eq!( + std::fs::read(pkg.join("index.js")).unwrap(), + BEFORE, + "a patch that cannot fully apply must not be applied partway" + ); + assert!(!pkg.join("new.js").exists()); + assert!( + stderr.contains("1 patch has no local source and --offline is set:") + && stderr.contains(PURL) + && stderr.contains("socket-patch repair"), + "the offline gate names the patch and the remedy; stderr={stderr}" + ); +} + +#[tokio::test] +async fn online_apply_with_a_cached_diff_archive_fetches_the_created_files_blob() { + let mock = MockServer::start().await; + mount_blob(&mock, CREATED).await; + + let tmp = tempfile::tempdir().unwrap(); + let pkg = seed_project(tmp.path()); + seed_cached_diff_archive(tmp.path()); + + let (code, stdout, stderr) = run_cli(tmp.path(), &["apply"], Some(&mock.uri())); + + assert_eq!(code, 0, "stdout={stdout}\nstderr={stderr}"); + assert_fully_patched(&pkg); + let requested: Vec = mock + .received_requests() + .await + .unwrap_or_default() + .iter() + .map(|r| r.url.path().to_string()) + .collect(); + assert!( + !requested + .iter() + .any(|p| p.contains("/patches/blob/") && p.ends_with(&git_sha256(AFTER))), + "the modified file's blob is not needed: its delta applies; requested={requested:?}" + ); +} + +#[tokio::test] +async fn default_repair_downloads_the_created_files_blob_for_offline_apply() { + let mock = MockServer::start().await; + Mock::given(method("GET")) + .and(path(format!("/v0/orgs/{ORG_SLUG}/patches/diff/{UUID}"))) + .respond_with(ResponseTemplate::new(200).set_body_bytes(diff_archive())) + .mount(&mock) + .await; + mount_blob(&mock, CREATED).await; + + let tmp = tempfile::tempdir().unwrap(); + let pkg = seed_project(tmp.path()); + + let (code, stdout, stderr) = run_cli(tmp.path(), &["repair"], Some(&mock.uri())); + assert_eq!(code, 0, "repair: stdout={stdout}\nstderr={stderr}"); + let blobs = tmp.path().join(".socket").join("blobs"); + assert!( + blobs.join(git_sha256(CREATED)).exists(), + "repair must cache the created file's blob; stdout={stdout}\nstderr={stderr}" + ); + assert!( + !blobs.join(git_sha256(AFTER)).exists(), + "the modified file's delta covers it; its blob is not downloaded" + ); + + let (code, stdout, stderr) = run_cli(tmp.path(), &["apply", "--offline"], None); + assert_eq!(code, 0, "apply: stdout={stdout}\nstderr={stderr}"); + assert_fully_patched(&pkg); +} From 4b16a4e1c187202055058814d1bcd3605758908b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 11:54:43 +0000 Subject: [PATCH 2/8] Batch vendor package-reference requests A vendor run asked the patch service for each package's download reference in its own request, though the endpoint takes 500 uuids at once: N round trips and N quota units for N packages. The run's download plan now resolves every planned uuid in one request, sent by the first planned call in place of its own and with the same retries, so an outage costs what it did before. Each package takes its answer from that batch at its turn; one still building is asked again then, as before. Hosted scan's reference lookup is chunked at the endpoint's 500-uuid cap, which it used to exceed with a 400. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01WzdTEhubve9yWfqBE7vAsB --- crates/socket-patch-core/src/api/client.rs | 126 +++++++-- .../src/api/vendor_prefetch.rs | 259 +++++++++++++++++- 2 files changed, 360 insertions(+), 25 deletions(-) diff --git a/crates/socket-patch-core/src/api/client.rs b/crates/socket-patch-core/src/api/client.rs index 5f716c02..7d2603cf 100644 --- a/crates/socket-patch-core/src/api/client.rs +++ b/crates/socket-patch-core/src/api/client.rs @@ -332,6 +332,9 @@ fn retry_after_secs(headers: &HeaderMap) -> Option { /// without a hint). type VendorAttemptError = (ApiError, Option>); +/// Most UUIDs the package-reference endpoint takes in one request. +pub(crate) const MAX_REFERENCE_BATCH: usize = 500; + /// Body payload for the batch search POST endpoint. #[derive(Serialize)] struct BatchSearchBody { @@ -740,8 +743,9 @@ impl ApiClient { /// Resolve hosted-patch references for a set of published-patch UUIDs /// (hosted-mode `scan`, the default for a bare `scan`). Uses the authenticated /// `POST /v0/orgs/{org}/patches/package` when a token+org are set, else the - /// public proxy `POST /patch/package` (free patches only). Returns a - /// UUID → reference map (missing/404 → empty). + /// public proxy `POST /patch/package` (free patches only), in requests + /// of at most [`MAX_REFERENCE_BATCH`] UUIDs (the endpoint rejects more). + /// Returns a UUID → reference map (missing/404 → empty). /// /// Uses the client's configured org slug; see /// [`Self::fetch_registry_references_for_org`] for a per-call override. @@ -765,14 +769,18 @@ impl ApiClient { return Ok(std::collections::HashMap::new()); } let path = self.patches_path(org_slug, "package"); - let body = PackageVendorRequest { - uuids: uuids.to_vec(), - free_only: None, - }; - let resp = self - .post_json::(&path, &body) - .await?; - Ok(resp.map(|r| r.results).unwrap_or_default()) + let mut results = std::collections::HashMap::new(); + for chunk in uuids.chunks(MAX_REFERENCE_BATCH) { + let body = PackageVendorRequest { + uuids: chunk.to_vec(), + free_only: None, + }; + let resp = self + .post_json::(&path, &body) + .await?; + results.extend(resp.map(|r| r.results).unwrap_or_default()); + } + Ok(results) } /// Internal: POST the batch search to the public proxy's @@ -1398,21 +1406,54 @@ impl ApiClient { /// Step 1 of [`Self::fetch_vendor_package`], retried per the client's /// [`VendorRetryPolicy`]. `Err` carries whether the final failure was a - /// retryable (availability) one. + /// retryable (availability) one. A uuid the attached plan names is + /// answered from the plan's one reference batch (see + /// [`VendorPrefetch::reference`]). async fn request_vendor_package( &self, uuid: &str, free_only: bool, vendor_url: Option<&str>, ) -> Result { + let plan = self + .vendor_prefetch + .lock() + .ok() + .and_then(|slot| slot.clone()); + if let Some(plan) = plan { + if let Some(result) = plan.reference(self, uuid, free_only, vendor_url).await { + return result; + } + } + let mut results = self + .request_vendor_references(&[uuid.to_string()], free_only, vendor_url) + .await?; + results.remove(uuid).ok_or_else(|| { + ( + ApiError::Other(format!("package response missing a result for {uuid}")), + false, + ) + }) + } + + /// One package-reference request for `uuids` (at most + /// [`MAX_REFERENCE_BATCH`]), retried per the client's + /// [`VendorRetryPolicy`]: the per-uuid results. `Err` carries whether + /// the final failure was a retryable (availability) one. + pub(crate) async fn request_vendor_references( + &self, + uuids: &[String], + free_only: bool, + vendor_url: Option<&str>, + ) -> Result, (ApiError, bool)> { let attempts = self.vendor_retry.attempts.max(1); let mut attempt = 1; loop { match self - .request_vendor_package_once(uuid, free_only, vendor_url) + .request_vendor_references_once(uuids, free_only, vendor_url) .await { - Ok(result) => return Ok(result), + Ok(results) => return Ok(results), Err((e, Some(retry_after))) if attempt < attempts => { debug_log(&format!( "vendor package request attempt {attempt} failed: {e}" @@ -1425,15 +1466,15 @@ impl ApiClient { } } - /// One package-reference POST: the single requested UUID's result. - async fn request_vendor_package_once( + /// One package-reference POST: the requested UUIDs' results. + async fn request_vendor_references_once( &self, - uuid: &str, + uuids: &[String], free_only: bool, vendor_url: Option<&str>, - ) -> Result { + ) -> Result, VendorAttemptError> { let body = PackageVendorRequest { - uuids: vec![uuid.to_string()], + uuids: uuids.to_vec(), // Only send freeOnly when forcing it (the public-proxy contract); // the authenticated endpoint defaults to false. free_only: free_only.then_some(true), @@ -1481,12 +1522,7 @@ impl ApiClient { hint, ) })?; - return parsed.results.get(uuid).cloned().ok_or_else(|| { - ( - ApiError::Other(format!("package response missing a result for {uuid}")), - None, - ) - }); + return Ok(parsed.results); } // 429 classifies as RateLimited but is still retried (the hint); // 401/403 carry no hint. @@ -4583,6 +4619,48 @@ mod vendor_package_tests { assert_eq!(map[UUID].status, "granted"); } + /// The endpoint rejects more than [`MAX_REFERENCE_BATCH`] uuids per + /// request (400), so a larger scan goes out in capped chunks whose + /// results are merged. + #[tokio::test] + async fn fetch_registry_references_chunks_at_the_endpoint_cap() { + struct EchoGranted; + impl wiremock::Respond for EchoGranted { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let uuids = body["uuids"].as_array().unwrap(); + if uuids.len() > MAX_REFERENCE_BATCH { + return ResponseTemplate::new(400); + } + let results: serde_json::Map = uuids + .iter() + .map(|u| { + ( + u.as_str().unwrap().to_string(), + json!({ "status": "granted", "url": null, "artifacts": [] }), + ) + }) + .collect(); + ResponseTemplate::new(200).set_body_json(json!({ "results": results })) + } + } + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/patch/package")) + .respond_with(EchoGranted) + .expect(2) + .mount(&server) + .await; + let uuids: Vec = (0..=MAX_REFERENCE_BATCH) + .map(|i| format!("{i:08x}-0000-4000-8000-{i:012x}")) + .collect(); + let map = proxy_client(server.uri()) + .fetch_registry_references(&uuids) + .await + .expect("chunked resolution succeeds"); + assert_eq!(map.len(), uuids.len()); + } + /// The package-reference route honors a per-call org override: /// `Some(slug)` beats the client's configured `acme`, and the one-arg /// wrapper keeps using `acme`. diff --git a/crates/socket-patch-core/src/api/vendor_prefetch.rs b/crates/socket-patch-core/src/api/vendor_prefetch.rs index ff57163c..057336dc 100644 --- a/crates/socket-patch-core/src/api/vendor_prefetch.rs +++ b/crates/socket-patch-core/src/api/vendor_prefetch.rs @@ -86,6 +86,15 @@ //! download, removed as soon as their package is passed over or the loop //! ends) and by the [`crate::vendor::prestage`] pool, not by size. //! +//! The package-reference half is batched: the plan's first call sends +//! one request naming every planned uuid (see +//! [`VendorPrefetch::reference`]) in place of its own, and later calls +//! take their granted reference from it. That request grants the whole +//! plan up front, which the plan's exactness makes safe. The bounds +//! above then limit the archive downloads. A package the batch reports +//! still building, or leaves out, makes its own request at its turn, as +//! before. +//! //! A planned download may name a secondary artifact (the gem stub //! gemspec) its backend fetches right after a verified archive; the task //! fetches it along with the archive, under the backend's own conditions, @@ -107,8 +116,10 @@ use futures_util::StreamExt; use super::client::{ hold_back_debug, ApiClient, HeldBack, PlannedDownload, PrefetchedSecondary, - VendorServiceOutcome, VENDOR_BREAKER_THRESHOLD, + VendorServiceOutcome, MAX_REFERENCE_BATCH, VENDOR_BREAKER_THRESHOLD, }; +use super::types::PackageVendorResult; +use crate::api::client::ApiError; use crate::vendor::lock_inventory::LockIntegrity; use crate::vendor::prestage::PrestageRecipe; use crate::vendor::registry_fetch::{artifact_matches_integrity, verify_go_h1}; @@ -153,6 +164,9 @@ pub(crate) struct VendorPrefetch { /// What the task is allowed to request, shared with it. look: Arc, state: tokio::sync::Mutex, + /// The plan's package references, resolved in one batch by the first + /// planned call (see [`Self::reference`]) and taken uuid by uuid. + references: tokio::sync::OnceCell>>, } /// The window of plan positions the task may request: `[at, at + reach)`, @@ -354,7 +368,84 @@ impl VendorPrefetch { window: window.max(1), look: Arc::new(Lookahead::new(byte_budget)), state: tokio::sync::Mutex::new(PrefetchState::default()), + references: tokio::sync::OnceCell::new(), + } + } + + /// Step 1 (the package-reference request) for a planned `uuid`, or + /// `None` to make the live request (not planned, other parameters, or + /// the batch did not answer it). + /// + /// The first planned call sends ONE request naming the whole plan, its + /// own uuid first, in place of its single-uuid request and with the + /// same retry ladder: the endpoint takes up to [`MAX_REFERENCE_BATCH`] + /// uuids, and one request per package paid a round trip and a quota + /// unit each. Its failure is that call's own failure, so an outage + /// costs what the serial loop paid, and every later call makes its + /// live request. A package still building (`pending_build`) is asked + /// again at its own turn, as the serial loop did, since it may be + /// ready by then; every other answer is final and is reused. + pub(crate) async fn reference( + &self, + client: &ApiClient, + uuid: &str, + free_only: bool, + vendor_url: Option<&str>, + ) -> Option> { + if free_only != self.free_only + || vendor_url != self.vendor_url.as_deref() + || !self.planned.iter().any(|planned| planned == uuid) + { + return None; + } + let mut own = None; + let cache = self + .references + .get_or_init(|| async { + let mut order = vec![uuid.to_string()]; + for planned in &self.planned { + if !order.contains(planned) { + order.push(planned.clone()); + } + } + let mut kept = HashMap::new(); + for (i, chunk) in order.chunks(MAX_REFERENCE_BATCH).enumerate() { + match client + .request_vendor_references(chunk, free_only, vendor_url) + .await + { + Ok(mut results) => { + if i == 0 { + own = Some(results.remove(uuid).ok_or_else(|| { + ( + ApiError::Other(format!( + "package response missing a result for {uuid}" + )), + false, + ) + })); + } + kept.extend( + results + .into_iter() + .filter(|(_, r)| r.status != "pending_build"), + ); + } + Err(e) if i == 0 => { + own = Some(Err(e)); + break; + } + // The chunk's uuids make their live requests. + Err(_) => break, + } + } + std::sync::Mutex::new(kept) + }) + .await; + if own.is_some() { + return own; } + cache.lock().ok()?.remove(uuid).map(Ok) } /// The prefetched outcome of the loop's call for `uuid`, or `None` to @@ -1353,4 +1444,170 @@ mod tests { planned_requests.sort(); assert_eq!(planned_requests, serial_requests); } + + /// A service that answers every uuid a package-reference request names, + /// as the real endpoint does: the request's first uuid decides a + /// whole-request failure (503 / 403), and the others it cannot grant + /// are left out of the results. + struct BatchService { + base: String, + scripts: HashMap, + } + + impl wiremock::Respond for BatchService { + fn respond(&self, request: &wiremock::Request) -> ResponseTemplate { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let uuids: Vec = body["uuids"] + .as_array() + .unwrap() + .iter() + .map(|u| u.as_str().unwrap().to_string()) + .collect(); + match self.scripts.get(&uuids[0]) { + Some(Script::Down) => return ResponseTemplate::new(503), + Some(Script::Forbidden) => return ResponseTemplate::new(403), + _ => {} + } + let mut results = serde_json::Map::new(); + for u in uuids { + let status = match self.scripts.get(&u) { + Some(Script::Granted(_)) => "granted", + Some(Script::Pending) => "pending_build", + Some(Script::NotFound) => "not_found", + _ => continue, + }; + let url = format!("{}/serve/{u}.tgz", self.base); + let sri = format!( + "sha512-{}", + base64::engine::general_purpose::STANDARD.encode(Sha512::digest(u.as_bytes())) + ); + let artifacts = if status == "granted" { + serde_json::json!([{ "kind": "tarball", "url": url, + "integrity": { "sha512": sri } }]) + } else { + serde_json::json!([]) + }; + results.insert( + u.clone(), + serde_json::json!({ "status": status, "url": url, "artifacts": artifacts }), + ); + } + ResponseTemplate::new(200).set_body_json(serde_json::json!({ "results": results })) + } + } + + async fn serve_batches(scripts: &[Script]) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(POST_PATH)) + .respond_with(BatchService { + base: server.uri(), + scripts: scripts + .iter() + .enumerate() + .map(|(i, s)| (uuid(i), *s)) + .collect(), + }) + .mount(&server) + .await; + for i in 0..scripts.len() { + let u = uuid(i); + Mock::given(method("GET")) + .and(path(format!("/serve/{u}.tgz"))) + .respond_with(ResponseTemplate::new(200).set_body_bytes(u.into_bytes())) + .mount(&server) + .await; + } + server + } + + /// The `(POSTs, GETs)` the server has seen. + async fn request_counts(server: &MockServer) -> (usize, usize) { + let log = request_log(server).await; + let posts = log.iter().filter(|r| r.starts_with("POST")).count(); + (posts, log.len() - posts) + } + + /// One package-reference request resolves the whole plan: the serial + /// loop's outcomes, one POST instead of one per package, and the same + /// downloads. + #[tokio::test] + async fn one_reference_request_serves_the_whole_plan() { + let scripts: Vec