From 3632fbdc973051208ab94b5f39cde32ce747faf8 Mon Sep 17 00:00:00 2001 From: Lucas Kent Date: Fri, 2 Feb 2024 11:31:56 +1100 Subject: [PATCH] Add AwsBuilder::expose_ports_to_internet (#43) --- Cargo.lock | 37 +++++++++--------- aws-throwaway/Cargo.toml | 3 +- aws-throwaway/src/backend/cli/mod.rs | 54 +++++++++++++++++---------- aws-throwaway/src/backend/sdk/aws.rs | 54 +++++++++++++++++++-------- aws-throwaway/src/backend/sdk/tags.rs | 2 +- aws-throwaway/src/lib.rs | 14 +++++++ 6 files changed, 109 insertions(+), 55 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f5558a5..1630e17 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -483,6 +483,7 @@ dependencies = [ "aws-sdk-iam", "base64", "clap", + "futures", "russh", "russh-keys", "serde", @@ -1021,9 +1022,9 @@ dependencies = [ [[package]] name = "futures" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0290714b38af9b4a7b094b8a37086d1b4e61f2df9122c3cad2577669145335" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" dependencies = [ "futures-channel", "futures-core", @@ -1036,9 +1037,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff4dd66668b557604244583e3e1e1eada8c5c2e96a6d0d6653ede395b78bbacb" +checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" dependencies = [ "futures-core", "futures-sink", @@ -1046,15 +1047,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb1d22c66e66d9d72e1758f0bd7d4fd0bee04cad842ee34587d68c07e45d088c" +checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" [[package]] name = "futures-executor" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f4fb8693db0cf099eadcca0efe2a5a22e4550f98ed16aba6c48700da29597bc" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" dependencies = [ "futures-core", "futures-task", @@ -1063,15 +1064,15 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bf34a163b5c4c52d0478a4d757da8fb65cabef42ba90515efee0f6f9fa45aaa" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" [[package]] name = "futures-macro" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53b153fd91e4b0147f4aced87be237c98248656bb01050b96bf3ee89220a8ddb" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" dependencies = [ "proc-macro2", "quote", @@ -1080,21 +1081,21 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e36d3378ee38c2a36ad710c5d30c2911d752cb941c00c72dbabfb786a7970817" +checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" [[package]] name = "futures-task" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efd193069b0ddadc69c46389b740bbccdd97203899b48d09c5f7969591d6bae2" +checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" [[package]] name = "futures-util" -version = "0.3.29" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a19526d624e703a3179b3d322efec918b6246ea0fa51d41124525f00f1cc8104" +checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ "futures-channel", "futures-core", diff --git a/aws-throwaway/Cargo.toml b/aws-throwaway/Cargo.toml index 49a8025..96e5ef1 100644 --- a/aws-throwaway/Cargo.toml +++ b/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"] } diff --git a/aws-throwaway/src/backend/cli/mod.rs b/aws-throwaway/src/backend/cli/mod.rs index 7f8b3f4..3637b58 100644 --- a/aws-throwaway/src/backend/cli/mod.rs +++ b/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, security_group_id: Option, + 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::>>>::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) { diff --git a/aws-throwaway/src/backend/sdk/aws.rs b/aws-throwaway/src/backend/sdk/aws.rs index dfa1cac..6a84b62 100644 --- a/aws-throwaway/src/backend/sdk/aws.rs +++ b/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, security_group_id: Option, + 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::>>>::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,7 +202,13 @@ 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) @@ -194,14 +216,16 @@ impl Aws { .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( diff --git a/aws-throwaway/src/backend/sdk/tags.rs b/aws-throwaway/src/backend/sdk/tags.rs index 4aa6259..6d30c02 100644 --- a/aws-throwaway/src/backend/sdk/tags.rs +++ b/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) diff --git a/aws-throwaway/src/lib.rs b/aws-throwaway/src/lib.rs index 5d3a58c..788423f 100644 --- a/aws-throwaway/src/lib.rs +++ b/aws-throwaway/src/lib.rs @@ -19,6 +19,7 @@ pub struct AwsBuilder { subnet_id: Option, placement_strategy: PlacementStrategy, security_group_id: Option, + expose_ports_to_internet: Vec, } /// 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) -> 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 } }