AWS SDK

/*

 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 * SPDX-License-Identifier: Apache-2.0

 */

use aws_smithy_eventstream::frame::{

    DecodedFrame, MessageFrameDecoder, UnmarshallMessage, UnmarshalledMessage,

};

use aws_smithy_runtime_api::client::result::{ConnectorError, SdkError};

use aws_smithy_types::body::SdkBody;

use aws_smithy_types::event_stream::{Message, RawMessage};

use bytes::Buf;

use bytes::Bytes;

use bytes_utils::SegmentedBuf;

use std::error::Error as StdError;

use std::fmt;

use std::marker::PhantomData;

use std::mem;

use tracing::trace;

/// Wrapper around SegmentedBuf that tracks the state of the stream.

#[derive(Debug)]

enum RecvBuf {

    /// Nothing has been buffered yet.

    Empty,

    /// Some data has been buffered.

    /// The SegmentedBuf will automatically purge when it reads off the end of a chunk boundary.

    Partial(SegmentedBuf<Bytes>),

    /// The end of the stream has been reached, but there may still be some buffered data.

    EosPartial(SegmentedBuf<Bytes>),

    /// An exception terminated this stream.

    Terminated,

}

impl RecvBuf {

    /// Returns true if there's more buffered data.

    fn has_data(&self) -> bool {

        match self {

            RecvBuf::Empty | RecvBuf::Terminated => false,

            RecvBuf::Partial(segments) | RecvBuf::EosPartial(segments) => segments.remaining() > 0,

        }

    }

    /// Returns true if the stream has ended.

    fn is_eos(&self) -> bool {

        matches!(self, RecvBuf::EosPartial(_) | RecvBuf::Terminated)

    }

    /// Returns a mutable reference to the underlying buffered data.

    fn buffered(&mut self) -> &mut SegmentedBuf<Bytes> {

        match self {

            RecvBuf::Empty => panic!("buffer must be populated before reading; this is a bug"),

            RecvBuf::Partial(segmented) => segmented,

            RecvBuf::EosPartial(segmented) => segmented,

            RecvBuf::Terminated => panic!("buffer has been terminated; this is a bug"),

        }

    }

    /// Returns a new `RecvBuf` with additional data buffered. This will only allocate

    /// if the `RecvBuf` was previously empty.

    fn with_partial(self, partial: Bytes) -> Self {

        match self {

            RecvBuf::Empty => {

                let mut segmented = SegmentedBuf::new();

                segmented.push(partial);

                RecvBuf::Partial(segmented)

            }

            RecvBuf::Partial(mut segmented) => {

                segmented.push(partial);

                RecvBuf::Partial(segmented)

            }

            RecvBuf::EosPartial(_) | RecvBuf::Terminated => {

                panic!("cannot buffer more data after the stream has ended or been terminated; this is a bug")

            }

        }

    }

    /// Returns a `RecvBuf` that has reached end of stream.

    fn ended(self) -> Self {

        match self {

            RecvBuf::Empty => RecvBuf::EosPartial(SegmentedBuf::new()),

            RecvBuf::Partial(segmented) => RecvBuf::EosPartial(segmented),

            RecvBuf::EosPartial(_) => panic!("already end of stream; this is a bug"),

            RecvBuf::Terminated => panic!("stream terminated; this is a bug"),

        }

    }

}

#[derive(Debug)]

enum ReceiverErrorKind {

    /// The stream ended before a complete message frame was received.

    UnexpectedEndOfStream,

}

/// An error that occurs within an event stream receiver.

#[derive(Debug)]

pub struct ReceiverError {

    kind: ReceiverErrorKind,

}

impl fmt::Display for ReceiverError {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        match self.kind {

            ReceiverErrorKind::UnexpectedEndOfStream => write!(f, "unexpected end of stream"),

        }

    }

}

impl StdError for ReceiverError {}

/// Receives Smithy-modeled messages out of an Event Stream.

#[derive(Debug)]

pub struct Receiver<T, E> {

    unmarshaller: Box<dyn UnmarshallMessage<Output = T, Error = E> + Send + Sync>,

