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, PinnedFile, PinnedState, RecordedFileDigest, RestoreStats, WrapperSpan,
53 WrapperTiming, digest_file, is_task_identity, merge_task_action_predictions,
54 task_manifest_actions,
55};
56pub use client::BlockingAgentClient;
57pub use local::{LocalActionCache, LocalCas};
58pub use mbx_cache_protocol::{
59 ACTION_PROMISE_MEDIA_TYPE, ACTION_RESULT_BATCH_MEDIA_TYPE, ACTION_RESULT_MEDIA_TYPE,
60 ActionPrediction, ActionPromiseCompletion, ActionPromiseJoin, ActionPromiseState,
61 ActionResult as RemoteActionResult, BLOB_MEDIA_TYPE, BLOB_PACK_BLOBS_HEADER,
62 BLOB_PACK_BYTES_HEADER, BLOB_PACK_HEADER_BYTES, BLOB_PACK_MAGIC, BLOB_PACK_MEDIA_TYPE,
63 BLOB_PACK_RECEIPT_MEDIA_TYPE, CLIENT_METADATA_MEDIA_TYPE, Capabilities, CapabilityFeatures,
64 CapabilityLimits, CapabilityProtocol, CcMetadata, DIGEST_LIST_MEDIA_TYPE, DIRECTORY_MEDIA_TYPE,
65 Digest as CacheDigest, DigestAlgorithm, Directory as CacheDirectory,
66 DirectoryNode as CacheDirectoryNode, FileNode as CacheFileNode, MAX_ACTION_PREDICTION_PAYLOAD,
67 MAX_ACTION_PROMISE_CLAIM_BYTES, NAMESPACE_HEADER, PROTOCOL_HEADER, PROTOCOL_VERSION,
68 RustcMetadata, SymlinkNode as CacheSymlinkNode, TASK_ACTION_MANIFEST_MEDIA_TYPE,
69 TaskActionManifest,
70};
71pub use path_mapping::{
72 PathAliases, PathMapping, PathNormalizationError, normalize_mapped_path,
73 normalize_resolved_mapped_path, normalize_resolved_mapped_path_with, resolve_path_mappings,
74};
75use remote_http::HttpRemoteCache;
76#[cfg(feature = "fuzzing")]
77#[doc(hidden)]
78pub use remote_http::fuzz_decode_blob_pack;
79pub(crate) use remote_http::{BlobPackLimits, blob_pack_chunk};
80use remote_s3::S3RemoteCache;
81pub use remote_s3::{S3ConditionalWrites, S3RemoteCacheConfig};
82pub use sigv4::S3Credentials;
83const MAX_REMOTE_JSON_BYTES: u64 = 16 * 1024 * 1024;
90const MAX_ETAG_BYTES: usize = 256;
95const MAX_REMOTE_BLOB_BYTES: u64 = 5 * 1024 * 1024 * 1024;
98const MAX_STAGED_BLOB_PACK_BYTES: u64 = 256 * 1024 * 1024;
99const MAX_STAGED_BLOB_PACK_ITEMS: usize = 2 * 1024;
100const BLOB_PACK_TIMEOUT_BYTES_PER_UNIT: u64 = MAX_STAGED_BLOB_PACK_BYTES / 4;
101const BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT: usize = MAX_STAGED_BLOB_PACK_ITEMS / 4;
102const PACK_STREAM_CHUNK_BYTES: usize = 64 * 1024;
104const MAX_ACTION_BATCH_ITEMS: usize = 256;
110const MAX_ACTION_RESULT_BYTES: u64 = 64 * 1024;
112
113pub fn canonical_json(value: &impl Serialize) -> Result<Vec<u8>> {
118 Ok(mbx_cache_protocol::canonical_json(value)?)
119}
120
121#[derive(
122 Debug,
123 Clone,
124 Copy,
125 Serialize,
126 Deserialize,
127 Default,
128 strum::EnumString,
129 strum::Display,
130 PartialEq,
131 Eq,
132)]
133#[serde(rename_all = "kebab-case")]
134#[strum(serialize_all = "kebab-case")]
135pub enum RemoteCacheMode {
137 #[default]
139 ReadWrite,
140 ReadOnly,
142 WriteOnly,
144}
145
146impl RemoteCacheMode {
147 pub fn reads(self) -> bool {
149 matches!(self, Self::ReadWrite | Self::ReadOnly)
150 }
151
152 pub fn writes(self) -> bool {
154 matches!(self, Self::ReadWrite | Self::WriteOnly)
155 }
156}
157
158pub struct RemoteCacheConfig {
160 pub base_url: Url,
162 pub namespace: String,
164 pub token: Option<String>,
166 pub token_file: Option<PathBuf>,
168 pub oidc_audience: Option<String>,
170 pub connect_timeout: Duration,
172 pub read_timeout: Duration,
174 pub download_timeout: Duration,
183 pub retries: i64,
185}
186
187pub enum BlobSource {
189 Bytes(Vec<u8>),
191 File(tempfile::NamedTempFile),
193 Path(PathBuf),
195}
196
197pub struct BlobUpload {
199 pub digest: CacheDigest,
201 pub source: BlobSource,
203}
204
205pub struct RemoteActionManifest {
207 pub bytes: Vec<u8>,
209 pub etag: String,
211}
212
213pub struct RemoteBlobPack {
215 _directory: tempfile::TempDir,
216 pub blobs: Vec<(CacheDigest, PathBuf)>,
218 pub requests: u64,
220 pub requested: Vec<CacheDigest>,
222 pub blob_count: u64,
224 pub payload_bytes: u64,
226 pub framed_bytes: u64,
228}
229
230#[derive(Debug, Clone, Copy, Deserialize)]
232pub struct BlobPackReceipt {
233 #[serde(default)]
235 pub created: u64,
236 #[serde(default)]
238 pub existing: u64,
239}
240
241#[derive(Debug, Clone, Copy, PartialEq, Eq)]
242pub enum ManifestPutOutcome {
244 Stored,
246 PreconditionFailed,
248}
249
250pub struct RemoteCacheClient {
260 backend: Backend,
261}
262
263enum Backend {
269 Http(HttpRemoteCache),
270 S3(S3RemoteCache),
271}
272
273impl RemoteCacheClient {
274 pub fn new(config: RemoteCacheConfig) -> Result<Self> {
276 Ok(Self {
277 backend: Backend::Http(HttpRemoteCache::new(config)?),
278 })
279 }
280
281 pub fn with_read_stall_budget(mut self, budget: Duration) -> Self {
292 if let Backend::Http(client) = &mut self.backend {
293 client.set_read_stall_budget(budget);
294 }
295 self
296 }
297
298 pub fn new_s3(config: S3RemoteCacheConfig) -> Result<Self> {
303 Ok(Self {
304 backend: Backend::S3(S3RemoteCache::new(config)?),
305 })
306 }
307
308 pub async fn check_connection(&self) -> Result<()> {
314 match &self.backend {
315 Backend::Http(client) => client.check_connection().await,
316 Backend::S3(store) => store.check_connection().await,
317 }
318 }
319
320 pub async fn get_blob_pack(
326 &self,
327 digests: &[CacheDigest],
328 staging_dir: &Path,
329 ) -> Result<Option<RemoteBlobPack>> {
330 match &self.backend {
331 Backend::Http(client) => client.get_blob_pack(digests, staging_dir).await,
332 Backend::S3(store) => store.get_blob_pack(digests, staging_dir).await,
333 }
334 }
335
336 pub(crate) async fn get_blob_pack_with_limit(
337 &self,
338 digests: &[CacheDigest],
339 staging_dir: &Path,
340 max_bytes: u64,
341 ) -> Result<Option<RemoteBlobPack>> {
342 match &self.backend {
343 Backend::Http(client) => {
344 client
345 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
346 .await
347 }
348 Backend::S3(store) => {
349 store
350 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
351 .await
352 }
353 }
354 }
355
356 pub(crate) async fn blob_pack_limits(&self) -> Result<Option<BlobPackLimits>> {
358 match &self.backend {
359 Backend::Http(client) => client.blob_pack_limits().await,
360 Backend::S3(_) => Ok(None),
361 }
362 }
363
364 pub async fn get_action_result(
366 &self,
367 action: &CacheDigest,
368 ) -> Result<Option<RemoteActionResult>> {
369 match &self.backend {
370 Backend::Http(client) => client.get_action_result(action).await,
371 Backend::S3(store) => store.get_action_result(action).await,
372 }
373 }
374
375 pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
379 match &self.backend {
380 Backend::Http(client) => client.action_batch_limit().await,
381 Backend::S3(store) => store.action_batch_limit().await,
382 }
383 }
384
385 pub async fn get_action_results(
392 &self,
393 actions: &[CacheDigest],
394 ) -> Result<Option<Vec<RemoteActionResult>>> {
395 match &self.backend {
396 Backend::Http(client) => client.get_action_results(actions).await,
397 Backend::S3(store) => store.get_action_results(actions).await,
398 }
399 }
400
401 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
403 match &self.backend {
404 Backend::Http(client) => client.put_action_result(result).await,
405 Backend::S3(store) => store.put_action_result(result).await,
406 }
407 }
408
409 pub async fn join_action_promise(
413 &self,
414 invocation: &CacheDigest,
415 adapter: &str,
416 ) -> Result<Option<ActionPromiseState>> {
417 match &self.backend {
418 Backend::Http(client) => client.join_action_promise(invocation, adapter).await,
419 Backend::S3(_) => Ok(None),
420 }
421 }
422
423 pub async fn complete_action_promise(
427 &self,
428 invocation: &CacheDigest,
429 completion: &ActionPromiseCompletion,
430 ) -> Result<bool> {
431 match &self.backend {
432 Backend::Http(client) => client.complete_action_promise(invocation, completion).await,
433 Backend::S3(_) => Ok(false),
434 }
435 }
436
437 pub async fn get_action_manifest(
439 &self,
440 key: &CacheDigest,
441 ) -> Result<Option<RemoteActionManifest>> {
442 match &self.backend {
443 Backend::Http(client) => client.get_action_manifest(key).await,
444 Backend::S3(store) => store.get_action_manifest(key).await,
445 }
446 }
447
448 pub async fn put_action_manifest(
450 &self,
451 key: &CacheDigest,
452 bytes: &[u8],
453 expected_etag: Option<&str>,
454 ) -> Result<ManifestPutOutcome> {
455 match &self.backend {
456 Backend::Http(client) => client.put_action_manifest(key, bytes, expected_etag).await,
457 Backend::S3(store) => store.put_action_manifest(key, bytes, expected_etag).await,
458 }
459 }
460
461 pub async fn get_blob(
463 &self,
464 digest: &CacheDigest,
465 media_type: &'static str,
466 ) -> Result<Vec<u8>> {
467 match &self.backend {
468 Backend::Http(client) => client.get_blob(digest, media_type).await,
469 Backend::S3(store) => store.get_blob(digest, media_type).await,
470 }
471 }
472
473 pub async fn get_blob_file(
475 &self,
476 digest: &CacheDigest,
477 staging_dir: &Path,
478 ) -> Result<tempfile::NamedTempFile> {
479 match &self.backend {
480 Backend::Http(client) => client.get_blob_file(digest, staging_dir).await,
481 Backend::S3(store) => store.get_blob_file(digest, staging_dir).await,
482 }
483 }
484
485 pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<BlobPackLimits>> {
489 match &self.backend {
490 Backend::Http(client) => client.blob_pack_upload_limits().await,
491 Backend::S3(store) => store.blob_pack_upload_limits().await,
492 }
493 }
494
495 pub async fn put_blob_pack(&self, uploads: &[BlobUpload]) -> Result<Option<BlobPackReceipt>> {
503 match &self.backend {
504 Backend::Http(client) => client.put_blob_pack(uploads).await,
505 Backend::S3(store) => store.put_blob_pack(uploads).await,
506 }
507 }
508
509 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
511 match &self.backend {
512 Backend::Http(client) => client.put_blob(upload).await,
513 Backend::S3(store) => store.put_blob(upload).await,
514 }
515 }
516}
517
518async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
524 read_json_within(response, what, MAX_REMOTE_JSON_BYTES).await
525}
526
527async fn read_json_within(response: reqwest::Response, what: &str, limit: u64) -> Result<Vec<u8>> {
528 if let Some(length) = response.content_length()
529 && length > limit
530 {
531 bail!("remote cache {what} declared {length} bytes, over the {limit} byte limit");
532 }
533 let mut response = response;
534 let mut bytes = Vec::new();
535 while let Some(chunk) = response.chunk().await? {
536 if bytes.len() as u64 + chunk.len() as u64 > limit {
537 bail!("remote cache {what} exceeded the {limit} byte limit");
538 }
539 bytes.extend_from_slice(&chunk);
540 }
541 Ok(bytes)
542}
543fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
552 let value = value
553 .and_then(|value| value.to_str().ok())
554 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
555 if value.starts_with("W/") {
556 bail!("remote action manifest response has a weak ETag");
559 }
560 let etag = value
561 .strip_prefix('"')
562 .and_then(|value| value.strip_suffix('"'))
563 .filter(|value| is_entity_tag(value))
564 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
565 Ok(etag.to_owned())
566}
567
568fn quoted_etag(etag: &str) -> Result<HeaderValue> {
569 if !is_entity_tag(etag) {
570 bail!("invalid remote action manifest ETag");
571 }
572 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
573}
574
575fn is_entity_tag(value: &str) -> bool {
580 !value.is_empty()
581 && value.len() <= MAX_ETAG_BYTES
582 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
583}
584fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
585 [200u64, 1_000, 4_000, 15_000]
586 .into_iter()
587 .chain(std::iter::repeat(15_000))
588 .map(Duration::from_millis)
589 .map(|duration| {
590 let factor = 0.5 + rand::random::<f64>() * 0.5;
591 Duration::from_secs_f64(duration.as_secs_f64() * factor)
592 })
593 .take(retries.max(0) as usize)
594}
595
596fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
600 let mut current = Some(error);
601 while let Some(source) = current {
602 if source.to_string() == "dns error" {
603 return true;
604 }
605 current = source.source();
606 }
607 false
608}
609
610#[derive(Debug)]
617pub(crate) struct TransientRequest(pub(crate) &'static str);
618
619impl std::fmt::Display for TransientRequest {
620 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
621 formatter.write_str(self.0)
622 }
623}
624
625impl std::error::Error for TransientRequest {}
626
627fn is_transient(error: &eyre::Report) -> bool {
628 if is_dns_error(error.as_ref()) {
631 return false;
632 }
633 error.chain().any(|source| {
634 if source.downcast_ref::<TransientRequest>().is_some() {
635 return true;
636 }
637 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
638 return false;
639 };
640 if error.is_timeout() || error.is_connect() || error.is_body() {
641 return true;
642 }
643 error.status().is_some_and(|status| {
644 let status = status.as_u16();
645 status == 408 || status == 429 || (500..600).contains(&status)
646 })
647 })
648}
649
650async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
651where
652 F: FnMut() -> Fut,
653 Fut: std::future::Future<Output = Result<T>>,
654{
655 let mut delays = retry_delays(retries);
656 let mut attempt = 1;
657 loop {
658 let started_at = Instant::now();
659 match operation().await {
660 Ok(value) => return Ok(value),
661 Err(error) if is_transient(&error) => {
662 let Some(delay) = delays.next() else {
663 return Err(error);
664 };
665 warn!(
666 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
667 started_at.elapsed()
668 );
669 tokio::time::sleep(delay).await;
670 attempt += 1;
671 }
672 Err(error) => return Err(error),
673 }
674 }
675}