1use std::convert::TryFrom;
2use std::io::{self, IoSlice};
3use std::io::{Read, Write};
4
5pub type MethodId = u8;
6
7pub const REQUEST_HEADER_LEN: usize = 5;
8pub const RESPONSE_HEADER_LEN: usize = 5;
9pub(crate) const MAX_PAYLOAD_LEN: usize = 2 * 1024;
10pub(crate) const SOCKET_PERMISSIONS: u32 = 0o766;
11
12#[repr(u8)]
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum UmbralStatus {
15 Ok = 0,
16 MethodNotFound = 1,
17 PayloadTooLarge = 2,
18 InvalidFrame = 3,
19 HandlerError = 4,
20}
21
22impl TryFrom<u8> for UmbralStatus {
23 type Error = io::Error;
24
25 fn try_from(value: u8) -> io::Result<Self> {
26 match value {
27 0 => Ok(Self::Ok),
28 1 => Ok(Self::MethodNotFound),
29 2 => Ok(Self::PayloadTooLarge),
30 3 => Ok(Self::InvalidFrame),
31 4 => Ok(Self::HandlerError),
32 _ => Err(io::Error::new(
33 io::ErrorKind::InvalidData,
34 "unknown umbral status",
35 )),
36 }
37 }
38}
39
40fn payload_too_large() -> io::Error {
41 io::Error::new(io::ErrorKind::InvalidData, "payload too large")
42}
43
44fn payload_len_from_header(header: &[u8], max_payload_len: usize) -> io::Result<usize> {
45 let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
46 if len > max_payload_len {
47 return Err(payload_too_large());
48 }
49 Ok(len)
50}
51
52fn ensure_encodable_payload(payload: &[u8]) -> io::Result<()> {
53 if payload.len() > u32::MAX as usize {
54 return Err(io::Error::new(
55 io::ErrorKind::InvalidInput,
56 "payload length exceeds u32",
57 ));
58 }
59 Ok(())
60}
61
62#[cfg(test)]
63fn read_request_sync<R: Read>(
64 reader: &mut R,
65 max_payload_len: usize,
66) -> io::Result<(MethodId, Vec<u8>)> {
67 let mut payload = Vec::new();
68 let method = read_request_into_sync(reader, max_payload_len, &mut payload)?;
69 Ok((method, payload))
70}
71
72#[cfg(test)]
73fn read_request_into_sync<R: Read>(
74 reader: &mut R,
75 max_payload_len: usize,
76 payload: &mut Vec<u8>,
77) -> io::Result<MethodId> {
78 let mut header = [0u8; REQUEST_HEADER_LEN];
79 reader.read_exact(&mut header)?;
80 let len = payload_len_from_header(&header, max_payload_len)?;
81 payload.clear();
82 payload.resize(len, 0);
83 reader.read_exact(payload.as_mut_slice())?;
84 Ok(header[0])
85}
86
87pub(crate) fn write_request_sync<W: Write>(
88 writer: &mut W,
89 method: MethodId,
90 payload: &[u8],
91) -> io::Result<()> {
92 ensure_encodable_payload(payload)?;
93 let len = payload.len() as u32;
94 let mut header = [0u8; REQUEST_HEADER_LEN];
95 header[0] = method;
96 header[1..].copy_from_slice(&len.to_be_bytes());
97 write_frame_sync(writer, &header, payload)
98}
99
100pub(crate) fn read_response_sync<R: Read>(
101 reader: &mut R,
102 max_payload_len: usize,
103) -> io::Result<(UmbralStatus, Vec<u8>)> {
104 let mut header = [0u8; RESPONSE_HEADER_LEN];
105 reader.read_exact(&mut header)?;
106 let status = UmbralStatus::try_from(header[0])?;
107 let len = payload_len_from_header(&header, max_payload_len)?;
108 let mut payload = vec![0u8; len];
109 reader.read_exact(&mut payload)?;
110 Ok((status, payload))
111}
112
113#[cfg(test)]
114fn read_response_into_sync<R: Read>(
115 reader: &mut R,
116 max_payload_len: usize,
117 payload: &mut Vec<u8>,
118) -> io::Result<UmbralStatus> {
119 let mut header = [0u8; RESPONSE_HEADER_LEN];
120 reader.read_exact(&mut header)?;
121 let status = UmbralStatus::try_from(header[0])?;
122 let len = payload_len_from_header(&header, max_payload_len)?;
123 payload.clear();
124 payload.resize(len, 0);
125 reader.read_exact(payload.as_mut_slice())?;
126 Ok(status)
127}
128
129#[cfg(test)]
130fn write_response_sync<W: Write>(
131 writer: &mut W,
132 status: UmbralStatus,
133 payload: &[u8],
134) -> io::Result<()> {
135 ensure_encodable_payload(payload)?;
136 let len = payload.len() as u32;
137 let mut header = [0u8; RESPONSE_HEADER_LEN];
138 header[0] = status as u8;
139 header[1..].copy_from_slice(&len.to_be_bytes());
140 write_frame_sync(writer, &header, payload)
141}
142
143pub(crate) fn write_frame_sync<W: Write>(
144 writer: &mut W,
145 header: &[u8],
146 payload: &[u8],
147) -> io::Result<()> {
148 if payload.is_empty() {
149 return writer.write_all(header);
150 }
151
152 let mut header_offset = 0;
153 let mut payload_offset = 0;
154
155 while header_offset < header.len() || payload_offset < payload.len() {
156 let written = if header_offset == header.len() {
157 writer.write(&payload[payload_offset..])?
158 } else if payload_offset == payload.len() {
159 writer.write(&header[header_offset..])?
160 } else {
161 writer.write_vectored(&[
162 IoSlice::new(&header[header_offset..]),
163 IoSlice::new(&payload[payload_offset..]),
164 ])?
165 };
166
167 if written == 0 {
168 return Err(io::Error::new(
169 io::ErrorKind::WriteZero,
170 "failed to write umbral frame",
171 ));
172 }
173
174 let remaining_header = header.len() - header_offset;
175 if written < remaining_header {
176 header_offset += written;
177 } else {
178 header_offset = header.len();
179 payload_offset += written - remaining_header;
180 }
181 }
182
183 Ok(())
184}
185
186#[cfg(test)]
187mod tests {
188 use std::io::Cursor;
189
190 use super::*;
191
192 #[test]
193 fn encode_decode_empty_request() {
194 let mut buffer = Vec::new();
195 write_request_sync(&mut buffer, 7, b"").unwrap();
196
197 let (method, payload) =
198 read_request_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN).unwrap();
199
200 assert_eq!(method, 7);
201 assert!(payload.is_empty());
202 }
203
204 #[test]
205 fn encode_decode_small_request() {
206 let mut buffer = Vec::new();
207 write_request_sync(&mut buffer, 9, b"abc").unwrap();
208
209 let (method, payload) =
210 read_request_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN).unwrap();
211
212 assert_eq!(method, 9);
213 assert_eq!(payload, b"abc");
214 }
215
216 #[test]
217 fn encode_decode_ok_response() {
218 let mut buffer = Vec::new();
219 write_response_sync(&mut buffer, UmbralStatus::Ok, b"done").unwrap();
220
221 let (status, payload) =
222 read_response_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN).unwrap();
223
224 assert_eq!(status, UmbralStatus::Ok);
225 assert_eq!(payload, b"done");
226 }
227
228 #[test]
229 fn encode_decode_method_not_found_response() {
230 let mut buffer = Vec::new();
231 write_response_sync(&mut buffer, UmbralStatus::MethodNotFound, b"").unwrap();
232
233 let (status, payload) =
234 read_response_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN).unwrap();
235
236 assert_eq!(status, UmbralStatus::MethodNotFound);
237 assert!(payload.is_empty());
238 }
239
240 #[test]
241 fn unknown_status_returns_invalid_data() {
242 let buffer = [99, 0, 0, 0, 0];
243
244 let err = read_response_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN)
245 .expect_err("unknown status must fail");
246
247 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
248 }
249
250 #[test]
251 fn payload_above_max_returns_invalid_data() {
252 let buffer = [1, 0, 0, 0, 2];
253
254 let err = read_request_sync(&mut Cursor::new(buffer), 1)
255 .expect_err("payload above max must fail");
256
257 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
258 }
259
260 #[test]
261 fn incomplete_header_returns_unexpected_eof() {
262 let buffer = [1, 0, 0];
263
264 let err = read_request_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN)
265 .expect_err("incomplete header must fail");
266
267 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
268 }
269
270 #[test]
271 fn incomplete_payload_returns_unexpected_eof() {
272 let buffer = [1, 0, 0, 0, 3, b'a'];
273
274 let err = read_request_sync(&mut Cursor::new(buffer), MAX_PAYLOAD_LEN)
275 .expect_err("incomplete payload must fail");
276
277 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
278 }
279
280 #[test]
281 fn concatenated_frames_are_read_sequentially() {
282 let mut buffer = Vec::new();
283 write_request_sync(&mut buffer, 1, b"one").unwrap();
284 write_request_sync(&mut buffer, 2, b"two").unwrap();
285 let mut cursor = Cursor::new(buffer);
286
287 let first = read_request_sync(&mut cursor, MAX_PAYLOAD_LEN).unwrap();
288 let second = read_request_sync(&mut cursor, MAX_PAYLOAD_LEN).unwrap();
289
290 assert_eq!(first, (1, b"one".to_vec()));
291 assert_eq!(second, (2, b"two".to_vec()));
292 }
293
294 #[test]
295 fn read_request_into_sync_reuses_vec_capacity() {
296 let mut buffer = Vec::new();
297 write_request_sync(&mut buffer, 12, b"abc").unwrap();
298 let mut cursor = Cursor::new(buffer);
299 let mut payload = Vec::with_capacity(128);
300 let capacity = payload.capacity();
301
302 read_request_into_sync(&mut cursor, MAX_PAYLOAD_LEN, &mut payload).unwrap();
303
304 assert_eq!(payload, b"abc");
305 assert_eq!(payload.capacity(), capacity);
306 }
307
308 #[test]
309 fn read_request_into_sync_rejects_oversize_before_reading_payload() {
310 let mut buffer = Vec::new();
311 write_request_sync(&mut buffer, 13, b"12345").unwrap();
312 let mut cursor = Cursor::new(buffer);
313 let mut payload = Vec::new();
314
315 let err = read_request_into_sync(&mut cursor, 4, &mut payload)
316 .expect_err("oversize payload must fail");
317
318 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
319 assert!(payload.is_empty());
320 assert_eq!(cursor.position(), REQUEST_HEADER_LEN as u64);
321 }
322
323 #[test]
324 fn read_response_into_sync_reads_status_and_payload() {
325 let mut buffer = Vec::new();
326 write_response_sync(&mut buffer, UmbralStatus::Ok, b"done").unwrap();
327 let mut cursor = Cursor::new(buffer);
328 let mut payload = Vec::new();
329
330 let status = read_response_into_sync(&mut cursor, MAX_PAYLOAD_LEN, &mut payload).unwrap();
331
332 assert_eq!(status, UmbralStatus::Ok);
333 assert_eq!(payload, b"done");
334 }
335
336 #[test]
337 fn read_response_into_sync_rejects_oversize() {
338 let mut buffer = Vec::new();
339 write_response_sync(&mut buffer, UmbralStatus::Ok, b"12345").unwrap();
340 let mut cursor = Cursor::new(buffer);
341 let mut payload = Vec::new();
342
343 let err = read_response_into_sync(&mut cursor, 4, &mut payload)
344 .expect_err("oversize response must fail");
345
346 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
347 assert!(payload.is_empty());
348 assert_eq!(cursor.position(), RESPONSE_HEADER_LEN as u64);
349 }
350}