clusterflux-public/crates/disasmer-core/src/capability.rs

use std::collections::BTreeSet;

use serde::{Deserialize, Serialize};
use thiserror::Error;

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Capability {
    Command,
    Containers,
    RootlessPodman,
    SourceFilesystem,
    SourceGit,
    HostFilesystem,
    Network,
    Secrets,
    InboundPorts,
    ArbitrarySyscalls,
    VfsArtifacts,
    Wasmtime,
    WindowsCommandDev,
    QuicDirect,
}

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum EnvironmentBackend {
    Container,
    NixFlake,
    WindowsCommandDev,
}

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Os {
    Linux,
    Windows,
    Macos,
    Other(String),
}

impl Os {
    pub fn current() -> Self {
        match std::env::consts::OS {
            "linux" => Self::Linux,
            "windows" => Self::Windows,
            "macos" => Self::Macos,
            other => Self::Other(other.to_owned()),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeCapabilities {
    pub os: Os,
    pub arch: String,
    pub capabilities: BTreeSet<Capability>,
    pub environment_backends: BTreeSet<EnvironmentBackend>,
    pub source_providers: BTreeSet<String>,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum CapabilityReportError {
    #[error("node architecture `{0}` is invalid")]
    InvalidArchitecture(String),
    #[error("node OS label `{0}` is invalid")]
    InvalidOsLabel(String),
    #[error("source provider id `{0}` is invalid")]
    InvalidSourceProvider(String),
}

impl NodeCapabilities {
    pub fn detect_current() -> Self {
        let os = Os::current();
        let mut capabilities = BTreeSet::from([
            Capability::Command,
            Capability::SourceFilesystem,
            Capability::VfsArtifacts,
            Capability::Wasmtime,
        ]);
        let mut environment_backends = BTreeSet::new();

        match os {
            Os::Linux => {
                capabilities.insert(Capability::Containers);
                capabilities.insert(Capability::RootlessPodman);
                environment_backends.insert(EnvironmentBackend::Container);
            }
            Os::Windows => {
                capabilities.insert(Capability::WindowsCommandDev);
                environment_backends.insert(EnvironmentBackend::WindowsCommandDev);
            }
            Os::Macos | Os::Other(_) => {}
        }

        Self {
            os,
            arch: std::env::consts::ARCH.to_owned(),
            capabilities,
            environment_backends,
            source_providers: BTreeSet::from(["filesystem".to_owned(), "git".to_owned()]),
        }
    }

    pub fn with_capability(mut self, capability: Capability) -> Self {
        self.capabilities.insert(capability);
        self
    }

    pub fn has_all(&self, required: &BTreeSet<Capability>) -> bool {
        required
            .iter()
            .all(|capability| self.capabilities.contains(capability))
    }

    pub fn validate_public_report(&self) -> Result<(), CapabilityReportError> {
        if !valid_capability_label(&self.arch) {
            return Err(CapabilityReportError::InvalidArchitecture(
                self.arch.clone(),
            ));
        }
        if let Os::Other(label) = &self.os {
            if !valid_capability_label(label) {
                return Err(CapabilityReportError::InvalidOsLabel(label.clone()));
            }
        }
        for provider in &self.source_providers {
            if !valid_source_provider_id(provider) {
                return Err(CapabilityReportError::InvalidSourceProvider(
                    provider.clone(),
                ));
            }
        }
        Ok(())
    }
}

fn valid_capability_label(label: &str) -> bool {
    !label.is_empty()
        && label.len() <= 64
        && label.bytes().all(
            |byte| matches!(byte, b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'-' | b'_' | b'.'),
        )
}

fn valid_source_provider_id(provider: &str) -> bool {
    !provider.is_empty()
        && provider.len() <= 64
        && provider
            .bytes()
            .all(|byte| matches!(byte, b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.'))
}

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

    fn capabilities() -> NodeCapabilities {
        NodeCapabilities {
            os: Os::Linux,
            arch: "x86_64".to_owned(),
            capabilities: BTreeSet::from([Capability::Command]),
            environment_backends: BTreeSet::new(),
            source_providers: BTreeSet::from(["filesystem".to_owned(), "git".to_owned()]),
        }
    }

    #[test]
    fn capability_reports_validate_hostile_strings() {
        assert!(capabilities().validate_public_report().is_ok());

        let mut invalid_arch = capabilities();
        invalid_arch.arch = "x86_64\nmalicious".to_owned();
        assert_eq!(
            invalid_arch.validate_public_report(),
            Err(CapabilityReportError::InvalidArchitecture(
                "x86_64\nmalicious".to_owned()
            ))
        );

        let mut invalid_provider = capabilities();
        invalid_provider
            .source_providers
            .insert("../checkout".to_owned());
        assert_eq!(
            invalid_provider.validate_public_report(),
            Err(CapabilityReportError::InvalidSourceProvider(
                "../checkout".to_owned()
            ))
        );
    }
}