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

async_profiler_agent/metadata/
aws.rs
