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(crate) async fn blob_pack_limits(&self) -> Result<Option<BlobPackLimits>> {
356 match &self.backend {
357 Backend::Http(client) => client.blob_pack_limits().await,
358 Backend::S3(_) => Ok(None),
359 }
360 }
361
362 pub async fn get_action_result(
364 &self,
365 action: &CacheDigest,
366 ) -> Result<Option<RemoteActionResult>> {
367 match &self.backend {
368 Backend::Http(client) => client.get_action_result(action).await,
369 Backend::S3(store) => store.get_action_result(action).await,
370 }
371 }
372
373 pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
377 match &self.backend {
378 Backend::Http(client) => client.action_batch_limit().await,
379 Backend::S3(store) => store.action_batch_limit().await,
380 }
381 }
382
383 pub async fn get_action_results(
390 &self,
391 actions: &[CacheDigest],
392 ) -> Result<Option<Vec<RemoteActionResult>>> {
393 match &self.backend {
394 Backend::Http(client) => client.get_action_results(actions).await,
395 Backend::S3(store) => store.get_action_results(actions).await,
396 }
397 }
398
399 pub async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
401 match &self.backend {
402 Backend::Http(client) => client.put_action_result(result).await,
403 Backend::S3(store) => store.put_action_result(result).await,
404 }
405 }
406
407 pub async fn join_action_promise(
411 &self,
412 invocation: &CacheDigest,
413 adapter: &str,
414 ) -> Result<Option<ActionPromiseState>> {
415 match &self.backend {
416 Backend::Http(client) => client.join_action_promise(invocation, adapter).await,
417 Backend::S3(_) => Ok(None),
418 }
419 }
420
421 pub async fn complete_action_promise(
425 &self,
426 invocation: &CacheDigest,
427 completion: &ActionPromiseCompletion,
428 ) -> Result<bool> {
429 match &self.backend {
430 Backend::Http(client) => client.complete_action_promise(invocation, completion).await,
431 Backend::S3(_) => Ok(false),
432 }
433 }
434
435 pub async fn get_action_manifest(
437 &self,
438 key: &CacheDigest,
439 ) -> Result<Option<RemoteActionManifest>> {
440 match &self.backend {
441 Backend::Http(client) => client.get_action_manifest(key).await,
442 Backend::S3(store) => store.get_action_manifest(key).await,
443 }
444 }
445
446 pub async fn put_action_manifest(
448 &self,
449 key: &CacheDigest,
450 bytes: &[u8],
451 expected_etag: Option<&str>,
452 ) -> Result<ManifestPutOutcome> {
453 match &self.backend {
454 Backend::Http(client) => client.put_action_manifest(key, bytes, expected_etag).await,
455 Backend::S3(store) => store.put_action_manifest(key, bytes, expected_etag).await,
456 }
457 }
458
459 pub async fn get_blob(
461 &self,
462 digest: &CacheDigest,
463 media_type: &'static str,
464 ) -> Result<Vec<u8>> {
465 match &self.backend {
466 Backend::Http(client) => client.get_blob(digest, media_type).await,
467 Backend::S3(store) => store.get_blob(digest, media_type).await,
468 }
469 }
470
471 pub async fn get_blob_file(
473 &self,
474 digest: &CacheDigest,
475 staging_dir: &Path,
476 ) -> Result<tempfile::NamedTempFile> {
477 match &self.backend {
478 Backend::Http(client) => client.get_blob_file(digest, staging_dir).await,
479 Backend::S3(store) => store.get_blob_file(digest, staging_dir).await,
480 }
481 }
482
483 pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<BlobPackLimits>> {
487 match &self.backend {
488 Backend::Http(client) => client.blob_pack_upload_limits().await,
489 Backend::S3(store) => store.blob_pack_upload_limits().await,
490 }
491 }
492
493 pub async fn put_blob_pack(&self, uploads: &[BlobUpload]) -> Result<Option<BlobPackReceipt>> {
501 match &self.backend {
502 Backend::Http(client) => client.put_blob_pack(uploads).await,
503 Backend::S3(store) => store.put_blob_pack(uploads).await,
504 }
505 }
506
507 pub async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
509 match &self.backend {
510 Backend::Http(client) => client.put_blob(upload).await,
511 Backend::S3(store) => store.put_blob(upload).await,
512 }
513 }
514}
515
516async fn read_bounded_json(response: reqwest::Response, what: &str) -> Result<Vec<u8>> {
522 read_json_within(response, what, MAX_REMOTE_JSON_BYTES).await
523}
524
525async fn read_json_within(response: reqwest::Response, what: &str, limit: u64) -> Result<Vec<u8>> {
526 if let Some(length) = response.content_length()
527 && length > limit
528 {
529 bail!("remote cache {what} declared {length} bytes, over the {limit} byte limit");
530 }
531 let mut response = response;
532 let mut bytes = Vec::new();
533 while let Some(chunk) = response.chunk().await? {
534 if bytes.len() as u64 + chunk.len() as u64 > limit {
535 bail!("remote cache {what} exceeded the {limit} byte limit");
536 }
537 bytes.extend_from_slice(&chunk);
538 }
539 Ok(bytes)
540}
541fn parse_strong_etag(value: Option<&HeaderValue>) -> Result<String> {
550 let value = value
551 .and_then(|value| value.to_str().ok())
552 .ok_or_else(|| eyre!("remote action manifest response is missing an ETag"))?;
553 if value.starts_with("W/") {
554 bail!("remote action manifest response has a weak ETag");
557 }
558 let etag = value
559 .strip_prefix('"')
560 .and_then(|value| value.strip_suffix('"'))
561 .filter(|value| is_entity_tag(value))
562 .ok_or_else(|| eyre!("remote action manifest response has an invalid ETag"))?;
563 Ok(etag.to_owned())
564}
565
566fn quoted_etag(etag: &str) -> Result<HeaderValue> {
567 if !is_entity_tag(etag) {
568 bail!("invalid remote action manifest ETag");
569 }
570 Ok(HeaderValue::from_str(&format!("\"{etag}\""))?)
571}
572
573fn is_entity_tag(value: &str) -> bool {
578 !value.is_empty()
579 && value.len() <= MAX_ETAG_BYTES
580 && value.bytes().all(|byte| matches!(byte, 0x21 | 0x23..=0x7e))
581}
582fn retry_delays(retries: i64) -> impl Iterator<Item = Duration> {
583 [200u64, 1_000, 4_000, 15_000]
584 .into_iter()
585 .chain(std::iter::repeat(15_000))
586 .map(Duration::from_millis)
587 .map(|duration| {
588 let factor = 0.5 + rand::random::<f64>() * 0.5;
589 Duration::from_secs_f64(duration.as_secs_f64() * factor)
590 })
591 .take(retries.max(0) as usize)
592}
593
594fn is_dns_error(error: &(dyn std::error::Error + 'static)) -> bool {
598 let mut current = Some(error);
599 while let Some(source) = current {
600 if source.to_string() == "dns error" {
601 return true;
602 }
603 current = source.source();
604 }
605 false
606}
607
608#[derive(Debug)]
615pub(crate) struct TransientRequest(pub(crate) &'static str);
616
617impl std::fmt::Display for TransientRequest {
618 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
619 formatter.write_str(self.0)
620 }
621}
622
623impl std::error::Error for TransientRequest {}
624
625fn is_transient(error: &eyre::Report) -> bool {
626 if is_dns_error(error.as_ref()) {
629 return false;
630 }
631 error.chain().any(|source| {
632 if source.downcast_ref::<TransientRequest>().is_some() {
633 return true;
634 }
635 let Some(error) = source.downcast_ref::<reqwest::Error>() else {
636 return false;
637 };
638 if error.is_timeout() || error.is_connect() || error.is_body() {
639 return true;
640 }
641 error.status().is_some_and(|status| {
642 let status = status.as_u16();
643 status == 408 || status == 429 || (500..600).contains(&status)
644 })
645 })
646}
647
648async fn retry_async<F, Fut, T>(verb: &str, url: &Url, retries: i64, mut operation: F) -> Result<T>
649where
650 F: FnMut() -> Fut,
651 Fut: std::future::Future<Output = Result<T>>,
652{
653 let mut delays = retry_delays(retries);
654 let mut attempt = 1;
655 loop {
656 let started_at = Instant::now();
657 match operation().await {
658 Ok(value) => return Ok(value),
659 Err(error) if is_transient(&error) => {
660 let Some(delay) = delays.next() else {
661 return Err(error);
662 };
663 warn!(
664 "HTTP {verb} {url} attempt {attempt} failed after {:?} (transient): {error}; retrying in {delay:?}",
665 started_at.elapsed()
666 );
667 tokio::time::sleep(delay).await;
668 attempt += 1;
669 }
670 Err(error) => return Err(error),
671 }
672 }
673}