lodestone-proto 0.0.16

//! Pod-bound workload identity and provider-observed topology codecs.

use std::error::Error;
use std::fmt::{self, Display, Formatter};

use lodestone_types::{AwsBucketKind, AwsPartition, KubernetesObjectUidDigest, NodeRole};
use serde::{Deserialize, Serialize};

pub const KUBERNETES_IDENTITY_FORMAT_VERSION: u16 = 1;
pub const NODE_TOPOLOGY_FORMAT_VERSION: u16 = 1;
pub const MAX_KUBERNETES_IDENTITY_BYTES: usize = 4 * 1024;
pub const MAX_NODE_TOPOLOGY_BYTES: usize = 256;
const MAX_IDENTITY_STRING_BYTES: usize = 256;
const MAX_KUBERNETES_DNS_LABEL_BYTES: usize = 63;
const MAX_KUBERNETES_DNS_SUBDOMAIN_BYTES: usize = 253;

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct PodIdentityV1 {
    pub format_version: u16,
    pub cluster_uid_digest: KubernetesObjectUidDigest,
    pub namespace: String,
    pub service_account: String,
    pub pod_uid_digest: KubernetesObjectUidDigest,
    pub requested_role: NodeRole,
    pub public_key_digest: [u8; 32],
}

impl PodIdentityV1 {
    pub fn validate(&self) -> Result<(), KubernetesIdentityError> {
        if self.format_version != KUBERNETES_IDENTITY_FORMAT_VERSION {
            return Err(KubernetesIdentityError::UnsupportedIdentityVersion(self.format_version));
        }
        if !valid_dns_label(&self.namespace)
            || !valid_dns_subdomain(&self.service_account)
            || self.public_key_digest == [0; 32]
        {
            return Err(KubernetesIdentityError::InvalidPodIdentity);
        }
        Ok(())
    }

    pub fn provider_subject_digest(&self) -> Result<[u8; 32], KubernetesIdentityError> {
        self.validate()?;
        let bytes = serde_json::to_vec(self).map_err(|error| KubernetesIdentityError::Codec(error.to_string()))?;
        let mut hasher = blake3::Hasher::new();
        hasher.update(b"lodestone/p7/kubernetes-provider-subject/v1");
        hasher.update(&bytes);
        Ok(*hasher.finalize().as_bytes())
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "provider", rename_all = "snake_case", deny_unknown_fields)]
pub enum ProviderComputeTopologyV1 {
    Aws {
        partition: AwsPartition,
        region: String,
        availability_zone_id: String,
    },
    StaticDevelopment {
        deployment_identity: String,
        zone: String,
    },
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct NodeTopologyV1 {
    pub format_version: u16,
    pub provider_observation: ProviderComputeTopologyV1,
    pub kubernetes_zone_name: String,
}

impl NodeTopologyV1 {
    pub fn validate(&self) -> Result<(), KubernetesIdentityError> {
        if self.format_version != NODE_TOPOLOGY_FORMAT_VERSION {
            return Err(KubernetesIdentityError::UnsupportedTopologyVersion(self.format_version));
        }
        if !valid_dns_label(&self.kubernetes_zone_name) {
            return Err(KubernetesIdentityError::InvalidTopology);
        }
        match &self.provider_observation {
            ProviderComputeTopologyV1::Aws {
                region,
                availability_zone_id,
                ..
            } => {
                if !valid_provider_value(region) || !valid_aws_availability_zone_id(availability_zone_id) {
                    return Err(KubernetesIdentityError::InvalidTopology);
                }
            }
            ProviderComputeTopologyV1::StaticDevelopment {
                deployment_identity,
                zone,
            } => {
                if !valid_provider_value(deployment_identity) || !valid_provider_value(zone) {
                    return Err(KubernetesIdentityError::InvalidTopology);
                }
            }
        }
        Ok(())
    }

