From f9e755ced0ff81945bbc05a1b9271b276f6b3d59 Mon Sep 17 00:00:00 2001
From: "Patrick J.P. Culp" <jpculp@amazon.com>
Date: Mon, 8 Mar 2021 22:12:11 +0000
Subject: [PATCH 1/4] Add imdsclient library for querying IMDS

This adds an imdsclient library for tools and setting generators to
leverage for fetching user-data, meta-data, and dynamic data from the
AWS Instance Metadata Service.
---
 packages/os/Cargo.toml        |   1 +
 sources/Cargo.lock            |  63 +++
 sources/Cargo.toml            |   2 +
 sources/imdsclient/Cargo.toml |  28 ++
 sources/imdsclient/README.md  |  13 +
 sources/imdsclient/README.tpl |   9 +
 sources/imdsclient/build.rs   |  32 ++
 sources/imdsclient/src/lib.rs | 736 ++++++++++++++++++++++++++++++++++
 8 files changed, 884 insertions(+)
 create mode 100644 sources/imdsclient/Cargo.toml
 create mode 100644 sources/imdsclient/README.md
 create mode 100644 sources/imdsclient/README.tpl
 create mode 100644 sources/imdsclient/build.rs
 create mode 100644 sources/imdsclient/src/lib.rs

diff --git a/packages/os/Cargo.toml b/packages/os/Cargo.toml
index 98babd27486..a6e032b601b 100644
--- a/packages/os/Cargo.toml
+++ b/packages/os/Cargo.toml
@@ -18,6 +18,7 @@ source-groups = [
     "webpki-roots-shim",
     "logdog",
     "models",
+    "imdsclient"
 ]
 
 [lib]
diff --git a/sources/Cargo.lock b/sources/Cargo.lock
index f35563444ed..d7ed655adee 100644
--- a/sources/Cargo.lock
+++ b/sources/Cargo.lock
@@ -328,6 +328,27 @@ version = "0.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "781f336cc9826dbaddb9754cb5db61e64cab4f69668bd19dcc4a0394a86f4cb1"
 
+[[package]]
+name = "async-stream"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a26cb53174ddd320edfff199a853f93d571f48eeb4dde75e67a9a3dbb7b7e5e"
+dependencies = [
+ "async-stream-impl",
+ "futures-core",
+]
+
+[[package]]
+name = "async-stream-impl"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "db134ba52475c060f3329a8ef0f8786d6b872ed01515d4b79c162e5798da1340"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "async-trait"
 version = "0.1.42"
@@ -1399,6 +1420,24 @@ dependencies = [
  "unicode-normalization",
 ]
 
+[[package]]
+name = "imdsclient"
+version = "0.1.0"
+dependencies = [
+ "cargo-readme",
+ "http",
+ "httptest",
+ "log",
+ "reqwest",
+ "serde",
+ "serde_json",
+ "simplelog",
+ "snafu",
+ "tokio",
+ "tokio-test",
+ "url",
+]
+
 [[package]]
 name = "indexmap"
 version = "1.6.2"
@@ -3251,6 +3290,30 @@ dependencies = [
  "webpki",
 ]
 
+[[package]]
+name = "tokio-stream"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e177a5d8c3bf36de9ebe6d58537d8879e964332f93fb3339e43f618c81361af0"
+dependencies = [
+ "futures-core",
+ "pin-project-lite",
+ "tokio",
+]
+
+[[package]]
+name = "tokio-test"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f58403903e94d4bc56805e46597fced893410b2e753e229d3f7f22423ea03f67"
+dependencies = [
+ "async-stream",
+ "bytes",
+ "futures-core",
+ "tokio",
+ "tokio-stream",
+]
+
 [[package]]
 name = "tokio-util"
 version = "0.6.5"
