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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
use std::io::Write;
use std::panic::UnwindSafe;

use libc::c_void;

use crate::nxt_unit::{
    self, nxt_unit_buf_send, nxt_unit_buf_t, nxt_unit_request_info_t, nxt_unit_response_buf_alloc,
};

use crate::error::{IntoUnitResult, UnitError, UnitResult};
use crate::request::Request;

/// A buffer for constructing an initial response.
///
/// This object is created by calling the
/// [`create_response()`](Request::create_response) method on a [`Request`].
///
/// Dropping this object will _not_ send the response; it must be manually sent
/// with the [`Response::send()`] method.
pub struct Response<'a> {
    pub(crate) request: &'a Request<'a>,
}

impl<'a> Response<'a> {
    pub fn add_field<N: AsRef<[u8]>, V: AsRef<[u8]>>(&self, name: N, value: V) -> UnitResult<()> {
        // SAFETY: Unit's C API does state and buffer size checks internally.
        // This structure is not Send nor Sync, so sharing it is fine.
        unsafe {
            nxt_unit::nxt_unit_response_add_field(
                self.request.nxt_request,
                name.as_ref().as_ptr() as *const libc::c_char,
                name.as_ref().len() as u8,
                value.as_ref().as_ptr() as *const libc::c_char,
                value.as_ref().len() as u32,
            )
            .into_unit_result()
        }
    }

    pub fn add_content<C: AsRef<[u8]>>(&self, content: C) -> UnitResult<()> {
        // SAFETY: Unit's C API does state and buffer size checks internally.
        // This structure is not Send nor Sync, so sharing it is fine.
        unsafe {
            nxt_unit::nxt_unit_response_add_content(
                self.request.nxt_request,
                content.as_ref().as_ptr() as *const c_void,
                content.as_ref().len() as u32,
            )
            .into_unit_result()
        }
    }

    pub fn realloc(&self, max_fields_count: usize, max_fields_size: usize) -> UnitResult<()> {
        // SAFETY: Unit's C API does state and buffer size checks internally.
        // This structure is not Send nor Sync, so sharing it is fine.
        unsafe {
            nxt_unit::nxt_unit_response_realloc(
                self.request.nxt_request,
                max_fields_count as u32,
                max_fields_size as u32,
            )
            .into_unit_result()
        }
    }

    pub fn send(&self) -> UnitResult<()> {
        // SAFETY: Unit's C API does state and buffer size checks internally.
        // This structure is not Send nor Sync, so sharing it is fine.
        unsafe { nxt_unit::nxt_unit_response_send(self.request.nxt_request).into_unit_result() }
    }
}

/// A writer that writes to a Unit shared memory response buffer.
///
/// This object is created using [`Request::write_chunks()`] or
/// [`Request::send_chunks_with_writer()`].
///
/// A chunk will be immediately sent to the client once the writer's memory
/// buffer reaches `chunk_size`, or [`flush()`](std::io::Write::flush) is
/// called on the writer, and a new shared memory buffer will be allocated.
///
/// The writer will also flush when dropped, but any errors that happen during
/// a drop will panic.
pub struct BodyWriter<'a> {
    _lifetime: std::marker::PhantomData<&'a mut ()>,
    nxt_request: *mut nxt_unit_request_info_t,
    response_buffer: *mut nxt_unit_buf_t,
    chunk_cursor: *mut u8,
    chunk_size: usize,
    bytes_remaining: usize,
}

impl UnwindSafe for BodyWriter<'_> {}

impl<'a> BodyWriter<'a> {
    pub(crate) fn new(request: &'a Request<'a>, chunk_size: usize) -> std::io::Result<Self> {
        let mut writer = BodyWriter {
            _lifetime: Default::default(),
            nxt_request: request.nxt_request,
            response_buffer: std::ptr::null_mut(),
            chunk_cursor: std::ptr::null_mut(),
            chunk_size,
            bytes_remaining: 0,
        };
        writer.allocate_buffer()?;
        Ok(writer)
    }

    fn allocate_buffer(&mut self) -> std::io::Result<()> {
        unsafe {
            let buf = nxt_unit_response_buf_alloc(self.nxt_request, self.chunk_size as u32);

            if buf.is_null() {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Other,
                    "Could not allocate response buffer in Unit's shared memory",
                ));
            }

            self.response_buffer = buf;
            self.chunk_cursor = (*buf).start as *mut u8;
            self.bytes_remaining = self.chunk_size;
        }

        Ok(())
    }

    /// Copy from a reader to this writer without using an intermediary buffer.
    ///
    /// Normally the [`Write`](std::io::Write) trait receives an input buffer to
    /// copy from, and the `ResponseWriter` writer will copy from it into Unit's
    /// shared memory.
    ///
    /// This method will instead give Unit's shared memory buffer directly to
    /// the [`Read`](std::io::Read) trait in order to skip copying to a third
    /// temporary buffer (such as when using [`std::io::copy`]).
    pub fn copy_from_reader<R: std::io::Read>(&mut self, mut r: R) -> std::io::Result<()> {
        loop {
            if self.bytes_remaining == 0 {
                self.flush()?;
                self.allocate_buffer()?;
            }

            // SAFETY: Allocated by Unit and fully initialized with memset.
            // TODO: The memset is unnecessary, use std::io::ReadBuf once that
            // is stabilized.
            let write_buffer = unsafe {
                libc::memset(self.chunk_cursor as *mut c_void, 0, self.bytes_remaining);
                std::slice::from_raw_parts_mut(self.chunk_cursor, self.bytes_remaining)
            };

            let bytes = r.read(write_buffer)?;

            self.chunk_cursor = unsafe { self.chunk_cursor.add(bytes) };
            self.bytes_remaining -= bytes;

            if bytes == 0 {
                break;
            }
        }

        return Ok(());
    }
}

impl std::io::Write for BodyWriter<'_> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let buf = if buf.len() >= self.bytes_remaining && !buf.is_empty() {
            if self.bytes_remaining == 0 {
                self.flush()?;
                self.allocate_buffer()?;
            }

            &buf[..buf.len().min(self.bytes_remaining)]
        } else {
            buf
        };

        // SAFETY: The target region is not initialized and never made
        // available to the user until it is, so it cannot overlap.
        // The buffer length is truncated above to fit the target's limit.
        unsafe {
            std::ptr::copy_nonoverlapping(buf.as_ptr(), self.chunk_cursor, buf.len());
            self.chunk_cursor = self.chunk_cursor.add(buf.len());
        }
        self.bytes_remaining -= buf.len();

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        if self.response_buffer.is_null() || self.bytes_remaining == self.chunk_size {
            return Ok(());
        }

        unsafe {
            (*self.response_buffer).free = (*self.response_buffer)
                .start
                .add(self.chunk_size - self.bytes_remaining);
            nxt_unit_buf_send(self.response_buffer)
                .into_unit_result()
                .map_err(|UnitError(_)| {
                    std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "Could not send response buffer to Unit server",
                    )
                })?;
        }

        self.response_buffer = std::ptr::null_mut();
        self.bytes_remaining = 0;

        Ok(())
    }
}

impl Drop for BodyWriter<'_> {
    fn drop(&mut self) {
        if !self.chunk_cursor.is_null() {
            if std::thread::panicking() {
                unsafe {
                    nxt_unit::nxt_unit_buf_free(self.response_buffer);
                }
            } else {
                if let Err(err) = self.flush() {
                    // Prevent a double-panic, which causes an abort and hides
                    // details on the initial panic.
                    if !std::thread::panicking() {
                        panic!("Error while dropping ResponseWriter: {}", err);
                    }
                }
            }
        }
    }
}