|
| 1 | +/* |
| 2 | + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +//! Adapters to use http-body 1.0 bodies with SdkBody & ByteStream |
| 7 | +
|
| 8 | +use std::pin::Pin; |
| 9 | +use std::task::{ready, Context, Poll}; |
| 10 | + |
| 11 | +use bytes::Bytes; |
| 12 | +use http_body_util::BodyExt; |
| 13 | +use pin_project_lite::pin_project; |
| 14 | + |
| 15 | +use crate::body::{BoxBody, Error, Inner, SdkBody}; |
| 16 | + |
| 17 | +impl SdkBody { |
| 18 | + /// Construct an `SdkBody` from a type that implements [`http_body_1_0::Body<Data = Bytes>`](http_body_1_0::Body). |
| 19 | + /// |
| 20 | + /// _Note: This is only available with `http-body-1-0` enabled._ |
| 21 | + pub fn from_body_1_0<T, E>(body: T) -> Self |
| 22 | + where |
| 23 | + T: http_body_1_0::Body<Data = Bytes, Error = E> + Send + Sync + 'static, |
| 24 | + E: Into<Error> + 'static, |
| 25 | + { |
| 26 | + Self { |
| 27 | + inner: Inner::Dyn { |
| 28 | + inner: BoxBody::HttpBody04(http_body_0_4::combinators::BoxBody::new( |
| 29 | + Http1toHttp04::new(body.map_err(Into::into)), |
| 30 | + )), |
| 31 | + }, |
| 32 | + rebuild: None, |
| 33 | + bytes_contents: None, |
| 34 | + } |
| 35 | + } |
| 36 | +} |
| 37 | + |
| 38 | +pin_project! { |
| 39 | + struct Http1toHttp04<B> { |
| 40 | + #[pin] |
| 41 | + inner: B, |
| 42 | + trailers: Option<http_1x::HeaderMap>, |
| 43 | + } |
| 44 | +} |
| 45 | + |
| 46 | +impl<B> Http1toHttp04<B> { |
| 47 | + fn new(inner: B) -> Self { |
| 48 | + Self { |
| 49 | + inner, |
| 50 | + trailers: None, |
| 51 | + } |
| 52 | + } |
| 53 | +} |
| 54 | + |
| 55 | +impl<B> http_body_0_4::Body for Http1toHttp04<B> |
| 56 | +where |
| 57 | + B: http_body_1_0::Body, |
| 58 | +{ |
| 59 | + type Data = B::Data; |
| 60 | + type Error = B::Error; |
| 61 | + |
| 62 | + /// |
| 63 | + /// |
| 64 | + /// # Arguments |
| 65 | + /// |
| 66 | + /// * `cx`: |
| 67 | + /// |
| 68 | + /// returns: Poll<Option<Result<<Http1toHttp04<B> as Body>::Data, <Http1toHttp04<B> as Body>::Error>>> |
| 69 | + /// |
| 70 | + /// # Examples |
| 71 | + /// |
| 72 | + /// ``` |
| 73 | + /// |
| 74 | + /// ``` |
| 75 | + fn poll_data( |
| 76 | + mut self: Pin<&mut Self>, |
| 77 | + cx: &mut Context<'_>, |
| 78 | + ) -> Poll<Option<Result<Self::Data, Self::Error>>> { |
| 79 | + loop { |
| 80 | + let this = self.as_mut().project(); |
| 81 | + match ready!(this.inner.poll_frame(cx)) { |
| 82 | + Some(Ok(frame)) => { |
| 83 | + let frame = match frame.into_data() { |
| 84 | + Ok(data) => return Poll::Ready(Some(Ok(data))), |
| 85 | + Err(frame) => frame, |
| 86 | + }; |
| 87 | + // when we get a trailers frame, store the trailers for the next poll |
| 88 | + if let Ok(trailers) = frame.into_trailers() { |
| 89 | + this.trailers.replace(trailers); |
| 90 | + return Poll::Ready(None); |
| 91 | + }; |
| 92 | + // if the frame type was unknown, discard it. the next one might be something |
| 93 | + // useful |
| 94 | + } |
| 95 | + Some(Err(e)) => return Poll::Ready(Some(Err(e))), |
| 96 | + None => return Poll::Ready(None), |
| 97 | + } |
| 98 | + } |
| 99 | + } |
| 100 | + |
| 101 | + fn poll_trailers( |
| 102 | + self: Pin<&mut Self>, |
| 103 | + _cx: &mut Context<'_>, |
| 104 | + ) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> { |
| 105 | + // all of the polling happens in poll_data, once we get to the trailers we've actually |
| 106 | + // already read everything |
| 107 | + let this = self.project(); |
| 108 | + match this.trailers.take() { |
| 109 | + Some(headers) => Poll::Ready(Ok(Some(convert_header_map(headers)))), |
| 110 | + None => Poll::Ready(Ok(None)), |
| 111 | + } |
| 112 | + } |
| 113 | + |
| 114 | + fn is_end_stream(&self) -> bool { |
| 115 | + self.inner.is_end_stream() |
| 116 | + } |
| 117 | + |
| 118 | + fn size_hint(&self) -> http_body_0_4::SizeHint { |
| 119 | + let mut size_hint = http_body_0_4::SizeHint::new(); |
| 120 | + let inner_hint = self.inner.size_hint(); |
| 121 | + if let Some(exact) = inner_hint.exact() { |
| 122 | + size_hint.set_exact(exact); |
| 123 | + } else { |
| 124 | + size_hint.set_lower(inner_hint.lower()); |
| 125 | + if let Some(upper) = inner_hint.upper() { |
| 126 | + size_hint.set_upper(upper); |
| 127 | + } |
| 128 | + } |
| 129 | + size_hint |
| 130 | + } |
| 131 | +} |
| 132 | + |
| 133 | +fn convert_header_map(input: http_1x::HeaderMap) -> http::HeaderMap { |
| 134 | + let mut map = http::HeaderMap::with_capacity(input.capacity()); |
| 135 | + let mut mem: Option<http_1x::HeaderName> = None; |
| 136 | + for (k, v) in input.into_iter() { |
| 137 | + let name = k.or_else(|| mem.clone()).unwrap(); |
| 138 | + map.append( |
| 139 | + http::HeaderName::from_bytes(name.as_str().as_bytes()).expect("already validated"), |
| 140 | + http::HeaderValue::from_bytes(v.as_bytes()).expect("already validated"), |
| 141 | + ); |
| 142 | + mem = Some(name); |
| 143 | + } |
| 144 | + map |
| 145 | +} |
| 146 | + |
| 147 | +#[cfg(test)] |
| 148 | +mod test { |
| 149 | + use crate::body::http_body_1_x::convert_header_map; |
| 150 | + use crate::body::{Error, SdkBody}; |
| 151 | + use crate::byte_stream::ByteStream; |
| 152 | + use bytes::Bytes; |
| 153 | + use http::header::{CONTENT_LENGTH as CL0, CONTENT_TYPE as CT0}; |
| 154 | + use http_1x::header::{CONTENT_LENGTH as CL1, CONTENT_TYPE as CT1}; |
| 155 | + use http_1x::{HeaderMap, HeaderName, HeaderValue}; |
| 156 | + use http_body_1_0::Frame; |
| 157 | + use std::collections::VecDeque; |
| 158 | + use std::pin::Pin; |
| 159 | + use std::task::{Context, Poll}; |
| 160 | + |
| 161 | + struct TestBody { |
| 162 | + chunks: VecDeque<Chunk>, |
| 163 | + } |
| 164 | + |
| 165 | + enum Chunk { |
| 166 | + Data(&'static str), |
| 167 | + Error(&'static str), |
| 168 | + Trailers(HeaderMap), |
| 169 | + } |
| 170 | + |
| 171 | + impl http_body_1_0::Body for TestBody { |
| 172 | + type Data = Bytes; |
| 173 | + type Error = Error; |
| 174 | + |
| 175 | + fn poll_frame( |
| 176 | + mut self: Pin<&mut Self>, |
| 177 | + _cx: &mut Context<'_>, |
| 178 | + ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> { |
| 179 | + let next = self.chunks.pop_front(); |
| 180 | + let mk = |v: Frame<Bytes>| Poll::Ready(Some(Ok(v))); |
| 181 | + |
| 182 | + match next { |
| 183 | + Some(Chunk::Data(s)) => mk(Frame::data(Bytes::from_static(s.as_bytes()))), |
| 184 | + Some(Chunk::Trailers(headers)) => mk(Frame::trailers(headers)), |
| 185 | + Some(Chunk::Error(err)) => Poll::Ready(Some(Err(err.into()))), |
| 186 | + None => Poll::Ready(None), |
| 187 | + } |
| 188 | + } |
| 189 | + } |
| 190 | + |
| 191 | + fn trailers() -> HeaderMap { |
| 192 | + let mut map = HeaderMap::new(); |
| 193 | + map.insert( |
| 194 | + HeaderName::from_static("x-test"), |
| 195 | + HeaderValue::from_static("x-test-value"), |
| 196 | + ); |
| 197 | + map.append( |
| 198 | + HeaderName::from_static("x-test"), |
| 199 | + HeaderValue::from_static("x-test-value-2"), |
| 200 | + ); |
| 201 | + map.append( |
| 202 | + HeaderName::from_static("y-test"), |
| 203 | + HeaderValue::from_static("y-test-value-2"), |
| 204 | + ); |
| 205 | + map |
| 206 | + } |
| 207 | + |
| 208 | + #[tokio::test] |
| 209 | + async fn test_body_with_trailers() { |
| 210 | + let body = TestBody { |
| 211 | + chunks: vec![ |
| 212 | + Chunk::Data("123"), |
| 213 | + Chunk::Data("456"), |
| 214 | + Chunk::Data("789"), |
| 215 | + Chunk::Trailers(trailers()), |
| 216 | + ] |
| 217 | + .into(), |
| 218 | + }; |
| 219 | + let body = SdkBody::from_body_1_0(body); |
| 220 | + let data = ByteStream::new(body); |
| 221 | + assert_eq!(data.collect().await.unwrap().to_vec(), b"123456789"); |
| 222 | + } |
| 223 | + |
| 224 | + #[tokio::test] |
| 225 | + async fn test_read_trailers() { |
| 226 | + let body = TestBody { |
| 227 | + chunks: vec![ |
| 228 | + Chunk::Data("123"), |
| 229 | + Chunk::Data("456"), |
| 230 | + Chunk::Data("789"), |
| 231 | + Chunk::Trailers(trailers()), |
| 232 | + ] |
| 233 | + .into(), |
| 234 | + }; |
| 235 | + let mut body = SdkBody::from_body_1_0(body); |
| 236 | + while let Some(_data) = http_body_0_4::Body::data(&mut body).await {} |
| 237 | + assert_eq!( |
| 238 | + http_body_0_4::Body::trailers(&mut body).await.unwrap(), |
| 239 | + Some(convert_header_map(trailers())) |
| 240 | + ); |
| 241 | + } |
| 242 | + |
| 243 | + #[tokio::test] |
| 244 | + async fn test_errors() { |
| 245 | + let body = TestBody { |
| 246 | + chunks: vec![ |
| 247 | + Chunk::Data("123"), |
| 248 | + Chunk::Data("456"), |
| 249 | + Chunk::Data("789"), |
| 250 | + Chunk::Error("errors!"), |
| 251 | + ] |
| 252 | + .into(), |
| 253 | + }; |
| 254 | + |
| 255 | + let body = SdkBody::from_body_1_0(body); |
| 256 | + let body = ByteStream::new(body); |
| 257 | + body.collect().await.expect_err("body returned an error"); |
| 258 | + } |
| 259 | + #[tokio::test] |
| 260 | + async fn test_no_trailers() { |
| 261 | + let body = TestBody { |
| 262 | + chunks: vec![Chunk::Data("123"), Chunk::Data("456"), Chunk::Data("789")].into(), |
| 263 | + }; |
| 264 | + |
| 265 | + let body = SdkBody::from_body_1_0(body); |
| 266 | + let body = ByteStream::new(body); |
| 267 | + assert_eq!(body.collect().await.unwrap().to_vec(), b"123456789"); |
| 268 | + } |
| 269 | + |
| 270 | + #[test] |
| 271 | + fn test_convert_headers() { |
| 272 | + let mut http1_headermap = http_1x::HeaderMap::new(); |
| 273 | + http1_headermap.append(CT1, HeaderValue::from_static("a")); |
| 274 | + http1_headermap.append(CT1, HeaderValue::from_static("b")); |
| 275 | + http1_headermap.append(CT1, HeaderValue::from_static("c")); |
| 276 | + |
| 277 | + http1_headermap.insert(CL1, HeaderValue::from_static("1234")); |
| 278 | + |
| 279 | + let mut expect = http::HeaderMap::new(); |
| 280 | + expect.append(CT0, http::HeaderValue::from_static("a")); |
| 281 | + expect.append(CT0, http::HeaderValue::from_static("b")); |
| 282 | + expect.append(CT0, http::HeaderValue::from_static("c")); |
| 283 | + |
| 284 | + expect.insert(CL0, http::HeaderValue::from_static("1234")); |
| 285 | + |
| 286 | + assert_eq!(convert_header_map(http1_headermap), expect); |
| 287 | + } |
| 288 | +} |
0 commit comments