lazaret/crates/lazaret-core/src/protocol/session.rs

use std::collections::BTreeSet;
use std::fmt;

use ed25519_dalek::{SigningKey, VerifyingKey};
use serde::Serialize;

use crate::identity::{
    CapabilityId, CertificateDigest, ProtocolCredentialId, RequestDigest, RequestId,
    TrustManifestDigest,
};
use crate::records::{
    DeveloperDecision, GenerationIdentity, MailboxCancellation, MailboxRequest, MailboxResult,
    MailboxResultAck, MailboxResultChunk, ProtocolCredential, ProtocolCredentialRole,
    ProtocolSubjectId, ProtocolVersion, SessionBinding, SessionRequestCapability, TlsCredential,
    TlsRole, TlsSubjectId, TrustManifest,
};

use super::signing::{
    ProtocolCodecError, ProtocolRecordKind, decode_signed, encode_payload, sha256, sign_payload,
    trust_manifest_digest, verify_payload,
};
use super::window::MAX_UNACKNOWLEDGED_CHUNKS;
use super::{MAX_RESULT_CHUNK_BYTES, request_digest};

#[derive(Debug)]
pub enum ProtocolValidationError {
    Codec(ProtocolCodecError),
    InvalidPublicKey,
    ManifestDigestMismatch,
    CertificateDigestMismatch,
    CredentialDigestMismatch,
    CapabilityDigestMismatch,
    GenerationMismatch,
    RoleMismatch,
    SubjectMismatch,
    SessionBindingMismatch,
    CredentialMismatch,
    CapabilityMismatch,
    RequestMismatch,
    RequestReplay,
    ChunkTooLarge { length: usize },
}

impl fmt::Display for ProtocolValidationError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Codec(error) => error.fmt(formatter),
            Self::InvalidPublicKey => formatter.write_str("protocol public key is invalid"),
            Self::ManifestDigestMismatch => formatter.write_str("Trust Manifest digest mismatch"),
            Self::CertificateDigestMismatch => {
                formatter.write_str("TLS certificate digest mismatch")
            }
            Self::CredentialDigestMismatch => {
                formatter.write_str("protocol credential digest mismatch")
            }
            Self::CapabilityDigestMismatch => {
                formatter.write_str("Session Request Capability digest mismatch")
            }
            Self::GenerationMismatch => {
                formatter.write_str("protocol generation identity mismatch")
            }
            Self::RoleMismatch => formatter.write_str("protocol credential role mismatch"),
            Self::SubjectMismatch => formatter.write_str("protocol subject identity mismatch"),
            Self::SessionBindingMismatch => formatter.write_str("session binding mismatch"),
            Self::CredentialMismatch => formatter.write_str("protocol credential mismatch"),
            Self::CapabilityMismatch => formatter.write_str("Session Request Capability mismatch"),
            Self::RequestMismatch => formatter.write_str("escalation request correlation mismatch"),
            Self::RequestReplay => {
                formatter.write_str("escalation request ID was already admitted")
            }
            Self::ChunkTooLarge { length } => write!(
                formatter,
                "result chunk length {length} exceeds {MAX_RESULT_CHUNK_BYTES}"
            ),
        }
    }
}

impl std::error::Error for ProtocolValidationError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Codec(error) => Some(error),
            _ => None,
        }
    }
}

impl From<ProtocolCodecError> for ProtocolValidationError {
    fn from(error: ProtocolCodecError) -> Self {
        Self::Codec(error)
    }
}

#[derive(Serialize)]
struct TrustManifestClaims<'a> {
    protocol_version: ProtocolVersion,
    generation_identity: &'a GenerationIdentity,
    tls_ca_sha256: CertificateDigest,
    protocol_issuer_public_key: &'a [u8],
}

#[derive(Serialize)]
struct ProtocolCredentialClaims<'a> {
    protocol_version: ProtocolVersion,
    generation_identity: &'a GenerationIdentity,
    role: ProtocolCredentialRole,
    subject_id: ProtocolSubjectId,
    public_key: &'a [u8],
}

#[derive(Serialize)]
struct SessionCapabilityClaims<'a> {
    protocol_version: ProtocolVersion,
    session_binding: &'a SessionBinding,
    coordinator_credential_id: ProtocolCredentialId,
    sandbox_credential_id: ProtocolCredentialId,
}

