1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107

use async_std::prelude::*;
use async_std::io::{Read, Write};
use crate::{Error, relay_exact};

pub async fn write_all<O>(output: &mut O, data: &[u8]) -> Result<usize, Error>
    where
    O: Write + Unpin,
{
    match output.write(data).await {
        Ok(size) => Ok(size),
        Err(_) => Err(Error::UnableToWrite),
    }
}

pub async fn write_line<O>(output: &mut O, data: &[u8]) -> Result<usize, Error>
    where
    O: Write + Unpin,
{
    write_all(output, data).await?;
    write_all(output, b"\r\n").await
}

pub async fn flush_write<O>(output: &mut O) -> Result<(), Error>
    where
    O: Write + Unpin,
{
    match output.flush().await {
        Ok(_) => Ok(()),
        Err(_) => Err(Error::UnableToWrite),
    }
}

pub async fn write_exact<O, I>(output: &mut O, input: &mut I, length: usize) -> Result<usize, Error>
    where
    O: Write + Unpin,
    I: Read + Unpin,
{
    relay_exact(input, output, length).await
}

pub async fn write_chunks<O, I>(output: &mut O, input: &mut I, limits: (Option<usize>, Option<usize>)) -> Result<usize, Error>
    where
    O: Write + Unpin,
    I: Read + Unpin,
{
    let (chunklimit, datalimit) = limits;
    let chunksize = match chunklimit {
        Some(chunksize) => chunksize,
        None => 1024,
    };
    let mut total = 0; // all written bytes
    let mut length = 0; // data written bytes
    
    loop {
        let chunksize = match datalimit {
            Some(datalimit) => match length + chunksize > datalimit {
                true => datalimit - length,
                false => chunksize,
            },
            None => chunksize,
        };

        let mut bytes = vec![0u8; chunksize];
        let size = match input.read(&mut bytes).await {
            Ok(size) => size,
            Err(_) => return Err(Error::UnableToRead),
        };
        bytes = bytes[0..size].to_vec();
        length += size;

        total += write_all(output, format!("{:x}\r\n", size).as_bytes()).await?;
        total += write_all(output, &bytes).await?;
        total += write_all(output, b"\r\n").await?;
        flush_write(output).await?;

        if size == 0 {
            break;
        }
    }

    Ok(total)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[async_std::test]
    async fn writes_exact() {
        let mut output = Vec::new();
        let size = write_exact(&mut output, &mut "0123456789".as_bytes(), 5).await.unwrap();
        assert_eq!(size, 5);
        assert_eq!(output, b"01234");
    }

    #[async_std::test]
    async fn writes_chunks() {
        let mut output = Vec::new();
        let size = write_chunks(&mut output, &mut "0123456789".as_bytes(), (Some(3), None)).await.unwrap();
        assert_eq!(size, 35);
        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());
        let mut output = Vec::new();
        let size = write_chunks(&mut output, &mut "0123456789".as_bytes(), (Some(3), Some(4))).await.unwrap();
        assert_eq!(size, 19);
        assert_eq!(output, "3\r\n012\r\n1\r\n3\r\n0\r\n\r\n".as_bytes());
    }
}