Skip to main content

h2ts_server/
lib.rs

1//! Make a WebSocket carry a raw byte stream, and serve/proxy HTTP/2 over it.
2//!
3//! [`accept`] performs the server-side WebSocket handshake on a hyper request
4//! (pluggable into any hyper/axum route) and yields the upgraded
5//! connection as a byte stream. [`bridge`] then pumps bytes full-duplex between
6//! that stream and any `AsyncRead + AsyncWrite` peer (a TCP upstream, an
7//! in-process h2c server, …). WebSocket message *payloads* become a continuous
8//! byte stream, so h2c framing rides straight through.
9//!
10//! Framing is done by [`wslay`](https://github.com/tatsuhiro-t/wslay) (vendored
11//! C, via the `wslay-sys` crate). Driven through its event API with buffering
12//! off, wslay streams each frame's payload **incrementally** — it never holds a
13//! whole frame in memory, no matter how large — and auto-handles ping/close.
14//!
15//! Three entry points sit on top of [`bridge`]:
16//! - [`WsByteStream`] — the WebSocket as an `AsyncRead + AsyncWrite` handle.
17//! - [`serve_h2`] — run any hyper `Service` as HTTP/2 over the tunnel.
18//! - the `h2ts-proxy` binary — a standalone WS→upstream-h2c proxy.
19//!
20//! Part of [h2ts](https://github.com/debdattabasu/h2ts): this server pairs with
21//! the TypeScript [`@debdattabasu/h2ts`](https://www.npmjs.com/package/@debdattabasu/h2ts)
22//! and Rust/WASM [`h2ts-client`](https://crates.io/crates/h2ts-client) frontend
23//! clients — both behavior-mirrored against it by the shared conformance suite —
24//! or any client that forwards a byte stream over a WebSocket.
25
26use std::io;
27use std::pin::Pin;
28use std::task::{Context, Poll};
29
30use hyper::body::{Body, Incoming};
31use hyper::server::conn::http2;
32use hyper::service::Service;
33use hyper::{Request, Response};
34use hyper_util::rt::{TokioExecutor, TokioIo};
35use tokio::io::{AsyncRead, AsyncWrite, DuplexStream, ReadBuf};
36
37mod handshake;
38mod wslay;
39
40pub use handshake::{
41    accept, accept_with, accept_with_options, is_upgrade_request, offered_protocols, AcceptOptions,
42    UpgradedIo, WebSocketError, DEFAULT_SUBPROTOCOL,
43};
44pub use wslay::{
45    bridge, bridge_with, control_channel, BridgeConfig, CloseFrame, CloseHook, ControlHook,
46    ControlReceiver, KeepAlive, WsControl,
47};
48
49/// Size of the in-memory duplex between the WebSocket pump and the app side.
50const DUPLEX_BUF: usize = 64 * 1024;
51
52/// A WebSocket presented as a raw byte duplex (`AsyncRead + AsyncWrite`) — item
53/// 1 in its "looks like a TCP stream" form.
54///
55/// A [`bridge`] runs on a spawned task pumping the WebSocket to one end of an
56/// in-memory [`tokio::io::duplex`]; this handle is the other end. That makes it
57/// usable anywhere a TCP stream is — most importantly, handed straight to an
58/// h2c server via `serve_connection` to run an in-process HTTP/2 service over
59/// the tunnel (item 2).
60pub struct WsByteStream {
61    inner: DuplexStream,
62}
63
64impl WsByteStream {
65    /// Wrap an upgraded WebSocket byte stream (from [`accept`]), spawning the
66    /// bridge pump onto the current tokio runtime.
67    pub fn new<S>(ws_io: S) -> Self
68    where
69        S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
70    {
71        Self::with_config(ws_io, BridgeConfig::default())
72    }
73
74    /// Like [`WsByteStream::new`], but with control-frame configuration and hooks
75    /// ([`BridgeConfig`]) — e.g. a [`control_channel`] to send control frames, or
76    /// `on_close` to observe the peer's close reason.
77    pub fn with_config<S>(ws_io: S, config: BridgeConfig) -> Self
78    where
79        S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
80    {
81        let (app_side, ws_side) = tokio::io::duplex(DUPLEX_BUF);
82        tokio::spawn(async move {
83            let _ = bridge_with(ws_io, ws_side, config).await;
84        });
85        Self { inner: app_side }
86    }
87}
88
89impl AsyncRead for WsByteStream {
90    fn poll_read(
91        self: Pin<&mut Self>,
92        cx: &mut Context<'_>,
93        buf: &mut ReadBuf<'_>,
94    ) -> Poll<io::Result<()>> {
95        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
96    }
97}
98
99impl AsyncWrite for WsByteStream {
100    fn poll_write(
101        self: Pin<&mut Self>,
102        cx: &mut Context<'_>,
103        buf: &[u8],
104    ) -> Poll<io::Result<usize>> {
105        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
106    }
107
108    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
109        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
110    }
111
112    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
113        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
114    }
115}
116
117/// Serve **any** hyper HTTP/2 service over a WebSocket tunnel (item 2).
118///
119/// `service` is any `hyper::service::Service<Request<Incoming>>` — a
120/// `service_fn`, an `axum::Router`, a `tower` service via hyper's compat, etc.
121/// The WebSocket is presented as a byte stream via [`WsByteStream`] and the
122/// request/response traffic is real HTTP/2 (h2c, prior-knowledge) on top.
123///
124/// The tunnel runs server-initiated WebSocket keepalive **on by default**
125/// ([`BridgeConfig::default`]) so a silently-dead client can't leak the
126/// connection; use [`serve_h2_with`] with a custom [`BridgeConfig`] to tune or
127/// disable it (`keepalive: None`).
128///
129/// Typical use in a route handler:
130/// ```ignore
131/// let (response, ws_fut) = h2ts_server::accept(&mut req)?;
132/// tokio::spawn(async move {
133///     if let Ok(ws) = ws_fut.await {
134///         let _ = h2ts_server::serve_h2(ws, my_service).await;
135///     }
136/// });
137/// Ok(response) // send the 101 back
138/// ```
139///
140/// For custom HTTP/2 settings, build the connection yourself over
141/// [`WsByteStream::new`] instead.
142pub async fn serve_h2<S, Svc, B>(ws_io: S, service: Svc) -> hyper::Result<()>
143where
144    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
145    Svc: Service<Request<Incoming>, Response = Response<B>> + Send + 'static,
146    Svc::Future: Send + 'static,
147    Svc::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
148    B: Body + Send + 'static,
149    B::Data: Send,
150    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
151{
152    serve_h2_with(ws_io, service, BridgeConfig::default()).await
153}
154
155/// Like [`serve_h2`], but with control-frame configuration and hooks
156/// ([`BridgeConfig`]) applied to the underlying WebSocket bridge — send control
157/// frames via a [`control_channel`], observe the peer's close reason, etc.
158pub async fn serve_h2_with<S, Svc, B>(
159    ws_io: S,
160    service: Svc,
161    config: BridgeConfig,
162) -> hyper::Result<()>
163where
164    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
165    Svc: Service<Request<Incoming>, Response = Response<B>> + Send + 'static,
166    Svc::Future: Send + 'static,
167    Svc::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
168    B: Body + Send + 'static,
169    B::Data: Send,
170    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
171{
172    let io = TokioIo::new(WsByteStream::with_config(ws_io, config));
173    http2::Builder::new(TokioExecutor::new())
174        .serve_connection(io, service)
175        .await
176}