    decoder: MessageFrameDecoder,

    buffer: RecvBuf,

    body: SdkBody,

    /// Event Stream has optional initial response frames an with `:message-type` of

    /// `initial-response`. If `try_recv_initial()` is called and the next message isn't an

    /// initial response, then the message will be stored in `buffered_message` so that it can

    /// be returned with the next call of `recv()`.

    buffered_message: Option<Message>,

    _phantom: PhantomData<E>,

}

// Used by `Receiver::try_recv_initial`, hence this enum is also doc hidden

#[doc(hidden)]

#[non_exhaustive]

pub enum InitialMessageType {

    Request,

    Response,

}

impl InitialMessageType {

    fn as_str(&self) -> &'static str {

        match self {

            InitialMessageType::Request => "initial-request",

            InitialMessageType::Response => "initial-response",

        }

    }

}

impl<T, E> Receiver<T, E> {

    /// Creates a new `Receiver` with the given message unmarshaller and SDK body.

    pub fn new(

        unmarshaller: impl UnmarshallMessage<Output = T, Error = E> + Send + Sync + 'static,

        body: SdkBody,

    ) -> Self {

        Receiver {

            unmarshaller: Box::new(unmarshaller),

            decoder: MessageFrameDecoder::new(),

            buffer: RecvBuf::Empty,

            body,

            buffered_message: None,

            _phantom: Default::default(),

        }

    }

    fn unmarshall(&self, message: Message) -> Result<Option<T>, SdkError<E, RawMessage>> {

        match self.unmarshaller.unmarshall(&message) {

            Ok(unmarshalled) => match unmarshalled {

                UnmarshalledMessage::Event(event) => Ok(Some(event)),

                UnmarshalledMessage::Error(err) => {

                    Err(SdkError::service_error(err, RawMessage::Decoded(message)))

                }

            },

            Err(err) => Err(SdkError::response_error(err, RawMessage::Decoded(message))),

        }

    }

    async fn buffer_next_chunk(&mut self) -> Result<(), SdkError<E, RawMessage>> {

        use http_body_04x::Body;

        if !self.buffer.is_eos() {

            let next_chunk = self

                .body

                .data()

                .await

                .transpose()

                .map_err(|err| SdkError::dispatch_failure(ConnectorError::io(err)))?;

            let buffer = mem::replace(&mut self.buffer, RecvBuf::Empty);

            if let Some(chunk) = next_chunk {

                self.buffer = buffer.with_partial(chunk);

            } else {

                self.buffer = buffer.ended();

            }

        }

        Ok(())

    }

    async fn next_message(&mut self) -> Result<Option<Message>, SdkError<E, RawMessage>> {

        while !self.buffer.is_eos() {

            if self.buffer.has_data() {

                if let DecodedFrame::Complete(message) = self

                    .decoder

                    .decode_frame(self.buffer.buffered())

                    .map_err(|err| {

                        SdkError::response_error(

                            err,

                            // the buffer has been consumed

                            RawMessage::Invalid(None),

                        )

                    })?

                {

                    trace!(message = ?message, "received complete event stream message");

                    return Ok(Some(message));

                }

            }

            self.buffer_next_chunk().await?;

        }

        if self.buffer.has_data() {

            trace!(remaining_data = ?self.buffer, "data left over in the event stream response stream");

            let buf = self.buffer.buffered();

            return Err(SdkError::response_error(

                ReceiverError {

                    kind: ReceiverErrorKind::UnexpectedEndOfStream,

                },

                RawMessage::invalid(Some(buf.copy_to_bytes(buf.remaining()))),

            ));

        }

        Ok(None)

    }

    /// Tries to receive the initial response message that has `:event-type` of a given `message_type`.

    /// If a different event type is received, then it is buffered and `Ok(None)` is returned.

    #[doc(hidden)]

