add aws builder with use_public_addresses (#24)

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(

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()

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

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<IpAddr>,
     private_ip: IpAddr,
     client_private_key: String,
     host_public_key_bytes: Vec<u8>,
@@ -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<IpAddr> {
         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<IpAddr>,
         private_ip: IpAddr,
         host_public_key_bytes: Vec<u8>,
         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,

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<u8>,
-    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<u8>,
+    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,23 +467,24 @@ 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| {
                             InstanceNetworkInterfaceSpecification::builder()
                                 .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())
                                 .build()
                         })
                         .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(),

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;
 ```