1#![deny(missing_docs)]
30
31use eyre::{Result, bail, eyre};
32use log::warn;
33use reqwest::header::HeaderValue;
34use serde::{Deserialize, Serialize};
35use std::path::{Path, PathBuf};
36use std::time::{Duration, Instant};
37use url::Url;
38
39mod agent;
40mod client;
41mod local;
42mod path_mapping;
43mod remote_http;
44mod remote_s3;
45mod sigv4;
46mod uploads;
47
48pub use agent::{
49 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRemoteCache,
50 AgentRequest, AgentResponse, AgentStats, CacheAgent, CompilerStats, FileDigestCache,
51 FileDigestResolution, FileDigestScope, FileIdentity, FileObjectIdentity, FileSnapshot,
52 NoFileDigestCache, RecordedFileDigest, RestoreStats, digest_file, is_task_identity,
53 task_manifest_actions,
54};
55pub use client::BlockingAgentClient;
56pub use local::{LocalActionCache, LocalCas};
57pub use mbx_cache_protocol::{
58 ACTION_PROMISE_MEDIA_TYPE, ACTION_RESULT_BATCH_MEDIA_TYPE, ACTION_RESULT_MEDIA_TYPE,
59 ActionPrediction, ActionPromiseCompletion, ActionPromiseJoin, ActionPromiseState,
60 ActionResult as RemoteActionResult, BLOB_MEDIA_TYPE, BLOB_PACK_BLOBS_HEADER,
61 BLOB_PACK_BYTES_HEADER, BLOB_PACK_HEADER_BYTES, BLOB_PACK_MAGIC, BLOB_PACK_MEDIA_TYPE,
62 BLOB_PACK_RECEIPT_MEDIA_TYPE, CLIENT_METADATA_MEDIA_TYPE, Capabilities, CapabilityFeatures,
63 CapabilityLimits, CapabilityProtocol, CcMetadata, DIGEST_LIST_MEDIA_TYPE, DIRECTORY_MEDIA_TYPE,
64 Digest as CacheDigest, DigestAlgorithm, Directory as CacheDirectory,
65 DirectoryNode as CacheDirectoryNode, FileNode as CacheFileNode, MAX_ACTION_PREDICTION_PAYLOAD,
66 MAX_ACTION_PROMISE_CLAIM_BYTES, NAMESPACE_HEADER, PROTOCOL_HEADER, PROTOCOL_VERSION,
67 RustcMetadata, SymlinkNode as CacheSymlinkNode, TASK_ACTION_MANIFEST_MEDIA_TYPE,
68 TaskActionManifest,
69};
70pub use path_mapping::{
71 PathMapping, PathNormalizationError, normalize_mapped_path, normalize_resolved_mapped_path,
72 resolve_path_mappings,
73};
74use remote_http::HttpRemoteCache;
75#[cfg(feature = "fuzzing")]
76#[doc(hidden)]
77pub use remote_http::fuzz_decode_blob_pack;
78pub(crate) use remote_http::{BlobPackLimits, blob_pack_chunk};
79use remote_s3::S3RemoteCache;
80pub use remote_s3::{S3ConditionalWrites, S3RemoteCacheConfig};
81pub use sigv4::S3Credentials;
82const MAX_REMOTE_JSON_BYTES: u64 = 16 * 1024 * 1024;
89const MAX_ETAG_BYTES: usize = 256;
94const MAX_REMOTE_BLOB_BYTES: u64 = 5 * 1024 * 1024 * 1024;
97const MAX_STAGED_BLOB_PACK_BYTES: u64 = 256 * 1024 * 1024;
98const MAX_STAGED_BLOB_PACK_ITEMS: usize = 2 * 1024;
99const BLOB_PACK_TIMEOUT_BYTES_PER_UNIT: u64 = MAX_STAGED_BLOB_PACK_BYTES / 4;
100const BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT: usize = MAX_STAGED_BLOB_PACK_ITEMS / 4;
101const PACK_STREAM_CHUNK_BYTES: usize = 64 * 1024;
103const MAX_ACTION_BATCH_ITEMS: usize = 256;
109const MAX_ACTION_RESULT_BYTES: u64 = 64 * 1024;
111
112pub fn canonical_json(value: &impl Serialize) -> Result<Vec<u8>> {
117 Ok(mbx_cache_protocol::canonical_json(value)?)
118}
119
120#[derive(
121 Debug,
122 Clone,
123 Copy,
124 Serialize,
125 Deserialize,
126 Default,
127 strum::EnumString,
128 strum::Display,
129 PartialEq,
130 Eq,
131)]
132#[serde(rename_all = "kebab-case")]
133#[strum(serialize_all = "kebab-case")]
134pub enum RemoteCacheMode {
136 #[default]
138 ReadWrite,
139 ReadOnly,
141 WriteOnly,
143}
144
145impl RemoteCacheMode {
146 pub fn reads(self) -> bool {
148 matches!(self, Self::ReadWrite | Self::ReadOnly)
149 }
150
151 pub fn writes(self) -> bool {
153 matches!(self, Self::ReadWrite | Self::WriteOnly)
154 }
155}
156
157pub struct RemoteCacheConfig {
159 pub base_url: Url,
161 pub namespace: String,
163 pub token: Option<String>,
165 pub token_file: Option<PathBuf>,
167 pub oidc_audience: Option<String>,
169 pub connect_timeout: Duration,
171 pub read_timeout: Duration,
173 pub download_timeout: Duration,
182 pub retries: i64,
184}
185
186pub enum BlobSource {
188 Bytes(Vec<u8>),
190 File(tempfile::NamedTempFile),
192 Path(PathBuf),
194}
195
196pub struct BlobUpload {
198 pub digest: CacheDigest,
200 pub source: BlobSource,
202}
203
204pub struct RemoteActionManifest {
206 pub bytes: Vec<u8>,
208 pub etag: String,
210}
211
212pub struct RemoteBlobPack {
214 _directory: tempfile::TempDir,
215 pub blobs: Vec<(CacheDigest, PathBuf)>,
217 pub requests: u64,
219 pub requested: Vec<CacheDigest>,
221 pub blob_count: u64,
223 pub payload_bytes: u64,
225 pub framed_bytes: u64,
227}
228
229#[derive(Debug, Clone, Copy, Deserialize)]
231pub struct BlobPackReceipt {
232 #[serde(default)]
234 pub created: u64,
235 #[serde(default)]
237 pub existing: u64,
238}
239
240#[derive(Debug, Clone, Copy, PartialEq, Eq)]
241pub enum ManifestPutOutcome {
243 Stored,
245 PreconditionFailed,
247}
248
249pub struct RemoteCacheClient {
259 backend: Backend,
260}
261
262enum Backend {
268 Http(HttpRemoteCache),
269 S3(S3RemoteCache),
270}
271
272impl RemoteCacheClient {
273 pub fn new(config: RemoteCacheConfig) -> Result<Self> {
275 Ok(Self {
276 backend: Backend::Http(HttpRemoteCache::new(config)?),
277 })
278 }
279
280 pub fn with_read_stall_budget(mut self, budget: Duration) -> Self {
291 if let Backend::Http(client) = &mut self.backend {
292 client.set_read_stall_budget(budget);
293 }
294 self
295 }
296
297 pub fn new_s3(config: S3RemoteCacheConfig) -> Result<Self> {
302 Ok(Self {
303 backend: Backend::S3(S3RemoteCache::new(config)?),
304 })
305 }
306
307 pub async fn check_connection(&self) -> Result<()> {
313 match &self.backend {
314 Backend::Http(client) => client.check_connection().await,
315 Backend::S3(store) => store.check_connection().await,
316 }
317 }
318
319 pub async fn get_blob_pack(
325 &self,
326 digests: &[CacheDigest],
327 staging_dir: &Path,
328 ) -> Result<Option<RemoteBlobPack>> {
329 match &self.backend {
330 Backend::Http(client) => client.get_blob_pack(digests, staging_dir).await,
331 Backend::S3(store) => store.get_blob_pack(digests, staging_dir).await,
332 }
333 }
334
335 pub(crate) async fn get_blob_pack_with_limit(
336 &self,
337 digests: &[CacheDigest],
338 staging_dir: &Path,
339 max_bytes: u64,
340 ) -> Result<Option<RemoteBlobPack>> {
341 match &self.backend {
342 Backend::Http(client) => {
343 client
344 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
345 .await
346 }
347 Backend::S3(store) => {
348 store
349 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
350 .await
351 }
352 }
353 }
354
355 pub(crate) async fn blob_pack_limits(&self) -> Result<Option<BlobPackLimits>> {
357 match &self.backend {
358 Backend::Http(client) => client.blob_pack_limits().await,
359 Backend::S3(_) => Ok(None),
360 }
361 }
362
363 pub async fn get_action_result(
365 &self,
366 action: &CacheDigest,
367 ) -> Result<Option<RemoteActionResult>> {
368 match &self.backend {
369 Backend::Http(client) => client.get_action_result(action).await,
370 Backend::S3(store) => store.get_action_result(action).await,
371 }
372 }
373
374 pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
378 match &self.backend {
379 Backend::Http(client) => client.action_batch_limit().await,
380 Backend::S3(store) => store.action_batch_limit().await,
381 }
382 }
383
384 pub async fn get_action_results(
391 &self,
392 actions: &[CacheDigest],
393 ) -> Result<Option<Vec<RemoteActionResult>>> {
394 match &self.backend {
395 Backend::Http(client) => client.get_action_results(actions).await,
396 Backend::S3(store) => store.get_action_results(actions).await,
397 }
398 }
399
400 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
402 match &self.backend {
403 Backend::Http(client) => client.put_action_result(result).await,
404 Backend::S3(store) => store.put_action_result(result).await,
405 }
406 }
407
408 pub async fn join_action_promise(
412 &self,
413 invocation: &CacheDigest,
414 adapter: &str,
415 ) -> Result<Option<ActionPromiseState>> {
416 match &self.backend {
417 Backend::Http(client) => client.join_action_promise(invocation, adapter).await,
418 Backend::S3(_) => Ok(None),
419 }
420 }
421
422 pub async fn complete_action_promise(
426 &self,
427 invocation: &CacheDigest,
428 completion: &ActionPromiseCompletion,
429 ) -> Result<bool> {
430 match &self.backend {
431 Backend::Http(client) => client.complete_action_promise(invocation, completion).await,
432 Backend::S3(_) => Ok(false),
433 }
434 }
435
436 pub async fn get_action_manifest(
438 &self,
439 key: &CacheDigest,
440 ) -> Result<Option<RemoteActionManifest>> {
441 match &self.backend {
442 Backend::Http(client) => client.get_action_manifest(key).await,
443 Backend::S3(store) => store.get_action_manifest(key).await,
444 }
445 }
446
447 pub async fn put_action_manifest(
449 &self,
450 key: &CacheDigest,
451 bytes: &[u8],
452 expected_etag: Option<&str>,
453 ) -> Result<ManifestPutOutcome> {
454 match &self.backend {
455 Backend::Http(client) => client.put_action_manifest(key, bytes, expected_etag).await,
456 Backend::S3(store) => store.put_action_manifest(key, bytes, expected_etag).await,
457 }
458 }
459
460 pub async fn get_blob(
462 &self,
463 digest: &CacheDigest,
464 media_type: &'static str,
465 ) -> Result<Vec<u8>> {
466 match &self.backend {
467 Backend::Http(client) => client.get_blob(digest, media_type).await,
468 Backend::S3(store) => store.get_blob(digest, media_type).await,
469 }
470 }
471
472 pub async fn get_blob_file(
474 &self,
475 digest: &CacheDigest,
476 staging_dir: &Path,
477 ) -> Result<tempfile::NamedTempFile> {
478 match &self.backend {
479 Backend::Http(client) => client.get_blob_file(digest, staging_dir).await,
480 Backend::S3(store) => store.get_blob_file(digest, staging_dir).await,
481 }
482 }
483
484 pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<BlobPackLimits>> {
488 match &self.backend {
489 Backend::Http(client) => client.blob_pack_upload_limits().await,
490 Backend::S3(store) => store.blob_pack_upload_limits().await,
491 }
492 }
493
494 pub async fn put_blob_pack(&self, uploads: &[BlobUpload]) -> Result<Option<BlobPackReceipt>> {
502 match &self.backend {
503 Backend::Http(client) => client.put_blob_pack(uploads).await,
504 Backend::S3(store) => store.put_blob_pack(uploads).await,
505 }
506 }
507
508 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
510 match &self.backend {
511 Backend::Http(client) => client.put_blob(upload).await,
512 Backend::S3(store) => store.put_blob(upload).await,
513 }
514 }
515}
516
517async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
523 read_json_within(response, what, MAX_REMOTE_JSON_BYTES).await
524}
525
526async fn read_json_within(response: reqwest::Response, what: &str, limit: u64) -> Result<Vec<u8>> {
527 if let Some(length) = response.content_length()
528 && length > limit
529 {
530 bail!("remote cache {what} declared {length} bytes, over the {limit} byte limit");
531 }
532 let mut response = response;
533 let mut bytes = Vec::new();
534 while let Some(chunk) = response.chunk().await? {
535 if bytes.len() as u64 + chunk.len() as u64 > limit {
536 bail!("remote cache {what} exceeded the {limit} byte limit");
537 }
538 bytes.extend_from_slice(&chunk);
539 }
540 Ok(bytes)
541}
542fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
551 let value = value
552 .and_then(|value| value.to_str().ok())
553 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
554 if value.starts_with("W/") {
555 bail!("remote action manifest response has a weak ETag");
558 }
559 let etag = value
560 .strip_prefix('"')
561 .and_then(|value| value.strip_suffix('"'))
562 .filter(|value| is_entity_tag(value))
563 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
564 Ok(etag.to_owned())
565}
566
567fn quoted_etag(etag: &str) -> Result<HeaderValue> {
568 if !is_entity_tag(etag) {
569 bail!("invalid remote action manifest ETag");
570 }
571 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
572}
573
574fn is_entity_tag(value: &str) -> bool {
579 !value.is_empty()
580 && value.len() <= MAX_ETAG_BYTES
581 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
582}
583fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
584 [200u64, 1_000, 4_000, 15_000]
585 .into_iter()
586 .chain(std::iter::repeat(15_000))
587 .map(Duration::from_millis)
588 .map(|duration| {
589 let factor = 0.5 + rand::random::<f64>() * 0.5;
590 Duration::from_secs_f64(duration.as_secs_f64() * factor)
591 })
592 .take(retries.max(0) as usize)
593}
594
595fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
599 let mut current = Some(error);
600 while let Some(source) = current {
601 if source.to_string() == "dns error" {
602 return true;
603 }
604 current = source.source();
605 }
606 false
607}
608
609#[derive(Debug)]
616pub(crate) struct TransientRequest(pub(crate) &'static str);
617
618impl std::fmt::Display for TransientRequest {
619 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
620 formatter.write_str(self.0)
621 }
622}
623
624impl std::error::Error for TransientRequest {}
625
626fn is_transient(error: &eyre::Report) -> bool {
627 if is_dns_error(error.as_ref()) {
630 return false;
631 }
632 error.chain().any(|source| {
633 if source.downcast_ref::<TransientRequest>().is_some() {
634 return true;
635 }
636 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
637 return false;
638 };
639 if error.is_timeout() || error.is_connect() || error.is_body() {
640 return true;
641 }
642 error.status().is_some_and(|status| {
643 let status = status.as_u16();
644 status == 408 || status == 429 || (500..600).contains(&status)
645 })
646 })
647}
648
649async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
650where
651 F: FnMut() -> Fut,
652 Fut: std::future::Future<Output = Result<T>>,
653{
654 let mut delays = retry_delays(retries);
655 let mut attempt = 1;
656 loop {
657 let started_at = Instant::now();
658 match operation().await {
659 Ok(value) => return Ok(value),
660 Err(error) if is_transient(&error) => {
661 let Some(delay) = delays.next() else {
662 return Err(error);
663 };
664 warn!(
665 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
666 started_at.elapsed()
667 );
668 tokio::time::sleep(delay).await;
669 attempt += 1;
670 }
671 Err(error) => return Err(error),
672 }
673 }
674}