lodestone-client 0.0.17

//! Dedicated current-only cluster-control client.
//!
//! The client has no Collection profile, API-key provider, public session, or
//! connection-pool reuse. Each call opens one TLS connection with the isolated
//! ALPN and sends one bounded canonical cluster-control frame.

use std::error::Error;
use std::fmt::{self, Display, Formatter};
use std::sync::Arc;
use std::time::Duration;

use lodestone_proto::{
    CLUSTER_CONTROL_REQUEST_FORMAT_VERSION, CLUSTER_CONTROL_SESSION_FORMAT_VERSION, ClusterControlActionV1,
    ClusterControlCredentialKindV1, ClusterControlRequestV1, ClusterControlResponseDispositionV1,
    ClusterControlResponseOutcomeV1, ClusterControlSessionV1, MAX_CLUSTER_CONTROL_FRAME_BYTES,
    decode_cluster_control_response_v1, encode_cluster_control_request_v1,
};
use rustls::pki_types::{CertificateDer, ServerName};
use serde::Serialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

pub const CLUSTER_CONTROL_ALPN: &[u8] = b"lodestone-cluster-control-v1";
const MAX_CLUSTER_CONTROL_CALL_MILLIS: u64 = 120_000;
const MAX_CLUSTER_CONTROL_ENDPOINT_BYTES: usize = 512;

pub struct ClusterControlCredential {
    kind: ClusterControlCredentialKindV1,
    bytes: Vec<u8>,
}

impl ClusterControlCredential {
    pub fn new(kind: ClusterControlCredentialKindV1, bytes: Vec<u8>) -> Result<Self, ClusterControlClientError> {
        if lodestone_proto::validate_cluster_control_credential_v1(&bytes).is_err() {
            return Err(ClusterControlClientError::InvalidConfiguration);
        }
        Ok(Self { kind, bytes })
    }

    #[must_use]
    pub const fn kind(&self) -> ClusterControlCredentialKindV1 {
        self.kind
    }
}

impl fmt::Debug for ClusterControlCredential {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ClusterControlCredential")
            .field("kind", &self.kind)
            .field("bytes", &"[REDACTED]")
            .finish()
    }
}

impl Drop for ClusterControlCredential {
    fn drop(&mut self) {
        self.bytes.fill(0);
        std::hint::black_box(&mut self.bytes);
    }
}

pub struct ClusterControlClientConfig {
    endpoint: String,
    server_name: String,
    trust_certificates_der: Vec<Vec<u8>>,
    credential: ClusterControlCredential,
    timeout: Duration,
}

impl ClusterControlClientConfig {
    pub fn new(
        endpoint: String,
        server_name: String,
        trust_certificates_der: Vec<Vec<u8>>,
        credential: ClusterControlCredential,
        timeout: Duration,
    ) -> Result<Self, ClusterControlClientError> {
        if endpoint.is_empty()
            || endpoint.len() > MAX_CLUSTER_CONTROL_ENDPOINT_BYTES
            || endpoint.parse::<std::net::SocketAddr>().is_err()
            || !matches!(ServerName::try_from(server_name.clone()), Ok(ServerName::DnsName(_)))
            || trust_certificates_der.is_empty()
            || timeout.is_zero()
            || timeout > Duration::from_millis(MAX_CLUSTER_CONTROL_CALL_MILLIS)
        {
            return Err(ClusterControlClientError::InvalidConfiguration);
        }
        Ok(Self {
            endpoint,
            server_name,
            trust_certificates_der,
            credential,
            timeout,
        })
    }
}

pub struct ClusterControlClient {
    config: ClusterControlClientConfig,
    tls: Arc<rustls::ClientConfig>,
}

impl ClusterControlClient {
    pub fn new(config: ClusterControlClientConfig) -> Result<Self, ClusterControlClientError> {
        let tls = build_tls(&config)?;
        Ok(Self {
            config,
            tls: Arc::new(tls),
        })
    }

