1pub mod response;
104
105pub use response::{
106 decode_response, encode_response_fragment, pack_response, split_response, unpack_response,
107 FragmentOutcome, ResponseDecoder, ResponseHeader, ResponseSymbols, RpcResponse,
108 MAX_RESPONSE_BODY, MAX_RESPONSE_FRAGMENT, MAX_RESPONSE_FRAGMENTS, MAX_RESPONSE_HEADER_BLOCK,
109 RESPONSE_HEADERS_FLAG, RPC_REPAIR_SYMBOLS, RPC_RESPONSE_OVERHEAD_BASE,
110};
111
112use crate::aux_mux::AUX_MAX_PAYLOAD;
113use crate::node_status::MAX_DEVICE_ID;
114
115pub const RPC_REQUEST_FRAGMENT_OVERHEAD: usize = 10;
118
119pub const MAX_REQUEST_CHUNK: usize =
122 AUX_MAX_PAYLOAD - RPC_REQUEST_FRAGMENT_OVERHEAD - MAX_DEVICE_ID;
123
124pub const MAX_REQUEST_BODY: usize = 1024 * 1024;
127
128pub const MAX_REQUEST_TICKET: usize = 1024;
130
131pub const MAX_REQUEST_OBJECT: usize =
133 MAX_REQUEST_BODY + 2 + u16::MAX as usize + 1 + u8::MAX as usize + 2 + MAX_REQUEST_TICKET;
134
135pub const MAX_REQUEST_FRAGMENTS: usize = MAX_REQUEST_OBJECT.div_ceil(MAX_REQUEST_CHUNK);
137
138pub const REQUEST_REASSEMBLY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
142
143pub const MAX_PARTIAL_REQUESTS: usize = 8;
145
146pub const MAX_PARTIAL_REQUEST_BYTES: usize = 4 * MAX_REQUEST_OBJECT;
148
149#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154#[repr(u8)]
155pub enum RpcMethod {
156 Get = 1,
157 Post = 2,
158 Put = 3,
159 Delete = 4,
160 Patch = 5,
161}
162
163impl RpcMethod {
164 pub fn from_u8(v: u8) -> Option<Self> {
167 match v {
168 1 => Some(Self::Get),
169 2 => Some(Self::Post),
170 3 => Some(Self::Put),
171 4 => Some(Self::Delete),
172 5 => Some(Self::Patch),
173 _ => None,
174 }
175 }
176
177 pub fn from_http_method(method: &str) -> Option<Self> {
179 match method {
180 "GET" => Some(Self::Get),
181 "POST" => Some(Self::Post),
182 "PUT" => Some(Self::Put),
183 "DELETE" => Some(Self::Delete),
184 "PATCH" => Some(Self::Patch),
185 _ => None,
186 }
187 }
188
189 pub fn as_http_method(&self) -> &'static str {
191 match self {
192 Self::Get => "GET",
193 Self::Post => "POST",
194 Self::Put => "PUT",
195 Self::Delete => "DELETE",
196 Self::Patch => "PATCH",
197 }
198 }
199}
200
201#[derive(Debug, Clone, PartialEq, Eq)]
203pub struct RpcRequest {
204 pub id: u32,
205 pub method: RpcMethod,
206 pub target: Vec<u8>,
209 pub path: Vec<u8>,
211 pub content_type: Vec<u8>,
213 pub ticket: Vec<u8>,
216 pub body: Vec<u8>,
217}
218
219#[derive(Debug, Clone, Copy, PartialEq, Eq)]
221pub struct RequestFragment<'a> {
222 pub id: u32,
223 pub method: RpcMethod,
224 pub target: &'a [u8],
225 pub seq: u16,
226 pub total: u16,
227 pub chunk: &'a [u8],
228}
229
230#[derive(Debug, Clone, Copy, PartialEq, Eq)]
233pub enum RpcCodecError {
234 TooShort,
236 BadMethod(u8),
238 LengthMismatch { declared: usize, actual: usize },
240 BadFragmentIndex { index: u16, total: u16 },
243 BadSenderLen(u8),
245}
246
247#[derive(Debug, Clone, Copy, PartialEq, Eq)]
249pub enum RequestEncodeError {
250 BodyTooLarge,
252 FieldTooLong,
254}
255
256pub fn encode_request(
262 method: RpcMethod,
263 id: u32,
264 target: &[u8],
265 request: &RequestParts<'_>,
266) -> Result<Vec<Vec<u8>>, RequestEncodeError> {
267 if request.body.len() > MAX_REQUEST_BODY {
268 return Err(RequestEncodeError::BodyTooLarge);
269 }
270 if target.len() > MAX_DEVICE_ID
271 || request.path.len() > u16::MAX as usize
272 || request.content_type.len() > u8::MAX as usize
273 || request.ticket.len() > MAX_REQUEST_TICKET
274 {
275 return Err(RequestEncodeError::FieldTooLong);
276 }
277 let mut object = Vec::with_capacity(
278 5 + request.path.len()
279 + request.content_type.len()
280 + request.ticket.len()
281 + request.body.len(),
282 );
283 object.extend_from_slice(&(request.path.len() as u16).to_be_bytes());
284 object.extend_from_slice(request.path);
285 object.push(request.content_type.len() as u8);
286 object.extend_from_slice(request.content_type);
287 object.extend_from_slice(&(request.ticket.len() as u16).to_be_bytes());
288 object.extend_from_slice(request.ticket);
289 object.extend_from_slice(request.body);
290
291 let chunks: Vec<&[u8]> = object.chunks(MAX_REQUEST_CHUNK).collect();
292 let total = chunks.len() as u16;
293 Ok(chunks
294 .iter()
295 .enumerate()
296 .map(|(seq, chunk)| {
297 let mut out =
298 Vec::with_capacity(RPC_REQUEST_FRAGMENT_OVERHEAD + target.len() + chunk.len());
299 out.push(method as u8);
300 out.extend_from_slice(&id.to_be_bytes());
301 out.push(target.len() as u8);
302 out.extend_from_slice(target);
303 out.extend_from_slice(&(seq as u16).to_be_bytes());
304 out.extend_from_slice(&total.to_be_bytes());
305 out.extend_from_slice(chunk);
306 out
307 })
308 .collect())
309}
310
311#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
313pub struct RequestParts<'a> {
314 pub path: &'a [u8],
315 pub content_type: &'a [u8],
316 pub ticket: &'a [u8],
317 pub body: &'a [u8],
318}
319
320pub fn decode_request_fragment(payload: &[u8]) -> Result<RequestFragment<'_>, RpcCodecError> {
322 if payload.len() < RPC_REQUEST_FRAGMENT_OVERHEAD {
323 return Err(RpcCodecError::TooShort);
324 }
325 let method = RpcMethod::from_u8(payload[0]).ok_or(RpcCodecError::BadMethod(payload[0]))?;
326 let id = u32::from_be_bytes([payload[1], payload[2], payload[3], payload[4]]);
327 let target_len = payload[5] as usize;
328 let seq_offset = 6 + target_len;
329 if seq_offset + 4 > payload.len() {
330 return Err(RpcCodecError::LengthMismatch {
331 declared: seq_offset + 4,
332 actual: payload.len(),
333 });
334 }
335 let target = &payload[6..seq_offset];
336 let seq = u16::from_be_bytes([payload[seq_offset], payload[seq_offset + 1]]);
337 let total = u16::from_be_bytes([payload[seq_offset + 2], payload[seq_offset + 3]]);
338 if total == 0 || seq >= total || total as usize > MAX_REQUEST_FRAGMENTS {
339 return Err(RpcCodecError::BadFragmentIndex { index: seq, total });
340 }
341 Ok(RequestFragment {
342 id,
343 method,
344 target,
345 seq,
346 total,
347 chunk: &payload[seq_offset + 4..],
348 })
349}
350
351fn decode_request_object(
353 id: u32,
354 method: RpcMethod,
355 target: &[u8],
356 object: &[u8],
357) -> Result<RpcRequest, RpcCodecError> {
358 let mut at = 0usize;
359 let take = |at: &mut usize, n: usize| -> Result<std::ops::Range<usize>, RpcCodecError> {
360 let end = at.checked_add(n).filter(|end| *end <= object.len()).ok_or(
361 RpcCodecError::LengthMismatch {
362 declared: at.saturating_add(n),
363 actual: object.len(),
364 },
365 )?;
366 let range = *at..end;
367 *at = end;
368 Ok(range)
369 };
370 let r = take(&mut at, 2)?;
371 let path_len = u16::from_be_bytes([object[r.start], object[r.start + 1]]) as usize;
372 let path = take(&mut at, path_len)?;
373 let r = take(&mut at, 1)?;
374 let ct_len = object[r.start] as usize;
375 let content_type = take(&mut at, ct_len)?;
376 let r = take(&mut at, 2)?;
377 let ticket_len = u16::from_be_bytes([object[r.start], object[r.start + 1]]) as usize;
378 if ticket_len > MAX_REQUEST_TICKET {
379 return Err(RpcCodecError::LengthMismatch {
380 declared: ticket_len,
381 actual: MAX_REQUEST_TICKET,
382 });
383 }
384 let ticket = take(&mut at, ticket_len)?;
385 let body = &object[at..];
386 if body.len() > MAX_REQUEST_BODY {
387 return Err(RpcCodecError::LengthMismatch {
388 declared: body.len(),
389 actual: MAX_REQUEST_BODY,
390 });
391 }
392 Ok(RpcRequest {
393 id,
394 method,
395 target: target.to_vec(),
396 path: object[path].to_vec(),
397 content_type: object[content_type].to_vec(),
398 ticket: object[ticket].to_vec(),
399 body: body.to_vec(),
400 })
401}
402
403#[derive(Debug, Clone, PartialEq, Eq)]
405pub enum ReassemblyOutcome {
406 Complete(RpcRequest),
408 Pending,
410 Duplicate,
412 Rejected,
415}
416
417struct PartialRequest {
418 method: RpcMethod,
419 target: Vec<u8>,
420 total: u16,
421 chunks: Vec<Option<Vec<u8>>>,
422 received: u16,
423 bytes: usize,
424 started: std::time::Instant,
425}
426
427#[derive(Default)]
433pub struct RequestReassembler {
434 partials: std::collections::HashMap<u32, PartialRequest>,
435 buffered: usize,
436}
437
438impl RequestReassembler {
439 pub fn new() -> Self {
440 Self::default()
441 }
442
443 pub fn push(
445 &mut self,
446 fragment: &RequestFragment<'_>,
447 now: std::time::Instant,
448 ) -> ReassemblyOutcome {
449 self.expire(now);
450 if fragment.total == 1 {
451 return match decode_request_object(
452 fragment.id,
453 fragment.method,
454 fragment.target,
455 fragment.chunk,
456 ) {
457 Ok(request) => ReassemblyOutcome::Complete(request),
458 Err(_) => ReassemblyOutcome::Rejected,
459 };
460 }
461 if fragment.chunk.len() > MAX_REQUEST_CHUNK {
462 return ReassemblyOutcome::Rejected;
463 }
464 let conflicting = self.partials.get(&fragment.id).is_some_and(|p| {
468 p.method != fragment.method || p.total != fragment.total || p.target != fragment.target
469 });
470 if conflicting {
471 self.drop_partial(fragment.id);
472 }
473 if !self.partials.contains_key(&fragment.id) {
474 if self.partials.len() >= MAX_PARTIAL_REQUESTS {
475 return ReassemblyOutcome::Rejected;
476 }
477 self.partials.insert(
478 fragment.id,
479 PartialRequest {
480 method: fragment.method,
481 target: fragment.target.to_vec(),
482 total: fragment.total,
483 chunks: vec![None; fragment.total as usize],
484 received: 0,
485 bytes: 0,
486 started: now,
487 },
488 );
489 }
490 let buffered = self.buffered;
491 let Some(partial) = self.partials.get_mut(&fragment.id) else {
492 return ReassemblyOutcome::Rejected;
493 };
494 if partial.chunks[fragment.seq as usize].is_some() {
495 return ReassemblyOutcome::Duplicate;
496 }
497 if buffered + fragment.chunk.len() > MAX_PARTIAL_REQUEST_BYTES
498 || partial.bytes + fragment.chunk.len() > MAX_REQUEST_OBJECT
499 {
500 self.drop_partial(fragment.id);
501 return ReassemblyOutcome::Rejected;
502 }
503 partial.chunks[fragment.seq as usize] = Some(fragment.chunk.to_vec());
504 partial.received += 1;
505 partial.bytes += fragment.chunk.len();
506 self.buffered += fragment.chunk.len();
507 if partial.received < partial.total {
508 return ReassemblyOutcome::Pending;
509 }
510 let Some(done) = self.partials.remove(&fragment.id) else {
511 return ReassemblyOutcome::Rejected;
512 };
513 self.buffered = self.buffered.saturating_sub(done.bytes);
514 let mut object = Vec::with_capacity(done.bytes);
515 for chunk in done.chunks.into_iter().flatten() {
516 object.extend_from_slice(&chunk);
517 }
518 match decode_request_object(fragment.id, done.method, &done.target, &object) {
519 Ok(request) => ReassemblyOutcome::Complete(request),
520 Err(_) => ReassemblyOutcome::Rejected,
521 }
522 }
523
524 pub fn partial_count(&self) -> usize {
526 self.partials.len()
527 }
528
529 fn drop_partial(&mut self, id: u32) {
530 if let Some(p) = self.partials.remove(&id) {
531 self.buffered = self.buffered.saturating_sub(p.bytes);
532 }
533 }
534
535 fn expire(&mut self, now: std::time::Instant) {
536 let stale: Vec<u32> = self
537 .partials
538 .iter()
539 .filter(|(_, p)| now.saturating_duration_since(p.started) > REQUEST_REASSEMBLY_TIMEOUT)
540 .map(|(id, _)| *id)
541 .collect();
542 for id in stale {
543 self.drop_partial(id);
544 }
545 }
546}
547
548#[cfg(test)]
549mod tests {
550 use super::*;
551 use std::time::{Duration, Instant};
552
553 fn parts<'a>(path: &'a [u8], body: &'a [u8]) -> RequestParts<'a> {
554 RequestParts {
555 path,
556 content_type: b"",
557 ticket: b"",
558 body,
559 }
560 }
561
562 fn round_trip(method: RpcMethod, id: u32, target: &[u8], req: &RequestParts<'_>) -> RpcRequest {
564 let frags = encode_request(method, id, target, req).unwrap();
565 let mut r = RequestReassembler::new();
566 let now = Instant::now();
567 let mut done = None;
568 for f in &frags {
569 if let ReassemblyOutcome::Complete(req) =
570 r.push(&decode_request_fragment(f).unwrap(), now)
571 {
572 done = Some(req);
573 }
574 }
575 done.expect("the request completes")
576 }
577
578 #[test]
579 fn a_small_request_is_one_fragment_and_round_trips() {
580 let req = RequestParts {
581 path: b"/api/config",
582 content_type: b"application/json",
583 ticket: b"v2|tok",
584 body: br#"{"key":"agent.name","value":"example-drone"}"#,
585 };
586 let frags = encode_request(RpcMethod::Put, 7, b"0a1b2c3d4e5f", &req).unwrap();
587 assert_eq!(frags.len(), 1);
588 assert!(frags[0].len() <= AUX_MAX_PAYLOAD);
589 let got = round_trip(RpcMethod::Put, 7, b"0a1b2c3d4e5f", &req);
590 assert_eq!(got.id, 7);
591 assert_eq!(got.method, RpcMethod::Put);
592 assert_eq!(got.target, b"0a1b2c3d4e5f");
593 assert_eq!(got.path, req.path);
594 assert_eq!(got.content_type, req.content_type);
595 assert_eq!(got.ticket, req.ticket);
596 assert_eq!(got.body, req.body);
597 }
598
599 #[test]
600 fn a_body_far_past_one_frame_fragments_and_reassembles_out_of_order() {
601 let body: Vec<u8> = (0..200_000u32).map(|i| (i % 251) as u8).collect();
603 let req = RequestParts {
604 path: b"/api/plugins/com.example.tool/x/upload",
605 content_type: b"application/octet-stream",
606 ticket: b"",
607 body: &body,
608 };
609 let frags = encode_request(RpcMethod::Patch, 9, b"d1", &req).unwrap();
610 assert!(frags.len() > 100);
611 assert!(frags.iter().all(|f| f.len() <= AUX_MAX_PAYLOAD));
612 let mut r = RequestReassembler::new();
613 let now = Instant::now();
614 let mut order: Vec<usize> = (0..frags.len()).rev().collect();
615 order.insert(5, order[3]);
617 let mut done = None;
618 for i in order {
619 match r.push(&decode_request_fragment(&frags[i]).unwrap(), now) {
620 ReassemblyOutcome::Complete(req) => done = Some(req),
621 ReassemblyOutcome::Pending | ReassemblyOutcome::Duplicate => {}
622 ReassemblyOutcome::Rejected => panic!("rejected"),
623 }
624 }
625 let got = done.expect("complete");
626 assert_eq!(got.method, RpcMethod::Patch);
627 assert_eq!(got.body, body);
628 assert_eq!(got.content_type, b"application/octet-stream");
629 assert_eq!(r.partial_count(), 0);
630 }
631
632 #[test]
633 fn a_body_over_one_mebibyte_is_refused_at_encode() {
634 let body = vec![0u8; MAX_REQUEST_BODY + 1];
635 assert_eq!(
636 encode_request(RpcMethod::Post, 1, b"d1", &parts(b"/api/x", &body)),
637 Err(RequestEncodeError::BodyTooLarge)
638 );
639 let body = vec![0u8; MAX_REQUEST_BODY];
640 assert!(encode_request(RpcMethod::Post, 1, b"d1", &parts(b"/api/x", &body)).is_ok());
641 }
642
643 #[test]
644 fn a_partial_request_is_dropped_after_the_reassembly_window() {
645 let body = vec![1u8; 5_000];
646 let frags = encode_request(RpcMethod::Post, 3, b"d1", &parts(b"/api/x", &body)).unwrap();
647 let mut r = RequestReassembler::new();
648 let t0 = Instant::now();
649 assert_eq!(
650 r.push(&decode_request_fragment(&frags[0]).unwrap(), t0),
651 ReassemblyOutcome::Pending
652 );
653 let later = t0 + REQUEST_REASSEMBLY_TIMEOUT + Duration::from_secs(1);
654 for f in &frags[1..] {
655 assert_eq!(
656 r.push(&decode_request_fragment(f).unwrap(), later),
657 ReassemblyOutcome::Pending,
658 "the stale first fragment is gone, so this cannot complete"
659 );
660 }
661 }
662
663 #[test]
664 fn concurrent_partials_are_bounded() {
665 let body = vec![1u8; 3_000];
666 let mut r = RequestReassembler::new();
667 let now = Instant::now();
668 for id in 0..MAX_PARTIAL_REQUESTS as u32 {
669 let frags =
670 encode_request(RpcMethod::Post, id, b"d1", &parts(b"/api/x", &body)).unwrap();
671 assert_eq!(
672 r.push(&decode_request_fragment(&frags[0]).unwrap(), now),
673 ReassemblyOutcome::Pending
674 );
675 }
676 let frags = encode_request(RpcMethod::Post, 999, b"d1", &parts(b"/api/x", &body)).unwrap();
677 assert_eq!(
678 r.push(&decode_request_fragment(&frags[0]).unwrap(), now),
679 ReassemblyOutcome::Rejected
680 );
681 }
682
683 #[test]
684 fn rejects_a_target_longer_than_a_device_id() {
685 let too_long = vec![b'a'; MAX_DEVICE_ID + 1];
686 assert_eq!(
687 encode_request(RpcMethod::Get, 1, &too_long, &parts(b"/api/x", b"")),
688 Err(RequestEncodeError::FieldTooLong)
689 );
690 let fits = vec![b'a'; MAX_DEVICE_ID];
691 assert_eq!(
692 round_trip(RpcMethod::Get, 1, &fits, &parts(b"/api/x", b"")).target,
693 fits
694 );
695 }
696
697 #[test]
698 fn rejects_a_truncated_fragment() {
699 let frags =
700 encode_request(RpcMethod::Post, 9, b"abc", &parts(b"/api/x", b"payload")).unwrap();
701 assert_eq!(
702 decode_request_fragment(&frags[0][..5]).unwrap_err(),
703 RpcCodecError::TooShort
704 );
705 }
706
707 #[test]
708 fn an_impossible_sequence_number_is_refused() {
709 let mut frag =
710 encode_request(RpcMethod::Get, 1, b"", &parts(b"/api/x", b"")).unwrap()[0].clone();
711 frag[6..8].copy_from_slice(&3u16.to_be_bytes());
713 assert_eq!(
714 decode_request_fragment(&frag).unwrap_err(),
715 RpcCodecError::BadFragmentIndex { index: 3, total: 1 }
716 );
717 }
718
719 #[test]
720 fn an_object_whose_lengths_lie_is_rejected() {
721 let mut frag =
722 encode_request(RpcMethod::Get, 1, b"", &parts(b"/api/x", b"")).unwrap()[0].clone();
723 frag[10..12].copy_from_slice(&500u16.to_be_bytes());
725 let mut r = RequestReassembler::new();
726 assert_eq!(
727 r.push(&decode_request_fragment(&frag).unwrap(), Instant::now()),
728 ReassemblyOutcome::Rejected
729 );
730 }
731
732 #[test]
733 fn rejects_an_unknown_method_byte() {
734 let mut frag =
735 encode_request(RpcMethod::Get, 1, b"", &parts(b"/api/x", b"")).unwrap()[0].clone();
736 frag[0] = 0xFF;
737 assert_eq!(
738 decode_request_fragment(&frag).unwrap_err(),
739 RpcCodecError::BadMethod(0xFF)
740 );
741 }
742
743 #[test]
744 fn method_numbers_are_pinned() {
745 assert_eq!(RpcMethod::Get as u8, 1);
747 assert_eq!(RpcMethod::Post as u8, 2);
748 assert_eq!(RpcMethod::Put as u8, 3);
749 assert_eq!(RpcMethod::Delete as u8, 4);
750 assert_eq!(RpcMethod::Patch as u8, 5);
751 assert_eq!(RpcMethod::from_u8(5), Some(RpcMethod::Patch));
752 assert_eq!(RpcMethod::from_u8(0), None);
753 for m in [
754 RpcMethod::Get,
755 RpcMethod::Post,
756 RpcMethod::Put,
757 RpcMethod::Delete,
758 RpcMethod::Patch,
759 ] {
760 assert_eq!(RpcMethod::from_http_method(m.as_http_method()), Some(m));
761 }
762 assert_eq!(RpcMethod::from_http_method("OPTIONS"), None);
763 }
764}