laser-sdk 0.8.1

use crate::error::LaserError;
use async_trait::async_trait;
use laser_wire::agent::BodyRef;
use laser_wire::content::ContentType;
use laser_wire::framing::{decode_named, encode_named};
use sha2::{Digest, Sha256};

/// Where claim-checked bodies live. The seam rule applies: no default store
/// ships, the deployment chooses (an S3-compatible bucket, the kv surface for
/// small overflow, a filesystem in dev), and the same trait shape as
/// [`Deduplicator`](crate::agent::Deduplicator) keeps it `dyn`-usable behind
/// a plain reference.
#[async_trait]
pub trait BlobStore: Send + Sync {
    /// Store `payload` and return the reference a [`BodyRef`] will carry
    /// (a URI, object key, or kv key, bounded by the capsule's cap).
    async fn put(&self, payload: Vec<u8>) -> Result<String, LaserError>;
    /// Fetch the bytes behind `reference`.
    async fn get(&self, reference: &str) -> Result<Vec<u8>, LaserError>;
}

/// Claim-check `payload` against `store` when it is at or over
/// `threshold_bytes`: the bytes are `put`, hashed, and replaced by the encoded
/// [`BodyRef`] capsule with content-type `ref`. Under the threshold the
/// payload comes back untouched, byte-identical to no claim check at all.
/// Returns the (possibly replaced) payload and the content type to stamp when
/// the check fired.
pub async fn check_in(
    store: &dyn BlobStore,
    threshold_bytes: usize,
    payload: Vec<u8>,
) -> Result<(Vec<u8>, Option<ContentType>), LaserError> {
    if payload.len() < threshold_bytes {
        return Ok((payload, None));
    }
    let sha256: [u8; 32] = Sha256::digest(&payload).into();
    let size_bytes = payload.len() as u64;
    let reference = store.put(payload).await?;
    let capsule = BodyRef::new(reference, size_bytes, sha256);
    let encoded = encode_named(&capsule)
        .map_err(|error| LaserError::Codec(format!("encode body ref: {error}")))?;
    Ok((encoded, Some(ContentType::Ref)))
}

/// Resolve a claim-checked body: decode the [`BodyRef`] capsule from
/// `payload`, fetch the referenced bytes, and verify their SHA-256 against the
/// capsule's digest. A mismatch is the typed integrity refusal, never
/// unverified payload: the capsule's digest is the reader's proof the store
/// returned what the producer externalized.
pub async fn resolve_body(store: &dyn BlobStore, payload: &[u8]) -> Result<Vec<u8>, LaserError> {
    let capsule: BodyRef = decode_named(payload)
        .map_err(|error| LaserError::Codec(format!("decode body ref: {error}")))?;
    // The capsule arrives off the log, so its reference is untrusted input to
    // whatever store the caller supplied. Enforce the wire caps before handing
    // it over.
    capsule
        .validate()
        .map_err(|error| LaserError::Invalid(format!("body ref: {error}")))?;
    let payload = store.get(&capsule.reference).await?;
    // The capsule states the size it externalized. A length that disagrees is
    // already a failed claim check, caught before hashing the bytes.
    if payload.len() as u64 != capsule.size_bytes {
        return Err(LaserError::Integrity {
            reference: capsule.reference,
        });
    }
    let digest: [u8; 32] = Sha256::digest(&payload).into();
    if digest.as_slice() != capsule.sha256.as_slice() {
        return Err(LaserError::Integrity {
            reference: capsule.reference,
        });
    }
    Ok(payload)
}

