Skip to main content

aws_runtime/content_encoding/
options.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use aws_smithy_types::config_bag::{Storable, StoreReplace};
7
8use super::{
9    CHUNK_SIGNATURE_BEGIN, CHUNK_TERMINATOR, CRLF, DEFAULT_CHUNK_SIZE_BYTE, SIGNATURE_LENGTH,
10};
11
12/// Options used when constructing an [`AwsChunkedBody`](super::AwsChunkedBody).
13#[derive(Clone, Debug, Default)]
14#[non_exhaustive]
15pub struct AwsChunkedBodyOptions {
16    /// The total size of the stream.
17    pub(crate) stream_length: u64,
18    /// The length of each trailer sent within an `AwsChunkedBody`. Necessary in
19    /// order to correctly calculate the total size of the body accurately.
20    pub(crate) trailer_lengths: Vec<u64>,
21    /// Whether the aws-chunked encoding is disabled. This could occur, for instance,
22    /// if a user specifies a custom checksum, rendering aws-chunked encoding unnecessary.
23    pub(crate) disabled: bool,
24    /// Whether chunks and trailer are signed.
25    pub(crate) is_signed: bool,
26    /// The size of each chunk in bytes.
27    /// None means use default (64 KiB)
28    pub(crate) chunk_size: Option<usize>,
29}
30
31impl Storable for AwsChunkedBodyOptions {
32    type Storer = StoreReplace<Self>;
33}
34
35impl AwsChunkedBodyOptions {
36    /// Create a new [`AwsChunkedBodyOptions`].
37    pub fn new(stream_length: u64, trailer_lengths: Vec<u64>) -> Self {
38        Self {
39            stream_length,
40            trailer_lengths,
41            disabled: false,
42            is_signed: false,
43            chunk_size: None,
44        }
45    }
46
47    /// Set the chunk size for aws-chunked encoding.
48    ///
49    /// This allows customizing the size of each chunk when using aws-chunked encoding.
50    /// The chunk size is validated by the interceptor (minimum 8 KiB).
51    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
52        self.chunk_size = Some(chunk_size);
53        self
54    }
55
56    /// Get the chunk size that will be used for aws-chunked encoding.
57    ///
58    /// Returns the configured chunk size, or the default if not set.
59    pub fn chunk_size(&self) -> usize {
60        self.chunk_size.unwrap_or(DEFAULT_CHUNK_SIZE_BYTE)
61    }
62
63    pub(super) fn total_trailer_length(&self) -> u64 {
64        self.trailer_lengths.iter().sum::<u64>()
65            // We need to account for a CRLF after each trailer name/value pair
66            + (self.trailer_lengths.len() * CRLF.len()) as u64
67    }
68
69    /// Set the stream length in the options
70    pub fn with_stream_length(mut self, stream_length: u64) -> Self {
71        self.stream_length = stream_length;
72        self
73    }
74
75    /// Append a trailer length to the options
76    pub fn with_trailer_len(mut self, trailer_len: u64) -> Self {
77        self.trailer_lengths.push(trailer_len);
78        self
79    }
80
81    /// Return whether there are no trailers
82    pub fn is_trailer_empty(&self) -> bool {
83        self.trailer_lengths.is_empty()
84    }
85
86    /// Create a new [`AwsChunkedBodyOptions`] with aws-chunked encoding disabled.
87    ///
88    /// When the option is disabled, the body must not be wrapped in an `AwsChunkedBody`.
89    pub fn disable_chunked_encoding() -> Self {
90        Self {
91            disabled: true,
92            ..Default::default()
93        }
94    }
95
96    /// Return whether aws-chunked encoding is disabled.
97    pub fn disabled(&self) -> bool {
98        self.disabled
99    }
100
101    /// Set whether to use signed chunked encoding
102    pub fn signed_chunked_encoding(mut self, is_signed: bool) -> Self {
103        self.is_signed = is_signed;
104        self
105    }
106
107    /// Return the length of the body after `aws-chunked` encoding is applied
108    pub fn encoded_length(&self) -> u64 {
109        if self.is_signed {
110            self.signed_encoded_length()
111        } else {
112            self.unsigned_encoded_length()
113        }
114    }
115
116    fn signed_encoded_length(&self) -> u64 {
117        let number_of_data_chunks = self.stream_length / self.chunk_size() as u64;
118        let remaining_data_chunk = self.stream_length % self.chunk_size() as u64;
119
120        let mut length = number_of_data_chunks
121            * get_signed_chunk_bytes_length(self.chunk_size() as u64)
122            + if remaining_data_chunk > 0 {
123                get_signed_chunk_bytes_length(remaining_data_chunk)
124            } else {
125                0
126            };
127
128        // End chunk
129        length += get_signed_chunk_bytes_length(0);
130
131        length -= CRLF.len() as u64; // The last CRLF is not needed for 0-sized signed chunk
132
133        // Trailers
134        for len in self.trailer_lengths.iter() {
135            length += len + CRLF.len() as u64;
136        }
137
138        // Encoding terminator
139        length += CRLF.len() as u64;
140
141        length
142    }
143
144    fn unsigned_encoded_length(&self) -> u64 {
145        let number_of_data_chunks = self.stream_length / self.chunk_size() as u64;
146        let remaining_data_chunk = self.stream_length % self.chunk_size() as u64;
147
148        let mut length = number_of_data_chunks
149            * get_unsigned_chunk_bytes_length(self.chunk_size() as u64)
150            + if remaining_data_chunk > 0 {
151                get_unsigned_chunk_bytes_length(remaining_data_chunk)
152            } else {
153                0
154            };
155
156        // End chunk
157        length += CHUNK_TERMINATOR.len() as u64;
158
159        // Trailers
160        for len in self.trailer_lengths.iter() {
161            length += len + CRLF.len() as u64;
162        }
163
164        // Encoding terminator
165        length += CRLF.len() as u64;
166
167        length
168    }
169}
170
171fn int_log16<T>(mut i: T) -> u64
172where
173    T: std::ops::DivAssign + PartialOrd + From<u8> + Copy,
174{
175    let mut len = 0;
176    let zero = T::from(0);
177    let sixteen = T::from(16);
178
179    // Handle an edge case where 0 is passed in, which still requires 1 hex digit to represent
180    if i == zero {
181        return 1;
182    }
183
184    while i > zero {
185        i /= sixteen;
186        len += 1;
187    }
188
189    len
190}
191
192// Return the length of a signed chunk:
193//
194// A signed chunk looks like:
195// 10000;chunk-signature=b474d8862b1487a5145d686f57f013e54db672cee1c953b3010fb58501ef5aa2\r\n
196// <65536-bytes>\r\n
197fn get_signed_chunk_bytes_length(payload_length: u64) -> u64 {
198    let hex_repr_len = int_log16(payload_length);
199    hex_repr_len
200        + CHUNK_SIGNATURE_BEGIN.len() as u64
201        + SIGNATURE_LENGTH as u64
202        + CRLF.len() as u64
203        + payload_length
204        + CRLF.len() as u64
205}
206
207// Return the length of an unsigned chunk:
208//
209// An unsigned chunk looks like:
210// 10000\r\n
211// <65536-bytes>\r\n
212fn get_unsigned_chunk_bytes_length(payload_length: u64) -> u64 {
213    let hex_repr_len = int_log16(payload_length);
214    hex_repr_len + CRLF.len() as u64 + payload_length + CRLF.len() as u64
215}
216
217#[cfg(test)]
218mod tests {
219    use super::int_log16;
220
221    #[test]
222    fn test_int_log16() {
223        assert_eq!(int_log16(0u64), 1); // 0x0
224        assert_eq!(int_log16(1u64), 1); // 0x1
225        assert_eq!(int_log16(15u64), 1); // 0xF
226        assert_eq!(int_log16(16u64), 2); // 0x10
227        assert_eq!(int_log16(255u64), 2); // 0xFF
228        assert_eq!(int_log16(256u64), 3); // 0x100
229        assert_eq!(int_log16(4095u64), 3); // 0xFFF
230        assert_eq!(int_log16(4096u64), 4); // 0x1000
231        assert_eq!(int_log16(65535u64), 4); // 0xFFFF
232        assert_eq!(int_log16(65536u64), 5); // 0x10000
233        assert_eq!(int_log16(1048575u64), 5); // 0xFFFFF
234        assert_eq!(int_log16(1048576u64), 6); // 0x100000
235        assert_eq!(int_log16(u64::MAX), 16); // 0xFFFFFFFFFFFFFFFF
236    }
237}