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 FileDigestScope, FileIdentity, NoFileDigestCache, RecordedFileDigest, RestoreStats,
52 is_task_identity, task_manifest_actions,
53};
54pub use client::BlockingAgentClient;
55pub use local::{LocalActionCache, LocalCas};
56pub use mbx_cache_protocol::{
57 ACTION_PROMISE_MEDIA_TYPE, ACTION_RESULT_BATCH_MEDIA_TYPE, ACTION_RESULT_MEDIA_TYPE,
58 ActionPrediction, ActionPromiseCompletion, ActionPromiseJoin, ActionPromiseState,
59 ActionResult as RemoteActionResult, BLOB_MEDIA_TYPE, BLOB_PACK_BLOBS_HEADER,
60 BLOB_PACK_BYTES_HEADER, BLOB_PACK_HEADER_BYTES, BLOB_PACK_MAGIC, BLOB_PACK_MEDIA_TYPE,
61 BLOB_PACK_RECEIPT_MEDIA_TYPE, CLIENT_METADATA_MEDIA_TYPE, Capabilities, CapabilityFeatures,
62 CapabilityLimits, CapabilityProtocol, CcMetadata, DIGEST_LIST_MEDIA_TYPE, DIRECTORY_MEDIA_TYPE,
63 Digest as CacheDigest, DigestAlgorithm, Directory as CacheDirectory,
64 DirectoryNode as CacheDirectoryNode, FileNode as CacheFileNode, MAX_ACTION_PREDICTION_PAYLOAD,
65 MAX_ACTION_PROMISE_CLAIM_BYTES, NAMESPACE_HEADER, PROTOCOL_HEADER, PROTOCOL_VERSION,
66 RustcMetadata, SymlinkNode as CacheSymlinkNode, TASK_ACTION_MANIFEST_MEDIA_TYPE,
67 TaskActionManifest,
68};
69pub use path_mapping::{
70 PathMapping, PathNormalizationError, normalize_mapped_path, normalize_resolved_mapped_path,
71 resolve_path_mappings,
72};
73use remote_http::HttpRemoteCache;
74#[cfg(feature = "fuzzing")]
75#[doc(hidden)]
76pub use remote_http::fuzz_decode_blob_pack;
77pub(crate) use remote_http::{BlobPackLimits, blob_pack_chunk};
78use remote_s3::S3RemoteCache;
79pub use remote_s3::{S3ConditionalWrites, S3RemoteCacheConfig};
80pub use sigv4::S3Credentials;
81const MAX_REMOTE_JSON_BYTES: u64 = 16 * 1024 * 1024;
88const MAX_ETAG_BYTES: usize = 256;
93const MAX_REMOTE_BLOB_BYTES: u64 = 5 * 1024 * 1024 * 1024;
96const MAX_STAGED_BLOB_PACK_BYTES: u64 = 256 * 1024 * 1024;
97const MAX_STAGED_BLOB_PACK_ITEMS: usize = 2 * 1024;
98const BLOB_PACK_TIMEOUT_BYTES_PER_UNIT: u64 = MAX_STAGED_BLOB_PACK_BYTES / 4;
99const BLOB_PACK_TIMEOUT_ITEMS_PER_UNIT: usize = MAX_STAGED_BLOB_PACK_ITEMS / 4;
100const PACK_STREAM_CHUNK_BYTES: usize = 64 * 1024;
102const MAX_ACTION_BATCH_ITEMS: usize = 256;
108const MAX_ACTION_RESULT_BYTES: u64 = 64 * 1024;
110
111pub fn canonical_json(value: &impl Serialize) -> Result<Vec<u8>> {
116 Ok(mbx_cache_protocol::canonical_json(value)?)
117}
118
119#[derive(
120 Debug,
121 Clone,
122 Copy,
123 Serialize,
124 Deserialize,
125 Default,
126 strum::EnumString,
127 strum::Display,
128 PartialEq,
129 Eq,
130)]
131#[serde(rename_all = "kebab-case")]
132#[strum(serialize_all = "kebab-case")]
133pub enum RemoteCacheMode {
135 #[default]
137 ReadWrite,
138 ReadOnly,
140 WriteOnly,
142}
143
144impl RemoteCacheMode {
145 pub fn reads(self) -> bool {
147 matches!(self, Self::ReadWrite | Self::ReadOnly)
148 }
149
150 pub fn writes(self) -> bool {
152 matches!(self, Self::ReadWrite | Self::WriteOnly)
153 }
154}
155
156pub struct RemoteCacheConfig {
158 pub base_url: Url,
160 pub namespace: String,
162 pub token: Option<String>,
164 pub token_file: Option<PathBuf>,
166 pub oidc_audience: Option<String>,
168 pub connect_timeout: Duration,
170 pub read_timeout: Duration,
172 pub download_timeout: Duration,
181 pub retries: i64,
183}
184
185pub enum BlobSource {
187 Bytes(Vec<u8>),
189 File(tempfile::NamedTempFile),
191 Path(PathBuf),
193}
194
195pub struct BlobUpload {
197 pub digest: CacheDigest,
199 pub source: BlobSource,
201}
202
203pub struct RemoteActionManifest {
205 pub bytes: Vec<u8>,
207 pub etag: String,
209}
210
211pub struct RemoteBlobPack {
213 _directory: tempfile::TempDir,
214 pub blobs: Vec<(CacheDigest, PathBuf)>,
216 pub requests: u64,
218 pub requested: Vec<CacheDigest>,
220 pub blob_count: u64,
222 pub payload_bytes: u64,
224 pub framed_bytes: u64,
226}
227
228#[derive(Debug, Clone, Copy, Deserialize)]
230pub struct BlobPackReceipt {
231 #[serde(default)]
233 pub created: u64,
234 #[serde(default)]
236 pub existing: u64,
237}
238
239#[derive(Debug, Clone, Copy, PartialEq, Eq)]
240pub enum ManifestPutOutcome {
242 Stored,
244 PreconditionFailed,
246}
247
248pub struct RemoteCacheClient {
258 backend: Backend,
259}
260
261enum Backend {
267 Http(HttpRemoteCache),
268 S3(S3RemoteCache),
269}
270
271impl RemoteCacheClient {
272 pub fn new(config: RemoteCacheConfig) -> Result<Self> {
274 Ok(Self {
275 backend: Backend::Http(HttpRemoteCache::new(config)?),
276 })
277 }
278
279 pub fn with_read_stall_budget(mut self, budget: Duration) -> Self {
290 if let Backend::Http(client) = &mut self.backend {
291 client.set_read_stall_budget(budget);
292 }
293 self
294 }
295
296 pub fn new_s3(config: S3RemoteCacheConfig) -> Result<Self> {
301 Ok(Self {
302 backend: Backend::S3(S3RemoteCache::new(config)?),
303 })
304 }
305
306 pub async fn check_connection(&self) -> Result<()> {
312 match &self.backend {
313 Backend::Http(client) => client.check_connection().await,
314 Backend::S3(store) => store.check_connection().await,
315 }
316 }
317
318 pub async fn get_blob_pack(
324 &self,
325 digests: &[CacheDigest],
326 staging_dir: &Path,
327 ) -> Result<Option<RemoteBlobPack>> {
328 match &self.backend {
329 Backend::Http(client) => client.get_blob_pack(digests, staging_dir).await,
330 Backend::S3(store) => store.get_blob_pack(digests, staging_dir).await,
331 }
332 }
333
334 pub(crate) async fn get_blob_pack_with_limit(
335 &self,
336 digests: &[CacheDigest],
337 staging_dir: &Path,
338 max_bytes: u64,
339 ) -> Result<Option<RemoteBlobPack>> {
340 match &self.backend {
341 Backend::Http(client) => {
342 client
343 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
344 .await
345 }
346 Backend::S3(store) => {
347 store
348 .get_blob_pack_with_limit(digests, staging_dir, max_bytes)
349 .await
350 }
351 }
352 }
353
354 pub async fn get_action_result(
356 &self,
357 action: &CacheDigest,
358 ) -> Result<Option<RemoteActionResult>> {
359 match &self.backend {
360 Backend::Http(client) => client.get_action_result(action).await,
361 Backend::S3(store) => store.get_action_result(action).await,
362 }
363 }
364
365 pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
369 match &self.backend {
370 Backend::Http(client) => client.action_batch_limit().await,
371 Backend::S3(store) => store.action_batch_limit().await,
372 }
373 }
374
375 pub async fn get_action_results(
382 &self,
383 actions: &[CacheDigest],
384 ) -> Result<Option<Vec<RemoteActionResult>>> {
385 match &self.backend {
386 Backend::Http(client) => client.get_action_results(actions).await,
387 Backend::S3(store) => store.get_action_results(actions).await,
388 }
389 }
390
391 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
393 match &self.backend {
394 Backend::Http(client) => client.put_action_result(result).await,
395 Backend::S3(store) => store.put_action_result(result).await,
396 }
397 }
398
399 pub async fn join_action_promise(
403 &self,
404 invocation: &CacheDigest,
405 adapter: &str,
406 ) -> Result<Option<ActionPromiseState>> {
407 match &self.backend {
408 Backend::Http(client) => client.join_action_promise(invocation, adapter).await,
409 Backend::S3(_) => Ok(None),
410 }
411 }
412
413 pub async fn complete_action_promise(
417 &self,
418 invocation: &CacheDigest,
419 completion: &ActionPromiseCompletion,
420 ) -> Result<bool> {
421 match &self.backend {
422 Backend::Http(client) => client.complete_action_promise(invocation, completion).await,
423 Backend::S3(_) => Ok(false),
424 }
425 }
426
427 pub async fn get_action_manifest(
429 &self,
430 key: &CacheDigest,
431 ) -> Result<Option<RemoteActionManifest>> {
432 match &self.backend {
433 Backend::Http(client) => client.get_action_manifest(key).await,
434 Backend::S3(store) => store.get_action_manifest(key).await,
435 }
436 }
437
438 pub async fn put_action_manifest(
440 &self,
441 key: &CacheDigest,
442 bytes: &[u8],
443 expected_etag: Option<&str>,
444 ) -> Result<ManifestPutOutcome> {
445 match &self.backend {
446 Backend::Http(client) => client.put_action_manifest(key, bytes, expected_etag).await,
447 Backend::S3(store) => store.put_action_manifest(key, bytes, expected_etag).await,
448 }
449 }
450
451 pub async fn get_blob(
453 &self,
454 digest: &CacheDigest,
455 media_type: &'static str,
456 ) -> Result<Vec<u8>> {
457 match &self.backend {
458 Backend::Http(client) => client.get_blob(digest, media_type).await,
459 Backend::S3(store) => store.get_blob(digest, media_type).await,
460 }
461 }
462
463 pub async fn get_blob_file(
465 &self,
466 digest: &CacheDigest,
467 staging_dir: &Path,
468 ) -> Result<tempfile::NamedTempFile> {
469 match &self.backend {
470 Backend::Http(client) => client.get_blob_file(digest, staging_dir).await,
471 Backend::S3(store) => store.get_blob_file(digest, staging_dir).await,
472 }
473 }
474
475 pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<BlobPackLimits>> {
479 match &self.backend {
480 Backend::Http(client) => client.blob_pack_upload_limits().await,
481 Backend::S3(store) => store.blob_pack_upload_limits().await,
482 }
483 }
484
485 pub async fn put_blob_pack(&self, uploads: &[BlobUpload]) -> Result<Option<BlobPackReceipt>> {
493 match &self.backend {
494 Backend::Http(client) => client.put_blob_pack(uploads).await,
495 Backend::S3(store) => store.put_blob_pack(uploads).await,
496 }
497 }
498
499 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
501 match &self.backend {
502 Backend::Http(client) => client.put_blob(upload).await,
503 Backend::S3(store) => store.put_blob(upload).await,
504 }
505 }
506}
507
508async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
514 read_json_within(response, what, MAX_REMOTE_JSON_BYTES).await
515}
516
517async fn read_json_within(response: reqwest::Response, what: &str, limit: u64) -> Result<Vec<u8>> {
518 if let Some(length) = response.content_length()
519 && length > limit
520 {
521 bail!("remote cache {what} declared {length} bytes, over the {limit} byte limit");
522 }
523 let mut response = response;
524 let mut bytes = Vec::new();
525 while let Some(chunk) = response.chunk().await? {
526 if bytes.len() as u64 + chunk.len() as u64 > limit {
527 bail!("remote cache {what} exceeded the {limit} byte limit");
528 }
529 bytes.extend_from_slice(&chunk);
530 }
531 Ok(bytes)
532}
533fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
542 let value = value
543 .and_then(|value| value.to_str().ok())
544 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
545 if value.starts_with("W/") {
546 bail!("remote action manifest response has a weak ETag");
549 }
550 let etag = value
551 .strip_prefix('"')
552 .and_then(|value| value.strip_suffix('"'))
553 .filter(|value| is_entity_tag(value))
554 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
555 Ok(etag.to_owned())
556}
557
558fn quoted_etag(etag: &str) -> Result<HeaderValue> {
559 if !is_entity_tag(etag) {
560 bail!("invalid remote action manifest ETag");
561 }
562 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
563}
564
565fn is_entity_tag(value: &str) -> bool {
570 !value.is_empty()
571 && value.len() <= MAX_ETAG_BYTES
572 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
573}
574fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
575 [200u64, 1_000, 4_000, 15_000]
576 .into_iter()
577 .chain(std::iter::repeat(15_000))
578 .map(Duration::from_millis)
579 .map(|duration| {
580 let factor = 0.5 + rand::random::<f64>() * 0.5;
581 Duration::from_secs_f64(duration.as_secs_f64() * factor)
582 })
583 .take(retries.max(0) as usize)
584}
585
586fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
590 let mut current = Some(error);
591 while let Some(source) = current {
592 if source.to_string() == "dns error" {
593 return true;
594 }
595 current = source.source();
596 }
597 false
598}
599
600#[derive(Debug)]
607pub(crate) struct TransientRequest(pub(crate) &'static str);
608
609impl std::fmt::Display for TransientRequest {
610 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
611 formatter.write_str(self.0)
612 }
613}
614
615impl std::error::Error for TransientRequest {}
616
617fn is_transient(error: &eyre::Report) -> bool {
618 if is_dns_error(error.as_ref()) {
621 return false;
622 }
623 error.chain().any(|source| {
624 if source.downcast_ref::<TransientRequest>().is_some() {
625 return true;
626 }
627 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
628 return false;
629 };
630 if error.is_timeout() || error.is_connect() || error.is_body() {
631 return true;
632 }
633 error.status().is_some_and(|status| {
634 let status = status.as_u16();
635 status == 408 || status == 429 || (500..600).contains(&status)
636 })
637 })
638}
639
640async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
641where
642 F: FnMut() -> Fut,
643 Fut: std::future::Future<Output = Result<T>>,
644{
645 let mut delays = retry_delays(retries);
646 let mut attempt = 1;
647 loop {
648 let started_at = Instant::now();
649 match operation().await {
650 Ok(value) => return Ok(value),
651 Err(error) if is_transient(&error) => {
652 let Some(delay) = delays.next() else {
653 return Err(error);
654 };
655 warn!(
656 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
657 started_at.elapsed()
658 );
659 tokio::time::sleep(delay).await;
660 attempt += 1;
661 }
662 Err(error) => return Err(error),
663 }
664 }
665}