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

use std::fmt;

use ed25519_dalek::{Signature, Signer, SigningKey, VerifyingKey};
use serde::Serialize;
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};

use crate::identity::{RequestDigest, StreamDigest, TrustManifestDigest};
use crate::records::{
    DeveloperDecision, MailboxCancellation, MailboxRequest, MailboxResult, MailboxResultAck,
    MailboxResultChunk, SignedProtocolRecord,
};

pub const MAX_SIGNED_RECORD_BYTES: usize = 65_536;
pub const MAX_RESULT_CHUNK_BYTES: usize = 49_152;
const SIGNATURE_BYTES: usize = 64;
const LENGTH_BYTES: usize = size_of::<u64>();

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ProtocolRecordKind {
    TrustManifest,
    ProtocolCredential,
    SessionRequestCapability,
    MailboxRequest,
    MailboxCancellation,
    MailboxResultAck,
    DeveloperDecision,
    MailboxResult,
    MailboxResultChunk,
}

impl ProtocolRecordKind {
    #[must_use]
    pub const fn signing_context(self) -> &'static [u8] {
        match self {
            Self::TrustManifest => b"lazaret/v1/trust-manifest",
            Self::ProtocolCredential => b"lazaret/v1/protocol-credential",
            Self::SessionRequestCapability => b"lazaret/v1/session-request-capability",
            Self::MailboxRequest => b"lazaret/v1/mailbox-request",
            Self::MailboxCancellation => b"lazaret/v1/mailbox-cancellation",
            Self::MailboxResultAck => b"lazaret/v1/mailbox-result-ack",
            Self::DeveloperDecision => b"lazaret/v1/developer-decision",
            Self::MailboxResult => b"lazaret/v1/mailbox-result",
            Self::MailboxResultChunk => b"lazaret/v1/mailbox-result-chunk",
        }
    }
}

pub trait ProtocolPayload: DeserializeOwned + Serialize {
    const RECORD_KIND: ProtocolRecordKind;
}

macro_rules! protocol_payload {
    ($($type:ty => $kind:ident),+ $(,)?) => {
        $(
            impl ProtocolPayload for $type {
                const RECORD_KIND: ProtocolRecordKind = ProtocolRecordKind::$kind;
            }
        )+
    };
}

protocol_payload!(
    MailboxRequest => MailboxRequest,
    MailboxCancellation => MailboxCancellation,
    MailboxResultAck => MailboxResultAck,
    DeveloperDecision => DeveloperDecision,
    MailboxResult => MailboxResult,
    MailboxResultChunk => MailboxResultChunk,
);

#[derive(Debug)]
pub enum ProtocolCodecError {
    RecordTooLarge { length: usize },
    LengthOverflow,
    Malformed(postcard::Error),
    TrailingBytes,
    NonMinimalEncoding,
    InvalidSignatureLength { length: usize },
    InvalidSignature,
}

impl fmt::Display for ProtocolCodecError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::RecordTooLarge { length } => write!(
                formatter,
                "signed protocol record length {length} exceeds {MAX_SIGNED_RECORD_BYTES}"
            ),
            Self::LengthOverflow => formatter.write_str("protocol payload length exceeds u64"),
            Self::Malformed(error) => write!(formatter, "malformed protocol record: {error}"),
            Self::TrailingBytes => formatter.write_str("protocol record has trailing bytes"),
            Self::NonMinimalEncoding => {
                formatter.write_str("protocol record does not use minimal Postcard encoding")
            }
            Self::InvalidSignatureLength { length } => {
                write!(
                    formatter,
                    "protocol signature length is {length}, expected 64"
                )
            }
            Self::InvalidSignature => formatter.write_str("protocol signature is invalid"),
        }
    }
}

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

pub(crate) fn encode_payload<T: Serialize>(payload: &T) -> Result<Vec<u8>, ProtocolCodecError> {
    postcard::to_stdvec(payload).map_err(ProtocolCodecError::Malformed)
}

