Skip to main content

base_http1/
output.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4use crate::{Error, Exchange, wire};
5use base_endpoint::stream::WritableStream;
6use base_http::{
7    InterimHandler, Response, ResponseMetadata, Status, Version,
8    body::{BodyFrame, ReadableBody},
9};
10use base_scheduler::Cancellation;
11use std::{
12    future::{Future, poll_fn},
13    io,
14    pin::pin,
15    task::Poll,
16};
17
18async fn wait<F: Future>(future: F, cancellation: &Cancellation) -> Result<F::Output, Error> {
19    if cancellation.is_never() {
20        return Ok(future.await);
21    }
22    let mut future = pin!(future);
23    let mut cancelled = pin!(cancellation.cancelled());
24    poll_fn(|cx| {
25        if cancelled.as_mut().poll(cx).is_ready() {
26            return Poll::Ready(Err(Error::Cancelled));
27        }
28        future.as_mut().poll(cx).map(Ok)
29    })
30    .await
31}
32
33// Bound copied output independently of the number or size of body chunks.
34const OUTPUT_BATCH_LIMIT: usize = 16 * 1024;
35
36async fn send_all<W: WritableStream>(
37    writer: &mut W,
38    mut bytes: &[u8],
39    cancellation: &Cancellation,
40) -> Result<(), Error> {
41    while !bytes.is_empty() {
42        let n = writer.write(bytes, cancellation).await?;
43        if n == 0 {
44            return Err(
45                io::Error::new(io::ErrorKind::WriteZero, "HTTP output made no progress").into(),
46            );
47        }
48        bytes = &bytes[n..];
49    }
50    Ok(())
51}
52
53/// The response lane. A partially sent or abandoned response permanently
54/// poisons admission. The owner must still await `close` to retire the backend.
55pub struct ResponseWriter<W> {
56    writer: W,
57    pending: Vec<u8>,
58    usable: bool,
59    continue_sent: bool,
60}
61
62impl<W: WritableStream> ResponseWriter<W> {
63    pub fn new(writer: W) -> Self {
64        Self {
65            writer,
66            pending: Vec::new(),
67            usable: true,
68            continue_sent: false,
69        }
70    }
71    /// Whether another message can be written. A failed write which leaves output usable wrote none of its message, so the owner can still send a different response.
72    pub fn is_usable(&self) -> bool {
73        self.usable
74    }
75    pub fn interim(&mut self, version: Version) -> InterimWriter<'_, W> {
76        InterimWriter {
77            output: self,
78            version,
79        }
80    }
81
82    /// Best-effort rejection for invalid input, followed by connection closure.
83    pub async fn reject(
84        &mut self,
85        status: Status,
86        cancellation: &Cancellation,
87    ) -> Result<(), Error> {
88        cancellation.check()?;
89        if !self.usable {
90            return Err(Error::Invalid("response output is not reusable"));
91        }
92        self.usable = false;
93        self.write_start_line(status, Version::Http11);
94        self.buffer(b"Content-Length: 0\r\nConnection: close\r\n\r\n");
95        self.flush(cancellation).await
96    }
97
98    pub(crate) async fn write_all(
99        &mut self,
100        bytes: &[u8],
101        cancellation: &Cancellation,
102    ) -> Result<(), Error> {
103        cancellation.check()?;
104        if bytes.len() > OUTPUT_BATCH_LIMIT - self.pending.len() {
105            self.send_pending(cancellation).await?;
106        }
107        if bytes.len() >= OUTPUT_BATCH_LIMIT {
108            // Large chunks retain their existing owner and are not copied.
109            return send_all(&mut self.writer, bytes, cancellation).await;
110        }
111        self.buffer(bytes);
112        Ok(())
113    }
114
115    fn reserve(&mut self, additional: usize) {
116        let required = self.pending.len() + additional;
117        if required > self.pending.capacity() {
118            let capacity = required
119                .max(self.pending.capacity().saturating_mul(2))
120                .clamp(256, OUTPUT_BATCH_LIMIT);
121            self.pending.reserve_exact(capacity - self.pending.len());
122        }
123    }
124
125    async fn send_pending(&mut self, cancellation: &Cancellation) -> Result<(), Error> {
126        // The lane is already poisoned while sending. On error/cancellation the
127        // retained prefix must never be retried by a subsequent message or close.
128        send_all(&mut self.writer, &self.pending, cancellation).await?;
129        self.pending.clear();
130        Ok(())
131    }
132
133    fn buffer(&mut self, bytes: &[u8]) {
134        self.reserve(bytes.len());
135        self.pending.extend_from_slice(bytes);
136    }
137
138    // Every message starts with an empty batch, since the previous message or interim head was flushed, so its start line always fits.
139    fn write_start_line(&mut self, status: Status, version: Version) {
140        if status == Status::OK && version == Version::Http11 {
141            self.buffer(b"HTTP/1.1 200 OK\r\n");
142            return;
143        }
144        let version = match version {
145            Version::Http10 => b"HTTP/1.0 ",
146            Version::Http11 => b"HTTP/1.1 ",
147            _ => unreachable!("HTTP/1 exchanges and interim writers are HTTP/1.0 or HTTP/1.1"),
148        };
149        let status_code = status.as_u16();
150        let code = [
151            b'0' + (status_code / 100) as u8,
152            b'0' + (status_code / 10 % 10) as u8,
153            b'0' + (status_code % 10) as u8,
154            b' ',
155        ];
156        let reason = status.description().unwrap_or("").as_bytes();
157        self.reserve(version.len() + code.len() + reason.len() + 2);
158        self.pending.extend_from_slice(version);
159        self.pending.extend_from_slice(&code);
160        self.pending.extend_from_slice(reason);
161        self.pending.extend_from_slice(b"\r\n");
162    }
163
164    async fn write_fields(
165        &mut self,
166        fields: &[base_http::Field],
167        cancellation: &Cancellation,
168    ) -> Result<(), Error> {
169        for field in fields {
170            self.write_field(field.name(), field.value(), cancellation)
171                .await?;
172        }
173        Ok(())
174    }
175
176    async fn write_field(
177        &mut self,
178        name: &str,
179        value: &[u8],
180        cancellation: &Cancellation,
181    ) -> Result<(), Error> {
182        let size = name.len().saturating_add(value.len()).saturating_add(4);
183        if size <= OUTPUT_BATCH_LIMIT - self.pending.len() {
184            self.reserve(size);
185            self.pending.extend_from_slice(name.as_bytes());
186            self.pending.extend_from_slice(b": ");
187            self.pending.extend_from_slice(value);
188            self.pending.extend_from_slice(b"\r\n");
189            return Ok(());
190        }
191        self.write_all(name.as_bytes(), cancellation).await?;
192        self.write_all(b": ", cancellation).await?;
193        self.write_all(value, cancellation).await?;
194        self.write_all(b"\r\n", cancellation).await
195    }
196
197    async fn write_plan(
198        &mut self,
199        plan: &Plan<'_>,
200        cancellation: &Cancellation,
201    ) -> Result<(), Error> {
202        self.write_start_line(plan.response.status, plan.version);
203        for field in plan.response.fields.pairs() {
204            if (plan.upgrade.is_some() && field.name().eq_ignore_ascii_case("upgrade"))
205                || (matches!(plan.framing, Framing::Chunked(_))
206                    && field.name().eq_ignore_ascii_case("content-length"))
207            {
208                continue;
209            }
210            self.write_field(field.name(), field.value(), cancellation)
211                .await?;
212        }
213        if let Some(protocol) = plan.upgrade {
214            self.write_field("Upgrade", protocol, cancellation).await?;
215        }
216        if plan.connection_upgrade {
217            self.write_field("Connection", b"upgrade", cancellation)
218                .await?;
219        }
220        if let Some(length) = plan.content_length {
221            let mut digits = [0; 20];
222            let bytes = decimal(length, &mut digits);
223            self.write_field("Content-Length", bytes, cancellation)
224                .await?;
225        }
226        if plan.chunked {
227            self.write_field("Transfer-Encoding", b"chunked", cancellation)
228                .await?;
229        }
230        if let Some(option) = plan.connection {
231            self.write_field("Connection", option, cancellation).await?;
232        }
233        self.write_all(b"\r\n", cancellation).await
234    }
235
236    /// Write one final response. `true` means output permits persistence; the
237    /// connection owner must also finish input and join concurrent work before
238    /// admitting another request. Discarded input disables persistence.
239    ///
240    /// A response refused before any of it is written, such as one with an informational status, conflicting framing or a `Content-Length` field its body contradicts, leaves output usable.
241    ///
242    /// A response to an HTTP/1.0 request has no chunked framing, so any trailers its body produces are dropped.
243    pub async fn write_response<B: ReadableBody>(
244        &mut self,
245        exchange: &Exchange,
246        response: &mut Response<B>,
247        cancellation: &Cancellation,
248    ) -> Result<bool, Error> {
249        let result = self
250            .write_response_inner(exchange, response, cancellation)
251            .await;
252        response.body.close(None);
253        result
254    }
255
256    async fn write_response_inner<B: ReadableBody>(
257        &mut self,
258        exchange: &Exchange,
259        response: &mut Response<B>,
260        cancellation: &Cancellation,
261    ) -> Result<bool, Error> {
262        cancellation.check()?;
263        if !self.usable {
264            return Err(Error::Invalid("response output is not reusable"));
265        }
266        let (framing, persistent, tunnel) = {
267            let plan = plan(
268                exchange,
269                &response.metadata,
270                response.body.length(),
271                response.body.may_have_trailers(),
272            )?;
273            if plan.tunnel
274                && response.metadata.status == Status::SWITCHING_PROTOCOLS
275                && !self.continue_sent
276                && exchange.expects_continue()
277            {
278                self.interim(Version::Http11)
279                    .on_interim_response(ResponseMetadata::new(Status::CONTINUE), cancellation)
280                    .await?;
281            }
282            self.usable = false;
283            self.write_plan(&plan, cancellation).await?;
284            if plan.tunnel {
285                // Commit the handshake before the first potentially blocking input read.
286                self.flush(cancellation).await?;
287                exchange.switch();
288            }
289            (plan.framing, plan.persistent, plan.tunnel)
290        };
291        let encoding = BodyEncoder::new(
292            framing,
293            !tunnel && response.metadata.status == Status::RESET_CONTENT,
294        )
295        .discard_trailers(!tunnel && exchange.version() == Version::Http10);
296        self.write_body(&mut response.body, encoding, cancellation)
297            .await?;
298        self.usable = persistent;
299        self.continue_sent = false;
300        Ok(persistent)
301    }
302
303    pub(crate) fn begin(&mut self) -> Result<(), Error> {
304        if !self.usable {
305            return Err(Error::Invalid("output is not reusable"));
306        }
307        self.usable = false;
308        Ok(())
309    }
310
311    pub(crate) fn completed(&mut self, persistent: bool) {
312        self.usable = persistent;
313    }
314
315    pub(crate) async fn flush(&mut self, cancellation: &Cancellation) -> Result<(), Error> {
316        self.send_pending(cancellation).await?;
317        Ok(self.writer.flush(cancellation).await?)
318    }
319
320    pub(crate) async fn write_body<B: ReadableBody>(
321        &mut self,
322        body: &mut B,
323        mut encoding: BodyEncoder,
324        cancellation: &Cancellation,
325    ) -> Result<(), Error> {
326        if !matches!(encoding.framing, Framing::Absent) {
327            loop {
328                if !body.is_ready() {
329                    self.flush(cancellation).await?;
330                }
331                let frame = {
332                    let reading = pin!(body.read(cancellation));
333                    let mut read = pin!(wait(reading, cancellation));
334                    let ready = poll_fn(|cx| Poll::Ready(read.as_mut().poll(cx))).await;
335                    match ready {
336                        Poll::Ready(result) => result,
337                        Poll::Pending => {
338                            self.flush(cancellation).await?;
339                            read.await
340                        }
341                    }
342                }?
343                .map_err(|e| {
344                    if cancellation.is_cancelled() {
345                        Error::Cancelled
346                    } else {
347                        Error::Body(Box::new(e))
348                    }
349                })?;
350                let end = frame.is_end();
351                encoding.write_frame(self, frame, cancellation).await?;
352                if end {
353                    break;
354                }
355            }
356        }
357        self.flush(cancellation).await?;
358        Ok(())
359    }
360
361    pub async fn close(&mut self) -> io::Result<()> {
362        self.usable = false;
363        self.pending.clear();
364        self.writer.close().await
365    }
366}
367
368/// Statically dispatched interim-response hook borrowing the output lane.
369pub struct InterimWriter<'a, W> {
370    output: &'a mut ResponseWriter<W>,
371    version: Version,
372}
373
374impl<W: WritableStream> InterimHandler for InterimWriter<'_, W> {
375    type Error = Error;
376    async fn on_interim_response(
377        &mut self,
378        response: ResponseMetadata,
379        cancellation: &Cancellation,
380    ) -> Result<(), Error> {
381        cancellation.check()?;
382        if !self.output.usable {
383            return Err(Error::Invalid("response output is not reusable"));
384        }
385        if !response.status.is_informational() || response.status == Status::SWITCHING_PROTOCOLS {
386            return Err(Error::Invalid("invalid interim status"));
387        }
388        if self.version != Version::Http11 {
389            return Err(Error::Invalid("HTTP/1.0 interim response"));
390        }
391        if response.fields.contains("content-length")
392            || response.fields.contains("transfer-encoding")
393            || response.content_length.is_some()
394            || response.protocol.is_some()
395        {
396            return Err(Error::Invalid("invalid interim framing"));
397        }
398        wire::output_fields(response.fields.pairs())?;
399        self.output.usable = false;
400        self.output.write_start_line(response.status, self.version);
401        self.output
402            .write_fields(response.fields.pairs(), cancellation)
403            .await?;
404        self.output.write_all(b"\r\n", cancellation).await?;
405        self.output.flush(cancellation).await?;
406        if response.status == Status::CONTINUE {
407            self.output.continue_sent = true;
408        }
409        self.output.usable = true;
410        Ok(())
411    }
412}
413
414#[derive(Clone, Copy)]
415pub(crate) enum Framing {
416    Absent,
417    Fixed(u64),
418    Chunked(Option<u64>),
419    Raw,
420}
421/// Incremental body framing shared by pull and imperative output. The owner
422/// poisons its lane before entering this state and restores reuse only at EOF.
423pub(crate) struct BodyEncoder {
424    framing: Framing,
425    total: u64,
426    trailers: Option<base_http::Fields>,
427    forbid_content: bool,
428    discard_trailers: bool,
429}
430impl BodyEncoder {
431    pub(crate) fn end_requires_output(&self) -> bool {
432        matches!(self.framing, Framing::Chunked(_))
433    }
434    pub(crate) fn new(framing: Framing, forbid_content: bool) -> Self {
435        Self {
436            framing,
437            total: 0,
438            trailers: None,
439            forbid_content,
440            discard_trailers: false,
441        }
442    }
443    /// Accept trailers that the framing cannot carry and drop them. HTTP/1.0 has no chunked framing, and HTTP semantics permit a recipient to discard trailers, so its fixed-length and close-delimited bodies validate and then omit them.
444    pub(crate) fn discard_trailers(mut self, discard: bool) -> Self {
445        self.discard_trailers = discard;
446        self
447    }
448    pub(crate) async fn write_frame<W: WritableStream, D: AsRef<[u8]>>(
449        &mut self,
450        output: &mut ResponseWriter<W>,
451        frame: BodyFrame<D>,
452        cancellation: &Cancellation,
453    ) -> Result<(), Error> {
454        let end = frame.is_end();
455        match frame {
456            BodyFrame::Data { data, .. } => {
457                if self.trailers.is_some() {
458                    return Err(Error::Invalid("data after trailing headers"));
459                }
460                let bytes = data.as_ref();
461                // No transport carries `u64::MAX` bytes, so saturation only guards the arithmetic.
462                self.total = self.total.saturating_add(bytes.len() as u64);
463                if self.forbid_content && self.total != 0 {
464                    return Err(Error::Invalid("content on 205"));
465                }
466                if let Framing::Fixed(length) | Framing::Chunked(Some(length)) = self.framing {
467                    if self.total > length {
468                        return Err(Error::LengthMismatch);
469                    }
470                }
471                if !bytes.is_empty() {
472                    if matches!(self.framing, Framing::Chunked(_)) {
473                        output
474                            .write_all(format!("{:x}\r\n", bytes.len()).as_bytes(), cancellation)
475                            .await?;
476                        output.write_all(bytes, cancellation).await?;
477                        output.write_all(b"\r\n", cancellation).await?;
478                    } else {
479                        output.write_all(bytes, cancellation).await?;
480                    }
481                }
482            }
483            BodyFrame::Headers { fields, .. } => {
484                if self.trailers.is_some() {
485                    return Err(Error::Invalid("repeated trailing headers"));
486                }
487                if !matches!(self.framing, Framing::Chunked(_)) && !self.discard_trailers {
488                    return Err(Error::Invalid("trailers require chunked framing"));
489                }
490                validate_trailer_fields(fields.pairs())?;
491                wire::output_fields(fields.pairs())?;
492                self.trailers = Some(fields);
493            }
494            BodyFrame::End => {}
495        }
496        if end {
497            if let Framing::Fixed(length) | Framing::Chunked(Some(length)) = self.framing {
498                if self.total != length {
499                    return Err(Error::LengthMismatch);
500                }
501            }
502            if matches!(self.framing, Framing::Chunked(_)) {
503                output.write_all(b"0\r\n", cancellation).await?;
504                if let Some(trailers) = &self.trailers {
505                    output.write_fields(trailers.pairs(), cancellation).await?;
506                }
507                output.write_all(b"\r\n", cancellation).await?;
508            }
509        }
510        Ok(())
511    }
512}
513
514// Borrow semantic fields and describe only the framing additions. This avoids
515// constructing/cloning a second semantic message just to serialize a wire head.
516struct Plan<'a> {
517    response: &'a ResponseMetadata,
518    version: Version,
519    framing: Framing,
520    persistent: bool,
521    tunnel: bool,
522    content_length: Option<u64>,
523    chunked: bool,
524    upgrade: Option<&'a [u8]>,
525    connection_upgrade: bool,
526    connection: Option<&'static [u8]>,
527}
528
529fn plan<'a>(
530    exchange: &Exchange,
531    response: &'a ResponseMetadata,
532    length: Option<u64>,
533    may_have_trailers: bool,
534) -> Result<Plan<'a>, Error> {
535    if !response.status.is_final() {
536        return Err(Error::Invalid("final response required"));
537    }
538    let fields = &response.fields;
539    wire::output_fields(fields.pairs())?;
540    let version = exchange.version();
541    if response.version.is_some_and(|v| v != version) {
542        return Err(Error::Invalid("response version differs from request"));
543    }
544    // Derive framing and connection facts together from the initial section.
545    let facts = wire::response_framing(fields)?;
546    let connect = exchange.is_connect() && response.status.is_success();
547    let upgrade = response.status == Status::SWITCHING_PROTOCOLS;
548    let tunnel = connect || upgrade;
549    let status_bodyless =
550        response.status == Status::NO_CONTENT || response.status == Status::NOT_MODIFIED;
551    let absent = !tunnel && (exchange.is_head() || status_bodyless);
552    // Raw controls are written deliberately, so they stay strict: no declared length may accompany a raw Transfer-Encoding, and the body must agree with a raw Content-Length.
553    if facts.chunked && (response.content_length.is_some() || facts.length.is_some()) {
554        return Err(Error::Invalid("content-length with transfer-encoding"));
555    }
556    // A body of known length describes the data a response sends, so a typed declaration it contradicts is ignored; a response without data or framing is described by its declaration.
557    let typed = if tunnel || absent {
558        response.content_length
559    } else {
560        response.declared_length(length)
561    };
562    let declared = wire::declared_length(typed, facts.length)?;
563    let chunked = facts.chunked;
564    let mut persistent = exchange.is_persistent()
565        && !exchange.is_abandoned()
566        && response.persistent != Some(false)
567        && !facts.close
568        && !tunnel;
569    let mut content_length = typed.filter(|_| facts.length.is_none());
570    let mut generated_chunked = false;
571    let mut selected_upgrade = None;
572    let mut connection_upgrade = false;
573    let mut connection = None;
574    let framing;
575    if tunnel {
576        if !exchange.is_complete() {
577            return Err(Error::Invalid("switch with unread request body"));
578        }
579        if declared.is_some() || chunked {
580            return Err(Error::Invalid("framing fields on tunnel response"));
581        }
582        if upgrade {
583            if version != Version::Http11 || !exchange.has_upgrade() {
584                return Err(Error::Invalid("unsolicited upgrade"));
585            }
586            if facts.duplicate_upgrade {
587                return Err(Error::Invalid("duplicate singleton field"));
588            }
589            let selected = response
590                .protocol
591                .as_ref()
592                .map(|p| p.as_str().as_bytes())
593                .or(facts.upgrade)
594                .ok_or(Error::Invalid("missing selected protocol"))?;
595            if !wire::token(selected) || !exchange.accepts_upgrade(selected) {
596                return Err(Error::Invalid("upgrade protocol was not offered"));
597            }
598            selected_upgrade = Some(selected);
599            connection_upgrade = !facts.connection_upgrade;
600        } else if response.protocol.is_some() {
601            return Err(Error::Invalid("protocol on traditional CONNECT"));
602        }
603        framing = Framing::Raw;
604    } else {
605        if response.protocol.is_some() {
606            return Err(Error::Invalid("protocol without upgrade status"));
607        }
608        if facts.upgrade.is_some() {
609            connection_upgrade = !facts.connection_upgrade;
610        }
611        if chunked && version != Version::Http11 {
612            return Err(Error::Invalid("HTTP/1.0 transfer-encoding"));
613        }
614        if absent {
615            if chunked {
616                return Err(Error::Invalid(
617                    "trailers or transfer-encoding on bodyless response",
618                ));
619            }
620            if response.status == Status::NO_CONTENT {
621                if declared.is_some() {
622                    return Err(Error::Invalid("content-length on 204"));
623                }
624            } else if declared.is_none() {
625                content_length = length;
626            }
627            framing = Framing::Absent;
628        } else {
629            if let (Some(actual), Some(declared)) = (length, declared) {
630                if actual != declared {
631                    return Err(Error::LengthMismatch);
632                }
633            }
634            if response.status == Status::RESET_CONTENT
635                && (length.is_some_and(|n| n != 0) || declared.is_some_and(|n| n != 0))
636            {
637                return Err(Error::Invalid("content on 205"));
638            }
639            // Only HTTP/1.1 can frame trailers; HTTP/1.0 drops them. A Trailer field is only a hint, so the body alone decides.
640            let trailers = may_have_trailers && version == Version::Http11;
641            if response.status == Status::RESET_CONTENT
642                && declared.is_none()
643                && !chunked
644                && !trailers
645            {
646                content_length = Some(0);
647                framing = Framing::Fixed(0);
648            } else if chunked || trailers {
649                content_length = None;
650                generated_chunked = !chunked;
651                framing = Framing::Chunked(declared.or(length));
652            } else if let Some(length) = declared.or(length) {
653                if declared.is_none() {
654                    content_length = Some(length);
655                }
656                framing = Framing::Fixed(length);
657            } else if version == Version::Http11 {
658                generated_chunked = true;
659                framing = Framing::Chunked(declared.or(length));
660            } else {
661                persistent = false;
662                framing = Framing::Raw;
663            }
664        }
665        if !persistent {
666            if !facts.close {
667                connection = Some(&b"close"[..]);
668            }
669        } else if version == Version::Http10 && !facts.keep_alive {
670            connection = Some(&b"keep-alive"[..]);
671        }
672    }
673    Ok(Plan {
674        response,
675        version,
676        framing,
677        persistent,
678        tunnel,
679        content_length,
680        chunked: generated_chunked,
681        upgrade: selected_upgrade,
682        connection_upgrade,
683        connection,
684    })
685}
686
687fn decimal(mut value: u64, digits: &mut [u8; 20]) -> &[u8] {
688    let mut start = digits.len();
689    loop {
690        start -= 1;
691        digits[start] = b'0' + (value % 10) as u8;
692        value /= 10;
693        if value == 0 {
694            return &digits[start..];
695        }
696    }
697}
698
699pub(crate) fn validate_trailer_fields(fields: &[base_http::Field]) -> Result<(), Error> {
700    let policy = base_http::headers::Policy::default();
701    for field in fields {
702        if !policy.get(field.name()).trailer
703            || ["keep-alive", "proxy-connection"]
704                .iter()
705                .any(|name| field.name().eq_ignore_ascii_case(name))
706        {
707            return Err(base_http::FieldError::InvalidTrailer(field.name().to_owned()).into());
708        }
709    }
710    Ok(())
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use crate::{Options, RequestReader};
717    use base_endpoint::stream::ReadableStream;
718    use base_http::{Fields, Protocol, body::BufferedBody};
719    use base_scheduler::OperationResult;
720    use futures_lite::future::block_on;
721
722    struct Bytes(&'static [u8]);
723    impl ReadableStream for Bytes {
724        async fn read_into(
725            &mut self,
726            buffer: &mut [u8],
727            _: &Cancellation,
728        ) -> OperationResult<usize> {
729            self.try_read_into(buffer).map_err(Into::into)
730        }
731        fn try_read_into(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
732            let n = self.0.len().min(buffer.len());
733            buffer[..n].copy_from_slice(&self.0[..n]);
734            self.0 = &self.0[n..];
735            Ok(n)
736        }
737        async fn close(&mut self) -> io::Result<()> {
738            Ok(())
739        }
740    }
741
742    // Accepts no bytes, so the first transport write fails and reveals where a batch boundary fell.
743    struct Refusing;
744    impl WritableStream for Refusing {
745        async fn write(&mut self, _: &[u8], _: &Cancellation) -> OperationResult<usize> {
746            Err(io::Error::from(io::ErrorKind::BrokenPipe).into())
747        }
748        async fn flush(&mut self, _: &Cancellation) -> OperationResult<()> {
749            Ok(())
750        }
751        async fn close(&mut self) -> io::Result<()> {
752            Ok(())
753        }
754    }
755
756    fn exchange(bytes: &'static [u8]) -> Exchange {
757        let mut input = RequestReader::new(Bytes(bytes), Options::default()).unwrap();
758        let exchange = block_on(input.read_request(&Cancellation::never()))
759            .unwrap()
760            .unwrap()
761            .1;
762        block_on(input.close()).unwrap();
763        exchange
764    }
765
766    // A field occupying exactly `size` bytes of the batch.
767    fn filler(size: usize) -> (&'static str, Vec<u8>) {
768        ("x", vec![b'v'; size - 5])
769    }
770
771    fn is_broken_pipe(result: Result<impl Sized, Error>) -> bool {
772        matches!(result, Err(Error::Io(error)) if error.kind() == io::ErrorKind::BrokenPipe)
773    }
774
775    #[test]
776    fn a_head_field_crossing_the_batch_limit_sends_the_batch_first() {
777        const GET: &[u8] = b"GET / HTTP/1.1\r\nHost: x\r\n\r\n";
778        const UPGRADE: &[u8] =
779            b"GET / HTTP/1.1\r\nHost: x\r\nConnection: upgrade\r\nUpgrade: echo\r\n\r\n";
780        const OK: usize = b"HTTP/1.1 200 OK\r\n".len();
781        const SWITCHING: usize = b"HTTP/1.1 101 Switching Protocols\r\n".len();
782        // Each response leaves one free byte in the batch before the generated field.
783        let full = |start: usize| filler(OUTPUT_BATCH_LIMIT - 1 - start);
784        type Case = (&'static [u8], Response<BufferedBody<Vec<u8>>>);
785        let mut cases: Vec<Case> = Vec::new();
786
787        let mut upgrade = Response::new(Status::SWITCHING_PROTOCOLS, BufferedBody::default());
788        upgrade.protocol = Some(Protocol::new("echo").unwrap());
789        let (name, value) = full(SWITCHING);
790        upgrade.fields.add(name, value).unwrap();
791        cases.push((UPGRADE, upgrade));
792
793        let mut advertised = Response::new(Status::OK, BufferedBody::default());
794        advertised.fields.add("Upgrade", "h2c").unwrap();
795        let (name, value) = full(OK + "Upgrade: h2c\r\n".len());
796        advertised.fields.add(name, value).unwrap();
797        cases.push((GET, advertised));
798
799        let mut length = Response::new(Status::OK, BufferedBody::from(b"hi".to_vec()));
800        let (name, value) = full(OK);
801        length.fields.add(name, value).unwrap();
802        cases.push((GET, length));
803
804        let mut close = Response::new(Status::OK, BufferedBody::default());
805        close.persistent = Some(false);
806        close.fields.add("Content-Length", "0").unwrap();
807        let (name, value) = full(OK + "Content-Length: 0\r\n".len());
808        close.fields.add(name, value).unwrap();
809        cases.push((GET, close));
810
811        // A large value, a name filling the batch, and a name leaving no room for its separator.
812        let mut value = Response::new(Status::OK, BufferedBody::default());
813        value
814            .fields
815            .add("x", vec![b'v'; OUTPUT_BATCH_LIMIT + 1])
816            .unwrap();
817        cases.push((GET, value));
818        let mut name = Response::new(Status::OK, BufferedBody::default());
819        name.fields
820            .add("x".repeat(OUTPUT_BATCH_LIMIT), "v")
821            .unwrap();
822        cases.push((GET, name));
823        let mut separator = Response::new(Status::OK, BufferedBody::default());
824        separator
825            .fields
826            .add("x".repeat(OUTPUT_BATCH_LIMIT - 1 - OK), "v")
827            .unwrap();
828        cases.push((GET, separator));
829
830        for (index, (request, mut response)) in cases.into_iter().enumerate() {
831            let mut output = ResponseWriter::new(Refusing);
832            let result = block_on(output.write_response(
833                &exchange(request),
834                &mut response,
835                &Cancellation::never(),
836            ));
837            assert!(is_broken_pipe(result), "case {index}");
838            assert!(!output.is_usable());
839        }
840
841        // A body that may have trailers generates chunked framing.
842        let (body, _writer) = base_http::body::BodyReader::<Vec<u8>>::bounded(1, 1);
843        let mut chunked = Response::new(Status::OK, body);
844        let (name, value) = full(OK);
845        chunked.fields.add(name, value).unwrap();
846        let mut output = ResponseWriter::new(Refusing);
847        let result =
848            block_on(output.write_response(&exchange(GET), &mut chunked, &Cancellation::never()));
849        assert!(is_broken_pipe(result));
850        assert!(!output.is_usable());
851    }
852
853    #[test]
854    fn interim_heads_crossing_the_batch_limit_send_the_batch_first() {
855        const EARLY_HINTS: usize = b"HTTP/1.1 103 Early Hints\r\n".len();
856        let large = ("link", vec![b'v'; OUTPUT_BATCH_LIMIT + 1]);
857        let full = filler(OUTPUT_BATCH_LIMIT - 1 - EARLY_HINTS);
858        for (name, value) in [large, full] {
859            let mut output = ResponseWriter::new(Refusing);
860            let mut response = ResponseMetadata::new(Status::new(103).unwrap());
861            response.fields.add(name, value).unwrap();
862            let result = block_on(
863                output
864                    .interim(Version::Http11)
865                    .on_interim_response(response, &Cancellation::never()),
866            );
867            assert!(is_broken_pipe(result));
868            assert!(!output.is_usable());
869        }
870    }
871
872    #[test]
873    fn body_frames_crossing_the_batch_limit_send_the_batch_first() {
874        let data = |bytes: &[u8]| BodyFrame::Data {
875            data: bytes.to_vec(),
876            end: false,
877        };
878        let mut trailers = Fields::new();
879        trailers
880            .add("x-digest", vec![b'v'; OUTPUT_BATCH_LIMIT + 1])
881            .unwrap();
882        let large = vec![b'v'; OUTPUT_BATCH_LIMIT];
883        // Each case batches `pending` bytes from earlier frames before writing its frame.
884        let cases = [
885            (Framing::Chunked(None), OUTPUT_BATCH_LIMIT - 2, data(b"abc")),
886            (Framing::Chunked(None), 1, data(&large)),
887            (Framing::Chunked(None), OUTPUT_BATCH_LIMIT - 7, data(b"abc")),
888            (Framing::Fixed(OUTPUT_BATCH_LIMIT as u64), 1, data(&large)),
889            (
890                Framing::Chunked(None),
891                OUTPUT_BATCH_LIMIT - 2,
892                BodyFrame::End,
893            ),
894            (
895                Framing::Chunked(None),
896                OUTPUT_BATCH_LIMIT - 4,
897                BodyFrame::End,
898            ),
899            (
900                Framing::Chunked(None),
901                1,
902                BodyFrame::Headers {
903                    fields: trailers,
904                    end: true,
905                },
906            ),
907        ];
908        for (index, (framing, pending, frame)) in cases.into_iter().enumerate() {
909            let never = Cancellation::never();
910            let mut output = ResponseWriter::new(Refusing);
911            block_on(output.write_all(&vec![b'.'; pending], &never)).unwrap();
912            let mut encoder = BodyEncoder::new(framing, false);
913            let result = block_on(encoder.write_frame(&mut output, frame, &never));
914            assert!(is_broken_pipe(result), "case {index}");
915        }
916    }
917
918    #[test]
919    fn empty_fixed_data_needs_no_transport_write() {
920        let never = Cancellation::never();
921        let mut output = ResponseWriter::new(Refusing);
922        let mut encoder = BodyEncoder::new(Framing::Fixed(0), false);
923        let frame = BodyFrame::Data {
924            data: Vec::new(),
925            end: true,
926        };
927        block_on(encoder.write_frame(&mut output, frame, &never)).unwrap();
928        block_on(output.flush(&never)).unwrap();
929        block_on(output.close()).unwrap();
930    }
931}