async_httplib/
read.rs

1use std::io::{Error, ErrorKind};
2use async_std::prelude::*;
3use async_std::io::{Read};
4
5pub async fn read_first_line<I>(input: &mut I, data: (&mut Vec<u8>, &mut Vec<u8>, &mut Vec<u8>), limit: Option<usize>) -> Result<usize, Error>
6    where
7    I: Read + Unpin,
8{
9    let mut part = 0;
10    let mut length = 0;
11    let mut stage = 0; // 0..data, 1..\r, 2..\n
12
13    loop {
14        let mut bytes = [0u8];
15        let size = input.read(&mut bytes).await?;
16        length += size;
17
18        if size == 0 {
19            break;
20        } else if limit.is_some() && limit.unwrap() < length { // method + url + version = 2065
21            return Err(Error::new(ErrorKind::InvalidData, "The operation hit the limit of {} bytes while reading the HTTP first line."));
22        } else if bytes[0] == 32 { // space
23            part += 1;
24            continue;
25        } else if bytes[0] == 13 { // \r
26            stage = 1;
27            continue;
28        } else if bytes[0] == 10 { // \n
29            if stage == 1 {
30                break;
31            } else {
32                return Err(Error::new(ErrorKind::InvalidData, "The data is not a valid HTTP first line."));
33            }
34        }
35
36        match part {
37            0 => data.0.push(bytes[0]),
38            1 => data.1.push(bytes[0]),
39            _ => data.2.push(bytes[0]),
40        };
41    }
42
43    Ok(length)
44}
45
46pub async fn read_header_line<I>(input: &mut I, data: (&mut Vec<u8>, &mut Vec<u8>), limit: Option<usize>) -> Result<usize, Error>
47    where
48    I: Read + Unpin,
49{
50    let mut length = 0;
51    let mut stage = 0; // 0..name, 1..:, 2..space, 3..value, 4..\r, 5..\n
52
53    loop {
54        let mut bytes = [0u8];
55        let size = input.read(&mut bytes).await?;
56        length += size;
57
58        if size == 0 {
59            break;
60        } else if limit.is_some() && limit.unwrap() < length {
61            return Err(Error::new(ErrorKind::InvalidData, format!("The operation hit the limit of {} bytes while reading the HTTP header line.", limit.unwrap())));
62        } else if stage == 0 && bytes[0] == 58 { // first :
63            stage = 1;
64            continue;
65        } else if stage == 1 && bytes[0] == 32 { // first space
66            stage = 2;
67            continue;
68        } else if bytes[0] == 13 { // first/second \r
69            if stage == 0 || stage == 2 {
70                stage = 3;
71                continue;
72            } else {
73                return Err(Error::new(ErrorKind::InvalidData, "The data is not a valid HTTP header line."));
74            }
75        } else if bytes[0] == 10 { // first/second \n
76            if stage == 3 {
77                break;
78            } else {
79                return Err(Error::new(ErrorKind::InvalidData, "The data is not a valid HTTP header line."));
80            }
81        }
82
83        if stage == 0 {
84            data.0.push(bytes[0]);
85        } else if stage == 2 {
86            data.1.push(bytes[0]);
87        }
88    }
89
90    Ok(length)
91}
92
93pub async fn read_exact<I>(input: &mut I, data: &mut Vec<u8>, length: usize) -> Result<usize, Error>
94    where
95    I: Read + Unpin,
96{
97    let mut bytes = vec![0u8; length];
98    input.read_exact(&mut bytes).await?;
99    data.append(&mut bytes);
100
101    Ok(length)
102}
103
104pub async fn read_chunk_line<I>(input: &mut I, data: (&mut Vec<u8>, &mut Vec<u8>), limit: Option<usize>) -> Result<usize, Error>
105    where
106    I: Read + Unpin,
107{
108    let mut length = 0;
109    let mut stage = 0; // 0..number, 1..extension 2..\r, 3=\n
110
111    loop {
112        let mut bytes = [0u8];
113        let size = input.read(&mut bytes).await?;
114        length += size;
115
116        if size == 0 { // end of data
117            break;
118        } else if limit.is_some() && limit.unwrap() < length {
119            return Err(Error::new(ErrorKind::InvalidData, format!("The operation hit the limit of {} bytes while reading the HTTP body chunk line.", limit.unwrap())));
120        } else if stage == 0 && bytes[0] == 59 { // char ;
121            stage = 1;
122            continue;
123        } else if bytes[0] == 13 { // char \r
124            if stage == 0 || stage == 1 {
125                stage = 2;
126                continue;
127            } else {
128                return Err(Error::new(ErrorKind::InvalidData, "The data is not a valid HTTP chunk line."));
129            }
130        } else if bytes[0] == 10 { // char \n
131            break;
132        }
133        match stage {
134            0 => data.0.push(bytes[0]),
135            1 => data.1.push(bytes[0]),
136            _ => (),
137        };
138    }
139
140    Ok(length)
141}
142
143pub async fn read_chunks<I>(input: &mut I, data: &mut Vec<u8>, limits: (Option<usize>, Option<usize>)) -> Result<usize, Error>
144    where
145    I: Read + Unpin,
146{
147    let (chunklimit, datalimit) = limits;
148    let mut length = 0;
149    let mut total = 0; // actual data size
150
151    loop {
152        let (mut size, mut ext) = (vec![], vec![]);
153        length += read_chunk_line(input, (&mut size, &mut ext), chunklimit).await?;
154        let size = match String::from_utf8(size) {
155            Ok(length) => match i64::from_str_radix(&length, 16) {
156                Ok(length) => length as usize,
157                Err(e) => return Err(Error::new(ErrorKind::InvalidData, e.to_string())),
158            },
159            Err(e) => return Err(Error::new(ErrorKind::InvalidData, e.to_string())),
160        };
161        total += size;
162        if size == 0 {
163            length += read_exact(input, &mut Vec::new(), 2).await?;
164            break; // last chunk
165        } else if datalimit.is_some() && total + size > datalimit.unwrap() {
166            return Err(Error::new(ErrorKind::InvalidData, format!("The operation hit the limit of {} bytes while reading the HTTP body chunk data.", datalimit.unwrap())));
167        } else {
168            length += read_exact(input, data, size).await?;
169            length += read_exact(input, &mut Vec::new(), 2).await?;
170        }
171    }
172
173    Ok(length)
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[async_std::test]
181    async fn reads_first_line() {
182        let (mut a, mut b, mut c) = (vec![], vec![], vec![]);
183        let size = read_first_line(&mut "OPTIONS /path HTTP/1.1\r\n".as_bytes(), (&mut a, &mut b, &mut c), None).await.unwrap();
184        assert_eq!(size, 24);
185        assert_eq!(a, b"OPTIONS");
186        assert_eq!(b, b"/path");
187        assert_eq!(c, b"HTTP/1.1");
188        let (mut a, mut b, mut c) = (vec![], vec![], vec![]);
189        let exceeded = read_first_line(&mut "OPTI\r\n".as_bytes(), (&mut a, &mut b, &mut c), Some(1)).await;
190        assert!(exceeded.is_err());
191    }
192
193    #[async_std::test]
194    async fn reads_header() {
195        let (mut name, mut value) = (vec![], vec![]);
196        let size = read_header_line(&mut "Foo: foo\r\nBar: bar\r\n".as_bytes(), (&mut name, &mut value), None).await.unwrap();
197        assert_eq!(size, 10);
198        assert_eq!(name, b"Foo");
199        assert_eq!(value, b"foo");
200        let (mut name, mut value) = (vec![], vec![]);
201        let size = read_header_line(&mut "\r\n".as_bytes(), (&mut name, &mut value), None).await.unwrap();
202        assert_eq!(size, 2);
203        assert_eq!(name, b"");
204        assert_eq!(value, b"");
205        let exceeded = read_header_line(&mut "Foo".as_bytes(), (&mut name, &mut value), Some(1)).await;
206        assert!(exceeded.is_err());
207    }
208
209    #[async_std::test]
210    async fn reads_exact() {
211        let mut output = Vec::new();
212        read_exact(&mut "0123456789".as_bytes(), &mut output, 5).await.unwrap();
213        assert_eq!(String::from_utf8(output).unwrap(), "01234");
214    }
215
216    #[async_std::test]
217    async fn reads_chunk_line() {
218        let (mut number, mut ext) = (vec![], vec![]);
219        let size = read_chunk_line(&mut "6;ex;ex\r\n".as_bytes(), (&mut number, &mut ext), None).await.unwrap();
220        assert_eq!(size, 9);
221        assert_eq!(String::from_utf8(number).unwrap(), "6");
222        assert_eq!(String::from_utf8(ext).unwrap(), "ex;ex");
223        let (mut number, mut ext) = (vec![], vec![]);
224        let exceeded = read_chunk_line(&mut "6\r\n".as_bytes(), (&mut number, &mut ext), Some(1)).await;
225        assert!(exceeded.is_err());
226    }
227
228    #[async_std::test]
229    async fn reads_chunks() {
230        let mut output = Vec::new();
231        let size = read_chunks(&mut "6\r\nHello \r\n6;ex=fo\r\nWorld!\r\n0\r\n\r\nTrail: er\r\n\r\n".as_bytes(), &mut output, (None, None)).await.unwrap(); // with extension `ex=fo` and trailer `Trail: er`
232        assert_eq!(size, 33);
233        assert_eq!(String::from_utf8(output).unwrap(), "Hello World!");
234        let mut output = Vec::new();
235        let exceeded = read_chunks(&mut "6\r\nHello 0\r\n\r\n".as_bytes(), &mut output, (Some(1), None)).await;
236        assert!(exceeded.is_err());
237    }
238}