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..dc2d0beac40 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" @@ -836,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", @@ -1399,6 +1421,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" @@ -2105,8 +2145,8 @@ version = "0.1.0" dependencies = [ "apiclient", "cargo-readme", + "imdsclient", "models", - "reqwest", "rusoto_core", "rusoto_eks", "serde_json", @@ -2742,12 +2782,13 @@ version = "0.1.0" dependencies = [ "base64", "cargo-readme", + "imdsclient", "log", - "reqwest", "serde", "serde_json", "simplelog", "snafu", + "tokio", ] [[package]] @@ -3251,6 +3292,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/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, Box>; + async fn platform_data( + &self, + ) -> std::result::Result, Box>; } 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 { - 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>> { - 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 { - "".to_string() - } else if let Ok(s) = String::from_utf8(response.into()) { - s - } else { - "".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> { - 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> { + 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> { + async fn identity_document(client: &mut ImdsClient) -> Result> { 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, Box> { + async fn platform_data( + &self, + ) -> std::result::Result, Box> { 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, Box> { + async fn platform_data(&self) -> std::result::Result, Box> { 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, Box> { + async fn platform_data(&self) -> std::result::Result, Box> { let mut output = Vec::new(); // Look at the CD-ROM for user data first, and then... 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 = std::result::Result; -async fn get_text_from_imds(client: &Client, uri: &str, session_token: &str) -> Result { - 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 { + 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 { - 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 { /// 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 { +async fn get_cluster_dns_ip(client: &mut ImdsClient) -> Result { // 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 Result { /// 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 { - 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 { + // 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 { - get_text_from_imds(&client, IMDS_NODE_IPV4_ENDPOINT, session_token).await +async fn get_node_ip(client: &mut ImdsClient) -> Result { + 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)), 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 { - 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> { - 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> { - 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 { - 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::().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> { + 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 { } 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 }, 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 "] +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::new_impl(BASE_URI.to_string()).await + } + + async fn new_impl(imds_base_uri: String) -> Result { + 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> { + 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 { + 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> { + 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> { + // 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 { + 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> { + 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(&mut self, end_target: S) -> Result> + where + S: AsRef, + { + 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(&mut self, end_target: S) -> Result + where + S: AsRef, + { + 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(&mut self, schema_version: S1, target: S2) -> Result> + where + S1: AsRef, + S2: AsRef, + { + 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!("{}", &s[0..2034]) + } + } else { + "".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 { + 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::().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 { + 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 = std::result::Result; + +#[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 = "".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(""); + 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() + ); + } +}