haem_frames/embedding/
codec.rs

1use std::fmt;
2use std::io::{Read, Write};
3use std::marker::PhantomData;
4use std::num::NonZeroU32;
5
6use serde::{
7    Deserialize, Deserializer, Serialize,
8    de::{DeserializeOwned, MapAccess, SeqAccess, Visitor, value::MapAccessDeserializer},
9};
10
11use super::{
12    Embed, Error, FAILED_MESSAGE_CHARS, Failed, FailedCode, MINIMUM_FRAME_BYTES, Ready, Start, Tag,
13    Vectors, check_ready, check_start,
14};
15
16struct Object<T>(T);
17
18impl<'de, T: Deserialize<'de>> Deserialize<'de> for Object<T> {
19    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
20        object(deserializer).map(Self)
21    }
22}
23
24pub(super) fn object<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
25    deserializer: D,
26) -> Result<T, D::Error> {
27    struct ObjectVisitor<T>(PhantomData<T>);
28    impl<'de, T: Deserialize<'de>> Visitor<'de> for ObjectVisitor<T> {
29        type Value = T;
30        fn expecting(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
31            output.write_str("a JSON object")
32        }
33        fn visit_map<M: MapAccess<'de>>(self, map: M) -> Result<T, M::Error> {
34            T::deserialize(MapAccessDeserializer::new(map))
35        }
36    }
37    deserializer.deserialize_map(ObjectVisitor(PhantomData))
38}
39
40pub(super) fn objects<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
41    deserializer: D,
42) -> Result<Vec<T>, D::Error> {
43    struct Objects<T>(PhantomData<T>);
44    impl<'de, T: Deserialize<'de>> Visitor<'de> for Objects<T> {
45        type Value = Vec<T>;
46        fn expecting(&self, output: &mut fmt::Formatter<'_>) -> fmt::Result {
47            output.write_str("an array of JSON objects")
48        }
49        fn visit_seq<S: SeqAccess<'de>>(self, mut sequence: S) -> Result<Vec<T>, S::Error> {
50            let mut items = Vec::new();
51            while let Some(Object(item)) = sequence.next_element::<Object<T>>()? {
52                items.push(item);
53            }
54            Ok(items)
55        }
56    }
57    deserializer.deserialize_seq(Objects(PhantomData))
58}
59
60#[derive(Debug)]
61pub enum Message {
62    Start(Start),
63    Ready(Ready),
64    Embed(Embed),
65    Vectors(Vectors),
66    Failed(Failed),
67}
68
69impl Message {
70    pub const fn tag(&self) -> Tag {
71        match self {
72            Self::Start(_) => Tag::Start,
73            Self::Ready(_) => Tag::Ready,
74            Self::Embed(_) => Tag::Embed,
75            Self::Vectors(_) => Tag::Vectors,
76            Self::Failed(_) => Tag::Failed,
77        }
78    }
79}
80
81fn protocol(error: impl fmt::Display) -> Error {
82    Error::new(FailedCode::EmbedProtocol, error.to_string())
83}
84
85fn parse<T: DeserializeOwned>(payload: &[u8]) -> Result<T, Error> {
86    let mut deserializer = serde_json::Deserializer::from_slice(payload);
87    let value = object(&mut deserializer).map_err(protocol)?;
88    deserializer.end().map_err(protocol)?;
89    Ok(value)
90}
91
92fn check_failed(failed: &Failed) -> Result<(), Error> {
93    if failed.message.chars().count() > FAILED_MESSAGE_CHARS {
94        return Err(protocol("Failed message exceeds 1024 characters"));
95    }
96    Ok(())
97}
98
99pub fn decode_message(tag: u8, payload: &[u8]) -> Result<Message, Error> {
100    match Tag::try_from(tag)? {
101        Tag::Start => {
102            if payload
103                .len()
104                .checked_add(1)
105                .is_none_or(|size| size > MINIMUM_FRAME_BYTES as usize)
106            {
107                return Err(protocol("Start frame exceeds its fixed limit"));
108            }
109            let start = parse(payload)?;
110            check_start(&start)?;
111            Ok(Message::Start(start))
112        }
113        Tag::Ready => {
114            let ready = parse(payload)?;
115            check_ready(&ready)?;
116            Ok(Message::Ready(ready))
117        }
118        Tag::Embed => parse(payload).map(Message::Embed),
119        Tag::Vectors => parse(payload).map(Message::Vectors),
120        Tag::Failed => {
121            let failed = parse(payload)?;
122            check_failed(&failed)?;
123            Ok(Message::Failed(failed))
124        }
125    }
126}
127
128pub fn read_message(input: &mut impl Read, limit: NonZeroU32) -> Result<Option<Message>, Error> {
129    let Some(frame) = crate::frame::read(input, limit).map_err(protocol)? else {
130        return Ok(None);
131    };
132    let (tag, payload) = frame
133        .split_first()
134        .ok_or_else(|| protocol("frame has no tag"))?;
135    decode_message(*tag, payload).map(Some)
136}
137
138/// Read the first exchange with its fixed bound before allocating the payload.
139pub fn read_start(input: &mut impl Read) -> Result<Option<Start>, Error> {
140    let limit = NonZeroU32::new(MINIMUM_FRAME_BYTES)
141        .ok_or_else(|| protocol("fixed frame limit is zero"))?;
142    match read_message(input, limit)? {
143        Some(Message::Start(start)) => Ok(Some(start)),
144        None => Ok(None),
145        Some(_) => Err(protocol("first embedding frame must be Start")),
146    }
147}
148
149fn payload(value: &impl Serialize) -> Result<Vec<u8>, Error> {
150    serde_json::to_vec(value).map_err(protocol)
151}
152
153pub fn write_message(
154    output: &mut impl Write,
155    message: &Message,
156    limit: NonZeroU32,
157) -> Result<(), Error> {
158    let bytes = match message {
159        Message::Start(start) => {
160            check_start(start)?;
161            payload(start)?
162        }
163        Message::Ready(ready) => {
164            check_ready(ready)?;
165            payload(ready)?
166        }
167        Message::Embed(embed) => payload(embed)?,
168        Message::Vectors(vectors) => payload(vectors)?,
169        Message::Failed(failed) => {
170            check_failed(failed)?;
171            payload(failed)?
172        }
173    };
174    let limit = if matches!(message, Message::Start(_)) {
175        limit.min(
176            NonZeroU32::new(MINIMUM_FRAME_BYTES)
177                .ok_or_else(|| protocol("fixed frame limit is zero"))?,
178        )
179    } else {
180        limit
181    };
182    crate::frame::write(output, message.tag() as u8, &bytes, limit).map_err(protocol)
183}
haem_frames/embedding/codec.rs

haem_frames/embedding/
codec.rs