    /// AWS storage admission uses only an AWS provider observation. The
    /// Kubernetes zone name is a scheduling hint and is intentionally absent.
    pub fn admits_bucket_kind(&self, bucket_kind: &AwsBucketKind) -> bool {
        match (&self.provider_observation, bucket_kind) {
            (ProviderComputeTopologyV1::Aws { .. }, AwsBucketKind::GeneralPurpose) => true,
            (
                ProviderComputeTopologyV1::Aws {
                    availability_zone_id: observed,
                    ..
                },
                AwsBucketKind::Directory {
                    availability_zone_id: expected,
                },
            ) => observed == expected,
            (ProviderComputeTopologyV1::StaticDevelopment { .. }, _) => false,
        }
    }
}

pub fn encode_pod_identity(identity: &PodIdentityV1) -> Result<Vec<u8>, KubernetesIdentityError> {
    identity.validate()?;
    encode_bounded(identity, MAX_KUBERNETES_IDENTITY_BYTES)
}

pub fn decode_pod_identity(bytes: &[u8]) -> Result<PodIdentityV1, KubernetesIdentityError> {
    let identity: PodIdentityV1 = decode_bounded(bytes, MAX_KUBERNETES_IDENTITY_BYTES)?;
    identity.validate()?;
    Ok(identity)
}

pub fn encode_node_topology(topology: &NodeTopologyV1) -> Result<Vec<u8>, KubernetesIdentityError> {
    topology.validate()?;
    encode_bounded(topology, MAX_NODE_TOPOLOGY_BYTES)
}

pub fn decode_node_topology(bytes: &[u8]) -> Result<NodeTopologyV1, KubernetesIdentityError> {
    let topology: NodeTopologyV1 = decode_bounded(bytes, MAX_NODE_TOPOLOGY_BYTES)?;
    topology.validate()?;
    Ok(topology)
}

fn encode_bounded<T: Serialize>(value: &T, limit: usize) -> Result<Vec<u8>, KubernetesIdentityError> {
    let bytes = serde_json::to_vec(value).map_err(|error| KubernetesIdentityError::Codec(error.to_string()))?;
    if bytes.len() > limit {
        return Err(KubernetesIdentityError::MessageTooLarge(bytes.len()));
    }
    Ok(bytes)
}

fn decode_bounded<T>(bytes: &[u8], limit: usize) -> Result<T, KubernetesIdentityError>
where
    T: for<'de> Deserialize<'de> + Serialize,
{
    if bytes.len() > limit {
        return Err(KubernetesIdentityError::MessageTooLarge(bytes.len()));
    }
    let value = serde_json::from_slice(bytes).map_err(|error| KubernetesIdentityError::Codec(error.to_string()))?;
    let canonical = serde_json::to_vec(&value).map_err(|error| KubernetesIdentityError::Codec(error.to_string()))?;
    if canonical != bytes {
        return Err(KubernetesIdentityError::NonCanonical);
    }
    Ok(value)
}

fn valid_dns_label(value: &str) -> bool {
    !value.is_empty()
        && value.len() <= MAX_KUBERNETES_DNS_LABEL_BYTES
        && value.as_bytes().first().is_some_and(u8::is_ascii_alphanumeric)
        && value.as_bytes().last().is_some_and(u8::is_ascii_alphanumeric)
        && value
            .bytes()
            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
}

fn valid_dns_subdomain(value: &str) -> bool {
    !value.is_empty() && value.len() <= MAX_KUBERNETES_DNS_SUBDOMAIN_BYTES && value.split('.').all(valid_dns_label)
}

fn valid_provider_value(value: &str) -> bool {
    !value.is_empty() && value.len() <= MAX_IDENTITY_STRING_BYTES && value.bytes().all(|byte| byte.is_ascii_graphic())
}

fn valid_aws_availability_zone_id(value: &str) -> bool {
    let Some((region_code, suffix)) = value.rsplit_once("-az") else {
        return false;
    };
    !region_code.is_empty()
        && region_code.len() <= 32
        && region_code
            .bytes()
            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit())
        && !suffix.is_empty()
        && suffix.bytes().all(|byte| byte.is_ascii_digit())
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum KubernetesIdentityError {
    Codec(String),
    InvalidPodIdentity,
    InvalidTopology,
    MessageTooLarge(usize),
    NonCanonical,
    UnsupportedIdentityVersion(u16),
    UnsupportedTopologyVersion(u16),
}

impl Display for KubernetesIdentityError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        write!(formatter, "{self:?}")
    }
}

