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
26const TOKEN_TTL_HEADER: &str = "x-aws-ec2-metadata-token-ttl-seconds";
27const TOKEN_HEADER: &str = "x-aws-ec2-metadata-token";
28
29/// Temporary credentials and the instant they stop working.
30pub struct TemporaryCredentials {
31    /// The credentials to sign with.
32    pub credentials: S3Credentials,
33    /// When the service will start refusing them.
34    pub expires_at: SystemTime,
35}
36
37/// A source of credentials for the EC2 instance role attached to this machine.
38pub struct InstanceRoleCredentials {
39    client: reqwest::Client,
40    endpoint: Url,
41}
42
43impl InstanceRoleCredentials {
44    /// A provider for the metadata service at `endpoint`.
45    pub fn new(endpoint: Url) -> Result<Self> {
46        let client = reqwest::Client::builder()
47            .connect_timeout(REQUEST_TIMEOUT)
48            .timeout(REQUEST_TIMEOUT * 2)
49            .redirect(reqwest::redirect::Policy::none())
50            // A link-local address is never behind a proxy, and sending role
51            // credentials through one would hand them to it.
52            .no_proxy()
53            .build()?;
54        Ok(Self { client, endpoint })
55    }
56
57    /// A provider configured the way the AWS SDKs read it, or `None` when
58    /// `AWS_EC2_METADATA_DISABLED` turns instance role lookup off.
59    ///
60    /// `AWS_EC2_METADATA_SERVICE_ENDPOINT` overrides the address, for a proxy
61    /// in front of the service or a test double.
62    pub fn from_env() -> Result<Option<Self>> {
63        Self::from_vars(|name| std::env::var(name).ok())
64    }
65
66    fn from_vars(var: impl Fn(&str) -> Option<String>) -> Result<Option<Self>> {
67        if var("AWS_EC2_METADATA_DISABLED")
68            .is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
69        {
70            return Ok(None);
71        }
72        let endpoint = var("AWS_EC2_METADATA_SERVICE_ENDPOINT")
73            .map(|value| value.trim().to_string())
74            .filter(|value| !value.is_empty());
75        let endpoint = endpoint.as_deref().unwrap_or(DEFAULT_ENDPOINT);
76        let endpoint = endpoint
77            .parse()
78            .wrap_err("invalid AWS_EC2_METADATA_SERVICE_ENDPOINT")?;
79        Self::new(endpoint).map(Some)
80    }
81
82    /// Ask the metadata service for the role's current credentials.
83    ///
84    /// The service hands out the same credentials until about five minutes
85    /// before they expire, so calling this early returns what a caller already
86    /// holds.
87    pub async fn fetch(&self) -> Result<TemporaryCredentials> {
88        let token = self
89            .token()
90            .await
91            .wrap_err("could not get an IMDSv2 session token")?;
92        let roles = self
93            .get(&token, "latest/meta-data/iam/security-credentials/")
94            .await
95            .wrap_err("could not list the instance's IAM roles")?;
96        let role = roles
97            .lines()
98            .map(str::trim)
99            .find(|line| !line.is_empty())
100            .ok_or_else(|| eyre!("the instance has no IAM role attached"))?;
101        let document = self
102            .get(
103                &token,
104                &format!("latest/meta-data/iam/security-credentials/{role}"),
105            )
106            .await
107            .wrap_err_with(|| format!("could not read credentials for role {role}"))?;
108        parse_role_document(&document)
109    }
110
111    /// The region the instance runs in, for a machine that names none itself.
112    ///
113    /// The metadata service answers this without any role, so it is a separate
114    /// lookup that only runs when the region would otherwise be missing.
115    pub async fn region(&self) -> Result<String> {
116        let token = self
117            .token()
118            .await
119            .wrap_err("could not get an IMDSv2 session token")?;
120        let region = self
121            .get(&token, "latest/meta-data/placement/region")
122            .await
123            .wrap_err("could not read the instance's region")?;
124        let region = region.trim();
125        if region.is_empty() {
126            bail!("the metadata service returned an empty region");
127        }
128        Ok(region.to_string())
129    }
130
131    async fn token(&self) -> Result<String> {
132        let response = self
133            .client
134            .put(self.endpoint.join("latest/api/token")?)
135            .header(TOKEN_TTL_HEADER, TOKEN_TTL_SECONDS)
136            .send()
137            .await?;
138        expect_success(response.status())?;
139        Ok(response.text().await?.trim().to_string())
140    }
141
142    async fn get(&self, token: &str, path: &str) -> Result<String> {
143        let response = self
144            .client
145            .get(self.endpoint.join(path)?)
146            .header(TOKEN_HEADER, token)
147            .send()
148            .await?;
149        expect_success(response.status())?;
150        Ok(response.text().await?)
151    }
152}
153
154fn expect_success(status: StatusCode) -> Result<()> {
155    if status.is_success() {
156        Ok(())
157    } else {
158        bail!("the metadata service answered {status}")
159    }
160}
161
162/// The JSON document under `security-credentials/<role>`.
163///
164/// Every field is optional because a failing lookup returns a document with a
165/// `Code` and none of the credentials.
166#[derive(Deserialize)]
167#[serde(rename_all = "PascalCase")]
168struct RoleDocument {
169    code: Option<String>,
170    access_key_id: Option<String>,
171    secret_access_key: Option<String>,
172    token: Option<String>,
173    expiration: Option<String>,
174}
175
176fn parse_role_document(document: &str) -> Result<TemporaryCredentials> {
177    // The body is never included in an error: it carries the secret.
178    let document: RoleDocument = serde_json::from_str(document).map_err(|_| {
179        eyre!("the metadata service returned a credentials document mbx cannot read")
180    })?;
181    if let Some(code) = document.code.as_deref().filter(|code| *code != "Success") {
182        bail!("the metadata service reported {code} for the instance role");
183    }
184    let (Some(access_key_id), Some(secret_access_key), Some(expiration)) = (
185        document.access_key_id,
186        document.secret_access_key,
187        document.expiration,
188    ) else {
189        bail!("the instance role's credentials document is missing fields");
190    };
191    Ok(TemporaryCredentials {
192        credentials: S3Credentials {
193            access_key_id,
194            secret_access_key,
195            session_token: document.token.filter(|token| !token.is_empty()),
196        },
197        expires_at: parse_timestamp(&expiration)?,
198    })
199}
200
201/// Parse the `2026-09-29T18:45:12Z` form the metadata service uses.
202///
203/// A fractional second is accepted and dropped. A numeric UTC offset is not,
204/// since the service always answers in UTC.
205fn parse_timestamp(text: &str) -> Result<SystemTime> {
206    let invalid = || eyre!("the instance role's expiration is not a UTC timestamp");
207    let text = text.trim().strip_suffix('Z').ok_or_else(invalid)?;
208    let (date, time) = text.split_once('T').ok_or_else(invalid)?;
209    let time = time.split_once('.').map_or(time, |(whole, _)| whole);
210    let number = |part: &str| part.parse::<i64>().map_err(|_| invalid());
211    let date = date.split('-').map(number).collect::<Result<Vec<_>>>()?;
212    let time = time.split(':').map(number).collect::<Result<Vec<_>>>()?;
213    let [year, month, day] = date[..].try_into().map_err(|_| invalid())?;
214    let [hour, minute, second] = time[..].try_into().map_err(|_| invalid())?;
215    if !(1..=12).contains(&month)
216        || !(1..=31).contains(&day)
217        || !(0..24).contains(&hour)
218        || !(0..60).contains(&minute)
219        || !(0..=60).contains(&second)
220    {
221        return Err(invalid());
222    }
223    let seconds = days_from_civil(year, month, day) * 86_400 + hour * 3_600 + minute * 60 + second;
224    u64::try_from(seconds)
225        .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds))
226        .map_err(|_| invalid())
227}
228
229/// The inverse of `civil_from_days` in `sigv4`, from the same source.
230fn days_from_civil(year: i64, month: i64, day: i64) -> i64 {
231    let year = if month <= 2 { year - 1 } else { year };
232    let era = if year >= 0 { year } else { year - 399 } / 400;
233    let year_of_era = year - era * 400;
234    let shifted_month = if month > 2 { month - 3 } else { month + 9 };
235    let day_of_year = (153 * shifted_month + 2) / 5 + day - 1;
236    let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year;
237    era * 146_097 + day_of_era - 719_468
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    fn provider(server: &mockito::ServerGuard) -> InstanceRoleCredentials {
245        InstanceRoleCredentials::new(server.url().parse().unwrap()).unwrap()
246    }
247
248    fn role_document(expiration: &str) -> String {
249        format!(
250            r#"{{"Code":"Success","LastUpdated":"2026-09-29T12:00:00Z","Type":"AWS-HMAC","AccessKeyId":"ASIAROLE","SecretAccessKey":"role-secret","Token":"role-token","Expiration":"{expiration}"}}"#
251        )
252    }
253
254    #[test]
255    fn a_timestamp_parses_to_the_instant_it_names() {
256        let parsed = parse_timestamp("2026-09-29T18:45:12Z").unwrap();
257        assert_eq!(parsed, UNIX_EPOCH + Duration::from_secs(1_790_707_512));
258        // The leap day and the epoch itself.
259        assert_eq!(parse_timestamp("1970-01-01T00:00:00Z").unwrap(), UNIX_EPOCH);
260        assert_eq!(
261            parse_timestamp("2024-02-29T00:00:00Z").unwrap(),
262            UNIX_EPOCH + Duration::from_secs(1_709_164_800)
263        );
264        assert_eq!(parse_timestamp("2026-09-29T18:45:12.500Z").unwrap(), parsed);
265    }
266
267    #[test]
268    fn a_malformed_timestamp_is_refused() {
269        for text in [
270            "",
271            "2026-09-29",
272            "2026-09-29T18:45:12",
273            "2026-09-29T18:45:12+02:00",
274            "2026-13-29T18:45:12Z",
275            "2026-09-29T25:45:12Z",
276            "2026-09-29T18:45Z",
277            "2026-09-29-01T18:45:12Z",
278            "yesterday",
279        ] {
280            assert!(parse_timestamp(text).is_err(), "{text:?} should be refused");
281        }
282    }
283
284    #[test]
285    fn a_role_document_becomes_credentials() {
286        let parsed = parse_role_document(&role_document("2026-09-29T18:45:12Z")).unwrap();
287        assert_eq!(parsed.credentials.access_key_id, "ASIAROLE");
288        assert_eq!(parsed.credentials.secret_access_key, "role-secret");
289        assert_eq!(
290            parsed.credentials.session_token.as_deref(),
291            Some("role-token")
292        );
293        assert_eq!(
294            parsed.expires_at,
295            UNIX_EPOCH + Duration::from_secs(1_790_707_512)
296        );
297    }
298
299    #[test]
300    fn a_failure_document_names_its_code_and_nothing_secret() {
301        let error = parse_role_document(r#"{"Code":"AssumeRoleUnauthorizedAccess"}"#)
302            .err()
303            .unwrap()
304            .to_string();
305        assert!(error.contains("AssumeRoleUnauthorizedAccess"));
306
307        let error = parse_role_document("not json with SECRET inside")
308            .err()
309            .unwrap()
310            .to_string();
311        assert!(!error.contains("SECRET"));
312    }
313
314    #[test]
315    fn the_environment_can_disable_or_redirect_the_lookup() {
316        let vars = |pairs: &'static [(&'static str, &'static str)]| {
317            move |name: &str| {
318                pairs
319                    .iter()
320                    .find(|(key, _)| *key == name)
321                    .map(|(_, value)| (*value).to_string())
322            }
323        };
324
325        assert!(
326            InstanceRoleCredentials::from_vars(vars(&[("AWS_EC2_METADATA_DISABLED", "TRUE")]))
327                .unwrap()
328                .is_none()
329        );
330        // Only "true" disables it, as in the SDKs.
331        assert!(
332            InstanceRoleCredentials::from_vars(vars(&[("AWS_EC2_METADATA_DISABLED", "false")]))
333                .unwrap()
334                .is_some()
335        );
336        let redirected = InstanceRoleCredentials::from_vars(vars(&[(
337            "AWS_EC2_METADATA_SERVICE_ENDPOINT",
338            "http://127.0.0.1:1338",
339        )]))
340        .unwrap()
341        .unwrap();
342        assert_eq!(redirected.endpoint.as_str(), "http://127.0.0.1:1338/");
343        assert!(
344            InstanceRoleCredentials::from_vars(vars(&[(
345                "AWS_EC2_METADATA_SERVICE_ENDPOINT",
346                "not a url",
347            )]))
348            .is_err()
349        );
350        let default = InstanceRoleCredentials::from_vars(vars(&[]))
351            .unwrap()
352            .unwrap();
353        assert_eq!(default.endpoint.as_str(), "http://169.254.169.254/");
354    }
355
356    #[tokio::test]
357    async fn the_three_requests_yield_the_roles_credentials() {
358        let mut server = mockito::Server::new_async().await;
359        let token = server
360            .mock("PUT", "/latest/api/token")
361            .match_header(TOKEN_TTL_HEADER, TOKEN_TTL_SECONDS)
362            .with_body("session-token\n")
363            .create_async()
364            .await;
365        let roles = server
366            .mock("GET", "/latest/meta-data/iam/security-credentials/")
367            .match_header(TOKEN_HEADER, "session-token")
368            .with_body("build-runner\n")
369            .create_async()
370            .await;
371        let credentials = server
372            .mock(
373                "GET",
374                "/latest/meta-data/iam/security-credentials/build-runner",
375            )
376            .match_header(TOKEN_HEADER, "session-token")
377            .with_body(role_document("2026-09-29T18:45:12Z"))
378            .create_async()
379            .await;
380
381        let fetched = provider(&server).fetch().await.unwrap();
382
383        assert_eq!(fetched.credentials.access_key_id, "ASIAROLE");
384        token.assert_async().await;
385        roles.assert_async().await;
386        credentials.assert_async().await;
387    }
388
389    #[tokio::test]
390    async fn the_instances_region_is_read_from_placement_metadata() {
391        let mut server = mockito::Server::new_async().await;
392        server
393            .mock("PUT", "/latest/api/token")
394            .with_body("session-token")
395            .create_async()
396            .await;
397        server
398            .mock("GET", "/latest/meta-data/placement/region")
399            .match_header(TOKEN_HEADER, "session-token")
400            .with_body("eu-west-1\n")
401            .create_async()
402            .await;
403
404        assert_eq!(provider(&server).region().await.unwrap(), "eu-west-1");
405    }
406
407    #[tokio::test]
408    async fn a_service_that_refuses_the_token_is_reported() {
409        let mut server = mockito::Server::new_async().await;
410        server
411            .mock("PUT", "/latest/api/token")
412            .with_status(403)
413            .create_async()
414            .await;
415
416        let error = format!("{:#}", provider(&server).fetch().await.err().unwrap());
417
418        assert!(error.contains("session token"), "{error}");
419        assert!(error.contains("403"), "{error}");
420    }
421
422    #[tokio::test]
423    async fn an_instance_without_a_role_is_reported() {
424        let mut server = mockito::Server::new_async().await;
425        server
426            .mock("PUT", "/latest/api/token")
427            .with_body("session-token")
428            .create_async()
429            .await;
430        server
431            .mock("GET", "/latest/meta-data/iam/security-credentials/")
432            .with_body("")
433            .create_async()
434            .await;
435
436        let error = format!("{:#}", provider(&server).fetch().await.err().unwrap());
437
438        assert!(error.contains("no IAM role"), "{error}");
439    }
440
441    #[tokio::test]
442    async fn nothing_listening_fails_within_the_timeout() {
443        // Bind and drop to get a port nothing answers on.
444        let port = std::net::TcpListener::bind("127.0.0.1:0")
445            .unwrap()
446            .local_addr()
447            .unwrap()
448            .port();
449        let provider =
450            InstanceRoleCredentials::new(format!("http://127.0.0.1:{port}").parse().unwrap())
451                .unwrap();
452
453        let started = std::time::Instant::now();
454        assert!(provider.fetch().await.is_err());
455        assert!(started.elapsed() < Duration::from_secs(3));
456    }
457}