Skip to main content

reqsign_aws_core/provide_credential/
imds.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::Credential;
19use crate::provide_credential::utils::parse_imds_error;
20use bytes::Bytes;
21use http::Method;
22use http::header::CONTENT_LENGTH;
23use reqsign_core::time::Timestamp;
24use reqsign_core::{Context, Error, ProvideCredential, Result};
25use serde::Deserialize;
26use std::sync::{Arc, Mutex};
27use std::time::Duration;
28
29#[derive(Debug, Clone)]
30pub struct IMDSv2CredentialProvider {
31    endpoint: Option<String>,
32    token: Arc<Mutex<(String, Timestamp)>>,
33}
34
35impl Default for IMDSv2CredentialProvider {
36    fn default() -> Self {
37        Self {
38            endpoint: None,
39            token: Arc::new(Mutex::new((String::new(), Timestamp::default()))),
40        }
41    }
42}
43
44impl IMDSv2CredentialProvider {
45    /// Create a new `IMDSv2CredentialProvider` instance.
46    pub fn new() -> Self {
47        Self::default()
48    }
49
50    /// Set the endpoint for the metadata service.
51    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
52        self.endpoint = Some(endpoint.into());
53        self
54    }
55}
56
57impl IMDSv2CredentialProvider {
58    fn get_endpoint(&self, ctx: &Context) -> String {
59        // First check configured endpoint, then environment, then default
60        self.endpoint.clone().unwrap_or_else(|| {
61            ctx.env_vars()
62                .get("AWS_EC2_METADATA_SERVICE_ENDPOINT")
63                .cloned()
64                .unwrap_or_else(|| "http://169.254.169.254".into())
65        })
66    }
67
68    async fn load_ec2_metadata_token(&self, ctx: &Context) -> Result<String> {
69        {
70            let (token, expires_in) = self.token.lock().expect("lock poisoned").clone();
71            if expires_in > Timestamp::now() {
72                return Ok(token);
73            }
74        }
75
76        let endpoint = self.get_endpoint(ctx);
77        let url = format!("{endpoint}/latest/api/token");
78        let req = http::Request::builder()
79            .uri(&url)
80            .method(Method::PUT)
81            .header(CONTENT_LENGTH, "0")
82            // 21600s (6h) is recommended by AWS.
83            .header("x-aws-ec2-metadata-token-ttl-seconds", "21600")
84            .body(Bytes::new())
85            .map_err(|e| {
86                Error::request_invalid("failed to build IMDS token request")
87                    .with_source(e)
88                    .with_context(format!("url: {url}"))
89            })?;
90
91        let resp = ctx.http_send_as_string(req).await.map_err(|e| {
92            Error::unexpected("failed to connect to IMDS")
93                .with_source(e)
94                .with_context("endpoint: {endpoint}")
95                .with_context("hint: check if running on EC2 instance")
96                .set_retryable(true)
97        })?;
98
99        if resp.status() != http::StatusCode::OK {
100            return Err(parse_imds_error(
101                "fetch_imds_token",
102                resp.status(),
103                resp.body(),
104            ));
105        }
106        let ec2_token = resp.into_body();
107        // Set expires_in to 10 minutes to enforce re-read.
108        let expires_in = Timestamp::now() + Duration::from_secs(21600) - Duration::from_secs(600);
109
110        {
111            *self.token.lock().expect("lock poisoned") = (ec2_token.clone(), expires_in);
112        }
113
114        Ok(ec2_token)
115    }
116}
117impl ProvideCredential for IMDSv2CredentialProvider {
118    type Credential = Credential;
119
120    async fn provide_credential(&self, ctx: &Context) -> Result<Option<Self::Credential>> {
121        // Check if disabled via environment
122        let disabled_env = ctx
123            .env_vars()
124            .get("AWS_EC2_METADATA_DISABLED")
125            .map(|v| v == "true")
126            .unwrap_or(false);
127
128        if disabled_env {
129            return Ok(None);
130        }
131
132        let token = self.load_ec2_metadata_token(ctx).await?;
133
134        // List all credentials that node has.
135        let endpoint = self.get_endpoint(ctx);
136        let url = format!("{endpoint}/latest/meta-data/iam/security-credentials/");
137        let req = http::Request::builder()
138            .uri(&url)
139            .method(Method::GET)
140            // 21600s (6h) is recommended by AWS.
141            .header("x-aws-ec2-metadata-token", &token)
142            .body(Bytes::new())
143            .map_err(|e| {
144                Error::request_invalid("failed to build IMDS credentials list request")
145                    .with_source(e)
146                    .with_context("url: {url}")
147            })?;
148
149        let resp = ctx.http_send_as_string(req).await.map_err(|e| {
150            Error::unexpected("failed to list IMDS credentials")
151                .with_source(e)
152                .with_context("operation: list_instance_profiles")
153                .set_retryable(true)
154        })?;
155
156        if resp.status() != http::StatusCode::OK {
157            return Err(parse_imds_error(
158                "list_instance_profiles",
159                resp.status(),
160                resp.body(),
161            ));
162        }
163
164        let profile_name = resp.into_body();
165
166        if profile_name.is_empty() {
167            return Err(
168                Error::config_invalid("no IAM role attached to EC2 instance")
169                    .with_context("hint: attach an IAM role to your EC2 instance"),
170            );
171        }
172
173        // Get the credentials via role_name.
174        let endpoint = self.get_endpoint(ctx);
175        let url = format!("{endpoint}/latest/meta-data/iam/security-credentials/{profile_name}");
176        let req = http::Request::builder()
177            .uri(url)
178            .method(Method::GET)
179            // 21600s (6h) is recommended by AWS.
180            .header("x-aws-ec2-metadata-token", &token)
181            .body(Bytes::new())
182            .map_err(|e| {
183                Error::request_invalid("failed to build IMDS credentials fetch request")
184                    .with_source(e)
185                    .with_context(format!("profile: {profile_name}"))
186            })?;
187
188        let resp = ctx.http_send_as_string(req).await.map_err(|e| {
189            Error::unexpected("failed to fetch IMDS credentials")
190                .with_source(e)
191                .with_context(format!("profile: {profile_name}"))
192                .set_retryable(true)
193        })?;
194
195        if resp.status() != http::StatusCode::OK {
196            return Err(
197                parse_imds_error("fetch_credentials", resp.status(), resp.body())
198                    .with_context(format!("profile: {profile_name}")),
199            );
200        }
201
202        let content = resp.into_body();
203        let resp: Ec2MetadataIamSecurityCredentials =
204            serde_json::from_str(&content).map_err(|e| {
205                Error::unexpected("failed to parse IMDS credentials response")
206                    .with_source(e)
207                    .with_context(format!("response_length: {}", content.len()))
208                    .with_context(format!("profile: {profile_name}"))
209            })?;
210
211        // Check for specific error codes
212        match resp.code.as_str() {
213            "Success" => {} // Continue processing
214            "AssumeRoleUnauthorizedAccess" => {
215                return Err(Error::permission_denied(format!(
216                    "EC2 instance not authorized to assume role: {}",
217                    resp.message
218                ))
219                .with_context(format!("error_code: {}", resp.code))
220                .with_context(format!("profile: {profile_name}"))
221                .with_context("hint: check if the IAM role has a trust relationship with EC2"));
222            }
223            code if code.contains("Expired") => {
224                return Err(Error::credential_invalid(format!(
225                    "IMDS credentials expired: {}",
226                    resp.message
227                ))
228                .with_context(format!("error_code: {}", resp.code))
229                .with_context(format!("profile: {profile_name}")));
230            }
231            _ => {
232                return Err(Error::unexpected(format!(
233                    "IMDS returned error: [{}] {}",
234                    resp.code, resp.message
235                ))
236                .with_context(format!("profile: {profile_name}")));
237            }
238        }
239
240        let cred = Credential {
241            access_key_id: resp.access_key_id,
242            secret_access_key: resp.secret_access_key,
243            session_token: Some(resp.token),
244            expires_in: Some(resp.expiration.parse().map_err(|e| {
245                Error::unexpected("failed to parse IMDS credential expiration time")
246                    .with_source(e)
247                    .with_context(format!("expiration_value: {}", resp.expiration))
248            })?),
249        };
250
251        Ok(Some(cred))
252    }
253}
254
255#[derive(Default, Debug, Deserialize)]
256#[serde(default, rename_all = "PascalCase")]
257struct Ec2MetadataIamSecurityCredentials {
258    access_key_id: String,
259    secret_access_key: String,
260    token: String,
261    expiration: String,
262
263    code: String,
264    message: String,
265}