pluto: add regional map of pause container source accounts

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"

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<String> {
+    // Look up the pause container account, or fall back to the default ID and region
+    let (region, account) = match PAUSE_CONTAINER_ACCOUNT.get(&region) {
+        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
+            )
+        );
+    }
+}