1use crate::buf::count::CountBuf;
9use crate::buf::crc::{CrcBuf, CrcBufMut};
10use crate::error::{Error, ErrorKind};
11pub use aws_smithy_types::event_stream::DeferredSignerSender;
12use aws_smithy_types::event_stream::{DeferredSignerReceiver, Header, HeaderValue, Message};
13use aws_smithy_types::str_bytes::StrBytes;
14use aws_smithy_types::DateTime;
15use bytes::{Buf, BufMut};
16use std::fmt;
17use std::mem::size_of;
18
19const PRELUDE_LENGTH_BYTES: u32 = 3 * size_of::<u32>() as u32;
20const PRELUDE_LENGTH_BYTES_USIZE: usize = PRELUDE_LENGTH_BYTES as usize;
21const MESSAGE_CRC_LENGTH_BYTES: u32 = size_of::<u32>() as u32;
22const MAX_HEADER_NAME_LEN: usize = 255;
23const MAX_HEADER_VALUE_LEN: usize = i16::MAX as usize;
26const MIN_HEADER_LEN: usize = 2;
27
28pub(crate) const TYPE_TRUE: u8 = 0;
29pub(crate) const TYPE_FALSE: u8 = 1;
30pub(crate) const TYPE_BYTE: u8 = 2;
31pub(crate) const TYPE_INT16: u8 = 3;
32pub(crate) const TYPE_INT32: u8 = 4;
33pub(crate) const TYPE_INT64: u8 = 5;
34pub(crate) const TYPE_BYTE_ARRAY: u8 = 6;
35pub(crate) const TYPE_STRING: u8 = 7;
36pub(crate) const TYPE_TIMESTAMP: u8 = 8;
37pub(crate) const TYPE_UUID: u8 = 9;
38
39pub use aws_smithy_types::event_stream::{SignMessage, SignMessageError};
40
41#[derive(Debug)]
57pub struct DeferredSigner {
58 rx: Option<DeferredSignerReceiver>,
59 signer: Option<Box<dyn SignMessage + Send + Sync>>,
60}
61
62impl DeferredSigner {
63 pub fn new() -> (Self, DeferredSignerSender) {
64 let (rx, sender) = DeferredSignerSender::new();
65 (
66 Self {
67 rx: Some(rx),
68 signer: None,
69 },
70 sender,
71 )
72 }
73
74 fn acquire(&mut self) -> &mut (dyn SignMessage + Send + Sync) {
75 if self.signer.is_none() {
76 self.signer = Some(
77 self.rx
78 .take()
79 .expect("only taken once")
80 .recv::<Box<dyn SignMessage + Send + Sync>>()
81 .unwrap_or_else(|_| Box::new(NoOpSigner {}) as _),
84 );
85 }
86 self.signer.as_mut().unwrap().as_mut()
87 }
88}
89
90impl SignMessage for DeferredSigner {
91 fn sign(&mut self, message: Message) -> Result<Message, SignMessageError> {
92 self.acquire().sign(message)
93 }
94
95 fn sign_empty(&mut self) -> Option<Result<Message, SignMessageError>> {
96 self.acquire().sign_empty()
97 }
98}
99
100#[derive(Debug)]
101pub struct NoOpSigner {}
102impl SignMessage for NoOpSigner {
103 fn sign(&mut self, message: Message) -> Result<Message, SignMessageError> {
104 Ok(message)
105 }
106
107 fn sign_empty(&mut self) -> Option<Result<Message, SignMessageError>> {
108 None
109 }
110}
111
112pub trait MarshallMessage: fmt::Debug {
114 type Input;
116
117 fn marshall(&self, input: Self::Input) -> Result<Message, Error>;
118}
119
120#[derive(Debug)]
122pub enum UnmarshalledMessage<T, E> {
123 Event(T),
124 Error(E),
125}
126
127pub trait UnmarshallMessage: fmt::Debug {
129 type Output;
131 type Error;
133
134 fn unmarshall(
135 &self,
136 message: &Message,
137 ) -> Result<UnmarshalledMessage<Self::Output, Self::Error>, Error>;
138}
139
140macro_rules! read_value {
141 ($buf:ident, $typ:ident, $size_typ:ident, $read_fn:ident) => {
142 if $buf.remaining() >= size_of::<$size_typ>() {
143 Ok(HeaderValue::$typ($buf.$read_fn()))
144 } else {
145 Err(ErrorKind::InvalidHeaderValue.into())
146 }
147 };
148}
149
150fn read_header_value_from<B: Buf>(mut buffer: B) -> Result<HeaderValue, Error> {
151 let value_type = buffer.get_u8();
152 match value_type {
153 TYPE_TRUE => Ok(HeaderValue::Bool(true)),
154 TYPE_FALSE => Ok(HeaderValue::Bool(false)),
155 TYPE_BYTE => read_value!(buffer, Byte, i8, get_i8),
156 TYPE_INT16 => read_value!(buffer, Int16, i16, get_i16),
157 TYPE_INT32 => read_value!(buffer, Int32, i32, get_i32),
158 TYPE_INT64 => read_value!(buffer, Int64, i64, get_i64),
159 TYPE_BYTE_ARRAY | TYPE_STRING => {
160 if buffer.remaining() > size_of::<u16>() {
161 let len = buffer.get_u16() as usize;
162 if buffer.remaining() < len {
163 return Err(ErrorKind::InvalidHeaderValue.into());
164 }
165 let bytes = buffer.copy_to_bytes(len);
166 if value_type == TYPE_STRING {
167 Ok(HeaderValue::String(
168 bytes.try_into().map_err(|_| ErrorKind::InvalidUtf8String)?,
169 ))
170 } else {
171 Ok(HeaderValue::ByteArray(bytes))
172 }
173 } else {
174 Err(ErrorKind::InvalidHeaderValue.into())
175 }
176 }
177 TYPE_TIMESTAMP => {
178 if buffer.remaining() >= size_of::<i64>() {
179 let epoch_millis = buffer.get_i64();
180 Ok(HeaderValue::Timestamp(DateTime::from_millis(epoch_millis)))
181 } else {
182 Err(ErrorKind::InvalidHeaderValue.into())
183 }
184 }
185 TYPE_UUID => read_value!(buffer, Uuid, u128, get_u128),
186 _ => Err(ErrorKind::InvalidHeaderValueType(value_type).into()),
187 }
188}
189
190fn write_header_value_to<B: BufMut>(value: &HeaderValue, mut buffer: B) -> Result<(), Error> {
191 use HeaderValue::*;
192 match value {
193 Bool(val) => buffer.put_u8(if *val { TYPE_TRUE } else { TYPE_FALSE }),
194 Byte(val) => {
195 buffer.put_u8(TYPE_BYTE);
196 buffer.put_i8(*val);
197 }
198 Int16(val) => {
199 buffer.put_u8(TYPE_INT16);
200 buffer.put_i16(*val);
201 }
202 Int32(val) => {
203 buffer.put_u8(TYPE_INT32);
204 buffer.put_i32(*val);
205 }
206 Int64(val) => {
207 buffer.put_u8(TYPE_INT64);
208 buffer.put_i64(*val);
209 }
210 ByteArray(val) => {
211 buffer.put_u8(TYPE_BYTE_ARRAY);
212 buffer.put_u16(header_value_len(val.len())?);
213 buffer.put_slice(&val[..]);
214 }
215 String(val) => {
216 buffer.put_u8(TYPE_STRING);
217 buffer.put_u16(header_value_len(val.as_bytes().len())?);
218 buffer.put_slice(&val.as_bytes()[..]);
219 }
220 Timestamp(time) => {
221 buffer.put_u8(TYPE_TIMESTAMP);
222 buffer.put_i64(
223 time.to_millis()
224 .map_err(|_| ErrorKind::TimestampValueTooLarge(*time))?,
225 );
226 }
227 Uuid(val) => {
228 buffer.put_u8(TYPE_UUID);
229 buffer.put_u128(*val);
230 }
231 _ => {
232 panic!("matched on unexpected variant in `aws_smithy_types::event_stream::HeaderValue`")
233 }
234 }
235 Ok(())
236}
237
238fn read_header_from<B: Buf>(mut buffer: B) -> Result<(Header, usize), Error> {
240 if buffer.remaining() < MIN_HEADER_LEN {
241 return Err(ErrorKind::InvalidHeadersLength.into());
242 }
243
244 let mut counting_buf = CountBuf::new(&mut buffer);
245 let name_len = counting_buf.get_u8();
246 if name_len as usize >= counting_buf.remaining() {
247 return Err(ErrorKind::InvalidHeaderNameLength.into());
248 }
249
250 let name: StrBytes = counting_buf
251 .copy_to_bytes(name_len as usize)
252 .try_into()
253 .map_err(|_| ErrorKind::InvalidUtf8String)?;
254 let value = read_header_value_from(&mut counting_buf)?;
255 Ok((Header::new(name, value), counting_buf.into_count()))
256}
257
258fn write_header_to<B: BufMut>(header: &Header, mut buffer: B) -> Result<(), Error> {
260 if header.name().as_bytes().len() > MAX_HEADER_NAME_LEN {
261 return Err(ErrorKind::InvalidHeaderNameLength.into());
262 }
263
264 buffer.put_u8(u8::try_from(header.name().as_bytes().len()).expect("bounds check above"));
265 buffer.put_slice(&header.name().as_bytes()[..]);
266 write_header_value_to(header.value(), buffer)
267}
268
269pub fn write_headers_to<B: BufMut>(headers: &[Header], mut buffer: B) -> Result<(), Error> {
273 let mut staged = Vec::with_capacity(headers.len());
274 for header in headers {
275 write_header_to(header, &mut staged)?;
276 }
277 buffer.put_slice(&staged);
278 Ok(())
279}
280
281fn read_prelude_from<B: Buf>(mut buffer: B) -> Result<(u32, u32), Error> {
283 let mut crc_buffer = CrcBuf::new(&mut buffer);
284
285 let total_len = crc_buffer.get_u32();
287 if crc_buffer.remaining() + size_of::<u32>() < total_len as usize {
288 return Err(ErrorKind::InvalidMessageLength.into());
289 }
290
291 let header_len = crc_buffer.get_u32();
293 let (expected_crc, prelude_crc) = (crc_buffer.into_crc(), buffer.get_u32());
294 if expected_crc != prelude_crc {
295 return Err(ErrorKind::PreludeChecksumMismatch(expected_crc, prelude_crc).into());
296 }
297 if header_len == 1 || header_len > max_header_len(total_len)? {
299 return Err(ErrorKind::InvalidHeadersLength.into());
300 }
301 Ok((total_len, header_len))
302}
303
304pub fn read_message_from<B: Buf>(mut buffer: B) -> Result<Message, Error> {
307 if buffer.remaining() < PRELUDE_LENGTH_BYTES_USIZE {
308 return Err(ErrorKind::InvalidMessageLength.into());
309 }
310
311 let mut crc_buffer = CrcBuf::new(&mut buffer);
313 let (total_len, header_len) = read_prelude_from(&mut crc_buffer)?;
314
315 let remaining_len = total_len
317 .checked_sub(PRELUDE_LENGTH_BYTES)
318 .ok_or_else(|| Error::from(ErrorKind::InvalidMessageLength))?;
319 if crc_buffer.remaining() < remaining_len as usize {
320 return Err(ErrorKind::InvalidMessageLength.into());
321 }
322
323 let mut header_bytes_read = 0;
325 let mut headers = Vec::new();
326 while header_bytes_read < header_len as usize {
327 let (header, bytes_read) = read_header_from(&mut crc_buffer)?;
328 header_bytes_read += bytes_read;
329 if header_bytes_read > header_len as usize {
330 return Err(ErrorKind::InvalidHeaderValue.into());
331 }
332 headers.push(header);
333 }
334
335 let payload_len = payload_len(total_len, header_len)?;
337 let payload = crc_buffer.copy_to_bytes(payload_len as usize);
338
339 let expected_crc = crc_buffer.into_crc();
340 let message_crc = buffer.get_u32();
341 if expected_crc != message_crc {
342 return Err(ErrorKind::MessageChecksumMismatch(expected_crc, message_crc).into());
343 }
344
345 Ok(Message::new_from_parts(headers, payload))
346}
347
348pub fn write_message_to(message: &Message, buffer: &mut dyn BufMut) -> Result<(), Error> {
350 let mut headers = Vec::new();
351 for header in message.headers() {
352 write_header_to(header, &mut headers)?;
353 }
354
355 let headers_len = checked(headers.len(), ErrorKind::HeadersTooLong.into())?;
356 let payload_len = checked(message.payload().len(), ErrorKind::PayloadTooLong.into())?;
357 let message_len = [
358 PRELUDE_LENGTH_BYTES,
359 headers_len,
360 payload_len,
361 MESSAGE_CRC_LENGTH_BYTES,
362 ]
363 .iter()
364 .try_fold(0u32, |acc, v| {
365 acc.checked_add(*v)
366 .ok_or_else(|| Error::from(ErrorKind::MessageTooLong))
367 })?;
368
369 let mut crc_buffer = CrcBufMut::new(buffer);
370 crc_buffer.put_u32(message_len);
371 crc_buffer.put_u32(headers_len);
372 crc_buffer.put_crc();
373 crc_buffer.put(&headers[..]);
374 crc_buffer.put(&message.payload()[..]);
375 crc_buffer.put_crc();
376 Ok(())
377}
378
379fn checked<T: TryFrom<U>, U>(from: U, err: Error) -> Result<T, Error> {
380 T::try_from(from).map_err(|_| err)
381}
382
383fn header_value_len(len: usize) -> Result<u16, Error> {
389 if len > MAX_HEADER_VALUE_LEN {
390 return Err(ErrorKind::HeaderValueTooLong.into());
391 }
392 Ok(len as u16)
394}
395
396fn max_header_len(total_len: u32) -> Result<u32, Error> {
397 total_len
398 .checked_sub(PRELUDE_LENGTH_BYTES + MESSAGE_CRC_LENGTH_BYTES)
399 .ok_or_else(|| Error::from(ErrorKind::InvalidMessageLength))
400}
401
402fn payload_len(total_len: u32, header_len: u32) -> Result<u32, Error> {
403 total_len
404 .checked_sub(
405 header_len
406 .checked_add(PRELUDE_LENGTH_BYTES + MESSAGE_CRC_LENGTH_BYTES)
407 .ok_or_else(|| Error::from(ErrorKind::InvalidHeadersLength))?,
408 )
409 .ok_or_else(|| Error::from(ErrorKind::InvalidMessageLength))
410}
411
412#[cfg(test)]
413mod message_tests {
414 use super::{
415 read_message_from, MAX_HEADER_NAME_LEN, MAX_HEADER_VALUE_LEN, MESSAGE_CRC_LENGTH_BYTES,
416 PRELUDE_LENGTH_BYTES, TYPE_BYTE, TYPE_INT16, TYPE_INT32, TYPE_INT64, TYPE_STRING,
417 TYPE_TIMESTAMP, TYPE_UUID,
418 };
419 use crate::error::ErrorKind;
420 use crate::frame::{write_headers_to, write_message_to, Header, HeaderValue, Message};
421 use aws_smithy_types::str_bytes::StrBytes;
422 use aws_smithy_types::DateTime;
423 use bytes::Bytes;
424
425 macro_rules! read_message_expect_err {
426 ($bytes:expr, $err:pat) => {
427 let result = read_message_from(&mut Bytes::from_static($bytes));
428 let result = result.as_ref();
429 assert!(result.is_err(), "Expected error, got {:?}", result);
430 assert!(
431 matches!(result.err().unwrap().kind(), $err),
432 "Expected {}, got {:?}",
433 stringify!($err),
434 result
435 );
436 };
437 }
438
439 #[test]
440 fn invalid_messages() {
441 read_message_expect_err!(
442 include_bytes!("../test_data/invalid_header_string_value_length"),
443 ErrorKind::InvalidHeaderValue
444 );
445 read_message_expect_err!(
446 include_bytes!("../test_data/invalid_header_string_length_cut_off"),
447 ErrorKind::InvalidHeaderValue
448 );
449 read_message_expect_err!(
450 include_bytes!("../test_data/invalid_header_value_type"),
451 ErrorKind::InvalidHeaderValueType(0x60)
452 );
453 read_message_expect_err!(
454 include_bytes!("../test_data/invalid_header_name_length"),
455 ErrorKind::InvalidHeaderNameLength
456 );
457 read_message_expect_err!(
458 include_bytes!("../test_data/invalid_headers_length"),
459 ErrorKind::InvalidHeadersLength
460 );
461 read_message_expect_err!(
462 include_bytes!("../test_data/invalid_prelude_checksum"),
463 ErrorKind::PreludeChecksumMismatch(0x8BB495FB, 0xDEADBEEF)
464 );
465 read_message_expect_err!(
466 include_bytes!("../test_data/invalid_message_checksum"),
467 ErrorKind::MessageChecksumMismatch(0x01a05860, 0xDEADBEEF)
468 );
469 read_message_expect_err!(
470 include_bytes!("../test_data/invalid_header_name_length_too_long"),
471 ErrorKind::InvalidUtf8String
472 );
473 }
474
475 #[test]
476 fn read_message_no_headers() {
477 let data: &'static [u8] = &[
480 0x00, 0x00, 0x00, 0x1D, 0x00, 0x00, 0x00, 0x00, 0xfd, 0x52, 0x8c, 0x5a, 0x7b, 0x27,
481 0x66, 0x6f, 0x6f, 0x27, 0x3a, 0x27, 0x62, 0x61, 0x72, 0x27, 0x7d, 0xc3, 0x65, 0x39,
482 0x36,
483 ];
484
485 let result = read_message_from(&mut Bytes::from_static(data)).unwrap();
486 assert_eq!(result.headers(), Vec::new());
487
488 let expected_payload = b"{'foo':'bar'}";
489 assert_eq!(expected_payload, result.payload().as_ref());
490 }
491
492 #[test]
493 fn read_message_one_header() {
494 let data: &'static [u8] = &[
497 0x00, 0x00, 0x00, 0x3D, 0x00, 0x00, 0x00, 0x20, 0x07, 0xFD, 0x83, 0x96, 0x0C, b'c',
498 b'o', b'n', b't', b'e', b'n', b't', b'-', b't', b'y', b'p', b'e', 0x07, 0x00, 0x10,
499 b'a', b'p', b'p', b'l', b'i', b'c', b'a', b't', b'i', b'o', b'n', b'/', b'j', b's',
500 b'o', b'n', 0x7b, 0x27, 0x66, 0x6f, 0x6f, 0x27, 0x3a, 0x27, 0x62, 0x61, 0x72, 0x27,
501 0x7d, 0x8D, 0x9C, 0x08, 0xB1,
502 ];
503
504 let result = read_message_from(&mut Bytes::from_static(data)).unwrap();
505 assert_eq!(
506 result.headers(),
507 vec![Header::new(
508 "content-type",
509 HeaderValue::String("application/json".into())
510 )]
511 );
512
513 let expected_payload = b"{'foo':'bar'}";
514 assert_eq!(expected_payload, result.payload().as_ref());
515 }
516
517 #[test]
518 fn read_all_headers_and_payload() {
519 let message = include_bytes!("../test_data/valid_with_all_headers_and_payload");
520 let result = read_message_from(&mut Bytes::from_static(message)).unwrap();
521 assert_eq!(
522 result.headers(),
523 vec![
524 Header::new("true", HeaderValue::Bool(true)),
525 Header::new("false", HeaderValue::Bool(false)),
526 Header::new("byte", HeaderValue::Byte(50)),
527 Header::new("short", HeaderValue::Int16(20_000)),
528 Header::new("int", HeaderValue::Int32(500_000)),
529 Header::new("long", HeaderValue::Int64(50_000_000_000)),
530 Header::new(
531 "bytes",
532 HeaderValue::ByteArray(Bytes::from(&b"some bytes"[..]))
533 ),
534 Header::new("str", HeaderValue::String("some str".into())),
535 Header::new(
536 "time",
537 HeaderValue::Timestamp(DateTime::from_secs(5_000_000))
538 ),
539 Header::new(
540 "uuid",
541 HeaderValue::Uuid(0xb79bc914_de21_4e13_b8b2_bc47e85b7f0b)
542 ),
543 ]
544 );
545
546 assert_eq!(b"some payload", result.payload().as_ref());
547 }
548
549 #[test]
550 fn round_trip_all_headers_payload() {
551 let message = Message::new(&b"some payload"[..])
552 .add_header(Header::new("true", HeaderValue::Bool(true)))
553 .add_header(Header::new("false", HeaderValue::Bool(false)))
554 .add_header(Header::new("byte", HeaderValue::Byte(50)))
555 .add_header(Header::new("short", HeaderValue::Int16(20_000)))
556 .add_header(Header::new("int", HeaderValue::Int32(500_000)))
557 .add_header(Header::new("long", HeaderValue::Int64(50_000_000_000)))
558 .add_header(Header::new(
559 "bytes",
560 HeaderValue::ByteArray((&b"some bytes"[..]).into()),
561 ))
562 .add_header(Header::new("str", HeaderValue::String("some str".into())))
563 .add_header(Header::new(
564 "time",
565 HeaderValue::Timestamp(DateTime::from_secs(5_000_000)),
566 ))
567 .add_header(Header::new(
568 "uuid",
569 HeaderValue::Uuid(0xb79bc914_de21_4e13_b8b2_bc47e85b7f0b),
570 ));
571
572 let mut actual = Vec::new();
573 write_message_to(&message, &mut actual).unwrap();
574
575 let expected = include_bytes!("../test_data/valid_with_all_headers_and_payload").to_vec();
576 assert_eq!(expected, actual);
577
578 let result = read_message_from(&mut Bytes::from(actual)).unwrap();
579 assert_eq!(message.headers(), result.headers());
580 assert_eq!(message.payload().as_ref(), result.payload().as_ref());
581 }
582
583 macro_rules! write_message_expect_err {
585 ($message:expr, $err:pat) => {
586 let mut buffer = Vec::new();
587 let result = write_message_to(&$message, &mut buffer);
588 assert!(result.is_err(), "Expected error, got {:?}", result);
589 assert!(
590 matches!(result.as_ref().err().unwrap().kind(), $err),
591 "Expected {}, got {:?}",
592 stringify!($err),
593 result
594 );
595 assert!(
596 buffer.is_empty(),
597 "rejected message must not emit a partial frame, wrote {} bytes",
598 buffer.len()
599 );
600 };
601 }
602
603 fn message_with_header(name: impl Into<StrBytes>, value: HeaderValue) -> Message {
604 Message::new(&b"payload"[..]).add_header(Header::new(name, value))
605 }
606
607 #[test]
608 fn header_name_length_boundaries() {
609 let mut buffer = Vec::new();
610 let message = message_with_header("n".repeat(MAX_HEADER_NAME_LEN), HeaderValue::Bool(true));
611 write_message_to(&message, &mut buffer).unwrap();
612 assert_eq!(
613 message.headers(),
614 read_message_from(&mut Bytes::from(buffer))
615 .unwrap()
616 .headers()
617 );
618
619 write_message_expect_err!(
620 message_with_header("n".repeat(MAX_HEADER_NAME_LEN + 1), HeaderValue::Bool(true)),
621 ErrorKind::InvalidHeaderNameLength
622 );
623 }
624
625 #[test]
626 fn header_name_length_is_validated_in_bytes_not_chars() {
627 let name = "€".repeat(MAX_HEADER_NAME_LEN / 2);
629 assert!(name.chars().count() < MAX_HEADER_NAME_LEN);
630 assert!(name.len() > MAX_HEADER_NAME_LEN);
631 write_message_expect_err!(
632 message_with_header(name, HeaderValue::Bool(true)),
633 ErrorKind::InvalidHeaderNameLength
634 );
635 }
636
637 #[test]
638 fn string_header_value_length_boundaries() {
639 let longest = "a".repeat(MAX_HEADER_VALUE_LEN);
640 let mut buffer = Vec::new();
641 let message = message_with_header("h", HeaderValue::String(longest.clone().into()));
642 write_message_to(&message, &mut buffer).unwrap();
643
644 let decoded = read_message_from(&mut Bytes::from(buffer)).unwrap();
648 assert_eq!(1, decoded.headers().len());
649 assert_eq!(message.headers(), decoded.headers());
650
651 write_message_expect_err!(
652 message_with_header(
653 "h",
654 HeaderValue::String("a".repeat(MAX_HEADER_VALUE_LEN + 1).into())
655 ),
656 ErrorKind::HeaderValueTooLong
657 );
658
659 write_message_expect_err!(
662 message_with_header("h", HeaderValue::String("a".repeat(1 << 16).into())),
663 ErrorKind::HeaderValueTooLong
664 );
665 }
666
667 #[test]
668 fn byte_array_header_value_length_boundaries() {
669 let mut buffer = Vec::new();
670 let message = message_with_header(
671 "h",
672 HeaderValue::ByteArray(Bytes::from(vec![0u8; MAX_HEADER_VALUE_LEN])),
673 );
674 write_message_to(&message, &mut buffer).unwrap();
675 let decoded = read_message_from(&mut Bytes::from(buffer)).unwrap();
676 assert_eq!(1, decoded.headers().len());
677 assert_eq!(message.headers(), decoded.headers());
678
679 write_message_expect_err!(
680 message_with_header(
681 "h",
682 HeaderValue::ByteArray(Bytes::from(vec![0u8; MAX_HEADER_VALUE_LEN + 1]))
683 ),
684 ErrorKind::HeaderValueTooLong
685 );
686
687 write_message_expect_err!(
690 message_with_header("h", HeaderValue::ByteArray(Bytes::from(vec![0u8; 1 << 16]))),
691 ErrorKind::HeaderValueTooLong
692 );
693 }
694
695 #[test]
696 fn string_header_value_length_is_validated_in_bytes_not_chars() {
697 let value = "€".repeat(MAX_HEADER_VALUE_LEN / 2);
699 assert!(value.chars().count() < MAX_HEADER_VALUE_LEN);
700 assert!(value.len() > MAX_HEADER_VALUE_LEN);
701 write_message_expect_err!(
702 message_with_header("h", HeaderValue::String(value.into())),
703 ErrorKind::HeaderValueTooLong
704 );
705 }
706
707 #[test]
712 fn oversized_header_value_cannot_smuggle_a_header() {
713 let mut smuggled = Vec::new();
714 smuggled.push(b":event-type".len() as u8);
715 smuggled.extend_from_slice(b":event-type");
716 smuggled.push(TYPE_STRING);
717 smuggled.extend_from_slice(&(b"forged".len() as u16).to_be_bytes());
718 smuggled.extend_from_slice(b"forged");
719 smuggled.resize(1 << 16, b'a');
720
721 write_message_expect_err!(
722 message_with_header("DataType", HeaderValue::ByteArray(Bytes::from(smuggled))),
723 ErrorKind::HeaderValueTooLong
724 );
725 }
726
727 #[test]
730 fn write_message_rejects_timestamp_too_large() {
731 write_message_expect_err!(
732 message_with_header(
733 "time",
734 HeaderValue::Timestamp(DateTime::from_secs(i64::MAX))
735 ),
736 ErrorKind::TimestampValueTooLarge(_)
737 );
738 }
739
740 #[test]
744 fn truncated_fixed_width_header_values_are_rejected() {
745 let cases = [
747 (TYPE_BYTE, 0, 1),
748 (TYPE_INT16, 1, 2),
749 (TYPE_INT32, 3, 4),
750 (TYPE_INT64, 7, 8),
751 (TYPE_TIMESTAMP, 7, 8),
752 (TYPE_UUID, 15, 16),
753 ];
754 for (value_type, supplied, width) in cases {
755 let mut headers = vec![1u8, b'h', value_type];
756 headers.extend_from_slice(&vec![0u8; supplied]);
757
758 let headers_len = u32::try_from(headers.len()).unwrap();
759 let total_len = PRELUDE_LENGTH_BYTES + headers_len + MESSAGE_CRC_LENGTH_BYTES;
760 let mut frame = Vec::new();
761 frame.extend_from_slice(&total_len.to_be_bytes());
762 frame.extend_from_slice(&headers_len.to_be_bytes());
763 frame.extend_from_slice(&crc32fast::hash(&frame).to_be_bytes());
764 frame.extend_from_slice(&headers);
765 frame.extend_from_slice(&crc32fast::hash(&frame).to_be_bytes());
766
767 let result = read_message_from(&mut Bytes::from(frame));
768 assert!(
769 result.is_err(),
770 "value type {value_type} with {supplied} of {width} bytes must be rejected, \
771 got {result:?}"
772 );
773 }
774 }
775
776 #[test]
780 fn write_headers_to_leaves_no_partial_header() {
781 let headers = [
782 Header::new("ok", HeaderValue::String("fine".into())),
783 Header::new(
784 "toolong",
785 HeaderValue::String("a".repeat(MAX_HEADER_VALUE_LEN + 1).into()),
786 ),
787 ];
788 let mut buffer = b"PREVIOUS-CONTENT".to_vec();
790 let result = write_headers_to(&headers, &mut buffer);
791 assert!(result.is_err(), "Expected error, got {result:?}");
792 assert_eq!(
793 b"PREVIOUS-CONTENT".to_vec(),
794 buffer,
795 "no bytes may be written when a header is rejected"
796 );
797 }
798}
799
800#[cfg(test)]
803mod independent_decoder_tests {
804 use super::write_message_to;
805 use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
806 use aws_smithy_types::DateTime;
807 use bytes::Bytes;
808
809 #[derive(Debug, PartialEq)]
810 enum Value {
811 Bool(bool),
812 Byte(i8),
813 Int16(i16),
814 Int32(i32),
815 Int64(i64),
816 Bytes(Vec<u8>),
817 Str(String),
818 Timestamp(i64),
819 Uuid(u128),
820 }
821
822 fn u16_at(b: &[u8], i: usize) -> usize {
823 u16::from_be_bytes([b[i], b[i + 1]]) as usize
824 }
825 fn u32_at(b: &[u8], i: usize) -> usize {
826 u32::from_be_bytes([b[i], b[i + 1], b[i + 2], b[i + 3]]) as usize
827 }
828
829 type Decoded = (Vec<(String, Value)>, Vec<u8>);
831
832 fn decode(frame: &[u8]) -> Result<Decoded, String> {
835 if frame.len() < 16 {
836 return Err("shorter than prelude plus message crc".into());
837 }
838 if u32_at(frame, 0) != frame.len() {
839 return Err(format!(
840 "total_length {} != actual {}",
841 u32_at(frame, 0),
842 frame.len()
843 ));
844 }
845 if crc32fast::hash(&frame[..8]) != u32_at(frame, 8) as u32 {
846 return Err("prelude crc mismatch".into());
847 }
848 let crc_start = frame.len() - 4;
849 if crc32fast::hash(&frame[..crc_start]) != u32_at(frame, crc_start) as u32 {
850 return Err("message crc mismatch".into());
851 }
852
853 let headers_len = u32_at(frame, 4);
854 if 12 + headers_len + 4 > frame.len() {
855 return Err("headers_length overruns the frame".into());
856 }
857 let block = &frame[12..12 + headers_len];
858
859 let mut headers = Vec::new();
860 let mut i = 0;
861 while i < block.len() {
862 let name_len = block[i] as usize;
863 i += 1;
864 if name_len == 0 || i + name_len > block.len() {
865 return Err(format!("bad header name length {name_len}"));
866 }
867 let name = String::from_utf8(block[i..i + name_len].to_vec())
868 .map_err(|_| "header name is not utf-8")?;
869 i += name_len;
870 if i >= block.len() {
871 return Err("header has no value type".into());
872 }
873 let value_type = block[i];
874 i += 1;
875 let mut take = |n: usize| -> Result<&[u8], String> {
877 if i + n > block.len() {
878 return Err(format!(
879 "type {value_type} needs {n} bytes, {} left",
880 block.len() - i
881 ));
882 }
883 let out = &block[i..i + n];
884 i += n;
885 Ok(out)
886 };
887 let value = match value_type {
888 0 => Value::Bool(true),
889 1 => Value::Bool(false),
890 2 => Value::Byte(take(1)?[0] as i8),
891 3 => Value::Int16(i16::from_be_bytes(take(2)?.try_into().unwrap())),
892 4 => Value::Int32(i32::from_be_bytes(take(4)?.try_into().unwrap())),
893 5 => Value::Int64(i64::from_be_bytes(take(8)?.try_into().unwrap())),
894 6 | 7 => {
895 let len = u16_at(take(2)?, 0);
896 let raw = take(len)?.to_vec();
897 if value_type == 7 {
898 Value::Str(String::from_utf8(raw).map_err(|_| "value is not utf-8")?)
899 } else {
900 Value::Bytes(raw)
901 }
902 }
903 8 => Value::Timestamp(i64::from_be_bytes(take(8)?.try_into().unwrap())),
904 9 => Value::Uuid(u128::from_be_bytes(take(16)?.try_into().unwrap())),
905 other => return Err(format!("unknown value type {other}")),
906 };
907 headers.push((name, value));
908 }
909 if i != block.len() {
911 return Err("header block was not consumed exactly".into());
912 }
913 Ok((headers, frame[12 + headers_len..crc_start].to_vec()))
914 }
915
916 #[test]
919 fn agrees_with_the_crt_reference_message() {
920 let data: &[u8] = &[
921 0x00, 0x00, 0x00, 0x3D, 0x00, 0x00, 0x00, 0x20, 0x07, 0xFD, 0x83, 0x96, 0x0C, b'c',
922 b'o', b'n', b't', b'e', b'n', b't', b'-', b't', b'y', b'p', b'e', 0x07, 0x00, 0x10,
923 b'a', b'p', b'p', b'l', b'i', b'c', b'a', b't', b'i', b'o', b'n', b'/', b'j', b's',
924 b'o', b'n', 0x7b, 0x27, 0x66, 0x6f, 0x6f, 0x27, 0x3a, 0x27, 0x62, 0x61, 0x72, 0x27,
925 0x7d, 0x8D, 0x9C, 0x08, 0xB1,
926 ];
927 let (headers, payload) = decode(data).expect("the reference message must decode");
928 assert_eq!(
929 vec![(
930 "content-type".to_string(),
931 Value::Str("application/json".into())
932 )],
933 headers
934 );
935 assert_eq!(b"{'foo':'bar'}".to_vec(), payload);
936 }
937
938 #[test]
939 fn every_header_type_survives_a_round_trip() {
940 let message = Message::new(&b"some payload"[..])
941 .add_header(Header::new("true", HeaderValue::Bool(true)))
942 .add_header(Header::new("false", HeaderValue::Bool(false)))
943 .add_header(Header::new("byte", HeaderValue::Byte(i8::MIN)))
944 .add_header(Header::new("short", HeaderValue::Int16(i16::MIN)))
945 .add_header(Header::new("int", HeaderValue::Int32(i32::MIN)))
946 .add_header(Header::new("long", HeaderValue::Int64(i64::MIN)))
947 .add_header(Header::new(
948 "bytes",
949 HeaderValue::ByteArray((&b"some bytes"[..]).into()),
950 ))
951 .add_header(Header::new("str", HeaderValue::String("some str".into())))
952 .add_header(Header::new(
953 "time",
954 HeaderValue::Timestamp(DateTime::from_secs(5_000_000)),
955 ))
956 .add_header(Header::new("uuid", HeaderValue::Uuid(u128::MAX)));
957
958 let mut buffer = Vec::new();
959 write_message_to(&message, &mut buffer).unwrap();
960 let (headers, payload) = decode(&buffer).expect("our output must satisfy the spec");
961
962 assert_eq!(
963 vec![
964 ("true".to_string(), Value::Bool(true)),
965 ("false".to_string(), Value::Bool(false)),
966 ("byte".to_string(), Value::Byte(i8::MIN)),
967 ("short".to_string(), Value::Int16(i16::MIN)),
968 ("int".to_string(), Value::Int32(i32::MIN)),
969 ("long".to_string(), Value::Int64(i64::MIN)),
970 ("bytes".to_string(), Value::Bytes(b"some bytes".to_vec())),
971 ("str".to_string(), Value::Str("some str".into())),
972 ("time".to_string(), Value::Timestamp(5_000_000_000)),
973 ("uuid".to_string(), Value::Uuid(u128::MAX)),
974 ],
975 headers
976 );
977 assert_eq!(b"some payload".to_vec(), payload);
978 }
979
980 #[test]
984 fn longest_permitted_value_decodes_as_one_header() {
985 let value = "a".repeat(super::MAX_HEADER_VALUE_LEN);
986 let message = Message::new(&b""[..])
987 .add_header(Header::new("h", HeaderValue::String(value.clone().into())));
988 let mut buffer = Vec::new();
989 write_message_to(&message, &mut buffer).unwrap();
990
991 let (headers, _) = decode(&buffer).expect("the longest permitted value must decode");
992 assert_eq!(vec![("h".to_string(), Value::Str(value))], headers);
993 }
994
995 #[test]
997 fn longest_permitted_byte_array_decodes_as_one_header() {
998 let value = vec![b'z'; super::MAX_HEADER_VALUE_LEN];
999 let message = Message::new(&b""[..]).add_header(Header::new(
1000 "h",
1001 HeaderValue::ByteArray(Bytes::from(value.clone())),
1002 ));
1003 let mut buffer = Vec::new();
1004 write_message_to(&message, &mut buffer).unwrap();
1005
1006 let (headers, _) = decode(&buffer).expect("the longest permitted value must decode");
1007 assert_eq!(vec![("h".to_string(), Value::Bytes(value))], headers);
1008 }
1009}
1010
1011#[derive(Debug)]
1013pub enum DecodedFrame {
1014 Incomplete,
1016 Complete(Message),
1018}
1019
1020#[non_exhaustive]
1022#[derive(Default, Debug)]
1023pub struct MessageFrameDecoder {
1024 prelude: [u8; PRELUDE_LENGTH_BYTES_USIZE],
1025 prelude_read: bool,
1026}
1027
1028impl MessageFrameDecoder {
1029 pub fn new() -> Self {
1031 Default::default()
1032 }
1033
1034 fn remaining_bytes_if_frame_available<B: Buf>(
1039 &self,
1040 buffer: &B,
1041 ) -> Result<Option<usize>, Error> {
1042 if self.prelude_read {
1043 let remaining_len = (&self.prelude[..])
1044 .get_u32()
1045 .checked_sub(PRELUDE_LENGTH_BYTES)
1046 .ok_or_else(|| Error::from(ErrorKind::InvalidMessageLength))?;
1047 if buffer.remaining() >= remaining_len as usize {
1048 return Ok(Some(remaining_len as usize));
1049 }
1050 }
1051 Ok(None)
1052 }
1053
1054 fn reset(&mut self) {
1056 self.prelude_read = false;
1057 self.prelude = [0u8; PRELUDE_LENGTH_BYTES_USIZE];
1058 }
1059
1060 pub fn decode_frame<B: Buf>(&mut self, mut buffer: B) -> Result<DecodedFrame, Error> {
1069 if !self.prelude_read && buffer.remaining() >= PRELUDE_LENGTH_BYTES_USIZE {
1070 buffer.copy_to_slice(&mut self.prelude);
1071 self.prelude_read = true;
1072 }
1073
1074 if let Some(remaining_len) = self.remaining_bytes_if_frame_available(&buffer)? {
1075 let mut message_buf = (&self.prelude[..]).chain(buffer.take(remaining_len));
1076 let result = read_message_from(&mut message_buf).map(DecodedFrame::Complete);
1077 self.reset();
1078 return result;
1079 }
1080
1081 Ok(DecodedFrame::Incomplete)
1082 }
1083}
1084
1085#[cfg(test)]
1086mod message_frame_decoder_tests {
1087 use super::{DecodedFrame, MessageFrameDecoder};
1088 use crate::frame::read_message_from;
1089 use bytes::Bytes;
1090 use bytes_utils::SegmentedBuf;
1091
1092 #[test]
1093 fn single_streaming_message() {
1094 let message = include_bytes!("../test_data/valid_with_all_headers_and_payload");
1095
1096 let mut decoder = MessageFrameDecoder::new();
1097 let mut segmented = SegmentedBuf::new();
1098 for i in 0..(message.len() - 1) {
1099 segmented.push(&message[i..(i + 1)]);
1100 if let DecodedFrame::Complete(_) = decoder.decode_frame(&mut segmented).unwrap() {
1101 panic!("incomplete frame shouldn't result in message");
1102 }
1103 }
1104
1105 segmented.push(&message[(message.len() - 1)..]);
1106 match decoder.decode_frame(&mut segmented).unwrap() {
1107 DecodedFrame::Incomplete => panic!("frame should be complete now"),
1108 DecodedFrame::Complete(actual) => {
1109 let expected = read_message_from(&mut Bytes::from_static(message)).unwrap();
1110 assert_eq!(expected, actual);
1111 }
1112 }
1113 }
1114
1115 fn multiple_streaming_messages_chunk_size(chunk_size: usize) {
1116 let message1 = include_bytes!("../test_data/valid_with_all_headers_and_payload");
1117 let message2 = include_bytes!("../test_data/valid_empty_payload");
1118 let message3 = include_bytes!("../test_data/valid_no_headers");
1119 let mut repeated = message1.to_vec();
1120 repeated.extend_from_slice(message2);
1121 repeated.extend_from_slice(message3);
1122
1123 let mut decoder = MessageFrameDecoder::new();
1124 let mut segmented = SegmentedBuf::new();
1125 let mut decoded = Vec::new();
1126 for window in repeated.chunks(chunk_size) {
1127 segmented.push(window);
1128 match dbg!(decoder.decode_frame(&mut segmented)).unwrap() {
1129 DecodedFrame::Incomplete => {}
1130 DecodedFrame::Complete(message) => {
1131 decoded.push(message);
1132 }
1133 }
1134 }
1135
1136 let expected1 = read_message_from(&mut Bytes::from_static(message1)).unwrap();
1137 let expected2 = read_message_from(&mut Bytes::from_static(message2)).unwrap();
1138 let expected3 = read_message_from(&mut Bytes::from_static(message3)).unwrap();
1139 assert_eq!(3, decoded.len());
1140 assert_eq!(expected1, decoded[0]);
1141 assert_eq!(expected2, decoded[1]);
1142 assert_eq!(expected3, decoded[2]);
1143 }
1144
1145 #[test]
1146 fn multiple_streaming_messages() {
1147 for chunk_size in 1..=11 {
1148 println!("chunk size: {chunk_size}");
1149 multiple_streaming_messages_chunk_size(chunk_size);
1150 }
1151 }
1152}
1153
1154#[cfg(test)]
1155mod deferred_signer_tests {
1156 use crate::frame::{DeferredSigner, Header, HeaderValue, Message, SignMessage};
1157 use bytes::Bytes;
1158
1159 fn check_send_sync<T: Send + Sync>(value: T) -> T {
1160 value
1161 }
1162
1163 #[test]
1164 fn deferred_signer() {
1165 #[derive(Default, Debug)]
1166 struct TestSigner {
1167 call_num: i32,
1168 }
1169 impl SignMessage for TestSigner {
1170 fn sign(
1171 &mut self,
1172 message: Message,
1173 ) -> Result<Message, crate::frame::SignMessageError> {
1174 self.call_num += 1;
1175 Ok(message.add_header(Header::new("call_num", HeaderValue::Int32(self.call_num))))
1176 }
1177
1178 fn sign_empty(&mut self) -> Option<Result<Message, crate::frame::SignMessageError>> {
1179 None
1180 }
1181 }
1182
1183 let (mut signer, sender) = check_send_sync(DeferredSigner::new());
1184
1185 sender
1186 .send(Box::<TestSigner>::default() as Box<dyn SignMessage + Send + Sync>)
1187 .expect("success");
1188
1189 let message = signer.sign(Message::new(Bytes::new())).expect("success");
1190 assert_eq!(1, message.headers()[0].value().as_int32().unwrap());
1191
1192 let message = signer.sign(Message::new(Bytes::new())).expect("success");
1193 assert_eq!(2, message.headers()[0].value().as_int32().unwrap());
1194
1195 assert!(signer.sign_empty().is_none());
1196 }
1197
1198 #[test]
1199 fn deferred_signer_defaults_to_noop_signer() {
1200 let (mut signer, _sender) = DeferredSigner::new();
1201 assert_eq!(
1202 Message::new(Bytes::new()),
1203 signer.sign(Message::new(Bytes::new())).unwrap()
1204 );
1205 assert!(signer.sign_empty().is_none());
1206 }
1207}