diff --git a/sources/Cargo.toml b/sources/Cargo.toml
index 21442c155b9..6d54b11fea0 100644
--- a/sources/Cargo.toml
+++ b/sources/Cargo.toml
@@ -34,6 +34,8 @@ members = [
 
     "bottlerocket-release",
 
+    "imdsclient",
+
     "ghostdog",
 
     "growpart",
diff --git a/sources/imdsclient/Cargo.toml b/sources/imdsclient/Cargo.toml
new file mode 100644
index 00000000000..c867935c4b3
--- /dev/null
+++ b/sources/imdsclient/Cargo.toml
@@ -0,0 +1,28 @@
+[package]
+name = "imdsclient"
+version = "0.1.0"
+authors = ["Patrick J.P. Culp <jpculp@amazon.com>"]
+license = "Apache-2.0 OR MIT"
+edition = "2018"
+publish = false
+build = "build.rs"
+# Don't rebuild crate just because of changes to README.
+exclude = ["README.md"]
+
+[dependencies]
+http = "0.2"
+log = "0.4"
+reqwest = { version = "0.11.1", default-features = false }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1"
+simplelog = "0.10"
+snafu = "0.6"
+tokio = { version = "1", default-features = false, features = ["macros", "rt-multi-thread", "time"] }
+url = "2.1.1"
+
+[build-dependencies]
+cargo-readme = "3.1"
+
+[dev-dependencies]
+httptest = "0.15"
+tokio-test = "0.4.1"
diff --git a/sources/imdsclient/README.md b/sources/imdsclient/README.md
new file mode 100644
index 00000000000..ee3889ca57a
--- /dev/null
+++ b/sources/imdsclient/README.md
@@ -0,0 +1,13 @@
+# imdsclient
+
+Current version: 0.1.0
+
+The imdsclient library provides high-level methods to interact with the AWS Instance Metadata Service.
+The high-level methods provided are [`fetch_dynamic`], [`fetch_metadata`], and [`fetch_userdata`].
+
+For more control, and to query IMDS without high-level wrappers, there is also a [`fetch_imds`] method.
+This method is useful for specifying things like a pinned date for the IMDS schema version.
+
+## Colophon
+
+This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`.
\ No newline at end of file
diff --git a/sources/imdsclient/README.tpl b/sources/imdsclient/README.tpl
new file mode 100644
index 00000000000..91fb62910c8
--- /dev/null
+++ b/sources/imdsclient/README.tpl
@@ -0,0 +1,9 @@
+# {{crate}}
+
+Current version: {{version}}
+
+{{readme}}
+
+## Colophon
+
+This text was generated from `README.tpl` using [cargo-readme](https://crates.io/crates/cargo-readme), and includes the rustdoc from `src/lib.rs`.
diff --git a/sources/imdsclient/build.rs b/sources/imdsclient/build.rs
new file mode 100644
index 00000000000..86c7bbc026e
--- /dev/null
+++ b/sources/imdsclient/build.rs
@@ -0,0 +1,32 @@
+// Automatically generate README.md from rustdoc.
+
+use std::env;
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+
+fn main() {
+    // Check for environment variable "SKIP_README". If it is set,
+    // skip README generation
+    if env::var_os("SKIP_README").is_some() {
+        return;
+    }
+
+    let mut source = File::open("src/lib.rs").unwrap();
+    let mut template = File::open("README.tpl").unwrap();
+
+    let content = cargo_readme::generate_readme(
+        &PathBuf::from("."), // root
+        &mut source,         // source
+        Some(&mut template), // template
+        // The "add x" arguments don't apply when using a template.
+        true,  // add title
+        false, // add badges
+        false, // add license
+        true,  // indent headings
+    )
+    .unwrap();
+
+    let mut readme = File::create("README.md").unwrap();
+    readme.write_all(content.as_bytes()).unwrap();
+}
diff --git a/sources/imdsclient/src/lib.rs b/sources/imdsclient/src/lib.rs
new file mode 100644
index 00000000000..63824aa8a5c
--- /dev/null
+++ b/sources/imdsclient/src/lib.rs
@@ -0,0 +1,736 @@
+/*!
+The imdsclient library provides high-level methods to interact with the AWS Instance Metadata Service.
+The high-level methods provided are [`fetch_dynamic`], [`fetch_metadata`], and [`fetch_userdata`].
+
+For more control, and to query IMDS without high-level wrappers, there is also a [`fetch_imds`] method.
+This method is useful for specifying things like a pinned date for the IMDS schema version.
+*/
+
+#![deny(rust_2018_idioms)]
+
+use http::StatusCode;
+use log::{debug, info, trace, warn};
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use snafu::{ensure, ResultExt};
+use std::time::Duration;
+use tokio::time;
+
+const BASE_URI: &str = "http://169.254.169.254";
+const PINNED_SCHEMA: &str = "2021-01-03";
+
+// Currently only able to get fetch session tokens from `latest`
+const SESSION_TARGET: &str = "latest/api/token";
+
+/// A client for making IMDSv2 queries.
+/// It obtains a session token when it is first instantiated and is reused between helper functions.
+pub struct ImdsClient {
+    client: Client,
+    imds_base_uri: String,
+    session_token: String,
+}
+
+/// This is the return type when querying for the IMDS identity document, which contains information
+/// such as region and instance_type. We only include the fields that we are using in Bottlerocket.
+#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct IdentityDocument {
+    region: String,
+    instance_type: String,
+}
+
+impl IdentityDocument {
+    pub fn region(&self) -> &str {
+        self.region.as_str()
+    }
+
+    pub fn instance_type(&self) -> &str {
+        self.instance_type.as_str()
+    }
+}
+
+impl ImdsClient {
+    pub async fn new() -> Result<Self> {
+        Self::new_impl(BASE_URI.to_string()).await
+    }
+
+    async fn new_impl(imds_base_uri: String) -> Result<Self> {
+        let client = Client::new();
+        let session_token = fetch_token(&client, &imds_base_uri).await?;
+        Ok(Self {
+            client,
+            imds_base_uri,
+            session_token,
+        })
+    }
+
+    /// Gets `user-data` from IMDS. The user-data may be either a UTF-8 string or compressed bytes.
+    pub async fn fetch_userdata(&mut self) -> Result<Vec<u8>> {
+        self.fetch_imds(PINNED_SCHEMA, "user-data").await
+    }
+
+    /// Returns the 'identity document' with fields like region and instance_type.
+    pub async fn fetch_identity_document(&mut self) -> Result<IdentityDocument> {
+        let target = "dynamic/instance-identity/document";
+        let response = self.fetch_bytes(target).await?;
+        let identity_document: IdentityDocument =
+            serde_json::from_slice(&response).context(error::Serde)?;
+        Ok(identity_document)
+    }
+
+    /// Returns the list of network interface mac addresses.
+    pub async fn fetch_mac_addresses(&mut self) -> Result<Vec<String>> {
+        let macs_target = "meta-data/network/interfaces/macs";
+        let macs = self.fetch_string(&macs_target).await?;
+        Ok(macs.split('\n').map(|s| s.to_string()).collect())
+    }
+
+    /// Gets the list of CIDR blocks for a given network interface `mac` address.
+    pub async fn fetch_cidr_blocks_for_mac(&mut self, mac: &str) -> Result<Vec<String>> {
+        // Infer the cluster DNS based on our CIDR blocks.
+        let mac_cidr_blocks_target = format!(
+            "meta-data/network/interfaces/macs/{}/vpc-ipv4-cidr-blocks",
+            mac
+        );
+        let cidr_blocks = self.fetch_string(&mac_cidr_blocks_target).await?;
+        Ok(cidr_blocks.split('\n').map(|s| s.to_string()).collect())
+    }
+
+    /// Gets the local IPV4 address from instance metadata.
+    pub async fn fetch_local_ipv4_address(&mut self) -> Result<String> {
+        let node_ip_target = "meta-data/local-ipv4";
+        self.fetch_string(&node_ip_target).await
+    }
+
+    /// Returns a list of public ssh keys skipping any keys that do not start with 'ssh'.
+    pub async fn fetch_public_ssh_keys(&mut self) -> Result<Vec<String>> {
+        info!("Fetching list of available public keys from IMDS");
+        // Returns a list of available public keys as '0=my-public-key'
+        let public_key_list = match self.fetch_string("meta-data/public-keys").await {
+            Err(error::Error::NotFound { uri: _ }) => {
+                // this is OK, it just means there are no keys
+                debug!("no available public keys");
+                return Ok(Vec::new());
+            }
+            Err(e) => {
+                return Err(e);
+            }
+            Ok(value) => value,
+        };
+
+        debug!("available public keys '{}'", &public_key_list);
+        info!("Generating targets to fetch text of available public keys");
+        let public_key_targets = build_public_key_targets(&public_key_list);
+
+        let mut public_keys = Vec::new();
+        let target_count: u32 = 0;
+        for target in &public_key_targets {
+            let target_count = target_count + 1;
+            info!(
+                "Fetching public key ({}/{})",
+                target_count,
+                &public_key_targets.len()
+            );
+
+            let public_key_text = self.fetch_string(&target).await?;
+            let public_key = public_key_text.trim_end();
+            // Simple check to see if the text is probably an ssh key.
+            if public_key.starts_with("ssh") {
+                debug!("{}", &public_key);
+                public_keys.push(public_key.to_string())
+            } else {
+                warn!(
+                    "'{}' does not appear to be a valid key. Skipping...",
+                    &public_key
+                );
+                continue;
+            }
+        }
+        if public_keys.is_empty() {
+            warn!("No valid keys found");
+        }
+        Ok(public_keys)
+    }
+
+    /// Helper to fetch bytes from IMDS using the pinned schema version.
+    async fn fetch_bytes<S>(&mut self, end_target: S) -> Result<Vec<u8>>
+    where
+        S: AsRef<str>,
+    {
+        self.fetch_imds(PINNED_SCHEMA, end_target.as_ref()).await
+    }
+
+    /// Helper to fetch a string from IMDS using the pinned schema version.
+    async fn fetch_string<S>(&mut self, end_target: S) -> Result<String>
+    where
+        S: AsRef<str>,
+    {
+        let response_body = self.fetch_imds(PINNED_SCHEMA, end_target).await?;
+        Ok(String::from_utf8(response_body).context(error::NonUtf8Response)?)
+    }
+
+    /// Fetch data from IMDS.
+    async fn fetch_imds<S1, S2>(&mut self, schema_version: S1, target: S2) -> Result<Vec<u8>>
+    where
+        S1: AsRef<str>,
+        S2: AsRef<str>,
+    {
+        let uri = format!(
+            "{}/{}/{}",
+            self.imds_base_uri,
+            schema_version.as_ref(),
+            target.as_ref()
+        );
+        debug!("Requesting {}", &uri);
+        let mut attempt: u8 = 0;
+        let max_attempts: u8 = 3;
+        loop {
+            attempt += 1;
+            if attempt > 1 {
+                time::sleep(Duration::from_millis(100)).await;
+            }
+            ensure!(attempt <= max_attempts, error::FailedFetch { attempt });
+            let response = self
+                .client
+                .get(&uri)
+                .header("X-aws-ec2-metadata-token", &self.session_token)
+                .send()
+                .await
+                .context(error::Request {
+                    method: "GET",
+                    uri: &uri,
+                })?;
+            trace!("IMDS response: {:?}", &response);
+
+            match response.status() {
+                code @ StatusCode::OK => {
+                    info!("Received {}", target.as_ref());
+                    let response_body = response
+                        .bytes()
+                        .await
+                        .context(error::ResponseBody {
+                            method: "GET",
+                            uri: &uri,
+                            code,
+                        })?
+                        .to_vec();
+
+                    let response_str = printable_string(&response_body);
+                    trace!("Response: {:?}", response_str);
+
+                    return Ok(response_body);
+                }
+
+                // IMDS returns 404 if no user data is given, or if IMDS is disabled
+                StatusCode::NOT_FOUND => return Err(error::Error::NotFound { uri }),
+
+                // IMDS returns 401 if the session token is expired or invalid
+                StatusCode::UNAUTHORIZED => {
+                    info!("Session token is invalid or expired");
+                    self.refresh_token().await?;
+                    info!("Refreshed session token");
+                    continue;
+                }
+
+                StatusCode::REQUEST_TIMEOUT => {
+                    info!("Retrying request");
+                    continue;
+                }
+
+                code => {
+                    let response_body = response
+                        .bytes()
+                        .await
+                        .context(error::ResponseBody {
+                            method: "GET",
+                            uri: &uri,
+                            code,
+                        })?
+                        .to_vec();
+
+                    let response_str = printable_string(&response_body);
+
+                    trace!("Response: {:?}", response_str);
+
+                    return error::Response {
+                        method: "GET",
+                        uri: &uri,
+                        code,
+                        response_body: response_str,
+                    }
+                    .fail();
+                }
+            }
+        }
+    }
+
+    /// Fetches a new session token and adds it to the current ImdsClient.
+    async fn refresh_token(&mut self) -> Result<()> {
+        self.session_token = fetch_token(&self.client, &self.imds_base_uri).await?;
+        Ok(())
+    }
+}
+
+/// Converts `bytes` to a `String` if it is a UTF-8 encoded string.
+/// Truncates the string if it is too long for printing.
+fn printable_string(bytes: &[u8]) -> String {
+    if let Ok(s) = String::from_utf8(bytes.into()) {
+        if s.len() < 2048 {
+            s
+        } else {
+            format!("{}<truncated...>", &s[0..2034])
+        }
+    } else {
+        "<binary>".to_string()
+    }
+}
+
+/// Returns a list of public keys available in IMDS. Since IMDS returns the list of keys as
+/// '0=my-public-key', we need to strip the index and insert it into the public key target.
+fn build_public_key_targets(public_key_list: &str) -> Vec<String> {
+    let mut public_key_targets = Vec::new();
+    for available_key in public_key_list.lines() {
+        let f: Vec<&str> = available_key.split('=').collect();
+        // If f[0] isn't a number, then it isn't a valid index.
+        if f[0].parse::<u32>().is_ok() {
+            let public_key_target = format!("meta-data/public-keys/{}/openssh-key", f[0]);
+            public_key_targets.push(public_key_target);
+        } else {
+            warn!(
+                "'{}' does not appear to be a valid index. Skipping...",
+                &f[0]
+            );
+            continue;
+        }
+    }
+    if public_key_targets.is_empty() {
+        warn!("No valid key targets found");
+    }
+    public_key_targets
+}
+
+/// Helper to fetch an IMDSv2 session token that is valid for 60 seconds.
+async fn fetch_token(client: &Client, imds_base_uri: &str) -> Result<String> {
+    let uri = format!("{}/{}", imds_base_uri, SESSION_TARGET);
+    let response = client
+        .put(&uri)
+        .header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+        .send()
+        .await
+        .context(error::Request {
+            method: "PUT",
+            uri: &uri,
+        })?
+        .error_for_status()
+        .context(error::BadResponse { uri: &uri })?;
+    let code = response.status();
+    response.text().await.context(error::ResponseBody {
+        method: "PUT",
+        uri,
+        code,
+    })
+}
+
+mod error {
+    use http::StatusCode;
+    use snafu::Snafu;
+
+    // Extracts the status code from a reqwest::Error and converts it to a string to be displayed
+    fn get_status_code(source: &reqwest::Error) -> String {
+        source
+            .status()
+            .as_ref()
+            .map(|i| i.as_str())
+            .unwrap_or("Unknown")
+            .to_string()
+    }
+
+    #[derive(Debug, Snafu)]
+    #[snafu(visibility = "pub(super)")]
+
+    pub enum Error {
+        #[snafu(display("Response '{}' from '{}': {}", get_status_code(&source), uri, source))]
+        BadResponse { uri: String, source: reqwest::Error },
+
+        #[snafu(display("IMDS fetch failed after {} attempts", attempt))]
+        FailedFetch { attempt: u8 },
+
+        #[snafu(display("IMDS session failed: {}", source))]
+        FailedSession { source: reqwest::Error },
+
+        #[snafu(display("Response was not UTF-8: {}", source))]
+        NonUtf8Response { source: std::string::FromUtf8Error },
+
+        #[snafu(display("404 file not found fetching '{}'", uri))]
+        NotFound { uri: String },
+
+        #[snafu(display("Error {}ing '{}': {}", method, uri, source))]
+        Request {
+            method: String,
+            uri: String,
+            source: reqwest::Error,
+        },
+
+        #[snafu(display("Error {} when {}ing '{}': {}", code, method, uri, response_body))]
+        Response {
+            method: String,
+            uri: String,
+            code: StatusCode,
+            response_body: String,
+        },
+
+        #[snafu(display(
+            "Unable to read response body when {}ing '{}' (code {}) - {}",
+            method,
+            uri,
+            code,
+            source
+        ))]
+        ResponseBody {
+            method: String,
+            uri: String,
+            code: StatusCode,
+            source: reqwest::Error,
+        },
+
+        #[snafu(display("Deserialization error: {}", source))]
+        Serde { source: serde_json::Error },
+    }
+}
+
+pub use error::Error;
+pub type Result<T> = std::result::Result<T, Error>;
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use httptest::{matchers::*, responders::*, Expectation, Server};
+
+    #[tokio::test]
+    async fn new_imds_client() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        let imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        assert_eq!(imds_client.session_token, token);
+    }
+
+    #[tokio::test]
+    async fn fetch_imds() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let schema_version = "latest";
+        let target = "meta-data/instance-type";
+        let response_code = 200;
+        let response_body = "m5.large";
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", schema_version, target),
+            ))
+            .times(1)
+            .respond_with(
+                status_code(response_code)
+                    .append_header("X-aws-ec2-metadata-token", token)
+                    .body(response_body),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        let imds_data = imds_client
+            .fetch_imds(schema_version, target)
+            .await
+            .unwrap();
+        assert_eq!(imds_data, response_body.as_bytes().to_vec());
+    }
+
+    #[tokio::test]
+    async fn fetch_imds_notfound() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let schema_version = "latest";
+        let target = "meta-data/instance-type";
+        let response_code = 404;
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", schema_version, target),
+            ))
+            .times(1)
+            .respond_with(
+                status_code(response_code).append_header("X-aws-ec2-metadata-token", token),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        let result = imds_client.fetch_imds(schema_version, target).await;
+        assert!(matches!(result, Err(error::Error::NotFound { .. })));
+    }
+
+    #[tokio::test]
+    async fn fetch_imds_unauthorized() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let schema_version = "latest";
+        let target = "meta-data/instance-type";
+        let response_code = 401;
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(4)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", schema_version, target),
+            ))
+            .times(3)
+            .respond_with(
+                status_code(response_code).append_header("X-aws-ec2-metadata-token", token),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        assert!(imds_client
+            .fetch_imds(schema_version, target)
+            .await
+            .is_err());
+    }
+
+    #[tokio::test]
+    async fn fetch_imds_timeout() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let schema_version = "latest";
+        let target = "meta-data/instance-type";
+        let response_code = 408;
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", schema_version, target),
+            ))
+            .times(3)
+            .respond_with(
+                status_code(response_code).append_header("X-aws-ec2-metadata-token", token),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        assert!(imds_client
+            .fetch_imds(schema_version, target)
+            .await
+            .is_err());
+    }
+
+    #[tokio::test]
+    async fn fetch_string() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let end_target = "meta-data/instance-type";
+        let response_code = 200;
+        let response_body = "m5.large";
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", PINNED_SCHEMA, end_target),
+            ))
+            .times(1)
+            .respond_with(
+                status_code(response_code)
+                    .append_header("X-aws-ec2-metadata-token", token)
+                    .body(response_body),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        let imds_data = imds_client.fetch_string(end_target).await.unwrap();
+        assert_eq!(imds_data, response_body.to_string());
+    }
+
+    #[tokio::test]
+    async fn fetch_bytes() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let end_target = "dynamic/instance-identity/document";
+        let response_code = 200;
+        let response_body = r#"{"region" : "us-west-2"}"#;
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/{}", PINNED_SCHEMA, end_target),
+            ))
+            .times(1)
+            .respond_with(
+                status_code(response_code)
+                    .append_header("X-aws-ec2-metadata-token", token)
+                    .body(response_body),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        let imds_data = imds_client.fetch_bytes(end_target).await.unwrap();
+        assert_eq!(imds_data, response_body.as_bytes().to_vec());
+    }
+
+    #[tokio::test]
+    async fn fetch_userdata() {
+        let server = Server::run();
+        let port = server.addr().port();
+        let base_uri = format!("http://localhost:{}", port);
+        let token = "some+token";
+        let response_code = 200;
+        let response_body = r#"settings.motd = "Welcome to Bottlerocket!""#;
+        server.expect(
+            Expectation::matching(request::method_path("PUT", "/latest/api/token"))
+                .times(1)
+                .respond_with(
+                    status_code(200)
+                        .append_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
+                        .body(token),
+                ),
+        );
+        server.expect(
+            Expectation::matching(request::method_path(
+                "GET",
+                format!("/{}/user-data", PINNED_SCHEMA),
+            ))
+            .times(1)
+            .respond_with(
+                status_code(response_code)
+                    .append_header("X-aws-ec2-metadata-token", token)
+                    .body(response_body),
+            ),
+        );
+        let mut imds_client = ImdsClient::new_impl(base_uri).await.unwrap();
+        let imds_data = imds_client.fetch_userdata().await.unwrap();
+        assert_eq!(imds_data, response_body.as_bytes().to_vec());
+    }
+
+    #[test]
+    fn printable_string_short() {
+        let input = "Hello".as_bytes();
+        let expected = "Hello".to_string();
+        let actual = printable_string(input);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn printable_string_binary() {
+        let input: [u8; 5] = [0, 254, 1, 0, 4];
+        let expected = "<binary>".to_string();
+        let actual = printable_string(&input);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn printable_string_untruncated() {
+        let mut input = String::new();
+        for _ in 0..2047 {
+            input.push('.');
+        }
+        let expected = input.clone();
+        let actual = printable_string(input.as_bytes());
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn printable_string_truncated() {
+        let mut input = String::new();
+        for _ in 0..2048 {
+            input.push('.');
+        }
+        let mut expected = String::new();
+        for _ in 0..2034 {
+            expected.push('.');
+        }
+        expected.push_str("<truncated...>");
+        let actual = printable_string(input.as_bytes());
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn parse_public_key_list() {
+        let list = r#"0=zero
+1=one
+2=two"#;
+        let parsed_list = build_public_key_targets(list);
+        assert_eq!(3, parsed_list.len());
+        assert_eq!(
+            "meta-data/public-keys/0/openssh-key",
+            parsed_list.get(0).unwrap()
+        );
+        assert_eq!(
+            "meta-data/public-keys/1/openssh-key",
+            parsed_list.get(1).unwrap()
+        );
+        assert_eq!(
+            "meta-data/public-keys/2/openssh-key",
+            parsed_list.get(2).unwrap()
+        );
+    }
+}

From dfb7024e07d88ae1808b95e655eb49d7b4b2e338 Mon Sep 17 00:00:00 2001
From: "Patrick J.P. Culp" <jpculp@amazon.com>
Date: Fri, 21 May 2021 18:46:18 +0000
Subject: [PATCH 2/4] shibaken: update to use imdsclient library

This removes the stand-alone IMDS logic in favor of the shared
imdsclient library.

Log messages were also tweaked for more clarity.
---
 sources/Cargo.lock               |   3 +-
 sources/api/shibaken/Cargo.toml  |   3 +-
 sources/api/shibaken/src/main.rs | 167 ++++++-------------------------
 3 files changed, 32 insertions(+), 141 deletions(-)

diff --git a/sources/Cargo.lock b/sources/Cargo.lock
index d7ed655adee..b82062aebfe 100644
--- a/sources/Cargo.lock
+++ b/sources/Cargo.lock
@@ -2781,12 +2781,13 @@ version = "0.1.0"
 dependencies = [
  "base64",
  "cargo-readme",
+ "imdsclient",
  "log",
- "reqwest",
  "serde",
  "serde_json",
  "simplelog",
  "snafu",
+ "tokio",
 ]
 
 [[package]]
diff --git a/sources/api/shibaken/Cargo.toml b/sources/api/shibaken/Cargo.toml
index 36f817bc381..111895b6d00 100644
--- a/sources/api/shibaken/Cargo.toml
+++ b/sources/api/shibaken/Cargo.toml
@@ -11,12 +11,13 @@ exclude = ["README.md"]
 
 [dependencies]
 base64 = "0.13"
+imdsclient = { path = "../../imdsclient" }
 log = "0.4"
-reqwest = { version = "0.11.1", default-features = false, features = ["blocking"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 simplelog = "0.10"
 snafu = "0.6"
+tokio = { version = "1", default-features = false, features = ["macros", "rt-multi-thread"] }
 
 [build-dependencies]
 cargo-readme = "3.1"
diff --git a/sources/api/shibaken/src/main.rs b/sources/api/shibaken/src/main.rs
index dacdf02071f..90b9d3cd571 100644
--- a/sources/api/shibaken/src/main.rs
+++ b/sources/api/shibaken/src/main.rs
@@ -11,21 +11,14 @@ AWS instance metadata service (IMDS).
 
 #![deny(rust_2018_idioms)]
 
-use log::{debug, info, warn};
-use reqwest::blocking::Client;
+use imdsclient::ImdsClient;
+use log::{debug, info};
 use serde::Serialize;
 use simplelog::{ColorChoice, Config as LogConfig, LevelFilter, TermLogger, TerminalMode};
 use snafu::{OptionExt, ResultExt};
 use std::str::FromStr;
 use std::{env, process};
 
-// Instance Meta Data Service.
-//
-// Currently only able to get fetch session tokens from `latest`
-// FIXME Pin to a date version that supports IMDSv2 once such a date version is available.
-const IMDS_PUBLIC_KEY_BASE_URI: &str = "http://169.254.169.254/latest/meta-data/public-keys";
-const IMDS_SESSION_TOKEN_URI: &str = "http://169.254.169.254/latest/api/token";
-
 #[derive(Serialize)]
 struct UserData {
     ssh: Ssh,
@@ -46,108 +39,14 @@ impl UserData {
     }
 }
 
-/// Helper to fetch an IMDSv2 session token that is valid for 60 seconds.
-fn fetch_imds_session_token(client: &Client) -> Result<String> {
-    let uri = IMDS_SESSION_TOKEN_URI;
-    let imds_session_token = client
-        .put(uri)
-        .header("X-aws-ec2-metadata-token-ttl-seconds", "60")
-        .send()
-        .context(error::ImdsRequest { method: "PUT", uri })?
-        .error_for_status()
-        .context(error::ImdsResponse { uri })?
-        .text()
-        .context(error::ImdsText { uri })?;
-    Ok(imds_session_token)
-}
-
-/// Helper to fetch data from IMDS. Inspired by pluto.
-fn fetch_from_imds(client: &Client, uri: &str, session_token: &str) -> Result<Option<String>> {
-    let response = client
-        .get(uri)
-        .header("X-aws-ec2-metadata-token", session_token)
-        .send()
-        .context(error::ImdsRequest { method: "GET", uri })?;
-    if response.status().as_u16() == 404 {
-        return Ok(None);
-    }
-    Ok(Some(
-        response
-            .error_for_status()
-            .context(error::ImdsResponse { uri })?
-            .text()
-            .context(error::ImdsText { uri })?,
-    ))
-}
-
 /// Returns a list of public keys.
-fn fetch_public_keys_from_imds() -> Result<Vec<String>> {
-    info!("Fetching IMDS session token");
-    let client = Client::new();
-    let imds_session_token = fetch_imds_session_token(&client)?;
-
-    info!("Fetching list of available public keys from IMDS");
-    // Returns a list of available public keys as '0=my-public-key'
-    let public_key_list = if let Some(public_key_list) =
-        fetch_from_imds(&client, IMDS_PUBLIC_KEY_BASE_URI, &imds_session_token)?
-    {
-        debug!("available public keys '{}'", &public_key_list);
-        public_key_list
-    } else {
-        debug!("no available public keys");
-        return Ok(Vec::new());
-    };
-
-    info!("Generating uris to fetch text of available public keys");
-    let public_key_uris = build_public_key_uris(&public_key_list);
-
-    info!("Fetching public keys from IMDS");
-    let mut public_keys = Vec::new();
-    for uri in public_key_uris {
-        let public_key_text = fetch_from_imds(&client, &uri, &imds_session_token)?
-            .context(error::KeyNotFound { uri })?;
-        let public_key = public_key_text.trim_end();
-        // Simple check to see if the text is probably an ssh key.
-        if public_key.starts_with("ssh") {
-            debug!("{}", &public_key);
-            public_keys.push(public_key.to_string())
-        } else {
-            warn!(
-                "'{}' does not appear to be a valid key. Skipping...",
-                &public_key_text
-            );
-            continue;
-        }
-    }
-    if public_keys.is_empty() {
-        warn!("No valid keys found");
-    }
-    Ok(public_keys)
-}
-
-/// Returns a list of public key uris strings for the public keys in IMDS. Since IMDS returns the
-/// list of available public keys as '0=my-public-key', we need to strip the index from the list and
-/// insert it into the key uri.
-fn build_public_key_uris(public_key_list: &str) -> Vec<String> {
-    let mut public_key_uris = Vec::new();
-    for available_key in public_key_list.lines() {
-        let f: Vec<&str> = available_key.split('=').collect();
-        // If f[0] isn't a number, then it isn't a valid index.
-        if f[0].parse::<u32>().is_ok() {
-            let public_key_uri = format!("{}/{}/openssh-key", IMDS_PUBLIC_KEY_BASE_URI, f[0]);
-            public_key_uris.push(public_key_uri);
-        } else {
-            warn!(
-                "'{}' does not appear to be a valid index. Skipping...",
-                &f[0]
-            );
-            continue;
-        }
-    }
-    if public_key_uris.is_empty() {
-        warn!("No valid key uris found");
-    }
-    public_key_uris
+async fn fetch_public_keys_from_imds() -> Result<Vec<String>> {
+    info!("Connecting to IMDS");
+    let mut client = ImdsClient::new().await.context(error::ImdsClient)?;
+    client
+        .fetch_public_ssh_keys()
+        .await
+        .context(error::ImdsClient)
 }
 
 /// Store the args we receive on the command line.
@@ -192,11 +91,11 @@ fn parse_args(args: env::Args) -> Result<Args> {
     }
 
     Ok(Args {
-        log_level: log_level.unwrap_or_else(|| LevelFilter::Info),
+        log_level: log_level.unwrap_or(LevelFilter::Info),
     })
 }
 
-fn run() -> Result<()> {
+async fn run() -> Result<()> {
     let args = parse_args(env::args())?;
 
     // TerminalMode::Stderr will send all logs to stderr, as sundog only expects the json output of
@@ -211,7 +110,7 @@ fn run() -> Result<()> {
 
     info!("shibaken started");
 
-    let public_keys = fetch_public_keys_from_imds()?;
+    let public_keys = fetch_public_keys_from_imds().await?;
 
     let user_data = UserData::new(public_keys);
 
@@ -226,7 +125,7 @@ fn run() -> Result<()> {
     // user to bypass shibaken and use their own user-data if desired.
     let user_data_base64 = base64::encode(&user_data_json);
 
-    info!("Outputting user-data");
+    info!("Outputting base64-encoded user-data");
     // sundog expects JSON-serialized output so that many types can be represented, allowing the
     // API model to use more accurate types.
     let output = serde_json::to_string(&user_data_base64).context(error::SerializeJson)?;
@@ -239,8 +138,9 @@ fn run() -> Result<()> {
 // Returning a Result from main makes it print a Debug representation of the error, but with Snafu
 // we have nice Display representations of the error, so we wrap "main" (run) and print any error.
 // https://github.com/shepmaster/snafu/issues/110
-fn main() {
-    if let Err(e) = run() {
+#[tokio::main]
+async fn main() {
+    if let Err(e) = run().await {
         match e {
             error::Error::Usage { .. } => {
                 eprintln!("{}", e);
@@ -259,33 +159,22 @@ fn main() {
 
 mod error {
     use snafu::Snafu;
-    fn code(source: &reqwest::Error) -> String {
-        source
-            .status()
-            .as_ref()
-            .map(|i| i.as_str())
-            .unwrap_or("Unknown")
-            .to_string()
-    }
 
     #[derive(Debug, Snafu)]
     #[snafu(visibility = "pub(super)")]
     pub(super) enum Error {
-        #[snafu(display("Error {}ing '{}': {}", method, uri, source))]
-        ImdsRequest {
-            method: String,
-            uri: String,
-            source: reqwest::Error,
-        },
-
-        #[snafu(display("Error '{}' from '{}': {}", code(&source), uri, source))]
-        ImdsResponse { uri: String, source: reqwest::Error },
-
-        #[snafu(display("Error getting text response from {}: {}", uri, source))]
-        ImdsText { uri: String, source: reqwest::Error },
-
-        #[snafu(display("Error retrieving key from {}", uri))]
-        KeyNotFound { uri: String },
+        #[snafu(display("IMDS request failed: {}", source))]
+        ImdsRequest { source: imdsclient::Error },
+
+        #[snafu(display("IMDS client failed: {}", source))]
+        ImdsClient { source: imdsclient::Error },
+
+        #[snafu(display(
+            "IMDS client failed: Response '404' while fetching '{}' from '{}'",
+            target,
+            target_type,
+        ))]
+        ImdsData { target: String, target_type: String },
 
         #[snafu(display("Logger setup error: {}", source))]
         Logger { source: log::SetLoggerError },

From b0282c959313d4543746d72b491ec1a7d7201eb5 Mon Sep 17 00:00:00 2001
From: "Patrick J.P. Culp" <jpculp@amazon.com>
Date: Fri, 21 May 2021 18:47:29 +0000
Subject: [PATCH 3/4] pluto: update to use imdsclient library

This removes the stand-alone IMDS logic in favor of the shared
imdsclient library.
---
 sources/Cargo.lock            |   2 +-
 sources/api/pluto/Cargo.toml  |   2 +-
 sources/api/pluto/src/main.rs | 138 ++++++++++++----------------------
 3 files changed, 51 insertions(+), 91 deletions(-)

diff --git a/sources/Cargo.lock b/sources/Cargo.lock
index b82062aebfe..228e1e580e2 100644
--- a/sources/Cargo.lock
+++ b/sources/Cargo.lock
@@ -2144,8 +2144,8 @@ version = "0.1.0"
 dependencies = [
  "apiclient",
  "cargo-readme",
+ "imdsclient",
  "models",
- "reqwest",
  "rusoto_core",
  "rusoto_eks",
  "serde_json",
diff --git a/sources/api/pluto/Cargo.toml b/sources/api/pluto/Cargo.toml
index 5021dfec3fb..dfc96167f0b 100644
--- a/sources/api/pluto/Cargo.toml
+++ b/sources/api/pluto/Cargo.toml
@@ -11,8 +11,8 @@ exclude = ["README.md"]
 
 [dependencies]
 apiclient = { path = "../apiclient" }
+imdsclient = { path = "../../imdsclient" }
 models = { path = "../../models" }
-reqwest = { version = "0.11.1", default-features = false }
 rusoto_core = { version = "0.46", default-features = false, features = ["rustls"] }
 rusoto_eks = { version = "0.46", default-features = false, features = ["rustls"] }
 serde_json = "1"
diff --git a/sources/api/pluto/src/main.rs b/sources/api/pluto/src/main.rs
index 548cd354901..816f5f64be4 100644
--- a/sources/api/pluto/src/main.rs
+++ b/sources/api/pluto/src/main.rs
@@ -36,7 +36,7 @@ reasonable default is available.
 mod api;
 mod eks;
 
-use reqwest::Client;
+use imdsclient::ImdsClient;
 use snafu::{ensure, OptionExt, ResultExt};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
@@ -48,51 +48,26 @@ const DEFAULT_DNS_CLUSTER_IP: &str = "10.100.0.10";
 // If our CIDR block begins with "10." this is our DNS.
 const DEFAULT_10_RANGE_DNS_CLUSTER_IP: &str = "172.20.0.10";
 
-// Instance Meta Data Service
-const IMDS_BASE_URL: &str = "http://169.254.169.254/2018-09-24";
-// Currently only able to get fetch session tokens from `latest`
-// FIXME Pin to a date version that supports IMDSv2 once such a date version is available.
-const IMDS_SESSION_TOKEN_ENDPOINT: &str = "http://169.254.169.254/latest/api/token";
-const IMDS_NODE_IPV4_ENDPOINT: &str = "http://169.254.169.254/2018-09-24/meta-data/local-ipv4";
-const IMDS_MAC_ENDPOINT: &str =
-    "http://169.254.169.254/2018-09-24/meta-data/network/interfaces/macs";
-const IMDS_INSTANCE_TYPE_ENDPOINT: &str =
-    "http://169.254.169.254/2018-09-24/meta-data/instance-type";
-
 const ENI_MAX_PODS_PATH: &str = "/usr/share/eks/eni-max-pods";
 
 mod error {
     use crate::eks;
     use snafu::Snafu;
 
-    // Taken from sundog.
-    fn code(source: &reqwest::Error) -> String {
-        source
-            .status()
-            .as_ref()
-            .map(|i| i.as_str())
-            .unwrap_or("Unknown")
-            .to_string()
-    }
-
     #[derive(Debug, Snafu)]
     #[snafu(visibility = "pub(super)")]
     pub(super) enum PlutoError {
         #[snafu(display("Unable to parse CIDR '{}': {}", cidr, reason))]
         CidrParse { cidr: String, reason: String },
 
-        #[snafu(display("Error {}ing '{}': {}", method, uri, source))]
-        ImdsRequest {
-            method: String,
-            uri: String,
-            source: reqwest::Error,
-        },
+        #[snafu(display("IMDS request failed: {}", source))]
+        ImdsRequest { source: imdsclient::Error },
 
-        #[snafu(display("Error '{}' from '{}': {}", code(&source), uri, source))]
-        ImdsResponse { uri: String, source: reqwest::Error },
+        #[snafu(display("IMDS client failed: {}", source))]
+        ImdsClient { source: imdsclient::Error },
 
-        #[snafu(display("Error getting text response from {}: {}", uri, source))]
-        ImdsText { uri: String, source: reqwest::Error },
+        #[snafu(display("IMDS request failed: No '{}' found", what))]
+        ImdsNone { what: String },
 
         #[snafu(display("Error deserializing response into JSON from {}: {}", uri, source))]
         ImdsJson {
@@ -110,12 +85,6 @@ mod error {
             source: serde_json::error::Error,
         },
 
-        #[snafu(display("Missing MAC address from IMDS: {}", uri))]
-        MissingMac { uri: String },
-
-        #[snafu(display("Invalid machine architecture, not one of 'x86_64' or 'aarch64'"))]
-        UnknownArchitecture,
-
         #[snafu(display("{}", source))]
         EksError { source: eks::Error },
 
@@ -146,23 +115,14 @@ use error::PlutoError;
 
 type Result<T> = std::result::Result<T, PlutoError>;
 
-async fn get_text_from_imds(client: &Client, uri: &str, session_token: &str) -> Result<String> {
-    client
-        .get(uri)
-        .header("X-aws-ec2-metadata-token", session_token)
-        .send()
-        .await
-        .context(error::ImdsRequest { method: "GET", uri })?
-        .error_for_status()
-        .context(error::ImdsResponse { uri })?
-        .text()
+async fn get_max_pods(client: &mut ImdsClient) -> Result<String> {
+    let instance_type = client
+        .fetch_identity_document()
         .await
-        .context(error::ImdsText { uri })
-}
+        .context(error::ImdsRequest)?
+        .instance_type()
+        .to_string();
 
-async fn get_max_pods(client: &Client, session_token: &str) -> Result<String> {
-    let instance_type =
-        get_text_from_imds(&client, IMDS_INSTANCE_TYPE_ENDPOINT, session_token).await?;
     // Find the corresponding maximum number of pods supported by this instance type
     let file = BufReader::new(
         File::open(ENI_MAX_PODS_PATH).context(error::EniMaxPodsFile {
@@ -187,7 +147,7 @@ async fn get_max_pods(client: &Client, session_token: &str) -> Result<String> {
 /// the `serviceIPv4CIDR`. If that works, it returns the expected cluster DNS IP address which is
 /// obtained by substituting `10` for the last octet. If the EKS call is not successful, it falls
 /// back to using IMDS MAC CIDR blocks to return one of two default addresses.
-async fn get_cluster_dns_ip(client: &Client, session_token: &str) -> Result<String> {
+async fn get_cluster_dns_ip(client: &mut ImdsClient) -> Result<String> {
     // try calling eks describe-cluster to figure out the dns cluster ip
     if let Some(dns_ip) = get_dns_from_eks().await {
         // we were able to calculate the dns ip from the cidr range we received from eks
@@ -196,7 +156,7 @@ async fn get_cluster_dns_ip(client: &Client, session_token: &str) -> Result<Stri
 
     // we were unable to obtain or parse the cidr range from eks, fallback to one of two default
     // values based on the cidr range of our primary network interface
-    get_cluster_dns_from_imds_mac(client, session_token).await
+    get_cluster_dns_from_imds_mac(client).await
 }
 
 /// Gets the Service IPV4 CIDR setting from EKS and parses it to calculate the cluster DNS IP.
@@ -242,20 +202,31 @@ fn get_dns_from_cidr(cidr: &str) -> Result<String> {
 
 /// Gets gets the the first VPC IPV4 CIDR block from IMDS. If it starts with `10`, returns
 /// `10.100.0.10`, otherwise returns `172.20.0.10`
-async fn get_cluster_dns_from_imds_mac(client: &Client, session_token: &str) -> Result<String> {
-    let uri = IMDS_MAC_ENDPOINT;
-    let macs = get_text_from_imds(&client, uri, session_token).await?;
-    // Take the first (primary) MAC address. Others will exist from attached ENIs.
-    let mac = macs.split('\n').next().context(error::MissingMac { uri })?;
-
-    // Infer the cluster DNS based on our CIDR blocks.
-    let mac_cidr_blocks_uri = format!(
-        "{}/meta-data/network/interfaces/macs/{}/vpc-ipv4-cidr-blocks",
-        IMDS_BASE_URL, mac
-    );
-    let mac_cidr_blocks = get_text_from_imds(&client, &mac_cidr_blocks_uri, session_token).await?;
-
-    let dns = if mac_cidr_blocks.starts_with("10.") {
+async fn get_cluster_dns_from_imds_mac(client: &mut ImdsClient) -> Result<String> {
+    // Take the first (primary) MAC address. Others may exist from attached ENIs.
+    let mac = client
+        .fetch_mac_addresses()
+        .await
+        .context(error::ImdsRequest)?
+        .first()
+        .context(error::ImdsNone {
+            what: "mac addresses",
+        })?
+        .clone();
+
+    // Take the first CIDR block for the primary MAC.
+    let cidr_block = client
+        .fetch_cidr_blocks_for_mac(&mac)
+        .await
+        .context(error::ImdsRequest)?
+        .first()
+        .context(error::ImdsNone {
+            what: "CIDR blocks",
+        })?
+        .clone();
+
+    // Infer the cluster DNS based on the CIDR block.
+    let dns = if cidr_block.starts_with("10.") {
         DEFAULT_10_RANGE_DNS_CLUSTER_IP
     } else {
         DEFAULT_DNS_CLUSTER_IP
@@ -264,8 +235,11 @@ async fn get_cluster_dns_from_imds_mac(client: &Client, session_token: &str) ->
     Ok(dns)
 }
 
-async fn get_node_ip(client: &Client, session_token: &str) -> Result<String> {
-    get_text_from_imds(&client, IMDS_NODE_IPV4_ENDPOINT, session_token).await
+async fn get_node_ip(client: &mut ImdsClient) -> Result<String> {
+    client
+        .fetch_local_ipv4_address()
+        .await
+        .context(error::ImdsRequest)
 }
 
 /// Print usage message.
@@ -285,28 +259,14 @@ fn parse_args(mut args: env::Args) -> String {
 
 async fn run() -> Result<()> {
     let setting_name = parse_args(env::args());
-    let client = Client::new();
-
-    // Use IMDSv2 for accessing instance metadata
-    let uri = IMDS_SESSION_TOKEN_ENDPOINT;
-    let imds_session_token = client
-        .put(uri)
-        .header("X-aws-ec2-metadata-token-ttl-seconds", "60")
-        .send()
-        .await
-        .context(error::ImdsRequest { method: "PUT", uri })?
-        .error_for_status()
-        .context(error::ImdsResponse { uri })?
-        .text()
-        .await
-        .context(error::ImdsText { uri })?;
+    let mut client = ImdsClient::new().await.context(error::ImdsClient)?;
 
     let setting = match setting_name.as_ref() {
-        "cluster-dns-ip" => get_cluster_dns_ip(&client, &imds_session_token).await,
-        "node-ip" => get_node_ip(&client, &imds_session_token).await,
+        "cluster-dns-ip" => get_cluster_dns_ip(&mut client).await,
+        "node-ip" => get_node_ip(&mut client).await,
         // If we want to specify a reasonable default in a template, we can exit 2 to tell
         // sundog to skip this setting.
-        "max-pods" => get_max_pods(&client, &imds_session_token)
+        "max-pods" => get_max_pods(&mut client)
             .await
             .map_err(|_| process::exit(2)),
 

From 27707734c822b56769f299dbe6df77986abc91de Mon Sep 17 00:00:00 2001
From: "Patrick J.P. Culp" <jpculp@amazon.com>
Date: Fri, 21 May 2021 18:48:07 +0000
Subject: [PATCH 4/4] early-boot-config: update to use imdsclient library

This removes the stand-alone IMDS logic in favor of the shared
imdsclient library.
---
 sources/Cargo.lock                            |   3 +-
 sources/api/early-boot-config/Cargo.toml      |   3 +-
 sources/api/early-boot-config/src/main.rs     |   1 +
 sources/api/early-boot-config/src/provider.rs |   6 +-
 .../api/early-boot-config/src/provider/aws.rs | 240 ++++--------------
 .../src/provider/local_file.rs                |   4 +-
 .../early-boot-config/src/provider/vmware.rs  |   4 +-
 7 files changed, 69 insertions(+), 192 deletions(-)

diff --git a/sources/Cargo.lock b/sources/Cargo.lock
index 228e1e580e2..dc2d0beac40 100644
--- a/sources/Cargo.lock
+++ b/sources/Cargo.lock
@@ -857,14 +857,15 @@ name = "early-boot-config"
 version = "0.1.0"
 dependencies = [
  "apiclient",
+ "async-trait",
  "base64",
  "cargo-readme",
  "flate2",
  "hex-literal",
  "http",
+ "imdsclient",
  "lazy_static",
  "log",
- "reqwest",
  "serde",
  "serde-xml-rs",
  "serde_json",
diff --git a/sources/api/early-boot-config/Cargo.toml b/sources/api/early-boot-config/Cargo.toml
index e2ea5aa2ffc..9b5f1e0a30d 100644
--- a/sources/api/early-boot-config/Cargo.toml
+++ b/sources/api/early-boot-config/Cargo.toml
@@ -11,11 +11,12 @@ exclude = ["README.md"]
 
 [dependencies]
 apiclient = { path = "../apiclient" }
+async-trait = "0.1.36"
 base64 = "0.13"
 flate2 = { version = "1.0", default-features = false, features = ["rust_backend"] }
 http = "0.2"
+imdsclient = { path = "../../imdsclient" }
 log = "0.4"
-reqwest = { version = "0.11", default-features = false, features = ["blocking"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 serde_plain = "0.3"
diff --git a/sources/api/early-boot-config/src/main.rs b/sources/api/early-boot-config/src/main.rs
index 3c4d670b200..3c3b039e2ea 100644
--- a/sources/api/early-boot-config/src/main.rs
+++ b/sources/api/early-boot-config/src/main.rs
@@ -142,6 +142,7 @@ async fn run() -> Result<()> {
     let method = "PATCH";
     for settings_json in data_provider
         .platform_data()
+        .await
         .context(error::ProviderError)?
     {
         // Don't send an empty request to the API
diff --git a/sources/api/early-boot-config/src/provider.rs b/sources/api/early-boot-config/src/provider.rs
index b223725d596..33529c61bba 100644
--- a/sources/api/early-boot-config/src/provider.rs
+++ b/sources/api/early-boot-config/src/provider.rs
@@ -1,6 +1,7 @@
 //! The provider module owns the `PlatformDataProvider` trait
 
 use crate::settings::SettingsJson;
+use async_trait::async_trait;
 
 #[cfg(any(bottlerocket_platform = "aws", bottlerocket_platform = "aws-dev"))]
 pub(crate) mod aws;
@@ -12,11 +13,14 @@ pub(crate) mod local_file;
 pub(crate) mod vmware;
 
 /// Support for new platforms can be added by implementing this trait.
+#[async_trait]
 pub(crate) trait PlatformDataProvider {
     /// You should return a list of SettingsJson, representing the settings changes you want to
     /// send to the API.
     ///
     /// This is a list so that handling multiple data sources within a platform can feel more
     /// natural; you can also send all changes in one entry if you like.
-    fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>>;
+    async fn platform_data(
+        &self,
+    ) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>>;
 }
diff --git a/sources/api/early-boot-config/src/provider/aws.rs b/sources/api/early-boot-config/src/provider/aws.rs
index ca65ee060df..a60e3a323f0 100644
--- a/sources/api/early-boot-config/src/provider/aws.rs
+++ b/sources/api/early-boot-config/src/provider/aws.rs
@@ -2,8 +2,8 @@
 
 use super::{PlatformDataProvider, SettingsJson};
 use crate::compression::expand_slice_maybe;
-use http::StatusCode;
-use reqwest::blocking::Client;
+use async_trait::async_trait;
+use imdsclient::ImdsClient;
 use serde_json::json;
 use snafu::{OptionExt, ResultExt};
 use std::fs;
@@ -13,123 +13,12 @@ use std::path::Path;
 pub(crate) struct AwsDataProvider;
 
 impl AwsDataProvider {
-    // Currently only able to get fetch session tokens from `latest`
-    // FIXME Pin to a date version that supports IMDSv2 once such a date version is available.
-    const IMDS_TOKEN_ENDPOINT: &'static str = "http://169.254.169.254/latest/api/token";
-
-    const USER_DATA_ENDPOINT: &'static str = "http://169.254.169.254/2018-09-24/user-data";
     const IDENTITY_DOCUMENT_FILE: &'static str = "/etc/early-boot-config/identity-document";
-    const IDENTITY_DOCUMENT_ENDPOINT: &'static str =
-        "http://169.254.169.254/2018-09-24/dynamic/instance-identity/document";
-
-    /// Helper to fetch an IMDSv2 session token that is valid for 60 seconds.
-    fn fetch_imds_session_token(client: &Client) -> Result<String> {
-        let uri = Self::IMDS_TOKEN_ENDPOINT;
-        let response = client
-            .put(uri)
-            .header("X-aws-ec2-metadata-token-ttl-seconds", "60")
-            .send()
-            .context(error::Request { method: "PUT", uri })?
-            .error_for_status()
-            .context(error::BadResponse { uri })?;
-        let code = response.status();
-        response.text().context(error::ResponseBody {
-            method: "PUT",
-            uri,
-            code,
-        })
-    }
-
-    /// Helper to fetch data from IMDS, preferring an override file if present.
-    ///
-    /// IMDS returns a 404 if no user data was given, for example; we return Ok(None) to represent
-    /// this, otherwise Ok(Some(body)) with the response body.
-    fn fetch_imds(
-        client: &Client,
-        session_token: &str,
-        uri: &str,
-        description: &str,
-    ) -> Result<Option<Vec<u8>>> {
-        debug!("Requesting {} from {}", description, uri);
-        let response = client
-            .get(uri)
-            .header("X-aws-ec2-metadata-token", session_token)
-            .send()
-            .context(error::Request { method: "GET", uri })?;
-        trace!("IMDS response: {:?}", &response);
-
-        // IMDS data can be larger than we'd want to log (50k+ compressed) so we don't necessarily
-        // want to show the whole thing, and don't want to show binary data.
-        fn response_string(response: &[u8]) -> String {
-            // arbitrary max len; would be nice to print the start of the data if it's
-            // uncompressed, but we'd need to break slice at a safe point for UTF-8, and without
-            // reading in the whole thing like String::from_utf8.
-            if response.len() > 2048 {
-                "<very long>".to_string()
-            } else if let Ok(s) = String::from_utf8(response.into()) {
-                s
-            } else {
-                "<binary>".to_string()
-            }
-        }
-
-        match response.status() {
-            code @ StatusCode::OK => {
-                info!("Received {}", description);
-                let response_body = response
-                    .bytes()
-                    .context(error::ResponseBody {
-                        method: "GET",
-                        uri,
-                        code,
-                    })?
-                    .to_vec();
-
-                let response_str = response_string(&response_body);
-                trace!("Response: {:?}", response_str);
-
-                Ok(Some(response_body))
-            }
-
-            // IMDS returns 404 if no user data is given, or if IMDS is disabled, for example
-            StatusCode::NOT_FOUND => Ok(None),
-
-            code @ _ => {
-                let response_body = response
-                    .bytes()
-                    .context(error::ResponseBody {
-                        method: "GET",
-                        uri,
-                        code,
-                    })?
-                    .to_vec();
-
-                let response_str = response_string(&response_body);
-
-                trace!("Response: {:?}", response_str);
-
-                error::Response {
-                    method: "GET",
-                    uri,
-                    code,
-                    response_body: response_str,
-                }
-                .fail()
-            }
-        }
-    }
 
     /// Fetches user data, which is expected to be in TOML form and contain a `[settings]` section,
     /// returning a SettingsJson representing the inside of that section.
-    fn user_data(client: &Client, session_token: &str) -> Result<Option<SettingsJson>> {
-        let desc = "user data";
-        let uri = Self::USER_DATA_ENDPOINT;
-
-        let user_data_raw = match Self::fetch_imds(client, session_token, uri, desc) {
-            Err(e) => return Err(e),
-            Ok(None) => return Ok(None),
-            Ok(Some(s)) => s,
-        };
+    async fn user_data(client: &mut ImdsClient) -> Result<Option<SettingsJson>> {
+        let user_data_raw = client.fetch_userdata().await.context(error::ImdsRequest)?;
         let user_data_str = expand_slice_maybe(&user_data_raw)
             .context(error::Decompression { what: "user data" })?;
         trace!("Received user data: {}", user_data_str);
@@ -144,31 +33,36 @@ impl AwsDataProvider {
 
     /// Fetches the instance identity, returning a SettingsJson representing the values from the
     /// document which we'd like to send to the API - currently just region.
-    fn identity_document(client: &Client, session_token: &str) -> Result<Option<SettingsJson>> {
+    async fn identity_document(client: &mut ImdsClient) -> Result<Option<SettingsJson>> {
         let desc = "instance identity document";
-        let uri = Self::IDENTITY_DOCUMENT_ENDPOINT;
         let file = Self::IDENTITY_DOCUMENT_FILE;
 
-        let iid_str = if Path::new(file).exists() {
+        let region = if Path::new(file).exists() {
             info!("{} found at {}, using it", desc, file);
-            fs::read_to_string(file).context(error::InputFileRead { path: file })?
+            let data = fs::read_to_string(file).context(error::InputFileRead { path: file })?;
+            let iid: serde_json::Value =
+                serde_json::from_str(&data).context(error::DeserializeJson)?;
+            iid.get("region")
+                .context(error::IdentityDocMissingData { missing: "region" })?
+                .as_str()
+                .context(error::WrongType {
+                    field_name: "region",
+                    expected_type: "string",
+                })?
+                .to_owned()
         } else {
-            match Self::fetch_imds(client, session_token, uri, desc) {
-                Err(e) => return Err(e),
-                Ok(None) => return Ok(None),
-                Ok(Some(raw)) => {
-                    expand_slice_maybe(&raw).context(error::Decompression { what: "user data" })?
-                }
-            }
+            client
+                .fetch_identity_document()
+                .await
+                .context(error::ImdsRequest)?
+                .region()
+                .to_owned()
         };
-        trace!("Received instance identity document: {}", iid_str);
+        trace!(
+            "Retrieved region from instance identity document: {}",
+            region
+        );
 
-        // Grab region from instance identity document.
-        let iid: serde_json::Value =
-            serde_json::from_str(&iid_str).context(error::DeserializeJson)?;
-        let region = iid
-            .get("region")
-            .context(error::IdentityDocMissingData { missing: "region" })?;
         let val = json!({ "aws": {"region": region} });
 
         let json = SettingsJson::from_val(&val, desc).context(error::SettingsToJSON {
@@ -178,26 +72,26 @@ impl AwsDataProvider {
     }
 }
 
+#[async_trait]
 impl PlatformDataProvider for AwsDataProvider {
     /// Return settings changes from the instance identity document and user data.
-    fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
+    async fn platform_data(
+        &self,
+    ) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
         let mut output = Vec::new();
-        let client = Client::new();
 
-        let session_token = Self::fetch_imds_session_token(&client)?;
+        let mut client = ImdsClient::new().await.context(error::ImdsClient)?;
 
         // Instance identity doc first, so the user has a chance to override
-        match Self::identity_document(&client, &session_token) {
-            Err(e) => return Err(e).map_err(Into::into),
-            Ok(None) => warn!("No instance identity document found."),
-            Ok(Some(s)) => output.push(s),
+        match Self::identity_document(&mut client).await? {
+            None => warn!("No instance identity document found."),
+            Some(s) => output.push(s),
         }
 
         // Optional user-specified configuration / overrides
-        match Self::user_data(&client, &session_token) {
-            Err(e) => return Err(e).map_err(Into::into),
-            Ok(None) => warn!("No user data found."),
-            Ok(Some(s)) => output.push(s),
+        match Self::user_data(&mut client).await? {
+            None => warn!("No user data found."),
+            Some(s) => output.push(s),
         }
 
         Ok(output)
@@ -205,28 +99,13 @@ impl PlatformDataProvider for AwsDataProvider {
 }
 
 mod error {
-    use http::StatusCode;
     use snafu::Snafu;
     use std::io;
     use std::path::PathBuf;
 
-    // Taken from pluto.
-    // Extracts the status code from a reqwest::Error and converts it to a string to be displayed
-    fn get_bad_status_code(source: &reqwest::Error) -> String {
-        source
-            .status()
-            .as_ref()
-            .map(|i| i.as_str())
-            .unwrap_or("Unknown")
-            .to_string()
-    }
-
     #[derive(Debug, Snafu)]
     #[snafu(visibility = "pub(super)")]
     pub(crate) enum Error {
-        #[snafu(display("Response '{}' from '{}': {}", get_bad_status_code(&source), uri, source))]
-        BadResponse { uri: String, source: reqwest::Error },
-
         #[snafu(display("Failed to decompress {}: {}", what, source))]
         Decompression { what: String, source: io::Error },
 
@@ -236,43 +115,30 @@ mod error {
         #[snafu(display("Instance identity document missing {}", missing))]
         IdentityDocMissingData { missing: String },
 
+        #[snafu(display("IMDS client failed: {}", source))]
+        ImdsClient { source: imdsclient::Error },
+
         #[snafu(display("Unable to read input file '{}': {}", path.display(), source))]
         InputFileRead { path: PathBuf, source: io::Error },
 
-        #[snafu(display("Error {}ing '{}': {}", method, uri, source))]
-        Request {
-            method: String,
-            uri: String,
-            source: reqwest::Error,
-        },
-
-        #[snafu(display("Error {} when {}ing '{}': {}", code, method, uri, response_body))]
-        Response {
-            method: String,
-            uri: String,
-            code: StatusCode,
-            response_body: String,
-        },
-
-        #[snafu(display(
-            "Unable to read response body when {}ing '{}' (code {}) - {}",
-            method,
-            uri,
-            code,
-            source
-        ))]
-        ResponseBody {
-            method: String,
-            uri: String,
-            code: StatusCode,
-            source: reqwest::Error,
-        },
+        #[snafu(display("IMDS request failed: {}", source))]
+        ImdsRequest { source: imdsclient::Error },
 
         #[snafu(display("Unable to serialize settings from {}: {}", from, source))]
         SettingsToJSON {
             from: String,
             source: crate::settings::Error,
         },
+
+        #[snafu(display(
+            "Wrong type while deserializing, expected '{}' to be type '{}'",
+            field_name,
+            expected_type
+        ))]
+        WrongType {
+            field_name: &'static str,
+            expected_type: &'static str,
+        },
     }
 }
 
diff --git a/sources/api/early-boot-config/src/provider/local_file.rs b/sources/api/early-boot-config/src/provider/local_file.rs
index 9db23235100..69dcb45cb53 100644
--- a/sources/api/early-boot-config/src/provider/local_file.rs
+++ b/sources/api/early-boot-config/src/provider/local_file.rs
@@ -3,6 +3,7 @@
 
 use super::{PlatformDataProvider, SettingsJson};
 use crate::compression::expand_file_maybe;
+use async_trait::async_trait;
 use snafu::ResultExt;
 
 pub(crate) struct LocalFileDataProvider;
@@ -11,8 +12,9 @@ impl LocalFileDataProvider {
     pub(crate) const USER_DATA_FILE: &'static str = "/etc/early-boot-config/user-data";
 }
 
+#[async_trait]
 impl PlatformDataProvider for LocalFileDataProvider {
-    fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
+    async fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
         let mut output = Vec::new();
         info!("'{}' exists, using it", Self::USER_DATA_FILE);
 
diff --git a/sources/api/early-boot-config/src/provider/vmware.rs b/sources/api/early-boot-config/src/provider/vmware.rs
index 47ed3f207a8..eeb8529cfc8 100644
--- a/sources/api/early-boot-config/src/provider/vmware.rs
+++ b/sources/api/early-boot-config/src/provider/vmware.rs
@@ -3,6 +3,7 @@
 
 use super::{PlatformDataProvider, SettingsJson};
 use crate::compression::{expand_file_maybe, expand_slice_maybe, OptionalCompressionReader};
+use async_trait::async_trait;
 use serde::Deserialize;
 use snafu::{ensure, ResultExt};
 use std::ffi::OsStr;
@@ -235,8 +236,9 @@ impl VmwareDataProvider {
     }
 }
 
+#[async_trait]
 impl PlatformDataProvider for VmwareDataProvider {
-    fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
+    async fn platform_data(&self) -> std::result::Result<Vec<SettingsJson>, Box<dyn std::error::Error>> {
         let mut output = Vec::new();
 
         // Look at the CD-ROM for user data first, and then...