Skip to main content

mbx_cache_core/
imds.rs

1//! Credentials from an EC2 instance role, through IMDSv2.
2//!
3//! Three small requests against the instance metadata service: a session
4//! token, the name of the attached role, and that role's temporary
5//! credentials. The AWS SDKs and CLI do the same as the last step of their
6//! credential chain, and a runner on EC2 gets working credentials this way
7//! without anything exporting them first.
8
9use crate::sigv4::S3Credentials;
10use eyre::{Context as _, Result, bail, eyre};
11use reqwest::StatusCode;
12use serde::Deserialize;
13use std::time::{Duration, SystemTime, UNIX_EPOCH};
14use url::Url;
15
16/// Where the metadata service lives on every EC2 instance.
17const DEFAULT_ENDPOINT: &str = "http://169.254.169.254";
18/// Lifetime requested for the session token. It is used within a moment and
19/// never kept, so this only has to be a value the service accepts.
20const TOKEN_TTL_SECONDS: &str = "21600";
21/// How long one request may take. The address is link-local, so anything
22/// slower than this means no metadata service is listening: a machine that is
23/// not on EC2 fails a lookup in about this long.
24const REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
25/// Attempts per request. The service throttles per instance, and AWS asks
26/// clients to retry with exponential backoff, which matters when several jobs
27/// start on one runner at once.
28const ATTEMPTS: u32 = 3;
29/// Wait before the second attempt; each later wait doubles.
30const RETRY_BACKOFF: Duration = Duration::from_millis(200);
31/// Appended when the token request times out after connecting, which is what a
32/// hop limit of 1 looks like from inside a container.
33const HOP_LIMIT_HINT: &str = "the metadata service accepted the connection but never answered. \
34    From a container, the instance's metadata hop limit may be 1. Raising \
35    HttpPutResponseHopLimit to 2 with `aws ec2 modify-instance-metadata-options` fixes \
36    that for every container on the instance, so first consider whether they should all \
37    be able to use its role";
38
39const TOKEN_TTL_HEADER: &str = "x-aws-ec2-metadata-token-ttl-seconds";
40const TOKEN_HEADER: &str = "x-aws-ec2-metadata-token";
41
42/// Temporary credentials and the instant they stop working.
43pub struct TemporaryCredentials {
44    /// The credentials to sign with.
45    pub credentials: S3Credentials,
46    /// When the service will start refusing them.
47    pub expires_at: SystemTime,
48}
49
50/// A source of credentials for the EC2 instance role attached to this machine.
51pub struct InstanceRoleCredentials {
52    client: reqwest::Client,
53    endpoint: Url,
54    backoff: Duration,
55}
56
57impl InstanceRoleCredentials {
58    /// A provider for the metadata service at `endpoint`.
59    pub fn new(endpoint: Url) -> Result<Self> {
60        Self::with_timing(endpoint, REQUEST_TIMEOUT, RETRY_BACKOFF)
61    }
62
63    fn with_timing(endpoint: Url, timeout: Duration, backoff: Duration) -> Result<Self> {
64        let client = reqwest::Client::builder()
65            .connect_timeout(timeout)
66            .timeout(timeout * 2)
67            .redirect(reqwest::redirect::Policy::none())
68            // A link-local address is never behind a proxy, and sending role
69            // credentials through one would hand them to it.
70            .no_proxy()
71            .build()?;
72        Ok(Self {
73            client,
74            endpoint,
75            backoff,
76        })
77    }
78
79    /// A provider configured the way the AWS SDKs read it, or `None` when
80    /// `AWS_EC2_METADATA_DISABLED` turns instance role lookup off.
81    ///
82    /// `AWS_EC2_METADATA_SERVICE_ENDPOINT` overrides the address, for a proxy
83    /// in front of the service or a test double.
84    pub fn from_env() -> Result<Option<Self>> {
85        Self::from_vars(|name| std::env::var(name).ok())
86    }
87
88    fn from_vars(var: impl Fn(&str) -> Option<String>) -> Result<Option<Self>> {
89        if var("AWS_EC2_METADATA_DISABLED")
90            .is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
91        {
92            return Ok(None);
93        }
94        let endpoint = var("AWS_EC2_METADATA_SERVICE_ENDPOINT")
95            .map(|value| value.trim().to_string())
96            .filter(|value| !value.is_empty());
97        let endpoint = endpoint.as_deref().unwrap_or(DEFAULT_ENDPOINT);
98        let endpoint = endpoint
99            .parse()
100            .wrap_err("invalid AWS_EC2_METADATA_SERVICE_ENDPOINT")?;
101        Self::new(endpoint).map(Some)
102    }
103
104    /// Ask the metadata service for the role's current credentials.
105    ///
106    /// The service hands out the same credentials until about five minutes
107    /// before they expire, so calling this early returns what a caller already
108    /// holds.
109    pub async fn fetch(&self) -> Result<TemporaryCredentials> {
110        let token = self
111            .token()
112            .await
113            .wrap_err("could not get an IMDSv2 session token")?;
114        let roles = self
115            .get(&token, "latest/meta-data/iam/security-credentials/")
116            .await
117            .wrap_err("could not list the instance's IAM roles")?;
118        let role = roles
119            .lines()
120            .map(str::trim)
121            .find(|line| !line.is_empty())
122            .ok_or_else(|| eyre!("the instance has no IAM role attached"))?;
123        let document = self
124            .get(
125                &token,
126                &format!("latest/meta-data/iam/security-credentials/{role}"),
127            )
128            .await
129            .wrap_err_with(|| format!("could not read credentials for role {role}"))?;
130        parse_role_document(&document)
131    }
132
133    /// The region the instance runs in, for a machine that names none itself.
134    ///
135    /// The metadata service answers this without any role, so it is a separate
136    /// lookup that only runs when the region would otherwise be missing.
137    pub async fn region(&self) -> Result<String> {
138        let token = self
139            .token()
140            .await
141            .wrap_err("could not get an IMDSv2 session token")?;
142        let region = self
143            .get(&token, "latest/meta-data/placement/region")
144            .await
145            .wrap_err("could not read the instance's region")?;
146        let region = region.trim();
147        if region.is_empty() {
148            bail!("the metadata service returned an empty region");
149        }
150        Ok(region.to_string())
151    }
152
153    async fn token(&self) -> Result<String> {
154        let request = self
155            .client
156            .put(self.endpoint.join("latest/api/token")?)
157            .header(TOKEN_TTL_HEADER, TOKEN_TTL_SECONDS);
158        let token = self.send(request).await.map_err(|error| {
159            let unanswered = error
160                .downcast_ref::<reqwest::Error>()
161                .is_some_and(|error| error.is_timeout() && !error.is_connect());
162            if unanswered {
163                error.wrap_err(HOP_LIMIT_HINT)
164            } else {
165                error
166            }
167        })?;
168        Ok(token.trim().to_string())
169    }
170
171    async fn get(&self, token: &str, path: &str) -> Result<String> {
172        let request = self
173            .client
174            .get(self.endpoint.join(path)?)
175            .header(TOKEN_HEADER, token);
176        self.send(request).await
177    }
178
179    /// Send a request and return the body of a successful answer.
180    ///
181    /// Throttling, server errors, and a request that connected but got no
182    /// answer are tried again. A refusal to connect is not: nothing is
183    /// listening, and waiting would only slow the answer on a machine that is
184    /// not on EC2.
185    async fn send(&self, request: reqwest::RequestBuilder) -> Result<String> {
186        let mut attempt = 1;
187        loop {
188            let request = request
189                .try_clone()
190                .expect("a request without a body can be cloned");
191            let (error, transient) = match request.send().await {
192                Ok(response) if response.status().is_success() => match response.text().await {
193                    Ok(body) => return Ok(body),
194                    Err(error) => (eyre::Report::new(error), true),
195                },
196                Ok(response) => {
197                    let status = response.status();
198                    let transient =
199                        status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error();
200                    (eyre!("the metadata service answered {status}"), transient)
201                }
202                Err(error) => {
203                    let transient = !error.is_connect();
204                    (eyre::Report::new(error), transient)
205                }
206            };
207            if !transient || attempt == ATTEMPTS {
208                return Err(error);
209            }
210            tokio::time::sleep(self.backoff * 2u32.pow(attempt - 1)).await;
211            attempt += 1;
212        }
213    }
214}
215
216/// The JSON document under `security-credentials/<role>`.
217///
218/// Every field is optional because a failing lookup returns a document with a
219/// `Code` and none of the credentials.
220#[derive(Deserialize)]
221#[serde(rename_all = "PascalCase")]
222struct RoleDocument {
223    code: Option<String>,
224    access_key_id: Option<String>,
225    secret_access_key: Option<String>,
226    token: Option<String>,
227    expiration: Option<String>,
228}
229
230fn parse_role_document(document: &str) -> Result<TemporaryCredentials> {
231    // The body is never included in an error: it carries the secret.
232    let document: RoleDocument = serde_json::from_str(document).map_err(|_| {
233        eyre!("the metadata service returned a credentials document mbx cannot read")
234    })?;
235    if let Some(code) = document.code.as_deref().filter(|code| *code != "Success") {
236        bail!("the metadata service reported {code} for the instance role");
237    }
238    let (Some(access_key_id), Some(secret_access_key), Some(expiration)) = (
239        document.access_key_id,
240        document.secret_access_key,
241        document.expiration,
242    ) else {
243        bail!("the instance role's credentials document is missing fields");
244    };
245    Ok(TemporaryCredentials {
246        credentials: S3Credentials {
247            access_key_id,
248            secret_access_key,
249            session_token: document.token.filter(|token| !token.is_empty()),
250        },
251        expires_at: parse_timestamp(&expiration)?,
252    })
253}
254
255/// Parse the `2026-09-29T18:45:12Z` form the metadata service uses.
256///
257/// A fractional second is accepted and dropped. A numeric UTC offset is not,
258/// since the service always answers in UTC.
259fn parse_timestamp(text: &str) -> Result<SystemTime> {
260    let invalid = || eyre!("the instance role's expiration is not a UTC timestamp");
261    let text = text.trim().strip_suffix('Z').ok_or_else(invalid)?;
262    let (date, time) = text.split_once('T').ok_or_else(invalid)?;
263    let time = time.split_once('.').map_or(time, |(whole, _)| whole);
264    let number = |part: &str| part.parse::<i64>().map_err(|_| invalid());
265    let date = date.split('-').map(number).collect::<Result<Vec<_>>>()?;
266    let time = time.split(':').map(number).collect::<Result<Vec<_>>>()?;
267    let [year, month, day] = date[..].try_into().map_err(|_| invalid())?;
268    let [hour, minute, second] = time[..].try_into().map_err(|_| invalid())?;
269    if !(1..=12).contains(&month)
270        || !(1..=31).contains(&day)
271        || !(0..24).contains(&hour)
272        || !(0..60).contains(&minute)
273        || !(0..=60).contains(&second)
274    {
275        return Err(invalid());
276    }
277    let seconds = days_from_civil(year, month, day) * 86_400 + hour * 3_600 + minute * 60 + second;
278    u64::try_from(seconds)
279        .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds))
280        .map_err(|_| invalid())
281}
282
283/// The inverse of `civil_from_days` in `sigv4`, from the same source.
284fn days_from_civil(year: i64, month: i64, day: i64) -> i64 {
285    let year = if month <= 2 { year - 1 } else { year };
286    let era = if year >= 0 { year } else { year - 399 } / 400;
287    let year_of_era = year - era * 400;
288    let shifted_month = if month > 2 { month - 3 } else { month + 9 };
289    let day_of_year = (153 * shifted_month + 2) / 5 + day - 1;
290    let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year;
291    era * 146_097 + day_of_era - 719_468
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    fn provider(server: &mockito::ServerGuard) -> InstanceRoleCredentials {
299        InstanceRoleCredentials::with_timing(
300            server.url().parse().unwrap(),
301            Duration::from_secs(1),
302            Duration::from_millis(1),
303        )
304        .unwrap()
305    }
306
307    fn role_document(expiration: &str) -> String {
308        format!(
309            r#"{{"Code":"Success","LastUpdated":"2026-09-29T12:00:00Z","Type":"AWS-HMAC","AccessKeyId":"ASIAROLE","SecretAccessKey":"role-secret","Token":"role-token","Expiration":"{expiration}"}}"#
310        )
311    }
312
313    #[test]
314    fn a_timestamp_parses_to_the_instant_it_names() {
315        let parsed = parse_timestamp("2026-09-29T18:45:12Z").unwrap();
316        assert_eq!(parsed, UNIX_EPOCH + Duration::from_secs(1_790_707_512));
317        // The leap day and the epoch itself.
318        assert_eq!(parse_timestamp("1970-01-01T00:00:00Z").unwrap(), UNIX_EPOCH);
319        assert_eq!(
320            parse_timestamp("2024-02-29T00:00:00Z").unwrap(),
321            UNIX_EPOCH + Duration::from_secs(1_709_164_800)
322        );
323        assert_eq!(parse_timestamp("2026-09-29T18:45:12.500Z").unwrap(), parsed);
324    }
325
326    #[test]
327    fn a_malformed_timestamp_is_refused() {
328        for text in [
329            "",
330            "2026-09-29",
331            "2026-09-29T18:45:12",
332            "2026-09-29T18:45:12+02:00",
333            "2026-13-29T18:45:12Z",
334            "2026-09-29T25:45:12Z",
335            "2026-09-29T18:45Z",
336            "2026-09-29-01T18:45:12Z",
337            "yesterday",
338        ] {
339            assert!(parse_timestamp(text).is_err(), "{text:?} should be refused");
340        }
341    }
342
343    #[test]
344    fn a_role_document_becomes_credentials() {
345        let parsed = parse_role_document(&role_document("2026-09-29T18:45:12Z")).unwrap();
346        assert_eq!(parsed.credentials.access_key_id, "ASIAROLE");
347        assert_eq!(parsed.credentials.secret_access_key, "role-secret");
348        assert_eq!(
349            parsed.credentials.session_token.as_deref(),
350            Some("role-token")
351        );
352        assert_eq!(
353            parsed.expires_at,
354            UNIX_EPOCH + Duration::from_secs(1_790_707_512)
355        );
356    }
357
358    #[test]
359    fn a_failure_document_names_its_code_and_nothing_secret() {
360        let error = parse_role_document(r#"{"Code":"AssumeRoleUnauthorizedAccess"}"#)
361            .err()
362            .unwrap()
363            .to_string();
364        assert!(error.contains("AssumeRoleUnauthorizedAccess"));
365
366        let error = parse_role_document("not json with SECRET inside")
367            .err()
368            .unwrap()
369            .to_string();
370        assert!(!error.contains("SECRET"));
371    }
372
373    #[test]
374    fn the_environment_can_disable_or_redirect_the_lookup() {
375        let vars = |pairs: &'static [(&'static str, &'static str)]| {
376            move |name: &str| {
377                pairs
378                    .iter()
379                    .find(|(key, _)| *key == name)
380                    .map(|(_, value)| (*value).to_string())
381            }
382        };
383
384        assert!(
385            InstanceRoleCredentials::from_vars(vars(&[("AWS_EC2_METADATA_DISABLED", "TRUE")]))
386                .unwrap()
387                .is_none()
388        );
389        // Only "true" disables it, as in the SDKs.
390        assert!(
391            InstanceRoleCredentials::from_vars(vars(&[("AWS_EC2_METADATA_DISABLED", "false")]))
392                .unwrap()
393                .is_some()
394        );
395        let redirected = InstanceRoleCredentials::from_vars(vars(&[(
396            "AWS_EC2_METADATA_SERVICE_ENDPOINT",
397            "http://127.0.0.1:1338",
398        )]))
399        .unwrap()
400        .unwrap();
401        assert_eq!(redirected.endpoint.as_str(), "http://127.0.0.1:1338/");
402        assert!(
403            InstanceRoleCredentials::from_vars(vars(&[(
404                "AWS_EC2_METADATA_SERVICE_ENDPOINT",
405                "not a url",
406            )]))
407            .is_err()
408        );
409        let default = InstanceRoleCredentials::from_vars(vars(&[]))
410            .unwrap()
411            .unwrap();
412        assert_eq!(default.endpoint.as_str(), "http://169.254.169.254/");
413    }
414
415    #[tokio::test]
416    async fn the_three_requests_yield_the_roles_credentials() {
417        let mut server = mockito::Server::new_async().await;
418        let token = server
419            .mock("PUT", "/latest/api/token")
420            .match_header(TOKEN_TTL_HEADER, TOKEN_TTL_SECONDS)
421            .with_body("session-token\n")
422            .create_async()
423            .await;
424        let roles = server
425            .mock("GET", "/latest/meta-data/iam/security-credentials/")
426            .match_header(TOKEN_HEADER, "session-token")
427            .with_body("build-runner\n")
428            .create_async()
429            .await;
430        let credentials = server
431            .mock(
432                "GET",
433                "/latest/meta-data/iam/security-credentials/build-runner",
434            )
435            .match_header(TOKEN_HEADER, "session-token")
436            .with_body(role_document("2026-09-29T18:45:12Z"))
437            .create_async()
438            .await;
439
440        let fetched = provider(&server).fetch().await.unwrap();
441
442        assert_eq!(fetched.credentials.access_key_id, "ASIAROLE");
443        token.assert_async().await;
444        roles.assert_async().await;
445        credentials.assert_async().await;
446    }
447
448    #[tokio::test]
449    async fn the_instances_region_is_read_from_placement_metadata() {
450        let mut server = mockito::Server::new_async().await;
451        server
452            .mock("PUT", "/latest/api/token")
453            .with_body("session-token")
454            .create_async()
455            .await;
456        server
457            .mock("GET", "/latest/meta-data/placement/region")
458            .match_header(TOKEN_HEADER, "session-token")
459            .with_body("eu-west-1\n")
460            .create_async()
461            .await;
462
463        assert_eq!(provider(&server).region().await.unwrap(), "eu-west-1");
464    }
465
466    #[tokio::test]
467    async fn a_service_that_refuses_the_token_is_reported() {
468        let mut server = mockito::Server::new_async().await;
469        let token = server
470            .mock("PUT", "/latest/api/token")
471            .with_status(403)
472            .expect(1)
473            .create_async()
474            .await;
475
476        let error = format!("{:#}", provider(&server).fetch().await.err().unwrap());
477
478        assert!(error.contains("session token"), "{error}");
479        assert!(error.contains("403"), "{error}");
480        // A refusal is an answer, so it is not asked again.
481        token.assert_async().await;
482    }
483
484    #[tokio::test]
485    async fn throttling_and_server_errors_are_retried() {
486        for status in [429, 503] {
487            let mut server = mockito::Server::new_async().await;
488            let throttled = server
489                .mock("PUT", "/latest/api/token")
490                .with_status(status)
491                .expect(2)
492                .create_async()
493                .await;
494            let token = server
495                .mock("PUT", "/latest/api/token")
496                .with_body("session-token")
497                .expect(1)
498                .create_async()
499                .await;
500            server
501                .mock("GET", "/latest/meta-data/placement/region")
502                .with_body("eu-west-1")
503                .create_async()
504                .await;
505
506            assert_eq!(provider(&server).region().await.unwrap(), "eu-west-1");
507
508            throttled.assert_async().await;
509            token.assert_async().await;
510        }
511    }
512
513    #[tokio::test]
514    async fn retries_stop_after_the_last_attempt() {
515        let mut server = mockito::Server::new_async().await;
516        let token = server
517            .mock("PUT", "/latest/api/token")
518            .with_status(503)
519            .expect(ATTEMPTS as usize)
520            .create_async()
521            .await;
522
523        let error = format!("{:#}", provider(&server).fetch().await.err().unwrap());
524
525        assert!(error.contains("503"), "{error}");
526        token.assert_async().await;
527    }
528
529    #[tokio::test]
530    async fn a_token_request_that_connects_but_gets_no_answer_mentions_the_hop_limit() {
531        // Accepts connections and never replies, as a hop limit of 1 does.
532        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
533        let address = listener.local_addr().unwrap();
534        let held = tokio::spawn(async move {
535            let mut connections = Vec::new();
536            while let Ok((connection, _)) = listener.accept().await {
537                connections.push(connection);
538            }
539        });
540        let provider = InstanceRoleCredentials::with_timing(
541            format!("http://{address}").parse().unwrap(),
542            Duration::from_millis(50),
543            Duration::from_millis(1),
544        )
545        .unwrap();
546
547        let error = format!("{:#}", provider.fetch().await.err().unwrap());
548        held.abort();
549
550        assert!(error.contains("HttpPutResponseHopLimit"), "{error}");
551    }
552
553    #[tokio::test]
554    async fn an_instance_without_a_role_is_reported() {
555        let mut server = mockito::Server::new_async().await;
556        server
557            .mock("PUT", "/latest/api/token")
558            .with_body("session-token")
559            .create_async()
560            .await;
561        server
562            .mock("GET", "/latest/meta-data/iam/security-credentials/")
563            .with_body("")
564            .create_async()
565            .await;
566
567        let error = format!("{:#}", provider(&server).fetch().await.err().unwrap());
568
569        assert!(error.contains("no IAM role"), "{error}");
570    }
571
572    #[tokio::test]
573    async fn nothing_listening_fails_within_the_timeout() {
574        // Bind and drop to get a port nothing answers on.
575        let port = std::net::TcpListener::bind("127.0.0.1:0")
576            .unwrap()
577            .local_addr()
578            .unwrap()
579            .port();
580        let provider =
581            InstanceRoleCredentials::new(format!("http://127.0.0.1:{port}").parse().unwrap())
582                .unwrap();
583
584        let started = std::time::Instant::now();
585        assert!(provider.fetch().await.is_err());
586        assert!(started.elapsed() < Duration::from_secs(3));
587    }
588}