From 06670822bd26b83fb79de457765134519592c287 Mon Sep 17 00:00:00 2001 From: Lucas Kent Date: Fri, 20 Oct 2023 09:03:59 +1100 Subject: [PATCH] User specified VPC --- .gitignore | 1 + .../aws-throwaway-test-multiple-instances.rs | 2 +- aws-throwaway/examples/aws-throwaway-test.rs | 2 +- aws-throwaway/examples/create-instance.rs | 8 +- aws-throwaway/src/ec2_instance_definition.rs | 13 ++ aws-throwaway/src/lib.rs | 158 +++++++++++------- 6 files changed, 125 insertions(+), 59 deletions(-) diff --git a/.gitignore b/.gitignore index ea8c4bf..1541b63 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /target +some_local_file diff --git a/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs b/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs index b50a39f..3749646 100644 --- a/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs +++ b/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs @@ -10,7 +10,7 @@ async fn main() { .init(); println!("Creating instances"); - let aws = Aws::new(CleanupResources::AllResources).await; + let aws = Aws::new(CleanupResources::AllResources, None, None, None).await; let (instance1, instance2) = tokio::join!( aws.create_ec2_instance(Ec2InstanceDefinition::new(InstanceType::T2Small)), aws.create_ec2_instance( diff --git a/aws-throwaway/examples/aws-throwaway-test.rs b/aws-throwaway/examples/aws-throwaway-test.rs index c67e1ee..050f0f1 100644 --- a/aws-throwaway/examples/aws-throwaway-test.rs +++ b/aws-throwaway/examples/aws-throwaway-test.rs @@ -10,7 +10,7 @@ async fn main() { .with_writer(non_blocking) .init(); - let aws = Aws::new(CleanupResources::AllResources).await; + let aws = Aws::new(CleanupResources::AllResources, None, None, None).await; let instance = aws .create_ec2_instance(Ec2InstanceDefinition::new(InstanceType::T2Micro)) .await; diff --git a/aws-throwaway/examples/create-instance.rs b/aws-throwaway/examples/create-instance.rs index 5ee049a..0be79e4 100644 --- a/aws-throwaway/examples/create-instance.rs +++ b/aws-throwaway/examples/create-instance.rs @@ -21,7 +21,13 @@ async fn main() { } else if let Some(instance_type) = args.instance_type { println!("Creating instance of type {instance_type}"); - let aws = Aws::new(CleanupResources::WithAppTag(AWS_THROWAWAY_TAG.to_owned())).await; + let aws = Aws::new( + CleanupResources::WithAppTag(AWS_THROWAWAY_TAG.to_owned()), + None, + None, + None, + ) + .await; let instance_type = InstanceType::from_str(&instance_type).unwrap(); let network_interface_count = args.network_interfaces; let instance = aws diff --git a/aws-throwaway/src/ec2_instance_definition.rs b/aws-throwaway/src/ec2_instance_definition.rs index f43ac41..848381b 100644 --- a/aws-throwaway/src/ec2_instance_definition.rs +++ b/aws-throwaway/src/ec2_instance_definition.rs @@ -6,6 +6,7 @@ pub struct Ec2InstanceDefinition { pub(crate) volume_size_gb: u32, pub(crate) network_interface_count: u32, pub(crate) os: InstanceOs, + pub(crate) ami: Option, } impl Ec2InstanceDefinition { @@ -16,6 +17,7 @@ impl Ec2InstanceDefinition { volume_size_gb: 8, network_interface_count: 1, os: InstanceOs::Ubuntu22_04, + ami: None, } } @@ -43,6 +45,17 @@ impl Ec2InstanceDefinition { self.os = os; self } + + /// Override the AMI used. + /// When used together with the `os` setting, the os setting is used to determine how to configure the instance while the specified AMI used as the instances image. + /// Defaults to None which indicates that the appropriate AMI should be looked up via SSM. + /// + /// This option is useful when you have custom variation of the configured OS or if your user does not have access to SSM. + /// AMI's are region specific so be careful in picking your AMI. + pub fn override_ami(mut self, ami: Option) -> Self { + self.ami = ami; + self + } } /// aws-throwaway needs to manually support each OS, so the only OS's you can use are those listed in this enum. diff --git a/aws-throwaway/src/lib.rs b/aws-throwaway/src/lib.rs index d1a63bd..ea016c7 100644 --- a/aws-throwaway/src/lib.rs +++ b/aws-throwaway/src/lib.rs @@ -41,7 +41,7 @@ pub struct Aws { host_private_key: String, security_group_id: String, placement_group_name: String, - default_subnet_id: String, + subnet_id: String, tags: Tags, } @@ -50,7 +50,24 @@ impl Aws { /// /// Before returning the [`Aws`], all preexisting resources conforming to the specified [`CleanupResources`] approach are destroyed. /// The specified [`CleanupResources`] is then also used by the [`Aws::cleanup_resources`] method. - pub async fn new(cleanup: CleanupResources) -> Self { + /// + /// All resources will be created in us-east-1c. + /// This is hardcoded so that aws-throawaway only has to look into one region when cleaning up. + /// All instances are created in a single spread placement group in a single AZ to ensure consistent latency between instances. + /// + /// If vpc_id is: + /// * Some(_) => all resources will go into that vpc + /// * None => all resources will go into the default vpc + /// If subnet_id is: + /// * Some(_) => all instances will go into that subnet + /// * None => all instances will go into the default subnet for the specified or default vpc + pub async fn new( + cleanup: CleanupResources, + // TODO: make these into a builder + vpc_id: Option, + subnet_id: Option, + security_group_id: Option, + ) -> Self { let config = config().await; let user_name = iam::user_name(&config).await; let keyname = format!("aws-throwaway-{user_name}-{}", Uuid::new_v4()); @@ -68,9 +85,15 @@ impl Aws { let (client_private_key, security_group_id, _, default_subnet_id) = tokio::join!( Aws::create_key_pair(&client, &tags, &keyname), - Aws::create_security_group(&client, &tags, &security_group_name), + Aws::create_security_group( + &client, + &tags, + &security_group_name, + &vpc_id, + security_group_id + ), Aws::create_placement_group(&client, &tags, &placement_group_name), - Aws::get_default_subnet_id(&client) + Aws::get_default_subnet_id(&client, subnet_id) ); let key = PrivateKey::random(&mut OsRng {}, ssh_key::Algorithm::Ed25519).unwrap(); @@ -87,7 +110,7 @@ impl Aws { host_private_key, security_group_id, placement_group_name, - default_subnet_id, + subnet_id: default_subnet_id, tags, } } @@ -111,26 +134,35 @@ impl Aws { client: &aws_sdk_ec2::Client, tags: &Tags, name: &str, + vpc_id: &Option, + security_group_id: Option, ) -> String { - let security_group_id = client - .create_security_group() - .group_name(name) - .description("aws-throwaway security group") - .tag_specifications(tags.create_tags(ResourceType::SecurityGroup, "aws-throwaway")) - .send() - .await - .map_err(|e| e.into_service_error()) - .unwrap() - .group_id - .unwrap(); - tracing::info!("created security group"); - - tokio::join!( - Aws::create_ingress_rule_internal(client, tags, name), - Aws::create_ingress_rule_ssh(client, tags, name), - ); - - security_group_id + match security_group_id { + Some(id) => id, + None => { + let security_group_id = client + .create_security_group() + .group_name(name) + .set_vpc_id(vpc_id.clone()) + .description("aws-throwaway security group") + .tag_specifications( + tags.create_tags(ResourceType::SecurityGroup, "aws-throwaway"), + ) + .send() + .await + .map_err(|e| e.into_service_error()) + .unwrap() + .group_id + .unwrap(); + tracing::info!("created security group"); + + tokio::join!( + Aws::create_ingress_rule_internal(client, tags, name), + Aws::create_ingress_rule_ssh(client, tags, name), + ); + security_group_id + } + } } async fn create_ingress_rule_internal( @@ -187,31 +219,37 @@ impl Aws { tracing::info!("created placement group"); } - async fn get_default_subnet_id(client: &aws_sdk_ec2::Client) -> String { - client - .describe_subnets() - .filters( - Filter::builder() - .name("default-for-az") - .values("true") - .build(), - ) - .filters( - Filter::builder() - .name("availability-zone") - .values(AZ) - .build(), - ) - .send() - .await - .map_err(|e| e.into_service_error()) - .unwrap() - .subnets - .unwrap() - .pop() - .unwrap() - .subnet_id - .unwrap() + async fn get_default_subnet_id( + client: &aws_sdk_ec2::Client, + subnet_id: Option, + ) -> String { + match subnet_id { + Some(subnet_id) => subnet_id, + None => client + .describe_subnets() + .filters( + Filter::builder() + .name("default-for-az") + .values("true") + .build(), + ) + .filters( + Filter::builder() + .name("availability-zone") + .values(AZ) + .build(), + ) + .send() + .await + .map_err(|e| e.into_service_error()) + .unwrap() + .subnets + .unwrap() + .pop() + .unwrap() + .subnet_id + .unwrap(), + } } /// Call before dropping [`Aws`] @@ -315,7 +353,7 @@ impl Aws { err.into_service_error().meta().message() ) } else { - tracing::info!("security group {id:?} was succesfully deleted",) + tracing::info!("security group {id:?} was succesfully deleted") } } } @@ -353,7 +391,7 @@ impl Aws { async fn delete_keypairs(client: &aws_sdk_ec2::Client, tags: &Tags) { for id in Self::get_all_throwaway_tags(client, tags, "key-pair").await { - client + if let Err(err) = client .delete_key_pair() .key_pair_id(&id) .send() @@ -362,8 +400,11 @@ impl Aws { anyhow::anyhow!(e.into_service_error()) .context(format!("Failed to delete keypair {id:?}")) }) - .unwrap(); - tracing::info!("keypair {id:?} was succesfully deleted"); + { + tracing::error!("keypair {id:?} could not be deleted: {err}"); + } else { + tracing::info!("keypair {id:?} was succesfully deleted"); + } } } @@ -398,11 +439,11 @@ impl Aws { InstanceOs::Ubuntu20_04 => "20.04", InstanceOs::Ubuntu22_04 => "22.04", }; - let image_id = format!( + let image_id = definition.ami.unwrap_or_else(|| format!( "resolve:ssm:/aws/service/canonical/ubuntu/server/{}/stable/current/{}/hvm/ebs-gp2/ami-id", ubuntu_version, cpu_arch::get_arch_of_instance_type(definition.instance_type.clone()).get_ubuntu_arch_identifier() - ); + )); let result = self .client .run_instances() @@ -413,6 +454,11 @@ impl Aws { .availability_zone(AZ) .build(), )) + .set_subnet_id(if elastic_ip.is_some() { + None + } else { + Some(self.subnet_id.to_owned()) + }) .min_count(1) .max_count(1) .block_device_mappings( @@ -437,7 +483,7 @@ impl Aws { .device_index(i as i32) .groups(&self.security_group_id) .associate_public_ip_address(false) - .subnet_id(&self.default_subnet_id) + .subnet_id(&self.subnet_id) .description(i.to_string()) .build() })