    pub async fn try_recv_initial(

        &mut self,

        message_type: InitialMessageType,

    ) -> Result<Option<Message>, SdkError<E, RawMessage>> {

        if let Some(message) = self.next_message().await? {

            if let Some(event_type) = message

                .headers()

                .iter()

                .find(|h| h.name().as_str() == ":event-type")

            {

                if event_type

                    .value()

                    .as_string()

                    .map(|s| s.as_str() == message_type.as_str())

                    .unwrap_or(false)

                {

                    return Ok(Some(message));

                }

            }

            // Buffer the message so that it can be returned by the next call to `recv()`

            self.buffered_message = Some(message);

        }

        Ok(None)

    }

    /// Asynchronously tries to receive a message from the stream. If the stream has ended,

    /// it returns an `Ok(None)`. If there is a transport layer error, it will return

    /// `Err(SdkError::DispatchFailure)`. Service-modeled errors will be a part of the returned

    /// messages.

    pub async fn recv(&mut self) -> Result<Option<T>, SdkError<E, RawMessage>> {

        if let Some(buffered) = self.buffered_message.take() {

            return match self.unmarshall(buffered) {

                Ok(message) => Ok(message),

                Err(error) => {

                    self.buffer = RecvBuf::Terminated;

                    Err(error)

                }

            };

        }

        if let Some(message) = self.next_message().await? {

            match self.unmarshall(message) {

                Ok(message) => Ok(message),

                Err(error) => {

                    self.buffer = RecvBuf::Terminated;

                    Err(error)

                }

            }

        } else {

            Ok(None)

        }

    }

}

#[cfg(test)]

mod tests {

    use super::{InitialMessageType, Receiver, UnmarshallMessage};

    use aws_smithy_eventstream::error::Error as EventStreamError;

    use aws_smithy_eventstream::frame::{write_message_to, UnmarshalledMessage};

    use aws_smithy_runtime_api::client::result::SdkError;

    use aws_smithy_types::body::SdkBody;

    use aws_smithy_types::event_stream::{Header, HeaderValue, Message};

    use bytes::Bytes;

    use hyper::body::Body;

    use std::error::Error as StdError;

    use std::io::{Error as IOError, ErrorKind};

    fn encode_initial_response() -> Bytes {

        let mut buffer = Vec::new();

        let message = Message::new(Bytes::new())

            .add_header(Header::new(

                ":message-type",

                HeaderValue::String("event".into()),

            ))

            .add_header(Header::new(

                ":event-type",

                HeaderValue::String("initial-response".into()),

            ));

        write_message_to(&message, &mut buffer).unwrap();

        buffer.into()

    }

    fn encode_message(message: &str) -> Bytes {

        let mut buffer = Vec::new();

        let message = Message::new(Bytes::copy_from_slice(message.as_bytes()));

        write_message_to(&message, &mut buffer).unwrap();

        buffer.into()

    }

    #[derive(Debug)]

    struct FakeError;

    impl std::fmt::Display for FakeError {

        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

