early-boot-config: update to use imdsclient library

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

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"

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

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>>;
 }

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,55 +72,40 @@ 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)
     }
 }
 
 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,
+        },
     }
 }
 

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);