tls.rs

//! Certificate management for mTLS authentication.
//!
//! Provides a self-signed Certificate Authority that can issue server and client
//! certificates.  The CA is generated once during `relay-runner init` and stored
//! on disk.  Client certificates are signed by the CA and can be verified by the
//! server during TLS handshake (mutual TLS).

use std::fmt;
use std::fs;
use std::path::{Path, PathBuf};

use rcgen::{
    BasicConstraints, CertificateParams, DistinguishedName, DnType, ExtendedKeyUsagePurpose, IsCa,
    Issuer, KeyPair, KeyUsagePurpose,
};
use sha2::{Digest, Sha256};
use time::OffsetDateTime;

// ---------------------------------------------------------------------------
// TlsError
// ---------------------------------------------------------------------------

/// Errors that can occur during certificate operations.
#[derive(Debug)]
pub enum TlsError {
    IoError(std::io::Error),
    CertGenError(String),
    InvalidCert(String),
}

impl fmt::Display for TlsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TlsError::IoError(e) => write!(f, "I/O error: {e}"),
            TlsError::CertGenError(msg) => write!(f, "certificate generation error: {msg}"),
            TlsError::InvalidCert(msg) => write!(f, "invalid certificate: {msg}"),
        }
    }
}

impl std::error::Error for TlsError {}

impl From<std::io::Error> for TlsError {
    fn from(e: std::io::Error) -> Self {
        TlsError::IoError(e)
    }
}

impl From<rcgen::Error> for TlsError {
    fn from(e: rcgen::Error) -> Self {
        TlsError::CertGenError(e.to_string())
    }
}

// ---------------------------------------------------------------------------
// ClientCert
// ---------------------------------------------------------------------------

/// A client certificate issued by the CA, ready to be exported.
pub struct ClientCert {
    pub cert_pem: String,
    pub key_pem: String,
    pub fingerprint: String,
    pub verification_code: String,
}

// ---------------------------------------------------------------------------
// CertificateAuthority
// ---------------------------------------------------------------------------

/// A self-signed CA that issues server and client certificates for mTLS.
pub struct CertificateAuthority {
    ca_cert_pem: String,
    ca_key_pem: String,
    // Stored for future use: cert rotation and re-issuance paths.
    #[allow(dead_code)]
    cert_path: PathBuf,
    #[allow(dead_code)]
    key_path: PathBuf,
}

impl CertificateAuthority {
    /// Well-known file names inside the state directory.
    const CA_CERT_FILE: &'static str = "ca.crt";
    const CA_KEY_FILE: &'static str = "ca.key";
    const SERVER_CERT_FILE: &'static str = "server.crt";
    const SERVER_KEY_FILE: &'static str = "server.key";
    const CLIENTS_DIR: &'static str = "clients";

    /// Load an existing CA from `state_dir`, or generate a new one if the files
    /// do not exist.
    pub fn load_or_generate(state_dir: &Path) -> Result<Self, TlsError> {
        let cert_path = state_dir.join(Self::CA_CERT_FILE);
        let key_path = state_dir.join(Self::CA_KEY_FILE);

        if cert_path.exists() && key_path.exists() {
            let ca_cert_pem = fs::read_to_string(&cert_path)?;
            let ca_key_pem = fs::read_to_string(&key_path)?;
            return Ok(Self {
                ca_cert_pem,
                ca_key_pem,
                cert_path,
                key_path,
            });
        }

        // Generate new CA.
        let ca_key = KeyPair::generate()?;

        let mut ca_params = CertificateParams::default();
        ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
        ca_params.key_usages.push(KeyUsagePurpose::KeyCertSign);
        ca_params.key_usages.push(KeyUsagePurpose::CrlSign);

        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, "Relay Runner CA");
        dn.push(DnType::OrganizationName, "Relay");
        ca_params.distinguished_name = dn;

