dynk8s-provisioner/src/k8s.rs

//! Kubernetes Integration
use std::collections::btree_map::BTreeMap;

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::{ListParams, Patch, PatchParams, PostParams};
use kube::{Api, Client};
use log::{debug, error, info};
use rand::seq::SliceRandom;

/// The set of characters allowed to appear in bootstrap tokens
const TOKEN_CHARS: [char; 36] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
    'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3',
    '4', '5', '6', '7', '8', '9',
];

/// Kubernetes Bootstrap Token
///
/// Bootstrap tokens are typically used to add new nodes to the cluster using
/// `kubeadm join`.  They are bearer tokens consisting of two parts: a token ID
/// and a secret.  For additional information, see [Authenticating with
/// Bootstrap Tokens][0].
///
/// The Dynk8s Provisioner allocates bootstrap tokens for ephemeral nodes.
/// Each token is assigned to a specific EC2 instance; the instance ID is
/// stored in the token Secret's metadata, using the
/// `dynk8s.du5t1n.me/ec2-instance-id` label.
///
/// [0]: https://kubernetes.io/docs/reference/access-authn-authz/bootstrap-tokens/
#[derive(Clone, Debug)]
struct BootstrapToken {
    /// The token ID (generated)
    token_id: String,
    /// The token secret (generated)
    secret: String,
    /// The date and time the token expires
    expiration: DateTime<Utc>,
    /// The EC2 instance ID to which the token is assigned
    instance_id: Option<String>,
}

#[allow(dead_code)]
impl BootstrapToken {
    /// Generate a new token
    ///
    /// Initially, the token is *not* assigned to an EC2 instance.  Use
    /// [`Self::instance_id`] to set the associated instance ID.
    pub fn new() -> Self {
        let mut rng = rand::thread_rng();
        let mut token_id = String::with_capacity(6);
        while token_id.len() < 6 {
            token_id.push(TOKEN_CHARS.choose(&mut rng).unwrap().clone());
        }
        let mut secret = String::with_capacity(16);
        while secret.len() < 16 {
            secret.push(TOKEN_CHARS.choose(&mut rng).unwrap().clone());
        }
        let expiration = Utc::now() + Duration::hours(1);
        Self {
            token_id,
            secret,
            expiration,
            instance_id: None,
        }
    }

    /// Return the expiration date/time
    pub fn expiration(&self) -> &DateTime<Utc> {
        &self.expiration
    }

    /// Set the ID of the EC2 instance associated with the token
    pub fn instance_id(mut self, instance_id: String) -> Self {
        self.instance_id = Some(instance_id);
        self
    }

    /// Return the secret part of the token
    pub fn secret(&self) -> &str {
        &self.secret
    }

    /// Set the expiration date/time from a time-to-live duration
    pub fn set_ttl(mut self, ttl: Duration) -> Self {
        self.expiration = Utc::now() + ttl;
        self
    }

    /// Set the expiration date/time
    pub fn set_expiration(mut self, expiration: DateTime<Utc>) -> Self {
        self.expiration = expiration;
        self
    }

    /// Return the ID part of the token
    pub fn token_id(&self) -> &str {
        &self.token_id
    }

    /// Return the token as a string
    pub fn token(&self) -> String {
        format!("{}.{}", self.token_id, self.secret)
    }
}

