aws_chunked.rs

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

#![allow(dead_code)]

use std::fmt;

use aws_runtime::{
    auth::PayloadSigningOverride,
    content_encoding::{header_value::AWS_CHUNKED, AwsChunkedBody, AwsChunkedBodyOptions},
};
use aws_smithy_runtime_api::{
    box_error::BoxError,
    client::{
        interceptors::{context::BeforeTransmitInterceptorContextMut, Intercept},
        runtime_components::RuntimeComponents,
    },
    http::Request,
};
use aws_smithy_types::{body::SdkBody, config_bag::ConfigBag, error::operation::BuildError};
use http::{header, HeaderValue};
use http_body::Body;

const X_AMZ_DECODED_CONTENT_LENGTH: &str = "x-amz-decoded-content-length";

/// Errors related to constructing aws-chunked encoded HTTP requests.
#[derive(Debug)]
enum Error {
    UnsizedRequestBody,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsizedRequestBody => write!(
                f,
                "Only request bodies with a known size can be aws-chunked encoded."
            ),
        }
    }
}

impl std::error::Error for Error {}

#[derive(Debug)]
pub(crate) struct AwsChunkedContentEncodingInterceptor;

impl Intercept for AwsChunkedContentEncodingInterceptor {
    fn name(&self) -> &'static str {
        "AwsChunkedContentEncodingInterceptor"
    }

    fn modify_before_signing(
        &self,
        context: &mut BeforeTransmitInterceptorContextMut<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        if must_not_use_chunked_encoding(context.request(), cfg) {
            tracing::debug!(
                "short-circuiting modify_before_signing because chunked encoding must not be used"
            );
            return Ok(());
        }

        let original_body_size = if let Some(size) = context
            .request()
            .headers()
            .get(header::CONTENT_LENGTH)
            .and_then(|s| s.parse::<u64>().ok())
            .or_else(|| context.request().body().size_hint().exact())
        {
            size
        } else {
            return Err(BuildError::other(Error::UnsizedRequestBody))?;
        };

        let chunked_body_options = if let Some(chunked_body_options) =
            cfg.get_mut_from_interceptor_state::<AwsChunkedBodyOptions>()
        {
            let chunked_body_options = std::mem::take(chunked_body_options);
            chunked_body_options.with_stream_length(original_body_size)
        } else {
            AwsChunkedBodyOptions::default().with_stream_length(original_body_size)
        };

        let request = context.request_mut();
        // For for aws-chunked encoding, `x-amz-decoded-content-length` must be set to the original body size.
        request.headers_mut().insert(
            header::HeaderName::from_static(X_AMZ_DECODED_CONTENT_LENGTH),
            HeaderValue::from(original_body_size),
        );
        // Other than `x-amz-decoded-content-length`, either `content-length` or `transfer-encoding`
        // must be set, but not both. For uses cases we support, we know the original body size and
        // can calculate the encoded size, so we set `content-length`.
        request.headers_mut().insert(
            header::CONTENT_LENGTH,
            HeaderValue::from(chunked_body_options.encoded_length()),
        );
        // Setting `content-length` above means we must unset `transfer-encoding`.
        request.headers_mut().remove(header::TRANSFER_ENCODING);
        request.headers_mut().append(
            header::CONTENT_ENCODING,
            HeaderValue::from_str(AWS_CHUNKED)
                .map_err(BuildError::other)
                .expect("\"aws-chunked\" will always be a valid HeaderValue"),
        );

        cfg.interceptor_state().store_put(chunked_body_options);
        cfg.interceptor_state()
            .store_put(PayloadSigningOverride::StreamingUnsignedPayloadTrailer);

        Ok(())
    }

    fn modify_before_transmit(
        &self,
        ctx: &mut BeforeTransmitInterceptorContextMut<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        if must_not_use_chunked_encoding(ctx.request(), cfg) {
            tracing::debug!(
                "short-circuiting modify_before_transmit because chunked encoding must not be used"
            );
            return Ok(());
        }

        let request = ctx.request_mut();

        let mut body = {
            let body = std::mem::replace(request.body_mut(), SdkBody::taken());
            let opt = cfg
                .get_mut_from_interceptor_state::<AwsChunkedBodyOptions>()
                .ok_or_else(|| {
                    BuildError::other("AwsChunkedBodyOptions missing from config bag")
                })?;
            let aws_chunked_body_options = std::mem::take(opt);
            body.map(move |body| {
                let body = AwsChunkedBody::new(body, aws_chunked_body_options.clone());
                SdkBody::from_body_0_4(body)
            })
        };

        std::mem::swap(request.body_mut(), &mut body);

        Ok(())
    }
}