    pub async fn initialize_collection(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::InitializeCollectionAuthority, request_id, body)
            .await
    }

    pub async fn activate_collection(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::ActivateCollectionAuthority, request_id, body)
            .await
    }

    pub async fn resolve_collection(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::ResolveCollectionAuthority, request_id, body)
            .await
    }

    pub async fn inspect_collection_retirement(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::InspectCollectionRetirement, request_id, body)
            .await
    }

    pub async fn retire_collection(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::RetireCollectionAuthority, request_id, body)
            .await
    }

    pub async fn observe_capacity(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::ObserveClusterCapacity,
            request_id,
            serde_json::json!({}),
        )
        .await
    }

    pub async fn observe_resource_diagnostics(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::ObserveResourceDiagnostics, request_id, body)
            .await
    }

    pub async fn reserve_credential_rotation(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::ReserveClusterControlCredentialRotation,
            request_id,
            body,
        )
        .await
    }

    pub async fn activate_credential_rotation(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::ActivateClusterControlCredentialRotation,
            request_id,
            body,
        )
        .await
    }

    pub async fn revoke_credential(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::RevokeClusterControlCredential, request_id, body)
            .await
    }

    pub async fn resolve_credential(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::ResolveClusterControlCredential,
            request_id,
            body,
        )
        .await
    }

    pub async fn list_backups(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::ListClusterBackups,
            request_id,
            serde_json::json!({}),
        )
        .await
    }

    pub async fn get_backup(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::GetClusterBackup, request_id, body)
            .await
    }

    pub async fn create_backup(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::CreateClusterBackup, request_id, body)
            .await
    }

    pub async fn resolve_backup(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::ResolveClusterBackup, request_id, body)
            .await
    }

    pub async fn restore_readiness(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV1::GetClusterRestoreReadiness,
            request_id,
            serde_json::json!({}),
        )
        .await
    }

    pub async fn prepare_offline_recovery(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV1::PrepareOfflineRecovery, request_id, body)
            .await
    }

    pub async fn call(
        &self,
        action: ClusterControlActionV1,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        if !self.config.credential.kind.admits(action) {
            return Err(ClusterControlClientError::ActionDenied);
        }
        let body = serde_json::to_value(body).map_err(|_| ClusterControlClientError::Codec)?;
        if !body.is_object() {
            return Err(ClusterControlClientError::Codec);
        }
        let mut correlation = [0_u8; 16];
        getrandom::fill(&mut correlation).map_err(|_| ClusterControlClientError::RandomUnavailable)?;
        let mut request = ClusterControlRequestV1 {
            format_version: CLUSTER_CONTROL_REQUEST_FORMAT_VERSION,
            session: ClusterControlSessionV1 {
                format_version: CLUSTER_CONTROL_SESSION_FORMAT_VERSION,
                credential_kind: self.config.credential.kind,
                credential: self.config.credential.bytes.clone(),
                request_id: request_id.to_owned(),
                action,
                request_correlation_id: correlation,
            },
            body,
        };
        let encoded = encode_cluster_control_request_v1(&request);
        request.session.credential.fill(0);
        std::hint::black_box(&mut request.session.credential);
        let mut payload = encoded.map_err(|_| ClusterControlClientError::Codec)?;
        let result = tokio::time::timeout(self.config.timeout, self.call_once(&payload, request_id, correlation)).await;
        payload.fill(0);
        std::hint::black_box(&mut payload);
        result.map_err(|_| ClusterControlClientError::DeadlineExceeded)?
    }

    async fn call_once(
        &self,
        payload: &[u8],
        request_id: &str,
        correlation: [u8; 16],
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        let tcp = TcpStream::connect(&self.config.endpoint)
            .await
            .map_err(|_| ClusterControlClientError::Transport)?;
        let server_name = ServerName::try_from(self.config.server_name.clone())
            .map_err(|_| ClusterControlClientError::InvalidConfiguration)?;
        let connector = TlsConnector::from(Arc::clone(&self.tls));
        let mut tls = connector
            .connect(server_name, tcp)
            .await
            .map_err(|_| ClusterControlClientError::Transport)?;
        let length = u32::try_from(payload.len()).map_err(|_| ClusterControlClientError::Codec)?;
        tls.write_u32(length)
            .await
            .map_err(|_| ClusterControlClientError::Transport)?;
        tls.write_all(payload)
            .await
            .map_err(|_| ClusterControlClientError::Transport)?;
        let length = usize::try_from(tls.read_u32().await.map_err(|_| ClusterControlClientError::Transport)?)
            .map_err(|_| ClusterControlClientError::Codec)?;
        if length == 0 || length > MAX_CLUSTER_CONTROL_FRAME_BYTES {
            return Err(ClusterControlClientError::Codec);
        }
        let mut response = vec![0; length];
        tls.read_exact(&mut response)
            .await
            .map_err(|_| ClusterControlClientError::Transport)?;
        let response = decode_cluster_control_response_v1(&response).map_err(|_| ClusterControlClientError::Codec)?;
        if response.request_id != request_id || response.request_correlation_id != correlation {
            return Err(ClusterControlClientError::ResponseIdentityMismatch);
        }
        match response.outcome {
            ClusterControlResponseOutcomeV1::Succeeded { body } => Ok(body),
            ClusterControlResponseOutcomeV1::Refused { code, disposition } => {
                Err(ClusterControlClientError::Refused(ClusterControlRefusal {
                    code,
                    disposition,
                }))
            }
        }
    }
}