            write!(f, "FakeError")

        }

    }

    impl StdError for FakeError {}

    #[derive(Debug, Eq, PartialEq)]

    struct TestMessage(String);

    #[derive(Debug)]

    struct Unmarshaller;

    impl UnmarshallMessage for Unmarshaller {

        type Output = TestMessage;

        type Error = EventStreamError;

        fn unmarshall(

            &self,

            message: &Message,

        ) -> Result<UnmarshalledMessage<Self::Output, Self::Error>, EventStreamError> {

            Ok(UnmarshalledMessage::Event(TestMessage(

                std::str::from_utf8(&message.payload()[..]).unwrap().into(),

            )))

        }

    }

    #[tokio::test]

    async fn receive_success() {

        let chunks: Vec<Result<_, IOError>> =

            vec![Ok(encode_message("one")), Ok(encode_message("two"))];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("two".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(None, receiver.recv().await.unwrap());

    }

    #[tokio::test]

    async fn receive_last_chunk_empty() {

        let chunks: Vec<Result<_, IOError>> = vec![

            Ok(encode_message("one")),

            Ok(encode_message("two")),

            Ok(Bytes::from_static(&[])),

        ];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("two".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(None, receiver.recv().await.unwrap());

    }

    #[tokio::test]

    async fn receive_last_chunk_not_full_message() {

        let chunks: Vec<Result<_, IOError>> = vec![

            Ok(encode_message("one")),

            Ok(encode_message("two")),

            Ok(encode_message("three").split_to(10)),

        ];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("two".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert!(matches!(

            receiver.recv().await,

            Err(SdkError::ResponseError { .. }),

        ));

    }

    #[tokio::test]

    async fn receive_last_chunk_has_multiple_messages() {

        let chunks: Vec<Result<_, IOError>> = vec![

            Ok(encode_message("one")),

            Ok(encode_message("two")),

            Ok(Bytes::from(

                [encode_message("three"), encode_message("four")].concat(),

            )),

        ];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("two".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("three".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("four".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(None, receiver.recv().await.unwrap());

    }

    proptest::proptest! {

        #[test]

        fn receive_multiple_messages_split_unevenly_across_chunks(b1: usize, b2: usize) {

            let combined = Bytes::from([

                encode_message("one"),

                encode_message("two"),

                encode_message("three"),

                encode_message("four"),

                encode_message("five"),

                encode_message("six"),

                encode_message("seven"),

                encode_message("eight"),

            ].concat());

            let midpoint = combined.len() / 2;

            let (start, boundary1, boundary2, end) = (

                0,

                b1 % midpoint,

                midpoint + b2 % midpoint,

                combined.len()

            );

            println!("[{}, {}], [{}, {}], [{}, {}]", start, boundary1, boundary1, boundary2, boundary2, end);

            let rt = tokio::runtime::Runtime::new().unwrap();

            rt.block_on(async move {

                let chunks: Vec<Result<_, IOError>> = vec![

                    Ok(Bytes::copy_from_slice(&combined[start..boundary1])),

                    Ok(Bytes::copy_from_slice(&combined[boundary1..boundary2])),

                    Ok(Bytes::copy_from_slice(&combined[boundary2..end])),

                ];

                let chunk_stream = futures_util::stream::iter(chunks);

                let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

                let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

                for payload in &["one", "two", "three", "four", "five", "six", "seven", "eight"] {

                    assert_eq!(

                        TestMessage((*payload).into()),

                        receiver.recv().await.unwrap().unwrap()

                    );

                }

                assert_eq!(None, receiver.recv().await.unwrap());

            });

        }

    }

    #[tokio::test]

    async fn receive_network_failure() {

        let chunks: Vec<Result<_, IOError>> = vec![

            Ok(encode_message("one")),

            Err(IOError::new(ErrorKind::ConnectionReset, FakeError)),

        ];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert!(matches!(

            receiver.recv().await,

            Err(SdkError::DispatchFailure(_))

        ));

    }

    #[tokio::test]

    async fn receive_message_parse_failure() {

        let chunks: Vec<Result<_, IOError>> = vec![

            Ok(encode_message("one")),

            // A zero length message will be invalid. We need to provide a minimum of 12 bytes

            // for the MessageFrameDecoder to actually start parsing it.

            Ok(Bytes::from_static(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])),

        ];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert!(matches!(

            receiver.recv().await,

            Err(SdkError::ResponseError { .. })

        ));

    }

    #[tokio::test]

    async fn receive_initial_response() {

        let chunks: Vec<Result<_, IOError>> =

            vec![Ok(encode_initial_response()), Ok(encode_message("one"))];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert!(receiver

            .try_recv_initial(InitialMessageType::Response)

            .await

            .unwrap()

            .is_some());

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

    }

    #[tokio::test]

    async fn receive_no_initial_response() {

        let chunks: Vec<Result<_, IOError>> =

            vec![Ok(encode_message("one")), Ok(encode_message("two"))];

        let chunk_stream = futures_util::stream::iter(chunks);

        let body = SdkBody::from_body_0_4(Body::wrap_stream(chunk_stream));

        let mut receiver = Receiver::<TestMessage, EventStreamError>::new(Unmarshaller, body);

        assert!(receiver

            .try_recv_initial(InitialMessageType::Response)

            .await

            .unwrap()

            .is_none());

        assert_eq!(

            TestMessage("one".into()),

            receiver.recv().await.unwrap().unwrap()

        );

        assert_eq!(

            TestMessage("two".into()),

            receiver.recv().await.unwrap().unwrap()

        );

    }

    fn assert_send_and_sync<T: Send + Sync>() {}

    #[tokio::test]

    async fn receiver_is_send_and_sync() {

        assert_send_and_sync::<Receiver<(), ()>>();

    }

}

/*

 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 * SPDX-License-Identifier: Apache-2.0

 */

use aws_smithy_eventstream::frame::{write_message_to, MarshallMessage, SignMessage};

use aws_smithy_eventstream::message_size_hint::MessageSizeHint;

use aws_smithy_runtime_api::client::result::SdkError;

use aws_smithy_types::error::ErrorMetadata;

use bytes::Bytes;

use futures_core::Stream;

use std::error::Error as StdError;

use std::fmt;

use std::fmt::Debug;

use std::marker::PhantomData;

use std::pin::Pin;

use std::task::{Context, Poll};

use tracing::trace;

/// Input type for Event Streams.

pub struct EventStreamSender<T, E> {

    input_stream: Pin<Box<dyn Stream<Item = Result<T, E>> + Send + Sync>>,

}

impl<T, E> Debug for EventStreamSender<T, E> {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        let name_t = std::any::type_name::<T>();

        let name_e = std::any::type_name::<E>();

        write!(f, "EventStreamSender<{name_t}, {name_e}>")

    }

}

impl<T: Send + Sync + 'static, E: StdError + Send + Sync + 'static> EventStreamSender<T, E> {

    /// Creates an `EventStreamSender` from a single item.

    pub fn once(item: Result<T, E>) -> Self {

        Self::from(futures_util::stream::once(async move { item }))

    }

}

impl<T, E: StdError + Send + Sync + 'static> EventStreamSender<T, E> {

    #[doc(hidden)]

    pub fn into_body_stream(

        self,

        marshaller: impl MarshallMessage<Input = T> + Send + Sync + 'static,

        error_marshaller: impl MarshallMessage<Input = E> + Send + Sync + 'static,

        signer: impl SignMessage + Send + Sync + 'static,

    ) -> MessageStreamAdapter<T, E> {

        MessageStreamAdapter::new(marshaller, error_marshaller, signer, self.input_stream)

    }

}

impl<T, E, S> From<S> for EventStreamSender<T, E>

where

    S: Stream<Item = Result<T, E>> + Send + Sync + 'static,

{

    fn from(stream: S) -> Self {

        EventStreamSender {

            input_stream: Box::pin(stream),

        }

    }

}

/// An error that occurs within a message stream.

#[derive(Debug)]

pub struct MessageStreamError {

    kind: MessageStreamErrorKind,

    pub(crate) meta: ErrorMetadata,

}

#[derive(Debug)]

enum MessageStreamErrorKind {

    Unhandled(Box<dyn std::error::Error + Send + Sync + 'static>),

}

impl MessageStreamError {

    /// Creates the `MessageStreamError::Unhandled` variant from any error type.

    pub fn unhandled(err: impl Into<Box<dyn std::error::Error + Send + Sync + 'static>>) -> Self {

        Self {

            meta: Default::default(),

            kind: MessageStreamErrorKind::Unhandled(err.into()),

        }

    }

    /// Creates the `MessageStreamError::Unhandled` variant from an [`ErrorMetadata`].

    pub fn generic(err: ErrorMetadata) -> Self {

        Self {

            meta: err.clone(),

            kind: MessageStreamErrorKind::Unhandled(err.into()),

        }

    }

    /// Returns error metadata, which includes the error code, message,

    /// request ID, and potentially additional information.

    pub fn meta(&self) -> &ErrorMetadata {

        &self.meta

    }

}

impl StdError for MessageStreamError {

    fn source(&self) -> Option<&(dyn StdError + 'static)> {

        match &self.kind {

            MessageStreamErrorKind::Unhandled(source) => Some(source.as_ref() as _),

        }

    }

}

impl fmt::Display for MessageStreamError {

    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

        match &self.kind {

            MessageStreamErrorKind::Unhandled(_) => write!(f, "message stream error"),

        }

    }

}

/// Adapts a `Stream<SmithyMessageType>` to a signed `Stream<Bytes>` by using the provided

/// message marshaller and signer implementations.

///

/// This will yield an `Err(SdkError::ConstructionFailure)` if a message can't be

/// marshalled into an Event Stream frame, (e.g., if the message payload was too large).

#[allow(missing_debug_implementations)]

pub struct MessageStreamAdapter<T, E: StdError + Send + Sync + 'static> {

    marshaller: Box<dyn MarshallMessage<Input = T> + Send + Sync>,

    error_marshaller: Box<dyn MarshallMessage<Input = E> + Send + Sync>,

    signer: Box<dyn SignMessage + Send + Sync>,

    stream: Pin<Box<dyn Stream<Item = Result<T, E>> + Send>>,

    end_signal_sent: bool,

    _phantom: PhantomData<E>,

}

impl<T, E: StdError + Send + Sync + 'static> Unpin for MessageStreamAdapter<T, E> {}

impl<T, E: StdError + Send + Sync + 'static> MessageStreamAdapter<T, E> {

    /// Create a new `MessageStreamAdapter`.

    pub fn new(

        marshaller: impl MarshallMessage<Input = T> + Send + Sync + 'static,

        error_marshaller: impl MarshallMessage<Input = E> + Send + Sync + 'static,

        signer: impl SignMessage + Send + Sync + 'static,

        stream: Pin<Box<dyn Stream<Item = Result<T, E>> + Send>>,

    ) -> Self {

        MessageStreamAdapter {

            marshaller: Box::new(marshaller),

            error_marshaller: Box::new(error_marshaller),

            signer: Box::new(signer),

            stream,

            end_signal_sent: false,

            _phantom: Default::default(),

        }

    }

}

