trillium_http/
upgrade.rs

1use crate::{
2    Buffer, Conn, Headers, HttpContext, KnownHeaderName, Method, PeerGone, ProtocolSession,
3    ReceivedBody, Status, TypeSet, Version,
4    conn::ConnInner,
5    h2::H2Connection,
6    h3::{Frame, H3Connection},
7    headers::qpack::{FieldSection, PseudoHeaders},
8    received_body::{H3TrailerFuture, ReceivedBodyState, write_chunk},
9    util::encoding,
10};
11use encoding_rs::Encoding;
12use fieldwork::Fieldwork;
13use futures_lite::{
14    AsyncWriteExt,
15    io::{AsyncRead, AsyncWrite},
16};
17use std::{
18    borrow::Cow,
19    fmt::{self, Debug, Formatter},
20    io::{self, IoSlice, Write},
21    net::IpAddr,
22    pin::Pin,
23    str,
24    sync::Arc,
25    task::{Context, Poll, ready},
26    time::Instant,
27};
28
29/// Per-protocol outbound framing state for an [`Upgrade`], chosen at the upgrade
30/// transition.
31#[derive(Debug)]
32pub(crate) enum WriteState {
33    /// No framing on the `AsyncWrite` path. HTTP/1.1 without chunked encoding (raw
34    /// passthrough) and HTTP/2 (framed at the connection layer).
35    Raw,
36    /// HTTP/1.1 chunked transfer-encoding.
37    H1Chunked(H1ChunkedState),
38    /// HTTP/3 DATA-frame encoding.
39    H3Framed(H3FramedState),
40}
41
42#[derive(Debug, Default)]
43pub(crate) struct H1ChunkedState {
44    pub(crate) pending: Vec<u8>,
45    pub(crate) terminator_written: bool,
46}
47
48#[derive(Debug, Default)]
49pub(crate) struct H3FramedState {
50    pub(crate) pending: Vec<u8>,
51    pub(crate) terminator_written: bool,
52}
53
54/// Pick outbound framing from http version and the outbound headers' `Transfer-Encoding`.
55/// h3 is always DATA-framed; h1 chunks only when the headers request it; h2 is framed by
56/// the connection driver, so the `AsyncWrite` path stays raw.
57fn compute_write_state(version: Version, outbound_headers: &Headers) -> WriteState {
58    match version {
59        Version::Http1_0 | Version::Http1_1 if has_chunked_encoding(outbound_headers) => {
60            WriteState::H1Chunked(H1ChunkedState::default())
61        }
62        Version::Http3 => WriteState::H3Framed(H3FramedState::default()),
63        _ => WriteState::Raw,
64    }
65}
66
67/// True if `Transfer-Encoding` includes `chunked`. Tolerant of multi-codings like
68/// `gzip, chunked`; no ordering enforcement.
69fn has_chunked_encoding(headers: &Headers) -> bool {
70    headers
71        .token_iter(KnownHeaderName::TransferEncoding)
72        .any(|coding| coding.eq_ignore_ascii_case("chunked"))
73}
74
75/// Parse the inbound `Content-Length`. `None` for chunked, missing, or malformed.
76fn parse_content_length(inbound_headers: &Headers) -> Option<u64> {
77    if inbound_headers.has_header(KnownHeaderName::TransferEncoding) {
78        return None;
79    }
80    inbound_headers.content_length()
81}
82
83/// Drain `pending` to `transport`, returning `Pending` if the transport blocks.
84fn poll_drain_pending<T: AsyncWrite + Unpin>(
85    pending: &mut Vec<u8>,
86    cx: &mut Context<'_>,
87    transport: &mut T,
88) -> Poll<io::Result<()>> {
89    while !pending.is_empty() {
90        match Pin::new(&mut *transport).poll_write(cx, pending) {
91            Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
92            Poll::Ready(Ok(n)) => {
93                pending.drain(..n);
94            }
95            Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
96            Poll::Pending => return Poll::Pending,
97        }
98    }
99    Poll::Ready(Ok(()))
100}
101
102/// Drain `pending` until the transport blocks or `pending` is empty, without yielding
103/// `Pending`. The next call resumes the drain.
104fn best_effort_drain<T: AsyncWrite + Unpin>(
105    pending: &mut Vec<u8>,
106    cx: &mut Context<'_>,
107    transport: &mut T,
108) -> io::Result<()> {
109    while !pending.is_empty() {
110        match Pin::new(&mut *transport).poll_write(cx, pending) {
111            Poll::Ready(Ok(0)) => return Err(io::ErrorKind::WriteZero.into()),
112            Poll::Ready(Ok(n)) => {
113                pending.drain(..n);
114            }
115            Poll::Ready(Err(e)) => return Err(e),
116            Poll::Pending => break,
117        }
118    }
119    Ok(())
120}
121
122/// Append an HTTP/3 DATA frame header for `payload_len` bytes to `out`. Caller appends
123/// the payload immediately after.
124fn encode_h3_data_header(out: &mut Vec<u8>, payload_len: u64) {
125    let frame = Frame::Data(payload_len);
126    let header_len = frame.encoded_len();
127    let start = out.len();
128    out.resize(start + header_len, 0);
129    frame.encode(&mut out[start..]);
130}
131
132/// An HTTP upgrade — owns the underlying transport along with all the data from the
133/// originating [`Conn`].
134///
135/// **Reading the transport directly**: drain `buffer` first if it has bytes in it. Reading
136/// via the [`AsyncRead`] impl on `Upgrade` handles this automatically.
137#[derive(Fieldwork)]
138#[fieldwork(get, get_mut, set, with, take, into_field, rename_predicates)]
139pub struct Upgrade<Transport> {
140    /// Everything but the transport, boxed so that an `Upgrade`, and anything that might hold
141    /// one, is cheap to move.
142    #[field(delegate, skip)]
143    pub(crate) inner: Box<UpgradeInner>,
144
145    /// The underlying io (often a `TcpStream` or similar)
146    pub(crate) transport: Transport,
147}
148
149/// Everything in an [`Upgrade`] but its transport.
150#[derive(Fieldwork)]
151#[fieldwork(
152    get,
153    get_mut,
154    set,
155    with,
156    take,
157    into_field,
158    rename_predicates,
159    delegatable
160)]
161pub(crate) struct UpgradeInner {
162    /// The http headers the peer sent to us
163    #[field(deprecate(was = "request_headers", since = "1.3.0"))]
164    pub(crate) received_headers: Headers,
165
166    /// The http headers as set before the upgrade was negotiated and sent
167    /// to the peer.
168    #[field(deprecate(was = "response_headers", since = "1.3.0"))]
169    pub(crate) sent_headers: Headers,
170
171    /// The request path
172    #[field(get = false)]
173    pub(crate) path: Cow<'static, str>,
174
175    /// The http request method
176    #[field(copy)]
177    pub(crate) method: Method,
178
179    /// Any state that has been accumulated on the Conn before negotiating the upgrade
180    pub(crate) state: TypeSet,
181
182    /// Any bytes that have been read from the underlying transport already.
183    ///
184    /// It is your responsibility to process these bytes before reading directly from the
185    /// transport.
186    #[field(
187        deref = "[u8]",
188        into_field = false,
189        set = false,
190        with = false,
191        get_mut = false
192    )]
193    pub(crate) buffer: Buffer,
194
195    /// The [`HttpContext`] shared for this server
196    #[field(deref = false)]
197    pub(crate) context: Arc<HttpContext>,
198
199    /// the ip address of the connection, if available
200    #[field(copy)]
201    pub(crate) peer_ip: Option<IpAddr>,
202
203    /// the wall-clock time at which the underlying [`Conn`] was constructed
204    #[field(copy)]
205    pub(crate) start_time: Instant,
206
207    /// the :authority http/3 pseudo-header
208    pub(crate) authority: Option<Cow<'static, str>>,
209
210    /// the :scheme http/3 pseudo-header
211    pub(crate) scheme: Option<Cow<'static, str>>,
212
213    /// the [`ProtocolSession`] for this upgrade — h2/h3 connection driver + stream id
214    /// where applicable; `Http1` for upgrades from h1 or synthetic conns.
215    #[field = false]
216    pub(crate) protocol_session: ProtocolSession,
217
218    /// the :protocol http/3 pseudo-header
219    pub(crate) protocol: Option<Cow<'static, str>>,
220
221    /// the http version
222    #[field = "http_version"]
223    pub(crate) version: Version,
224
225    /// the http response status set on the underlying [`Conn`] before the upgrade
226    /// (typically `101 Switching Protocols`, or `200 OK` for CONNECT). `None` if unset.
227    #[field(copy)]
228    pub(crate) status: Option<Status>,
229
230    /// whether this connection was deemed secure by the handler stack
231    pub(crate) secure: bool,
232
233    /// Inbound framing state carried across the upgrade so the inbound state machine
234    /// resumes where the pre-upgrade handler left off. Request-body state on server
235    /// upgrades; response-body state on client upgrades.
236    #[field = false]
237    pub(crate) received_body_state: ReceivedBodyState,
238
239    /// Inbound trailers, populated either by a fully-consumed pre-upgrade body or by
240    /// the post-upgrade inbound state machine. `Some` only when non-empty.
241    #[field(get, get_mut, take, set = false, with = false, into_field = false)]
242    pub(crate) received_trailers: Option<Headers>,
243
244    /// Pre-parsed inbound `Content-Length`. `None` for chunked, missing, or malformed.
245    #[field = false]
246    pub(crate) content_length_in: Option<u64>,
247
248    /// Per-protocol outbound framing state. Decided at the upgrade transition.
249    #[field = false]
250    pub(crate) write_state: WriteState,
251
252    /// Charset of the inbound body, parsed from the inbound `Content-Type`'s `charset`
253    /// parameter at the upgrade transition.
254    #[field = false]
255    pub(crate) inbound_encoding: &'static Encoding,
256
257    /// In-flight QPACK trailer-decode future for inbound h3 trailing HEADERS. Held here
258    /// so its registered waker survives across `poll_read` calls — dropping the future
259    /// would drop the waker the QPACK decoder is parked on, hanging the reader.
260    #[field = false]
261    pub(crate) h3_trailer_decode_in: Option<H3TrailerFuture>,
262
263    /// Accumulator for inbound h3 trailing-HEADERS payload bytes pre-QPACK-decode.
264    /// Separate from [`Self::buffer`] so the inbound state machine doesn't recycle
265    /// accumulated trailer bytes back through the frame decoder and double-count them.
266    #[field = false]
267    pub(crate) h3_trailer_payload_in: Vec<u8>,
268
269    /// Resolves when the peer abandons this stream. HTTP/3 only; see [`PeerGone`].
270    #[field = false]
271    pub(crate) peer_gone: Option<PeerGone>,
272}
273
274impl<Transport> Upgrade<Transport> {
275    #[doc(hidden)]
276    pub fn new(
277        received_headers: Headers,
278        path: impl Into<Cow<'static, str>>,
279        method: Method,
280        transport: Transport,
281        buffer: Buffer,
282        version: Version,
283    ) -> Self {
284        Self {
285            inner: Box::new(UpgradeInner {
286                received_headers,
287                sent_headers: Headers::new(),
288                path: path.into(),
289                method,
290                buffer,
291                state: TypeSet::new(),
292                context: Arc::default(),
293                peer_ip: None,
294                start_time: Instant::now(),
295                authority: None,
296                scheme: None,
297                protocol_session: ProtocolSession::Http1,
298                protocol: None,
299                secure: false,
300                version,
301                status: None,
302                received_body_state: ReceivedBodyState::Raw { total: 0 },
303                received_trailers: None,
304                content_length_in: None,
305                write_state: WriteState::Raw,
306                inbound_encoding: encoding_rs::UTF_8,
307                h3_trailer_decode_in: None,
308                h3_trailer_payload_in: Vec::new(),
309                peer_gone: None,
310            }),
311            transport,
312        }
313    }
314
315    #[cfg(feature = "unstable")]
316    #[doc(hidden)]
317    #[allow(clippy::too_many_arguments)]
318    pub fn from_parts(
319        received_headers: Headers,
320        sent_headers: Headers,
321        path: Cow<'static, str>,
322        method: Method,
323        transport: Transport,
324        buffer: Buffer,
325        state: TypeSet,
326        context: Arc<HttpContext>,
327        peer_ip: Option<IpAddr>,
328        authority: Option<Cow<'static, str>>,
329        scheme: Option<Cow<'static, str>>,
330        protocol_session: ProtocolSession,
331        protocol: Option<Cow<'static, str>>,
332        version: Version,
333        status: Option<Status>,
334        secure: bool,
335        received_body_state: ReceivedBodyState,
336        received_trailers: Option<Headers>,
337    ) -> Self {
338        let write_state = compute_write_state(version, &sent_headers);
339        let content_length_in = parse_content_length(&received_headers);
340        let inbound_encoding = encoding(&received_headers);
341
342        Self {
343            inner: Box::new(UpgradeInner {
344                // Client-side upgrades have no peer-departure hook; on h3 that leaves
345                // `poll_closed` permanently pending. Client h3 resets surface as read errors.
346                peer_gone: None,
347                received_headers,
348                sent_headers,
349                path,
350                method,
351                state,
352                buffer,
353                context,
354                peer_ip,
355                start_time: Instant::now(),
356                authority,
357                scheme,
358                protocol_session,
359                protocol,
360                version,
361                status,
362                secure,
363                received_body_state,
364                received_trailers,
365                content_length_in,
366                write_state,
367                inbound_encoding,
368                h3_trailer_decode_in: None,
369                h3_trailer_payload_in: Vec::new(),
370            }),
371            transport,
372        }
373    }
374
375    /// the [`H2Connection`] driver for this upgrade, if it originated from an HTTP/2 stream
376    pub fn h2_connection(&self) -> Option<&Arc<H2Connection>> {
377        self.inner.protocol_session.h2_connection()
378    }
379
380    /// the h2 stream id for this upgrade, if it originated from an HTTP/2 stream
381    pub fn h2_stream_id(&self) -> Option<u32> {
382        self.inner.protocol_session.h2_stream_id()
383    }
384
385    /// the [`H3Connection`] driver for this upgrade, if it originated from an HTTP/3 stream
386    pub fn h3_connection(&self) -> Option<&Arc<H3Connection>> {
387        self.inner.protocol_session.h3_connection()
388    }
389
390    /// the h3 stream id for this upgrade, if it originated from an HTTP/3 stream
391    pub fn h3_stream_id(&self) -> Option<u64> {
392        self.inner.protocol_session.h3_stream_id()
393    }
394
395    /// Take any buffered bytes
396    pub fn take_buffer(&mut self) -> Vec<u8> {
397        std::mem::take(&mut self.inner.buffer).into()
398    }
399
400    /// Mutably borrow any bytes that have already been read from the underlying transport.
401    ///
402    /// It is your responsibility to process these bytes before reading directly from the
403    /// transport.
404    pub fn buffer_mut(&mut self) -> &mut [u8] {
405        self.inner.buffer.live_mut()
406    }
407
408    #[doc(hidden)]
409    pub fn buffer_and_transport_mut(&mut self) -> (&mut Buffer, &mut Transport) {
410        (&mut self.inner.buffer, &mut self.transport)
411    }
412
413    /// borrow the shared state [`TypeSet`] for this application
414    pub fn shared_state(&self) -> &TypeSet {
415        self.inner.context.shared_state()
416    }
417
418    /// the http request path up to but excluding any query component
419    pub fn path(&self) -> &str {
420        match self.inner.path.split_once('?') {
421            Some((path, _)) => path,
422            None => &self.inner.path,
423        }
424    }
425
426    /// retrieves the query component of the path
427    pub fn querystring(&self) -> &str {
428        self.inner
429            .path
430            .split_once('?')
431            .map_or_default(|(_, query)| query)
432    }
433
434    /// Modify the transport type of this upgrade.
435    ///
436    /// This is useful for boxing the transport in order to erase the type argument.
437    pub fn map_transport<T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static>(
438        self,
439        f: impl Fn(Transport) -> T,
440    ) -> Upgrade<T> {
441        Upgrade {
442            inner: self.inner,
443            transport: f(self.transport),
444        }
445    }
446}
447
448impl<Transport: AsyncRead + Unpin> Upgrade<Transport> {
449    /// Resolves when the peer has abandoned this upgrade — connection closed, stream
450    /// reset, or transport error.
451    ///
452    /// This is a liveness probe for upgrades the peer is not expected to speak into,
453    /// such as a server-sent event stream. If the peer may legitimately send data, read
454    /// it instead: on HTTP/1.x this probe detects closure *by reading*, treating inbound
455    /// bytes as incidental. Probed bytes accumulate on the internal buffer — where the
456    /// [`AsyncRead`] impl yields them before touching the transport — until the buffer
457    /// holds `read_allowance` bytes, after which the probe goes dormant, returning
458    /// `Poll::Pending` without scheduling a wake until reads drain the buffer.
459    ///
460    /// Per-protocol behavior:
461    /// - HTTP/1.x: reads the transport; end-of-file or a transport error resolves. A half-closed
462    ///   peer that shut down its write side but still reads is indistinguishable from a departed
463    ///   one and counts as closed.
464    /// - HTTP/2: resolves when the stream is reset or fully closed, or the connection is torn down.
465    ///   Never reads the transport; `read_allowance` is unused.
466    /// - HTTP/3: resolves on `STOP_SENDING`, stream reset, or connection loss. Never reads the
467    ///   transport; `read_allowance` is unused. Requires the runtime adapter to have supplied
468    ///   [`H3BidiRequest::with_peer_gone`][crate::h3::H3BidiRequest::with_peer_gone]; without it,
469    ///   this never resolves.
470    ///
471    /// A client that vanishes without signalling — a killed process, a severed network — is
472    /// only detected once the transport notices. On HTTP/3 that is QUIC's idle timeout, which
473    /// is always negotiated and typically under a minute. On HTTP/1.x and HTTP/2 it depends on
474    /// the TCP configuration and can be considerably longer.
475    pub fn poll_closed(
476        self: Pin<&mut Self>,
477        cx: &mut Context<'_>,
478        read_allowance: usize,
479    ) -> Poll<()> {
480        let Self { inner, transport } = self.get_mut();
481
482        crate::liveness::poll_peer_gone(
483            &inner.protocol_session,
484            inner.version,
485            &mut inner.buffer,
486            transport,
487            inner.peer_gone.as_mut(),
488            read_allowance,
489            cx,
490        )
491    }
492}
493
494impl<Transport: AsyncWrite + Unpin> Upgrade<Transport> {
495    /// Emit trailing headers and finish the outbound stream. Consumes `self`; further
496    /// writes are statically prevented.
497    ///
498    /// Per-protocol behavior:
499    /// - HTTP/1.1 with `Transfer-Encoding: chunked`: writes the last-chunk marker (`0\r\n`), the
500    ///   trailer section, and a final CRLF, then closes the transport.
501    /// - HTTP/2: enqueues a trailing `HEADERS` frame with `END_STREAM` via the connection driver
502    ///   and returns. The driver finishes the stream after draining any pending DATA frames.
503    /// - HTTP/3: encodes a trailing `HEADERS` frame via QPACK, writes it to the stream, then closes
504    ///   the stream (QUIC `FIN`).
505    /// - HTTP/1.1 without chunked encoding (raw upgrade, CONNECT tunnel, websocket-over-h1):
506    ///   trailers can't be expressed on the wire; dropped with a `log::warn!` and `Ok(())`
507    ///   returned.
508    ///
509    /// # Errors
510    ///
511    /// Returns the underlying [`io::Error`] when the wire write fails, `BrokenPipe` if
512    /// the stream has already been closed, and `NotConnected` if the carried
513    /// `ProtocolSession` is missing the expected driver for h2/h3.
514    pub async fn send_trailers(self, trailers: Headers) -> io::Result<()> {
515        let Self {
516            inner,
517            mut transport,
518        } = self;
519        let UpgradeInner {
520            mut write_state,
521            context,
522            protocol_session,
523            ..
524        } = *inner;
525
526        match &mut write_state {
527            WriteState::H1Chunked(state) => {
528                if state.terminator_written {
529                    return Err(io::ErrorKind::BrokenPipe.into());
530                }
531                state.pending.extend_from_slice(b"0\r\n");
532                crate::conn::write_headers_or_trailers(&mut state.pending, &trailers, &context)
533                    .map_err(io::Error::other)?;
534                state.pending.extend_from_slice(b"\r\n");
535                state.terminator_written = true;
536
537                transport.write_all(&state.pending).await?;
538                state.pending.clear();
539                transport.close().await
540            }
541            WriteState::H3Framed(state) => {
542                if state.terminator_written {
543                    return Err(io::ErrorKind::BrokenPipe.into());
544                }
545                let Some((h3, stream_id)) = protocol_session.as_h3() else {
546                    return Err(io::ErrorKind::NotConnected.into());
547                };
548                let field_section = FieldSection::new(PseudoHeaders::default(), &trailers);
549                h3.encode_field_section_framed(&field_section, &mut state.pending, stream_id)?;
550                state.terminator_written = true;
551
552                transport.write_all(&state.pending).await?;
553                state.pending.clear();
554                transport.close().await
555            }
556            WriteState::Raw => {
557                if let Some((h2, stream_id)) = protocol_session.as_h2() {
558                    h2.submit_trailers(stream_id, trailers)
559                } else {
560                    log::warn!(
561                        "Upgrade::send_trailers called on a raw upgrade with no per-stream \
562                         framing; trailers dropped. Set `Transfer-Encoding: chunked` on the \
563                         outbound headers if you intend to emit trailers over HTTP/1.1."
564                    );
565                    Ok(())
566                }
567            }
568        }
569    }
570}
571
572impl<Transport> Debug for Upgrade<Transport> {
573    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
574        f.debug_struct(&format!("Upgrade<{}>", std::any::type_name::<Transport>()))
575            .field("received_headers", &self.inner.received_headers)
576            .field("sent_headers", &self.inner.sent_headers)
577            .field("path", &self.inner.path)
578            .field("method", &self.inner.method)
579            .field("buffer", &self.inner.buffer)
580            .field("context", &self.inner.context)
581            .field("state", &self.inner.state)
582            .field("transport", &format_args!(".."))
583            .field("peer_ip", &self.inner.peer_ip)
584            .field("start_time", &self.inner.start_time)
585            .field("authority", &self.inner.authority)
586            .field("scheme", &self.inner.scheme)
587            .field("protocol_session", &self.inner.protocol_session)
588            .field("protocol", &self.inner.protocol)
589            .field("version", &self.inner.version)
590            .field("status", &self.inner.status)
591            .field("secure", &self.inner.secure)
592            .field("received_body_state", &self.inner.received_body_state)
593            .field("received_trailers", &self.inner.received_trailers)
594            .field("content_length_in", &self.inner.content_length_in)
595            .field("write_state", &self.inner.write_state)
596            .field("inbound_encoding", &self.inner.inbound_encoding.name())
597            .field(
598                "h3_trailer_decode_in",
599                &self
600                    .inner
601                    .h3_trailer_decode_in
602                    .as_ref()
603                    .map(|_| format_args!("..")),
604            )
605            .field(
606                "h3_trailer_payload_in_len",
607                &self.inner.h3_trailer_payload_in.len(),
608            )
609            .field(
610                "peer_gone",
611                &self.inner.peer_gone.as_ref().map(|_| format_args!("..")),
612            )
613            .finish()
614    }
615}
616
617impl<Transport> From<Conn<Transport>> for Upgrade<Transport> {
618    fn from(conn: Conn<Transport>) -> Self {
619        let Conn { inner, transport } = conn;
620        // Exhaustive destructure so new fields on `ConnInner` force a deliberate carry-vs-drop
621        // decision.
622        let ConnInner {
623            request_headers,
624            response_headers,
625            path,
626            method,
627            state,
628            buffer,
629            context,
630            peer_ip,
631            start_time,
632            authority,
633            scheme,
634            protocol_session,
635            protocol,
636            version,
637            status,
638            secure,
639            request_body_state,
640            request_trailers,
641            response_body,
642            // post-send hooks no longer apply; `upgrade` is the marker that brought us here
643            after_send: _,
644            upgrade: _,
645            // only meaningful to the h1 keepalive loop, which an upgrade leaves for good
646            spare_request_headers: _,
647            response_buffer: _,
648            peer_gone,
649        } = *inner;
650
651        if let Some(body) = &response_body
652            && !body.is_empty()
653        {
654            log::warn!(
655                "Conn::upgrade() and a non-empty response body are both set; body is being \
656                 discarded. The upgrade path is mutually exclusive with serving a response body."
657            );
658        }
659
660        // Server-side roles: outbound = response_headers, inbound = request_headers.
661        let write_state = compute_write_state(version, &response_headers);
662        let content_length_in = parse_content_length(&request_headers);
663        let inbound_encoding = encoding(&request_headers);
664        // An h1 request with no framing headers parses to `End` — correct for the request
665        // body, but inherited across the upgrade it would EOF the first read of a live raw
666        // stream (a browser websocket handshake is exactly this shape). Declared framing
667        // does carry over: a chunked request keeps chunked inbound framing.
668        let received_body_state = if matches!(version, Version::Http1_0 | Version::Http1_1)
669            && !request_headers.has_header(KnownHeaderName::TransferEncoding)
670            && !request_headers.has_header(KnownHeaderName::ContentLength)
671        {
672            ReceivedBodyState::Raw { total: 0 }
673        } else {
674            request_body_state
675        };
676        let received_trailers = request_trailers.filter(|t| !t.is_empty());
677
678        Self {
679            inner: Box::new(UpgradeInner {
680                received_headers: request_headers,
681                sent_headers: response_headers,
682                path,
683                method,
684                state,
685                buffer,
686                context: Arc::unwrap_or_clone(context),
687                peer_ip,
688                start_time,
689                authority,
690                scheme,
691                protocol_session,
692                protocol,
693                version,
694                status,
695                secure,
696                received_body_state,
697                received_trailers,
698                content_length_in,
699                write_state,
700                inbound_encoding,
701                h3_trailer_decode_in: None,
702                h3_trailer_payload_in: Vec::new(),
703                peer_gone,
704            }),
705            transport,
706        }
707    }
708}
709
710#[cfg(test)]
711mod tests;
712
713impl<Transport> AsyncRead for Upgrade<Transport>
714where
715    Transport: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
716{
717    fn poll_read(
718        mut self: Pin<&mut Self>,
719        cx: &mut Context<'_>,
720        buf: &mut [u8],
721    ) -> Poll<io::Result<usize>> {
722        let Self { inner, transport } = &mut *self;
723        let UpgradeInner {
724            buffer,
725            received_body_state,
726            content_length_in,
727            context,
728            protocol_session,
729            received_trailers,
730            h3_trailer_decode_in,
731            h3_trailer_payload_in,
732            inbound_encoding,
733            ..
734        } = &mut **inner;
735
736        let protocol_session = protocol_session.clone();
737        let mut body: ReceivedBody<'_, Transport> = ReceivedBody::new_with_config(
738            *content_length_in,
739            buffer,
740            transport,
741            received_body_state,
742            None,
743            inbound_encoding,
744            &context.config,
745        )
746        .with_trailers(received_trailers)
747        .with_protocol_session(protocol_session)
748        .with_h3_trailer_future(h3_trailer_decode_in)
749        .with_h3_trailer_payload_buffer(h3_trailer_payload_in);
750
751        Pin::new(&mut body).poll_read(cx, buf)
752    }
753}
754
755impl<Transport: AsyncWrite + Unpin> AsyncWrite for Upgrade<Transport> {
756    fn poll_write(
757        mut self: Pin<&mut Self>,
758        cx: &mut Context<'_>,
759        buf: &[u8],
760    ) -> Poll<io::Result<usize>> {
761        let Self { inner, transport } = &mut *self;
762        let write_state = &mut inner.write_state;
763        match write_state {
764            WriteState::Raw => Pin::new(transport).poll_write(cx, buf),
765            WriteState::H1Chunked(state) => {
766                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
767
768                // Empty buf must not become a chunk: `0\r\n` IS the last-chunk marker.
769                if buf.is_empty() {
770                    return Poll::Ready(Ok(0));
771                }
772
773                if state.terminator_written {
774                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
775                }
776
777                write_chunk(&mut state.pending, buf);
778                best_effort_drain(&mut state.pending, cx, transport)?;
779                Poll::Ready(Ok(buf.len()))
780            }
781            WriteState::H3Framed(state) => {
782                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
783
784                if buf.is_empty() {
785                    return Poll::Ready(Ok(0));
786                }
787
788                if state.terminator_written {
789                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
790                }
791
792                encode_h3_data_header(&mut state.pending, buf.len() as u64);
793                state.pending.extend_from_slice(buf);
794                best_effort_drain(&mut state.pending, cx, transport)?;
795                Poll::Ready(Ok(buf.len()))
796            }
797        }
798    }
799
800    fn poll_write_vectored(
801        mut self: Pin<&mut Self>,
802        cx: &mut Context<'_>,
803        bufs: &[IoSlice<'_>],
804    ) -> Poll<io::Result<usize>> {
805        let Self { inner, transport } = &mut *self;
806        let write_state = &mut inner.write_state;
807        match write_state {
808            WriteState::Raw => Pin::new(transport).poll_write_vectored(cx, bufs),
809            WriteState::H1Chunked(state) => {
810                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
811                let total: usize = bufs.iter().map(|b| b.len()).sum();
812                if total == 0 {
813                    return Poll::Ready(Ok(0));
814                }
815                if state.terminator_written {
816                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
817                }
818                // One chunk per vectored batch — the default impl would emit one chunk per
819                // iobuf, which is wasteful when the caller meant them as one logical write.
820                let _ = write!(state.pending, "{total:X}\r\n");
821                for b in bufs {
822                    state.pending.extend_from_slice(b);
823                }
824                state.pending.extend_from_slice(b"\r\n");
825                best_effort_drain(&mut state.pending, cx, transport)?;
826                Poll::Ready(Ok(total))
827            }
828            WriteState::H3Framed(state) => {
829                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
830                let total: usize = bufs.iter().map(|b| b.len()).sum();
831                if total == 0 {
832                    return Poll::Ready(Ok(0));
833                }
834                if state.terminator_written {
835                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
836                }
837                // One DATA frame per vectored batch — collapses `[length_prefix, payload]`
838                // pairs into a single frame.
839                encode_h3_data_header(&mut state.pending, total as u64);
840                for b in bufs {
841                    state.pending.extend_from_slice(b);
842                }
843                best_effort_drain(&mut state.pending, cx, transport)?;
844                Poll::Ready(Ok(total))
845            }
846        }
847    }
848
849    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
850        let Self { inner, transport } = &mut *self;
851        let write_state = &mut inner.write_state;
852        match write_state {
853            WriteState::Raw => Pin::new(transport).poll_flush(cx),
854            WriteState::H1Chunked(state) => {
855                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
856                Pin::new(transport).poll_flush(cx)
857            }
858            WriteState::H3Framed(state) => {
859                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
860                Pin::new(transport).poll_flush(cx)
861            }
862        }
863    }
864
865    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
866        let Self { inner, transport } = &mut *self;
867        let write_state = &mut inner.write_state;
868        match write_state {
869            WriteState::Raw => Pin::new(transport).poll_close(cx),
870            WriteState::H1Chunked(state) => {
871                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
872                if !state.terminator_written {
873                    state.pending.extend_from_slice(b"0\r\n\r\n");
874                    // Flag set before the drain so a re-poll after Pending doesn't re-append.
875                    state.terminator_written = true;
876                }
877                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
878                Pin::new(transport).poll_close(cx)
879            }
880            WriteState::H3Framed(state) => {
881                // h3 stream-end is the QUIC FIN — no separate terminator frame.
882                ready!(poll_drain_pending(&mut state.pending, cx, transport))?;
883                state.terminator_written = true;
884                Pin::new(transport).poll_close(cx)
885            }
886        }
887    }
888}
trillium_http/upgrade.rs

trillium_http/
upgrade.rs
