ureq_proto/server/
sendres.rs

1use std::io::Write;
2
3use http::{HeaderName, HeaderValue};
4
5use crate::util::Writer;
6use crate::Error;
7
8use super::state::SendResponse;
9use super::{do_write_send_line, Reply, ResponsePhase, SendResponseResult};
10
11impl Reply<SendResponse> {
12    /// Write the response headers to the output buffer.
13    ///
14    /// Writes the response status line and headers to the output buffer.
15    /// May need to be called multiple times if the output buffer isn't large enough.
16    ///
17    /// Returns the number of bytes written to the output buffer.
18    pub fn write(&mut self, output: &mut [u8]) -> Result<usize, Error> {
19        // unwrap is ok because we are not here without providing it
20        let response = self.inner.response.as_ref().unwrap();
21
22        let mut w = Writer::new(output);
23        try_write_prelude(response, &mut self.inner.phase, &mut w)?;
24
25        let output_used = w.len();
26
27        Ok(output_used)
28    }
29
30    /// Whether the response headers have been fully written.
31    ///
32    /// Returns true if all response headers have been written and the state
33    /// is ready to proceed to sending the response body.
34    pub fn is_finished(&self) -> bool {
35        !self.inner.phase.is_prelude()
36    }
37
38    /// Proceed to sending a response body or cleanup.
39    ///
40    /// Transitions to either:
41    /// - SendBody state if the response needs a body (based on status code and method)
42    /// - Cleanup state if no response body should be sent (e.g., HEAD requests)
43    ///
44    /// This is only possible when the response headers are fully written.
45    ///
46    /// Panics if the response headers have not been fully written.
47    pub fn proceed(self) -> SendResponseResult {
48        assert!(self.is_finished());
49
50        if let Some(writer) = self.inner.state.writer {
51            // Enter SendBody for both chunked and sized bodies, even when size is 0,
52            // to ensure consistent state progression for explicit body headers.
53            //
54            // TODO(martin): Do we actually want this API wise? Isn't it better to go straight to Cleanup?
55            if writer.is_chunked() || writer.left_to_send().is_some() {
56                return SendResponseResult::SendBody(Reply::wrap(self.inner));
57            }
58        }
59
60        SendResponseResult::Cleanup(Reply::wrap(self.inner))
61    }
62}
63
64fn try_write_prelude(
65    response: &super::amended::AmendedResponse,
66    phase: &mut ResponsePhase,
67    w: &mut Writer,
68) -> Result<(), Error> {
69    let at_start = w.len();
70
71    loop {
72        if try_write_prelude_part(response, phase, w) {
73            continue;
74        }
75
76        let written = w.len() - at_start;
77
78        if written > 0 || phase.is_body() {
79            return Ok(());
80        } else {
81            return Err(Error::OutputOverflow);
82        }
83    }
84}
85
86fn try_write_prelude_part(
87    response: &super::amended::AmendedResponse,
88    phase: &mut ResponsePhase,
89    w: &mut Writer,
90) -> bool {
91    match phase {
92        ResponsePhase::Status => {
93            let success = do_write_send_line(response.prelude(), w, false);
94            if success {
95                *phase = ResponsePhase::Headers(0);
96            }
97            success
98        }
99
100        ResponsePhase::Headers(index) => {
101            let header_count = response.headers_len();
102            let all = response.headers();
103            let skipped = all.skip(*index);
104
105            do_write_headers(skipped, index, w);
106            if *index == header_count && w.try_write(|w| write!(w, "\r\n")) {
107                *phase = ResponsePhase::Body;
108            }
109            false
110        }
111
112        // We're past the header.
113        _ => false,
114    }
115}
116
117fn do_write_headers<'a, I>(headers: I, index: &mut usize, w: &mut Writer)
118where
119    I: Iterator<Item = (&'a HeaderName, &'a HeaderValue)>,
120{
121    for h in headers {
122        let success = w.try_write(|w| {
123            write!(w, "{}: ", h.0)?;
124            w.write_all(h.1.as_bytes())?;
125            write!(w, "\r\n")?;
126
127            Ok(())
128        });
129
130        if success {
131            *index += 1;
132        } else {
133            break;
134        }
135    }
136}
ureq_proto/server/sendres.rs

ureq_proto/server/
sendres.rs
