async_profiler_agent/metadata/
aws.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Contains functions for getting host metadata from [IMDS]
5//!
6//! [IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
7
8use reqwest::Method;
9use serde::Deserialize;
10use thiserror::Error;
11
12use super::AgentMetadata;
13
14/// An error converting Fargate IMDS metadata to Agent metadata. This error
15/// should probably not happen except in case of a bug in either this crate or IMDS.
16#[derive(Error, Debug)]
17pub enum FargateMetadataToAgentMetadataError {
18    /// unable to parse task ARN as a valid ARN
19    #[error("unable to parse task ARN as a valid ARN")]
20    TaskArnInvalid(#[from] aws_arn::Error),
21    /// AWS account id not found in Fargate metadata
22    #[error("AWS account id not found in Fargate metadata")]
23    AccountIdNotFound,
24    /// AWS region not found in Fargate metadata
25    #[error("AWS region not found in Fargate metadata")]
26    AwsRegionNotFound,
27}
28
29/// An error getting IMDS metadata
30#[derive(Error, Debug)]
31#[error("profiler metadata error: {0}")]
32pub enum AwsProfilerMetadataError {
33    /// Internal IO error
34    #[error("failed to create profiler metadata file: {0}")]
35    FailedToCreateFile(#[from] std::io::Error),
36
37    /// Error parsing IMDS metadata. Should normally not happen except in case of a bug
38    #[error("failed fetching valid Fargate metadata: {0}")]
39    FargateMetadataToAgentMetadataError(#[from] FargateMetadataToAgentMetadataError),
40
41    /// Invalid endpoint URI in `ECS_CONTAINER_METADATA_URI_V4`
42    #[error("retrieved invalid endpoint URI from ECS_CONTAINER_METADATA_URI_V4: {0}")]
43    InvalidUri(String),
44
45    /// Failed to fetch metadata from FarGate endpoint
46    #[error("failed to fetch metadata from endpoint over HTTP: {0}")]
47    FailedToFetchMetadataFromEndpoint(reqwest::Error),
48
49    /// Failed to fetch metadata from IMDS
50    #[error("failed to fetch metadata from IMDS endpoint over HTTP: {0}")]
51    FailedToFetchMetadataFromImds(#[from] aws_config::imds::client::error::ImdsError),
52
53    /// Failed to parse metadata from IMDS - this indicates a bug in this crate
54    /// or in IMDS
55    #[error("failed to parse metadata as valid UTF-8 from endpoint over HTTP: {0}")]
56    FailedToParseMetadataFromEndpoint(reqwest::Error),
57
58    /// Failed to serialize metadata file - this indicates a bug in this crate
59    /// or in IMDS
60    #[error("failed to serialize metadata file: {0}")]
61    FailedToSerializeMetadataFile(#[from] serde_json::Error),
62}
63
64#[derive(Debug, Deserialize, PartialEq, Eq)]
65#[serde(rename_all = "camelCase")]
66struct ImdsEc2InstanceMetadata {
67    account_id: String,
68    region: String,
69    #[allow(dead_code)]
70    instance_type: String,
71    instance_id: String,
72}
73
74async fn read_ec2_metadata() -> Result<ImdsEc2InstanceMetadata, AwsProfilerMetadataError> {
75    let imds = aws_config::imds::Client::builder().build();
76    let imds_document = imds
77        .get("/latest/dynamic/instance-identity/document")
78        .await?;
79
80    Ok(serde_json::from_str(imds_document.as_ref())?)
81}
82
83#[derive(Deserialize, Debug, PartialEq, Eq)]
84struct FargateMetadata {
85    #[serde(rename = "Cluster")]
86    cluster: String,
87    #[serde(rename = "TaskARN")]
88    task_arn: String,
89}
90
91async fn read_fargate_metadata(
92    http_client: &reqwest::Client,
93) -> Result<FargateMetadata, AwsProfilerMetadataError> {
94    let Ok(md_uri) = std::env::var("ECS_CONTAINER_METADATA_URI_V4") else {
95        return Err(AwsProfilerMetadataError::InvalidUri(
96            "not running on fargate".into(),
97        ));
98    };
99    // Only available if running on Fargate. Something like
100    // `http://169.254.232.106/v4/5261e761e0e2a3d92da3f02c8e5bab1f-3356750833`.
101    let uri = format!("{md_uri}/task",);
102
103    let req = http_client
104        .request(Method::GET, uri.clone())
105        .build()
106        // The only thing that can be invalid about this request is necessarily the URI.
107        .map_err(|_e| AwsProfilerMetadataError::InvalidUri(uri))?;
108    let res = http_client
109        .execute(req)
110        .await
111        .map_err(AwsProfilerMetadataError::FailedToFetchMetadataFromEndpoint)?;
112    let body_str = res
113        .text()
114        .await
115        .map_err(AwsProfilerMetadataError::FailedToParseMetadataFromEndpoint)?;
116
117    Ok(serde_json::from_str(&body_str)?)
118}
119
120impl super::AgentMetadata {
121    fn from_imds_ec2_instance_metadata(
122        imds_ec2_instance_metadata: ImdsEc2InstanceMetadata,
123    ) -> Self {
124        Self::Ec2AgentMetadata {
125            aws_account_id: imds_ec2_instance_metadata.account_id,
126            aws_region_id: imds_ec2_instance_metadata.region,
127            ec2_instance_id: imds_ec2_instance_metadata.instance_id,
128        }
129    }
130
131    fn try_from_fargate_metadata(
132        fargate_metadata: FargateMetadata,
133    ) -> Result<Self, FargateMetadataToAgentMetadataError> {
134        let ecs_task_arn: aws_arn::ResourceName = fargate_metadata.task_arn.parse()?;
135
136        Ok(Self::FargateAgentMetadata {
137            aws_account_id: ecs_task_arn
138                .account_id
139                .ok_or(FargateMetadataToAgentMetadataError::AccountIdNotFound)?
140                .to_string(),
141            aws_region_id: ecs_task_arn
142                .region
143                .ok_or(FargateMetadataToAgentMetadataError::AwsRegionNotFound)?
144                .to_string(),
145            ecs_task_arn: fargate_metadata.task_arn,
146            ecs_cluster_arn: fargate_metadata.cluster,
147        })
148    }
149}
150
151/// Load agent metadata from [Fargate] or [IMDS].
152///
153/// This will return an error if this machine does not appear to be a [Fargate] or [EC2].
154///
155/// [Fargate]: https://aws.amazon.com/fargate
156/// [IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
157/// [EC2]: https://aws.amazon.com/ec2
158pub async fn load_agent_metadata() -> Result<AgentMetadata, AwsProfilerMetadataError> {
159    let agent_metadata: AgentMetadata = match read_ec2_metadata().await {
160        Ok(imds_ec2_instance_metadata) => {
161            AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata)
162        }
163        Err(_) => {
164            let http_client = reqwest::Client::new();
165            let fargate_metadata = read_fargate_metadata(&http_client).await?;
166
167            AgentMetadata::try_from_fargate_metadata(fargate_metadata)?
168        }
169    };
170    Ok(agent_metadata)
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    // these constants are "anonymized" (aka randomly generated in that format)
178
179    // AMI_ID = 52d4b310b0a459ff
180    // INSTANCE_ID = 92eba08c089f6325
181    // ACCOUNT_ID = 123456789012
182    // IP = 10.65.149.216
183    // DATE = 2025-03-20T16:41:24Z
184    // DATE1 = 2025-03-20T16:41:24.713942268Z
185    // DATE2 = 2025-03-20T16:41:25.623883595Z
186    // TASK_ARN = 5261e761e0e2a3d92da3f02c8e5bab1f
187    // DOCKER_ID = 3356750833
188    // PRIV_IP = 169.254.232.106
189    // IMAGE_ID = ad9d89a36c31afef34c79e05263b06087ad354796cfd90c66ced30f40ea2dbf4
190    // TASK_UUID = f4094744-1b40-4701-9f26-ad84ebb709d7
191    // CLOCK_ERROR = 0.3148955
192
193    #[test]
194    fn test_imds_ec2_metadata() {
195        let json_str = r#"
196{
197    "accountId" : "123456789012",
198    "architecture" : "x86_64",
199    "availabilityZone" : "eu-west-1b",
200    "billingProducts" : null,
201    "devpayProductCodes" : null,
202    "marketplaceProductCodes" : null,
203    "imageId" : "ami-052d4b310b0a459ff",
204    "instanceId" : "i-092eba08c089f6325",
205    "instanceType" : "c5.4xlarge",
206    "kernelId" : null,
207    "pendingTime" : "2025-03-20T16:41:24Z",
208    "privateIp" : "10.65.149.216",
209    "ramdiskId" : null,
210    "region" : "eu-west-1",
211    "version" : "2017-09-30"
212}"#;
213
214        let imds_ec2_instance_metadata: ImdsEc2InstanceMetadata =
215            serde_json::from_str(&json_str).unwrap();
216
217        assert_eq!(
218            imds_ec2_instance_metadata,
219            ImdsEc2InstanceMetadata {
220                account_id: "123456789012".to_owned(),
221                region: "eu-west-1".to_owned(),
222                instance_type: "c5.4xlarge".to_owned(),
223                instance_id: "i-092eba08c089f6325".to_owned(),
224            }
225        );
226
227        let agent_metadata =
228            AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata);
229
230        assert_eq!(
231            agent_metadata,
232            AgentMetadata::Ec2AgentMetadata {
233                aws_account_id: "123456789012".to_owned(),
234                aws_region_id: "eu-west-1".to_owned(),
235                ec2_instance_id: "i-092eba08c089f6325".to_owned(),
236            }
237        )
238    }
239
240    #[test]
241    fn test_fargate_metadata() {
242        let json_str = r#"
243{
244    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
245    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
246    "Family": "profiler-metadata",
247    "Revision": "1",
248    "DesiredStatus": "RUNNING",
249    "KnownStatus": "NONE",
250    "Limits": {
251        "CPU": 0.25,
252        "Memory": 2048
253    },
254    "PullStartedAt": "2025-03-20T16:41:24.713942268Z",
255    "PullStoppedAt": "2025-03-20T16:41:25.623883595Z",
256    "AvailabilityZone": "us-east-1f",
257    "LaunchType": "FARGATE",
258    "Containers": [
259        {
260            "DockerId": "5261e761e0e2a3d92da3f02c8e5bab1f-3356750833",
261            "Name": "profiler-metadata",
262            "DockerName": "profiler-metadata",
263            "Image": "123456789012.dkr.ecr.us-east-1.amazonaws.com/profiler-metadata",
264            "ImageID": "sha256:ad9d89a36c31afef34c79e05263b06087ad354796cfd90c66ced30f40ea2dbf4",
265            "Labels": {
266                "com.amazonaws.ecs.cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
267                "com.amazonaws.ecs.container-name": "profiler-metadata",
268                "com.amazonaws.ecs.task-arn": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
269                "com.amazonaws.ecs.task-definition-family": "profiler-metadata",
270                "com.amazonaws.ecs.task-definition-version": "1"
271            },
272            "DesiredStatus": "RUNNING",
273            "KnownStatus": "PULLED",
274            "Limits": {
275                "CPU": 0
276            },
277            "Type": "NORMAL",
278            "LogDriver": "awslogs",
279            "LogOptions": {
280                "awslogs-create-group": "true",
281                "awslogs-group": "/ecs/profiler-metadata",
282                "awslogs-region": "us-east-1",
283                "awslogs-stream": "ecs/profiler-metadata/5261e761e0e2a3d92da3f02c8e5bab1f",
284                "max-buffer-size": "25m",
285                "mode": "non-blocking"
286            },
287            "ContainerARN": "arn:aws:ecs:us-east-1:123456789012:container/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f/f4094744-1b40-4701-9f26-ad84ebb709d7",
288            "Networks": [
289                {
290                    "NetworkMode": "awsvpc",
291                    "IPv4Addresses": [
292                        "172.31.233.169"
293                    ],
294                    "AttachmentIndex": 0,
295                    "MACAddress": "16:ff:d6:e1:dc:99",
296                    "IPv4SubnetCIDRBlock": "172.31.192.0/20",
297                    "DomainNameServers": [
298                        "172.31.0.2"
299                    ],
300                    "DomainNameSearchList": [
301                        "ec2.internal"
302                    ],
303                    "PrivateDNSName": "ip-172-31-233-169.ec2.internal",
304                    "SubnetGatewayIpv4Address": "172.31.192.1/20"
305                }
306            ],
307            "Snapshotter": "overlayfs"
308        }
309    ],
310    "ClockDrift": {
311        "ClockErrorBound": 0.3148955,
312        "ReferenceTimestamp": "2025-03-20T16:41:24Z",
313        "ClockSynchronizationStatus": "SYNCHRONIZED"
314    },
315    "EphemeralStorageMetrics": {
316        "Utilized": 208,
317        "Reserved": 20496
318    }
319}
320"#;
321
322        let fargate_metadata: FargateMetadata = serde_json::from_str(&json_str).unwrap();
323
324        assert_eq!(
325            fargate_metadata,
326            FargateMetadata {
327                cluster: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster"
328                    .to_owned(),
329                task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
330            }
331        );
332
333        let agent_metadata = AgentMetadata::try_from_fargate_metadata(fargate_metadata).unwrap();
334
335        assert_eq!(
336            agent_metadata,
337            AgentMetadata::FargateAgentMetadata {
338                aws_account_id: "123456789012".to_owned(),
339                aws_region_id: "us-east-1".to_owned(),
340                ecs_task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
341                ecs_cluster_arn: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster".to_owned(),
342            }
343        )
344    }
345}