Merge pull request #3020 from mjsterckx/pubsys-ami-validation

pubsys: added ami validation

tools/pubsys/Cargo.toml

@@ -38,6 +38,7 @@ reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"
 semver = "1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+serde_plain = "1"
 simplelog = "0.12"
 snafu = "0.7"
 structopt = { version = "0.3", default-features = false }

tools/pubsys/src/aws/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod ami;
 pub(crate) mod promote_ssm;
 pub(crate) mod publish_ami;
 pub(crate) mod ssm;
+pub(crate) mod validate_ami;
 pub(crate) mod validate_ssm;
 
 /// Builds a Region from the given region name.

tools/pubsys/src/aws/validate_ami/ami.rs

@@ -0,0 +1,224 @@
+//! The ami module owns the describing of images in EC2.
+
+use aws_sdk_ec2::model::Image;
+use aws_sdk_ec2::{Client as Ec2Client, Region};
+use futures::future::{join, ready};
+use futures::stream::{FuturesUnordered, StreamExt};
+use log::{info, trace};
+use serde::{Deserialize, Serialize};
+use snafu::ResultExt;
+use std::collections::HashMap;
+
+use crate::aws::ami::launch_permissions::{get_launch_permissions, LaunchPermissionDef};
+
+/// Wrapper structure for the `ImageDef` struct, used during deserialization
+#[derive(Deserialize)]
+#[serde(untagged)]
+pub(crate) enum ImageData {
+    Image(ImageDef),
+    ImageList(Vec<ImageDef>),
+}
+
+impl ImageData {
+    pub(crate) fn images(&self) -> Vec<ImageDef> {
+        match self {
+            ImageData::Image(image) => vec![image.to_owned()],
+            ImageData::ImageList(images) => images.to_owned(),
+        }
+    }
+}
+
+/// Structure of the EC2 image fields that should be validated
+#[derive(Serialize, Deserialize, PartialEq, Eq, Hash, Debug, Clone)]
+pub(crate) struct ImageDef {
+    /// The ID of the EC2 image
+    pub(crate) id: String,
+
+    /// The name of the EC2 image
+    pub(crate) name: String,
+
+    /// Whether or not the EC2 image is public
+    #[serde(default)]
+    pub(crate) public: bool,
+
+    /// The launch permissions for the EC2 image.
+    pub(crate) launch_permissions: Option<Vec<LaunchPermissionDef>>,
+
+    /// Whether or not the EC2 image supports Elastic Network Adapter
+    #[serde(default = "default_ena_support")]
+    pub(crate) ena_support: bool,
+
+    /// The level of the EC2 image's Single Root I/O Virtualization support
+    #[serde(default = "default_sriov_net_support")]
+    pub(crate) sriov_net_support: String,
+}
+
+fn default_ena_support() -> bool {
+    true
+}
+
+fn default_sriov_net_support() -> String {
+    "simple".to_string()
+}
+
+impl From<(Image, Option<Vec<LaunchPermissionDef>>)> for ImageDef {
+    fn from(args: (Image, Option<Vec<LaunchPermissionDef>>)) -> Self {
+        Self {
+            id: args.0.image_id().unwrap_or_default().to_string(),
+            name: args.0.name().unwrap_or_default().to_string(),
+            public: args.0.public().unwrap_or_default(),
+            launch_permissions: args.1,
+            ena_support: args.0.ena_support().unwrap_or_default(),
+            sriov_net_support: args.0.sriov_net_support().unwrap_or_default().to_string(),
+        }
+    }
+}
+
+/// Fetches all images whose IDs are keys in `expected_images`. The map `expected_image_public` is
+/// used to determine if the launch permissions for the image should be fetched (only if the image is not
+/// public). The return value is a HashMap of Region to a Result, which is `Ok` if the request for
+/// that region was successful and `Err` if not. The Result contains a HashMap of `image_id` to
+/// `ImageDef`.
+pub(crate) async fn describe_images<'a>(
+    clients: &'a HashMap<Region, Ec2Client>,
+    expected_images: &HashMap<Region, Vec<ImageDef>>,
+) -> HashMap<&'a Region, Result<HashMap<String, ImageDef>>> {
+    // Build requests for images; we have to request with a regional client so we split them by
+    // region
+    let mut requests = Vec::with_capacity(clients.len());
+    clients.iter().for_each(|(region, ec2_client)| {
+        trace!("Requesting images in {}", region);
+        let get_future = describe_images_in_region(
+            region,
+            ec2_client,
+            expected_images
+                .get(region)
+                .map(|i| i.to_owned())
+                .unwrap_or_default()
+                .into_iter()
+                .map(|i| (i.id.clone(), i))
+                .collect::<HashMap<String, ImageDef>>(),
+        );
+
+        requests.push(join(ready(region), get_future));
+    });
+
+    // Send requests in parallel and wait for responses, collecting results into a list.
+    requests
+        .into_iter()
+        .collect::<FuturesUnordered<_>>()
+        .collect()
+        .await
+}
+
+/// Fetches the images whose IDs are keys in `expected_images`
+pub(crate) async fn describe_images_in_region(
+    region: &Region,
+    client: &Ec2Client,
+    expected_images: HashMap<String, ImageDef>,
+) -> Result<HashMap<String, ImageDef>> {
+    info!("Retrieving images in {}", region.to_string());
+    let mut images = HashMap::new();
+
+    // Send the request
+    let mut get_future = client
+        .describe_images()
+        .include_deprecated(true)
+        .set_image_ids(Some(Vec::from_iter(
+            expected_images.keys().map(|k| k.to_owned()),
+        )))
+        .into_paginator()
+        .send();
+
+    // Iterate over the retrieved images
+    while let Some(page) = get_future.next().await {
+        let retrieved_images = page
+            .context(error::DescribeImagesSnafu {
+                region: region.to_string(),
+            })?
+            .images()
+            .unwrap_or_default()
+            .to_owned();
+        for image in retrieved_images {
+            // Insert a new key-value pair into the map, with the key containing image ID
+            // and the value containing the ImageDef object created from the image
+            let image_id = image
+                .image_id()
+                .ok_or(error::Error::MissingField {
+                    missing: "image_id".to_string(),
+                })?
+                .to_string();
+            let expected_public = expected_images
+                .get(&image_id)
+                .ok_or(error::Error::MissingExpectedPublic {
+                    missing: image_id.clone(),
+                })?
+                .public;
+            // If the image is not expected to be public, retrieve the launch permissions
+            trace!(
+                "Retrieving launch permissions for {} in {}",
+                image_id,
+                region.as_ref()
+            );
+            let launch_permissions = if !expected_public {
+                Some(
+                    get_launch_permissions(client, region.as_ref(), &image_id)
+                        .await
+                        .context(error::GetLaunchPermissionsSnafu {
+                            region: region.as_ref(),
+                            image_id: image_id.clone(),
+                        })?,
+                )
+            } else {
+                None
+            };
+            let image_def = ImageDef::from((image.to_owned(), launch_permissions));
+            images.insert(image_id, image_def);
+        }
+    }
+
+    info!("Images in {} have been retrieved", region.to_string());
+    Ok(images)
+}
+
+pub(crate) mod error {
+    use aws_sdk_ec2::error::DescribeImagesError;
+    use aws_sdk_ssm::types::SdkError;
+    use aws_smithy_types::error::display::DisplayErrorContext;
+    use snafu::Snafu;
+
+    #[derive(Debug, Snafu)]
+    #[snafu(visibility(pub(super)))]
+    #[allow(clippy::large_enum_variant)]
+    pub(crate) enum Error {
+        #[snafu(display(
+            "Failed to describe images in {}: {}",
+            region,
+            DisplayErrorContext(source)
+        ))]
+        DescribeImages {
+            region: String,
+            source: SdkError<DescribeImagesError>,
+        },
+
+        #[snafu(display(
+            "Failed to retrieve launch permissions for image {} in region {}: {}",
+            image_id,
+            region,
+            source
+        ))]
+        GetLaunchPermissions {
+            region: String,
+            image_id: String,
+            source: crate::aws::ami::launch_permissions::Error,
+        },
+
+        #[snafu(display("Missing field in image: {}", missing))]
+        MissingField { missing: String },
+
+        #[snafu(display("Missing image ID in expected image publicity map: {}", missing))]
+        MissingExpectedPublic { missing: String },
+    }
+}
+
+pub(crate) type Result<T> = std::result::Result<T, error::Error>;