1 1 | /*
|
2 2 | * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
3 3 | * SPDX-License-Identifier: Apache-2.0
|
4 4 | */
|
5 5 |
|
6 + | use aws_smithy_types::config_bag::{Storable, StoreReplace};
|
6 7 | use bytes::{Bytes, BytesMut};
|
7 8 | use http_02x::{HeaderMap, HeaderValue};
|
8 9 | use http_body_04x::{Body, SizeHint};
|
9 10 | use pin_project_lite::pin_project;
|
10 11 |
|
11 12 | use std::pin::Pin;
|
12 13 | use std::task::{Context, Poll};
|
13 14 |
|
14 15 | const CRLF: &str = "\r\n";
|
15 16 | const CHUNK_TERMINATOR: &str = "0\r\n";
|
16 17 | const TRAILER_SEPARATOR: &[u8] = b":";
|
17 18 |
|
18 19 | /// Content encoding header value constants
|
19 20 | pub mod header_value {
|
20 21 | /// Header value denoting "aws-chunked" encoding
|
21 22 | pub const AWS_CHUNKED: &str = "aws-chunked";
|
22 23 | }
|
23 24 |
|
24 25 | /// Options used when constructing an [`AwsChunkedBody`].
|
25 - | #[derive(Debug, Default)]
|
26 + | #[derive(Clone, Debug, Default)]
|
26 27 | #[non_exhaustive]
|
27 28 | pub struct AwsChunkedBodyOptions {
|
28 29 | /// The total size of the stream. Because we only support unsigned encoding
|
29 30 | /// this implies that there will only be a single chunk containing the
|
30 31 | /// underlying payload.
|
31 32 | stream_length: u64,
|
32 33 | /// The length of each trailer sent within an `AwsChunkedBody`. Necessary in
|
33 34 | /// order to correctly calculate the total size of the body accurately.
|
34 35 | trailer_lengths: Vec<u64>,
|
36 + | /// Whether the aws-chunked encoding is disabled. This could occur, for instance,
|
37 + | /// if a user specifies a custom checksum, rendering aws-chunked encoding unnecessary.
|
38 + | disabled: bool,
|
39 + | }
|
40 + |
|
41 + | impl Storable for AwsChunkedBodyOptions {
|
42 + | type Storer = StoreReplace<Self>;
|
35 43 | }
|
36 44 |
|
37 45 | impl AwsChunkedBodyOptions {
|
38 46 | /// Create a new [`AwsChunkedBodyOptions`].
|
39 47 | pub fn new(stream_length: u64, trailer_lengths: Vec<u64>) -> Self {
|
40 48 | Self {
|
41 49 | stream_length,
|
42 50 | trailer_lengths,
|
51 + | disabled: false,
|
43 52 | }
|
44 53 | }
|
45 54 |
|
46 55 | fn total_trailer_length(&self) -> u64 {
|
47 56 | self.trailer_lengths.iter().sum::<u64>()
|
48 57 | // We need to account for a CRLF after each trailer name/value pair
|
49 58 | + (self.trailer_lengths.len() * CRLF.len()) as u64
|
50 59 | }
|
51 60 |
|
52 - | /// Set a trailer len
|
61 + | /// Set the stream length in the options
|
62 + | pub fn with_stream_length(mut self, stream_length: u64) -> Self {
|
63 + | self.stream_length = stream_length;
|
64 + | self
|
65 + | }
|
66 + |
|
67 + | /// Append a trailer length to the options
|
53 68 | pub fn with_trailer_len(mut self, trailer_len: u64) -> Self {
|
54 69 | self.trailer_lengths.push(trailer_len);
|
55 70 | self
|
56 71 | }
|
72 + |
|
73 + | /// Create a new [`AwsChunkedBodyOptions`] with aws-chunked encoding disabled.
|
74 + | ///
|
75 + | /// When the option is disabled, the body must not be wrapped in an `AwsChunkedBody`.
|
76 + | pub fn disable_chunked_encoding() -> Self {
|
77 + | Self {
|
78 + | disabled: true,
|
79 + | ..Default::default()
|
80 + | }
|
81 + | }
|
82 + |
|
83 + | /// Return whether aws-chunked encoding is disabled.
|
84 + | pub fn disabled(&self) -> bool {
|
85 + | self.disabled
|
86 + | }
|
87 + |
|
88 + | /// Return the length of the body after `aws-chunked` encoding is applied
|
89 + | pub fn encoded_length(&self) -> u64 {
|
90 + | let mut length = 0;
|
91 + | if self.stream_length != 0 {
|
92 + | length += get_unsigned_chunk_bytes_length(self.stream_length);
|
93 + | }
|
94 + |
|
95 + | // End chunk
|
96 + | length += CHUNK_TERMINATOR.len() as u64;
|
97 + |
|
98 + | // Trailers
|
99 + | for len in self.trailer_lengths.iter() {
|
100 + | length += len + CRLF.len() as u64;
|
101 + | }
|
102 + |
|
103 + | // Encoding terminator
|
104 + | length += CRLF.len() as u64;
|
105 + |
|
106 + | length
|
107 + | }
|
57 108 | }
|
58 109 |
|
59 110 | #[derive(Debug, PartialEq, Eq)]
|
60 111 | enum AwsChunkedBodyState {
|
61 112 | /// Write out the size of the chunk that will follow. Then, transition into the
|
62 113 | /// `WritingChunk` state.
|
63 114 | WritingChunkSize,
|
64 115 | /// Write out the next chunk of data. Multiple polls of the inner body may need to occur before
|
65 116 | /// all data is written out. Once there is no more data to write, transition into the
|
66 117 | /// `WritingTrailers` state.
|
67 118 | WritingChunk,
|
68 119 | /// Write out all trailers associated with this `AwsChunkedBody` and then transition into the
|
69 120 | /// `Closed` state.
|
70 121 | WritingTrailers,
|
71 122 | /// This is the final state. Write out the body terminator and then remain in this state.
|
72 123 | Closed,
|
73 124 | }
|
74 125 |
|
75 126 | pin_project! {
|
76 127 | /// A request body compatible with `Content-Encoding: aws-chunked`. This implementation is only
|
77 128 | /// capable of writing a single chunk and does not support signed chunks.
|
78 129 | ///
|
79 130 | /// Chunked-Body grammar is defined in [ABNF] as:
|
80 131 | ///
|
81 132 | /// ```txt
|
82 133 | /// Chunked-Body = *chunk
|
83 134 | /// last-chunk
|
84 135 | /// chunked-trailer
|
85 136 | /// CRLF
|
86 137 | ///
|
87 138 | /// chunk = chunk-size CRLF chunk-data CRLF
|
88 139 | /// chunk-size = 1*HEXDIG
|
89 140 | /// last-chunk = 1*("0") CRLF
|
90 141 | /// chunked-trailer = *( entity-header CRLF )
|
91 142 | /// entity-header = field-name ":" OWS field-value OWS
|
92 143 | /// ```
|
93 144 | /// For more info on what the abbreviations mean, see https://datatracker.ietf.org/doc/html/rfc7230#section-1.2
|
94 145 | ///
|
95 146 | /// [ABNF]:https://en.wikipedia.org/wiki/Augmented_Backus%E2%80%93Naur_form
|
96 147 | #[derive(Debug)]
|
97 148 | pub struct AwsChunkedBody<InnerBody> {
|
98 149 | #[pin]
|
99 150 | inner: InnerBody,
|
100 151 | #[pin]
|
101 152 | state: AwsChunkedBodyState,
|
102 153 | options: AwsChunkedBodyOptions,
|
103 154 | inner_body_bytes_read_so_far: usize,
|
104 155 | }
|
105 156 | }
|
106 157 |
|
107 158 | impl<Inner> AwsChunkedBody<Inner> {
|
108 159 | /// Wrap the given body in an outer body compatible with `Content-Encoding: aws-chunked`
|
109 160 | pub fn new(body: Inner, options: AwsChunkedBodyOptions) -> Self {
|
110 161 | Self {
|
111 162 | inner: body,
|
112 163 | state: AwsChunkedBodyState::WritingChunkSize,
|
113 164 | options,
|
114 165 | inner_body_bytes_read_so_far: 0,
|
115 166 | }
|
116 167 | }
|
117 - |
|
118 - | fn encoded_length(&self) -> u64 {
|
119 - | let mut length = 0;
|
120 - | if self.options.stream_length != 0 {
|
121 - | length += get_unsigned_chunk_bytes_length(self.options.stream_length);
|
122 - | }
|
123 - |
|
124 - | // End chunk
|
125 - | length += CHUNK_TERMINATOR.len() as u64;
|
126 - |
|
127 - | // Trailers
|
128 - | for len in self.options.trailer_lengths.iter() {
|
129 - | length += len + CRLF.len() as u64;
|
130 - | }
|
131 - |
|
132 - | // Encoding terminator
|
133 - | length += CRLF.len() as u64;
|
134 - |
|
135 - | length
|
136 - | }
|
137 168 | }
|
138 169 |
|
139 170 | fn get_unsigned_chunk_bytes_length(payload_length: u64) -> u64 {
|
140 171 | let hex_repr_len = int_log16(payload_length);
|
141 172 | hex_repr_len + CRLF.len() as u64 + payload_length + CRLF.len() as u64
|
142 173 | }
|
143 174 |
|
144 175 | /// Writes trailers out into a `string` and then converts that `String` to a `Bytes` before
|
145 176 | /// returning.
|
146 177 | ///
|