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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

use std::io::Result as IoResult;
use std::io::Write;

/// Splits the incoming data into HTTP chunks.
///
/// # Example
///
/// ```
/// use chunked_transfer::Encoder;
/// use std::io::Write;
///
/// let mut decoded = "hello world";
/// let mut encoded: Vec<u8> = vec![];
///
/// {
///     let mut encoder = Encoder::with_chunks_size(&mut encoded, 5);
///     encoder.write_all(decoded.as_bytes());
/// }
///
/// assert_eq!(encoded, b"5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
/// ```
pub struct Encoder<W>
where
    W: Write,
{
    // where to send the result
    output: W,

    // size of each chunk
    chunks_size: usize,

    // data waiting to be sent is stored here
    // This will always be at least 6 bytes long. The first 6 bytes
    // are reserved for the chunk size and \r\n.
    buffer: Vec<u8>,

    // Flushes the internal buffer after each write. This might be useful
    // if data should be sent immediately to downstream consumers
    flush_after_write: bool,
}

// const MAX_CHUNK_SIZE: usize = std::u32::MAX as usize;
// This accounts for four hex digits (enough to hold a u32) plus two bytes
// for the \r\n
const MAX_HEADER_SIZE: usize = 6;

impl<W> Encoder<W>
where
    W: Write,
{
    // pub fn new(output: W) -> Encoder<W> {
    //     Encoder::with_chunks_size(output, 8192)
    // }

    // pub fn with_chunks_size(output: W, chunks: usize) -> Encoder<W> {
    //     let chunks_size = chunks.min(MAX_CHUNK_SIZE);
    //     let mut encoder = Encoder {
    //         output,
    //         chunks_size,
    //         buffer: vec![0; MAX_HEADER_SIZE],
    //         flush_after_write: false,
    //     };
    //     encoder.reset_buffer();
    //     encoder
    // }

    pub fn with_flush_after_write(output: W) -> Encoder<W> {
        let mut encoder = Encoder {
            output,
            chunks_size: 8192,
            buffer: vec![0; MAX_HEADER_SIZE],
            flush_after_write: true,
        };
        encoder.reset_buffer();
        encoder
    }

    fn reset_buffer(&mut self) {
        // Reset buffer, still leaving space for the chunk size. That space
        // will be populated once we know the size of the chunk.
        self.buffer.truncate(MAX_HEADER_SIZE);
    }

    fn is_buffer_empty(&self) -> bool {
        self.buffer.len() == MAX_HEADER_SIZE
    }

    fn buffer_len(&self) -> usize {
        self.buffer.len() - MAX_HEADER_SIZE
    }

    fn send(&mut self) -> IoResult<()> {
        // Never send an empty buffer, because that would be interpreted
        // as the end of the stream, which we indicate explicitly on drop.
        if self.is_buffer_empty() {
            return Ok(());
        }
        // Prepend the length and \r\n to the buffer.
        let prelude = format!("{:x}\r\n", self.buffer_len());
        let prelude = prelude.as_bytes();

        // This should never happen because MAX_CHUNK_SIZE of u32::MAX
        // can always be encoded in 4 hex bytes.
        assert!(
            prelude.len() <= MAX_HEADER_SIZE,
            "invariant failed: prelude longer than MAX_HEADER_SIZE"
        );

        // Copy the prelude into the buffer. For small chunks, this won't necessarily
        // take up all the space that was reserved for the prelude.
        let offset = MAX_HEADER_SIZE - prelude.len();
        self.buffer[offset..MAX_HEADER_SIZE].clone_from_slice(prelude);

        // Append the chunk-finishing \r\n to the buffer.
        self.buffer.write_all(b"\r\n")?;

        self.output.write_all(&self.buffer[offset..])?;
        self.reset_buffer();

        Ok(())
    }
}

impl<W> Write for Encoder<W>
where
    W: Write,
{
    fn write(&mut self, data: &[u8]) -> IoResult<usize> {
        let remaining_buffer_space = self.chunks_size - self.buffer_len();
        let bytes_to_buffer = std::cmp::min(remaining_buffer_space, data.len());
        self.buffer.extend_from_slice(&data[0..bytes_to_buffer]);
        let more_to_write: bool = bytes_to_buffer < data.len();
        if self.flush_after_write || more_to_write {
            self.send()?;
        }

        // If we didn't write the whole thing, keep working on it.
        if more_to_write {
            self.write_all(&data[bytes_to_buffer..])?;
        }
        Ok(data.len())
    }

    fn flush(&mut self) -> IoResult<()> {
        self.send()
    }
}

impl<W> Drop for Encoder<W>
where
    W: Write,
{
    fn drop(&mut self) {
        self.flush().ok();
        write!(self.output, "0\r\n\r\n").ok();
    }
}