haem_frames/embedding/

codec.rs

use std::fmt;
use std::io::{Read, Write};
use std::marker::PhantomData;
use std::num::NonZeroU32;

use serde::{
    Deserialize, Deserializer, Serialize,
    de::{DeserializeOwned, MapAccess, SeqAccess, Visitor, value::MapAccessDeserializer},
};

use super::{
    Embed, Error, FAILED_MESSAGE_CHARS, Failed, FailedCode, MINIMUM_FRAME_BYTES, Ready, Start, Tag,
    Vectors, check_ready, check_start,
};

struct Object<T>(T);

impl<'de, T: Deserialize<'de>> Deserialize<'de> for Object<T> {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        object(deserializer).map(Self)
    }
}

pub(super) fn object<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
    deserializer: D,
) -> Result<T, D::Error> {
    struct ObjectVisitor<T>(PhantomData<T>);
    impl<'de, T: Deserialize<'de>> Visitor<'de> for ObjectVisitor<T> {
        type Value = T;
        fn expecting(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
            output.write_str("a JSON object")
        }
        fn visit_map<M: MapAccess<'de>>(self, map: M) -> Result<T, M::Error> {
            T::deserialize(MapAccessDeserializer::new(map))
        }
    }
    deserializer.deserialize_map(ObjectVisitor(PhantomData))
}

pub(super) fn objects<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
    deserializer: D,
) -> Result<Vec<T>, D::Error> {
    struct Objects<T>(PhantomData<T>);
    impl<'de, T: Deserialize<'de>> Visitor<'de> for Objects<T> {
        type Value = Vec<T>;
        fn expecting(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
            output.write_str("an array of JSON objects")
        }
        fn visit_seq<S: SeqAccess<'de>>(self, mut sequence: S) -> Result<Vec<T>, S::Error> {
            let mut items = Vec::new();
            while let Some(Object(item)) = sequence.next_element::<Object<T>>()? {
                items.push(item);
            }
            Ok(items)
        }
    }
    deserializer.deserialize_seq(Objects(PhantomData))
}

#[derive(Debug)]
pub enum Message {
    Start(Start),
    Ready(Ready),
    Embed(Embed),
    Vectors(Vectors),
    Failed(Failed),
}

impl Message {
    #[must_use]
    pub const fn tag(&self) -> Tag {
        match self {
            Self::Start(_) => Tag::Start,
            Self::Ready(_) => Tag::Ready,
            Self::Embed(_) => Tag::Embed,
            Self::Vectors(_) => Tag::Vectors,
            Self::Failed(_) => Tag::Failed,
        }
    }
}

fn protocol(error: impl fmt::Display) -> Error {
    Error::new(FailedCode::EmbedProtocol, error.to_string())
}

fn parse<T: DeserializeOwned>(payload: &[u8]) -> Result<T, Error> {
    let mut deserializer = serde_json::Deserializer::from_slice(payload);
    let value = object(&mut deserializer).map_err(protocol)?;
    deserializer.end().map_err(protocol)?;
    Ok(value)
}

fn check_failed(failed: &Failed) -> Result<(), Error> {
    if failed.message.chars().count() > FAILED_MESSAGE_CHARS {
        return Err(protocol("Failed message exceeds 1024 characters"));
    }
    Ok(())
}

/// Decode a tagged JSON object and validate its exchange envelope.
///
/// # Errors
/// Returns `embed_protocol` for unknown tags, malformed JSON, invalid Ready or
/// Failed envelopes, and oversized Start messages; Start limits may refuse
/// with `embed_limits`.
pub fn decode_message(tag: u8, payload: &[u8]) -> Result<Message, Error> {
    match Tag::try_from(tag)? {
        Tag::Start => {
            if payload
                .len()
                .checked_add(1)
                .is_none_or(|size| size > MINIMUM_FRAME_BYTES as usize)
            {
                return Err(protocol("Start frame exceeds its fixed limit"));
            }
            let start = parse(payload)?;
            check_start(&start)?;
            Ok(Message::Start(start))
        }
        Tag::Ready => {
            let ready = parse(payload)?;
            check_ready(&ready)?;
            Ok(Message::Ready(ready))
        }
        Tag::Embed => parse(payload).map(Message::Embed),
        Tag::Vectors => parse(payload).map(Message::Vectors),
        Tag::Failed => {
            let failed = parse(payload)?;
            check_failed(&failed)?;
            Ok(Message::Failed(failed))
        }
    }
}

/// Read and decode one bounded frame; clean EOF returns no message.
///
/// # Errors
/// Returns `embed_protocol` for framing or input failures; propagates the
/// envelope and limit refusals from [`decode_message`].
pub fn read_message(input: &mut impl Read, limit: NonZeroU32) -> Result<Option<Message>, Error> {
    let Some(frame) = crate::frame::read(input, limit).map_err(protocol)? else {
        return Ok(None);
    };
    let (tag, payload) = frame
        .split_first()
        .ok_or_else(|| protocol("frame has no tag"))?;
    decode_message(*tag, payload).map(Some)
}

/// Read the first exchange with its fixed bound before allocating the payload.
///
/// # Errors
/// Propagates framing and envelope refusals; returns `embed_protocol` when
/// the first message is not Start.
pub fn read_start(input: &mut impl Read) -> Result<Option<Start>, Error> {
    let limit = NonZeroU32::new(MINIMUM_FRAME_BYTES)
        .ok_or_else(|| protocol("fixed frame limit is zero"))?;
    match read_message(input, limit)? {
        Some(Message::Start(start)) => Ok(Some(start)),
        None => Ok(None),
        Some(_) => Err(protocol("first embedding frame must be Start")),
    }
}

fn payload(value: &impl Serialize) -> Result<Vec<u8>, Error> {
    serde_json::to_vec(value).map_err(protocol)
}

/// Validate an envelope and write its tagged JSON in one flushed frame.
///
/// # Errors
/// Propagates Start and Ready validation refusals; returns `embed_protocol`
/// for invalid Failed envelopes, encoding, framing, output or flush failures.
pub fn write_message(
    output: &mut impl Write,
    message: &Message,
    limit: NonZeroU32,
) -> Result<(), Error> {
    let bytes = match message {
        Message::Start(start) => {
            check_start(start)?;
            payload(start)?
        }
        Message::Ready(ready) => {
            check_ready(ready)?;
            payload(ready)?
        }
        Message::Embed(embed) => payload(embed)?,
        Message::Vectors(vectors) => payload(vectors)?,
        Message::Failed(failed) => {
            check_failed(failed)?;
            payload(failed)?
        }
    };
    let limit = if matches!(message, Message::Start(_)) {
        limit.min(
            NonZeroU32::new(MINIMUM_FRAME_BYTES)
                .ok_or_else(|| protocol("fixed frame limit is zero"))?,
        )
    } else {
        limit
    };
    crate::frame::write(output, message.tag() as u8, &bytes, limit).map_err(protocol)
}