From a5c1ced056410ed2ed3a14deefc469bf5f9d27ad Mon Sep 17 00:00:00 2001
From: Russell Cohen <rcoh@amazon.com>
Date: Tue, 19 Sep 2023 10:38:05 -0400
Subject: [PATCH] add map_immutable to allow signing to work with a wrapped
 SdkBody (#2567)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->
aws-sdk-rust#749

## Description
<!--- Describe your changes in detail -->
When wrapping an `SdkBody`, the resulting body is no longer signable.
Some services support this and others don't. This change adds a new
method for mapping an `SdkBody` called `map_readonly_body`. If called on
a body that is signable, it will produce a signable body. However, this
function should never when modifying the inner body is necessary, as
that will result in a signing error.

## Testing
<!--- Please describe in detail how you tested your changes -->
<!--- Include details of your testing environment, and the tests you ran
to -->
<!--- see how your change affects other areas of the code, etc. -->
Zelda wrote a test

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [x] I have updated `CHANGELOG.next.toml` if I made changes to the AWS
SDK, generated SDK code, or SDK runtime crates

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
---
 rust-runtime/aws-smithy-http/src/body.rs | 44 +++++++++++++++++++++---
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/rust-runtime/aws-smithy-http/src/body.rs b/rust-runtime/aws-smithy-http/src/body.rs
index 32ad205ed..5f8ac23a3 100644
--- a/rust-runtime/aws-smithy-http/src/body.rs
+++ b/rust-runtime/aws-smithy-http/src/body.rs
@@ -35,6 +35,7 @@ pin_project! {
         // In the event of retry, this function will be called to generate a new body. See
         // [`try_clone()`](SdkBody::try_clone)
         rebuild: Option<Arc<dyn (Fn() -> Inner) + Send + Sync>>,
+        bytes_contents: Option<Bytes>
     }
 }
 
@@ -94,6 +95,7 @@ impl SdkBody {
         Self {
             inner: Inner::Dyn { inner: body },
             rebuild: None,
+            bytes_contents: None,
         }
     }
 
@@ -110,6 +112,7 @@ impl SdkBody {
         SdkBody {
             inner: initial.inner,
             rebuild: Some(Arc::new(move || f().inner)),
+            bytes_contents: initial.bytes_contents,
         }
     }
 
@@ -119,6 +122,7 @@ impl SdkBody {
         Self {
             inner: Inner::Taken,
             rebuild: None,
+            bytes_contents: None,
         }
     }
 
@@ -127,6 +131,7 @@ impl SdkBody {
         Self {
             inner: Inner::Once { inner: None },
             rebuild: Some(Arc::new(|| Inner::Once { inner: None })),
+            bytes_contents: Some(Bytes::new()),
         }
     }
 
@@ -157,10 +162,9 @@ impl SdkBody {
     /// If this SdkBody is NOT streaming, this will return the byte slab
     /// If this SdkBody is streaming, this will return `None`
     pub fn bytes(&self) -> Option<&[u8]> {
-        match &self.inner {
-            Inner::Once { inner: Some(b) } => Some(b),
-            Inner::Once { inner: None } => Some(&[]),
-            _ => None,
+        match &self.bytes_contents {
+            Some(b) => Some(b),
+            None => None,
         }
     }
 
@@ -172,6 +176,7 @@ impl SdkBody {
             Self {
                 inner: next,
                 rebuild: self.rebuild.clone(),
+                bytes_contents: self.bytes_contents.clone(),
             }
         })
     }
@@ -191,6 +196,25 @@ impl SdkBody {
             f(self)
         }
     }
+
+    /// Update this `SdkBody` with `map`. **This function MUST NOT alert the data of the body.**
+    ///
+    /// This function is useful for adding metadata like progress tracking to an [`SdkBody`] that
+    /// does not alter the actual byte data. If your mapper alters the contents of the body, use [`SdkBody::map`]
+    /// instead.
+    pub fn map_preserve_contents(
+        self,
+        f: impl Fn(SdkBody) -> SdkBody + Sync + Send + 'static,
+    ) -> SdkBody {
+        let contents = self.bytes_contents.clone();
+        let mut out = if self.rebuild.is_some() {
+            SdkBody::retryable(move || f(self.try_clone().unwrap()))
+        } else {
+            f(self)
+        };
+        out.bytes_contents = contents;
+        out
+    }
 }
 
 impl From<&str> for SdkBody {
@@ -201,6 +225,7 @@ impl From<&str> for SdkBody {
 
 impl From<Bytes> for SdkBody {
     fn from(bytes: Bytes) -> Self {
+        let b = bytes.clone();
         SdkBody {
             inner: Inner::Once {
                 inner: Some(bytes.clone()),
@@ -208,6 +233,7 @@ impl From<Bytes> for SdkBody {
             rebuild: Some(Arc::new(move || Inner::Once {
                 inner: Some(bytes.clone()),
             })),
+            bytes_contents: Some(b),
         }
     }
 }
@@ -217,6 +243,7 @@ impl From<hyper::Body> for SdkBody {
         SdkBody {
             inner: Inner::Streaming { inner: body },
             rebuild: None,
+            bytes_contents: None,
         }
     }
 }
@@ -290,6 +317,15 @@ mod test {
         assert_eq!(SdkBody::from("").size_hint().exact(), Some(0));
     }
 
+    #[test]
+    fn map_preserve_preserves_bytes_hint() {
+        let initial = SdkBody::from("hello!");
+        assert_eq!(initial.bytes(), Some(b"hello!".as_slice()));
+
+        let new_body = initial.map_preserve_contents(|body| SdkBody::from_dyn(BoxBody::new(body)));
+        assert_eq!(new_body.bytes(), Some(b"hello!".as_slice()));
+    }
+
     #[allow(clippy::bool_assert_comparison)]
     #[test]
     fn valid_eos() {
-- 
GitLab