impl crate::agent::AgentMessage {
    /// The real body of a claim-checked message. It starts from
    /// [`body`](Self::body): the AGDX envelope's body for an AGDX record, the
    /// payload otherwise. When the content type is `ref`, that body is the
    /// capsule, which is decoded, fetched, and digest-verified through
    /// [`resolve_body`]. Any other content type returns the body as-is. The
    /// consume-side pairing of the publish builder's `claim_check`.
    pub async fn resolve_body(&self, store: &dyn BlobStore) -> Result<Vec<u8>, LaserError> {
        if self.content_type == Some(ContentType::Ref) {
            return resolve_body(store, self.body()).await;
        }
        Ok(self.body().to_vec())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::sync::Mutex;

    #[derive(Default)]
    struct MemoryStore {
        blobs: Mutex<HashMap<String, Vec<u8>>>,
    }

    #[async_trait]
    impl BlobStore for MemoryStore {
        async fn put(&self, payload: Vec<u8>) -> Result<String, LaserError> {
            let mut blobs = self.blobs.lock().expect("store lock");
            let reference = format!("blob-{}", blobs.len());
            blobs.insert(reference.clone(), payload);
            Ok(reference)
        }

        async fn get(&self, reference: &str) -> Result<Vec<u8>, LaserError> {
            self.blobs
                .lock()
                .expect("store lock")
                .get(reference)
                .cloned()
                .ok_or_else(|| LaserError::Invalid(format!("no blob at `{reference}`")))
        }
    }

    #[tokio::test]
    async fn given_a_small_body_when_checked_in_then_should_pass_through_byte_identical() {
        let store = MemoryStore::default();
        let payload = b"small".to_vec();
        let (out, content_type) = check_in(&store, 1024, payload.clone())
            .await
            .expect("passes through");
        assert_eq!(out, payload);
        assert!(content_type.is_none());
        assert!(store.blobs.lock().expect("store lock").is_empty());
    }

    #[tokio::test]
    async fn given_a_large_body_when_checked_in_then_should_round_trip_through_the_store() {
        let store = MemoryStore::default();
        let payload = vec![7u8; 4096];
        let (capsule, content_type) = check_in(&store, 1024, payload.clone())
            .await
            .expect("externalizes");
        assert_eq!(content_type, Some(ContentType::Ref));
        let resolved = resolve_body(&store, &capsule).await.expect("resolves");
        assert_eq!(resolved, payload);
    }

    #[tokio::test]
    async fn given_a_claim_checked_agdx_record_when_resolved_then_should_resolve_the_envelope_body()
    {
        use crate::types::MintUlid;
        use laser_wire::agent::{AgentEnvelope, ConversationId, RecordId};
        use std::str::FromStr;
        let store = MemoryStore::default();
        let body = vec![9u8; 4096];
        let (capsule, content_type) = check_in(&store, 1024, body.clone())
            .await
            .expect("externalizes");
        let envelope = AgentEnvelope::event(
            RecordId::mint(),
            ConversationId::mint(),
            laser_wire::agent::AgentId::from_str("worker").expect("a valid agent id"),
            capsule,
        );
        let mut message = crate::testing::agent_message(
            encode_named(&envelope).expect("the envelope encodes"),
            crate::provenance::Provenance::builder()
                .conversation_id(crate::types::ConversationId::new())
                .build(),
        );
        message.envelope = Some(envelope);
        message.content_type = content_type;
        assert_eq!(message.resolve_body(&store).await.expect("resolves"), body);
        message.content_type = None;
        assert_eq!(
            message.resolve_body(&store).await.expect("passes through"),
            message.body(),
            "a body that was not claim-checked comes back as the envelope body"
        );
    }

    #[tokio::test]
    async fn given_a_tampered_blob_when_resolved_then_should_refuse_with_the_integrity_error() {
        let store = MemoryStore::default();
        let payload = vec![7u8; 4096];
        let (capsule, _) = check_in(&store, 1024, payload).await.expect("externalizes");
        for blob in store.blobs.lock().expect("store lock").values_mut() {
            *blob = b"swapped".to_vec();
        }
        let refused = resolve_body(&store, &capsule).await;
        assert!(matches!(refused, Err(LaserError::Integrity { .. })));
    }
}