1use bytes::{BufMut, BytesMut};
4use prost::Message;
5
6use crate::heddle::api::v1alpha1::{CallContext, CallFailure};
7
8pub const MAX_METHOD_PATH: usize = 1024;
10pub const MAX_CALL_CONTEXT: usize = 64 * 1024;
12pub const MAX_CONTROL_BODY: usize = 8 * 1024 * 1024;
14pub const MAX_RAW_BODY: u64 = 64 * 1024 * 1024 * 1024;
16
17const RESPONSE_SUCCESS: u8 = 0;
18const RESPONSE_FAILURE: u8 = 1;
19const STREAM_MESSAGE: u8 = 0;
20const STREAM_FAILURE: u8 = 1;
21const STREAM_RAW_BODY: u8 = 2;
22const STREAM_HEADER: usize = 5;
23const STREAM_RAW_HEADER: usize = 9;
24
25#[derive(Debug, thiserror::Error)]
27pub enum FrameError {
28 #[error("invalid hosted-call frame: {0}")]
30 Invalid(String),
31 #[error("invalid hosted-call protobuf: {0}")]
33 Decode(#[from] prost::DecodeError),
34}
35
36#[derive(Debug)]
38pub struct RequestFrame<'a> {
39 pub method: &'a str,
41 pub context: CallContext,
43 pub body: &'a [u8],
45}
46
47#[derive(Debug)]
49pub struct RequestPrelude<'a> {
50 pub method: &'a str,
52 pub context: CallContext,
54}
55
56#[derive(Debug)]
58pub enum ResponseFrame<'a> {
59 Success(&'a [u8]),
61 Failure(CallFailure),
63}
64
65#[derive(Debug)]
67pub enum StreamFrame<'a> {
68 Message(&'a [u8]),
70 Failure(CallFailure),
72 RawBody { length: u64 },
74}
75
76pub fn encode_request_frame(
79 method: &str,
80 context: &CallContext,
81 body: &[u8],
82) -> Result<Vec<u8>, FrameError> {
83 validate_body(body)?;
84 let mut frame = encode_request_prelude(method, context)?;
85 frame.extend_from_slice(body);
86 Ok(frame)
87}
88
89pub fn encode_request_prelude(method: &str, context: &CallContext) -> Result<Vec<u8>, FrameError> {
91 validate_method(method)?;
92 let context = context.encode_to_vec();
93 if context.len() > MAX_CALL_CONTEXT {
94 return Err(FrameError::Invalid(format!(
95 "call context is {} bytes; maximum is {MAX_CALL_CONTEXT}",
96 context.len()
97 )));
98 }
99 let method_len = u16::try_from(method.len())
100 .map_err(|_| FrameError::Invalid("method path exceeds u16".to_string()))?;
101 let context_len = u32::try_from(context.len())
102 .map_err(|_| FrameError::Invalid("call context exceeds u32".to_string()))?;
103 let mut frame = Vec::with_capacity(6 + method.len() + context.len());
104 frame.extend_from_slice(&method_len.to_be_bytes());
105 frame.extend_from_slice(&context_len.to_be_bytes());
106 frame.extend_from_slice(method.as_bytes());
107 frame.extend_from_slice(&context);
108 Ok(frame)
109}
110
111pub fn decode_request_frame(frame: &[u8]) -> Result<RequestFrame<'_>, FrameError> {
113 let (prelude, body_start) = decode_request_prelude(frame)?.ok_or_else(|| {
114 FrameError::Invalid("request frame contains a truncated prelude".to_string())
115 })?;
116 let body = &frame[body_start..];
117 validate_body(body)?;
118 Ok(RequestFrame {
119 method: prelude.method,
120 context: prelude.context,
121 body,
122 })
123}
124
125pub fn decode_request_prelude(
127 frame: &[u8],
128) -> Result<Option<(RequestPrelude<'_>, usize)>, FrameError> {
129 if frame.len() < 6 {
130 return Ok(None);
131 }
132 let method_len = u16::from_be_bytes([frame[0], frame[1]]) as usize;
133 let context_len = u32::from_be_bytes([frame[2], frame[3], frame[4], frame[5]]) as usize;
134 if method_len == 0 || method_len > MAX_METHOD_PATH || context_len > MAX_CALL_CONTEXT {
135 return Err(FrameError::Invalid(
136 "request prelude declares an invalid length".to_string(),
137 ));
138 }
139 let context_start = 6_usize
140 .checked_add(method_len)
141 .ok_or_else(|| FrameError::Invalid("request length overflow".to_string()))?;
142 let consumed = context_start
143 .checked_add(context_len)
144 .ok_or_else(|| FrameError::Invalid("request length overflow".to_string()))?;
145 if frame.len() < consumed {
146 return Ok(None);
147 }
148 let method = std::str::from_utf8(&frame[6..context_start])
149 .map_err(|_| FrameError::Invalid("method path is not UTF-8".to_string()))?;
150 validate_method(method)?;
151 Ok(Some((
152 RequestPrelude {
153 method,
154 context: CallContext::decode(&frame[context_start..consumed])?,
155 },
156 consumed,
157 )))
158}
159
160pub fn encode_success_response(body: &[u8]) -> Result<Vec<u8>, FrameError> {
162 let mut frame = BytesMut::with_capacity(1 + body.len());
163 encode_success_response_into(&mut frame, body)?;
164 Ok(frame.to_vec())
165}
166
167pub fn encode_success_response_into(frame: &mut BytesMut, body: &[u8]) -> Result<(), FrameError> {
169 validate_body(body)?;
170 frame.clear();
171 frame.reserve(1 + body.len());
172 frame.put_u8(RESPONSE_SUCCESS);
173 frame.extend_from_slice(body);
174 Ok(())
175}
176
177pub fn encode_failure_response(failure: &CallFailure) -> Result<Vec<u8>, FrameError> {
179 let mut frame = BytesMut::with_capacity(1 + failure.encoded_len());
180 encode_failure_response_into(&mut frame, failure)?;
181 Ok(frame.to_vec())
182}
183
184pub fn encode_failure_response_into(
186 frame: &mut BytesMut,
187 failure: &CallFailure,
188) -> Result<(), FrameError> {
189 let body_len = failure.encoded_len();
190 validate_body_len(body_len)?;
191 frame.clear();
192 frame.reserve(1 + body_len);
193 frame.put_u8(RESPONSE_FAILURE);
194 failure
195 .encode(frame)
196 .expect("BytesMut reserves the exact protobuf failure size");
197 Ok(())
198}
199
200pub fn decode_response_frame(frame: &[u8]) -> Result<ResponseFrame<'_>, FrameError> {
202 let (&outcome, body) = frame
203 .split_first()
204 .ok_or_else(|| FrameError::Invalid("response frame is empty".to_string()))?;
205 validate_body(body)?;
206 match outcome {
207 RESPONSE_SUCCESS => Ok(ResponseFrame::Success(body)),
208 RESPONSE_FAILURE => Ok(ResponseFrame::Failure(CallFailure::decode(body)?)),
209 value => Err(FrameError::Invalid(format!(
210 "unknown response outcome {value}"
211 ))),
212 }
213}
214
215pub fn encode_stream_message(body: &[u8]) -> Result<Vec<u8>, FrameError> {
217 let mut frame = BytesMut::with_capacity(STREAM_HEADER + body.len());
218 encode_stream_message_into(&mut frame, body)?;
219 Ok(frame.to_vec())
220}
221
222pub fn encode_stream_message_into(frame: &mut BytesMut, body: &[u8]) -> Result<(), FrameError> {
224 encode_stream_item_into(frame, STREAM_MESSAGE, body)
225}
226
227pub fn encode_stream_failure(failure: &CallFailure) -> Result<Vec<u8>, FrameError> {
229 let mut frame = BytesMut::with_capacity(STREAM_HEADER + failure.encoded_len());
230 encode_stream_failure_into(&mut frame, failure)?;
231 Ok(frame.to_vec())
232}
233
234pub fn encode_stream_failure_into(
236 frame: &mut BytesMut,
237 failure: &CallFailure,
238) -> Result<(), FrameError> {
239 let body_len = failure.encoded_len();
240 validate_body_len(body_len)?;
241 let body_len = u32::try_from(body_len)
242 .map_err(|_| FrameError::Invalid("stream item exceeds u32".to_string()))?;
243 frame.clear();
244 frame.reserve(STREAM_HEADER + body_len as usize);
245 frame.put_u8(STREAM_FAILURE);
246 frame.extend_from_slice(&body_len.to_be_bytes());
247 failure
248 .encode(frame)
249 .expect("BytesMut reserves the exact protobuf failure size");
250 Ok(())
251}
252
253pub fn encode_stream_raw_body(length: u64) -> Result<Vec<u8>, FrameError> {
256 let mut frame = BytesMut::with_capacity(STREAM_RAW_HEADER);
257 encode_stream_raw_body_into(&mut frame, length)?;
258 Ok(frame.to_vec())
259}
260
261pub fn encode_stream_raw_body_into(frame: &mut BytesMut, length: u64) -> Result<(), FrameError> {
263 if length == 0 || length > MAX_RAW_BODY {
264 return Err(FrameError::Invalid(format!(
265 "raw stream body is {length} bytes; range is 1..={MAX_RAW_BODY}"
266 )));
267 }
268 frame.clear();
269 frame.reserve(STREAM_RAW_HEADER);
270 frame.put_u8(STREAM_RAW_BODY);
271 frame.extend_from_slice(&length.to_be_bytes());
272 Ok(())
273}
274
275pub fn decode_stream_frame(buffer: &[u8]) -> Result<Option<(StreamFrame<'_>, usize)>, FrameError> {
280 if buffer.len() < STREAM_HEADER {
281 return Ok(None);
282 }
283 let kind = buffer[0];
284 if kind == STREAM_RAW_BODY {
285 if buffer.len() < STREAM_RAW_HEADER {
286 return Ok(None);
287 }
288 let length = u64::from_be_bytes(
289 buffer[1..STREAM_RAW_HEADER]
290 .try_into()
291 .expect("fixed raw header width"),
292 );
293 if length == 0 || length > MAX_RAW_BODY {
294 return Err(FrameError::Invalid(format!(
295 "raw stream body is {length} bytes; range is 1..={MAX_RAW_BODY}"
296 )));
297 }
298 return Ok(Some((StreamFrame::RawBody { length }, STREAM_RAW_HEADER)));
299 }
300 let body_len = u32::from_be_bytes([buffer[1], buffer[2], buffer[3], buffer[4]]) as usize;
301 if body_len > MAX_CONTROL_BODY {
302 return Err(FrameError::Invalid(format!(
303 "stream item is {body_len} bytes; maximum is {MAX_CONTROL_BODY}"
304 )));
305 }
306 let consumed = STREAM_HEADER
307 .checked_add(body_len)
308 .ok_or_else(|| FrameError::Invalid("stream item length overflow".to_string()))?;
309 if buffer.len() < consumed {
310 return Ok(None);
311 }
312 let body = &buffer[STREAM_HEADER..consumed];
313 let frame = match kind {
314 STREAM_MESSAGE => StreamFrame::Message(body),
315 STREAM_FAILURE => StreamFrame::Failure(CallFailure::decode(body)?),
316 value => {
317 return Err(FrameError::Invalid(format!(
318 "unknown stream item kind {value}"
319 )));
320 }
321 };
322 Ok(Some((frame, consumed)))
323}
324
325fn encode_stream_item_into(frame: &mut BytesMut, kind: u8, body: &[u8]) -> Result<(), FrameError> {
326 validate_body(body)?;
327 let body_len = u32::try_from(body.len())
328 .map_err(|_| FrameError::Invalid("stream item exceeds u32".to_string()))?;
329 frame.clear();
330 frame.reserve(STREAM_HEADER + body.len());
331 frame.put_u8(kind);
332 frame.extend_from_slice(&body_len.to_be_bytes());
333 frame.extend_from_slice(body);
334 Ok(())
335}
336
337fn validate_method(method: &str) -> Result<(), FrameError> {
338 if method.is_empty() || !method.starts_with('/') || method.len() > MAX_METHOD_PATH {
339 return Err(FrameError::Invalid(
340 "method path must begin with '/' and fit the method-path limit".to_string(),
341 ));
342 }
343 Ok(())
344}
345
346fn validate_body(body: &[u8]) -> Result<(), FrameError> {
347 validate_body_len(body.len())
348}
349
350fn validate_body_len(body_len: usize) -> Result<(), FrameError> {
351 if body_len > MAX_CONTROL_BODY {
352 return Err(FrameError::Invalid(format!(
353 "control body is {body_len} bytes; maximum is {MAX_CONTROL_BODY}"
354 )));
355 }
356 Ok(())
357}