1 files changed, 83 insertions, 0 deletions
diff --git a/src/ec2.rs b/src/ec2.rs
new file mode 100644
index 0000000..28169e7
--- /dev/null
+++ b/src/ec2.rs
@@ -0,0 +1,83 @@
+use std::collections::HashSet;
+
+use aws_sdk_ec2 as ec2;
+use aws_sdk_ec2::types::Filter;
+use aws_sdk_route53::types::{ResourceRecordSet, RrType};
+use tokio_stream::StreamExt;
+use trust_dns_proto::rr::Name;
+
+use crate::dns::absolute;
+use crate::hashable::Hashable;
+use crate::result::Result;
+use crate::route53::recordset;
+
+pub trait Ec2 {
+    fn ec2(&self) -> &ec2::Client;
+}
+
+pub async fn instance_recordsets<C>(
+    asg_name: &str,
+    dns_suffix: &Name,
+    dns_ttl: i64,
+    live_instance_ids: &[String],
+    aws_context: &C,
+) -> Result<HashSet<Hashable<ResourceRecordSet>>>
+where
+    C: Ec2,
+{
+    // If there's nothing running, then (a) we don't need to ask AWS about
+    // running instances, and (b) we can't anyways as the API call requires at
+    // least one instance ID. Abort here.
+    if live_instance_ids.is_empty() {
+        return Ok(HashSet::new());
+    }
+
+    let asg_filter = Filter::builder()
+        .name("tag:aws:autoscaling:groupName")
+        .values(asg_name)
+        .build();
+
+    let mut apex_ip4 = HashSet::new();
+    let mut apex_ip6 = HashSet::new();
+
+    let mut instances_paginator = aws_context
+        .ec2()
+        .describe_instances()
+        .set_instance_ids(Some(live_instance_ids.to_owned()))
+        .filters(asg_filter)
+        .into_paginator()
+        .items()
+        .send();
+
+    while let Some(reservation) = instances_paginator.try_next().await? {
+        let instances = reservation.instances().unwrap_or(&[]);
+        for instance in instances {
+            // Mild abuse of the fact that optional values are also iterable
+            apex_ip4.extend(instance.public_ip_address().map(String::from));
+
+            let instance_interfaces = instance.network_interfaces().unwrap_or(&[]);
+            let instance_ip6: Vec<_> = instance_interfaces
+                .iter()
+                .flat_map(|interface| interface.ipv6_addresses().unwrap_or(&[]))
+                // Flatmap here to drop the None values, unwrap the Some values
+                .flat_map(|ipv6| ipv6.ipv6_address())
+                .map(String::from)
+                .collect();
+
+            apex_ip6.extend(instance_ip6.iter().map(ToOwned::to_owned).map(String::from));
+        }
+    }
+
+    let apex_hostname = absolute(dns_suffix.clone())?;
+    let apex_hostname = apex_hostname.to_ascii();
+
+    let mut asg_recordsets = HashSet::new();
+    if !apex_ip4.is_empty() {
+        asg_recordsets.insert(recordset(&apex_hostname, dns_ttl, RrType::A, apex_ip4).into());
+    }
+    if !apex_ip6.is_empty() {
+        asg_recordsets.insert(recordset(&apex_hostname, dns_ttl, RrType::Aaaa, apex_ip6).into());
+    }
+
+    Ok(asg_recordsets)
+}