lodestone-proto 0.0.31

//! Strict bounded public-session and workload-credential wire values.

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

use lodestone_types::{CollectionId, CredentialGeneration, IndexName, NodeToken, WorkloadClassV1};
use serde::{Deserialize, Serialize};

pub const PUBLIC_SESSION_FORMAT_VERSION: u16 = 6;
pub const WORKLOAD_SESSION_FORMAT_VERSION: u16 = 1;
pub const MAX_PUBLIC_SESSION_BYTES: usize = 1_024;
pub const MAX_API_KEY_CREDENTIAL_BYTES: usize = 512;
pub const MAX_WORKLOAD_SESSION_BYTES: usize = 512;

/// Derives the stable diagnostic identity carried by both a native response
/// and its durable audit event. The public session request identity is checked
/// against the request envelope before execution, so neither side can choose
/// a different correlation key after admission.
#[must_use]
pub fn native_response_diagnostic_identity(request_id: &str, correlation: [u8; 16]) -> [u8; 16] {
    let mut diagnostic = blake3::Hasher::new();
    diagnostic.update(b"lodestone/native-response-diagnostic/v1\0");
    diagnostic.update(&correlation);
    diagnostic.update(&(request_id.len() as u64).to_be_bytes());
    diagnostic.update(request_id.as_bytes());
    let digest = diagnostic.finalize();
    let mut diagnostic_id = [0_u8; 16];
    diagnostic_id.copy_from_slice(&digest.as_bytes()[..16]);
    if diagnostic_id == [0; 16] {
        diagnostic_id[15] = 1;
    }
    diagnostic_id
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum PublicSessionActionV1 {
    IdentityDiscovery,
    DataRead,
    DataWrite,
    Append,
    ResolveAppend,
    IndexAdmin,
    CatalogRead,
    TaskRead,
    TaskAdmin,
    SubscriptionRead,
    DiagnosticsRead,
    DiagnosticsCreate,
    SecurityRead,
    SecurityAdmin,
    AuditRead,
    QuotaRead,
    QuotaAdmin,
    BindingRead,
    LifecycleAdmin,
    AdmissionRead,
    AdmissionAdmin,
    UsageRead,
    OperationRead,
    OperationCancel,
    CollectionExport,
    CollectionRestore,
    CollectionTransferExport,
    CollectionTransferRestore,
    IndexExport,
    IndexImport,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct PublicSessionV1 {
    pub format_version: u16,
    pub api_key_credential: Vec<u8>,
    pub collection_id: CollectionId,
    pub request_id: String,
    pub action: PublicSessionActionV1,
    pub workload_class: WorkloadClassV1,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub index_name: Option<IndexName>,
    pub request_correlation_id: [u8; 16],
}

impl PublicSessionV1 {
    pub fn validate(&self) -> Result<(), PublicSessionWireError> {
        if self.format_version != PUBLIC_SESSION_FORMAT_VERSION {
            return Err(PublicSessionWireError::UnsupportedVersion(self.format_version));
        }
        if self.api_key_credential.is_empty()
            || self.api_key_credential.len() > MAX_API_KEY_CREDENTIAL_BYTES
            || !self.api_key_credential.is_ascii()
        {
            return Err(PublicSessionWireError::InvalidCredentialShape);
        }
        if self.request_id.is_empty()
            || self.request_id.len() > 128
            || self.request_id.bytes().any(|byte| byte.is_ascii_control())
        {
            return Err(PublicSessionWireError::InvalidRequestIdentity);
        }
        if !self.action.accepts_index_scope(self.index_name.is_some()) {
            return Err(PublicSessionWireError::InvalidCollectionIndexRef);
        }
        Ok(())
    }
}

impl PublicSessionActionV1 {
    #[must_use]
    pub const fn default_workload_class(self) -> WorkloadClassV1 {
        match self {
            Self::DataWrite | Self::Append | Self::ResolveAppend => WorkloadClassV1::StreamingIngest,
            Self::DataRead | Self::SubscriptionRead => WorkloadClassV1::InteractiveSearch,
            Self::IdentityDiscovery
            | Self::IndexAdmin
            | Self::CatalogRead
            | Self::TaskRead
            | Self::TaskAdmin
            | Self::DiagnosticsRead
            | Self::DiagnosticsCreate
            | Self::SecurityRead
            | Self::SecurityAdmin
            | Self::AuditRead
            | Self::QuotaRead
            | Self::QuotaAdmin
            | Self::BindingRead
            | Self::LifecycleAdmin
            | Self::AdmissionRead
            | Self::AdmissionAdmin
            | Self::UsageRead
            | Self::OperationRead
            | Self::OperationCancel
            | Self::CollectionExport
            | Self::CollectionRestore
            | Self::CollectionTransferExport
            | Self::CollectionTransferRestore
            | Self::IndexExport
            | Self::IndexImport => WorkloadClassV1::BatchAnalytics,
        }
    }

    const fn accepts_index_scope(self, index_present: bool) -> bool {
        match self {
            // Data reads may name one index in the prelude or defer the exact
            // Collection-local vector to a bounded search/PIT request body.
            Self::DataRead | Self::UsageRead => true,
            Self::DataWrite
            | Self::Append
            | Self::ResolveAppend
            | Self::IndexAdmin
            | Self::SubscriptionRead
            | Self::IndexExport
            | Self::IndexImport => index_present,
            _ => !index_present,
        }
    }

    #[must_use]
    pub const fn accepts_workload_class(self, workload_class: WorkloadClassV1) -> bool {
        match self {
            Self::DataWrite | Self::Append | Self::ResolveAppend => matches!(
                workload_class,
                WorkloadClassV1::StreamingIngest | WorkloadClassV1::BulkBackfill
            ),
            Self::DataRead | Self::SubscriptionRead => matches!(
                workload_class,
                WorkloadClassV1::InteractiveSearch | WorkloadClassV1::BatchAnalytics
            ),
            Self::IdentityDiscovery
            | Self::IndexAdmin
            | Self::CatalogRead
            | Self::TaskRead
            | Self::TaskAdmin
            | Self::DiagnosticsRead
            | Self::DiagnosticsCreate
            | Self::SecurityRead
            | Self::SecurityAdmin
            | Self::AuditRead
            | Self::QuotaRead
            | Self::QuotaAdmin
            | Self::BindingRead
            | Self::LifecycleAdmin
            | Self::AdmissionRead
            | Self::AdmissionAdmin
            | Self::UsageRead
            | Self::OperationRead
            | Self::OperationCancel
            | Self::CollectionExport
            | Self::CollectionRestore
            | Self::CollectionTransferExport
            | Self::CollectionTransferRestore
            | Self::IndexExport
            | Self::IndexImport => matches!(workload_class, WorkloadClassV1::BatchAnalytics),
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WorkloadSessionV1 {
    pub format_version: u16,
    pub node_token: NodeToken,
    pub credential_generation: CredentialGeneration,
    pub leaf_certificate_sha256: [u8; 32],
    pub issuer_sha256: [u8; 32],
}

impl WorkloadSessionV1 {
    pub fn validate(&self) -> Result<(), PublicSessionWireError> {
        if self.format_version != WORKLOAD_SESSION_FORMAT_VERSION {
            return Err(PublicSessionWireError::UnsupportedVersion(self.format_version));
        }
        if self.credential_generation.get() == 0
            || self.leaf_certificate_sha256 == [0; 32]
            || self.issuer_sha256 == [0; 32]
        {
            return Err(PublicSessionWireError::InvalidWorkloadBinding);
        }
        Ok(())
    }
}

pub fn encode_workload_session_v1(session: &WorkloadSessionV1) -> Result<Vec<u8>, PublicSessionWireError> {
    session.validate()?;
    let bytes = serde_json::to_vec(session).map_err(|error| PublicSessionWireError::Codec(error.to_string()))?;
    if bytes.len() > MAX_WORKLOAD_SESSION_BYTES {
        return Err(PublicSessionWireError::FrameTooLarge(bytes.len()));
    }
    Ok(bytes)
}

pub fn decode_workload_session_v1(bytes: &[u8]) -> Result<WorkloadSessionV1, PublicSessionWireError> {
    if bytes.is_empty() || bytes.len() > MAX_WORKLOAD_SESSION_BYTES {
        return Err(PublicSessionWireError::FrameTooLarge(bytes.len()));
    }
    let session: WorkloadSessionV1 =
        serde_json::from_slice(bytes).map_err(|error| PublicSessionWireError::Codec(error.to_string()))?;
    session.validate()?;
    if serde_json::to_vec(&session).map_err(|error| PublicSessionWireError::Codec(error.to_string()))? != bytes {
        return Err(PublicSessionWireError::NonCanonical);
    }
    Ok(session)
}

pub fn encode_public_session_v1(session: &PublicSessionV1) -> Result<Vec<u8>, PublicSessionWireError> {
    session.validate()?;
    let bytes = serde_json::to_vec(session).map_err(|error| PublicSessionWireError::Codec(error.to_string()))?;
    if bytes.len() > MAX_PUBLIC_SESSION_BYTES {
        return Err(PublicSessionWireError::FrameTooLarge(bytes.len()));
    }
    Ok(bytes)
}

pub fn decode_public_session_v1(bytes: &[u8]) -> Result<PublicSessionV1, PublicSessionWireError> {
    if bytes.is_empty() || bytes.len() > MAX_PUBLIC_SESSION_BYTES {
        return Err(PublicSessionWireError::FrameTooLarge(bytes.len()));
    }
    let session: PublicSessionV1 =
        serde_json::from_slice(bytes).map_err(|error| PublicSessionWireError::Codec(error.to_string()))?;
    session.validate()?;
    if serde_json::to_vec(&session).map_err(|error| PublicSessionWireError::Codec(error.to_string()))? != bytes {
        return Err(PublicSessionWireError::NonCanonical);
    }
    Ok(session)
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PublicSessionWireError {
    UnsupportedVersion(u16),
    InvalidCredentialShape,
    InvalidRequestIdentity,
    InvalidCollectionIndexRef,
    InvalidWorkloadBinding,
    FrameTooLarge(usize),
    NonCanonical,
    Codec(String),
}

impl Display for PublicSessionWireError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsupportedVersion(version) => write!(formatter, "unsupported security wire version {version}"),
            Self::InvalidCredentialShape => formatter.write_str("invalid API-key credential shape"),
            Self::InvalidRequestIdentity => formatter.write_str("invalid public request identity"),
            Self::InvalidCollectionIndexRef => formatter.write_str("public session index scope is invalid"),
            Self::InvalidWorkloadBinding => formatter.write_str("invalid workload credential binding"),
            Self::FrameTooLarge(size) => write!(formatter, "security frame size {size} exceeds the bound"),
            Self::NonCanonical => formatter.write_str("security frame is not canonical"),
            Self::Codec(error) => write!(formatter, "security frame codec failed: {error}"),
        }
    }
}

impl Error for PublicSessionWireError {}

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

    #[test]
    fn public_session_is_strict_and_bounded() -> Result<(), Box<dyn Error>> {
        let session = PublicSessionV1 {
            format_version: PUBLIC_SESSION_FORMAT_VERSION,
            api_key_credential: b"1:AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA".to_vec(),
            collection_id: CollectionId::new(1)?,
            request_id: "proto-request".to_owned(),
            action: PublicSessionActionV1::DataRead,
            workload_class: WorkloadClassV1::InteractiveSearch,
            index_name: Some(IndexName::new("logs-test")?),
            request_correlation_id: [7; 16],
        };
        let bytes = encode_public_session_v1(&session)?;
        assert_eq!(decode_public_session_v1(&bytes)?, session);
        assert!(
            session
                .action
                .accepts_workload_class(WorkloadClassV1::InteractiveSearch)
        );
        assert!(session.action.accepts_workload_class(WorkloadClassV1::BatchAnalytics));
        assert!(!session.action.accepts_workload_class(WorkloadClassV1::StreamingIngest));
        assert!(
            !session
                .action
                .accepts_workload_class(WorkloadClassV1::InternalMaintenance)
        );
        let mut newer = session.clone();
        newer.format_version += 1;
        assert_eq!(
            decode_public_session_v1(&serde_json::to_vec(&newer)?),
            Err(PublicSessionWireError::UnsupportedVersion(
                PUBLIC_SESSION_FORMAT_VERSION + 1
            ))
        );
        let mut audit_read = session.clone();
        audit_read.action = PublicSessionActionV1::AuditRead;
        audit_read.index_name = None;
        assert_eq!(
            decode_public_session_v1(&encode_public_session_v1(&audit_read)?)?,
            audit_read
        );
        let mut unknown: serde_json::Value = serde_json::from_slice(&bytes)?;
        unknown["secret"] = serde_json::Value::String("must-not-pass".to_owned());
        assert!(decode_public_session_v1(&serde_json::to_vec(&unknown)?).is_err());
        let mut collection_read = session.clone();
        collection_read.index_name = None;
        assert_eq!(
            decode_public_session_v1(&encode_public_session_v1(&collection_read)?)?,
            collection_read
        );
        let mut exact_usage_read = collection_read.clone();
        exact_usage_read.action = PublicSessionActionV1::UsageRead;
        exact_usage_read.workload_class = WorkloadClassV1::BatchAnalytics;
        exact_usage_read.index_name = Some(IndexName::new("logs-test")?);
        assert_eq!(
            decode_public_session_v1(&encode_public_session_v1(&exact_usage_read)?)?,
            exact_usage_read
        );
        let mut collection_usage_read = exact_usage_read.clone();
        collection_usage_read.index_name = None;
        assert_eq!(
            decode_public_session_v1(&encode_public_session_v1(&collection_usage_read)?)?,
            collection_usage_read
        );
        let mut missing_index = session;
        missing_index.action = PublicSessionActionV1::DataWrite;
        missing_index.index_name = None;
        assert_eq!(
            decode_public_session_v1(&serde_json::to_vec(&missing_index)?),
            Err(PublicSessionWireError::InvalidCollectionIndexRef)
        );
        Ok(())
    }

    #[test]
    fn public_session_decoder_refuses_malformed_corpus_without_panicking() {
        for size in 0..=MAX_PUBLIC_SESSION_BYTES + 2 {
            let input = vec![0xff; size];
            let _ = decode_public_session_v1(&input);
        }
    }

    #[test]
    fn workload_session_is_strict_and_canonical() -> Result<(), Box<dyn Error>> {
        use lodestone_types::{NodeId, NodeIncarnation, RecoveryEra};

        let session = WorkloadSessionV1 {
            format_version: WORKLOAD_SESSION_FORMAT_VERSION,
            node_token: NodeToken::new(RecoveryEra::new(1)?, NodeId::new(2), NodeIncarnation::new(3)?)?,
            credential_generation: CredentialGeneration::new(1),
            leaf_certificate_sha256: [1; 32],
            issuer_sha256: [2; 32],
        };
        let bytes = encode_workload_session_v1(&session)?;
        assert_eq!(decode_workload_session_v1(&bytes)?, session);
        let mut noncanonical = bytes;
        noncanonical.push(b' ');
        assert_eq!(
            decode_workload_session_v1(&noncanonical),
            Err(PublicSessionWireError::NonCanonical)
        );
        Ok(())
    }
}