lodestone-client 0.0.29

//! 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,
    COLLECTION_RETIREMENT_STATUS_FORMAT_VERSION, ClusterControlActionV2, ClusterControlCredentialKindV1,
    ClusterControlRequestV2, ClusterControlResponseDispositionV2, ClusterControlResponseOutcomeV2,
    ClusterControlSessionV2, CollectionRetirementOutcomeV1, CollectionRetirementRequestV1,
    MAX_CLUSTER_CONTROL_FRAME_BYTES, decode_cluster_control_response_v2, encode_cluster_control_request_v2,
};
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-v2";
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(ClusterControlActionV2::InitializeCollectionAuthority, request_id, body)
            .await
    }

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

    pub async fn resolve_collection(
        &self,
        request_id: &str,
        body: impl Serialize,
    ) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(ClusterControlActionV2::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(ClusterControlActionV2::InspectCollectionRetirement, request_id, body)
            .await
    }

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

    /// Submits or advances one exact provisioner-owned Collection retirement.
    /// Only an authoritative `NotSubmitted` outcome is safe to submit again;
    /// `Indeterminate` retains the exact request required by the resolver.
    pub async fn submit_collection_retirement(
        &self,
        request_id: &str,
        request: &CollectionRetirementRequestV1,
    ) -> Result<CollectionRetirementOutcomeV1, ClusterControlClientError> {
        request.validate().map_err(|_| ClusterControlClientError::Codec)?;
        self.call_collection_retirement(
            ClusterControlActionV2::RetireCollectionAuthority,
            request_id,
            request,
            request,
        )
        .await
    }

    /// Resolves one exact Collection-retirement request without inventing a
    /// replacement operation identity after response loss or process restart.
    pub async fn resolve_collection_retirement(
        &self,
        request_id: &str,
        request: &CollectionRetirementRequestV1,
    ) -> Result<CollectionRetirementOutcomeV1, ClusterControlClientError> {
        request.validate().map_err(|_| ClusterControlClientError::Codec)?;
        let body = CollectionRetirementResolutionBodyV1::Retire {
            request: request.clone(),
        };
        self.call_collection_retirement(
            ClusterControlActionV2::ResolveCollectionAuthority,
            request_id,
            body,
            request,
        )
        .await
    }

    async fn call_collection_retirement(
        &self,
        action: ClusterControlActionV2,
        request_id: &str,
        body: impl Serialize,
        request: &CollectionRetirementRequestV1,
    ) -> Result<CollectionRetirementOutcomeV1, ClusterControlClientError> {
        classify_collection_retirement_call(self.call(action, request_id, body).await, request)
    }

    pub async fn observe_capacity(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV2::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(ClusterControlActionV2::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(
            ClusterControlActionV2::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(
            ClusterControlActionV2::ActivateClusterControlCredentialRotation,
            request_id,
            body,
        )
        .await
    }

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

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

    pub async fn list_backups(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV2::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(ClusterControlActionV2::GetClusterBackup, request_id, body)
            .await
    }

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

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

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

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

    pub async fn restore_readiness(&self, request_id: &str) -> Result<serde_json::Value, ClusterControlClientError> {
        self.call(
            ClusterControlActionV2::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(ClusterControlActionV2::PrepareOfflineRecovery, request_id, body)
            .await
    }

    pub async fn call(
        &self,
        action: ClusterControlActionV2,
        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 = ClusterControlRequestV2 {
            format_version: CLUSTER_CONTROL_REQUEST_FORMAT_VERSION,
            session: ClusterControlSessionV2 {
                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_v2(&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_v2(&response).map_err(|_| ClusterControlClientError::Codec)?;
        if response.request_id != request_id || response.request_correlation_id != correlation {
            return Err(ClusterControlClientError::ResponseIdentityMismatch);
        }
        match response.outcome {
            ClusterControlResponseOutcomeV2::Succeeded { body } => Ok(body),
            ClusterControlResponseOutcomeV2::Refused { code, disposition } => {
                Err(ClusterControlClientError::Refused(ClusterControlRefusal {
                    code,
                    disposition,
                }))
            }
        }
    }
}

fn classify_collection_retirement_call(
    result: Result<serde_json::Value, ClusterControlClientError>,
    request: &CollectionRetirementRequestV1,
) -> Result<CollectionRetirementOutcomeV1, ClusterControlClientError> {
    match result {
        Ok(body) => match serde_json::from_value::<CollectionRetirementOutcomeV1>(body) {
            Ok(outcome)
                if outcome.validate().is_ok() && collection_retirement_outcome_matches_request(&outcome, request)? =>
            {
                Ok(outcome)
            }
            Ok(_) | Err(_) => Ok(CollectionRetirementOutcomeV1::Indeterminate {
                format_version: COLLECTION_RETIREMENT_STATUS_FORMAT_VERSION,
                request: Box::new(request.clone()),
            }),
        },
        Err(ClusterControlClientError::Refused(refusal))
            if refusal.disposition == ClusterControlResponseDispositionV2::RetrySameRequest =>
        {
            Ok(CollectionRetirementOutcomeV1::NotSubmitted {
                format_version: COLLECTION_RETIREMENT_STATUS_FORMAT_VERSION,
                request_digest: request
                    .canonical_request_digest()
                    .map_err(|_| ClusterControlClientError::Codec)?,
            })
        }
        Err(ClusterControlClientError::Transport)
        | Err(ClusterControlClientError::DeadlineExceeded)
        | Err(ClusterControlClientError::ResponseIdentityMismatch)
        | Err(ClusterControlClientError::Codec)
        | Err(ClusterControlClientError::Refused(ClusterControlRefusal {
            disposition: ClusterControlResponseDispositionV2::ResolveOperation,
            ..
        })) => Ok(CollectionRetirementOutcomeV1::Indeterminate {
            format_version: COLLECTION_RETIREMENT_STATUS_FORMAT_VERSION,
            request: Box::new(request.clone()),
        }),
        Err(error) => Err(error),
    }
}

fn collection_retirement_outcome_matches_request(
    outcome: &CollectionRetirementOutcomeV1,
    request: &CollectionRetirementRequestV1,
) -> Result<bool, ClusterControlClientError> {
    let request_digest = request
        .canonical_request_digest()
        .map_err(|_| ClusterControlClientError::Codec)?;
    let admitted_generation = request.expected_security_generation.get().checked_add(1);
    Ok(match outcome {
        CollectionRetirementOutcomeV1::NotSubmitted {
            request_digest: observed,
            ..
        } => *observed == request_digest,
        CollectionRetirementOutcomeV1::Blocked {
            request_digest: observed,
            admitted_security_generation,
            ..
        }
        | CollectionRetirementOutcomeV1::InProgress {
            request_digest: observed,
            admitted_security_generation,
            ..
        } => *observed == request_digest && admitted_generation == Some(admitted_security_generation.get()),
        CollectionRetirementOutcomeV1::Indeterminate { request: observed, .. } => **observed == *request,
        CollectionRetirementOutcomeV1::Resolved { receipt, .. } => {
            receipt.operation_id == request.operation_id
                && receipt.request_digest == request_digest
                && receipt.collection_id == request.collection_id
                && admitted_generation == Some(receipt.admitted_security_generation.get())
        }
        CollectionRetirementOutcomeV1::KnownRetired {
            request_digest: observed,
            collection_id,
            ..
        } => *observed == request_digest && *collection_id == request.collection_id,
        CollectionRetirementOutcomeV1::ExpiredUnknown {
            request_digest: observed,
            collection_id,
            operation_id,
            ..
        } => {
            *observed == request_digest
                && *collection_id == request.collection_id
                && *operation_id == request.operation_id
        }
    })
}

#[derive(Serialize)]
#[serde(tag = "operation", rename_all = "snake_case")]
enum CollectionRetirementResolutionBodyV1 {
    Retire { request: CollectionRetirementRequestV1 },
}

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: ClusterControlResponseDispositionV2,
}

#[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::*;
    use lodestone_proto::{
        COLLECTION_RETIREMENT_REQUEST_FORMAT_VERSION, CollectionId, CollectionRetirementClosurePolicyV1,
        CollectionSecurityGeneration,
    };

    fn retirement_request() -> Result<CollectionRetirementRequestV1, Box<dyn Error>> {
        Ok(CollectionRetirementRequestV1 {
            format_version: COLLECTION_RETIREMENT_REQUEST_FORMAT_VERSION,
            operation_id: [19; 32],
            collection_id: CollectionId::new(7)?,
            expected_security_generation: CollectionSecurityGeneration::new(3),
            closure_policy: CollectionRetirementClosurePolicyV1::RetireOwnedChildren,
        })
    }

    #[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(())
    }

    #[test]
    fn collection_retirement_call_classification_preserves_submission_authority() -> Result<(), Box<dyn Error>> {
        let request = retirement_request()?;
        let not_submitted = classify_collection_retirement_call(
            Err(ClusterControlClientError::Refused(ClusterControlRefusal {
                code: "temporarily_unavailable".to_owned(),
                disposition: ClusterControlResponseDispositionV2::RetrySameRequest,
            })),
            &request,
        )?;
        assert!(matches!(
            not_submitted,
            CollectionRetirementOutcomeV1::NotSubmitted { request_digest, .. }
                if request_digest == request.canonical_request_digest()?
        ));

        for ambiguous in [
            ClusterControlClientError::Transport,
            ClusterControlClientError::DeadlineExceeded,
            ClusterControlClientError::ResponseIdentityMismatch,
            ClusterControlClientError::Refused(ClusterControlRefusal {
                code: "submitted_indeterminate".to_owned(),
                disposition: ClusterControlResponseDispositionV2::ResolveOperation,
            }),
        ] {
            assert!(matches!(
                classify_collection_retirement_call(Err(ambiguous), &request)?,
                CollectionRetirementOutcomeV1::Indeterminate { request: retained, .. }
                    if *retained == request
            ));
        }
        assert!(matches!(
            classify_collection_retirement_call(Ok(serde_json::json!({"state": "unknown"})), &request)?,
            CollectionRetirementOutcomeV1::Indeterminate { .. }
        ));
        let mismatched = CollectionRetirementOutcomeV1::NotSubmitted {
            format_version: COLLECTION_RETIREMENT_STATUS_FORMAT_VERSION,
            request_digest: [99; 32],
        };
        assert!(matches!(
            classify_collection_retirement_call(Ok(serde_json::to_value(mismatched)?), &request)?,
            CollectionRetirementOutcomeV1::Indeterminate { request: retained, .. }
                if *retained == request
        ));
        assert!(matches!(
            classify_collection_retirement_call(
                Err(ClusterControlClientError::Refused(ClusterControlRefusal {
                    code: "conflict".to_owned(),
                    disposition: ClusterControlResponseDispositionV2::CorrectRequest,
                })),
                &request,
            ),
            Err(ClusterControlClientError::Refused(_))
        ));
        Ok(())
    }
}