pub(crate) fn decode_payload<T>(bytes: &[u8]) -> Result<T, ProtocolCodecError>
where
    T: DeserializeOwned + Serialize,
{
    let (value, remainder) =
        postcard::take_from_bytes(bytes).map_err(ProtocolCodecError::Malformed)?;
    if !remainder.is_empty() {
        return Err(ProtocolCodecError::TrailingBytes);
    }
    if encode_payload(&value)? != bytes {
        return Err(ProtocolCodecError::NonMinimalEncoding);
    }
    Ok(value)
}

fn signing_message(
    kind: ProtocolRecordKind,
    payload: &[u8],
) -> Result<Vec<u8>, ProtocolCodecError> {
    let length = u64::try_from(payload.len()).map_err(|_| ProtocolCodecError::LengthOverflow)?;
    let context = kind.signing_context();
    let mut message = Vec::with_capacity(context.len() + LENGTH_BYTES + payload.len());
    message.extend_from_slice(context);
    message.extend_from_slice(&length.to_be_bytes());
    message.extend_from_slice(payload);
    Ok(message)
}

pub(crate) fn sign_payload(
    kind: ProtocolRecordKind,
    payload: &[u8],
    signing_key: &SigningKey,
) -> Result<[u8; SIGNATURE_BYTES], ProtocolCodecError> {
    Ok(signing_key
        .sign(&signing_message(kind, payload)?)
        .to_bytes())
}

pub(crate) fn verify_payload(
    kind: ProtocolRecordKind,
    payload: &[u8],
    signature: &[u8],
    verifying_key: &VerifyingKey,
) -> Result<(), ProtocolCodecError> {
    let signature = Signature::from_slice(signature).map_err(|_| {
        ProtocolCodecError::InvalidSignatureLength {
            length: signature.len(),
        }
    })?;
    verifying_key
        .verify_strict(&signing_message(kind, payload)?, &signature)
        .map_err(|_| ProtocolCodecError::InvalidSignature)
}

/// Signs the minimal Postcard payload and returns one bounded outer record.
///
/// # Errors
///
/// Returns an error when serialization fails or the complete record exceeds 65,536 bytes.
pub fn encode_signed<T: ProtocolPayload>(
    payload: &T,
    signing_key: &SigningKey,
) -> Result<Vec<u8>, ProtocolCodecError> {
    let payload_bytes = encode_payload(payload)?;
    let signature = sign_payload(T::RECORD_KIND, &payload_bytes, signing_key)?;
    let record = SignedProtocolRecord {
        payload_bytes,
        signature: signature.to_vec(),
    };
    let encoded = encode_payload(&record)?;
    if encoded.len() > MAX_SIGNED_RECORD_BYTES {
        return Err(ProtocolCodecError::RecordTooLarge {
            length: encoded.len(),
        });
    }
    Ok(encoded)
}

pub(crate) fn decode_signed<T>(
    encoded: &[u8],
    verifying_key: &VerifyingKey,
) -> Result<(T, Vec<u8>), ProtocolCodecError>
where
    T: ProtocolPayload,
{
    if encoded.len() > MAX_SIGNED_RECORD_BYTES {
        return Err(ProtocolCodecError::RecordTooLarge {
            length: encoded.len(),
        });
    }
    let record: SignedProtocolRecord = decode_payload(encoded)?;
    verify_payload(
        T::RECORD_KIND,
        &record.payload_bytes,
        &record.signature,
        verifying_key,
    )?;
    let payload = decode_payload(&record.payload_bytes)?;
    Ok((payload, record.payload_bytes))
}

#[must_use]
pub(crate) fn sha256(bytes: &[u8]) -> [u8; 32] {
    Sha256::digest(bytes).into()
}

#[must_use]
pub fn request_digest(payload_bytes: &[u8]) -> RequestDigest {
    RequestDigest::new(sha256(payload_bytes))
}

#[must_use]
pub fn stream_digest(bytes: &[u8]) -> StreamDigest {
    StreamDigest::new(sha256(bytes))
}

#[must_use]
pub fn trust_manifest_digest(encoded_manifest: &[u8]) -> TrustManifestDigest {
    TrustManifestDigest::new(sha256(encoded_manifest))
}