impl<T, E: StdError + Send + Sync + 'static> Stream for MessageStreamAdapter<T, E> {

    type Item =

        Result<Bytes, SdkError<E, aws_smithy_runtime_api::client::orchestrator::HttpResponse>>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

        match self.stream.as_mut().poll_next(cx) {

            Poll::Ready(message_option) => {

                if let Some(message_result) = message_option {

                    let message = match message_result {

                        Ok(message) => self

                            .marshaller

                            .marshall(message)

                            .map_err(SdkError::construction_failure)?,

                        Err(message) => self

                            .error_marshaller

                            .marshall(message)

                            .map_err(SdkError::construction_failure)?,

                    };

                    trace!(unsigned_message = ?message, "signing event stream message");

                    let message = self

                        .signer

                        .sign(message)

                        .map_err(SdkError::construction_failure)?;

                    let mut buffer = Vec::with_capacity(message.size_hint());

                    write_message_to(&message, &mut buffer)

                        .map_err(SdkError::construction_failure)?;

                    trace!(signed_message = ?buffer, "sending signed event stream message");

                    Poll::Ready(Some(Ok(Bytes::from(buffer))))

                } else if !self.end_signal_sent {

                    self.end_signal_sent = true;

                    match self.signer.sign_empty() {

                        Some(sign) => {

                            let message = sign.map_err(SdkError::construction_failure)?;

                            let mut buffer = Vec::with_capacity(message.size_hint());

                            write_message_to(&message, &mut buffer)

                                .map_err(SdkError::construction_failure)?;

                            trace!(signed_message = ?buffer, "sending signed empty message to terminate the event stream");

                            Poll::Ready(Some(Ok(Bytes::from(buffer))))

                        }

                        None => Poll::Ready(None),

                    }

                } else {

                    Poll::Ready(None)

                }

            }

            Poll::Pending => Poll::Pending,

        }

    }

}