/// Creates a self-verifying manifest whose digest must still arrive through a trusted path.
///
/// # Errors
///
/// Returns an error when a claim cannot be encoded.
pub fn issue_trust_manifest(
    generation_identity: GenerationIdentity,
    tls_ca_sha256: CertificateDigest,
    protocol_issuer: &SigningKey,
) -> Result<(TrustManifest, TrustManifestDigest), ProtocolValidationError> {
    let public_key = protocol_issuer.verifying_key().to_bytes().to_vec();
    let claims = TrustManifestClaims {
        protocol_version: ProtocolVersion::V1,
        generation_identity: &generation_identity,
        tls_ca_sha256,
        protocol_issuer_public_key: &public_key,
    };
    let claims = encode_payload(&claims)?;
    let signature = sign_payload(ProtocolRecordKind::TrustManifest, &claims, protocol_issuer)?;
    let manifest = TrustManifest {
        protocol_version: ProtocolVersion::V1,
        generation_identity,
        tls_ca_sha256,
        protocol_issuer_public_key: public_key,
        issuer_signature: signature.to_vec(),
    };
    let encoded = encode_payload(&manifest)?;
    let digest = trust_manifest_digest(&encoded);
    Ok((manifest, digest))
}

/// Validates a pinned Trust Manifest and returns its protocol issuer key.
///
/// # Errors
///
/// Returns an error for a stale tuple, digest mismatch, malformed key, or invalid signature.
pub fn verify_trust_manifest(
    manifest: &TrustManifest,
    trusted_digest: TrustManifestDigest,
    expected_generation: &GenerationIdentity,
    expected_tls_ca_sha256: CertificateDigest,
) -> Result<VerifyingKey, ProtocolValidationError> {
    if &manifest.generation_identity != expected_generation {
        return Err(ProtocolValidationError::GenerationMismatch);
    }
    if manifest.tls_ca_sha256 != expected_tls_ca_sha256 {
        return Err(ProtocolValidationError::CertificateDigestMismatch);
    }
    let encoded = encode_payload(manifest)?;
    if trust_manifest_digest(&encoded) != trusted_digest {
        return Err(ProtocolValidationError::ManifestDigestMismatch);
    }
    let key_bytes: [u8; 32] = manifest
        .protocol_issuer_public_key
        .as_slice()
        .try_into()
        .map_err(|_| ProtocolValidationError::InvalidPublicKey)?;
    let key = VerifyingKey::from_bytes(&key_bytes)
        .map_err(|_| ProtocolValidationError::InvalidPublicKey)?;
    let claims = TrustManifestClaims {
        protocol_version: manifest.protocol_version,
        generation_identity: &manifest.generation_identity,
        tls_ca_sha256: manifest.tls_ca_sha256,
        protocol_issuer_public_key: &manifest.protocol_issuer_public_key,
    };
    verify_payload(
        ProtocolRecordKind::TrustManifest,
        &encode_payload(&claims)?,
        &manifest.issuer_signature,
        &key,
    )?;
    Ok(key)
}

/// Issues one generation- and role-bound protocol credential.
///
/// # Errors
///
/// Returns an error when a claim cannot be encoded.
pub fn issue_protocol_credential(
    generation_identity: GenerationIdentity,
    role: ProtocolCredentialRole,
    subject_id: ProtocolSubjectId,
    subject_key: &VerifyingKey,
    protocol_issuer: &SigningKey,
) -> Result<ProtocolCredential, ProtocolValidationError> {
    let public_key = subject_key.to_bytes().to_vec();
    let claims = ProtocolCredentialClaims {
        protocol_version: ProtocolVersion::V1,
        generation_identity: &generation_identity,
        role,
        subject_id,
        public_key: &public_key,
    };
    let claims = encode_payload(&claims)?;
    let credential_id = ProtocolCredentialId::new(sha256(&claims));
    let signature = sign_payload(
        ProtocolRecordKind::ProtocolCredential,
        &claims,
        protocol_issuer,
    )?;
    let credential = match (role, subject_id) {
        (
            ProtocolCredentialRole::SessionCoordinator,
            ProtocolSubjectId::SessionCoordinator(subject_id),
        ) => ProtocolCredential::session_coordinator(
            credential_id,
            generation_identity,
            subject_id,
            public_key,
            signature.to_vec(),
        ),
        (ProtocolCredentialRole::AgentSandbox, ProtocolSubjectId::AgentSandbox(subject_id)) => {
            ProtocolCredential::agent_sandbox(
                credential_id,
                generation_identity,
                subject_id,
                public_key,
                signature.to_vec(),
            )
        }
        _ => return Err(ProtocolValidationError::RoleMismatch),
    };
    Ok(credential)
}