fn build_tls(config: &ClusterControlClientConfig) -> Result<rustls::ClientConfig, ClusterControlClientError> {
    let mut roots = rustls::RootCertStore::empty();
    for certificate in &config.trust_certificates_der {
        roots
            .add(CertificateDer::from(certificate.clone()))
            .map_err(|_| ClusterControlClientError::InvalidConfiguration)?;
    }
    let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
    let mut tls = rustls::ClientConfig::builder_with_provider(provider)
        .with_protocol_versions(&[&rustls::version::TLS13])
        .map_err(|_| ClusterControlClientError::InvalidConfiguration)?
        .with_root_certificates(roots)
        .with_no_client_auth();
    tls.alpn_protocols = vec![CLUSTER_CONTROL_ALPN.to_vec()];
    Ok(tls)
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ClusterControlRefusal {
    pub code: String,
    pub disposition: ClusterControlResponseDispositionV1,
}

#[derive(Debug)]
pub enum ClusterControlClientError {
    InvalidConfiguration,
    ActionDenied,
    RandomUnavailable,
    Codec,
    Transport,
    DeadlineExceeded,
    ResponseIdentityMismatch,
    Refused(ClusterControlRefusal),
}

impl Display for ClusterControlClientError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::InvalidConfiguration => "cluster-control client configuration is invalid",
            Self::ActionDenied => "cluster-control credential kind cannot invoke this action",
            Self::RandomUnavailable => "cluster-control request identity entropy is unavailable",
            Self::Codec => "cluster-control frame is invalid",
            Self::Transport => "cluster-control transport failed",
            Self::DeadlineExceeded => "cluster-control call deadline elapsed",
            Self::ResponseIdentityMismatch => "cluster-control response identity differs from the request",
            Self::Refused(_) => "cluster-control request was refused",
        })
    }
}

impl Error for ClusterControlClientError {}

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

    #[test]
    fn debug_never_exposes_cluster_credential() -> Result<(), ClusterControlClientError> {
        let credential = ClusterControlCredential::new(
            ClusterControlCredentialKindV1::CapacityObserver,
            b"lcc1_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA".to_vec(),
        )?;
        let debug = format!("{credential:?}");
        assert!(!debug.contains("lcc1_AAA"));
        assert!(debug.contains("REDACTED"));
        Ok(())
    }
}