impl Error for KubernetesIdentityError {}

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

    fn uid(seed: u8) -> Result<KubernetesObjectUidDigest, Box<dyn Error>> {
        Ok(KubernetesObjectUidDigest::new([seed; 32])?)
    }

    #[test]
    fn shared_service_account_still_produces_distinct_pod_subjects() -> Result<(), Box<dyn Error>> {
        let base = PodIdentityV1 {
            format_version: KUBERNETES_IDENTITY_FORMAT_VERSION,
            cluster_uid_digest: uid(1)?,
            namespace: "lodestone".to_owned(),
            service_account: "indexer".to_owned(),
            pod_uid_digest: uid(2)?,
            requested_role: NodeRole::Indexer,
            public_key_digest: [3; 32],
        };
        let mut replacement = base.clone();
        replacement.pod_uid_digest = uid(4)?;
        replacement.public_key_digest = [5; 32];
        assert_ne!(base.provider_subject_digest()?, replacement.provider_subject_digest()?);
        assert_eq!(decode_pod_identity(&encode_pod_identity(&base)?)?, base);
        Ok(())
    }

    #[test]
    fn express_uses_provider_zone_id_not_scheduler_alias() -> Result<(), Box<dyn Error>> {
        let topology = NodeTopologyV1 {
            format_version: NODE_TOPOLOGY_FORMAT_VERSION,
            provider_observation: ProviderComputeTopologyV1::Aws {
                partition: AwsPartition::Aws,
                region: "ap-south-1".to_owned(),
                availability_zone_id: "aps1-az2".to_owned(),
            },
            kubernetes_zone_name: "ap-south-1a".to_owned(),
        };
        let expected = AwsBucketKind::Directory {
            availability_zone_id: "aps1-az1".to_owned(),
        };
        assert!(!topology.admits_bucket_kind(&expected));
        assert_eq!(decode_node_topology(&encode_node_topology(&topology)?)?, topology);
        Ok(())
    }

    #[test]
    fn static_development_topology_cannot_admit_an_aws_bucket() {
        let topology = NodeTopologyV1 {
            format_version: NODE_TOPOLOGY_FORMAT_VERSION,
            provider_observation: ProviderComputeTopologyV1::StaticDevelopment {
                deployment_identity: "local".to_owned(),
                zone: "zone-a".to_owned(),
            },
            kubernetes_zone_name: "zone-a".to_owned(),
        };
        assert!(!topology.admits_bucket_kind(&AwsBucketKind::GeneralPurpose));
        assert!(!topology.admits_bucket_kind(&AwsBucketKind::Directory {
            availability_zone_id: "aps1-az1".to_owned(),
        }));
    }

    #[test]
    fn kubernetes_identity_names_use_exact_dns_bounds() -> Result<(), Box<dyn Error>> {
        let mut identity = PodIdentityV1 {
            format_version: KUBERNETES_IDENTITY_FORMAT_VERSION,
            cluster_uid_digest: uid(1)?,
            namespace: "a".repeat(MAX_KUBERNETES_DNS_LABEL_BYTES + 1),
            service_account: "indexer".to_owned(),
            pod_uid_digest: uid(2)?,
            requested_role: NodeRole::Indexer,
            public_key_digest: [3; 32],
        };
        assert!(matches!(
            identity.validate(),
            Err(KubernetesIdentityError::InvalidPodIdentity)
        ));
        identity.namespace = "lodestone".to_owned();
        identity.service_account = "indexer..worker".to_owned();
        assert!(matches!(
            identity.validate(),
            Err(KubernetesIdentityError::InvalidPodIdentity)
        ));
        identity.service_account = "indexer.worker".to_owned();
        identity.validate()?;
        Ok(())
    }
}