1#![deny(missing_docs)]
30
31use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
32use eyre::{Result, bail, eyre};
33use futures_util::TryStreamExt as _;
34use log::warn;
35use reqwest::StatusCode;
36use reqwest::header::{
37 ACCEPT, AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, ETAG, HeaderMap,
38 HeaderValue, IF_MATCH, IF_NONE_MATCH,
39};
40use serde::{Deserialize, Serialize};
41use sha2::Digest as _;
42use std::collections::BTreeSet;
43use std::fs;
44use std::path::{Path, PathBuf};
45use std::sync::Arc;
46use std::sync::atomic::{AtomicBool, Ordering};
47use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
48use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
49use url::{Host, Url};
50
51mod agent;
52mod local;
53
54pub use agent::{
55 AGENT_PROTOCOL_VERSION, AgentRemoteCache, AgentRequest, AgentResponse, AgentStats, CacheAgent,
56 CompilerStats, RestoreStats, is_task_identity, task_manifest_actions,
57};
58pub use local::{LocalActionCache, LocalCas};
59pub use mbx_cache_protocol::{
60 ACTION_RESULT_MEDIA_TYPE, ActionPrediction, ActionResult as RemoteActionResult,
61 BLOB_MEDIA_TYPE, BLOB_PACK_BLOBS_HEADER, BLOB_PACK_BYTES_HEADER, BLOB_PACK_HEADER_BYTES,
62 BLOB_PACK_MAGIC, BLOB_PACK_MEDIA_TYPE, CLIENT_METADATA_MEDIA_TYPE, Capabilities,
63 CapabilityFeatures, CapabilityLimits, CapabilityProtocol, DIGEST_LIST_MEDIA_TYPE,
64 DIRECTORY_MEDIA_TYPE, Digest as CacheDigest, DigestAlgorithm, Directory as CacheDirectory,
65 DirectoryNode as CacheDirectoryNode, FileNode as CacheFileNode, NAMESPACE_HEADER,
66 PROTOCOL_HEADER, PROTOCOL_VERSION, RustcMetadata, SymlinkNode as CacheSymlinkNode,
67 TASK_ACTION_MANIFEST_MEDIA_TYPE, TaskActionManifest,
68};
69const MAX_REMOTE_JSON_BYTES: u64 = 16 * 1024 * 1024;
76const MAX_ETAG_BYTES: usize = 256;
81const MAX_REMOTE_BLOB_BYTES: u64 = 5 * 1024 * 1024 * 1024;
84const MAX_STAGED_BLOB_PACK_BYTES: u64 = 256 * 1024 * 1024;
85const MAX_STAGED_BLOB_PACK_ITEMS: usize = 2 * 1024;
86const BLOB_PACK_TIMEOUT_BYTES_PER_UNIT: u64 = MAX_STAGED_BLOB_PACK_BYTES / 4;
87const BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT: usize = MAX_STAGED_BLOB_PACK_ITEMS / 4;
88
89pub fn canonical_json(value: &impl Serialize) -> Result<Vec<u8>> {
94 Ok(mbx_cache_protocol::canonical_json(value)?)
95}
96
97#[derive(
98 Debug,
99 Clone,
100 Copy,
101 Serialize,
102 Deserialize,
103 Default,
104 strum::EnumString,
105 strum::Display,
106 PartialEq,
107 Eq,
108)]
109#[serde(rename_all = "kebab-case")]
110#[strum(serialize_all = "kebab-case")]
111pub enum RemoteCacheMode {
113 #[default]
115 ReadWrite,
116 ReadOnly,
118 WriteOnly,
120}
121
122impl RemoteCacheMode {
123 pub fn reads(self) -> bool {
125 matches!(self, Self::ReadWrite | Self::ReadOnly)
126 }
127
128 pub fn writes(self) -> bool {
130 matches!(self, Self::ReadWrite | Self::WriteOnly)
131 }
132}
133
134pub struct RemoteCacheConfig {
136 pub base_url: Url,
138 pub namespace: String,
140 pub token: Option<String>,
142 pub token_file: Option<PathBuf>,
144 pub oidc_audience: Option<String>,
146 pub connect_timeout: Duration,
148 pub read_timeout: Duration,
150 pub download_timeout: Duration,
152 pub retries: i64,
154}
155
156pub enum BlobSource {
158 Bytes(Vec<u8>),
160 File(tempfile::NamedTempFile),
162 Path(PathBuf),
164}
165
166pub struct BlobUpload {
168 pub digest: CacheDigest,
170 pub source: BlobSource,
172}
173
174pub struct RemoteActionManifest {
176 pub bytes: Vec<u8>,
178 pub etag: String,
180}
181
182pub struct RemoteBlobPack {
184 _directory: tempfile::TempDir,
185 pub blobs: Vec<(CacheDigest, PathBuf)>,
187 pub requests: u64,
189 pub requested: Vec<CacheDigest>,
191 pub blob_count: u64,
193 pub payload_bytes: u64,
195 pub framed_bytes: u64,
197}
198
199struct DownloadedBlobPack {
200 directory: tempfile::TempDir,
201 blobs: Vec<(CacheDigest, PathBuf)>,
202 metadata: BlobPackResponseStats,
203}
204
205#[derive(Debug, Clone, Copy, Default)]
206struct BlobPackResponseMetadata {
207 content_length: Option<u64>,
208 blob_count: Option<u64>,
209 payload_bytes: Option<u64>,
210}
211
212#[derive(Debug, Clone, Copy)]
213struct BlobPackResponseStats {
214 blob_count: u64,
215 payload_bytes: u64,
216 framed_bytes: u64,
217}
218
219impl BlobPackResponseMetadata {
220 fn from_headers(headers: &HeaderMap) -> Result<Self> {
221 Ok(Self {
222 content_length: optional_u64_header(headers, CONTENT_LENGTH.as_str())?,
223 blob_count: optional_u64_header(headers, BLOB_PACK_BLOBS_HEADER)?,
224 payload_bytes: optional_u64_header(headers, BLOB_PACK_BYTES_HEADER)?,
225 })
226 }
227
228 fn validate(self, decoded: BlobPackResponseStats) -> Result<BlobPackResponseStats> {
229 if let Some(content_length) = self.content_length
230 && content_length != decoded.framed_bytes
231 {
232 bail!(
233 "remote cache blob pack content length metadata mismatch: expected {}, decoded {}",
234 content_length,
235 decoded.framed_bytes
236 );
237 }
238 if let Some(blob_count) = self.blob_count
239 && blob_count != decoded.blob_count
240 {
241 bail!(
242 "remote cache blob pack blob count metadata mismatch: expected {}, decoded {}",
243 blob_count,
244 decoded.blob_count
245 );
246 }
247 if let Some(payload_bytes) = self.payload_bytes
248 && payload_bytes != decoded.payload_bytes
249 {
250 bail!(
251 "remote cache blob pack payload byte metadata mismatch: expected {}, decoded {}",
252 payload_bytes,
253 decoded.payload_bytes
254 );
255 }
256 Ok(BlobPackResponseStats {
257 blob_count: self.blob_count.unwrap_or(decoded.blob_count),
258 payload_bytes: self.payload_bytes.unwrap_or(decoded.payload_bytes),
259 framed_bytes: self.content_length.unwrap_or(decoded.framed_bytes),
260 })
261 }
262}
263
264fn optional_u64_header(headers: &HeaderMap, name: &str) -> Result<Option<u64>> {
265 let Some(value) = headers.get(name) else {
266 return Ok(None);
267 };
268 let value = value
269 .to_str()
270 .map_err(|_| eyre!("remote cache blob pack {name} header is not valid UTF-8"))?;
271 let value = value
272 .parse::<u64>()
273 .map_err(|_| eyre!("remote cache blob pack {name} header is not an unsigned integer"))?;
274 Ok(Some(value))
275}
276
277type RemoteCacheCapabilities = Capabilities;
278
279#[derive(Debug, Clone, Copy)]
280struct BlobPackLimits {
281 max_items: usize,
282 max_bytes: u64,
283}
284
285#[derive(Debug, Clone, Copy, Default)]
291struct NegotiatedCapabilities {
292 blob_packs: Option<BlobPackLimits>,
293 zstd_uploads: bool,
294}
295
296#[derive(Serialize)]
297struct DigestList<'a> {
298 digests: &'a [CacheDigest],
299}
300
301#[derive(Debug, Clone, Copy, PartialEq, Eq)]
302pub enum ManifestPutOutcome {
304 Stored,
306 PreconditionFailed,
308}
309
310pub struct RemoteCacheClient {
315 base_url: Url,
316 namespace: String,
317 client: reqwest::Client,
318 credential: RemoteCacheCredential,
319 download_timeout: Duration,
320 retries: i64,
321 capabilities: tokio::sync::OnceCell<NegotiatedCapabilities>,
322 blob_packs_disabled: AtomicBool,
323}
324
325impl RemoteCacheClient {
326 pub fn new(config: RemoteCacheConfig) -> Result<Self> {
328 let authenticated = config
329 .token
330 .as_deref()
331 .is_some_and(|token| !token.trim().is_empty())
332 || config.token_file.is_some()
333 || config
334 .oidc_audience
335 .as_deref()
336 .is_some_and(|audience| !audience.trim().is_empty());
337 validate_remote_url(&config.base_url, authenticated)?;
338 let client = reqwest::Client::builder()
339 .connect_timeout(config.connect_timeout)
340 .read_timeout(config.read_timeout)
341 .redirect(reqwest::redirect::Policy::none())
342 .build()?;
343 let credential = remote_credential(&config, client.clone())?;
344 Ok(Self {
345 base_url: normalized_base_url(config.base_url),
346 namespace: config.namespace,
347 client,
348 credential,
349 download_timeout: config.download_timeout,
350 retries: config.retries,
351 capabilities: tokio::sync::OnceCell::new(),
352 blob_packs_disabled: AtomicBool::new(false),
353 })
354 }
355
356 pub async fn check_connection(&self) -> Result<()> {
362 self.fetch_capabilities(false).await?;
363 Ok(())
364 }
365
366 fn action_result_endpoint(&self, action: &CacheDigest) -> Result<Url> {
367 action.validate()?;
368 if action.algorithm != "blake3" {
369 bail!("remote cache action keys must use blake3");
370 }
371 Ok(self.base_url.join(&format!(
372 "v{PROTOCOL_VERSION}/action-results/{}/{}/{}",
373 action.algorithm, action.hash, action.size
374 ))?)
375 }
376
377 fn blob_endpoint(&self, digest: &CacheDigest) -> Result<Url> {
378 digest.validate()?;
379 Ok(self.base_url.join(&format!(
380 "v{PROTOCOL_VERSION}/blobs/{}/{}/{}",
381 digest.algorithm, digest.hash, digest.size
382 ))?)
383 }
384
385 fn action_manifest_endpoint(&self, key: &CacheDigest) -> Result<Url> {
386 key.validate()?;
387 if key.algorithm != "blake3" {
388 bail!("remote action manifest keys must use blake3");
389 }
390 Ok(self.base_url.join(&format!(
391 "v{PROTOCOL_VERSION}/action-manifests/{}/{}/{}",
392 key.algorithm, key.hash, key.size
393 ))?)
394 }
395
396 fn capabilities_endpoint(&self) -> Result<Url> {
397 Ok(self
398 .base_url
399 .join(&format!("v{PROTOCOL_VERSION}/capabilities"))?)
400 }
401
402 fn blob_pack_endpoint(&self) -> Result<Url> {
403 Ok(self
404 .base_url
405 .join(&format!("v{PROTOCOL_VERSION}/blobs:pack"))?)
406 }
407
408 async fn request(
409 &self,
410 method: reqwest::Method,
411 url: Url,
412 media_type: &'static str,
413 ) -> Result<reqwest::RequestBuilder> {
414 let request = self
415 .client
416 .request(method, url)
417 .header(PROTOCOL_HEADER, u16::from(PROTOCOL_VERSION))
418 .header(NAMESPACE_HEADER, &self.namespace)
419 .header(ACCEPT, media_type);
420 if let Some(authorization) = self.credential.authorization().await? {
421 Ok(request.header(AUTHORIZATION, authorization))
422 } else {
423 Ok(request)
424 }
425 }
426
427 async fn blob_pack_limits(&self) -> Result<Option<BlobPackLimits>> {
428 Ok(self.negotiated_capabilities().await?.blob_packs)
429 }
430
431 async fn negotiated_capabilities(&self) -> Result<NegotiatedCapabilities> {
432 self.capabilities
433 .get_or_try_init(|| self.fetch_capabilities(true))
434 .await
435 .copied()
436 }
437
438 async fn fetch_capabilities(&self, allow_missing: bool) -> Result<NegotiatedCapabilities> {
439 let url = self.capabilities_endpoint()?;
440 let response = self
441 .request(reqwest::Method::GET, url, "application/json")
442 .await?
443 .send()
444 .await?;
445 if allow_missing
446 && matches!(
447 response.status(),
448 StatusCode::NOT_FOUND
449 | StatusCode::METHOD_NOT_ALLOWED
450 | StatusCode::NOT_IMPLEMENTED
451 )
452 {
453 return Ok(NegotiatedCapabilities::default());
454 }
455 let bytes = read_bounded_json(response.error_for_status()?, "capabilities").await?;
456 let capabilities: RemoteCacheCapabilities = serde_json::from_slice(&bytes)?;
457 if capabilities.protocol.major != PROTOCOL_VERSION {
458 bail!(
459 "remote cache capability protocol {} is incompatible with client protocol {PROTOCOL_VERSION}",
460 capabilities.protocol.major
461 );
462 }
463 let zstd_uploads = capabilities
467 .compressors
468 .iter()
469 .any(|compressor| compressor == "zstd");
470 let blob_packs = if capabilities.features.blob_packs {
471 let max_items = usize::try_from(capabilities.limits.max_batch_items)
472 .ok()
473 .filter(|limit| *limit > 0)
474 .ok_or_else(|| {
475 eyre!("remote cache blob packs require a positive max_batch_items limit")
476 })?;
477 if capabilities.limits.max_pack_bytes == 0 {
478 bail!("remote cache blob packs require a positive max_pack_bytes limit");
479 }
480 Some(BlobPackLimits {
481 max_items: max_items.min(MAX_STAGED_BLOB_PACK_ITEMS),
482 max_bytes: capabilities
483 .limits
484 .max_pack_bytes
485 .min(MAX_STAGED_BLOB_PACK_BYTES),
486 })
487 } else {
488 None
489 };
490 Ok(NegotiatedCapabilities {
491 blob_packs,
492 zstd_uploads,
493 })
494 }
495
496 pub async fn get_blob_pack(
502 &self,
503 digests: &[CacheDigest],
504 staging_dir: &Path,
505 ) -> Result<Option<RemoteBlobPack>> {
506 self.get_blob_pack_with_limit(digests, staging_dir, MAX_STAGED_BLOB_PACK_BYTES)
507 .await
508 }
509
510 pub(crate) async fn get_blob_pack_with_limit(
511 &self,
512 digests: &[CacheDigest],
513 staging_dir: &Path,
514 max_bytes: u64,
515 ) -> Result<Option<RemoteBlobPack>> {
516 if digests.is_empty() || self.blob_packs_disabled.load(Ordering::Relaxed) {
517 return Ok(None);
518 }
519 let Some(mut limits) = self.blob_pack_limits().await? else {
520 return Ok(None);
521 };
522 limits.max_bytes = limits.max_bytes.min(max_bytes);
523 if limits.max_bytes == 0 {
524 bail!("remote cache download budget is exhausted");
525 }
526 fs::create_dir_all(staging_dir)?;
527 let chunk = blob_pack_chunk(digests, limits)?;
528 if chunk.is_empty() {
529 return Ok(Some(RemoteBlobPack {
530 _directory: tempfile::tempdir_in(staging_dir)?,
531 blobs: Vec::new(),
532 requests: 0,
533 requested: Vec::new(),
534 blob_count: 0,
535 payload_bytes: 0,
536 framed_bytes: BLOB_PACK_MAGIC.len() as u64,
537 }));
538 }
539 match self.download_blob_pack_chunk(&chunk, staging_dir).await? {
540 Some(pack) => Ok(Some(RemoteBlobPack {
541 _directory: pack.directory,
542 blobs: pack.blobs,
543 requests: 1,
544 requested: chunk,
545 blob_count: pack.metadata.blob_count,
546 payload_bytes: pack.metadata.payload_bytes,
547 framed_bytes: pack.metadata.framed_bytes,
548 })),
549 None => {
550 self.blob_packs_disabled.store(true, Ordering::Relaxed);
551 Ok(None)
552 }
553 }
554 }
555
556 async fn download_blob_pack_chunk(
557 &self,
558 digests: &[CacheDigest],
559 staging_dir: &Path,
560 ) -> Result<Option<DownloadedBlobPack>> {
561 let url = self.blob_pack_endpoint()?;
562 let body = serde_json::to_vec(&DigestList { digests })?;
563 let download_timeout = blob_pack_download_timeout(self.download_timeout, digests);
564 let download = retry_async("POST", &url, self.retries, || async {
565 let response = self
566 .request(reqwest::Method::POST, url.clone(), BLOB_PACK_MEDIA_TYPE)
567 .await?
568 .header(CONTENT_TYPE, DIGEST_LIST_MEDIA_TYPE)
569 .body(body.clone())
570 .send()
571 .await?;
572 if matches!(
573 response.status(),
574 StatusCode::NOT_FOUND
575 | StatusCode::METHOD_NOT_ALLOWED
576 | StatusCode::NOT_IMPLEMENTED
577 ) {
578 return Ok(None);
579 }
580 let response = response.error_for_status()?;
581 let media_type = response
582 .headers()
583 .get(CONTENT_TYPE)
584 .and_then(|value| value.to_str().ok())
585 .and_then(|value| value.split(';').next())
586 .map(str::trim);
587 if media_type != Some(BLOB_PACK_MEDIA_TYPE) {
588 bail!("remote cache blob pack has an invalid content type");
589 }
590 Ok(Some(
591 decode_blob_pack(response, digests, staging_dir).await?,
592 ))
593 });
594 tokio::time::timeout(download_timeout, download)
595 .await
596 .map_err(|_| eyre!("remote cache blob pack download timed out for {url}"))?
597 }
598
599 pub async fn get_action_result(
601 &self,
602 action: &CacheDigest,
603 ) -> Result<Option<RemoteActionResult>> {
604 let url = self.action_result_endpoint(action)?;
605 let result = retry_async("GET", &url, self.retries, || async {
606 let response = self
607 .request(reqwest::Method::GET, url.clone(), ACTION_RESULT_MEDIA_TYPE)
608 .await?
609 .send()
610 .await?;
611 if response.status() == StatusCode::NOT_FOUND {
612 return Ok(None);
613 }
614 let bytes = read_bounded_json(response.error_for_status()?, "action result").await?;
615 Ok(Some(serde_json::from_slice::<RemoteActionResult>(&bytes)?))
616 })
617 .await?;
618 if let Some(result) = &result
619 && (result.version != 1 || result.action != *action)
620 {
621 bail!("remote action result does not match requested action");
622 }
623 Ok(result)
624 }
625
626 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
628 let url = self.action_result_endpoint(&result.action)?;
629 let body = serde_json::to_vec(result)?;
630 retry_async("PUT", &url, self.retries, || async {
631 let response = self
632 .request(reqwest::Method::PUT, url.clone(), ACTION_RESULT_MEDIA_TYPE)
633 .await?
634 .header(CONTENT_TYPE, ACTION_RESULT_MEDIA_TYPE)
635 .header(IF_NONE_MATCH, "*")
636 .body(body.clone())
637 .send()
638 .await?;
639 if response.status() != StatusCode::PRECONDITION_FAILED {
640 response.error_for_status()?;
641 }
642 Ok(())
643 })
644 .await
645 }
646
647 pub async fn get_action_manifest(
649 &self,
650 key: &CacheDigest,
651 ) -> Result<Option<RemoteActionManifest>> {
652 let url = self.action_manifest_endpoint(key)?;
653 retry_async("GET", &url, self.retries, || async {
654 let response = self
655 .request(
656 reqwest::Method::GET,
657 url.clone(),
658 TASK_ACTION_MANIFEST_MEDIA_TYPE,
659 )
660 .await?
661 .send()
662 .await?;
663 if response.status() == StatusCode::NOT_FOUND {
664 return Ok(None);
665 }
666 let response = response.error_for_status()?;
667 let etag = parse_strong_etag(response.headers().get(ETAG))?;
668 let bytes = read_bounded_json(response, "action manifest").await?;
669 Ok(Some(RemoteActionManifest { bytes, etag }))
670 })
671 .await
672 }
673
674 pub async fn put_action_manifest(
676 &self,
677 key: &CacheDigest,
678 bytes: &[u8],
679 expected_etag: Option<&str>,
680 ) -> Result<ManifestPutOutcome> {
681 let url = self.action_manifest_endpoint(key)?;
682 let body = bytes.to_vec();
683 let expected_etag = expected_etag.map(quoted_etag).transpose()?;
684 retry_async("PUT", &url, self.retries, || async {
685 let mut request = self
686 .request(
687 reqwest::Method::PUT,
688 url.clone(),
689 TASK_ACTION_MANIFEST_MEDIA_TYPE,
690 )
691 .await?
692 .header(CONTENT_TYPE, TASK_ACTION_MANIFEST_MEDIA_TYPE)
693 .body(body.clone());
694 request = if let Some(etag) = &expected_etag {
695 request.header(IF_MATCH, etag)
696 } else {
697 request.header(IF_NONE_MATCH, "*")
698 };
699 let response = request.send().await?;
700 if response.status() == StatusCode::PRECONDITION_FAILED {
701 return Ok(ManifestPutOutcome::PreconditionFailed);
702 }
703 response.error_for_status()?;
704 Ok(ManifestPutOutcome::Stored)
705 })
706 .await
707 }
708
709 pub async fn get_blob(
711 &self,
712 digest: &CacheDigest,
713 media_type: &'static str,
714 ) -> Result<Vec<u8>> {
715 digest.validate()?;
716 if digest.size > MAX_REMOTE_JSON_BYTES {
717 bail!(
718 "remote cache in-memory blob declared {} bytes, over the {} byte limit",
719 digest.size,
720 MAX_REMOTE_JSON_BYTES
721 );
722 }
723 let url = self.blob_endpoint(digest)?;
724 retry_async("GET", &url, self.retries, || async {
725 let mut response = self
726 .request(reqwest::Method::GET, url.clone(), media_type)
727 .await?
728 .send()
729 .await?
730 .error_for_status()?;
731 let mut bytes = Vec::new();
735 while let Some(chunk) = response.chunk().await? {
736 if bytes.len() as u64 + chunk.len() as u64 > digest.size {
737 bail!("remote cache blob exceeded the size of its digest");
738 }
739 bytes.extend_from_slice(&chunk);
740 }
741 if !digest.matches_bytes(&bytes)? {
742 bail!("remote cache blob failed digest verification");
743 }
744 Ok(bytes)
745 })
746 .await
747 }
748
749 pub async fn get_blob_file(
751 &self,
752 digest: &CacheDigest,
753 staging_dir: &Path,
754 ) -> Result<tempfile::NamedTempFile> {
755 digest.validate()?;
756 if digest.size > MAX_REMOTE_BLOB_BYTES {
757 bail!(
758 "remote cache blob declared {} bytes, over the {} byte limit",
759 digest.size,
760 MAX_REMOTE_BLOB_BYTES
761 );
762 }
763 let url = self.blob_endpoint(digest)?;
764 let download = retry_async("GET", &url, self.retries, || async {
765 let mut response = self
766 .request(reqwest::Method::GET, url.clone(), BLOB_MEDIA_TYPE)
767 .await?
768 .send()
769 .await?;
770 response.error_for_status_ref()?;
771 fs::create_dir_all(staging_dir)?;
772 let temporary = tempfile::NamedTempFile::new_in(staging_dir)?;
773 let mut output = tokio::fs::File::from_std(temporary.reopen()?);
774 let mut written = 0u64;
777 while let Some(chunk) = response.chunk().await? {
778 written += chunk.len() as u64;
779 if written > digest.size {
780 bail!("remote cache blob exceeded the size of its digest");
781 }
782 output.write_all(&chunk).await?;
783 }
784 output.flush().await?;
785 drop(output);
786 if !digest.matches_file(temporary.path())? {
787 bail!("remote cache blob failed digest verification");
788 }
789 Ok(temporary)
790 });
791 tokio::time::timeout(self.download_timeout, download)
792 .await
793 .map_err(|_| eyre!("remote cache blob download timed out for {url}"))?
794 }
795
796 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
798 let url = self.blob_endpoint(&upload.digest)?;
799 let compress = self
802 .negotiated_capabilities()
803 .await
804 .map(|capabilities| capabilities.zstd_uploads)
805 .unwrap_or(false);
806 retry_async("PUT", &url, self.retries, || async {
807 let request = self
808 .request(reqwest::Method::PUT, url.clone(), BLOB_MEDIA_TYPE)
809 .await?
810 .header(CONTENT_TYPE, BLOB_MEDIA_TYPE)
811 .header(IF_NONE_MATCH, "*");
812 let request = if compress {
813 let reader: Box<dyn tokio::io::AsyncRead + Send + Sync + Unpin> = match &upload
817 .source
818 {
819 BlobSource::Bytes(bytes) => Box::new(std::io::Cursor::new(bytes.clone())),
820 BlobSource::File(file) => Box::new(tokio::fs::File::open(file.path()).await?),
821 BlobSource::Path(path) => Box::new(tokio::fs::File::open(path).await?),
822 };
823 let encoder = async_compression::tokio::bufread::ZstdEncoder::new(
824 tokio::io::BufReader::new(reader),
825 );
826 request
827 .header(CONTENT_ENCODING, "zstd")
828 .body(reqwest::Body::wrap_stream(
829 tokio_util::io::ReaderStream::new(encoder),
830 ))
831 } else {
832 let (length, body) = match &upload.source {
833 BlobSource::Bytes(bytes) => {
834 (bytes.len() as u64, reqwest::Body::from(bytes.clone()))
835 }
836 BlobSource::File(file) => {
837 let file = tokio::fs::File::open(file.path()).await?;
838 let length = file.metadata().await?.len();
839 let stream = tokio_util::io::ReaderStream::new(file);
840 (length, reqwest::Body::wrap_stream(stream))
841 }
842 BlobSource::Path(path) => {
843 let file = tokio::fs::File::open(path).await?;
844 let length = file.metadata().await?.len();
845 let stream = tokio_util::io::ReaderStream::new(file);
846 (length, reqwest::Body::wrap_stream(stream))
847 }
848 };
849 request.header(CONTENT_LENGTH, length).body(body)
850 };
851 let response = request.send().await?;
852 if response.status() != StatusCode::PRECONDITION_FAILED {
853 response.error_for_status()?;
854 }
855 Ok(())
856 })
857 .await
858 }
859}
860
861fn blob_pack_chunk(digests: &[CacheDigest], limits: BlobPackLimits) -> Result<Vec<CacheDigest>> {
862 let mut seen = BTreeSet::new();
863 let mut chunk = Vec::new();
864 let mut chunk_bytes = 0_u64;
865 for digest in digests {
866 digest.validate()?;
867 if !seen.insert(digest.clone()) || digest.size > limits.max_bytes {
868 continue;
869 }
870 if chunk.len() == limits.max_items
871 || chunk_bytes.saturating_add(digest.size) > limits.max_bytes
872 {
873 break;
874 }
875 chunk_bytes = chunk_bytes.saturating_add(digest.size);
876 chunk.push(digest.clone());
877 }
878 Ok(chunk)
879}
880
881fn blob_pack_download_timeout(base: Duration, digests: &[CacheDigest]) -> Duration {
882 let bytes = digests
883 .iter()
884 .fold(0_u64, |total, digest| total.saturating_add(digest.size));
885 let byte_units = bytes.div_ceil(BLOB_PACK_TIMEOUT_BYTES_PER_UNIT);
886 let item_units = digests.len().div_ceil(BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT);
887 let item_units = u64::try_from(item_units).unwrap_or(u64::MAX);
888 let multiplier = byte_units.max(item_units).max(1);
889 base.saturating_mul(u32::try_from(multiplier).unwrap_or(u32::MAX))
890}
891
892async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
897 if let Some(length) = response.content_length()
898 && length > MAX_REMOTE_JSON_BYTES
899 {
900 bail!(
901 "remote cache {what} declared {length} bytes, over the {MAX_REMOTE_JSON_BYTES} byte limit"
902 );
903 }
904 let mut response = response;
905 let mut bytes = Vec::new();
906 while let Some(chunk) = response.chunk().await? {
907 if bytes.len() as u64 + chunk.len() as u64 > MAX_REMOTE_JSON_BYTES {
908 bail!("remote cache {what} exceeded the {MAX_REMOTE_JSON_BYTES} byte limit");
909 }
910 bytes.extend_from_slice(&chunk);
911 }
912 Ok(bytes)
913}
914
915async fn decode_blob_pack(
916 response: reqwest::Response,
917 requested: &[CacheDigest],
918 staging_dir: &Path,
919) -> Result<DownloadedBlobPack> {
920 let metadata = BlobPackResponseMetadata::from_headers(response.headers())?;
921 let stream = response.bytes_stream().map_err(std::io::Error::other);
922 let reader = tokio_util::io::StreamReader::new(stream);
923 decode_blob_pack_reader(reader, metadata, requested, staging_dir).await
924}
925
926async fn decode_blob_pack_reader<R>(
927 mut reader: R,
928 metadata: BlobPackResponseMetadata,
929 requested: &[CacheDigest],
930 staging_dir: &Path,
931) -> Result<DownloadedBlobPack>
932where
933 R: AsyncRead + Unpin,
934{
935 let requested = requested.iter().cloned().collect::<BTreeSet<_>>();
936 let mut magic = [0_u8; BLOB_PACK_MAGIC.len()];
937 reader.read_exact(&mut magic).await?;
938 if &magic != BLOB_PACK_MAGIC {
939 bail!("remote cache blob pack has invalid magic");
940 }
941
942 let directory = tempfile::tempdir_in(staging_dir)?;
943 let mut seen = BTreeSet::new();
944 let mut blobs = Vec::new();
945 let mut payload_bytes = 0_u64;
946 let mut framed_bytes = BLOB_PACK_MAGIC.len() as u64;
947 loop {
948 let mut algorithm = [0_u8; 1];
949 if reader.read(&mut algorithm).await? == 0 {
950 break;
951 }
952 let (algorithm, mut hasher) = match algorithm[0] {
953 1 => (
954 "blake3",
955 BlobPackHasher::Blake3(Box::new(blake3::Hasher::new())),
956 ),
957 2 => ("sha256", BlobPackHasher::Sha256(sha2::Sha256::new())),
958 _ => bail!("remote cache blob pack has an invalid digest algorithm"),
959 };
960 let mut hash = [0_u8; 32];
961 reader.read_exact(&mut hash).await?;
962 let mut size = [0_u8; 8];
963 reader.read_exact(&mut size).await?;
964 let digest = CacheDigest {
965 algorithm: algorithm.into(),
966 hash: hex::encode(hash),
967 size: u64::from_be_bytes(size),
968 };
969 if !requested.contains(&digest) {
970 bail!("remote cache blob pack returned an unrequested digest");
971 }
972 if !seen.insert(digest.clone()) {
973 bail!("remote cache blob pack returned a duplicate digest");
974 }
975 framed_bytes = framed_bytes
976 .checked_add(BLOB_PACK_HEADER_BYTES)
977 .and_then(|bytes| bytes.checked_add(digest.size))
978 .ok_or_else(|| eyre!("remote cache blob pack is too large"))?;
979 payload_bytes = payload_bytes
980 .checked_add(digest.size)
981 .ok_or_else(|| eyre!("remote cache blob pack payload is too large"))?;
982
983 let path = directory.path().join(blobs.len().to_string());
984 let mut output = tokio::fs::File::create(&path).await?;
985 let mut remaining = digest.size;
986 let mut buffer = [0_u8; 64 * 1024];
987 while remaining > 0 {
988 let limit = usize::try_from(remaining.min(buffer.len() as u64)).unwrap();
989 let count = reader.read(&mut buffer[..limit]).await?;
990 if count == 0 {
991 bail!("remote cache blob pack ended before a blob was complete");
992 }
993 output.write_all(&buffer[..count]).await?;
994 hasher.update(&buffer[..count]);
995 remaining -= count as u64;
996 }
997 output.flush().await?;
998 drop(output);
999 if !hasher.matches(&digest.hash) {
1000 bail!("remote cache blob pack failed digest verification");
1001 }
1002 blobs.push((digest, path));
1003 }
1004 let blob_count = blobs.len().try_into().unwrap_or(u64::MAX);
1005 let metadata = metadata.validate(BlobPackResponseStats {
1006 blob_count,
1007 payload_bytes,
1008 framed_bytes,
1009 })?;
1010 Ok(DownloadedBlobPack {
1011 directory,
1012 blobs,
1013 metadata,
1014 })
1015}
1016
1017#[cfg(feature = "fuzzing")]
1020#[doc(hidden)]
1021pub async fn fuzz_decode_blob_pack(
1022 bytes: &[u8],
1023 requested: &[CacheDigest],
1024 staging_dir: &Path,
1025) -> Result<()> {
1026 decode_blob_pack_reader(
1027 bytes,
1028 BlobPackResponseMetadata::default(),
1029 requested,
1030 staging_dir,
1031 )
1032 .await
1033 .map(drop)
1034}
1035
1036enum BlobPackHasher {
1037 Blake3(Box<blake3::Hasher>),
1038 Sha256(sha2::Sha256),
1039}
1040
1041impl BlobPackHasher {
1042 fn update(&mut self, bytes: &[u8]) {
1043 match self {
1044 Self::Blake3(hasher) => {
1045 hasher.update(bytes);
1046 }
1047 Self::Sha256(hasher) => {
1048 hasher.update(bytes);
1049 }
1050 }
1051 }
1052
1053 fn matches(self, expected: &str) -> bool {
1054 match self {
1055 Self::Blake3(hasher) => hasher.finalize().to_hex().as_str() == expected,
1056 Self::Sha256(hasher) => hex::encode(hasher.finalize()) == expected,
1057 }
1058 }
1059}
1060
1061fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
1070 let value = value
1071 .and_then(|value| value.to_str().ok())
1072 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
1073 if value.starts_with("W/") {
1074 bail!("remote action manifest response has a weak ETag");
1077 }
1078 let etag = value
1079 .strip_prefix('"')
1080 .and_then(|value| value.strip_suffix('"'))
1081 .filter(|value| is_entity_tag(value))
1082 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
1083 Ok(etag.to_owned())
1084}
1085
1086fn quoted_etag(etag: &str) -> Result<HeaderValue> {
1087 if !is_entity_tag(etag) {
1088 bail!("invalid remote action manifest ETag");
1089 }
1090 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
1091}
1092
1093fn is_entity_tag(value: &str) -> bool {
1098 !value.is_empty()
1099 && value.len() <= MAX_ETAG_BYTES
1100 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
1101}
1102
1103#[derive(Clone)]
1104enum RemoteCacheCredential {
1105 None,
1106 Static(HeaderValue),
1107 File(PathBuf),
1108 GithubActions(Arc<GithubActionsOidcCredential>),
1109}
1110
1111struct GithubActionsOidcCredential {
1112 audience: String,
1113 request_url: Url,
1114 request_token: HeaderValue,
1115 client: reqwest::Client,
1116 retries: i64,
1117 cached: tokio::sync::Mutex<Option<CachedOidcToken>>,
1118}
1119
1120struct CachedOidcToken {
1121 authorization: HeaderValue,
1122 expires_at: u64,
1123}
1124
1125#[derive(Deserialize)]
1126struct GithubActionsOidcResponse {
1127 value: String,
1128}
1129
1130#[derive(Deserialize)]
1131struct JwtExpiry {
1132 exp: u64,
1133}
1134
1135fn remote_credential(
1136 config: &RemoteCacheConfig,
1137 client: reqwest::Client,
1138) -> Result<RemoteCacheCredential> {
1139 if let Some(authorization) = authorization_header(config.token.as_deref())? {
1140 return Ok(RemoteCacheCredential::Static(authorization));
1141 }
1142 if let Some(path) = &config.token_file {
1143 return Ok(RemoteCacheCredential::File(path.clone()));
1144 }
1145 let Some(audience) = config
1146 .oidc_audience
1147 .as_deref()
1148 .map(str::trim)
1149 .filter(|audience| !audience.is_empty())
1150 else {
1151 return Ok(RemoteCacheCredential::None);
1152 };
1153 Ok(RemoteCacheCredential::GithubActions(Arc::new(
1154 GithubActionsOidcCredential::from_env(audience, client, config.retries)?,
1155 )))
1156}
1157
1158fn authorization_header(token: Option<&str>) -> Result<Option<HeaderValue>> {
1159 let Some(token) = token.map(str::trim).filter(|token| !token.is_empty()) else {
1160 return Ok(None);
1161 };
1162 let mut value = HeaderValue::from_str(&format!("Bearer {token}"))?;
1163 value.set_sensitive(true);
1164 Ok(Some(value))
1165}
1166
1167impl RemoteCacheCredential {
1168 async fn authorization(&self) -> Result<Option<HeaderValue>> {
1169 match self {
1170 Self::None => Ok(None),
1171 Self::Static(value) => Ok(Some(value.clone())),
1172 Self::File(path) => {
1173 let token = tokio::fs::read_to_string(path).await.map_err(|err| {
1174 eyre!(
1175 "failed to read remote cache token file {}: {err}",
1176 path.display()
1177 )
1178 })?;
1179 authorization_header(Some(&token))?
1180 .ok_or_else(|| eyre!("remote cache token file {} is empty", path.display()))
1181 .map(Some)
1182 }
1183 Self::GithubActions(credential) => credential.authorization().await.map(Some),
1184 }
1185 }
1186}
1187
1188impl GithubActionsOidcCredential {
1189 fn from_env(audience: &str, client: reqwest::Client, retries: i64) -> Result<Self> {
1190 let request_url = std::env::var("ACTIONS_ID_TOKEN_REQUEST_URL").map_err(|_| {
1191 eyre!(
1192 "remote cache OIDC audience requires GitHub Actions OIDC; \
1193 grant `id-token: write` or set MBX_REMOTE_TOKEN"
1194 )
1195 })?;
1196 let request_token = std::env::var("ACTIONS_ID_TOKEN_REQUEST_TOKEN").map_err(|_| {
1197 eyre!(
1198 "remote cache OIDC audience requires GitHub Actions OIDC; \
1199 ACTIONS_ID_TOKEN_REQUEST_TOKEN is missing"
1200 )
1201 })?;
1202 let request_url: Url = request_url
1203 .parse()
1204 .map_err(|err| eyre!("invalid GitHub Actions OIDC request URL: {err}"))?;
1205 Self::new(audience, request_url, &request_token, client, retries)
1206 }
1207
1208 fn new(
1209 audience: &str,
1210 mut request_url: Url,
1211 request_token: &str,
1212 client: reqwest::Client,
1213 retries: i64,
1214 ) -> Result<Self> {
1215 validate_oidc_request_url(&request_url)?;
1216 let query = request_url
1217 .query_pairs()
1218 .filter(|(key, _)| key != "audience")
1219 .map(|(key, value)| (key.into_owned(), value.into_owned()))
1220 .collect::<Vec<_>>();
1221 request_url.set_query(None);
1222 request_url
1223 .query_pairs_mut()
1224 .extend_pairs(query)
1225 .append_pair("audience", audience);
1226 let request_token = authorization_header(Some(request_token))?
1227 .ok_or_else(|| eyre!("GitHub Actions OIDC request token is empty"))?;
1228 Ok(Self {
1229 audience: audience.to_string(),
1230 request_url,
1231 request_token,
1232 client,
1233 retries,
1234 cached: tokio::sync::Mutex::new(None),
1235 })
1236 }
1237
1238 async fn authorization(&self) -> Result<HeaderValue> {
1239 const REFRESH_LEEWAY_SECONDS: u64 = 60;
1240 let mut cached = self.cached.lock().await;
1241 let now = unix_timestamp()?;
1242 if let Some(token) = cached.as_ref()
1243 && token.expires_at > now.saturating_add(REFRESH_LEEWAY_SECONDS)
1244 {
1245 return Ok(token.authorization.clone());
1246 }
1247 let response: GithubActionsOidcResponse =
1248 retry_async("GET", &self.request_url, self.retries, || async {
1249 Ok(self
1250 .client
1251 .get(self.request_url.clone())
1252 .header(AUTHORIZATION, self.request_token.clone())
1253 .send()
1254 .await?
1255 .error_for_status()?
1256 .json()
1257 .await?)
1258 })
1259 .await
1260 .map_err(|err| {
1261 eyre!(
1262 "failed to acquire GitHub Actions OIDC token for audience {:?}: {err}",
1263 self.audience
1264 )
1265 })?;
1266 let expires_at = jwt_expiry(&response.value)?;
1267 if expires_at <= now.saturating_add(REFRESH_LEEWAY_SECONDS) {
1268 bail!("GitHub Actions OIDC token expires too soon");
1269 }
1270 let authorization = authorization_header(Some(&response.value))?
1271 .ok_or_else(|| eyre!("GitHub Actions returned an empty OIDC token"))?;
1272 *cached = Some(CachedOidcToken {
1273 authorization: authorization.clone(),
1274 expires_at,
1275 });
1276 Ok(authorization)
1277 }
1278}
1279
1280fn jwt_expiry(token: &str) -> Result<u64> {
1281 let payload = token
1282 .split('.')
1283 .nth(1)
1284 .ok_or_else(|| eyre!("GitHub Actions returned a malformed OIDC token"))?;
1285 let payload = URL_SAFE_NO_PAD
1286 .decode(payload)
1287 .map_err(|_| eyre!("GitHub Actions returned a malformed OIDC token"))?;
1288 let claims: JwtExpiry = serde_json::from_slice(&payload)
1289 .map_err(|_| eyre!("GitHub Actions OIDC token is missing a valid expiry"))?;
1290 Ok(claims.exp)
1291}
1292
1293fn unix_timestamp() -> Result<u64> {
1294 Ok(SystemTime::now()
1295 .duration_since(UNIX_EPOCH)
1296 .map_err(|err| eyre!("system clock is before the Unix epoch: {err}"))?
1297 .as_secs())
1298}
1299
1300fn validate_oidc_request_url(url: &Url) -> Result<()> {
1301 if url.scheme() == "https"
1302 || url.scheme() == "http"
1303 && url.host().is_some_and(|host| match host {
1304 Host::Domain(host) => host.eq_ignore_ascii_case("localhost"),
1305 Host::Ipv4(address) => address.is_loopback(),
1306 Host::Ipv6(address) => address.is_loopback(),
1307 })
1308 {
1309 Ok(())
1310 } else {
1311 bail!("GitHub Actions OIDC request URL must use HTTPS")
1312 }
1313}
1314
1315fn validate_remote_url(base_url: &Url, authenticated: bool) -> Result<()> {
1316 if base_url.scheme() == "https" {
1317 return Ok(());
1318 }
1319 if base_url.scheme() != "http" {
1320 bail!("remote cache URL must use HTTPS");
1321 }
1322 let is_loopback = base_url.host().is_some_and(|host| match host {
1323 Host::Domain(host) => host.eq_ignore_ascii_case("localhost"),
1324 Host::Ipv4(address) => address.is_loopback(),
1325 Host::Ipv6(address) => address.is_loopback(),
1326 });
1327 if !is_loopback && authenticated {
1328 bail!("remote cache URL must use HTTPS except for loopback development servers");
1329 }
1330 if !is_loopback {
1331 warn!(
1332 "using an unauthenticated remote build cache over plain HTTP; cache traffic can be read \
1333 or modified in transit"
1334 );
1335 }
1336 Ok(())
1337}
1338
1339fn normalized_base_url(mut url: Url) -> Url {
1340 if !url.path().ends_with('/') {
1341 url.set_path(&format!("{}/", url.path()));
1342 }
1343 url
1344}
1345
1346fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
1347 [200u64, 1_000, 4_000, 15_000]
1348 .into_iter()
1349 .chain(std::iter::repeat(15_000))
1350 .map(Duration::from_millis)
1351 .map(|duration| {
1352 let factor = 0.5 + rand::random::<f64>() * 0.5;
1353 Duration::from_secs_f64(duration.as_secs_f64() * factor)
1354 })
1355 .take(retries.max(0) as usize)
1356}
1357
1358fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
1362 let mut current = Some(error);
1363 while let Some(source) = current {
1364 if source.to_string() == "dns error" {
1365 return true;
1366 }
1367 current = source.source();
1368 }
1369 false
1370}
1371
1372fn is_transient(error: &eyre::Report) -> bool {
1373 if is_dns_error(error.as_ref()) {
1376 return false;
1377 }
1378 error.chain().any(|source| {
1379 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
1380 return false;
1381 };
1382 if error.is_timeout() || error.is_connect() || error.is_body() {
1383 return true;
1384 }
1385 error.status().is_some_and(|status| {
1386 let status = status.as_u16();
1387 status == 408 || status == 429 || (500..600).contains(&status)
1388 })
1389 })
1390}
1391
1392async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
1393where
1394 F: FnMut() -> Fut,
1395 Fut: std::future::Future<Output = Result<T>>,
1396{
1397 let mut delays = retry_delays(retries);
1398 let mut attempt = 1;
1399 loop {
1400 let started_at = Instant::now();
1401 match operation().await {
1402 Ok(value) => return Ok(value),
1403 Err(error) if is_transient(&error) => {
1404 let Some(delay) = delays.next() else {
1405 return Err(error);
1406 };
1407 warn!(
1408 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
1409 started_at.elapsed()
1410 );
1411 tokio::time::sleep(delay).await;
1412 attempt += 1;
1413 }
1414 Err(error) => return Err(error),
1415 }
1416 }
1417}
1418
1419#[cfg(test)]
1420#[path = "core_tests.rs"]
1421mod tests;