/// Validates one protocol credential against its issuer and exact authority tuple.
///
/// # Errors
///
/// Returns an error for key, role, subject, generation, digest, or signature mismatch.
pub fn verify_protocol_credential(
    credential: &ProtocolCredential,
    protocol_issuer: &VerifyingKey,
    expected_generation: &GenerationIdentity,
    expected_role: ProtocolCredentialRole,
    expected_subject: ProtocolSubjectId,
) -> Result<VerifyingKey, ProtocolValidationError> {
    if &credential.generation_identity != expected_generation {
        return Err(ProtocolValidationError::GenerationMismatch);
    }
    if credential.role() != expected_role {
        return Err(ProtocolValidationError::RoleMismatch);
    }
    if credential.subject_id() != expected_subject {
        return Err(ProtocolValidationError::SubjectMismatch);
    }
    let claims = ProtocolCredentialClaims {
        protocol_version: credential.protocol_version,
        generation_identity: &credential.generation_identity,
        role: credential.role(),
        subject_id: credential.subject_id(),
        public_key: &credential.public_key,
    };
    let claims = encode_payload(&claims)?;
    if credential.credential_id != ProtocolCredentialId::new(sha256(&claims)) {
        return Err(ProtocolValidationError::CredentialDigestMismatch);
    }
    verify_payload(
        ProtocolRecordKind::ProtocolCredential,
        &claims,
        &credential.issuer_signature,
        protocol_issuer,
    )?;
    let key_bytes: [u8; 32] = credential
        .public_key
        .as_slice()
        .try_into()
        .map_err(|_| ProtocolValidationError::InvalidPublicKey)?;
    VerifyingKey::from_bytes(&key_bytes).map_err(|_| ProtocolValidationError::InvalidPublicKey)
}

/// Issues one capability for an exact live session and credential pair.
///
/// # Errors
///
/// Returns an error when its claims cannot be encoded.
pub fn issue_session_capability(
    session_binding: SessionBinding,
    coordinator_credential_id: ProtocolCredentialId,
    sandbox_credential_id: ProtocolCredentialId,
    coordinator_key: &SigningKey,
) -> Result<SessionRequestCapability, ProtocolValidationError> {
    let claims = SessionCapabilityClaims {
        protocol_version: ProtocolVersion::V1,
        session_binding: &session_binding,
        coordinator_credential_id,
        sandbox_credential_id,
    };
    let claims = encode_payload(&claims)?;
    let capability_id = CapabilityId::new(sha256(&claims));
    let signature = sign_payload(
        ProtocolRecordKind::SessionRequestCapability,
        &claims,
        coordinator_key,
    )?;
    Ok(SessionRequestCapability {
        protocol_version: ProtocolVersion::V1,
        capability_id,
        session_binding,
        coordinator_credential_id,
        sandbox_credential_id,
        coordinator_signature: signature.to_vec(),
    })
}

fn verify_session_capability(
    capability: &SessionRequestCapability,
    session_binding: &SessionBinding,
    coordinator_credential_id: ProtocolCredentialId,
    sandbox_credential_id: ProtocolCredentialId,
    coordinator_key: &VerifyingKey,
) -> Result<(), ProtocolValidationError> {
    if &capability.session_binding != session_binding {
        return Err(ProtocolValidationError::SessionBindingMismatch);
    }
    if capability.coordinator_credential_id != coordinator_credential_id
        || capability.sandbox_credential_id != sandbox_credential_id
    {
        return Err(ProtocolValidationError::CredentialMismatch);
    }
    let claims = SessionCapabilityClaims {
        protocol_version: capability.protocol_version,
        session_binding: &capability.session_binding,
        coordinator_credential_id: capability.coordinator_credential_id,
        sandbox_credential_id: capability.sandbox_credential_id,
    };
    let claims = encode_payload(&claims)?;
    if capability.capability_id != CapabilityId::new(sha256(&claims)) {
        return Err(ProtocolValidationError::CapabilityDigestMismatch);
    }
    verify_payload(
        ProtocolRecordKind::SessionRequestCapability,
        &claims,
        &capability.coordinator_signature,
        coordinator_key,
    )?;
    Ok(())
}

/// Checks the certificate bytes and exact TLS identity tuple carried by a TLS credential.
///
/// # Errors
///
/// Returns an error for a stale tuple, role/subject mismatch, or certificate digest mismatch.
pub fn verify_tls_credential(
    credential: &TlsCredential,
    expected_generation: &GenerationIdentity,
    expected_role: TlsRole,
    expected_subject: TlsSubjectId,
) -> Result<(), ProtocolValidationError> {
    if &credential.generation_identity != expected_generation {
        return Err(ProtocolValidationError::GenerationMismatch);
    }
    if credential.role() != expected_role {
        return Err(ProtocolValidationError::RoleMismatch);
    }
    if credential.subject_id() != expected_subject {
        return Err(ProtocolValidationError::SubjectMismatch);
    }
    if credential.certificate_sha256 != CertificateDigest::new(sha256(&credential.certificate_der))
    {
        return Err(ProtocolValidationError::CertificateDigestMismatch);
    }
    Ok(())
}

