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 imds;
42mod local;
43mod path_mapping;
44mod remote_http;
45mod remote_s3;
46mod sigv4;
47mod uploads;
48
49pub use agent::{
50 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRemoteCache,
51 AgentRequest, AgentResponse, AgentStats, CacheAgent, CompilerStats, FileDigestCache,
52 FileDigestResolution, FileDigestScope, FileIdentity, FileObjectIdentity, FileSnapshot,
53 NoFileDigestCache, PinnedFile, PinnedState, RecordedFileDigest, RestoreStats,
54 ValidatedFileDigest, WrapperSpan, WrapperTiming, digest_file, digest_file_validated,
55 is_task_identity, merge_task_action_predictions, task_manifest_actions,
56};
57pub use client::BlockingAgentClient;
58pub use imds::{InstanceRoleCredentials, TemporaryCredentials};
59pub use local::{LocalActionCache, LocalCas};
60pub use mbx_cache_protocol::{
61 ACTION_PROMISE_MEDIA_TYPE, ACTION_RESULT_BATCH_MEDIA_TYPE, ACTION_RESULT_MEDIA_TYPE,
62 ActionPrediction, ActionPromiseCompletion, ActionPromiseJoin, ActionPromiseState,
63 ActionResult as RemoteActionResult, BLOB_MEDIA_TYPE, BLOB_PACK_BLOBS_HEADER,
64 BLOB_PACK_BYTES_HEADER, BLOB_PACK_HEADER_BYTES, BLOB_PACK_MAGIC, BLOB_PACK_MEDIA_TYPE,
65 BLOB_PACK_RECEIPT_MEDIA_TYPE, CLIENT_METADATA_MEDIA_TYPE, Capabilities, CapabilityFeatures,
66 CapabilityLimits, CapabilityProtocol, CcMetadata, DIGEST_LIST_MEDIA_TYPE, DIRECTORY_MEDIA_TYPE,
67 Digest as CacheDigest, DigestAlgorithm, Directory as CacheDirectory,
68 DirectoryNode as CacheDirectoryNode, FileNode as CacheFileNode, MAX_ACTION_PREDICTION_PAYLOAD,
69 MAX_ACTION_PROMISE_CLAIM_BYTES, NAMESPACE_HEADER, PROTOCOL_HEADER, PROTOCOL_VERSION,
70 RustcMetadata, SymlinkNode as CacheSymlinkNode, TASK_ACTION_MANIFEST_MEDIA_TYPE,
71 TaskActionManifest,
72};
73pub use path_mapping::{
74 PathAliases, PathMapping, PathNormalizationError, normalize_mapped_path,
75 normalize_resolved_mapped_path, normalize_resolved_mapped_path_with, resolve_path_mappings,
76};
77use remote_http::HttpRemoteCache;
78#[cfg(feature = "fuzzing")]
79#[doc(hidden)]
80pub use remote_http::fuzz_decode_blob_pack;
81pub(crate) use remote_http::{BlobPackLimits, blob_pack_chunk};
82use remote_s3::S3RemoteCache;
83pub use remote_s3::{S3ConditionalWrites, S3RemoteCacheConfig};
84pub use sigv4::S3Credentials;
85const MAX_REMOTE_JSON_BYTES: u64 = 16 * 1024 * 1024;
92const MAX_ETAG_BYTES: usize = 256;
97const MAX_REMOTE_BLOB_BYTES: u64 = 5 * 1024 * 1024 * 1024;
100const MAX_STAGED_BLOB_PACK_BYTES: u64 = 256 * 1024 * 1024;
101const MAX_STAGED_BLOB_PACK_ITEMS: usize = 2 * 1024;
102const BLOB_PACK_TIMEOUT_BYTES_PER_UNIT: u64 = MAX_STAGED_BLOB_PACK_BYTES / 4;
103const BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT: usize = MAX_STAGED_BLOB_PACK_ITEMS / 4;
104const PACK_STREAM_CHUNK_BYTES: usize = 64 * 1024;
106const MAX_ACTION_BATCH_ITEMS: usize = 256;
112const MAX_ACTION_RESULT_BYTES: u64 = 64 * 1024;
114
115pub fn canonical_json(value: &impl Serialize) -> Result<Vec<u8>> {
120 Ok(mbx_cache_protocol::canonical_json(value)?)
121}
122
123#[derive(
124 Debug,
125 Clone,
126 Copy,
127 Serialize,
128 Deserialize,
129 Default,
130 strum::EnumString,
131 strum::Display,
132 PartialEq,
133 Eq,
134)]
135#[serde(rename_all = "kebab-case")]
136#[strum(serialize_all = "kebab-case")]
137pub enum RemoteCacheMode {
139 #[default]
141 ReadWrite,
142 ReadOnly,
144 WriteOnly,
146}
147
148impl RemoteCacheMode {
149 pub fn reads(self) -> bool {
151 matches!(self, Self::ReadWrite | Self::ReadOnly)
152 }
153
154 pub fn writes(self) -> bool {
156 matches!(self, Self::ReadWrite | Self::WriteOnly)
157 }
158}
159
160pub struct RemoteCacheConfig {
162 pub base_url: Url,
164 pub namespace: String,
166 pub token: Option<String>,
168 pub token_file: Option<PathBuf>,
170 pub oidc_audience: Option<String>,
172 pub connect_timeout: Duration,
174 pub read_timeout: Duration,
176 pub download_timeout: Duration,
185 pub retries: i64,
187}
188
189pub enum BlobSource {
191 Bytes(Vec<u8>),
193 File(tempfile::NamedTempFile),
195 Path(PathBuf),
197}
198
199pub struct BlobUpload {
201 pub digest: CacheDigest,
203 pub source: BlobSource,
205}
206
207pub struct RemoteActionManifest {
209 pub bytes: Vec<u8>,
211 pub etag: String,
213}
214
215pub struct RemoteBlobPack {
217 _directory: tempfile::TempDir,
218 pub blobs: Vec<(CacheDigest, PathBuf)>,
220 pub requests: u64,
222 pub requested: Vec<CacheDigest>,
224 pub blob_count: u64,
226 pub payload_bytes: u64,
228 pub framed_bytes: u64,
230}
231
232#[derive(Debug, Clone, Copy, Deserialize)]
234pub struct BlobPackReceipt {
235 #[serde(default)]
237 pub created: u64,
238 #[serde(default)]
240 pub existing: u64,
241}
242
243#[derive(Debug, Clone, Copy, PartialEq, Eq)]
244pub enum ManifestPutOutcome {
246 Stored,
248 PreconditionFailed,
250}
251
252pub struct RemoteCacheClient {
262 backend: Backend,
263}
264
265enum Backend {
271 Http(HttpRemoteCache),
272 S3(S3RemoteCache),
273}
274
275impl RemoteCacheClient {
276 pub fn new(config: RemoteCacheConfig) -> Result<Self> {
278 Ok(Self {
279 backend: Backend::Http(HttpRemoteCache::new(config)?),
280 })
281 }
282
283 pub fn with_read_stall_budget(mut self, budget: Duration) -> Self {
294 if let Backend::Http(client) = &mut self.backend {
295 client.set_read_stall_budget(budget);
296 }
297 self
298 }
299
300 pub fn new_s3(config: S3RemoteCacheConfig) -> Result<Self> {
305 Ok(Self {
306 backend: Backend::S3(S3RemoteCache::new(config)?),
307 })
308 }
309
310 pub fn new_s3_with_instance_role(
318 config: S3RemoteCacheConfig,
319 provider: InstanceRoleCredentials,
320 expires_at: std::time::SystemTime,
321 ) -> Result<Self> {
322 Ok(Self {
323 backend: Backend::S3(
324 S3RemoteCache::new(config)?.with_instance_role(provider, expires_at),
325 ),
326 })
327 }
328
329 pub async fn check_connection(&self) -> Result<()> {
335 match &self.backend {
336 Backend::Http(client) => client.check_connection().await,
337 Backend::S3(store) => store.check_connection().await,
338 }
339 }
340
341 pub async fn get_blob_pack(
347 &self,
348 digests: &[CacheDigest],
349 staging_dir: &Path,
350 ) -> Result<Option<RemoteBlobPack>> {
351 match &self.backend {
352 Backend::Http(client) => client.get_blob_pack(digests, staging_dir).await,
353 Backend::S3(store) => store.get_blob_pack(digests, staging_dir).await,
354 }
355 }
356
357 pub(crate) async fn get_blob_pack_with_limit(
358 &self,
359 digests: &[CacheDigest],
360 staging_dir: &Path,
361 max_bytes: u64,
362 ) -> Result<Option<RemoteBlobPack>> {
363 match &self.backend {
364 Backend::Http(client) => {
365 client
366 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
367 .await
368 }
369 Backend::S3(store) => {
370 store
371 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
372 .await
373 }
374 }
375 }
376
377 pub(crate) async fn blob_pack_limits(&self) -> Result<Option<BlobPackLimits>> {
379 match &self.backend {
380 Backend::Http(client) => client.blob_pack_limits().await,
381 Backend::S3(_) => Ok(None),
382 }
383 }
384
385 pub async fn get_action_result(
387 &self,
388 action: &CacheDigest,
389 ) -> Result<Option<RemoteActionResult>> {
390 match &self.backend {
391 Backend::Http(client) => client.get_action_result(action).await,
392 Backend::S3(store) => store.get_action_result(action).await,
393 }
394 }
395
396 pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
400 match &self.backend {
401 Backend::Http(client) => client.action_batch_limit().await,
402 Backend::S3(store) => store.action_batch_limit().await,
403 }
404 }
405
406 pub async fn get_action_results(
413 &self,
414 actions: &[CacheDigest],
415 ) -> Result<Option<Vec<RemoteActionResult>>> {
416 match &self.backend {
417 Backend::Http(client) => client.get_action_results(actions).await,
418 Backend::S3(store) => store.get_action_results(actions).await,
419 }
420 }
421
422 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
424 match &self.backend {
425 Backend::Http(client) => client.put_action_result(result).await,
426 Backend::S3(store) => store.put_action_result(result).await,
427 }
428 }
429
430 pub async fn join_action_promise(
434 &self,
435 invocation: &CacheDigest,
436 adapter: &str,
437 ) -> Result<Option<ActionPromiseState>> {
438 match &self.backend {
439 Backend::Http(client) => client.join_action_promise(invocation, adapter).await,
440 Backend::S3(_) => Ok(None),
441 }
442 }
443
444 pub async fn complete_action_promise(
448 &self,
449 invocation: &CacheDigest,
450 completion: &ActionPromiseCompletion,
451 ) -> Result<bool> {
452 match &self.backend {
453 Backend::Http(client) => client.complete_action_promise(invocation, completion).await,
454 Backend::S3(_) => Ok(false),
455 }
456 }
457
458 pub async fn get_action_manifest(
460 &self,
461 key: &CacheDigest,
462 ) -> Result<Option<RemoteActionManifest>> {
463 match &self.backend {
464 Backend::Http(client) => client.get_action_manifest(key).await,
465 Backend::S3(store) => store.get_action_manifest(key).await,
466 }
467 }
468
469 pub async fn put_action_manifest(
471 &self,
472 key: &CacheDigest,
473 bytes: &[u8],
474 expected_etag: Option<&str>,
475 ) -> Result<ManifestPutOutcome> {
476 match &self.backend {
477 Backend::Http(client) => client.put_action_manifest(key, bytes, expected_etag).await,
478 Backend::S3(store) => store.put_action_manifest(key, bytes, expected_etag).await,
479 }
480 }
481
482 pub async fn get_blob(
484 &self,
485 digest: &CacheDigest,
486 media_type: &'static str,
487 ) -> Result<Vec<u8>> {
488 match &self.backend {
489 Backend::Http(client) => client.get_blob(digest, media_type).await,
490 Backend::S3(store) => store.get_blob(digest, media_type).await,
491 }
492 }
493
494 pub async fn get_blob_file(
496 &self,
497 digest: &CacheDigest,
498 staging_dir: &Path,
499 ) -> Result<tempfile::NamedTempFile> {
500 match &self.backend {
501 Backend::Http(client) => client.get_blob_file(digest, staging_dir).await,
502 Backend::S3(store) => store.get_blob_file(digest, staging_dir).await,
503 }
504 }
505
506 pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<BlobPackLimits>> {
510 match &self.backend {
511 Backend::Http(client) => client.blob_pack_upload_limits().await,
512 Backend::S3(store) => store.blob_pack_upload_limits().await,
513 }
514 }
515
516 pub async fn put_blob_pack(&self, uploads: &[BlobUpload]) -> Result<Option<BlobPackReceipt>> {
524 match &self.backend {
525 Backend::Http(client) => client.put_blob_pack(uploads).await,
526 Backend::S3(store) => store.put_blob_pack(uploads).await,
527 }
528 }
529
530 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
532 match &self.backend {
533 Backend::Http(client) => client.put_blob(upload).await,
534 Backend::S3(store) => store.put_blob(upload).await,
535 }
536 }
537}
538
539async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
545 read_json_within(response, what, MAX_REMOTE_JSON_BYTES).await
546}
547
548async fn read_json_within(response: reqwest::Response, what: &str, limit: u64) -> Result<Vec<u8>> {
549 if let Some(length) = response.content_length()
550 && length > limit
551 {
552 bail!("remote cache {what} declared {length} bytes, over the {limit} byte limit");
553 }
554 let mut response = response;
555 let mut bytes = Vec::new();
556 while let Some(chunk) = response.chunk().await? {
557 if bytes.len() as u64 + chunk.len() as u64 > limit {
558 bail!("remote cache {what} exceeded the {limit} byte limit");
559 }
560 bytes.extend_from_slice(&chunk);
561 }
562 Ok(bytes)
563}
564fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
573 let value = value
574 .and_then(|value| value.to_str().ok())
575 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
576 if value.starts_with("W/") {
577 bail!("remote action manifest response has a weak ETag");
580 }
581 let etag = value
582 .strip_prefix('"')
583 .and_then(|value| value.strip_suffix('"'))
584 .filter(|value| is_entity_tag(value))
585 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
586 Ok(etag.to_owned())
587}
588
589fn quoted_etag(etag: &str) -> Result<HeaderValue> {
590 if !is_entity_tag(etag) {
591 bail!("invalid remote action manifest ETag");
592 }
593 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
594}
595
596fn is_entity_tag(value: &str) -> bool {
601 !value.is_empty()
602 && value.len() <= MAX_ETAG_BYTES
603 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
604}
605fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
606 [200u64, 1_000, 4_000, 15_000]
607 .into_iter()
608 .chain(std::iter::repeat(15_000))
609 .map(Duration::from_millis)
610 .map(|duration| {
611 let factor = 0.5 + rand::random::<f64>() * 0.5;
612 Duration::from_secs_f64(duration.as_secs_f64() * factor)
613 })
614 .take(retries.max(0) as usize)
615}
616
617fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
621 let mut current = Some(error);
622 while let Some(source) = current {
623 if source.to_string() == "dns error" {
624 return true;
625 }
626 current = source.source();
627 }
628 false
629}
630
631#[derive(Debug)]
638pub(crate) struct TransientRequest(pub(crate) &'static str);
639
640impl std::fmt::Display for TransientRequest {
641 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
642 formatter.write_str(self.0)
643 }
644}
645
646impl std::error::Error for TransientRequest {}
647
648fn is_transient(error: &eyre::Report) -> bool {
649 if is_dns_error(error.as_ref()) {
652 return false;
653 }
654 error.chain().any(|source| {
655 if source.downcast_ref::<TransientRequest>().is_some() {
656 return true;
657 }
658 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
659 return false;
660 };
661 if error.is_timeout() || error.is_connect() || error.is_body() {
662 return true;
663 }
664 error.status().is_some_and(|status| {
665 let status = status.as_u16();
666 status == 408 || status == 429 || (500..600).contains(&status)
667 })
668 })
669}
670
671async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
672where
673 F: FnMut() -> Fut,
674 Fut: std::future::Future<Output = Result<T>>,
675{
676 let mut delays = retry_delays(retries);
677 let mut attempt = 1;
678 loop {
679 let started_at = Instant::now();
680 match operation().await {
681 Ok(value) => return Ok(value),
682 Err(error) if is_transient(&error) => {
683 let Some(delay) = delays.next() else {
684 return Err(error);
685 };
686 warn!(
687 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
688 started_at.elapsed()
689 );
690 tokio::time::sleep(delay).await;
691 attempt += 1;
692 }
693 Err(error) => return Err(error),
694 }
695 }
696}