diff --git a/sources/Cargo.lock b/sources/Cargo.lock index d51adc3e32f..2238d2e9b0e 100644 --- a/sources/Cargo.lock +++ b/sources/Cargo.lock @@ -1805,6 +1805,7 @@ name = "pluto" version = "0.1.0" dependencies = [ "cargo-readme 3.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)", "reqwest 0.10.7 (registry+https://github.com/rust-lang/crates.io-index)", "serde_json 1.0.57 (registry+https://github.com/rust-lang/crates.io-index)", "snafu 0.6.8 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/sources/api/pluto/Cargo.toml b/sources/api/pluto/Cargo.toml index 0004839f36b..2e7603608af 100644 --- a/sources/api/pluto/Cargo.toml +++ b/sources/api/pluto/Cargo.toml @@ -10,6 +10,7 @@ build = "build.rs" exclude = ["README.md"] [dependencies] +lazy_static = "1.4" reqwest = { version = "0.10", default-features = false, features = ["blocking"]} serde_json = "1" snafu = "0.6" diff --git a/sources/api/pluto/src/main.rs b/sources/api/pluto/src/main.rs index 7dc34b03363..a858ee54130 100644 --- a/sources/api/pluto/src/main.rs +++ b/sources/api/pluto/src/main.rs @@ -14,7 +14,9 @@ It makes calls to IMDS to get meta data: - POD Infra Container Image */ +use lazy_static::lazy_static; use reqwest::blocking::Client; +use std::collections::HashMap; use std::fs::File; use std::io::{BufRead, BufReader}; use std::string::String; @@ -40,10 +42,46 @@ const IMDS_INSTANCE_TYPE_ENDPOINT: &str = const IMDS_INSTANCE_IDENTITY_DOCUMENT_ENDPOINT: &str = "http://169.254.169.254/2018-09-24/dynamic/instance-identity/document"; -const PAUSE_CONTAINER_ACCOUNT: &str = "602401143452"; +const ENI_MAX_PODS_PATH: &str = "/usr/share/eks/eni-max-pods"; + const PAUSE_CONTAINER_VERSION: &str = "3.1"; +lazy_static! { + /// A map to tell us which account to pull pause container images from for a given region. + static ref PAUSE_CONTAINER_ACCOUNT: HashMap<&'static str, &'static str> = { + let mut m = HashMap::new(); + m.insert("af-south-1", "877085696533"); + m.insert("ap-east-1", "800184023465"); + m.insert("ap-northeast-1", "602401143452"); + m.insert("ap-northeast-2", "602401143452"); + m.insert("ap-south-1", "602401143452"); + m.insert("ap-southeast-1", "602401143452"); + m.insert("ap-southeast-2", "602401143452"); + m.insert("ca-central-1", "602401143452"); + m.insert("cn-north-1", "918309763551"); + m.insert("cn-northwest-1", "961992271922"); + m.insert("eu-central-1", "602401143452"); + m.insert("eu-north-1", "602401143452"); + m.insert("eu-south-1", "590381155156"); + m.insert("eu-south-1", "590381155156"); + m.insert("eu-west-1", "602401143452"); + m.insert("eu-west-2", "602401143452"); + m.insert("eu-west-3", "602401143452"); + m.insert("me-south-1", "558608220178"); + m.insert("sa-east-1", "602401143452"); + m.insert("us-east-1", "602401143452"); + m.insert("us-east-2", "602401143452"); + m.insert("us-gov-east-1", "151742754352"); + m.insert("us-gov-west-1", "013241004608"); + m.insert("us-west-1", "602401143452"); + m.insert("us-west-2", "602401143452"); + m + }; +} -const ENI_MAX_PODS_PATH: &str = "/usr/share/eks/eni-max-pods"; +/// But if there is a region that does not exist in our map (for example a new +/// region is created or being tested), then we will fall back to this. +const PAUSE_FALLBACK_ACCOUNT: &str = "602401143452"; +const PAUSE_FALLBACK_REGION: &str = "us-east-1"; mod error { use snafu::Snafu; @@ -200,18 +238,35 @@ fn get_pod_infra_container_image(client: &Client, session_token: &str) -> Result .as_str() .context(error::MissingRegion { uri })?; - // Get machine architecture. - let arch = if cfg!(target_arch = "x86_64") { - "amd64" + pause_container_uri(region) +} + +/// Returns the machine architecture. +fn arch() -> Result<&'static str> { + if cfg!(target_arch = "x86_64") { + Ok("amd64") } else if cfg!(target_arch = "aarch64") { - "arm64" + Ok("arm64") } else { - return error::UnknownArchitecture.fail(); + error::UnknownArchitecture.fail() + } +} + +/// Constructs the URI of the pause container image for the given region. Returns a URI for the +/// default region/account if the region is not mapped. +fn pause_container_uri(region: &str) -> Result { + // Look up the pause container account, or fall back to the default ID and region + let (region, account) = match PAUSE_CONTAINER_ACCOUNT.get(®ion) { + Some(account) => (region, *account), + None => (PAUSE_FALLBACK_REGION, PAUSE_FALLBACK_ACCOUNT), }; Ok(format!( "{}.dkr.ecr.{}.amazonaws.com/eks/pause-{}:{}", - PAUSE_CONTAINER_ACCOUNT, region, arch, PAUSE_CONTAINER_VERSION + account, + region, + arch()?, + PAUSE_CONTAINER_VERSION )) } @@ -287,3 +342,44 @@ fn main() { process::exit(1); } } + +#[cfg(test)] +mod test_pause_container_account { + use super::{arch, pause_container_uri, PAUSE_CONTAINER_VERSION}; + + #[test] + fn url_eu_west_1() { + assert_eq!( + pause_container_uri("eu-west-1").unwrap(), + format!( + "602401143452.dkr.ecr.eu-west-1.amazonaws.com/eks/pause-{}:{}", + arch().unwrap(), + PAUSE_CONTAINER_VERSION + ) + ); + } + + #[test] + fn url_af_south_1() { + assert_eq!( + pause_container_uri("af-south-1").unwrap(), + format!( + "877085696533.dkr.ecr.af-south-1.amazonaws.com/eks/pause-{}:{}", + arch().unwrap(), + PAUSE_CONTAINER_VERSION + ) + ); + } + + #[test] + fn url_fallback() { + assert_eq!( + pause_container_uri("xy-ztown-1").unwrap(), + format!( + "602401143452.dkr.ecr.us-east-1.amazonaws.com/eks/pause-{}:{}", + arch().unwrap(), + PAUSE_CONTAINER_VERSION + ) + ); + } +}