aws_smithy_eventstream/
arbitrary.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Defines new-types wrapping inner types from `aws_smithy_types` to enable the `Arbitrary` trait
7//! for fuzz testing.
8
9use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
10use aws_smithy_types::str_bytes::StrBytes;
11use aws_smithy_types::DateTime;
12use bytes::Bytes;
13
14#[derive(Clone, Debug, PartialEq)]
15pub struct ArbHeaderValue(HeaderValue);
16
17impl<'a> arbitrary::Arbitrary<'a> for ArbHeaderValue {
18    fn arbitrary(unstruct: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
19        let value_type: u8 = unstruct.int_in_range(0..=9)?;
20        let header_value = match value_type {
21            crate::frame::TYPE_TRUE => HeaderValue::Bool(true),
22            crate::frame::TYPE_FALSE => HeaderValue::Bool(false),
23            crate::frame::TYPE_BYTE => HeaderValue::Byte(i8::arbitrary(unstruct)?),
24            crate::frame::TYPE_INT16 => HeaderValue::Int16(i16::arbitrary(unstruct)?),
25            crate::frame::TYPE_INT32 => HeaderValue::Int32(i32::arbitrary(unstruct)?),
26            crate::frame::TYPE_INT64 => HeaderValue::Int64(i64::arbitrary(unstruct)?),
27            crate::frame::TYPE_BYTE_ARRAY => {
28                HeaderValue::ByteArray(Bytes::from(Vec::<u8>::arbitrary(unstruct)?))
29            }
30            crate::frame::TYPE_STRING => {
31                HeaderValue::String(StrBytes::from(String::arbitrary(unstruct)?))
32            }
33            crate::frame::TYPE_TIMESTAMP => {
34                HeaderValue::Timestamp(DateTime::from_secs(i64::arbitrary(unstruct)?))
35            }
36            crate::frame::TYPE_UUID => HeaderValue::Uuid(u128::arbitrary(unstruct)?),
37            _ => unreachable!(),
38        };
39        Ok(ArbHeaderValue(header_value))
40    }
41}
42
43impl From<ArbHeaderValue> for HeaderValue {
44    fn from(header_value: ArbHeaderValue) -> Self {
45        header_value.0
46    }
47}
48
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub struct ArbStrBytes(StrBytes);
51
52#[cfg(feature = "derive-arbitrary")]
53impl<'a> arbitrary::Arbitrary<'a> for ArbStrBytes {
54    fn arbitrary(unstruct: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
55        Ok(ArbStrBytes(String::arbitrary(unstruct)?.into()))
56    }
57}
58
59impl From<ArbStrBytes> for StrBytes {
60    fn from(str_bytes: ArbStrBytes) -> Self {
61        str_bytes.0
62    }
63}
64
65#[derive(Clone, Debug, PartialEq, derive_arbitrary::Arbitrary)]
66pub struct ArbHeader {
67    name: ArbStrBytes,
68    value: ArbHeaderValue,
69}
70
71impl From<ArbHeader> for Header {
72    fn from(header: ArbHeader) -> Self {
73        Self::new(header.name, header.value)
74    }
75}
76
77#[derive(Clone, Debug, PartialEq)]
78pub struct ArbMessage(Message);
79
80impl<'a> arbitrary::Arbitrary<'a> for ArbMessage {
81    fn arbitrary(unstruct: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
82        let headers: Vec<ArbHeader> = unstruct
83            .arbitrary_iter()?
84            .collect::<arbitrary::Result<_>>()?;
85        let message = Message::new_from_parts(
86            headers.into_iter().map(Into::into).collect(),
87            Bytes::from(Vec::<u8>::arbitrary(unstruct)?),
88        );
89        Ok(ArbMessage(message))
90    }
91}
92
93impl From<ArbMessage> for Message {
94    fn from(message: ArbMessage) -> Self {
95        message.0
96    }
97}