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