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

use core::convert::Infallible;

use std::io::Write;

use crate::{
    bytes::{buf::Chain, Buf, BufMut, BufMutWriter, Bytes, BytesMut, EitherBuf},
    util::buffered::{BufWrite, ListWriteBuf, WriteBuf},
};

/// trait for add http/1 data to buffer that implement [BufWrite] trait.
pub trait H1BufWrite: BufWrite {
    /// write http response head(status code and reason line, header lines) to buffer with fallible
    /// closure. on error path the buffer is reverted back to state before method was called.
    #[inline]
    fn write_buf_head<F, T, E>(&mut self, func: F) -> Result<T, E>
    where
        F: FnOnce(&mut BytesMut) -> Result<T, E>,
    {
        self.write_buf(func)
    }

    /// write `&'static [u8]` to buffer.
    fn write_buf_static(&mut self, bytes: &'static [u8]) {
        let _ = self.write_buf(|buf| {
            buf.put_slice(bytes);
            Ok::<_, Infallible>(())
        });
    }

    /// write bytes to buffer as is.
    fn write_buf_bytes(&mut self, bytes: Bytes) {
        let _ = self.write_buf(|buf| {
            buf.put_slice(bytes.as_ref());
            Ok::<_, Infallible>(())
        });
    }

    /// write bytes to buffer as `transfer-encoding: chunked` encoded.
    fn write_buf_bytes_chunked(&mut self, bytes: Bytes) {
        let _ = self.write_buf(|buf| {
            write!(BufMutWriter(buf), "{:X}\r\n", bytes.len()).unwrap();
            buf.reserve(bytes.len() + 2);
            buf.put_slice(bytes.as_ref());
            buf.put_slice(b"\r\n");
            Ok::<_, Infallible>(())
        });
    }
}

impl H1BufWrite for BytesMut {}

impl<const BUF_LIMIT: usize> H1BufWrite for WriteBuf<BUF_LIMIT> {}

// as special type for eof chunk when using transfer-encoding: chunked
type Eof = Chain<Chain<Bytes, Bytes>, &'static [u8]>;

type EncodedBuf<B, B2> = EitherBuf<B, EitherBuf<B2, &'static [u8]>>;

impl<const BUF_LIMIT: usize> H1BufWrite for ListWriteBuf<EncodedBuf<Bytes, Eof>, BUF_LIMIT> {
    fn write_buf_head<F, T, E>(&mut self, func: F) -> Result<T, E>
    where
        F: FnOnce(&mut BytesMut) -> Result<T, E>,
    {
        // list buffer use BufWrite::write_buf for temporary response head storage.
        // after the head is successfully written we must move the bytes to list.
        self.write_buf(func).map(|t| {
            let bytes = self.split_buf().freeze();
            self.buffer(EitherBuf::Left(bytes));
            t
        })
    }

    #[inline]
    fn write_buf_static(&mut self, bytes: &'static [u8]) {
        self.buffer(EitherBuf::Right(EitherBuf::Right(bytes)));
    }

    #[inline]
    fn write_buf_bytes(&mut self, bytes: Bytes) {
        self.buffer(EitherBuf::Left(bytes));
    }

    #[inline]
    fn write_buf_bytes_chunked(&mut self, bytes: Bytes) {
        let chunk = Bytes::from(format!("{:X}\r\n", bytes.len()))
            .chain(bytes)
            .chain(b"\r\n" as &'static [u8]);
        self.buffer(EitherBuf::Right(EitherBuf::Left(chunk)));
    }
}

impl<L, R> H1BufWrite for EitherBuf<L, R>
where
    L: H1BufWrite,
    R: H1BufWrite,
{
    #[inline]
    fn write_buf_head<F, T, E>(&mut self, func: F) -> Result<T, E>
    where
        F: FnOnce(&mut BytesMut) -> Result<T, E>,
    {
        match *self {
            Self::Left(ref mut l) => l.write_buf_head(func),
            Self::Right(ref mut r) => r.write_buf_head(func),
        }
    }

    #[inline]
    fn write_buf_static(&mut self, bytes: &'static [u8]) {
        match *self {
            Self::Left(ref mut l) => l.write_buf_static(bytes),
            Self::Right(ref mut r) => r.write_buf_static(bytes),
        }
    }

    #[inline]
    fn write_buf_bytes(&mut self, bytes: Bytes) {
        match *self {
            Self::Left(ref mut l) => l.write_buf_bytes(bytes),
            Self::Right(ref mut r) => r.write_buf_bytes(bytes),
        }
    }

    #[inline]
    fn write_buf_bytes_chunked(&mut self, bytes: Bytes) {
        match *self {
            Self::Left(ref mut l) => l.write_buf_bytes_chunked(bytes),
            Self::Right(ref mut r) => r.write_buf_bytes_chunked(bytes),
        }
    }
}