ureq/
body.rs

1use crate::stream::Stream;
2use std::fmt;
3use std::io::{self, copy, empty, Cursor, Read, Write};
4
5#[cfg(feature = "charset")]
6use crate::response::DEFAULT_CHARACTER_SET;
7#[cfg(feature = "charset")]
8use encoding_rs::Encoding;
9
10/// The different kinds of bodies to send.
11///
12/// *Internal API*
13pub(crate) enum Payload<'a> {
14    Empty,
15    Text(&'a str, String),
16    Reader(Box<dyn Read + 'a>),
17    Bytes(&'a [u8]),
18}
19
20impl fmt::Debug for Payload<'_> {
21    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
22        match self {
23            Payload::Empty => write!(f, "Empty"),
24            Payload::Text(t, _) => write!(f, "{}", t),
25            Payload::Reader(_) => write!(f, "Reader"),
26            Payload::Bytes(v) => write!(f, "{:?}", v),
27        }
28    }
29}
30
31#[allow(clippy::derivable_impls)]
32impl Default for Payload<'_> {
33    fn default() -> Self {
34        Payload::Empty
35    }
36}
37
38/// The size of the body.
39///
40/// *Internal API*
41#[derive(Debug)]
42pub(crate) enum BodySize {
43    Empty,
44    Unknown,
45    Known(u64),
46}
47
48/// Payloads are turned into this type where we can hold both a size and the reader.
49///
50/// *Internal API*
51pub(crate) struct SizedReader<'a> {
52    pub size: BodySize,
53    pub reader: Box<dyn Read + 'a>,
54}
55
56impl fmt::Debug for SizedReader<'_> {
57    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
58        write!(f, "SizedReader[size={:?},reader]", self.size)
59    }
60}
61
62impl<'a> SizedReader<'a> {
63    fn new(size: BodySize, reader: Box<dyn Read + 'a>) -> Self {
64        SizedReader { size, reader }
65    }
66}
67
68impl<'a> Payload<'a> {
69    pub fn into_read(self) -> SizedReader<'a> {
70        match self {
71            Payload::Empty => SizedReader::new(BodySize::Empty, Box::new(empty())),
72            Payload::Text(text, _charset) => {
73                #[cfg(feature = "charset")]
74                let bytes = {
75                    let encoding = Encoding::for_label(_charset.as_bytes())
76                        .or_else(|| Encoding::for_label(DEFAULT_CHARACTER_SET.as_bytes()))
77                        .unwrap();
78                    encoding.encode(text).0
79                };
80                #[cfg(not(feature = "charset"))]
81                let bytes = text.as_bytes();
82                let len = bytes.len();
83                let cursor = Cursor::new(bytes);
84                SizedReader::new(BodySize::Known(len as u64), Box::new(cursor))
85            }
86            Payload::Reader(read) => SizedReader::new(BodySize::Unknown, read),
87            Payload::Bytes(bytes) => {
88                let len = bytes.len();
89                let cursor = Cursor::new(bytes);
90                SizedReader::new(BodySize::Known(len as u64), Box::new(cursor))
91            }
92        }
93    }
94}
95
96const CHUNK_MAX_SIZE: usize = 0x4000; // Maximum size of a TLS fragment
97const CHUNK_HEADER_MAX_SIZE: usize = 6; // four hex digits plus "\r\n"
98const CHUNK_FOOTER_SIZE: usize = 2; // "\r\n"
99const CHUNK_MAX_PAYLOAD_SIZE: usize = CHUNK_MAX_SIZE - CHUNK_HEADER_MAX_SIZE - CHUNK_FOOTER_SIZE;
100
101// copy_chunks() improves over chunked_transfer's Encoder + io::copy with the
102// following performance optimizations:
103// 1) It avoid copying memory.
104// 2) chunked_transfer's Encoder issues 4 separate write() per chunk. This is costly
105//    overhead. Instead, we do a single write() per chunk.
106// The measured benefit on a Linux machine is a 50% reduction in CPU usage on a https connection.
107fn copy_chunked<R: Read, W: Write>(reader: &mut R, writer: &mut W) -> io::Result<u64> {
108    // The chunk layout is:
109    // header:header_max_size | payload:max_payload_size | footer:footer_size
110    let mut chunk = Vec::with_capacity(CHUNK_MAX_SIZE);
111    let mut written = 0;
112    loop {
113        // We first read the payload
114        chunk.resize(CHUNK_HEADER_MAX_SIZE, 0);
115        let payload_size = reader
116            .take(CHUNK_MAX_PAYLOAD_SIZE as u64)
117            .read_to_end(&mut chunk)?;
118
119        // Then write the header
120        let header_str = format!("{:x}\r\n", payload_size);
121        let header = header_str.as_bytes();
122        assert!(header.len() <= CHUNK_HEADER_MAX_SIZE);
123        let start_index = CHUNK_HEADER_MAX_SIZE - header.len();
124        (&mut chunk[start_index..]).write_all(header).unwrap();
125
126        // And add the footer
127        chunk.extend_from_slice(b"\r\n");
128
129        // Finally Write the chunk
130        writer.write_all(&chunk[start_index..])?;
131        written += payload_size as u64;
132
133        // On EOF, we wrote a 0 sized chunk. This is what the chunked encoding protocol requires.
134        if payload_size == 0 {
135            return Ok(written);
136        }
137    }
138}
139
140/// Helper to send a body, either as chunked or not.
141pub(crate) fn send_body(
142    mut body: SizedReader,
143    do_chunk: bool,
144    stream: &mut Stream,
145) -> io::Result<()> {
146    if do_chunk {
147        copy_chunked(&mut body.reader, stream)?;
148    } else {
149        copy(&mut body.reader, stream)?;
150    };
151
152    Ok(())
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn test_copy_chunked() {
161        let mut source = Vec::<u8>::new();
162        source.resize(CHUNK_MAX_PAYLOAD_SIZE, 33);
163        source.extend_from_slice(b"hello world");
164
165        let mut dest = Vec::<u8>::new();
166        copy_chunked(&mut &source[..], &mut dest).unwrap();
167
168        let mut dest_expected = Vec::<u8>::new();
169        dest_expected.extend_from_slice(format!("{:x}\r\n", CHUNK_MAX_PAYLOAD_SIZE).as_bytes());
170        dest_expected.resize(dest_expected.len() + CHUNK_MAX_PAYLOAD_SIZE, 33);
171        dest_expected.extend_from_slice(b"\r\n");
172
173        dest_expected.extend_from_slice(b"b\r\nhello world\r\n");
174        dest_expected.extend_from_slice(b"0\r\n\r\n");
175
176        assert_eq!(dest, dest_expected);
177    }
178}