#[cfg(test)]

mod tests {

    use super::MarshallMessage;

    use crate::event_stream::{EventStreamSender, MessageStreamAdapter};

    use async_stream::stream;

    use aws_smithy_eventstream::error::Error as EventStreamError;

    use aws_smithy_eventstream::frame::{

        read_message_from, write_message_to, NoOpSigner, SignMessage, SignMessageError,

    };

    use aws_smithy_runtime_api::client::result::SdkError;

    use aws_smithy_types::event_stream::{Header, HeaderValue, Message};

    use bytes::Bytes;

    use futures_core::Stream;

    use futures_util::stream::StreamExt;

    use std::error::Error as StdError;

    #[derive(Debug, Eq, PartialEq)]

    struct TestMessage(String);

    #[derive(Debug)]

    struct Marshaller;

    impl MarshallMessage for Marshaller {

        type Input = TestMessage;

        fn marshall(&self, input: Self::Input) -> Result<Message, EventStreamError> {

            Ok(Message::new(input.0.as_bytes().to_vec()))

        }

    }

    #[derive(Debug)]

    struct ErrorMarshaller;

    impl MarshallMessage for ErrorMarshaller {

        type Input = TestServiceError;

        fn marshall(&self, _input: Self::Input) -> Result<Message, EventStreamError> {

            Err(read_message_from(&b""[..]).expect_err("this should always fail"))

        }

    }

