Use AWS_REGION and IMDS to determine region (awslabs#422)

* Add support for AWS_REGION to determine region

Determine the region using the following sources (in order):
  * `--region` flag,
  * `AWS_REGION` environment variables,
  * default region (us-east-1).

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

* Add get_identity_document to the IMDS client

Extend the IMDS client to query the identity document of the current EC2 instance. The identity document is returned as a json response, which is parsed (using serde) to extract  the region and the instance type. The new `get_identity_document()` is tested by comparing its result to those obtained via the SDK.

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

* Use AWS_REGION and IMDS to determine region

Determine the region using the following sources (in order):
  * `--region` flag,
  * `AWS_REGION` environment variables,
  * EC2 instance region (using the IMDS client),
  * default region (us-east-1).

In order to avoid querying the IMDS twice, this change also refactors the retrieval of the instance type (used in the calculation of the maximum network throughput). The new type `InstanceInfo` will lazily query the IMDS client for both value if and when required.

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

* Move InstanceInfo into a separate module

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

* Refactor all region handling logic into get_region

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

* Remove unused methods in ImdsCrtClient and ImdsClient

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

---------

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

mountpoint-s3-client/Cargo.toml

@@ -33,6 +33,7 @@ thiserror = "1.0.34"
 time = { version = "0.3.17", features = ["formatting", "parsing"] }
 tracing = { version = "0.1.35", default-features = false, features = ["std", "log"] }
 xmltree = "0.10.3"
+serde_json = "1.0.104"
 
 [dev-dependencies]
 anyhow = { version = "1.0.64", features = ["backtrace"] }

mountpoint-s3-client/src/imds_crt_client.rs

@@ -1,16 +1,11 @@
-use std::pin::Pin;
-use std::task::Context;
-use std::task::Poll;
-
 use futures::channel::oneshot;
-use futures::Future;
 use mountpoint_s3_crt::auth::imds_client::{ImdsClient, ImdsClientConfig};
 use mountpoint_s3_crt::common::allocator::Allocator;
 use mountpoint_s3_crt::common::error;
 use mountpoint_s3_crt::io::channel_bootstrap::{ClientBootstrap, ClientBootstrapOptions};
 use mountpoint_s3_crt::io::event_loop::EventLoopGroup;
 use mountpoint_s3_crt::io::host_resolver::{HostResolver, HostResolverDefaultOptions};
-use pin_project::pin_project;
+use serde_json::Value;
 use thiserror::Error;
 
 #[derive(Debug)]
@@ -55,34 +50,47 @@ impl ImdsCrtClient {
         })
     }
 