#[derive(Clone, Debug)]
pub struct VerifiedSession {
    session_binding: SessionBinding,
    capability: SessionRequestCapability,
    sandbox_credential_id: ProtocolCredentialId,
    coordinator_credential_id: ProtocolCredentialId,
    sandbox_key: VerifyingKey,
    coordinator_key: VerifyingKey,
    session_anchor_key: VerifyingKey,
}

impl VerifiedSession {
    /// Constructs an authority only after validating its manifest, credentials, and capability.
    ///
    /// # Errors
    ///
    /// Returns an error when any key or identity tuple fails validation.
    pub fn new(
        trusted_manifest_digest: TrustManifestDigest,
        tls_ca_sha256: CertificateDigest,
        manifest: &TrustManifest,
        session_binding: SessionBinding,
        coordinator_credential: &ProtocolCredential,
        sandbox_credential: &ProtocolCredential,
        capability: SessionRequestCapability,
    ) -> Result<Self, ProtocolValidationError> {
        let generation = &session_binding.generation_identity;
        let issuer =
            verify_trust_manifest(manifest, trusted_manifest_digest, generation, tls_ca_sha256)?;
        let coordinator_key = verify_protocol_credential(
            coordinator_credential,
            &issuer,
            generation,
            ProtocolCredentialRole::SessionCoordinator,
            ProtocolSubjectId::SessionCoordinator(generation.coordinator_id),
        )?;
        let sandbox_key = verify_protocol_credential(
            sandbox_credential,
            &issuer,
            generation,
            ProtocolCredentialRole::AgentSandbox,
            ProtocolSubjectId::AgentSandbox(generation.sandbox_id),
        )?;
        verify_session_capability(
            &capability,
            &session_binding,
            coordinator_credential.credential_id,
            sandbox_credential.credential_id,
            &coordinator_key,
        )?;
        let anchor_bytes: [u8; 32] = session_binding
            .session_anchor_public_key
            .as_slice()
            .try_into()
            .map_err(|_| ProtocolValidationError::InvalidPublicKey)?;
        let session_anchor_key = VerifyingKey::from_bytes(&anchor_bytes)
            .map_err(|_| ProtocolValidationError::InvalidPublicKey)?;
        Ok(Self {
            session_binding,
            capability,
            sandbox_credential_id: sandbox_credential.credential_id,
            coordinator_credential_id: coordinator_credential.credential_id,
            sandbox_key,
            coordinator_key,
            session_anchor_key,
        })
    }

    fn verify_sandbox_tuple(
        &self,
        session_binding: &SessionBinding,
        capability: &SessionRequestCapability,
        credential_id: ProtocolCredentialId,
    ) -> Result<(), ProtocolValidationError> {
        if session_binding != &self.session_binding {
            return Err(ProtocolValidationError::SessionBindingMismatch);
        }
        if capability != &self.capability {
            return Err(ProtocolValidationError::CapabilityMismatch);
        }
        if credential_id != self.sandbox_credential_id {
            return Err(ProtocolValidationError::CredentialMismatch);
        }
        Ok(())
    }

    /// Verifies an exact signed request and returns its digest over accepted payload bytes.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, or tuple mismatch.
    pub fn verify_request(
        &self,
        encoded: &[u8],
    ) -> Result<(MailboxRequest, RequestDigest), ProtocolValidationError> {
        let (request, payload_bytes) = decode_signed::<MailboxRequest>(encoded, &self.sandbox_key)?;
        self.verify_sandbox_tuple(
            request.session_binding(),
            request.session_capability(),
            request.sandbox_credential_id(),
        )?;
        Ok((request, request_digest(&payload_bytes)))
    }

    /// Verifies a signed cancellation for the exact request.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, tuple, or request mismatch.
    pub fn verify_cancellation(
        &self,
        encoded: &[u8],
        request_id: RequestId,
    ) -> Result<MailboxCancellation, ProtocolValidationError> {
        let (cancellation, _) = decode_signed::<MailboxCancellation>(encoded, &self.sandbox_key)?;
        self.verify_sandbox_tuple(
            &cancellation.session_binding,
            &cancellation.session_capability,
            cancellation.sandbox_credential_id,
        )?;
        if cancellation.request_id != request_id {
            return Err(ProtocolValidationError::RequestMismatch);
        }
        Ok(cancellation)
    }