    #[derive(Debug)]

    struct TestServiceError;

    impl std::fmt::Display for TestServiceError {

        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

            write!(f, "TestServiceError")

        }

    }

    impl StdError for TestServiceError {}

    #[derive(Debug)]

    struct TestSigner;

    impl SignMessage for TestSigner {

        fn sign(&mut self, message: Message) -> Result<Message, SignMessageError> {

            let mut buffer = Vec::new();

            write_message_to(&message, &mut buffer).unwrap();

            Ok(Message::new(buffer).add_header(Header::new("signed", HeaderValue::Bool(true))))

        }

        fn sign_empty(&mut self) -> Option<Result<Message, SignMessageError>> {

            Some(Ok(

                Message::new(&b""[..]).add_header(Header::new("signed", HeaderValue::Bool(true)))

            ))

        }

    }

    fn check_send_sync<T: Send + Sync>(value: T) -> T {

        value

    }

    #[test]

    fn event_stream_sender_send_sync() {

        check_send_sync(EventStreamSender::from(stream! {

            yield Result::<_, SignMessageError>::Ok(TestMessage("test".into()));

        }));

    }

    fn check_compatible_with_hyper_wrap_stream<S, O, E>(stream: S) -> S

    where

        S: Stream<Item = Result<O, E>> + Send + 'static,

        O: Into<Bytes> + 'static,

        E: Into<Box<dyn StdError + Send + Sync + 'static>> + 'static,

    {

        stream

    }

    #[tokio::test]

    async fn message_stream_adapter_success() {

        let stream = stream! {

            yield Ok(TestMessage("test".into()));

        };

        let mut adapter = check_compatible_with_hyper_wrap_stream(MessageStreamAdapter::<

            TestMessage,

            TestServiceError,

        >::new(

            Marshaller,

            ErrorMarshaller,

            TestSigner,

            Box::pin(stream),

        ));

        let mut sent_bytes = adapter.next().await.unwrap().unwrap();

        let sent = read_message_from(&mut sent_bytes).unwrap();

        assert_eq!("signed", sent.headers()[0].name().as_str());

        assert_eq!(&HeaderValue::Bool(true), sent.headers()[0].value());

        let inner = read_message_from(&mut (&sent.payload()[..])).unwrap();

        assert_eq!(&b"test"[..], &inner.payload()[..]);

        let mut end_signal_bytes = adapter.next().await.unwrap().unwrap();

        let end_signal = read_message_from(&mut end_signal_bytes).unwrap();

        assert_eq!("signed", end_signal.headers()[0].name().as_str());

        assert_eq!(&HeaderValue::Bool(true), end_signal.headers()[0].value());

        assert_eq!(0, end_signal.payload().len());

    }

    #[tokio::test]

    async fn message_stream_adapter_construction_failure() {

        let stream = stream! {

            yield Err(TestServiceError);

        };

        let mut adapter = check_compatible_with_hyper_wrap_stream(MessageStreamAdapter::<

            TestMessage,

            TestServiceError,

        >::new(

            Marshaller,

            ErrorMarshaller,

            NoOpSigner {},

            Box::pin(stream),

        ));

        let result = adapter.next().await.unwrap();

        assert!(result.is_err());

        assert!(matches!(

            result.err().unwrap(),

            SdkError::ConstructionFailure(_)

        ));

    }

    #[tokio::test]

    async fn event_stream_sender_once() {

        let sender = EventStreamSender::once(Ok(TestMessage("test".into())));

        let mut adapter = MessageStreamAdapter::<TestMessage, TestServiceError>::new(

            Marshaller,

            ErrorMarshaller,

            TestSigner,

            sender.input_stream,

        );

        let mut sent_bytes = adapter.next().await.unwrap().unwrap();

        let sent = read_message_from(&mut sent_bytes).unwrap();

        assert_eq!("signed", sent.headers()[0].name().as_str());

        let inner = read_message_from(&mut (&sent.payload()[..])).unwrap();

        assert_eq!(&b"test"[..], &inner.payload()[..]);

        // Should get end signal next

        let mut end_signal_bytes = adapter.next().await.unwrap().unwrap();

        let end_signal = read_message_from(&mut end_signal_bytes).unwrap();

        assert_eq!("signed", end_signal.headers()[0].name().as_str());

        assert_eq!(0, end_signal.payload().len());

        // Stream should be exhausted

        assert!(adapter.next().await.is_none());

    }

    // Verify the developer experience for this compiles

    #[allow(unused)]

    fn event_stream_input_ergonomics() {

        fn check(input: impl Into<EventStreamSender<TestMessage, TestServiceError>>) {

            let _: EventStreamSender<TestMessage, TestServiceError> = input.into();

        }

        check(stream! {

            yield Ok(TestMessage("test".into()));

        });

        check(stream! {

            yield Err(TestServiceError);

        });

    }

}

