Skip to main content

trillium_http/
conn.rs

1use crate::{
2    Body, Buffer, Headers, HttpContext,
3    KnownHeaderName::Host,
4    Method, ProtocolSession, ReceivedBody, Status, Swansong, TypeSet, Version,
5    after_send::{AfterSend, SendStatus},
6    h2::H2Connection,
7    h3::H3Connection,
8    liveness::{CancelOnDisconnect, LivenessFut},
9    received_body::ReceivedBodyState,
10    util::encoding,
11};
12use encoding_rs::Encoding;
13use futures_lite::{
14    future,
15    io::{AsyncRead, AsyncWrite},
16};
17use std::{
18    borrow::Cow,
19    fmt::{self, Debug, Formatter},
20    future::Future,
21    net::IpAddr,
22    pin::pin,
23    str,
24    sync::Arc,
25    time::Instant,
26};
27mod h1;
28#[cfg(test)]
29mod h1_tests;
30mod h2;
31mod h3;
32mod shared;
33pub(crate) use h1::{HeadError, write_headers_or_trailers};
34pub(crate) use h3::H3FirstFrame;
35pub(crate) use shared::ConnParts;
36
37/// An HTTP connection.
38///
39/// This struct represents both the request and the response, and holds the
40/// transport over which the response will be sent.
41#[derive(fieldwork::Fieldwork)]
42pub struct Conn<Transport> {
43    #[field(get)]
44    /// the shared [`HttpContext`]
45    pub(crate) context: Arc<HttpContext>,
46
47    /// request [headers](Headers)
48    #[field(get, get_mut)]
49    pub(crate) request_headers: Headers,
50
51    /// response [headers](Headers)
52    #[field(get, get_mut)]
53    pub(crate) response_headers: Headers,
54
55    pub(crate) path: Cow<'static, str>,
56
57    /// the http method for this conn's request
58    ///
59    /// ```
60    /// # use trillium_http::{Conn, Method};
61    /// let mut conn = Conn::new_synthetic(Method::Get, "/some/path?and&a=query", ());
62    /// assert_eq!(conn.method(), Method::Get);
63    /// ```
64    #[field(get, set, copy)]
65    pub(crate) method: Method,
66
67    /// the http status for this conn, if set
68    #[field(get, copy)]
69    pub(crate) status: Option<Status>,
70
71    /// The HTTP protocol version in use on this connection.
72    ///
73    /// ```
74    /// # use trillium_http::{Conn, Method, Version};
75    /// let conn = Conn::new_synthetic(Method::Get, "/", ());
76    /// assert_eq!(conn.http_version(), Version::Http1_1);
77    /// ```
78    #[field(get = http_version, copy)]
79    pub(crate) version: Version,
80
81    /// the [state typemap](TypeSet) for this conn
82    #[field(get, get_mut)]
83    pub(crate) state: TypeSet,
84
85    /// the response [body](Body)
86    ///
87    /// ```
88    /// # use trillium_testing::HttpTest;
89    /// HttpTest::new(|conn| async move { conn.with_response_body("hello") })
90    ///     .get("/")
91    ///     .block()
92    ///     .assert_body("hello");
93    ///
94    /// HttpTest::new(|conn| async move { conn.with_response_body(String::from("world")) })
95    ///     .get("/")
96    ///     .block()
97    ///     .assert_body("world");
98    ///
99    /// HttpTest::new(|conn| async move { conn.with_response_body(vec![99, 97, 116]) })
100    ///     .get("/")
101    ///     .block()
102    ///     .assert_body("cat");
103    /// ```
104    #[field(get, set, into, option_set_some, take, with)]
105    pub(crate) response_body: Option<Body>,
106
107    /// the transport
108    ///
109    /// This should only be used to call your own custom methods on the transport that do not read
110    /// or write any data. Calling any method that reads from or writes to the transport will
111    /// disrupt the HTTP protocol. If you're looking to transition from HTTP to another protocol,
112    /// use an HTTP upgrade.
113    #[field(get, get_mut)]
114    pub(crate) transport: Transport,
115
116    pub(crate) buffer: Buffer,
117
118    pub(crate) request_body_state: ReceivedBodyState,
119
120    pub(crate) after_send: AfterSend,
121
122    /// whether the connection is secure
123    ///
124    /// note that this does not necessarily indicate that the transport itself is secure, as it may
125    /// indicate that `trillium_http` is behind a trusted reverse proxy that has terminated tls and
126    /// provided appropriate headers to indicate this.
127    #[field(get, set, rename_predicates)]
128    pub(crate) secure: bool,
129
130    /// The [`Instant`] that the first header bytes for this conn were
131    /// received, before any processing or parsing has been performed.
132    #[field(get, copy)]
133    pub(crate) start_time: Instant,
134
135    /// The IP Address for the connection, if available
136    #[field(set, get, copy, into)]
137    pub(crate) peer_ip: Option<IpAddr>,
138
139    /// the `:authority` pseudo-header
140    #[field(set, get, into)]
141    pub(crate) authority: Option<Cow<'static, str>>,
142
143    /// the `:scheme` pseudo-header
144    #[field(set, get, into)]
145    pub(crate) scheme: Option<Cow<'static, str>>,
146
147    /// the [`ProtocolSession`] for this conn — the per-protocol session state
148    /// (h2/h3 connection driver and stream id) bundled into a single enum so the
149    /// "set together" invariant is enforced at the type level. `Http1` for
150    /// h1 / synthetic conns.
151    pub(crate) protocol_session: ProtocolSession,
152
153    /// the `:protocol` pseudo-header (extended CONNECT)
154    #[field(set, get, into)]
155    pub(crate) protocol: Option<Cow<'static, str>>,
156
157    /// request trailers, populated after the request body has been fully read
158    #[field(get, get_mut)]
159    pub(crate) request_trailers: Option<Headers>,
160
161    /// Marker set via [`Conn::upgrade`].
162    pub(crate) upgrade: bool,
163}
164
165impl<Transport> Debug for Conn<Transport> {
166    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
167        f.debug_struct("Conn")
168            .field("context", &self.context)
169            .field("request_headers", &self.request_headers)
170            .field("response_headers", &self.response_headers)
171            .field("path", &self.path)
172            .field("method", &self.method)
173            .field("status", &self.status)
174            .field("version", &self.version)
175            .field("state", &self.state)
176            .field("response_body", &self.response_body)
177            .field("transport", &format_args!(".."))
178            .field("buffer", &format_args!(".."))
179            .field("request_body_state", &self.request_body_state)
180            .field("secure", &self.secure)
181            .field("after_send", &format_args!(".."))
182            .field("start_time", &self.start_time)
183            .field("peer_ip", &self.peer_ip)
184            .field("authority", &self.authority)
185            .field("scheme", &self.scheme)
186            .field("protocol", &self.protocol)
187            .field("protocol_session", &self.protocol_session)
188            .field("request_trailers", &self.request_trailers)
189            .field("upgrade", &self.upgrade)
190            .finish()
191    }
192}
193
194impl<Transport> Conn<Transport>
195where
196    Transport: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
197{
198    /// Returns the shared state typemap for this conn.
199    pub fn shared_state(&self) -> &TypeSet {
200        &self.context.shared_state
201    }
202
203    /// sets the http status code from any `TryInto<Status>`.
204    ///
205    /// ```
206    /// # use trillium_http::Status;
207    /// # trillium_testing::HttpTest::new(|mut conn| async move {
208    /// assert!(conn.status().is_none());
209    ///
210    /// conn.set_status(200); // a status can be set as a u16
211    /// assert_eq!(conn.status().unwrap(), Status::Ok);
212    ///
213    /// conn.set_status(Status::ImATeapot); // or as a Status
214    /// assert_eq!(conn.status().unwrap(), Status::ImATeapot);
215    /// conn
216    /// # }).get("/").block().assert_status(Status::ImATeapot);
217    /// ```
218    pub fn set_status(&mut self, status: impl TryInto<Status>) -> &mut Self {
219        self.status = Some(status.try_into().unwrap_or_else(|_| {
220            log::error!("attempted to set an invalid status code");
221            Status::InternalServerError
222        }));
223        self
224    }
225
226    /// sets the http status code from any `TryInto<Status>`, returning Conn
227    #[must_use]
228    pub fn with_status(mut self, status: impl TryInto<Status>) -> Self {
229        self.set_status(status);
230        self
231    }
232
233    /// The status to send on the wire: the explicitly-set status, or a
234    /// method-appropriate default when a handler left it unset. Unhandled
235    /// requests default to `404 Not Found`, except CONNECT, which defaults to
236    /// `501 Not Implemented`: an origin server implements no tunnel, and 404's
237    /// resource model does not apply to CONNECT's authority-form target.
238    pub(crate) fn response_status(&self) -> Status {
239        self.status.unwrap_or(match self.method {
240            Method::Connect => Status::NotImplemented,
241            _ => Status::NotFound,
242        })
243    }
244
245    /// retrieves the path part of the request url, up to and excluding any query component
246    /// ```
247    /// # use trillium_testing::HttpTest;
248    /// HttpTest::new(|mut conn| async move {
249    ///     assert_eq!(conn.path(), "/some/path");
250    ///     conn.with_status(200)
251    /// })
252    /// .get("/some/path?and&a=query")
253    /// .block()
254    /// .assert_ok();
255    /// ```
256    pub fn path(&self) -> &str {
257        match self.path.split_once('?') {
258            Some((path, _)) => path,
259            None => &self.path,
260        }
261    }
262
263    /// retrieves the combined path and any query
264    pub fn path_and_query(&self) -> &str {
265        &self.path
266    }
267
268    /// retrieves the query component of the path, or an empty &str
269    ///
270    /// ```
271    /// # use trillium_testing::HttpTest;
272    /// let server = HttpTest::new(|conn| async move {
273    ///     let querystring = conn.querystring().to_string();
274    ///     conn.with_response_body(querystring).with_status(200)
275    /// });
276    ///
277    /// server
278    ///     .get("/some/path?and&a=query")
279    ///     .block()
280    ///     .assert_body("and&a=query");
281    ///
282    /// server.get("/some/path").block().assert_body("");
283    /// ```
284    pub fn querystring(&self) -> &str {
285        self.path
286            .split_once('?')
287            .map(|(_, query)| query)
288            .unwrap_or_default()
289    }
290
291    /// get the host for this conn, if it exists.
292    ///
293    /// On protocol versions where the equivalent of `Host` is `:authority`, this returns
294    /// `:authority`.
295    pub fn host(&self) -> Option<&str> {
296        self.request_headers
297            .get_str(Host)
298            .or_else(|| self.authority())
299    }
300
301    /// set the host for this conn
302    pub fn set_host(&mut self, host: String) -> &mut Self {
303        self.request_headers.insert(Host, host);
304        self
305    }
306
307    /// Cancels and drops the future if reading from the transport results in an error or empty read
308    ///
309    /// If the client disconnects from the conn's transport, this function will return None. If the
310    /// future completes without disconnection, this future will return Some containing the output
311    /// of the future.
312    ///
313    /// Disconnection is detected by reading from the transport, so any bytes the client sends
314    /// while the future runs — an unread request body, or pipelined requests — are buffered, up
315    /// to 16kb. A client that fills that allowance is considered alive for the remainder of the
316    /// future, even if it disconnects afterwards. If the request has a body, read it before
317    /// calling this.
318    ///
319    /// Note that the inner future cannot borrow conn, so you will need to clone or take any
320    /// information needed to execute the future prior to executing this method.
321    ///
322    /// # Example
323    ///
324    /// ```rust
325    /// # use futures_lite::{AsyncRead, AsyncWrite};
326    /// # use trillium_http::{Conn, Method};
327    /// async fn something_slow_and_cancel_safe() -> String {
328    ///     String::from("this was not actually slow")
329    /// }
330    /// async fn handler<T>(mut conn: Conn<T>) -> Conn<T>
331    /// where
332    ///     T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
333    /// {
334    ///     let Some(returned_body) = conn
335    ///         .cancel_on_disconnect(async { something_slow_and_cancel_safe().await })
336    ///         .await
337    ///     else {
338    ///         return conn;
339    ///     };
340    ///     conn.with_response_body(returned_body).with_status(200)
341    /// }
342    /// ```
343    pub async fn cancel_on_disconnect<'a, Fut>(&'a mut self, fut: Fut) -> Option<Fut::Output>
344    where
345        Fut: Future + Send + 'a,
346    {
347        CancelOnDisconnect(self, pin!(fut)).await
348    }
349
350    /// Check if the transport is connected by attempting to read from the transport
351    ///
352    /// Any bytes the client sends — an unread request body, or pipelined requests — are buffered,
353    /// up to 16kb, and count as evidence of liveness. A client that fills that allowance is
354    /// reported as connected until the buffered bytes are read, so read the request body before
355    /// polling this in a long-running handler.
356    ///
357    /// # Example
358    ///
359    /// This is best to use at appropriate points in a long-running handler, like:
360    ///
361    /// ```rust
362    /// # use futures_lite::{AsyncRead, AsyncWrite};
363    /// # use trillium_http::{Conn, Method};
364    /// # async fn something_slow_but_not_cancel_safe() {}
365    /// async fn handler<T>(mut conn: Conn<T>) -> Conn<T>
366    /// where
367    ///     T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
368    /// {
369    ///     for _ in 0..100 {
370    ///         if conn.is_disconnected().await {
371    ///             return conn;
372    ///         }
373    ///         something_slow_but_not_cancel_safe().await;
374    ///     }
375    ///     conn.with_status(200)
376    /// }
377    /// ```
378    pub async fn is_disconnected(&mut self) -> bool {
379        future::poll_once(LivenessFut::new(self)).await.is_some()
380    }
381
382    /// returns the [`encoding_rs::Encoding`] for this request, as determined from the mime-type
383    /// charset, if available
384    ///
385    /// ```
386    /// # use trillium_testing::HttpTest;
387    /// HttpTest::new(|mut conn| async move {
388    ///     assert_eq!(conn.request_encoding(), encoding_rs::UTF_8); // the default
389    ///
390    ///     conn.request_headers_mut()
391    ///         .insert("content-type", "text/plain;charset=utf-16");
392    ///     assert_eq!(conn.request_encoding(), encoding_rs::UTF_16LE);
393    ///
394    ///     conn.with_status(200)
395    /// })
396    /// .get("/")
397    /// .block()
398    /// .assert_ok();
399    /// ```
400    pub fn request_encoding(&self) -> &'static Encoding {
401        encoding(&self.request_headers)
402    }
403
404    /// returns the [`encoding_rs::Encoding`] for this response, as
405    /// determined from the mime-type charset, if available
406    ///
407    /// ```
408    /// # use trillium_testing::HttpTest;
409    /// HttpTest::new(|mut conn| async move {
410    ///     assert_eq!(conn.response_encoding(), encoding_rs::UTF_8); // the default
411    ///     conn.response_headers_mut()
412    ///         .insert("content-type", "text/plain;charset=utf-16");
413    ///
414    ///     assert_eq!(conn.response_encoding(), encoding_rs::UTF_16LE);
415    ///
416    ///     conn.with_status(200)
417    /// })
418    /// .get("/")
419    /// .block()
420    /// .assert_ok();
421    /// ```
422    pub fn response_encoding(&self) -> &'static Encoding {
423        encoding(&self.response_headers)
424    }
425
426    /// returns a [`ReceivedBody`] that references this conn. the conn
427    /// retains all data and holds the singular transport, but the
428    /// `ReceivedBody` provides an interface to read body content.
429    ///
430    /// If the request included an `Expect: 100-continue` header, the 100 Continue response is sent
431    /// lazily on the first read from the returned [`ReceivedBody`].
432    /// ```
433    /// # use trillium_testing::HttpTest;
434    /// let server = HttpTest::new(|mut conn| async move {
435    ///     let request_body = conn.request_body();
436    ///     assert_eq!(request_body.content_length(), Some(5));
437    ///     assert_eq!(request_body.read_string().await.unwrap(), "hello");
438    ///     conn.with_status(200)
439    /// });
440    ///
441    /// server.post("/").with_body("hello").block().assert_ok();
442    /// ```
443    pub fn request_body(&mut self) -> ReceivedBody<'_, Transport> {
444        let needs_100_continue = self.needs_100_continue();
445        let body = self.build_request_body();
446        if needs_100_continue {
447            body.with_send_100_continue()
448        } else {
449            body
450        }
451    }
452
453    /// returns a clone of the [`swansong::Swansong`] for this Conn. use
454    /// this to gracefully stop long-running futures and streams
455    /// inside of handler functions
456    pub fn swansong(&self) -> Swansong {
457        self.protocol_session
458            .h3_connection()
459            .map_or_else(|| self.context.swansong.clone(), |h| h.swansong().clone())
460    }
461
462    /// Registers a function to call after the http response has been
463    /// completely transferred.
464    ///
465    /// The callback is guaranteed to fire **exactly once** before the conn is
466    /// dropped. Either the codec's send path invokes it with the real outcome,
467    /// or — if the conn is dropped before send completes (handler panic,
468    /// transport error, mid-write disconnect) — the drop fallback invokes it
469    /// with a `SendStatus` whose `is_success()` returns false. Multiple
470    /// registrations on the same conn chain in registration order.
471    ///
472    /// Because firing is ordered by send-completion rather than handler return,
473    /// this is the right hook for instrumentation that wants to report what the
474    /// peer actually observed.
475    ///
476    /// This is a sync function and should be computationally lightweight. If
477    /// your _application_ needs additional async processing, use your runtime's
478    /// task spawn within this hook. If your _library_ needs additional async
479    /// processing in an `after_send` hook, please open an issue.
480    pub fn after_send<F>(&mut self, after_send: F)
481    where
482        F: FnOnce(SendStatus) + Send + Sync + 'static,
483    {
484        self.after_send.append(after_send);
485    }
486
487    /// applies a mapping function from one transport to another. This
488    /// is particularly useful for boxing the transport. unless you're
489    /// sure this is what you're looking for, you probably don't want
490    /// to be using this
491    pub fn map_transport<NewTransport>(
492        self,
493        f: impl Fn(Transport) -> NewTransport,
494    ) -> Conn<NewTransport>
495    where
496        NewTransport: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
497    {
498        // Manual respread: rustc treats `Conn<Transport>` and `Conn<NewTransport>` as
499        // disjoint types and rejects `..self` without the unstable
500        // `type_changing_struct_update` feature. If a new field is added to `Conn`,
501        // update this respread, `Upgrade::map_transport`, and `From<Conn> for Upgrade`
502        // (`upgrade.rs`) — they share this drift hazard.
503        Conn {
504            context: self.context,
505            request_headers: self.request_headers,
506            response_headers: self.response_headers,
507            method: self.method,
508            response_body: self.response_body,
509            path: self.path,
510            status: self.status,
511            version: self.version,
512            state: self.state,
513            transport: f(self.transport),
514            buffer: self.buffer,
515            request_body_state: self.request_body_state,
516            secure: self.secure,
517            after_send: self.after_send,
518            start_time: self.start_time,
519            peer_ip: self.peer_ip,
520            authority: self.authority,
521            scheme: self.scheme,
522            protocol: self.protocol,
523            protocol_session: self.protocol_session,
524            request_trailers: self.request_trailers,
525            upgrade: self.upgrade,
526        }
527    }
528
529    /// whether this conn is suitable for an http upgrade to another protocol
530    pub fn should_upgrade(&self) -> bool {
531        self.upgrade
532            || (self.method() == Method::Connect && self.status == Some(Status::Ok))
533            || self.status == Some(Status::SwitchingProtocols)
534    }
535
536    /// Mark this conn to be handed off as an upgrade once the response headers are sent.
537    /// Set the response status (typically `200`) and any headers describing the upgraded
538    /// byte stream before calling; the handler's `upgrade` method receives an [`Upgrade`]
539    /// with per-protocol framing applied on its `AsyncRead`/`AsyncWrite`.
540    #[doc(hidden)]
541    #[must_use]
542    pub fn upgrade(mut self) -> Self {
543        self.upgrade = true;
544        self
545    }
546
547    #[doc(hidden)]
548    pub fn finalize_headers(&mut self) {
549        if self.version == Version::Http3 {
550            self.finalize_response_headers_h3();
551        } else {
552            self.finalize_response_headers_1x();
553        }
554    }
555
556    /// the [`H2Connection`] driver for this conn, if this is an HTTP/2 request
557    pub fn h2_connection(&self) -> Option<&Arc<H2Connection>> {
558        self.protocol_session.h2_connection()
559    }
560
561    /// the h2 stream id for this conn, if this is an HTTP/2 request
562    pub fn h2_stream_id(&self) -> Option<u32> {
563        self.protocol_session.h2_stream_id()
564    }
565
566    /// the [`H3Connection`] driver for this conn, if this is an HTTP/3 request
567    pub fn h3_connection(&self) -> Option<&Arc<H3Connection>> {
568        self.protocol_session.h3_connection()
569    }
570
571    /// the h3 stream id for this conn, if this is an HTTP/3 request
572    pub fn h3_stream_id(&self) -> Option<u64> {
573        self.protocol_session.h3_stream_id()
574    }
575}