    /// Verifies a signed highest-contiguous acknowledgement for the exact request.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, tuple, or request mismatch.
    pub fn verify_result_ack(
        &self,
        encoded: &[u8],
        request_id: RequestId,
        digest: RequestDigest,
    ) -> Result<MailboxResultAck, ProtocolValidationError> {
        let (ack, _) = decode_signed::<MailboxResultAck>(encoded, &self.sandbox_key)?;
        self.verify_sandbox_tuple(
            &ack.session_binding,
            &ack.session_capability,
            ack.sandbox_credential_id,
        )?;
        if ack.request_id != request_id || ack.request_sha256 != digest {
            return Err(ProtocolValidationError::RequestMismatch);
        }
        Ok(ack)
    }

    /// Verifies one host-only Session Anchor decision for the exact request.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, tuple, or request mismatch.
    pub fn verify_developer_decision(
        &self,
        encoded: &[u8],
        request_id: RequestId,
        digest: RequestDigest,
    ) -> Result<DeveloperDecision, ProtocolValidationError> {
        let (decision, _) = decode_signed::<DeveloperDecision>(encoded, &self.session_anchor_key)?;
        if decision.session_binding != self.session_binding {
            return Err(ProtocolValidationError::SessionBindingMismatch);
        }
        if decision.request_id != request_id || decision.request_sha256 != digest {
            return Err(ProtocolValidationError::RequestMismatch);
        }
        Ok(decision)
    }

    /// Verifies one coordinator result chunk for the exact request.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, tuple, request, or size mismatch.
    pub fn verify_result_chunk(
        &self,
        encoded: &[u8],
        request_id: RequestId,
        digest: RequestDigest,
    ) -> Result<MailboxResultChunk, ProtocolValidationError> {
        let (chunk, _) = decode_signed::<MailboxResultChunk>(encoded, &self.coordinator_key)?;
        if chunk.bytes.len() > MAX_RESULT_CHUNK_BYTES {
            return Err(ProtocolValidationError::ChunkTooLarge {
                length: chunk.bytes.len(),
            });
        }
        self.verify_coordinator_tuple(
            &chunk.session_binding,
            chunk.request_id,
            chunk.request_sha256,
            chunk.coordinator_credential_id,
            request_id,
            digest,
        )?;
        Ok(chunk)
    }

    /// Verifies one coordinator terminal result for the exact request.
    ///
    /// # Errors
    ///
    /// Returns an error for signature, encoding, tuple, or request mismatch.
    pub fn verify_result(
        &self,
        encoded: &[u8],
        request_id: RequestId,
        digest: RequestDigest,
    ) -> Result<MailboxResult, ProtocolValidationError> {
        let (result, _) = decode_signed::<MailboxResult>(encoded, &self.coordinator_key)?;
        self.verify_coordinator_tuple(
            &result.session_binding,
            result.request_id,
            result.request_sha256,
            result.coordinator_credential_id,
            request_id,
            digest,
        )?;
        Ok(result)
    }

    fn verify_coordinator_tuple(
        &self,
        session_binding: &SessionBinding,
        actual_request_id: RequestId,
        actual_digest: RequestDigest,
        credential_id: ProtocolCredentialId,
        expected_request_id: RequestId,
        expected_digest: RequestDigest,
    ) -> Result<(), ProtocolValidationError> {
        if session_binding != &self.session_binding {
            return Err(ProtocolValidationError::SessionBindingMismatch);
        }
        if credential_id != self.coordinator_credential_id {
            return Err(ProtocolValidationError::CredentialMismatch);
        }
        if actual_request_id != expected_request_id || actual_digest != expected_digest {
            return Err(ProtocolValidationError::RequestMismatch);
        }
        Ok(())
    }
}

#[derive(Debug, Default)]
pub struct RequestReplayGuard {
    admitted: BTreeSet<RequestId>,
}

impl RequestReplayGuard {
    /// Verifies and records a request ID exactly once for this generation.
    ///
    /// # Errors
    ///
    /// Returns an error when validation fails or the request ID was previously admitted.
    pub fn admit(
        &mut self,
        session: &VerifiedSession,
        encoded: &[u8],
    ) -> Result<(MailboxRequest, RequestDigest), ProtocolValidationError> {
        let (request, digest) = session.verify_request(encoded)?;
        if !self.admitted.insert(request.request_id()) {
            return Err(ProtocolValidationError::RequestReplay);
        }
        Ok((request, digest))
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.admitted.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.admitted.is_empty()
    }
}

const _: () = assert!(MAX_UNACKNOWLEDGED_CHUNKS == 8);