async_httplib/
write.rs

1use async_std::prelude::*;
2use async_std::io::{Read, Write};
3use std::io::{Error, ErrorKind};
4use crate::{relay_exact};
5
6pub async fn write_slice<O>(output: &mut O, data: &[u8]) -> Result<usize, Error>
7    where
8    O: Write + Unpin,
9{
10    output.write(data).await
11}
12
13pub async fn flush_write<O>(output: &mut O) -> Result<(), Error>
14    where
15    O: Write + Unpin,
16{
17    output.flush().await
18}
19
20pub async fn write_exact<O, I>(output: &mut O, input: &mut I, length: usize) -> Result<usize, Error>
21    where
22    O: Write + Unpin,
23    I: Read + Unpin,
24{
25    relay_exact(input, output, length).await
26}
27
28pub async fn write_all<O, I>(output: &mut O, input: &mut I, limit: Option<usize>) -> Result<usize, Error>
29    where
30    O: Write + Unpin,
31    I: Read + Unpin,
32{
33    let mut total = 0; // all written bytes
34    let mut length = 0; // data written bytes
35    
36    loop {
37        let mut bytes = vec![0u8; 1024];
38        let size = input.read(&mut bytes).await?;
39        bytes = bytes[0..size].to_vec();
40        length += size;
41
42        if size == 0 {
43            break;
44        } else if limit.is_some() && length > limit.unwrap() {
45            return Err(Error::new(ErrorKind::InvalidData, format!("The operation hit the limit of {} bytes while writing.", limit.unwrap())));
46        }
47
48        total += output.write(&bytes).await?;
49        output.flush().await?;
50    }
51
52    Ok(total)
53}
54
55pub async fn write_chunks<O, I>(output: &mut O, input: &mut I, limits: (Option<usize>, Option<usize>)) -> Result<usize, Error>
56    where
57    O: Write + Unpin,
58    I: Read + Unpin,
59{
60    let (chunklimit, datalimit) = limits;
61    let chunksize = match chunklimit {
62        Some(chunksize) => chunksize,
63        None => 1024,
64    };
65    let mut total = 0; // all written bytes
66    let mut length = 0; // data written bytes
67    
68    loop {
69        let mut bytes = vec![0u8; chunksize];
70        let size = input.read(&mut bytes).await?;
71        bytes = bytes[0..size].to_vec();
72        length += size;
73
74        if datalimit.is_some() && length > datalimit.unwrap() {
75            return Err(Error::new(ErrorKind::InvalidData, format!("The operation hit the limit of {} bytes while writing chunked HTTP body.", datalimit.unwrap())));
76        }
77
78        total += output.write(format!("{:x}\r\n", size).as_bytes()).await?;
79        total += output.write(&bytes).await?;
80        total += output.write(b"\r\n").await?;
81        output.flush().await?;
82
83        if size == 0 {
84            break;
85        }
86    }
87
88    Ok(total)
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[async_std::test]
96    async fn writes_exact() {
97        let mut output = Vec::new();
98        let size = write_exact(&mut output, &mut "0123456789".as_bytes(), 5).await.unwrap();
99        assert_eq!(size, 5);
100        assert_eq!(output, b"01234");
101    }
102
103    #[async_std::test]
104    async fn writes_all() {
105        let mut output = Vec::new();
106        let size = write_all(&mut output, &mut "0123456789".as_bytes(), None).await.unwrap();
107        assert_eq!(size, 10);
108        assert_eq!(output, b"0123456789");
109        let mut output = Vec::new();
110        let exceeded = write_all(&mut output, &mut "012".as_bytes(), Some(2)).await;
111        assert!(exceeded.is_err());
112    }
113
114    #[async_std::test]
115    async fn writes_chunks() {
116        let mut output = Vec::new();
117        let size = write_chunks(&mut output, &mut "0123456789".as_bytes(), (Some(3), None)).await.unwrap();
118        assert_eq!(size, 35);
119        assert_eq!(output, "3\r\n012\r\n3\r\n345\r\n3\r\n678\r\n1\r\n9\r\n0\r\n\r\n".as_bytes());
120        let mut output = Vec::new();
121        let exceeded = write_chunks(&mut output, &mut "0123456789".as_bytes(), (Some(3), Some(4))).await;
122        assert!(exceeded.is_err());
123    }
124}
async_httplib/write.rs

async_httplib/
write.rs
