Assign WireGuard keys to EC2 instances

In order to join the on-premises Kubernetes cluster, EC2 instances will
need to first connect to the WireGuard VPN.  The *dynk8s* provisioner
will provide keys to instances to configure their WireGuard clients.

WireGuard keys must be pre-configured on the server and stored in
Kubernetes as *dynk8s.du5t1n.me/wireguard-key* Secret resources.  They
must also have a `dynk8s.du5t1n.me/ec2-instance-id` label.  If this
label is empty, the key is available to be assigned to an instance.

When an EventBridge event is received indicating an instance is now
running, a WireGuard key is assigned to that instance (by setting the
`dynk8s.du5t1n.me/ec2-instance-id` label).  Conversely, when an event is
received indicating that the instance is terminated, any WireGuard keys
assigned to that instance are freed.

src/events.rs

@@ -5,7 +5,9 @@
 //! Service.
 use log::{debug, error};
 
-use crate::k8s::create_bootstrap_token;
+use crate::k8s::{
+    assign_wireguard_key, create_bootstrap_token, unassign_wireguard_key,
+};
 use crate::model::events::*;
 
 /// Handle an EC2 instance state change event
@@ -14,16 +16,34 @@ use crate::model::events::*;
 /// associated with ephemeral nodes running as EC2 instances.
 ///
 /// When an instance starts:
-/// 1. A Kubernetes bootstrap token is generated, to be used by `kubeadm` to
+/// 1. A WireGuard key is assigned to the instance
+/// 2. A Kubernetes bootstrap token is generated, to be used by `kubeadm` to
 ///    add the node to the cluster.
+///
+/// When an instance is terminated:
+/// 1. Any WireGuard keys assigned to the instance are unassigned
 pub async fn on_ec2_instance_state_change(evt: Ec2InstanceStateChange) {
     debug!("EC2 instance {} is now {}", &evt.instance_id, &evt.state);
     if evt.state == "running" {
+        if let Err(e) = assign_wireguard_key(&evt.instance_id).await {
+            error!(
+                "Failed to assign WireGuard key to instnce {}: {}",
+                &evt.instance_id, e
+            );
+            return;
+        }
         if let Err(e) = create_bootstrap_token(&evt.instance_id).await {
             error!(
                 "Failed to create bootstrap token for instance {}: {}",
                 &evt.instance_id, e
             )
         };
+    } else if evt.state == "terminated" {
+        if let Err(e) = unassign_wireguard_key(&evt.instance_id).await {
+            error!(
+                "Failed to unassign WireGuard key from instance: {}: {}",
+                &evt.instance_id, e
+            );
+        }
     }
 }

src/k8s.rs

@@ -5,9 +5,9 @@ use chrono::offset::Utc;
 use chrono::{DateTime, Duration};
 use k8s_openapi::api::core::v1::Secret;
 use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
-use kube::core::params::PostParams;
+use kube::core::params::{ListParams, Patch, PatchParams, PostParams};
 use kube::{Api, Client};
-use log::info;
+use log::{debug, error, info};
 use rand::seq::SliceRandom;
 
 /// The set of characters allowed to appear in bootstrap tokens
@@ -157,6 +157,95 @@ impl Into<Secret> for BootstrapToken {
     }
 }
 
+/// Assign an existing WireGuard key to the specified EC2 instance
+///
+/// This function finds the first unused WireGuard key, stored as a Kubernetes
+/// Secret resource, and assigns it to the specified EC2 instance.  Keys are
+/// assigned by setting the `dynk8s.du5t1n.me/ec2-instance-id` label in the
+/// Secret resource's metadata.
+///
+/// Secret resources for WireGuard keys have a *type* of
+/// `dynk8s.du5t1n.me/wireguard-key`.  They must be created ahead of time and
+/// must refer to working keys already configured on the WireGuard server.
+pub async fn assign_wireguard_key(
+    instance_id: &str,
+) -> Result<(), kube::Error> {
+    let client = Client::try_default().await?;
+    let secrets: Api<Secret> = Api::default_namespaced(client);
+    debug!(
+        "Checking for WireGuard keys already assigned to instance {}",
+        instance_id
+    );
+    let lp = ListParams::default()
+        .fields("type=dynk8s.du5t1n.me/wireguard-key")
+        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", instance_id));
+    let res = secrets.list(&lp).await?;
+    if !res.items.is_empty() {
+        info!("WireGuard key already assigned to instance {}", instance_id);
+        return Ok(());
+    }
+    debug!("Looking for available WireGuard keys");
+    let lp = ListParams::default()
+        .fields("type=dynk8s.du5t1n.me/wireguard-key")
+        .labels("dynk8s.du5t1n.me/ec2-instance-id=");
+    let res = secrets.list(&lp).await?;
+    if res.items.is_empty() {
+        error!("No WireGuard keys available for instance {}", &instance_id);
+    } else {
+        if let Some(name) = &res.items[0].metadata.name {
+            let mut labels = BTreeMap::<String, String>::new();
+            labels.insert(
+                "dynk8s.du5t1n.me/ec2-instance-id".into(),
+                instance_id.into(),
+            );
+            let mut secret = Secret::default();
+            secret.metadata.labels = Some(labels);
+            let patch = Patch::Apply(&secret);
+            let pp = PatchParams::apply(env!("CARGO_PKG_NAME")).force();
+            secrets.patch(&name, &pp, &patch).await?;
+            info!(
+                "Assigned WireGuard key {} to instance {}",
+                name, &instance_id
+            );
+        }
+    }
+    Ok(())
+}
+
+/// Unassign all WireGuard keys from the specified EC2 instance
+///
+/// This function finds all WireGuard keys, stored as Kubernetes Secret
+/// resources, associated with the specified EC2 instance and unassigns them.
+/// Unassigned keys have the `dynk8s.du5t1n.me/ec2-instance-id` label set to
+/// the empty string.
+pub async fn unassign_wireguard_key(
+    instance_id: &str,
+) -> Result<(), kube::Error> {
+    let client = Client::try_default().await?;
+    let secrets: Api<Secret> = Api::default_namespaced(client);
+    let lp = ListParams::default()
+        .fields("type=dynk8s.du5t1n.me/wireguard-key")
+        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", instance_id));
+    info!("Unassigning WireGuard keys from instance {}", instance_id);
+    for s in secrets.list(&lp).await? {
+        if let Some(name) = &s.metadata.name {
+            let mut labels = BTreeMap::<String, String>::new();
+            labels
+                .insert("dynk8s.du5t1n.me/ec2-instance-id".into(), "".into());
+            let mut secret = Secret::default();
+            secret.metadata.labels = Some(labels);
+            let patch = Patch::Apply(&secret);
+            let pp = PatchParams::apply(env!("CARGO_PKG_NAME")).force();
+            secrets.patch(&name, &pp, &patch).await?;
+            info!(
+                "Unassigned WireGuard key {} from instance {}",
+                name, &instance_id
+            );
+        }
+    }
+    Ok(())
+}
+
 /// Generate and store a bootstrap token for the specified EC2 instance
 ///
 /// This function generates a new bootstrap token and stores it as a Kubernetes