From 75a80f44b4e8974563acf78f136007892aa64344 Mon Sep 17 00:00:00 2001 From: Lucas Kent Date: Mon, 23 Oct 2023 21:35:48 +1100 Subject: [PATCH] add aws builder with use_public_addresses (#24) --- .../aws-throwaway-test-multiple-instances.rs | 2 +- aws-throwaway/examples/aws-throwaway-test.rs | 10 +- aws-throwaway/examples/create-instance.rs | 4 +- aws-throwaway/src/ec2_instance.rs | 27 ++-- aws-throwaway/src/lib.rs | 118 +++++++++++++----- readme.md | 4 +- 6 files changed, 118 insertions(+), 47 deletions(-) diff --git a/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs b/aws-throwaway/examples/aws-throwaway-test-multiple-instances.rs index b50a39f..7947081 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::builder(CleanupResources::AllResources).build().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..bf18dc1 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::builder(CleanupResources::AllResources).build().await; let instance = aws .create_ec2_instance(Ec2InstanceDefinition::new(InstanceType::T2Micro)) .await; @@ -22,11 +22,17 @@ async fn main() { let result = instance.ssh().shell("xxd some_remote_file").await; println!("The bytes of the remote file:\n{}", result.stdout); + + // download a file and assert on its contents instance .ssh() .pull_file(Path::new("some_remote_file"), Path::new("some_local_file")) .await; - println!("Remote file copied locally to some_local_file"); + assert_eq!( + std::fs::read_to_string("some_local_file").unwrap(), + "some string\n" + ); + std::fs::remove_file("some_local_file").unwrap(); instance .ssh() diff --git a/aws-throwaway/examples/create-instance.rs b/aws-throwaway/examples/create-instance.rs index 5ee049a..d13e9b1 100644 --- a/aws-throwaway/examples/create-instance.rs +++ b/aws-throwaway/examples/create-instance.rs @@ -21,7 +21,9 @@ 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::builder(CleanupResources::WithAppTag(AWS_THROWAWAY_TAG.to_owned())) + .build() + .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.rs b/aws-throwaway/src/ec2_instance.rs index a5f6b73..2058dd7 100644 --- a/aws-throwaway/src/ec2_instance.rs +++ b/aws-throwaway/src/ec2_instance.rs @@ -5,7 +5,8 @@ use tokio::{net::TcpStream, time::Instant}; /// Represents a currently running EC2 instance and provides various methods for interacting with the instance. pub struct Ec2Instance { - public_ip: IpAddr, + connect_ip: IpAddr, + public_ip: Option, private_ip: IpAddr, client_private_key: String, host_public_key_bytes: Vec, @@ -21,7 +22,7 @@ pub struct NetworkInterface { impl Ec2Instance { /// Use this address to connect to this instance from outside of AWS - pub fn public_ip(&self) -> IpAddr { + pub fn public_ip(&self) -> Option { self.public_ip } @@ -48,7 +49,7 @@ impl Ec2Instance { /// Insert this into your known_hosts file to avoid errors due to unknown fingerprints pub fn openssh_known_hosts_line(&self) -> String { - format!("{} {}", &self.public_ip, &self.host_public_key) + format!("{} {}", &self.connect_ip, &self.host_public_key) } /// Returns an object that allows commands to be sent over ssh @@ -68,12 +69,14 @@ TERM=xterm ssh -i key ubuntu@{} -o "UserKnownHostsFile known_hosts" ```"#, self.client_private_key(), self.openssh_known_hosts_line(), - self.public_ip() + self.connect_ip ) } + /// It is gauranteed that public_ip will be Some if use_public_address is true pub(crate) async fn new( - public_ip: IpAddr, + connect_ip: IpAddr, + public_ip: Option, private_ip: IpAddr, host_public_key_bytes: Vec, host_public_key: String, @@ -85,24 +88,27 @@ TERM=xterm ssh -i key ubuntu@{} -o "UserKnownHostsFile known_hosts" // We retry many times before we are able to succesfully make an ssh connection. // Each error is expected and so is logged as a `info!` that describes the underlying startup process that is supposed to cause the error. // A numbered comment is left before each `info!` to demonstrate the order each error occurs in. - match tokio::time::timeout(Duration::from_secs(10), TcpStream::connect((public_ip, 22))) - .await + match tokio::time::timeout( + Duration::from_secs(10), + TcpStream::connect((connect_ip, 22)), + ) + .await { Err(_) => { // 1. - tracing::info!("Timed out connecting to {public_ip} over ssh, the host is probably not accessible yet, retrying"); + tracing::info!("Timed out connecting to {connect_ip} over ssh, the host is probably not accessible yet, retrying"); continue; } Ok(Err(e)) => { // 2. - tracing::info!("failed to connect to {public_ip}:22, the host probably hasnt started their ssh service yet, retrying, error was {e}"); + tracing::info!("failed to connect to {connect_ip}:22, the host probably hasnt started their ssh service yet, retrying, error was {e}"); tokio::time::sleep_until(start + Duration::from_secs(1)).await; continue; } Ok(Ok(stream)) => { match SshConnection::new( stream, - public_ip, + connect_ip, host_public_key_bytes.clone(), client_private_key, ) @@ -117,6 +123,7 @@ TERM=xterm ssh -i key ubuntu@{} -o "UserKnownHostsFile known_hosts" // 4. Then finally we have a working ssh connection. Ok(ssh) => { break Ec2Instance { + connect_ip, ssh, public_ip, private_ip, diff --git a/aws-throwaway/src/lib.rs b/aws-throwaway/src/lib.rs index d1a63bd..bb89d26 100644 --- a/aws-throwaway/src/lib.rs +++ b/aws-throwaway/src/lib.rs @@ -31,26 +31,34 @@ async fn config() -> SdkConfig { aws_config::from_env().region(region_provider).load().await } -/// Construct this type to create and cleanup aws resources. -pub struct Aws { - client: aws_sdk_ec2::Client, - keyname: String, - client_private_key: String, - host_public_key: String, - host_public_key_bytes: Vec, - host_private_key: String, - security_group_id: String, - placement_group_name: String, - default_subnet_id: String, - tags: Tags, +pub struct AwsBuilder { + cleanup: CleanupResources, + use_public_addresses: bool, } -impl Aws { - /// Construct a new [`Aws`] +/// The default configuration will succeed for an AMI user with sufficient access and unmodified default vpcs/subnets +/// Consider altering the configuration if: +/// * you want to reduce the amount of access required by the user +/// * you want to connect directly from within the VPC +/// * you have already created a specific VPC, subnet or security group that you want aws-throwaway to make use of. +// TODO: document minimum required access for default configuration. +impl AwsBuilder { + /// When set to: + /// * true - aws-throwaway will connect to the public ip of the instances that it creates. + /// + The subnet must have the property MapPublicIpOnLaunch set to true (the unmodified default subnet meets this requirement) + /// + Elastic IPs will be created for instances with multiple network interfaces because AWS does not assign a public IP in that scenario + /// * false - aws-throwaway will connect to the private ip of the instances that it creates. + /// + aws-throwaway must be running on a machine within the VPC used by aws-throwaaway or a VPN must be used to connect to the VPC or another similar setup. /// - /// 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 { + /// If the subnet used has MapPublicIpOnLaunch=true then all instances will be publically accessible regardless of this use_public_addresses field. + /// + /// The default is `true`. + pub fn use_public_addresses(mut self, use_public_addresses: bool) -> Self { + self.use_public_addresses = use_public_addresses; + self + } + + pub async fn build(self) -> Aws { let config = config().await; let user_name = iam::user_name(&config).await; let keyname = format!("aws-throwaway-{user_name}-{}", Uuid::new_v4()); @@ -60,11 +68,11 @@ impl Aws { let tags = Tags { user_name: user_name.clone(), - cleanup, + cleanup: self.cleanup, }; // Cleanup any resources that were previously failed to cleanup - Self::cleanup_resources_inner(&client, &tags).await; + Aws::cleanup_resources_inner(&client, &tags).await; let (client_private_key, security_group_id, _, default_subnet_id) = tokio::join!( Aws::create_key_pair(&client, &tags, &keyname), @@ -78,7 +86,10 @@ impl Aws { let host_public_key = key.public_key().to_openssh().unwrap(); let host_private_key = key.to_openssh(ssh_key::LineEnding::LF).unwrap().to_string(); + let use_public_addresses = self.use_public_addresses; + Aws { + use_public_addresses, client, keyname, client_private_key, @@ -91,6 +102,34 @@ impl Aws { tags, } } +} + +/// Construct this type to create and cleanup aws resources. +pub struct Aws { + client: aws_sdk_ec2::Client, + keyname: String, + client_private_key: String, + host_public_key: String, + host_public_key_bytes: Vec, + host_private_key: String, + security_group_id: String, + placement_group_name: String, + default_subnet_id: String, + use_public_addresses: bool, + tags: Tags, +} + +impl Aws { + /// Returns an [`AwsBuilder`] that will build a new [`Aws`]. + /// + /// Before building 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 fn builder(cleanup: CleanupResources) -> AwsBuilder { + AwsBuilder { + cleanup, + use_public_addresses: true, + } + } async fn create_key_pair(client: &aws_sdk_ec2::Client, tags: &Tags, name: &str) -> String { let keypair = client @@ -370,7 +409,7 @@ impl Aws { /// Creates a new EC2 instance as defined by [`Ec2InstanceDefinition`] pub async fn create_ec2_instance(&self, definition: Ec2InstanceDefinition) -> Ec2Instance { // elastic IP's are a limited resource so only create it if we truly need it. - let elastic_ip = if definition.network_interface_count > 1 { + let elastic_ip = if self.use_public_addresses && definition.network_interface_count > 1 { Some( self.client .allocate_address() @@ -388,10 +427,10 @@ impl Aws { }; // if we specify a list of network interfaces we cannot specify an instance level security group - let security_group_ids = if elastic_ip.is_some() { - None - } else { + let security_group_ids = if definition.network_interface_count == 1 { Some(vec![self.security_group_id.clone()]) + } else { + None }; let ubuntu_version = match definition.os { @@ -428,7 +467,9 @@ impl Aws { .build(), ) .set_security_group_ids(security_group_ids) - .set_network_interfaces(if elastic_ip.is_some() { + .set_network_interfaces(if definition.network_interface_count == 1 { + None + } else { Some( (0..definition.network_interface_count) .map(|i| { @@ -436,6 +477,7 @@ impl Aws { .delete_on_termination(true) .device_index(i as i32) .groups(&self.security_group_id) + // must be false when launching with multiple network interfaces .associate_public_ip_address(false) .subnet_id(&self.default_subnet_id) .description(i.to_string()) @@ -443,8 +485,6 @@ impl Aws { }) .collect(), ) - } else { - None }) .key_name(&self.keyname) .user_data(base64::engine::general_purpose::STANDARD.encode(format!( @@ -510,7 +550,7 @@ sudo systemctl start ssh // `The pending-instance-running instance to which 'eni-***' is attached is not in a valid state for this operation` if start.elapsed() > Duration::from_secs(120) { panic!( - "Received error while assosciating address after 120s retrying: {}", + "Received error while associating address after 120s retrying: {}", err.into_service_error() ); } else { @@ -524,8 +564,17 @@ sudo systemctl start ssh let mut public_ip = elastic_ip.map(|x| x.public_ip.unwrap().parse().unwrap()); let mut private_ip = None; - while public_ip.is_none() || private_ip.is_none() { - tokio::time::sleep(std::time::Duration::from_secs(1)).await; + // TODO: when we support custom subnets we will need the user to set this or query AWS for it. + let subnet_assigns_public_ips = true; + + let public_ip_expected = self.use_public_addresses || subnet_assigns_public_ips; + + if public_ip_expected { + tracing::info!("Waiting for instance private ip and public ip to be assigned"); + } else { + tracing::info!("Waiting for instance private ip to be assigned"); + } + while (public_ip_expected && public_ip.is_none()) || private_ip.is_none() { for reservation in self .client .describe_instances() @@ -544,12 +593,19 @@ sudo systemctl start ssh private_ip = instance.private_ip_address().map(|x| x.parse().unwrap()); } } + tokio::time::sleep(std::time::Duration::from_secs(1)).await; } - let public_ip = public_ip.unwrap(); + let private_ip = private_ip.unwrap(); - tracing::info!("created EC2 instance at: {public_ip}"); + let connect_ip = if self.use_public_addresses { + public_ip.unwrap() + } else { + private_ip + }; + tracing::info!("created EC2 instance at public:{public_ip:?} private:{private_ip}"); Ec2Instance::new( + connect_ip, public_ip, private_ip, self.host_public_key_bytes.clone(), diff --git a/readme.md b/readme.md index b60c454..38d06fd 100644 --- a/readme.md +++ b/readme.md @@ -12,7 +12,7 @@ It was developed for the use case of benchmarking. aws-throwaway makes it trivial to spin up an instance, interact with it, and then destroy it. ```rust -let aws = Aws::new(CleanupResources::AllResources).await; +let aws = Aws::builder(CleanupResources::AllResources).build().await; let instance = aws.create_ec2_instance(Ec2InstanceDefinition::new(InstanceType::T2Micro)).await; let output = instance.ssh().shell("echo 'Hello world!'").await; @@ -57,7 +57,7 @@ Rather than attempting to individually track each resource created like terrafor Consider this snippet from the example earlier: ```rust -let aws = Aws::new(CleanupResources::AllResources).await; +let aws = Aws::builder(CleanupResources::AllResources).build().await; let instance = aws.create_ec2_instance(Ec2InstanceDefinition::new(InstanceType::T2Micro)).await; ```