        // Valid for 10 years.
        ca_params.not_before = OffsetDateTime::now_utc();
        ca_params.not_after = OffsetDateTime::now_utc() + time::Duration::days(3650);

        let ca_cert = ca_params.self_signed(&ca_key)?;

        let ca_cert_pem = ca_cert.pem();
        let ca_key_pem = ca_key.serialize_pem();

        // Persist to disk.
        fs::create_dir_all(state_dir)?;
        fs::write(&cert_path, &ca_cert_pem)?;
        fs::write(&key_path, &ca_key_pem)?;

        // Restrict key file permissions on Unix.
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            fs::set_permissions(&key_path, fs::Permissions::from_mode(0o600))?;
        }

        Ok(Self {
            ca_cert_pem,
            ca_key_pem,
            cert_path,
            key_path,
        })
    }

    /// Generate a server certificate signed by this CA.
    ///
    /// The certificate is written to `state_dir/server.crt` and
    /// `state_dir/server.key`.  Returns `(cert_pem, key_pem)`.
    pub fn issue_server_cert(&self, state_dir: &Path) -> Result<(String, String), TlsError> {
        let server_key = KeyPair::generate()?;

        let mut params = CertificateParams::default();
        params.is_ca = IsCa::NoCa;

        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, "Relay Runner");
        dn.push(DnType::OrganizationName, "Relay");
        params.distinguished_name = dn;

        params
            .extended_key_usages
            .push(ExtendedKeyUsagePurpose::ServerAuth);

        // Include localhost + common LAN names as SANs.
        params.subject_alt_names = vec![
            rcgen::SanType::DnsName(
                "localhost"
                    .try_into()
                    .map_err(|e: rcgen::Error| TlsError::CertGenError(e.to_string()))?,
            ),
            rcgen::SanType::IpAddress(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)),
            rcgen::SanType::IpAddress(std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)),
        ];

        // Valid for 2 years.
        params.not_before = OffsetDateTime::now_utc();
        params.not_after = OffsetDateTime::now_utc() + time::Duration::days(730);

        let issuer = self.reconstruct_ca()?;
        let server_cert = params.signed_by(&server_key, &issuer)?;

        let cert_pem = server_cert.pem();
        let key_pem = server_key.serialize_pem();

        let cert_path = state_dir.join(Self::SERVER_CERT_FILE);
        let key_path = state_dir.join(Self::SERVER_KEY_FILE);

        fs::write(&cert_path, &cert_pem)?;
        fs::write(&key_path, &key_pem)?;

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            fs::set_permissions(&key_path, fs::Permissions::from_mode(0o600))?;
        }

        Ok((cert_pem, key_pem))
    }

    /// Generate a client certificate signed by this CA.
    ///
    /// The certificate is also saved to `state_dir/clients/<client_name>.crt`.
    pub fn issue_client_cert(
        &self,
        state_dir: &Path,
        client_name: &str,
    ) -> Result<ClientCert, TlsError> {
        let client_key = KeyPair::generate()?;

        let mut params = CertificateParams::default();
        params.is_ca = IsCa::NoCa;

        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, client_name);
        dn.push(DnType::OrganizationName, "Relay Client");
        params.distinguished_name = dn;

        params
            .extended_key_usages
            .push(ExtendedKeyUsagePurpose::ClientAuth);

        // Valid for 1 year.
        params.not_before = OffsetDateTime::now_utc();
        params.not_after = OffsetDateTime::now_utc() + time::Duration::days(365);

        let issuer = self.reconstruct_ca()?;
        let client_cert = params.signed_by(&client_key, &issuer)?;

        let cert_pem = client_cert.pem();
        let key_pem = client_key.serialize_pem();
        let fingerprint = Self::fingerprint(&cert_pem)?;
        let verification_code = Self::verification_code_from_fingerprint(&fingerprint);

        // Save client cert to clients/ subdirectory.
        let clients_dir = state_dir.join(Self::CLIENTS_DIR);
        fs::create_dir_all(&clients_dir)?;
        fs::write(clients_dir.join(format!("{client_name}.crt")), &cert_pem)?;
        fs::write(clients_dir.join(format!("{client_name}.key")), &key_pem)?;

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            fs::set_permissions(
                clients_dir.join(format!("{client_name}.key")),
                fs::Permissions::from_mode(0o600),
            )?;
        }

        Ok(ClientCert {
            cert_pem,
            key_pem,
            fingerprint,
            verification_code,
        })
    }

    /// Reconstruct an `Issuer` from the stored CA PEM data.
    ///
    /// `rcgen::Issuer::from_ca_cert_pem()` parses the CA certificate PEM and
    /// combines it with the original `KeyPair` to produce an `Issuer` that
    /// rcgen 0.14 uses as the signing authority for leaf certificates.
    fn reconstruct_ca(&self) -> Result<Issuer<'static, KeyPair>, TlsError> {
        let ca_key = KeyPair::from_pem(&self.ca_key_pem)?;

        Issuer::from_ca_cert_pem(&self.ca_cert_pem, ca_key)
            .map_err(|e| TlsError::CertGenError(format!("CA cert parse error: {e}")))
    }

    /// Return the CA certificate PEM (used for TLS client verification).
    pub fn ca_cert_pem(&self) -> &str {
        &self.ca_cert_pem
    }

    /// Return `(cert_pem, key_pem)` for the server identity, reading from the
    /// files on disk.
    #[allow(dead_code)]
    pub fn server_identity(&self, state_dir: &Path) -> Result<(String, String), TlsError> {
        let cert = fs::read_to_string(state_dir.join(Self::SERVER_CERT_FILE))?;
        let key = fs::read_to_string(state_dir.join(Self::SERVER_KEY_FILE))?;
        Ok((cert, key))
    }

    /// Compute the SHA-256 fingerprint of a PEM-encoded certificate.
    ///
    /// Returns a colon-separated hex string like `A1:B2:C3:...`.
    pub fn fingerprint(cert_pem: &str) -> Result<String, TlsError> {
        let der = pem::parse(cert_pem)
            .map_err(|e| TlsError::InvalidCert(format!("PEM parse error: {e}")))?;

        let digest = Sha256::digest(der.contents());
        let hex: Vec<String> = digest.iter().map(|b| format!("{b:02X}")).collect();
        Ok(hex.join(":"))
    }

    /// Generate a 6-digit verification code from a certificate PEM.
    ///
    /// The code is derived from the SHA-256 fingerprint so both the server
    /// operator and the client user can independently compute and compare it.
    pub fn verification_code(cert_pem: &str) -> Result<String, TlsError> {
        let fp = Self::fingerprint(cert_pem)?;
        Ok(Self::verification_code_from_fingerprint(&fp))
    }

    /// Rotate the CA: generate a new CA keypair and persist it.
    ///
    /// The old CA files are renamed with a `.old` suffix so they can be used
    /// during a transition period (the server can accept client certs signed
    /// by either CA).
    pub fn rotate(state_dir: &Path) -> Result<Self, TlsError> {
        let old_cert = state_dir.join(Self::CA_CERT_FILE);
        let old_key = state_dir.join(Self::CA_KEY_FILE);

        if old_cert.exists() {
            fs::rename(&old_cert, state_dir.join("ca.crt.old"))?;
        }
        if old_key.exists() {
            fs::rename(&old_key, state_dir.join("ca.key.old"))?;
        }

        Self::load_or_generate(state_dir)
    }

    // -- private helpers ------------------------------------------------------

    /// Derive a 6-digit code from a fingerprint string.
    fn verification_code_from_fingerprint(fingerprint: &str) -> String {
        // Take the first 8 bytes of the fingerprint hex, interpret as u64,
        // then modulo 1_000_000 to get a 6-digit code.
        let hex_chars: String = fingerprint
            .chars()
            .filter(char::is_ascii_hexdigit)
            .collect();
        let bytes = hex_chars
            .as_bytes()
            .chunks(2)
            .take(4)
            .fold(0u32, |acc, chunk| {
                let s = std::str::from_utf8(chunk).unwrap_or("00");
                let b = u32::from_str_radix(s, 16).unwrap_or(0);
                (acc << 8) | b
            });
        format!("{:06}", bytes % 1_000_000)
    }
}