-    /// Query the type of the EC2 instance this code is executed on.
-    pub fn make_instance_type_query(&self) -> Result<ImdsQueryRequest, ImdsQueryRequestError> {
+    /// Query the identity document of the EC2 instance this code is executed on.
+    pub async fn get_identity_document(&self) -> Result<IdentityDocument, ImdsQueryRequestError> {
+        const IDENTITY_DOCUMENT_PATH: &str = "/latest/dynamic/instance-identity/document";
+
         let (tx, rx) = oneshot::channel();
-        self.imds_client.get_instance_type(move |result| {
-            let _ = tx.send(result.map_err(ImdsQueryRequestError::CrtError));
+        self.imds_client.get_resource(IDENTITY_DOCUMENT_PATH, move |result| {
+            let _ = tx.send(result);
         })?;
 
-        Ok(ImdsQueryRequest { receiver: rx })
+        let json = match rx.await {
+            Ok(Ok(json)) => json,
+            Ok(Err(err)) => return Err(ImdsQueryRequestError::CrtError(err)),
+            Err(err) => return Err(ImdsQueryRequestError::InternalError(Box::new(err))),
+        };
+
+        let json = &json;
+        let parsed: Value =
+            serde_json::from_str(json).map_err(|_| ImdsQueryRequestError::InvalidResponse(json.to_owned()))?;
+        let instance_type = parsed
+            .get("instanceType")
+            .and_then(Value::as_str)
+            .ok_or_else(|| ImdsQueryRequestError::InvalidResponse(json.to_owned()))?
+            .to_owned();
+        let region = parsed
+            .get("region")
+            .and_then(Value::as_str)
+            .ok_or_else(|| ImdsQueryRequestError::InvalidResponse(json.to_owned()))?
+            .to_owned();
+
+        Ok(IdentityDocument { instance_type, region })
     }
 }
 
+/// Information about an EC2 instance.
 #[derive(Debug)]
-#[pin_project]
-/// ImdsQueryRequest is a data encapsulation wrapper for asnychronous instance metadata query.
-pub struct ImdsQueryRequest {
-    #[pin]
-    receiver: oneshot::Receiver<Result<String, ImdsQueryRequestError>>,
-}
-
-impl Future for ImdsQueryRequest {
-    type Output = Result<String, ImdsQueryRequestError>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = self.project();
-        this.receiver
-            .poll(cx)
-            .map(|result| result.unwrap_or_else(|err| Err(ImdsQueryRequestError::InternalError(Box::new(err)))))
-    }
+pub struct IdentityDocument {
+    // TODO: Add more fields as required.
+    /// The instance type of the instance.
+    pub instance_type: String,
+    /// The Region in which the instance is running.
+    pub region: String,
 }
 
 /// ImdsQueryRequestError is returned by an asynchronous query.
@@ -95,4 +103,8 @@ pub enum ImdsQueryRequestError {
     /// An internal error from within the AWS Common Runtime. The request may have been sent.
     #[error("Unknown CRT error")]
     CrtError(#[from] error::Error),
+
+    /// The response from Imds was not valid.
+    #[error("Invalid response from Imds: {0}")]
+    InvalidResponse(String),
 }

mountpoint-s3-client/src/lib.rs

@@ -9,7 +9,7 @@ mod s3_crt_client;
 mod util;
 
 pub use endpoint_config::{AddressingStyle, EndpointConfig};
-pub use imds_crt_client::ImdsCrtClient;
+pub use imds_crt_client::{IdentityDocument, ImdsCrtClient};
 pub use object_client::*;
 pub use s3_crt_client::head_bucket::HeadBucketError;
 pub use s3_crt_client::{S3ClientAuthConfig, S3ClientConfig, S3CrtClient, S3RequestError};

mountpoint-s3-client/tests/imds_crt_client.rs

@@ -0,0 +1,33 @@
+use std::assert_eq;
+
+use aws_config::imds::Client;
+use mountpoint_s3_client::ImdsCrtClient;
+use serde_json::Value;
+
+#[tokio::test]
+async fn test_get_identity_document() {
+    let sdk_client = Client::builder().build().await.unwrap();
+    match sdk_client.get("/latest/dynamic/instance-identity/document").await {
+        Ok(expected_json) => {
+            let expected_doc: Value = serde_json::from_str(&expected_json).unwrap();
+            let expected_region = expected_doc.get("region").unwrap().as_str().unwrap();
+            let expected_instance_type = expected_doc.get("instanceType").unwrap().as_str().unwrap();
+
+            let client = ImdsCrtClient::new().unwrap();
+            let identity_document = client
+                .get_identity_document()
+                .await
+                .expect("should return a valid document");
+
+            assert_eq!(identity_document.region, expected_region);
+            assert_eq!(identity_document.instance_type, expected_instance_type);
+        }
+        Err(_) => {
+            let client = ImdsCrtClient::new().unwrap();
+            let _ = client
+                .get_identity_document()
+                .await
+                .expect_err("should not return a valid document");
+        }
+    }
+}

mountpoint-s3-crt/src/auth/imds_client.rs

@@ -1,13 +1,13 @@
 //! A client for retrieving ec2 instance metadata.
 
-use crate::auth::auth_library_init;
 use crate::common::allocator::Allocator;
 use crate::common::error::Error;
 use crate::io::channel_bootstrap::ClientBootstrap;
 use crate::io::retry_strategy::RetryStrategy;
 use crate::CrtError;
+use crate::{auth::auth_library_init, ToAwsByteCursor};
 use mountpoint_s3_crt_sys::{
-    aws_byte_buf, aws_imds_client, aws_imds_client_get_instance_type, aws_imds_client_new, aws_imds_client_options,
+    aws_byte_buf, aws_imds_client, aws_imds_client_get_resource_async, aws_imds_client_new, aws_imds_client_options,
     aws_imds_client_release,
 };
 use std::ptr::NonNull;
@@ -66,8 +66,8 @@ impl ImdsClient {
         Ok(Self { inner })
     }
 
-    /// Get the EC2 instance type.
-    pub fn get_instance_type<F>(&self, callback: F) -> Result<(), Error>
+    /// Get the value for the given resource path.
+    pub fn get_resource<F>(&self, resource_path: &str, callback: F) -> Result<(), Error>
     where
         F: FnOnce(Result<String, Error>) + 'static,
     {
@@ -79,8 +79,13 @@ impl ImdsClient {
         // SAFETY: `self.inner` is a valid `aws_imds_client`. `get_resource_callback_fn_ptr` is leaked by [Box::into_raw]
         // and so will live until the `imds_client_get_resource_callback` function is invoked.
         unsafe {
-            aws_imds_client_get_instance_type(self.inner.as_ptr(), get_resource_callback_fn_ptr, callback_raw_ptr)
-                .ok_or_last_error()
+            aws_imds_client_get_resource_async(
+                self.inner.as_ptr(),
+                resource_path.as_aws_byte_cursor(),
+                get_resource_callback_fn_ptr,
+                callback_raw_ptr,
+            )
+            .ok_or_last_error()
         }
     }
 }