#[cfg(test)]
mod tests {
    use ed25519_dalek::SigningKey;
    use serde::Serialize;

    use crate::identity::{
        ClientId, ConfigurationFingerprint, CoordinatorId, DomainId, GenerationId,
        ProtocolCredentialId, RequestDigest, RequestId, SandboxId, SessionId,
    };
    use crate::records::{
        GenerationIdentity, MailboxResultChunk, ResultStream, SessionBinding, SignedProtocolRecord,
    };

    use super::{
        ProtocolCodecError, ProtocolPayload, ProtocolRecordKind, decode_signed, encode_payload,
        sign_payload,
    };

    #[test]
    fn signing_context_and_length_have_an_exact_vector() {
        let key = SigningKey::from_bytes(&[7; 32]);
        let signature = sign_payload(ProtocolRecordKind::MailboxRequest, &[0], &key).unwrap();
        assert_eq!(
            signature,
            [
                79, 132, 54, 199, 95, 126, 134, 255, 188, 61, 184, 49, 116, 41, 8, 91, 38, 77, 75,
                188, 184, 7, 164, 133, 194, 37, 215, 228, 53, 141, 99, 203, 254, 220, 210, 170,
                157, 52, 147, 199, 116, 247, 230, 77, 128, 126, 113, 16, 115, 207, 134, 21, 128,
                210, 193, 102, 45, 148, 252, 17, 87, 57, 218, 11,
            ]
        );
        let record = encode_payload(&SignedProtocolRecord {
            payload_bytes: vec![0],
            signature: signature.to_vec(),
        })
        .unwrap();
        assert_eq!(
            record,
            [
                1, 0, 64, 79, 132, 54, 199, 95, 126, 134, 255, 188, 61, 184, 49, 116, 41, 8, 91,
                38, 77, 75, 188, 184, 7, 164, 133, 194, 37, 215, 228, 53, 141, 99, 203, 254, 220,
                210, 170, 157, 52, 147, 199, 116, 247, 230, 77, 128, 126, 113, 16, 115, 207, 134,
                21, 128, 210, 193, 102, 45, 148, 252, 17, 87, 57, 218, 11,
            ]
        );
    }

    #[test]
    fn valid_signature_does_not_authorize_an_unknown_payload_version() {
        #[derive(Serialize)]
        struct UnknownVersionChunk {
            protocol_version: u8,
            session_binding: SessionBinding,
            request_id: RequestId,
            request_sha256: RequestDigest,
            coordinator_credential_id: ProtocolCredentialId,
            sequence: u64,
            stream: ResultStream,
            bytes: Vec<u8>,
        }

        let opaque = [1; 32];
        let generation_identity = GenerationIdentity {
            agent_profile: "agent".into(),
            domain_id: DomainId::new(opaque),
            generation_id: GenerationId::new(1),
            configuration_fingerprint: ConfigurationFingerprint::new(opaque),
            coordinator_id: CoordinatorId::new(opaque),
            sandbox_id: SandboxId::new(opaque),
        };
        let payload = encode_payload(&UnknownVersionChunk {
            protocol_version: 1,
            session_binding: SessionBinding {
                generation_identity,
                cli_id: ClientId::new(opaque),
                session_id: SessionId::new(opaque),
                session_anchor_public_key: vec![1; 32],
            },
            request_id: RequestId::new(opaque),
            request_sha256: RequestDigest::new(opaque),
            coordinator_credential_id: ProtocolCredentialId::new(opaque),
            sequence: 0,
            stream: ResultStream::Stdout,
            bytes: Vec::new(),
        })
        .unwrap();
        let key = SigningKey::from_bytes(&[8; 32]);
        let signature = sign_payload(MailboxResultChunk::RECORD_KIND, &payload, &key).unwrap();
        let encoded = encode_payload(&SignedProtocolRecord {
            payload_bytes: payload,
            signature: signature.to_vec(),
        })
        .unwrap();
        assert!(matches!(
            decode_signed::<MailboxResultChunk>(&encoded, &key.verifying_key()),
            Err(ProtocolCodecError::Malformed(_))
        ));
    }
}