async_profiler_agent/metadata/
aws.rs1use reqwest::Method;
9use serde::Deserialize;
10use thiserror::Error;
11
12use super::AgentMetadata;
13
14#[derive(Error, Debug)]
17pub enum FargateMetadataToAgentMetadataError {
18 #[error("unable to parse task ARN as a valid ARN")]
20 TaskArnInvalid(#[from] aws_arn::Error),
21 #[error("AWS account id not found in Fargate metadata")]
23 AccountIdNotFound,
24 #[error("AWS region not found in Fargate metadata")]
26 AwsRegionNotFound,
27}
28
29#[derive(Error, Debug)]
31#[error("profiler metadata error: {0}")]
32pub enum AwsProfilerMetadataError {
33 #[error("failed to create profiler metadata file: {0}")]
35 FailedToCreateFile(#[from] std::io::Error),
36
37 #[error("failed fetching valid Fargate metadata: {0}")]
39 FargateMetadataToAgentMetadataError(#[from] FargateMetadataToAgentMetadataError),
40
41 #[error("retrieved invalid endpoint URI from ECS_CONTAINER_METADATA_URI_V4: {0}")]
43 InvalidUri(String),
44
45 #[error("failed to fetch metadata from endpoint over HTTP: {0}")]
47 FailedToFetchMetadataFromEndpoint(reqwest::Error),
48
49 #[error("failed to fetch metadata from IMDS endpoint over HTTP: {0}")]
51 FailedToFetchMetadataFromImds(#[from] aws_config::imds::client::error::ImdsError),
52
53 #[error("failed to parse metadata as valid UTF-8 from endpoint over HTTP: {0}")]
56 FailedToParseMetadataFromEndpoint(reqwest::Error),
57
58 #[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 let uri = format!("{md_uri}/task",);
102
103 let req = http_client
104 .request(Method::GET, uri.clone())
105 .build()
106 .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
151pub 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 #[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}