Skip to main content

mbx_cache_core/
remote_s3.rs

1//! The remote cache backed directly by an S3-compatible object store.
2//!
3//! This is the other backend behind [`crate::RemoteCacheClient`]. Where the
4//! protocol server answers lookups, validates uploads, and advertises
5//! extensions, a bucket only stores objects. The v1 protocol is designed for
6//! that: blobs and action results are immutable and content-addressed, so they
7//! are written create-only and re-storing one is not an error, and the single
8//! record that is updated in place -- the task action manifest -- carries an
9//! entity tag that S3's conditional writes can honour.
10//!
11//! What a bucket cannot do it declines rather than approximates. Blob packs,
12//! batched lookups, and negotiated compression all report themselves absent,
13//! which is the same answer the client already handles from a server that does
14//! not implement them.
15
16use crate::imds::InstanceRoleCredentials;
17use crate::sigv4::{PayloadHash, S3Credentials, SigningContext, sign};
18use crate::{
19    BlobPackReceipt, BlobSource, BlobUpload, CacheDigest, MAX_REMOTE_BLOB_BYTES,
20    MAX_REMOTE_JSON_BYTES, ManifestPutOutcome, RemoteActionManifest, RemoteActionResult,
21    RemoteBlobPack, TransientRequest, parse_strong_etag, quoted_etag, read_bounded_json,
22    retry_async,
23};
24use eyre::{Result, bail, eyre};
25use log::warn;
26use reqwest::StatusCode;
27use reqwest::header::{CONTENT_LENGTH, ETAG, IF_MATCH, IF_NONE_MATCH};
28use std::fs;
29use std::path::Path;
30use std::sync::atomic::{AtomicBool, Ordering};
31use std::sync::{Arc, Mutex};
32use std::time::{Duration, Instant, SystemTime};
33use tokio::io::AsyncWriteExt;
34use url::Url;
35
36/// The object-key layout this client reads and writes.
37///
38/// Independent of the protocol version, which describes a wire format rather
39/// than a bucket layout, though today they are both 1.
40const LAYOUT_VERSION: u8 = 1;
41/// Object read to prove an endpoint answers, credentials work, and the bucket
42/// exists. It is never written, so a diagnostic stays read-only.
43const CONNECTIVITY_PROBE_KEY: &str = "connectivity-probe";
44/// Bytes of an S3 error document read before giving up on a diagnosis.
45const MAX_ERROR_BODY_BYTES: usize = 8 * 1024;
46/// How long before expiry instance role credentials are fetched again. EC2
47/// publishes the next set at least five minutes ahead, so asking any earlier
48/// would only return the set already held.
49const RENEWAL_MARGIN: Duration = Duration::from_secs(5 * 60);
50/// How long to wait before asking the metadata service again after it failed
51/// or had nothing newer, so a stalled service is not asked once per request.
52const RENEWAL_RETRY: Duration = Duration::from_secs(30);
53
54/// Whether conditional writes are required, refused, or tried and given up on.
55#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, strum::EnumString, strum::Display)]
56#[strum(serialize_all = "kebab-case")]
57pub enum S3ConditionalWrites {
58    /// Use conditional writes, and stop using them if the store rejects them.
59    #[default]
60    Auto,
61    /// Require conditional writes, failing if the store does not implement them.
62    Required,
63    /// Never send conditional headers.
64    Off,
65}
66
67/// Connection, addressing, and credential settings for an S3 remote cache.
68pub struct S3RemoteCacheConfig {
69    /// Bucket holding the cache.
70    pub bucket: String,
71    /// Key prefix within the bucket, empty for the bucket root.
72    pub prefix: String,
73    /// Namespace isolating one project's cache, used as a key prefix.
74    pub namespace: String,
75    /// Region used to sign requests.
76    pub region: String,
77    /// Endpoint override for a non-AWS store such as MinIO or R2.
78    pub endpoint: Option<Url>,
79    /// Force bucket-in-path addressing, overriding what the endpoint implies.
80    pub force_path_style: Option<bool>,
81    /// How to treat a store that does not implement conditional writes.
82    pub conditional_writes: S3ConditionalWrites,
83    /// Credentials used to sign requests.
84    pub credentials: S3Credentials,
85    /// Maximum time allowed to establish a connection.
86    pub connect_timeout: Duration,
87    /// Maximum time without response progress for ordinary requests.
88    pub read_timeout: Duration,
89    /// Deadline for one blob download, spanning every retry attempt and the
90    /// backoff between them rather than bounding a single attempt.
91    ///
92    /// A single stalled attempt is already bounded by `connect_timeout` and
93    /// `read_timeout`, so this budget exists to cap the total wall-clock one
94    /// logical download may spend: exhausting it fails the download even when
95    /// retries remain. Size it for the largest artifact worth waiting on, not
96    /// for one attempt at it.
97    pub download_timeout: Duration,
98    /// Number of attempts after the initial request for retryable failures.
99    pub retries: i64,
100}
101
102/// Kinds of object the cache stores, which are also their key prefixes.
103#[derive(Clone, Copy)]
104enum ObjectKind {
105    Blob,
106    ActionResult,
107    ActionManifest,
108}
109
110impl ObjectKind {
111    fn as_str(self) -> &'static str {
112        match self {
113            Self::Blob => "blobs",
114            Self::ActionResult => "action-results",
115            Self::ActionManifest => "action-manifests",
116        }
117    }
118}
119
120/// Where the credentials that sign requests come from.
121enum CredentialSource {
122    /// The same credentials for the whole session.
123    Fixed(S3Credentials),
124    /// Credentials that expire and are fetched again from the instance role.
125    InstanceRole(Arc<InstanceRole>),
126}
127
128struct InstanceRole {
129    provider: InstanceRoleCredentials,
130    /// Never held across an await, so signing a request never waits on the
131    /// metadata service while the held credentials are still good.
132    state: Mutex<RenewalState>,
133    /// Held for the length of a lookup, so callers that need new credentials
134    /// while one is running wait for it instead of making their own.
135    lookup: tokio::sync::Mutex<()>,
136    /// Whether a background renewal is already running.
137    renewing: AtomicBool,
138}
139
140struct RenewalState {
141    credentials: S3Credentials,
142    expires_at: SystemTime,
143    /// Earliest time to ask the metadata service again, after a lookup that
144    /// failed or returned credentials still inside the renewal margin.
145    retry_at: Option<Instant>,
146    /// Whether the last lookup failed, so a stalled service is reported once
147    /// rather than on every retry.
148    failing: bool,
149}
150
151impl RenewalState {
152    /// Whether the credentials are inside the renewal margin.
153    fn renewal_due(&self) -> bool {
154        self.expires_at
155            .checked_sub(RENEWAL_MARGIN)
156            .is_none_or(|due| due <= SystemTime::now())
157    }
158
159    /// Whether the credentials have already stopped working.
160    fn expired(&self) -> bool {
161        self.expires_at <= SystemTime::now()
162    }
163
164    /// Whether the last lookup was recent enough that another would be noise.
165    fn backing_off(&self) -> bool {
166        self.retry_at.is_some_and(|at| at > Instant::now())
167    }
168}
169
170impl InstanceRole {
171    /// Fetch new credentials unless another caller just did or the service was
172    /// asked too recently.
173    ///
174    /// A failed lookup keeps the credentials already held: they usually have
175    /// minutes left, and once they lapse S3 rejects them with an error that
176    /// says so.
177    async fn renew(&self) {
178        let _lookup = self.lookup.lock().await;
179        {
180            let state = self.state.lock().unwrap();
181            if !state.renewal_due() || state.backing_off() {
182                return;
183            }
184        }
185        let fetched = self.provider.fetch().await;
186        let mut state = self.state.lock().unwrap();
187        match fetched {
188            Ok(fresh) => {
189                state.credentials = fresh.credentials;
190                state.expires_at = fresh.expires_at;
191                state.failing = false;
192            }
193            Err(error) => {
194                if !state.failing {
195                    warn!(
196                        "could not renew the EC2 instance role credentials, \
197                         continuing with the current ones: {error:#}"
198                    );
199                }
200                state.failing = true;
201            }
202        }
203        state.retry_at = state.renewal_due().then(|| Instant::now() + RENEWAL_RETRY);
204    }
205}
206
207impl CredentialSource {
208    /// The credentials to sign the next request with.
209    ///
210    /// Instance role credentials are renewed in the background once they are
211    /// inside the renewal margin, and requests keep signing with the ones held
212    /// until the new set arrives. A metadata service that hangs therefore
213    /// costs nothing while the held set is valid. Only a set that has already
214    /// lapsed makes a request wait, since nothing valid is left to sign with.
215    async fn current(&self) -> S3Credentials {
216        let role = match self {
217            Self::Fixed(credentials) => return credentials.clone(),
218            Self::InstanceRole(role) => role,
219        };
220        let (due, expired) = {
221            let state = role.state.lock().unwrap();
222            (state.renewal_due() && !state.backing_off(), state.expired())
223        };
224        if due && expired {
225            role.renew().await;
226        } else if due && !role.renewing.swap(true, Ordering::AcqRel) {
227            let role = Arc::clone(role);
228            tokio::spawn(async move {
229                role.renew().await;
230                role.renewing.store(false, Ordering::Release);
231            });
232        }
233        role.state.lock().unwrap().credentials.clone()
234    }
235
236    /// What to check when the store rejects these credentials with `code`.
237    fn rejection_hint(&self, code: Option<&str>) -> &'static str {
238        match self {
239            Self::Fixed(_) => "Check AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY",
240            // Which set S3 saw is not knowable here: a renewal may have
241            // replaced it since the request was signed. So this names both
242            // causes rather than judging by the current state.
243            Self::InstanceRole(_)
244                if matches!(code, Some("ExpiredToken" | "TokenRefreshRequired")) =>
245            {
246                "S3 says the EC2 instance role's credentials have expired. If a renewal \
247                 warning was logged, the instance metadata service could not be reached; if \
248                 not, the two clocks disagree"
249            }
250            Self::InstanceRole(_) => {
251                "These are the EC2 instance role's credentials, so check the role's \
252                 permissions on the bucket"
253            }
254        }
255    }
256}
257
258pub(crate) struct S3RemoteCache {
259    client: reqwest::Client,
260    /// Bucket root, always ending in `/` so keys join onto it.
261    base_url: Url,
262    /// Key prefix covering the configured prefix, namespace, and layout version.
263    root: String,
264    region: String,
265    credentials: CredentialSource,
266    conditional_writes: S3ConditionalWrites,
267    /// Latched once a store has told us it does not implement conditional
268    /// writes, so the rest of the session stops asking.
269    conditionals_disabled: AtomicBool,
270    /// Latched once a read has been refused where an absent object would have
271    /// been, so the explanation is offered once rather than per lookup.
272    absence_is_ambiguous: AtomicBool,
273    download_timeout: Duration,
274    retries: i64,
275}
276
277impl S3RemoteCache {
278    pub(crate) fn new(config: S3RemoteCacheConfig) -> Result<Self> {
279        validate_bucket(&config.bucket)?;
280        let prefix = normalize_prefix(&config.prefix)?;
281        validate_key_path(&config.namespace, "remote cache namespace")?;
282        if config.region.trim().is_empty() {
283            bail!("an S3 remote cache needs a region");
284        }
285        let client = reqwest::Client::builder()
286            .connect_timeout(config.connect_timeout)
287            .read_timeout(config.read_timeout)
288            .redirect(reqwest::redirect::Policy::none())
289            .build()?;
290        Ok(Self {
291            client,
292            base_url: base_url(&config)?,
293            root: format!("{prefix}{}/v{LAYOUT_VERSION}/", config.namespace.trim()),
294            region: config.region.trim().to_string(),
295            credentials: CredentialSource::Fixed(config.credentials),
296            conditional_writes: config.conditional_writes,
297            conditionals_disabled: AtomicBool::new(false),
298            absence_is_ambiguous: AtomicBool::new(false),
299            download_timeout: config.download_timeout,
300            retries: config.retries,
301        })
302    }
303
304    /// Hand credentials to `provider` for renewal before they expire.
305    ///
306    /// `expires_at` is when the credentials the store was built with stop
307    /// working, as reported alongside them.
308    pub(crate) fn with_instance_role(
309        mut self,
310        provider: InstanceRoleCredentials,
311        expires_at: SystemTime,
312    ) -> Self {
313        let CredentialSource::Fixed(credentials) = self.credentials else {
314            return self;
315        };
316        self.credentials = CredentialSource::InstanceRole(Arc::new(InstanceRole {
317            provider,
318            state: Mutex::new(RenewalState {
319                credentials,
320                expires_at,
321                retry_at: None,
322                failing: false,
323            }),
324            lookup: tokio::sync::Mutex::new(()),
325            renewing: AtomicBool::new(false),
326        }));
327        self
328    }
329
330    fn object_url(&self, kind: ObjectKind, digest: &CacheDigest) -> Result<Url> {
331        digest.validate()?;
332        if matches!(kind, ObjectKind::ActionResult | ObjectKind::ActionManifest)
333            && digest.algorithm != "blake3"
334        {
335            bail!("remote cache action keys must use blake3");
336        }
337        self.key_url(&format!(
338            "{}/{}/{}/{}",
339            kind.as_str(),
340            digest.algorithm,
341            digest.hash,
342            digest.size
343        ))
344    }
345
346    fn key_url(&self, key: &str) -> Result<Url> {
347        Ok(self.base_url.join(&format!("{}{key}", self.root))?)
348    }
349
350    /// Build a request carrying a valid signature for this instant.
351    ///
352    /// Signing happens per attempt rather than once per operation: a retry
353    /// after a long backoff would otherwise present a stale `x-amz-date` and be
354    /// refused for clock skew.
355    async fn signed(
356        &self,
357        method: reqwest::Method,
358        url: &Url,
359        payload: &PayloadHash,
360    ) -> Result<reqwest::RequestBuilder> {
361        let credentials = self.credentials.current().await;
362        let context = SigningContext {
363            credentials: &credentials,
364            region: &self.region,
365            timestamp: SystemTime::now(),
366        };
367        let mut request = self.client.request(method.clone(), url.clone());
368        for (name, value) in sign(method.as_str(), url, &context, payload)? {
369            request = request.header(name, value);
370        }
371        Ok(request)
372    }
373
374    /// Whether a conditional header should be attached to the next write.
375    fn conditionals_enabled(&self) -> bool {
376        self.conditional_writes != S3ConditionalWrites::Off
377            && !self.conditionals_disabled.load(Ordering::Relaxed)
378    }
379
380    /// Whether the same write may be tried again without its condition.
381    ///
382    /// Under `required` it never may: a caller that asked for the guarantee gets
383    /// an error rather than a silent downgrade.
384    fn may_drop_conditionals(&self) -> bool {
385        self.conditional_writes == S3ConditionalWrites::Auto
386    }
387
388    /// Record that this store does not implement conditional writes.
389    ///
390    /// Called only once an unconditional write has actually succeeded, which is
391    /// what distinguishes a store that refuses conditions from one that refused
392    /// this request for some other reason and would have refused it anyway. A
393    /// `501` from an intermediary is not evidence about the store, and latching
394    /// on it would quietly turn every later manifest update into a
395    /// last-writer-wins one.
396    fn note_conditionals_unsupported(&self) {
397        if !self.conditionals_disabled.swap(true, Ordering::Relaxed) {
398            warn!(
399                "the remote object store does not implement conditional writes; \
400                 continuing without them. Blobs and action results are content-addressed, so \
401                 this is safe; concurrent task manifest updates can now lose predictions, \
402                 which costs prefetch coverage on later builds"
403            );
404        }
405    }
406
407    pub(crate) async fn check_connection(&self) -> Result<()> {
408        let url = self.key_url(CONNECTIVITY_PROBE_KEY)?;
409        // A GET rather than a HEAD, because S3 explains a refusal in the
410        // response body and a HEAD has none. The key is never written, so the
411        // expected answer is a 404 carrying an error document.
412        retry_async("GET", &url, self.retries, || async {
413            let response = self
414                .signed(reqwest::Method::GET, &url, &PayloadHash::empty())
415                .await?
416                .send()
417                .await?;
418            match response.status() {
419                // The bucket answered and authorized the request. Whether this
420                // one key exists is beside the point.
421                StatusCode::OK | StatusCode::NOT_FOUND => Ok(()),
422                StatusCode::BAD_REQUEST | StatusCode::FORBIDDEN => {
423                    let failure = FailedRequest::read(response).await;
424                    if failure.is_credentials_rejected() {
425                        bail!(
426                            "the remote object store rejected these credentials for {url}: {}. \
427                             {}, and that this machine's clock is correct",
428                            failure.code.as_deref().unwrap_or("forbidden"),
429                            self.credentials.rejection_hint(failure.code.as_deref())
430                        );
431                    }
432                    if failure.status != StatusCode::FORBIDDEN {
433                        return Err(failure.report("connect to", &url));
434                    }
435                    // The signature was accepted; something declined this one
436                    // object. The probe key is never written, so without
437                    // s3:ListBucket that is exactly what a working
438                    // configuration looks like -- and it is also what a
439                    // credential with no access to the prefix looks like.
440                    warn!(
441                        "the remote object store did not confirm access to {url}. That is \
442                         expected without s3:ListBucket on the bucket, where S3 refuses a \
443                         read rather than reporting the object absent; grant it to tell the \
444                         two apart. If the cache never hits, these credentials may not be \
445                         allowed to read the prefix"
446                    );
447                    Ok(())
448                }
449                StatusCode::MOVED_PERMANENTLY | StatusCode::TEMPORARY_REDIRECT => {
450                    let region = response
451                        .headers()
452                        .get("x-amz-bucket-region")
453                        .and_then(|value| value.to_str().ok())
454                        .unwrap_or("another region");
455                    bail!(
456                        "the bucket is in {region}, not {}; set the remote region to match",
457                        self.region
458                    )
459                }
460                _ => Err(FailedRequest::read(response)
461                    .await
462                    .report("connect to", &url)),
463            }
464        })
465        .await
466    }
467
468    pub(crate) async fn get_blob(
469        &self,
470        digest: &CacheDigest,
471        _media_type: &'static str,
472    ) -> Result<Vec<u8>> {
473        if digest.size > MAX_REMOTE_JSON_BYTES {
474            bail!(
475                "remote cache in-memory blob declared {} bytes, over the {} byte limit",
476                digest.size,
477                MAX_REMOTE_JSON_BYTES
478            );
479        }
480        let url = self.object_url(ObjectKind::Blob, digest)?;
481        retry_async("GET", &url, self.retries, || async {
482            let mut response = self.get(&url).await?;
483            // Stop reading as soon as the response outgrows the digest it claims
484            // to satisfy, exactly as the protocol client does.
485            let mut bytes = Vec::new();
486            while let Some(chunk) = response.chunk().await? {
487                if bytes.len() as u64 + chunk.len() as u64 > digest.size {
488                    bail!("remote cache blob exceeded the size of its digest");
489                }
490                bytes.extend_from_slice(&chunk);
491            }
492            if !digest.matches_bytes(&bytes)? {
493                bail!("remote cache blob failed digest verification");
494            }
495            Ok(bytes)
496        })
497        .await
498    }
499
500    pub(crate) async fn get_blob_file(
501        &self,
502        digest: &CacheDigest,
503        staging_dir: &Path,
504    ) -> Result<tempfile::NamedTempFile> {
505        if digest.size > MAX_REMOTE_BLOB_BYTES {
506            bail!(
507                "remote cache blob declared {} bytes, over the {} byte limit",
508                digest.size,
509                MAX_REMOTE_BLOB_BYTES
510            );
511        }
512        let url = self.object_url(ObjectKind::Blob, digest)?;
513        let download = retry_async("GET", &url, self.retries, || async {
514            let mut response = self.get(&url).await?;
515            fs::create_dir_all(staging_dir)?;
516            let temporary = tempfile::NamedTempFile::new_in(staging_dir)?;
517            let mut output = tokio::fs::File::from_std(temporary.reopen()?);
518            let mut written = 0u64;
519            while let Some(chunk) = response.chunk().await? {
520                written += chunk.len() as u64;
521                if written > digest.size {
522                    bail!("remote cache blob exceeded the size of its digest");
523                }
524                output.write_all(&chunk).await?;
525            }
526            output.flush().await?;
527            drop(output);
528            if !digest.matches_file(temporary.path())? {
529                bail!("remote cache blob failed digest verification");
530            }
531            Ok(temporary)
532        });
533        let download_timeout = self.download_timeout;
534        // Deliberately outside `retry_async`: `download_timeout` is a deadline
535        // for the whole download, not a per-attempt bound. A stalled attempt is
536        // already caught by the client's connect and read timeouts.
537        tokio::time::timeout(download_timeout, download)
538            .await
539            .map_err(|_| {
540                eyre!(
541                    "remote cache blob download for {url} exceeded its {download_timeout:?} budget across all attempts"
542                )
543            })?
544    }
545
546    /// Fetch an object that must exist, turning any other status into an error.
547    async fn get(&self, url: &Url) -> Result<reqwest::Response> {
548        let response = self
549            .signed(reqwest::Method::GET, url, &PayloadHash::empty())
550            .await?
551            .send()
552            .await?;
553        if response.status().is_success() {
554            Ok(response)
555        } else {
556            Err(FailedRequest::read(response).await.report("read", url))
557        }
558    }
559
560    /// Whether a failed read should be treated as the object not being there.
561    ///
562    /// S3 answers `403` rather than `404` for an absent object when the caller
563    /// may not list the bucket, so a miss under a least-privilege policy is
564    /// indistinguishable from a denial by status alone. The error code does
565    /// distinguish the one case that matters: credentials S3 itself rejected.
566    /// Anything else is treated as a miss, since reporting every cold lookup as
567    /// a failure would make such a policy unusable, and a warning explains it
568    /// once.
569    fn reads_as_absent(&self, failure: &FailedRequest) -> bool {
570        if failure.status == StatusCode::NOT_FOUND {
571            return true;
572        }
573        if failure.status != StatusCode::FORBIDDEN || failure.is_credentials_rejected() {
574            return false;
575        }
576        if !self.absence_is_ambiguous.swap(true, Ordering::Relaxed) {
577            warn!(
578                "the remote object store refused a read instead of reporting the object \
579                 absent, which is what S3 does without s3:ListBucket on the bucket. \
580                 Treating it as a cache miss. Grant s3:ListBucket so a miss is a miss; \
581                 if the cache never hits, these credentials may simply not be allowed to \
582                 read it"
583            );
584        }
585        true
586    }
587
588    pub(crate) async fn get_action_result(
589        &self,
590        action: &CacheDigest,
591    ) -> Result<Option<RemoteActionResult>> {
592        let url = self.object_url(ObjectKind::ActionResult, action)?;
593        let result = retry_async("GET", &url, self.retries, || async {
594            let response = self
595                .signed(reqwest::Method::GET, &url, &PayloadHash::empty())
596                .await?
597                .send()
598                .await?;
599            if !response.status().is_success() {
600                let failure = FailedRequest::read(response).await;
601                return if self.reads_as_absent(&failure) {
602                    Ok(None)
603                } else {
604                    Err(failure.report("read", &url))
605                };
606            }
607            let bytes = read_bounded_json(response, "action result").await?;
608            Ok(Some(serde_json::from_slice::<RemoteActionResult>(&bytes)?))
609        })
610        .await?;
611        if let Some(result) = &result
612            && (result.version != 1 || result.action != *action)
613        {
614            bail!("remote action result does not match requested action");
615        }
616        Ok(result)
617    }
618
619    pub(crate) async fn put_action_result(&self, result: &RemoteActionResult) -> Result<()> {
620        let url = self.object_url(ObjectKind::ActionResult, &result.action)?;
621        let body = serde_json::to_vec(result)?;
622        retry_async("PUT", &url, self.retries, || async {
623            // Storing the same record twice is not an error: the key is its
624            // content address, so an existing object already holds these bytes.
625            self.put_create(&url, &body).await.map(drop)
626        })
627        .await
628    }
629
630    pub(crate) async fn get_action_manifest(
631        &self,
632        key: &CacheDigest,
633    ) -> Result<Option<RemoteActionManifest>> {
634        let url = self.object_url(ObjectKind::ActionManifest, key)?;
635        retry_async("GET", &url, self.retries, || async {
636            let response = self
637                .signed(reqwest::Method::GET, &url, &PayloadHash::empty())
638                .await?
639                .send()
640                .await?;
641            if !response.status().is_success() {
642                let failure = FailedRequest::read(response).await;
643                return if self.reads_as_absent(&failure) {
644                    Ok(None)
645                } else {
646                    Err(failure.report("read", &url))
647                };
648            }
649            let etag = parse_strong_etag(response.headers().get(ETAG))?;
650            let bytes = read_bounded_json(response, "action manifest").await?;
651            Ok(Some(RemoteActionManifest { bytes, etag }))
652        })
653        .await
654    }
655
656    pub(crate) async fn put_action_manifest(
657        &self,
658        key: &CacheDigest,
659        bytes: &[u8],
660        expected_etag: Option<&str>,
661    ) -> Result<ManifestPutOutcome> {
662        let url = self.object_url(ObjectKind::ActionManifest, key)?;
663        let body = bytes.to_vec();
664        let expected_etag = expected_etag.map(quoted_etag).transpose()?;
665        retry_async("PUT", &url, self.retries, || async {
666            // Set once the store has refused a condition and the same write is
667            // being tried without one, so that success can be attributed.
668            let mut dropped_condition = false;
669            let outcome = loop {
670                let conditional = self.conditionals_enabled() && !dropped_condition;
671                let mut request = self
672                    .signed(reqwest::Method::PUT, &url, &PayloadHash::of(&body))
673                    .await?
674                    .header(CONTENT_LENGTH, body.len())
675                    .body(body.clone());
676                if conditional {
677                    request = match &expected_etag {
678                        Some(etag) => request.header(IF_MATCH, etag),
679                        None => request.header(IF_NONE_MATCH, "*"),
680                    };
681                }
682                let response = request.send().await?;
683                let status = response.status();
684                if status.is_success() {
685                    if dropped_condition {
686                        self.note_conditionals_unsupported();
687                    }
688                    break ManifestPutOutcome::Stored;
689                }
690                if conditional && status == StatusCode::PRECONDITION_FAILED {
691                    // Either the manifest moved under us or, for a create, one
692                    // already exists. Both mean re-reading and merging.
693                    break ManifestPutOutcome::PreconditionFailed;
694                }
695                if status == StatusCode::CONFLICT {
696                    // A concurrent conditional write on the same key. S3 asks
697                    // that this be retried rather than treated as a conflict of
698                    // content.
699                    return Err(conditional_request_conflict(&url));
700                }
701                let failure = FailedRequest::read(response).await;
702                if conditional && failure.is_not_implemented() {
703                    if self.may_drop_conditionals() {
704                        dropped_condition = true;
705                        continue;
706                    }
707                    return Err(failure.report("update", &url).wrap_err(
708                        "conditional writes are required but this store does not implement them",
709                    ));
710                }
711                return Err(failure.report("update", &url));
712            };
713            Ok(outcome)
714        })
715        .await
716    }
717
718    pub(crate) async fn put_blob(&self, upload: &BlobUpload) -> Result<()> {
719        upload.digest.validate()?;
720        let url = self.object_url(ObjectKind::Blob, &upload.digest)?;
721        retry_async("PUT", &url, self.retries, || async {
722            match &upload.source {
723                BlobSource::Bytes(bytes) => self.put_create(&url, bytes).await.map(drop),
724                BlobSource::File(file) => self.put_create_file(&url, file.path()).await,
725                BlobSource::Path(path) => self.put_create_file(&url, path).await,
726            }
727        })
728        .await
729    }
730
731    /// Store bytes under a key that is their content address.
732    ///
733    /// Returns whether this request is what created the object; an object that
734    /// was already there is success, since its key names these exact bytes.
735    async fn put_create(&self, url: &Url, body: &[u8]) -> Result<bool> {
736        let mut dropped_condition = false;
737        loop {
738            let conditional = self.conditionals_enabled() && !dropped_condition;
739            let mut request = self
740                .signed(reqwest::Method::PUT, url, &PayloadHash::of(body))
741                .await?
742                .header(CONTENT_LENGTH, body.len())
743                .body(body.to_vec());
744            if conditional {
745                request = request.header(IF_NONE_MATCH, "*");
746            }
747            match self
748                .finish_create(url, request.send().await?, conditional)
749                .await?
750            {
751                Some(created) => {
752                    if dropped_condition {
753                        self.note_conditionals_unsupported();
754                    }
755                    return Ok(created);
756                }
757                None => dropped_condition = true,
758            }
759        }
760    }
761
762    /// Store a file under a key that is its content address.
763    ///
764    /// The body is streamed with an explicit length rather than chunked, which
765    /// S3 requires, and is signed as an unsigned payload so that a multi-gigabyte
766    /// artifact is not read twice just to compute a hash TLS and the content
767    /// address already stand behind.
768    async fn put_create_file(&self, url: &Url, path: &Path) -> Result<()> {
769        let mut dropped_condition = false;
770        loop {
771            let conditional = self.conditionals_enabled() && !dropped_condition;
772            let file = tokio::fs::File::open(path).await?;
773            let length = file.metadata().await?.len();
774            let mut request = self
775                .signed(reqwest::Method::PUT, url, &PayloadHash::Unsigned)
776                .await?
777                .header(CONTENT_LENGTH, length)
778                .body(reqwest::Body::wrap_stream(
779                    tokio_util::io::ReaderStream::new(file),
780                ));
781            if conditional {
782                request = request.header(IF_NONE_MATCH, "*");
783            }
784            match self
785                .finish_create(url, request.send().await?, conditional)
786                .await?
787            {
788                Some(_) => {
789                    if dropped_condition {
790                        self.note_conditionals_unsupported();
791                    }
792                    return Ok(());
793                }
794                None => dropped_condition = true,
795            }
796        }
797    }
798
799    /// Interpret the response to a create-only write.
800    ///
801    /// `None` asks the caller to send the same request again without its
802    /// conditional header, the store having just refused one.
803    async fn finish_create(
804        &self,
805        url: &Url,
806        response: reqwest::Response,
807        conditional: bool,
808    ) -> Result<Option<bool>> {
809        let status = response.status();
810        if status.is_success() {
811            return Ok(Some(true));
812        }
813        // Only a condition this request actually carried can have failed. A 412
814        // against an unconditional write means something else refused it, and
815        // reporting that as a stored object would lose the write silently.
816        if conditional && status == StatusCode::PRECONDITION_FAILED {
817            return Ok(Some(false));
818        }
819        if status == StatusCode::CONFLICT {
820            return Err(conditional_request_conflict(url));
821        }
822        let failure = FailedRequest::read(response).await;
823        if conditional && failure.is_not_implemented() {
824            if self.may_drop_conditionals() {
825                return Ok(None);
826            }
827            return Err(failure.report("store", url).wrap_err(
828                "conditional writes are required but this store does not implement them",
829            ));
830        }
831        Err(failure.report("store", url))
832    }
833
834    // Everything below is an extension a protocol server offers and a bucket
835    // does not. Reporting absence is what routes the caller to the per-object
836    // requests that every version of the protocol has made.
837
838    pub(crate) async fn get_action_results(
839        &self,
840        actions: &[CacheDigest],
841    ) -> Result<Option<Vec<RemoteActionResult>>> {
842        Ok(actions.is_empty().then(Vec::new))
843    }
844
845    pub(crate) async fn action_batch_limit(&self) -> Result<Option<usize>> {
846        Ok(None)
847    }
848
849    pub(crate) async fn get_blob_pack(
850        &self,
851        _digests: &[CacheDigest],
852        _staging_dir: &Path,
853    ) -> Result<Option<RemoteBlobPack>> {
854        Ok(None)
855    }
856
857    pub(crate) async fn get_blob_pack_with_limit(
858        &self,
859        _digests: &[CacheDigest],
860        _staging_dir: &Path,
861        _max_bytes: u64,
862    ) -> Result<Option<RemoteBlobPack>> {
863        Ok(None)
864    }
865
866    pub(crate) async fn blob_pack_upload_limits(&self) -> Result<Option<crate::BlobPackLimits>> {
867        Ok(None)
868    }
869
870    pub(crate) async fn put_blob_pack(
871        &self,
872        uploads: &[BlobUpload],
873    ) -> Result<Option<BlobPackReceipt>> {
874        Ok(uploads.is_empty().then_some(BlobPackReceipt {
875            created: 0,
876            existing: 0,
877        }))
878    }
879}
880
881/// A concurrent conditional write, which S3 asks callers to retry.
882fn conditional_request_conflict(url: &Url) -> eyre::Report {
883    eyre::Report::new(TransientRequest("conditional request conflict")).wrap_err(format!(
884        "a concurrent conditional write to {url} conflicted"
885    ))
886}
887
888/// A request that failed, read far enough to say why.
889///
890/// The body is consumed here because deciding what a failure means can depend
891/// on the store's own error code, and a response cannot be read twice.
892struct FailedRequest {
893    status: StatusCode,
894    code: Option<String>,
895}
896
897impl FailedRequest {
898    async fn read(mut response: reqwest::Response) -> Self {
899        let status = response.status();
900        // Bounded like every other body this client reads. An error document is
901        // a diagnostic, and a store that streams without end must not be able to
902        // exhaust this process through one.
903        let mut body = Vec::new();
904        while body.len() < MAX_ERROR_BODY_BYTES {
905            match response.chunk().await {
906                Ok(Some(chunk)) => body.extend_from_slice(&chunk),
907                Ok(None) | Err(_) => break,
908            }
909        }
910        body.truncate(MAX_ERROR_BODY_BYTES);
911        Self {
912            status,
913            code: error_code(&String::from_utf8_lossy(&body)).map(str::to_string),
914        }
915    }
916
917    /// Whether the store is saying it does not implement what was asked of it.
918    ///
919    /// S3-compatible stores disagree about how to say this. AWS answers `501`
920    /// for an unimplemented feature, while some others answer `400` with
921    /// `NotImplemented` in the document. A bare `400` is not enough on its own:
922    /// it is also how a store reports a request it merely disliked.
923    fn is_not_implemented(&self) -> bool {
924        self.status == StatusCode::NOT_IMPLEMENTED
925            || (self.status == StatusCode::BAD_REQUEST
926                && self.code.as_deref() == Some("NotImplemented"))
927    }
928
929    /// Whether S3 rejected the credentials themselves, rather than declining
930    /// this one object.
931    ///
932    /// These codes say the request never authenticated, which no bucket policy
933    /// can explain away. `AccessDenied` deliberately is not among them: it is
934    /// what an absent object looks like without `s3:ListBucket`.
935    fn is_credentials_rejected(&self) -> bool {
936        matches!(
937            (self.status, self.code.as_deref()),
938            (
939                StatusCode::BAD_REQUEST,
940                Some("ExpiredToken" | "TokenRefreshRequired" | "InvalidToken")
941            ) | (
942                StatusCode::FORBIDDEN,
943                Some(
944                    "SignatureDoesNotMatch"
945                        | "InvalidAccessKeyId"
946                        | "InvalidSecurity"
947                        | "ExpiredToken"
948                        | "TokenRefreshRequired"
949                        | "RequestTimeTooSkewed"
950                )
951            )
952        )
953    }
954
955    /// Whether the store is asking to be tried again.
956    ///
957    /// `500` and `503` are how S3 reports an internal error and a prefix under
958    /// too much load, and it documents both as the client's to retry. `501` is
959    /// deliberately absent: a store that does not implement something will not
960    /// implement it a moment later.
961    fn is_retryable(&self) -> bool {
962        matches!(self.status.as_u16(), 408 | 429 | 500 | 502 | 503 | 504)
963    }
964
965    fn report(&self, verb: &str, url: &Url) -> eyre::Report {
966        let detail = match &self.code {
967            Some(code) => format!("failed to {verb} {url}: {} ({code})", self.status),
968            None => format!("failed to {verb} {url}: {}", self.status),
969        };
970        if self.is_retryable() {
971            // The protocol client gets this for free: `error_for_status` hands
972            // `is_transient` a `reqwest::Error` carrying the status. Nothing
973            // here produces one, so a retryable status has to say so itself,
974            // or a `503` under load would fail a build that asked for retries.
975            return eyre::Report::new(TransientRequest("the store asked to be retried"))
976                .wrap_err(detail);
977        }
978        eyre!(detail)
979    }
980}
981
982/// Extract the `<Code>` of an S3 error document.
983///
984/// S3 reports errors as XML, but only the code is ever acted on, so it is
985/// scanned for rather than parsed with an XML reader this crate would otherwise
986/// not need.
987fn error_code(body: &str) -> Option<&str> {
988    let start = body.find("<Code>")? + "<Code>".len();
989    let end = body[start..].find("</Code>")? + start;
990    Some(body[start..end].trim()).filter(|code| !code.is_empty())
991}
992
993/// Reject a bucket name that could not address the store it names.
994fn validate_bucket(bucket: &str) -> Result<()> {
995    let bucket = bucket.trim();
996    if bucket.is_empty() {
997        bail!("an S3 remote cache needs a bucket");
998    }
999    if bucket.contains('/') || bucket.starts_with('.') || bucket.ends_with('.') {
1000        bail!("invalid S3 bucket name {bucket:?}");
1001    }
1002    Ok(())
1003}
1004
1005/// Normalize a key prefix to empty, or to something ending in a single `/`.
1006fn normalize_prefix(prefix: &str) -> Result<String> {
1007    let prefix = prefix.trim().trim_matches('/');
1008    if prefix.is_empty() {
1009        return Ok(String::new());
1010    }
1011    validate_key_path(prefix, "remote cache prefix")?;
1012    Ok(format!("{prefix}/"))
1013}
1014
1015/// Reject anything that would not survive being spliced into an object key.
1016///
1017/// The protocol carries the namespace in a header, where a server decides what
1018/// it means. In a bucket it is part of the key, so it has to be something a key
1019/// can hold and something that cannot climb out of its own prefix.
1020fn validate_key_path(value: &str, what: &str) -> Result<()> {
1021    let value = value.trim();
1022    if value.is_empty() {
1023        bail!("{what} must not be empty");
1024    }
1025    if value.starts_with('/') || value.ends_with('/') {
1026        bail!("{what} {value:?} must not start or end with a slash");
1027    }
1028    for segment in value.split('/') {
1029        if segment.is_empty() {
1030            bail!("{what} {value:?} must not contain an empty path segment");
1031        }
1032        if segment == "." || segment == ".." {
1033            bail!("{what} {value:?} must not contain a relative path segment");
1034        }
1035        if !segment
1036            .bytes()
1037            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
1038        {
1039            bail!(
1040                "{what} {value:?} must use only letters, digits, '.', '_', '-', and '/' \
1041                 when the remote cache is an object store"
1042            );
1043        }
1044    }
1045    Ok(())
1046}
1047
1048/// Resolve the bucket root a key is joined onto.
1049///
1050/// A custom endpoint addresses the bucket in the path, which is what MinIO and
1051/// R2 expect and what a bucket whose name is not a valid DNS label requires.
1052/// AWS itself is addressed by host, since a bucket-named host is what its TLS
1053/// certificate covers -- unless the name contains a dot, which the wildcard
1054/// certificate does not match.
1055fn base_url(config: &S3RemoteCacheConfig) -> Result<Url> {
1056    let bucket = config.bucket.trim();
1057    let (mut url, path_style) = match &config.endpoint {
1058        Some(endpoint) => (endpoint.clone(), config.force_path_style.unwrap_or(true)),
1059        None => (
1060            format!("https://s3.{}.amazonaws.com", config.region.trim()).parse()?,
1061            config
1062                .force_path_style
1063                .unwrap_or_else(|| bucket.contains('.')),
1064        ),
1065    };
1066    if path_style {
1067        let path = url.path().trim_end_matches('/').to_string();
1068        url.set_path(&format!("{path}/{bucket}/"));
1069    } else {
1070        let host = url
1071            .host_str()
1072            .ok_or_else(|| eyre!("an S3 endpoint must have a host"))?;
1073        url.set_host(Some(&format!("{bucket}.{host}")))?;
1074        // An endpoint may carry a path, and moving the bucket into the host is
1075        // no reason to drop it: a gateway at /s3 still wants to be asked at /s3.
1076        let path = url.path().trim_end_matches('/').to_string();
1077        url.set_path(&format!("{path}/"));
1078    }
1079    Ok(url)
1080}
1081
1082#[cfg(test)]
1083#[path = "remote_s3_tests.rs"]
1084mod tests;