lazaret/crates/lazaret-core/src/records/credentials.rs

use super::{
    CapabilityId, CertificateDigest, ClientId, ConfigurationFingerprint, CoordinatorId,
    Deserialize, DomainId, GenerationId, ProtocolCredentialId, RequestDigest, RequestId, SandboxId,
    Serialize, SessionId, StreamDigest, TrustManifestDigest, fmt,
};

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ProtocolVersion {
    V1,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecordInvariantError {
    CredentialBundleRoles,
    ParticipantSnapshotState,
    SessionTerminalSnapshotState,
}

impl fmt::Display for RecordInvariantError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::CredentialBundleRoles => {
                formatter.write_str("credential bundle roles are invalid")
            }
            Self::ParticipantSnapshotState => {
                formatter.write_str("participant ownership snapshot is invalid for its role/state")
            }
            Self::SessionTerminalSnapshotState => {
                formatter.write_str("Session Terminal evidence is invalid for its state")
            }
        }
    }
}

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

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecordTransitionError {
    RecordIdentityChanged,
    StateRegression,
    EvidenceChanged,
    EvidenceLost,
    ProcessStageSkipped,
    PriorAcquisitionSnapshotRequired,
    AbsenceProofRequired,
    InvalidCleanTarget,
}

impl fmt::Display for RecordTransitionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::RecordIdentityChanged => formatter.write_str("record identity changed"),
            Self::StateRegression => formatter.write_str("lifecycle transition is invalid"),
            Self::EvidenceChanged => formatter.write_str("recorded evidence changed"),
            Self::EvidenceLost => formatter.write_str("recorded evidence was cleared"),
            Self::ProcessStageSkipped => {
                formatter.write_str("process ownership did not follow the staged barrier")
            }
            Self::PriorAcquisitionSnapshotRequired => {
                formatter.write_str("active state requires a prior acquisition snapshot")
            }
            Self::AbsenceProofRequired => formatter.write_str("cleanup requires absence proof"),
            Self::InvalidCleanTarget => formatter.write_str("Clean target is not fully cleared"),
        }
    }
}

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

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct GenerationIdentity {
    pub agent_profile: String,
    pub domain_id: DomainId,
    pub generation_id: GenerationId,
    pub configuration_fingerprint: ConfigurationFingerprint,
    pub coordinator_id: CoordinatorId,
    pub sandbox_id: SandboxId,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct TrustManifest {
    pub protocol_version: ProtocolVersion,
    pub generation_identity: GenerationIdentity,
    pub tls_ca_sha256: CertificateDigest,
    pub protocol_issuer_public_key: Vec<u8>,
    pub issuer_signature: Vec<u8>,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum TlsRole {
    SessionCoordinator,
    ClientCli,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum TlsSubjectId {
    SessionCoordinator(CoordinatorId),
    ClientCli(ClientId),
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct TlsCredential {
    pub generation_identity: GenerationIdentity,
    role: TlsRole,
    subject_id: TlsSubjectId,
    pub certificate_der: Vec<u8>,
    pub certificate_sha256: CertificateDigest,
}

impl TlsCredential {
    #[must_use]
    pub fn session_coordinator(
        generation_identity: GenerationIdentity,
        subject_id: CoordinatorId,
        certificate_der: Vec<u8>,
        certificate_sha256: CertificateDigest,
    ) -> Self {
        Self {
            generation_identity,
            role: TlsRole::SessionCoordinator,
            subject_id: TlsSubjectId::SessionCoordinator(subject_id),
            certificate_der,
            certificate_sha256,
        }
    }

    #[must_use]
    pub fn client_cli(
        generation_identity: GenerationIdentity,
        subject_id: ClientId,
        certificate_der: Vec<u8>,
        certificate_sha256: CertificateDigest,
    ) -> Self {
        Self {
            generation_identity,
            role: TlsRole::ClientCli,
            subject_id: TlsSubjectId::ClientCli(subject_id),
            certificate_der,
            certificate_sha256,
        }
    }

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

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

impl<'de> Deserialize<'de> for TlsCredential {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Fields {
            generation_identity: GenerationIdentity,
            role: TlsRole,
            subject_id: TlsSubjectId,
            certificate_der: Vec<u8>,
            certificate_sha256: CertificateDigest,
        }

        let fields = Fields::deserialize(deserializer)?;
        match (fields.role, fields.subject_id) {
            (TlsRole::SessionCoordinator, TlsSubjectId::SessionCoordinator(subject_id)) => {
                Ok(Self::session_coordinator(
                    fields.generation_identity,
                    subject_id,
                    fields.certificate_der,
                    fields.certificate_sha256,
                ))
            }
            (TlsRole::ClientCli, TlsSubjectId::ClientCli(subject_id)) => Ok(Self::client_cli(
                fields.generation_identity,
                subject_id,
                fields.certificate_der,
                fields.certificate_sha256,
            )),
            _ => Err(serde::de::Error::custom(
                "TLS role does not match its subject identity",
            )),
        }
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ProtocolPublicationIdentity {
    pub protocol_version: ProtocolVersion,
    pub generation_identity: GenerationIdentity,
    pub tls_ca_certificate_der: Vec<u8>,
    pub trust_manifest_sha256: TrustManifestDigest,
    pub coordinator_tls_certificate_sha256: CertificateDigest,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ClientEnrollmentRequest {
    pub protocol_version: ProtocolVersion,
    pub generation_identity: GenerationIdentity,
    pub cli_id: ClientId,
    pub certificate_signing_request_der: Vec<u8>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ClientEnrollmentResponse {
    pub protocol_version: ProtocolVersion,
    pub credential: TlsCredential,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SessionRegistrationRequest {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SessionRegistrationResponse {
    pub protocol_version: ProtocolVersion,
    pub session_capability: SessionRequestCapability,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ControlRequest {
    Enroll(ClientEnrollmentRequest),
    Register(SessionRegistrationRequest),
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ControlResponse {
    Enrolled(ClientEnrollmentResponse),
    Registered(SessionRegistrationResponse),
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ProtocolCredentialRole {
    SessionCoordinator,
    AgentSandbox,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ProtocolSubjectId {
    SessionCoordinator(CoordinatorId),
    AgentSandbox(SandboxId),
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct ProtocolCredential {
    pub protocol_version: ProtocolVersion,
    pub credential_id: ProtocolCredentialId,
    pub generation_identity: GenerationIdentity,
    role: ProtocolCredentialRole,
    subject_id: ProtocolSubjectId,
    pub public_key: Vec<u8>,
    pub issuer_signature: Vec<u8>,
}

impl ProtocolCredential {
    #[must_use]
    pub fn session_coordinator(
        credential_id: ProtocolCredentialId,
        generation_identity: GenerationIdentity,
        subject_id: CoordinatorId,
        public_key: Vec<u8>,
        issuer_signature: Vec<u8>,
    ) -> Self {
        Self {
            protocol_version: ProtocolVersion::V1,
            credential_id,
            generation_identity,
            role: ProtocolCredentialRole::SessionCoordinator,
            subject_id: ProtocolSubjectId::SessionCoordinator(subject_id),
            public_key,
            issuer_signature,
        }
    }

    #[must_use]
    pub fn agent_sandbox(
        credential_id: ProtocolCredentialId,
        generation_identity: GenerationIdentity,
        subject_id: SandboxId,
        public_key: Vec<u8>,
        issuer_signature: Vec<u8>,
    ) -> Self {
        Self {
            protocol_version: ProtocolVersion::V1,
            credential_id,
            generation_identity,
            role: ProtocolCredentialRole::AgentSandbox,
            subject_id: ProtocolSubjectId::AgentSandbox(subject_id),
            public_key,
            issuer_signature,
        }
    }

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

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

impl<'de> Deserialize<'de> for ProtocolCredential {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Fields {
            protocol_version: ProtocolVersion,
            credential_id: ProtocolCredentialId,
            generation_identity: GenerationIdentity,
            role: ProtocolCredentialRole,
            subject_id: ProtocolSubjectId,
            public_key: Vec<u8>,
            issuer_signature: Vec<u8>,
        }

        let fields = Fields::deserialize(deserializer)?;
        match fields.protocol_version {
            ProtocolVersion::V1 => {}
        }
        match (fields.role, fields.subject_id) {
            (
                ProtocolCredentialRole::SessionCoordinator,
                ProtocolSubjectId::SessionCoordinator(subject_id),
            ) => Ok(Self::session_coordinator(
                fields.credential_id,
                fields.generation_identity,
                subject_id,
                fields.public_key,
                fields.issuer_signature,
            )),
            (ProtocolCredentialRole::AgentSandbox, ProtocolSubjectId::AgentSandbox(subject_id)) => {
                Ok(Self::agent_sandbox(
                    fields.credential_id,
                    fields.generation_identity,
                    subject_id,
                    fields.public_key,
                    fields.issuer_signature,
                ))
            }
            _ => Err(serde::de::Error::custom(
                "protocol credential role does not match its subject identity",
            )),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct SandboxCredentialBundle {
    pub protocol_version: ProtocolVersion,
    pub generation_identity: GenerationIdentity,
    pub trust_manifest: TrustManifest,
    coordinator_protocol_credential: ProtocolCredential,
    sandbox_protocol_credential: ProtocolCredential,
    pub sandbox_private_key: Vec<u8>,
}

impl SandboxCredentialBundle {
    /// Constructs a bundle only when its two credentials have their required roles.
    ///
    /// # Errors
    ///
    /// Returns an error when either credential has the wrong role.
    pub fn new(
        generation_identity: GenerationIdentity,
        trust_manifest: TrustManifest,
        coordinator_protocol_credential: ProtocolCredential,
        sandbox_protocol_credential: ProtocolCredential,
        sandbox_private_key: Vec<u8>,
    ) -> Result<Self, RecordInvariantError> {
        if coordinator_protocol_credential.role() != ProtocolCredentialRole::SessionCoordinator
            || sandbox_protocol_credential.role() != ProtocolCredentialRole::AgentSandbox
        {
            return Err(RecordInvariantError::CredentialBundleRoles);
        }
        Ok(Self {
            protocol_version: ProtocolVersion::V1,
            generation_identity,
            trust_manifest,
            coordinator_protocol_credential,
            sandbox_protocol_credential,
            sandbox_private_key,
        })
    }

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

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

impl<'de> Deserialize<'de> for SandboxCredentialBundle {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Fields {
            protocol_version: ProtocolVersion,
            generation_identity: GenerationIdentity,
            trust_manifest: TrustManifest,
            coordinator_protocol_credential: ProtocolCredential,
            sandbox_protocol_credential: ProtocolCredential,
            sandbox_private_key: Vec<u8>,
        }

        let fields = Fields::deserialize(deserializer)?;
        match fields.protocol_version {
            ProtocolVersion::V1 => {}
        }
        Self::new(
            fields.generation_identity,
            fields.trust_manifest,
            fields.coordinator_protocol_credential,
            fields.sandbox_protocol_credential,
            fields.sandbox_private_key,
        )
        .map_err(serde::de::Error::custom)
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SessionBinding {
    pub generation_identity: GenerationIdentity,
    pub cli_id: ClientId,
    pub session_id: SessionId,
    pub session_anchor_public_key: Vec<u8>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SessionRequestCapability {
    pub protocol_version: ProtocolVersion,
    pub capability_id: CapabilityId,
    pub session_binding: SessionBinding,
    pub coordinator_credential_id: ProtocolCredentialId,
    pub sandbox_credential_id: ProtocolCredentialId,
    pub coordinator_signature: Vec<u8>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SignedProtocolRecord {
    pub payload_bytes: Vec<u8>,
    pub signature: Vec<u8>,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct MailboxRequest {
    protocol_version: ProtocolVersion,
    session_binding: SessionBinding,
    request_id: RequestId,
    session_capability: SessionRequestCapability,
    sandbox_credential_id: ProtocolCredentialId,
    command_argv: Vec<Vec<u8>>,
    justification: String,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum MailboxRequestError {
    EmptyCommand,
    InvalidArgument,
    NonAbsoluteExecutable,
}

impl fmt::Display for MailboxRequestError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptyCommand => formatter.write_str("command argument vector is empty"),
            Self::InvalidArgument => formatter.write_str("command argument is not NUL-free UTF-8"),
            Self::NonAbsoluteExecutable => {
                formatter.write_str("command executable is not an absolute path")
            }
        }
    }
}

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

impl MailboxRequest {
    /// # Errors
    /// The command must contain an absolute executable followed by NUL-free UTF-8 arguments.
    pub fn new(
        session_binding: SessionBinding,
        request_id: RequestId,
        session_capability: SessionRequestCapability,
        sandbox_credential_id: ProtocolCredentialId,
        command_argv: Vec<Vec<u8>>,
        justification: String,
    ) -> Result<Self, MailboxRequestError> {
        let Some(executable) = command_argv.first() else {
            return Err(MailboxRequestError::EmptyCommand);
        };
        if command_argv
            .iter()
            .any(|argument| argument.contains(&0) || std::str::from_utf8(argument).is_err())
        {
            return Err(MailboxRequestError::InvalidArgument);
        }
        if !executable.starts_with(b"/") {
            return Err(MailboxRequestError::NonAbsoluteExecutable);
        }
        Ok(Self {
            protocol_version: ProtocolVersion::V1,
            session_binding,
            request_id,
            session_capability,
            sandbox_credential_id,
            command_argv,
            justification,
        })
    }

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

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

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

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

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

    #[must_use]
    pub fn command_argv(&self) -> &[Vec<u8>] {
        &self.command_argv
    }

    #[must_use]
    pub fn justification(&self) -> &str {
        &self.justification
    }
}

impl<'de> Deserialize<'de> for MailboxRequest {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct Fields {
            protocol_version: ProtocolVersion,
            session_binding: SessionBinding,
            request_id: RequestId,
            session_capability: SessionRequestCapability,
            sandbox_credential_id: ProtocolCredentialId,
            command_argv: Vec<Vec<u8>>,
            justification: String,
        }

        let fields = Fields::deserialize(deserializer)?;
        if fields.protocol_version != ProtocolVersion::V1 {
            return Err(serde::de::Error::custom("unknown mailbox request version"));
        }
        Self::new(
            fields.session_binding,
            fields.request_id,
            fields.session_capability,
            fields.sandbox_credential_id,
            fields.command_argv,
            fields.justification,
        )
        .map_err(serde::de::Error::custom)
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MailboxCancellation {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
    pub request_id: RequestId,
    pub session_capability: SessionRequestCapability,
    pub sandbox_credential_id: ProtocolCredentialId,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MailboxResultAck {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
    pub request_id: RequestId,
    pub request_sha256: RequestDigest,
    pub session_capability: SessionRequestCapability,
    pub sandbox_credential_id: ProtocolCredentialId,
    pub highest_contiguous_sequence: u64,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum DeveloperDecisionKind {
    Approve,
    Deny,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct DeveloperDecision {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
    pub request_id: RequestId,
    pub request_sha256: RequestDigest,
    pub decision: DeveloperDecisionKind,
    pub approved_command_argv: Vec<Vec<u8>>,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum EscalationOutcome {
    Completed,
    Denied,
    TimedOut,
    Cancelled,
    Failed,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MailboxResult {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
    pub request_id: RequestId,
    pub request_sha256: RequestDigest,
    pub coordinator_credential_id: ProtocolCredentialId,
    pub outcome: EscalationOutcome,
    pub exit_code: Option<i32>,
    pub chunk_count: u64,
    pub stdout_sha256: StreamDigest,
    pub stderr_sha256: StreamDigest,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum ResultStream {
    Stdout,
    Stderr,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MailboxResultChunk {
    pub protocol_version: ProtocolVersion,
    pub session_binding: SessionBinding,
    pub request_id: RequestId,
    pub request_sha256: RequestDigest,
    pub coordinator_credential_id: ProtocolCredentialId,
    pub sequence: u64,
    pub stream: ResultStream,
    pub bytes: Vec<u8>,
}