aws_smithy_eventstream/
frame.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Event Stream message frame types and serialization/deserialization logic.
7
8use 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;
23// String and byte array values are prefixed with a 2-byte length, but the protocol limits them to
24// `i16::MAX` rather than the `u16::MAX` that field could hold (matching the CRT, Java, and Go).
25const 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/// Deferred event stream signer to allow a signer to be wired up later.
42///
43/// HTTP request signing takes place after serialization, and the event stream
44/// message stream body is established during serialization. Since event stream
45/// signing may need context from the initial HTTP signing operation, this
46/// [`DeferredSigner`] is needed to wire up the signer later in the request lifecycle.
47///
48/// This signer basically just establishes a MPSC channel so that the sender can
49/// be placed in the request's config. Then the HTTP signer implementation can
50/// retrieve the sender from that config and send an actual signing implementation
51/// with all the context needed.
52///
53/// When an event stream implementation needs to sign a message, the first call to
54/// sign will acquire a signing implementation off of the channel and cache it
55/// for the remainder of the operation.
56#[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                    // Fall back to NoOpSigner when no signer is sent (e.g., server-side
82                    // event streams or tests that don't configure signing).
83                    .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
112/// Converts a Smithy modeled Event Stream type into a [`Message`].
113pub trait MarshallMessage: fmt::Debug {
114    /// Smithy modeled input type to convert from.
115    type Input;
116
117    fn marshall(&self, input: Self::Input) -> Result<Message, Error>;
118}
119
120/// A successfully unmarshalled message that is either an `Event` or an `Error`.
121#[derive(Debug)]
122pub enum UnmarshalledMessage<T, E> {
123    Event(T),
124    Error(E),
125}
126
127/// Converts an Event Stream [`Message`] into a Smithy modeled type.
128pub trait UnmarshallMessage: fmt::Debug {
129    /// Smithy modeled type to convert into.
130    type Output;
131    /// Smithy modeled error to convert into.
132    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
238/// Reads a header from the given `buffer`.
239fn 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
258/// Writes the header to the given `buffer`.
259fn 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
269/// Writes the given `headers` to a `buffer`.
270///
271/// If any header fails to serialize, nothing is written to `buffer`.
272pub 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
281// Returns (total_len, header_len)
282fn read_prelude_from<B: Buf>(mut buffer: B) -> Result<(u32, u32), Error> {
283    let mut crc_buffer = CrcBuf::new(&mut buffer);
284
285    // If the buffer doesn't have the entire, then error
286    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    // Validate the prelude
292    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    // The header length can be 0 or >= 2, but must fit within the frame size
298    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
304/// Reads a message from the given `buffer`. For streaming use cases, use
305/// the [`MessageFrameDecoder`] instead of this.
306pub 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    // Calculate a CRC as we go and read the prelude
312    let mut crc_buffer = CrcBuf::new(&mut buffer);
313    let (total_len, header_len) = read_prelude_from(&mut crc_buffer)?;
314
315    // Verify we have the full frame before continuing
316    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    // Read headers
324    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    // Read payload
336    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
348/// Writes the `message` to the given `buffer`.
349pub 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
383/// Returns `len` as the length prefix for a string or byte array header value.
384///
385/// This must error rather than truncate: the frame's lengths and CRCs are computed from the bytes
386/// actually written, so a frame with a truncated value length would still be valid, and a decoder
387/// would read the rest of the value as further headers.
388fn header_value_len(len: usize) -> Result<u16, Error> {
389    if len > MAX_HEADER_VALUE_LEN {
390        return Err(ErrorKind::HeaderValueTooLong.into());
391    }
392    // Lossless: `MAX_HEADER_VALUE_LEN` fits in a `u16`.
393    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        // Test message taken from the CRT:
478        // https://github.com/awslabs/aws-c-event-stream/blob/main/tests/message_deserializer_test.c
479        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        // Test message taken from the CRT:
495        // https://github.com/awslabs/aws-c-event-stream/blob/main/tests/message_deserializer_test.c
496        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    /// Asserts that serialization fails with `$err` and that nothing was written to the buffer.
584    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        // Well under the limit in characters, over it in UTF-8 bytes.
628        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        // The max-length value must come back as exactly one header. If its length had been
645        // truncated into the 16-bit field, the decoder would resume inside the value and parse
646        // the overflow as additional headers.
647        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        // One past what the 2-byte length field itself can hold. This is the size that wraps to an
660        // encoded length of zero in the implementations that don't bound the value at all.
661        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        // One past what the 2-byte length field itself can hold, which is the size that wraps to an
688        // encoded length of zero when the value isn't bounded at all.
689        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        // Well under the limit in characters, over it in UTF-8 bytes.
698        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    /// A value whose bytes spell out a well-formed header must not be encodable. At exactly 1 << 16
708    /// bytes the length field would hold 0, so a decoder would read an empty value and then parse
709    /// these bytes as the next header. The forged entry goes first because that is where the
710    /// decoder would resume.
711    #[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    /// A timestamp whose millisecond representation overflows `i64` is the one fixed-width value
728    /// that can fail to serialize, so it must error rather than write a wrong-width field.
729    #[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    /// Fixed-width values are unrepresentable at the wrong width on the write side (they are
741    /// `i8`/`i16`/`i32`/`i64`/`u128`), so the only place a short one can appear is an inbound
742    /// frame. Assert the decoder rejects those rather than under-reading the next header.
743    #[test]
744    fn truncated_fixed_width_header_values_are_rejected() {
745        // (value type, bytes supplied, protocol width)
746        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    /// `write_headers_to` is public and writes a header's name and value type before the value
777    /// length is validated, so a rejected header must not leave that prefix behind: a decoder would
778    /// read the orphaned name and type tag as a header whose value is whatever follows it.
779    #[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        // Seeded so the assertion catches an append as well as a bare partial write.
789        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/// A decoder written against the wire specification rather than by reusing [`read_message_from`],
801/// so that what [`write_message_to`] produces is checked by a second implementation.
802#[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    /// The headers and payload of a successfully decoded frame.
830    type Decoded = (Vec<(String, Value)>, Vec<u8>);
831
832    /// Decodes a frame per the specification, returning the headers and payload, or the first
833    /// invariant it violates.
834    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            // Every fixed-width type must find exactly its protocol width remaining.
876            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        // A value length that didn't match its bytes shows up here as leftover or overrun.
910        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    /// Establish that the decoder above is correct, using the reference message from the CRT's own
917    /// test suite, before relying on it to judge our encoder.
918    #[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    /// The longest permitted value must decode as a single header. A truncated length would leave
981    /// the decoder reading the remainder of the value as further headers, which is the smuggling
982    /// primitive this bound exists to prevent.
983    #[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    /// The same check over a byte array, whose length field is the one the attack overflows.
996    #[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/// Return value from [`MessageFrameDecoder`].
1012#[derive(Debug)]
1013pub enum DecodedFrame {
1014    /// There wasn't enough data in the buffer to decode a full message.
1015    Incomplete,
1016    /// There was enough data in the buffer to decode.
1017    Complete(Message),
1018}
1019
1020/// Streaming decoder for decoding a [`Message`] from a stream.
1021#[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    /// Returns a new `MessageFrameDecoder`.
1030    pub fn new() -> Self {
1031        Default::default()
1032    }
1033
1034    /// Determines if the `buffer` has enough data in it to read a full frame.
1035    /// Returns `Ok(None)` if there's not enough data, or `Some(remaining)` where
1036    /// `remaining` is the number of bytes after the prelude that belong to the
1037    /// message that's in the buffer.
1038    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    /// Resets the decoder.
1055    fn reset(&mut self) {
1056        self.prelude_read = false;
1057        self.prelude = [0u8; PRELUDE_LENGTH_BYTES_USIZE];
1058    }
1059
1060    /// Attempts to decode a [`Message`] from the given `buffer`. This function expects
1061    /// to be called over and over again with more data in the buffer each time its called.
1062    /// When there's not enough data to decode a message, it returns `Ok(None)`.
1063    ///
1064    /// Once there is enough data to read a message prelude, then it will mutate the `Buf`
1065    /// position. The state from the reading of the prelude is stored in the decoder so that
1066    /// the next call will be able to decode the entire message, even though the prelude
1067    /// is no longer available in the `Buf`.
1068    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}
aws_smithy_eventstream/frame.rs

aws_smithy_eventstream/
frame.rs
