Add AwsBuilder::expose_ports_to_internet (#43)

aws-throwaway/Cargo.toml

@@ -25,10 +25,11 @@ anyhow = "1.0.42"
 uuid = { version = "1.0.0", features = ["serde", "v4"] }
 tracing = "0.1.15"
 async-trait = "0.1.30"
-# TODO: avoid pulling in these dependencies when use_sdk is disabled,
+# TODO: avoid pulling in these dependencies when use_sdk is enabled,
 #       will need to introduce a use_cli feature to do so
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = "1.0.111"
+futures = "0.3.30"
 
 [dev-dependencies]
 tracing-subscriber = { version = "0.3.1", features = ["env-filter", "json"] }

aws-throwaway/src/backend/cli/mod.rs

@@ -1,22 +1,24 @@
 mod instance_type;
-use std::{
-    net::IpAddr,
-    time::{Duration, Instant},
-};
-
-use base64::Engine;
-pub use instance_type::InstanceType;
-
 mod placement_strategy;
-pub use placement_strategy::PlacementStrategy;
-use ssh_key::{rand_core::OsRng, PrivateKey};
 
 use crate::{
     backend::cli::instance_type::get_arch_of_instance_type, AwsBuilder, CleanupResources,
     Ec2Instance, Ec2InstanceDefinition, InstanceOs, NetworkInterface, APP_TAG_NAME, USER_TAG_NAME,
 };
 use anyhow::{anyhow, Result};
+use base64::Engine;
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+pub use instance_type::InstanceType;
+pub use placement_strategy::PlacementStrategy;
 use serde::Deserialize;
+use ssh_key::{rand_core::OsRng, PrivateKey};
+use std::future::Future;
+use std::pin::Pin;
+use std::{
+    net::IpAddr,
+    time::{Duration, Instant},
+};
 use uuid::Uuid;
 
 #[derive(Debug, Deserialize)]
@@ -137,7 +139,8 @@ impl Aws {
                 &tags,
                 &security_group_name,
                 &builder.vpc_id,
-                builder.security_group_id
+                builder.security_group_id,
+                &builder.expose_ports_to_internet
             ),
             Aws::create_placement_group(&tags, &placement_group_name, builder.placement_strategy),
             Aws::get_subnet(builder.subnet_id)
@@ -195,6 +198,7 @@ impl Aws {
         name: &str,
         vpc_id: &Option<String>,
         security_group_id: Option<String>,
+        ports: &[u16],
     ) -> String {
         match security_group_id {
             Some(id) => id,
@@ -222,10 +226,19 @@ impl Aws {
                 let result: SecurityGroup = run_command(&command).await.unwrap();
                 tracing::info!("created security group");
 
-                tokio::join!(
-                    Aws::create_ingress_rule_internal(tags, name),
-                    Aws::create_ingress_rule_ssh(tags, name),
-                );
+                let mut futures = FuturesUnordered::<Pin<Box<dyn Future<Output = ()>>>>::new();
+                futures.push(Box::pin(Aws::create_ingress_rule_internal(tags, name)));
+                if !ports.contains(&22) {
+                    // SSH
+                    futures.push(Box::pin(Aws::create_ingress_rule_for_port(tags, name, 22)));
+                }
+                for port in ports {
+                    futures.push(Box::pin(Aws::create_ingress_rule_for_port(
+                        tags, name, *port,
+                    )));
+                }
+                while futures.next().await.is_some() {}
+
                 result.group_id
             }
         }
@@ -247,7 +260,8 @@ impl Aws {
         tracing::info!("created security group rule - internal");
     }
 
-    async fn create_ingress_rule_ssh(tags: &Tags, group_name: &str) {
+    async fn create_ingress_rule_for_port(tags: &Tags, group_name: &str, port: u16) {
+        let port = port.to_string();
         let _result: Ignore = run_command(&[
             "ec2",
             "authorize-security-group-ingress",
@@ -256,17 +270,17 @@ impl Aws {
             "--ip-protocol",
             "tcp",
             "--from-port",
-            "22",
+            &port,
             "--to-port",
-            "22",
+            &port,
             "--cidr-ip",
             "0.0.0.0/0",
             "--tag-specifications",
-            &tags.create_tags("security-group-rule", "ssh"),
+            &tags.create_tags("security-group-rule", &format!("port {port}")),
         ])
         .await
         .unwrap();
-        tracing::info!("created security group rule - ssh");
+        tracing::info!("created security group rule - port {port}");
     }
 
     async fn create_placement_group(tags: &Tags, name: &str, strategy: PlacementStrategy) {

aws-throwaway/src/backend/sdk/aws.rs

@@ -1,26 +1,28 @@
 use super::tags::Tags;
+use crate::ec2_instance::{Ec2Instance, NetworkInterface};
+use crate::ec2_instance_definition::{Ec2InstanceDefinition, InstanceOs};
+use crate::AwsBuilder;
+use crate::CleanupResources;
 use anyhow::anyhow;
 use aws_config::meta::region::RegionProviderChain;
 use aws_config::retry::ProvideErrorKind;
 use aws_config::{BehaviorVersion, SdkConfig};
 use aws_sdk_ec2::config::Region;
+use aws_sdk_ec2::types::PlacementStrategy;
 use aws_sdk_ec2::types::{
     BlockDeviceMapping, EbsBlockDevice, Filter, InstanceNetworkInterfaceSpecification, KeyType,
     Placement, ResourceType, Subnet, VolumeType,
 };
 use base64::Engine;
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
 use ssh_key::rand_core::OsRng;
 use ssh_key::PrivateKey;
+use std::future::Future;
+use std::pin::Pin;
 use std::time::{Duration, Instant};
 use uuid::Uuid;
 
-use crate::ec2_instance::{Ec2Instance, NetworkInterface};
-use crate::ec2_instance_definition::{Ec2InstanceDefinition, InstanceOs};
-use crate::AwsBuilder;
-use crate::CleanupResources;
-
-use aws_sdk_ec2::types::PlacementStrategy;
-
 const AZ: &str = "us-east-1c";
 
 async fn config() -> SdkConfig {
@@ -71,7 +73,8 @@ impl Aws {
                 &tags,
                 &security_group_name,
                 &builder.vpc_id,
-                builder.security_group_id
+                builder.security_group_id,
+                &builder.expose_ports_to_internet
             ),
             Aws::create_placement_group(
                 &client,
@@ -136,6 +139,7 @@ impl Aws {
         name: &str,
         vpc_id: &Option<String>,
         security_group_id: Option<String>,
+        ports: &[u16],
     ) -> String {
         match security_group_id {
             Some(id) => id,
@@ -156,10 +160,22 @@ impl Aws {
                     .unwrap();
                 tracing::info!("created security group");
 
-                tokio::join!(
-                    Aws::create_ingress_rule_internal(client, tags, name),
-                    Aws::create_ingress_rule_ssh(client, tags, name),
-                );
+                let mut futures = FuturesUnordered::<Pin<Box<dyn Future<Output = ()>>>>::new();
+                futures.push(Box::pin(Aws::create_ingress_rule_internal(
+                    client, tags, name,
+                )));
+                // SSH
+                if !ports.contains(&22) {
+                    futures.push(Box::pin(Aws::create_ingress_rule_for_port(
+                        client, tags, name, 22,
+                    )));
+                }
+                for port in ports {
+                    futures.push(Box::pin(Aws::create_ingress_rule_for_port(
+                        client, tags, name, *port,
+                    )));
+                }
+                while futures.next().await.is_some() {}
                 security_group_id
             }
         }
@@ -186,22 +202,30 @@ impl Aws {
         tracing::info!("created security group rule - internal");
     }
 
-    async fn create_ingress_rule_ssh(client: &aws_sdk_ec2::Client, tags: &Tags, group_name: &str) {
+    async fn create_ingress_rule_for_port(
+        client: &aws_sdk_ec2::Client,
+        tags: &Tags,
+        group_name: &str,
+        port: u16,
+    ) {
+        let port = port.to_string();
         assert!(client
             .authorize_security_group_ingress()
             .group_name(group_name)
             .ip_protocol("tcp")
             .from_port(22)
             .to_port(22)
             .cidr_ip("0.0.0.0/0")
-            .tag_specifications(tags.create_tags(ResourceType::SecurityGroupRule, "ssh"))
+            .tag_specifications(
+                tags.create_tags(ResourceType::SecurityGroupRule, &format!("port {port}"))
+            )
             .send()
             .await
             .map_err(|e| e.into_service_error())
             .unwrap()
             .r#return()
             .unwrap());
-        tracing::info!("created security group rule - ssh");
+        tracing::info!("created security group rule - port {port}");
     }
 
     async fn create_placement_group(

aws-throwaway/src/backend/sdk/tags.rs

@@ -12,7 +12,7 @@ impl Tags {
     pub fn create_tags(
         &self,
         resource_type: ResourceType,
-        resource_name: &'static str,
+        resource_name: &str,
     ) -> TagSpecification {
         let mut builder = TagSpecification::builder()
             .resource_type(resource_type)

aws-throwaway/src/lib.rs

@@ -19,6 +19,7 @@ pub struct AwsBuilder {
     subnet_id: Option<String>,
     placement_strategy: PlacementStrategy,
     security_group_id: Option<String>,
+    expose_ports_to_internet: Vec<u16>,
 }
 
 /// The default configuration will succeed for an AMI user with sufficient access and unmodified default vpcs/subnets
@@ -42,6 +43,7 @@ impl AwsBuilder {
             // I believe Spread is the best default since it is the easiest for amazon to fulfill and gives the most realistic results in benchmarks.
             placement_strategy: PlacementStrategy::Spread,
             security_group_id: None,
+            expose_ports_to_internet: vec![],
         }
     }
 
@@ -98,7 +100,19 @@ impl AwsBuilder {
         self
     }
 
+    /// Adds the provided ports as allowing traffic in+out to internet in the automatically generated security group.
+    pub fn expose_ports_to_internet(mut self, ports: Vec<u16>) -> Self {
+        self.expose_ports_to_internet = ports;
+        self
+    }
+
+    /// Builds the Aws instance.
+    ///
+    /// Will panic if both `expose_ports_to_internet` and `use_security_group_id` are enabled.
     pub async fn build(self) -> Aws {
+        if !self.expose_ports_to_internet.is_empty() && self.security_group_id.is_some() {
+            panic!("Both `use_security_group_id` and `expose_ports_to_internet` are set. Ensure only one of these options is set.")
+        }
         Aws::new(self).await
     }
 }