/*

 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 * SPDX-License-Identifier: Apache-2.0

 */

use aws_smithy_types::body::SdkBody;

use aws_smithy_types::byte_stream::error::Error as ByteStreamError;

use aws_smithy_types::byte_stream::ByteStream;

use bytes::Bytes;

use futures_core::stream::Stream;

use std::pin::Pin;

use std::task::{Context, Poll};

/// A new-type wrapper to enable the impl of the `futures_core::stream::Stream` trait

///

/// [`ByteStream`] no longer implements `futures_core::stream::Stream` so we wrap it in the

/// new-type to enable the trait when it is required.

///

/// This is meant to be used by codegen code, and users should not need to use it directly.

#[derive(Debug)]

pub struct FuturesStreamCompatByteStream(ByteStream);

impl FuturesStreamCompatByteStream {

    /// Creates a new `FuturesStreamCompatByteStream` by wrapping `stream`.

    pub fn new(stream: ByteStream) -> Self {

        Self(stream)

    }

    /// Returns [`SdkBody`] of the wrapped [`ByteStream`].

    pub fn into_inner(self) -> SdkBody {

        self.0.into_inner()

    }

}

impl Stream for FuturesStreamCompatByteStream {

    type Item = Result<Bytes, ByteStreamError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

        Pin::new(&mut self.0).poll_next(cx)

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use futures_core::stream::Stream;

    fn check_compatible_with_hyper_wrap_stream<S, O, E>(stream: S) -> S

    where

        S: Stream<Item = Result<O, E>> + Send + 'static,

        O: Into<Bytes> + 'static,

        E: Into<Box<dyn std::error::Error + Send + Sync + 'static>> + 'static,

    {

        stream

    }

    #[test]

    fn test_byte_stream_stream_can_be_made_compatible_with_hyper_wrap_stream() {

        let stream = ByteStream::from_static(b"Hello world");

        check_compatible_with_hyper_wrap_stream(FuturesStreamCompatByteStream::new(stream));

    }

}