// ---------------------------------------------------------------------------
// State directory resolution
// ---------------------------------------------------------------------------

/// Determine the state directory for relay-runner.
///
/// Prefers `/var/lib/relay` if it exists (or can be created with appropriate
/// permissions).  Falls back to `~/.relay`.
pub fn default_state_dir() -> PathBuf {
    let system_dir = PathBuf::from("/var/lib/relay");
    if system_dir.exists() || fs::create_dir_all(&system_dir).is_ok() {
        return system_dir;
    }

    dirs::home_dir().map_or_else(|| PathBuf::from(".relay"), |h| h.join(".relay"))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn load_or_generate_creates_ca_files() {
        let dir = TempDir::new().unwrap();
        let ca = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        assert!(dir.path().join("ca.crt").exists());
        assert!(dir.path().join("ca.key").exists());
        assert!(ca.ca_cert_pem().contains("BEGIN CERTIFICATE"));
    }

    #[test]
    fn load_or_generate_reloads_existing() {
        let dir = TempDir::new().unwrap();
        let ca1 = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let ca2 = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        assert_eq!(ca1.ca_cert_pem(), ca2.ca_cert_pem());
    }

    #[test]
    fn issue_server_cert_creates_files() {
        let dir = TempDir::new().unwrap();
        let ca = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let (cert, key) = ca.issue_server_cert(dir.path()).unwrap();
        assert!(cert.contains("BEGIN CERTIFICATE"));
        assert!(key.contains("BEGIN PRIVATE KEY"));
        assert!(dir.path().join("server.crt").exists());
        assert!(dir.path().join("server.key").exists());
    }

    #[test]
    fn issue_client_cert_produces_valid_cert() {
        let dir = TempDir::new().unwrap();
        let ca = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let client = ca.issue_client_cert(dir.path(), "test-mac").unwrap();
        assert!(client.cert_pem.contains("BEGIN CERTIFICATE"));
        assert!(client.key_pem.contains("BEGIN PRIVATE KEY"));
        assert!(!client.fingerprint.is_empty());
        assert_eq!(client.verification_code.len(), 6);
        assert!(dir.path().join("clients/test-mac.crt").exists());
    }

    #[test]
    fn fingerprint_is_deterministic() {
        let dir = TempDir::new().unwrap();
        let ca = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let fp1 = CertificateAuthority::fingerprint(ca.ca_cert_pem()).unwrap();
        let fp2 = CertificateAuthority::fingerprint(ca.ca_cert_pem()).unwrap();
        assert_eq!(fp1, fp2);
        assert!(fp1.contains(':'));
    }

    #[test]
    fn verification_code_is_6_digits() {
        let dir = TempDir::new().unwrap();
        let ca = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let code = CertificateAuthority::verification_code(ca.ca_cert_pem()).unwrap();
        assert_eq!(code.len(), 6);
        assert!(code.chars().all(|c| c.is_ascii_digit()));
    }

    #[test]
    fn rotate_creates_old_files() {
        let dir = TempDir::new().unwrap();
        let _ca1 = CertificateAuthority::load_or_generate(dir.path()).unwrap();
        let _ca2 = CertificateAuthority::rotate(dir.path()).unwrap();
        assert!(dir.path().join("ca.crt.old").exists());
        assert!(dir.path().join("ca.key.old").exists());
        assert!(dir.path().join("ca.crt").exists());
    }
}