// Determine if chunked encoding must not be used; returns true when any of the following is true:
// - If the body is in-memory
// - If chunked encoding is disabled via `AwsChunkedBodyOptions`
fn must_not_use_chunked_encoding(request: &Request, cfg: &ConfigBag) -> bool {
    match (request.body().bytes(), cfg.load::<AwsChunkedBodyOptions>()) {
        (Some(_), _) => true,
        (_, Some(options)) if options.disabled() => true,
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aws_smithy_runtime_api::client::interceptors::context::{
        BeforeTransmitInterceptorContextMut, Input, InterceptorContext,
    };
    use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
    use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder;
    use aws_smithy_types::byte_stream::ByteStream;
    use bytes::BytesMut;
    use http_body::Body;
    use tempfile::NamedTempFile;

    #[tokio::test]
    async fn test_aws_chunked_body_is_retryable() {
        use std::io::Write;
        let mut file = NamedTempFile::new().unwrap();

        for i in 0..10000 {
            let line = format!("This is a large file created for testing purposes {}", i);
            file.as_file_mut().write_all(line.as_bytes()).unwrap();
        }

        let stream_length = file.as_file().metadata().unwrap().len();
        let request = HttpRequest::new(
            ByteStream::read_from()
                .path(&file)
                .buffer_size(1024)
                .build()
                .await
                .unwrap()
                .into_inner(),
        );

        // ensure original SdkBody is retryable
        assert!(request.body().try_clone().is_some());

        let interceptor = AwsChunkedContentEncodingInterceptor;
        let mut cfg = ConfigBag::base();
        cfg.interceptor_state()
            .store_put(AwsChunkedBodyOptions::default().with_stream_length(stream_length));
        let runtime_components = RuntimeComponentsBuilder::for_tests().build().unwrap();
        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
        ctx.enter_serialization_phase();
        let _ = ctx.take_input();
        ctx.set_request(request);
        ctx.enter_before_transmit_phase();
        let mut ctx: BeforeTransmitInterceptorContextMut<'_> = (&mut ctx).into();
        interceptor
            .modify_before_transmit(&mut ctx, &runtime_components, &mut cfg)
            .unwrap();

        // ensure wrapped SdkBody is retryable
        let mut body = ctx.request().body().try_clone().expect("body is retryable");

        let mut body_data = BytesMut::new();
        while let Some(data) = body.data().await {
            body_data.extend_from_slice(&data.unwrap())
        }
        let body_str = std::str::from_utf8(&body_data).unwrap();

        let expected = "This is a large file created for testing purposes 9999\r\n0\r\n\r\n";
        assert!(
            body_str.ends_with(expected),
            "expected '{body_str}' to end with '{expected}'"
        );
    }

    #[tokio::test]
    async fn test_short_circuit_modify_before_signing() {
        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
        ctx.enter_serialization_phase();
        let _ = ctx.take_input();
        let request = HttpRequest::new(SdkBody::from(
            "in-memory body, must not use chunked encoding",
        ));
        ctx.set_request(request);
        ctx.enter_before_transmit_phase();
        let mut ctx: BeforeTransmitInterceptorContextMut<'_> = (&mut ctx).into();

        let runtime_components = RuntimeComponentsBuilder::for_tests().build().unwrap();

        let mut cfg = ConfigBag::base();
        cfg.interceptor_state()
            .store_put(AwsChunkedBodyOptions::default());

        let interceptor = AwsChunkedContentEncodingInterceptor;
        interceptor
            .modify_before_signing(&mut ctx, &runtime_components, &mut cfg)
            .unwrap();

        let request = ctx.request();
        assert!(request.headers().get(header::CONTENT_ENCODING).is_none());
        assert!(request
            .headers()
            .get(header::HeaderName::from_static(
                X_AMZ_DECODED_CONTENT_LENGTH
            ))
            .is_none());
    }

    #[tokio::test]
    async fn test_short_circuit_modify_before_transmit() {
        let mut ctx = InterceptorContext::new(Input::doesnt_matter());
        ctx.enter_serialization_phase();
        let _ = ctx.take_input();
        let request = HttpRequest::new(SdkBody::from(
            "in-memory body, must not use chunked encoding",
        ));
        ctx.set_request(request);
        ctx.enter_before_transmit_phase();
        let mut ctx: BeforeTransmitInterceptorContextMut<'_> = (&mut ctx).into();

        let runtime_components = RuntimeComponentsBuilder::for_tests().build().unwrap();

        let mut cfg = ConfigBag::base();
        // Don't need to set the stream length properly because we expect the body won't be wrapped by `AwsChunkedBody`.
        cfg.interceptor_state()
            .store_put(AwsChunkedBodyOptions::default());

        let interceptor = AwsChunkedContentEncodingInterceptor;
        interceptor
            .modify_before_transmit(&mut ctx, &runtime_components, &mut cfg)
            .unwrap();

        let mut body = ctx.request().body().try_clone().expect("body is retryable");

        let mut body_data = BytesMut::new();
        while let Some(data) = body.data().await {
            body_data.extend_from_slice(&data.unwrap())
        }
        let body_str = std::str::from_utf8(&body_data).unwrap();
        // Also implies that `assert!(!body_str.ends_with("0\r\n\r\n"));`, i.e., shouldn't see chunked encoding epilogue.
        assert_eq!("in-memory body, must not use chunked encoding", body_str);
    }

    #[test]
    fn test_must_not_use_chunked_encoding_with_in_memory_body() {
        let request = HttpRequest::new(SdkBody::from("test body"));
        let cfg = ConfigBag::base();

        assert!(must_not_use_chunked_encoding(&request, &cfg));
    }

    async fn streaming_body(path: impl AsRef<std::path::Path>) -> SdkBody {
        let file = path.as_ref();
        ByteStream::read_from()
            .path(&file)
            .build()
            .await
            .unwrap()
            .into_inner()
    }

    #[tokio::test]
    async fn test_must_not_use_chunked_encoding_with_disabled_option() {
        let file = NamedTempFile::new().unwrap();
        let request = HttpRequest::new(streaming_body(&file).await);
        let mut cfg = ConfigBag::base();
        cfg.interceptor_state()
            .store_put(AwsChunkedBodyOptions::disable_chunked_encoding());

        assert!(must_not_use_chunked_encoding(&request, &cfg));
    }

    #[tokio::test]
    async fn test_chunked_encoding_is_used() {
        let file = NamedTempFile::new().unwrap();
        let request = HttpRequest::new(streaming_body(&file).await);
        let cfg = ConfigBag::base();

        assert!(!must_not_use_chunked_encoding(&request, &cfg));
    }
}