yapi_ai/aws/
eventstream.rs

1//! The `application/vnd.amazon.eventstream` framing of AWS streaming
2//! responses: length-prefixed messages with typed headers and CRC32
3//! checksums, as `@smithy/eventstream-codec` reads them.
4
5use indexmap::IndexMap;
6
7/// One decoded message.
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub struct Message {
10    /// String-valued headers, such as `:event-type`. Other header types are
11    /// kept as their text form.
12    pub headers: IndexMap<String, String>,
13    /// The payload.
14    pub payload: Vec<u8>,
15}
16
17impl Message {
18    /// A header's value.
19    pub fn header(&self, name: &str) -> Option<&str> {
20        self.headers.get(name).map(String::as_str)
21    }
22}
23
24/// Splits a byte stream into messages.
25#[derive(Debug, Default)]
26pub struct Decoder {
27    buffer: Vec<u8>,
28}
29
30fn u32_at(bytes: &[u8], at: usize) -> u32 {
31    u32::from_be_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]])
32}
33
34impl Decoder {
35    /// Adds received bytes.
36    pub fn push(&mut self, bytes: &[u8]) {
37        self.buffer.extend_from_slice(bytes);
38    }
39
40    /// The next complete message, if one has arrived.
41    pub fn next_message(&mut self) -> Result<Option<Message>, String> {
42        if self.buffer.len() < 12 {
43            return Ok(None);
44        }
45        let total = u32_at(&self.buffer, 0) as usize;
46        let headers_len = u32_at(&self.buffer, 4) as usize;
47        if total < 16 || headers_len > total - 16 {
48            return Err(format!("Invalid event stream message length: {total}"));
49        }
50        if self.buffer.len() < total {
51            return Ok(None);
52        }
53        let message: Vec<u8> = self.buffer.drain(..total).collect();
54        let prelude_crc = u32_at(&message, 8);
55        if crc32fast::hash(&message[..8]) != prelude_crc {
56            return Err(format!(
57                "The prelude checksum specified in the message ({prelude_crc}) does not match the calculated CRC32 checksum ({})",
58                crc32fast::hash(&message[..8])
59            ));
60        }
61        let message_crc = u32_at(&message, total - 4);
62        if crc32fast::hash(&message[..total - 4]) != message_crc {
63            return Err(format!(
64                "The message checksum ({}) did not match the expected value of {message_crc}",
65                crc32fast::hash(&message[..total - 4])
66            ));
67        }
68        let headers = parse_headers(&message[12..12 + headers_len])?;
69        Ok(Some(Message {
70            headers,
71            payload: message[12 + headers_len..total - 4].to_vec(),
72        }))
73    }
74
75    /// Checks the end of the stream: a message whose length arrived but not
76    /// all of its bytes is an error, as in smithy's `getChunkedStream`.
77    pub fn finish(&self) -> Result<(), String> {
78        match self.buffer.len() {
79            0..4 => Ok(()),
80            _ => Err("Truncated event message received.".into()),
81        }
82    }
83}
84
85fn parse_headers(mut bytes: &[u8]) -> Result<IndexMap<String, String>, String> {
86    let invalid = || "Invalid event stream header".to_owned();
87    let mut headers = IndexMap::new();
88    while !bytes.is_empty() {
89        let name_len = usize::from(bytes[0]);
90        let name = bytes.get(1..1 + name_len).ok_or_else(invalid)?;
91        let name = String::from_utf8_lossy(name).into_owned();
92        bytes = &bytes[1 + name_len..];
93        let kind = *bytes.first().ok_or_else(invalid)?;
94        bytes = &bytes[1..];
95        let fixed = |len: usize, bytes: &mut &[u8]| -> Result<Vec<u8>, String> {
96            let value = bytes.get(..len).ok_or_else(invalid)?.to_vec();
97            *bytes = &bytes[len..];
98            Ok(value)
99        };
100        let value = match kind {
101            0 => "true".to_owned(),
102            1 => "false".to_owned(),
103            2 => (fixed(1, &mut bytes)?[0] as i8).to_string(),
104            3 => i16::from_be_bytes(fixed(2, &mut bytes)?.try_into().map_err(|_| invalid())?)
105                .to_string(),
106            4 => i32::from_be_bytes(fixed(4, &mut bytes)?.try_into().map_err(|_| invalid())?)
107                .to_string(),
108            5 | 8 => i64::from_be_bytes(fixed(8, &mut bytes)?.try_into().map_err(|_| invalid())?)
109                .to_string(),
110            6 | 7 => {
111                let len = bytes.get(..2).ok_or_else(invalid)?;
112                let len = usize::from(u16::from_be_bytes([len[0], len[1]]));
113                bytes = &bytes[2..];
114                String::from_utf8_lossy(&fixed(len, &mut bytes)?).into_owned()
115            }
116            9 => yapi_types::time::hex(&fixed(16, &mut bytes)?),
117            _ => return Err(format!("Unrecognized event stream header type: {kind}")),
118        };
119        headers.insert(name, value);
120    }
121    Ok(headers)
122}
123
124/// Encodes a message with string headers, for tests and mock servers.
125pub fn encode(headers: &[(&str, &str)], payload: &[u8]) -> Vec<u8> {
126    let mut header_bytes = Vec::new();
127    for (name, value) in headers {
128        header_bytes.push(name.len() as u8);
129        header_bytes.extend_from_slice(name.as_bytes());
130        header_bytes.push(7);
131        header_bytes.extend_from_slice(&(value.len() as u16).to_be_bytes());
132        header_bytes.extend_from_slice(value.as_bytes());
133    }
134    let total = (16 + header_bytes.len() + payload.len()) as u32;
135    let mut message = Vec::new();
136    message.extend_from_slice(&total.to_be_bytes());
137    message.extend_from_slice(&(header_bytes.len() as u32).to_be_bytes());
138    let prelude_crc = crc32fast::hash(&message);
139    message.extend_from_slice(&prelude_crc.to_be_bytes());
140    message.extend_from_slice(&header_bytes);
141    message.extend_from_slice(payload);
142    let crc = crc32fast::hash(&message);
143    message.extend_from_slice(&crc.to_be_bytes());
144    message
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn split_and_corrupt_messages() {
153        let first = encode(
154            &[(":event-type", "messageStart")],
155            br#"{"role":"assistant"}"#,
156        );
157        let second = encode(&[(":message-type", "exception")], b"{}");
158        let mut decoder = Decoder::default();
159        let mut stream = first.clone();
160        stream.extend_from_slice(&second);
161        decoder.push(&stream[..7]);
162        assert_eq!(decoder.next_message(), Ok(None));
163        decoder.push(&stream[7..]);
164        let message = decoder.next_message().unwrap().unwrap();
165        assert_eq!(message.header(":event-type"), Some("messageStart"));
166        assert_eq!(message.payload, br#"{"role":"assistant"}"#);
167        assert_eq!(
168            decoder
169                .next_message()
170                .unwrap()
171                .unwrap()
172                .header(":message-type"),
173            Some("exception")
174        );
175        assert!(decoder.buffer.is_empty());
176        assert_eq!(decoder.finish(), Ok(()));
177        decoder.push(&stream[..3]);
178        assert_eq!(decoder.finish(), Ok(()));
179        decoder.push(&stream[3..4]);
180        assert!(decoder.finish().unwrap_err().contains("Truncated"));
181        decoder.buffer.clear();
182        let mut corrupt = first;
183        let last = corrupt.len() - 5;
184        corrupt[last] ^= 1;
185        decoder.push(&corrupt);
186        assert!(
187            decoder
188                .next_message()
189                .unwrap_err()
190                .contains("message checksum")
191        );
192    }
193}
yapi_ai/aws/eventstream.rs

yapi_ai/aws/
eventstream.rs
