Implement generation-bound protocol cryptography

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

@@ -0,0 +1,339 @@
+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(_))
+        ));
+    }
+}