impl Into<Secret> for BootstrapToken {
    /// Create a [`Secret`] for the token
    ///
    /// Converting a [`BootstrapToken`] into a [`Secret`] populates the fields
    /// necessary to store the token in Kubernetes.  The `Secret` can be passed
    /// to e.g. [`kube::Api::create`].
    fn into(self) -> Secret {
        let mut data = BTreeMap::<String, String>::new();
        data.insert("token-id".into(), self.token_id.clone());
        data.insert("token-secret".into(), self.secret);
        data.insert("expiration".into(), self.expiration.to_rfc3339());
        data.insert("usage-bootstrap-authentication".into(), "true".into());
        data.insert("usage-bootstrap-signing".into(), "true".into());
        data.insert(
            "auth-extra-groups".into(),
            "system:bootstrappers:kubeadm:default-node-token".into(),
        );
        let mut labels = BTreeMap::<String, String>::new();
        if let Some(instance_id) = self.instance_id {
            labels.insert(
                "dynk8s.du5t1n.me/ec2-instance-id".into(),
                instance_id.into(),
            );
        }
        Secret {
            data: None,
            immutable: None,
            metadata: ObjectMeta {
                annotations: None,
                cluster_name: None,
                creation_timestamp: None,
                deletion_grace_period_seconds: None,
                deletion_timestamp: None,
                finalizers: None,
                generate_name: None,
                generation: None,
                labels: Some(labels),
                managed_fields: None,
                name: Some(format!("bootstrap-token-{}", self.token_id)),
                namespace: None,
                owner_references: None,
                resource_version: None,
                self_link: None,
                uid: None,
            },
            string_data: Some(data),
            type_: Some("bootstrap.kubernetes.io/token".into()),
        }
    }
}

/// 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
/// Secret resource.  The token is assigned to the given EC2 instance, and will
/// be provided to that instance when it is ready to join the cluster.
pub async fn create_bootstrap_token<I: AsRef<str>>(
    instance_id: I,
) -> Result<(), kube::Error> {
    let instance_id = instance_id.as_ref();
    info!("Creating bootstrap token for instance {}", instance_id);
    let token = BootstrapToken::new().instance_id(instance_id.into());
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::namespaced(client, "kube-system");
    let pp: PostParams = Default::default();
    let secret = secrets.create(&pp, &token.into()).await?;
    info!("Successfully created secret {:?}", &secret.metadata.name);
    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;
    use regex::Regex;

    #[test]
    fn test_bootstrap_token_new() {
        let token = BootstrapToken::new();
        let id_re = Regex::new(r"^[a-z0-9]{6}$").unwrap();
        let secret_re = Regex::new(r"^[a-z0-9]{16}$").unwrap();
        let token_re = Regex::new(r"[a-z0-9]{6}\.[a-z0-9]{16}$").unwrap();
        assert!(id_re.is_match(&token.token_id()));
        assert!(secret_re.is_match(&token.secret()));
        assert!(token_re.is_match(&token.token()));
    }

    #[test]
    fn test_bootstrap_token_into_secret() {
        let token = BootstrapToken::new();
        let secret: Secret = token.clone().into();
        let data = secret.string_data.unwrap();
        assert_eq!(
            &secret.metadata.name,
            &Some(format!("bootstrap-token-{}", token.token_id))
        );
        assert_eq!(data.get("token-id").unwrap(), &token.token_id);
        assert_eq!(data.get("token-secret").unwrap(), &token.secret);
        assert_eq!(
            data.get("expiration").unwrap(),
            &token.expiration.to_rfc3339()
        );
        assert!(secret
            .metadata
            .labels
            .unwrap()
            .get("dynk8s.du5t1n.me/ec2-instance-id")
            .is_none());
    }

    #[test]
    fn test_bootstrap_token_into_secret_instance_id() {
        let token =
            BootstrapToken::new().instance_id("i-0a1b2c3d4e5f6f7f8".into());
        let secret: Secret = token.clone().into();
        let data = secret.string_data.unwrap();
        assert_eq!(
            &secret.metadata.name,
            &Some(format!("bootstrap-token-{}", token.token_id))
        );
        assert_eq!(data.get("token-id").unwrap(), &token.token_id);
        assert_eq!(data.get("token-secret").unwrap(), &token.secret);
        assert_eq!(
            data.get("expiration").unwrap(),
            &token.expiration.to_rfc3339()
        );
        assert_eq!(
            secret
                .metadata
                .labels
                .unwrap()
                .get("dynk8s.du5t1n.me/ec2-instance-id")
                .unwrap(),
            "i-0a1b2c3d4e5f6f7f8"
        );
    }
}