Add support for WebSockets over HTTP/2 (#2894)

examples/websockets-http2/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "example-websockets-http2"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+axum = { path = "../../axum", features = ["ws", "http2"] }
+axum-server = { version = "0.6", features = ["tls-rustls"] }
+tokio = { version = "1", features = ["full"] }
+tower-http = { version = "0.5.0", features = ["fs"] }
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }

examples/websockets-http2/assets/index.html

@@ -0,0 +1,7 @@
+<p>Open this page in two windows and try sending some messages!</p>
+<form action="javascript:void(0)">
+    <input type="text" name="content" required>
+    <button>Send</button>
+</form>
+<div id="messages"></div>
+<script src='script.js'></script>

examples/websockets-http2/assets/script.js

@@ -0,0 +1,11 @@
+const socket = new WebSocket('wss://localhost:3000/ws');
+
+socket.addEventListener('message', e => {
+    document.getElementById("messages").append(e.data, document.createElement("br"));
+});
+
+const form = document.querySelector("form");
+form.addEventListener("submit", () => {
+    socket.send(form.elements.namedItem("content").value);
+    form.elements.namedItem("content").value = "";
+});

examples/websockets-http2/self_signed_certs/cert.pem

@@ -0,0 +1,22 @@
+-----BEGIN CERTIFICATE-----
+MIIDkzCCAnugAwIBAgIUXVYkRCrM/ge03DVymDtXCuybp7gwDQYJKoZIhvcNAQEL
+BQAwWTELMAkGA1UEBhMCVVMxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM
+GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDESMBAGA1UEAwwJbG9jYWxob3N0MB4X
+DTIxMDczMTE0MjIxMloXDTIyMDczMTE0MjIxMlowWTELMAkGA1UEBhMCVVMxEzAR
+BgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoMGEludGVybmV0IFdpZGdpdHMgUHR5
+IEx0ZDESMBAGA1UEAwwJbG9jYWxob3N0MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8A
+MIIBCgKCAQEA02V5ZjmqLB/VQwTarrz/35qsa83L+DbAoa0001+jVmmC+G9Nufi0
+daroFWj/Uicv2fZWETU8JoZKUrX4BK9og5cg5rln/CtBRWCUYIwRgY9R/CdBGPn4
+kp+XkSJaCw74ZIyLy/Zfux6h8ES1m9YRnBza+s7U+ImRBRf4MRPtXQ3/mqJxAZYq
+dOnKnvssRyD2qutgVTAxwMUvJWIivRhRYDj7WOpS4CEEeQxP1iH1/T5P7FdtTGdT
+bVBABCA8JhL96uFGPpOYHcM/7R5EIA3yZ5FNg931QzoDITjtXGtQ6y9/l/IYkWm6
+J67RWcN0IoTsZhz0WNU4gAeslVtJLofn8QIDAQABo1MwUTAdBgNVHQ4EFgQUzFnK
+NfS4LAYuKeWwHbzooER0yZ0wHwYDVR0jBBgwFoAUzFnKNfS4LAYuKeWwHbzooER0
+yZ0wDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAk4O+e9jia59W
+ZwetN4GU7OWcYhmOgSizRSs6u7mTfp62LDMt96WKU3THksOnZ44HnqWQxsSfdFVU
+XJD12tjvVU8Z4FWzQajcHeemUYiDze8EAh6TnxnUcOrU8IcwiKGxCWRY/908jnWg
++MMscfMCMYTRdeTPqD8fGzAlUCtmyzH6KLE3s4Oo/r5+NR+Uvrwpdvb7xe0MwwO9
+Q/zR4N8ep/HwHVEObcaBofE1ssZLksX7ZgCP9wMgXRWpNAtC5EWxMbxYjBfWFH24
+fDJlBMiGJWg8HHcxK7wQhFh+fuyNzE+xEWPsI9VL1zDftd9x8/QsOagyEOnY8Vxr
+AopvZ09uEQ==
+-----END CERTIFICATE-----

examples/websockets-http2/self_signed_certs/key.pem

@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDTZXlmOaosH9VD
+BNquvP/fmqxrzcv4NsChrTTTX6NWaYL4b025+LR1qugVaP9SJy/Z9lYRNTwmhkpS
+tfgEr2iDlyDmuWf8K0FFYJRgjBGBj1H8J0EY+fiSn5eRIloLDvhkjIvL9l+7HqHw
+RLWb1hGcHNr6ztT4iZEFF/gxE+1dDf+aonEBlip06cqe+yxHIPaq62BVMDHAxS8l
+YiK9GFFgOPtY6lLgIQR5DE/WIfX9Pk/sV21MZ1NtUEAEIDwmEv3q4UY+k5gdwz/t
+HkQgDfJnkU2D3fVDOgMhOO1ca1DrL3+X8hiRabonrtFZw3QihOxmHPRY1TiAB6yV
+W0kuh+fxAgMBAAECggEADltu8k1qTFLhJgsXWxTFAAe+PBgfCT2WuaRM2So+qqjB
+12Of0MieYPt5hbK63HaC3nfHgqWt7yPhulpXfOH45C8IcgMXl93MMg0MJr58leMI
++2ojFrIrerHSFm5R1TxwDEwrVm/mMowzDWFtQCc6zPJ8wNn5RuP48HKfTZ3/2fjw
+zEjSwPO2wFMfo1EJNTjlI303lFbdFBs67NaX6puh30M7Tn+gznHKyO5a7F57wkIt
+fkgnEy/sgMedQlwX7bRpUoD6f0fZzV8Qz4cHFywtYErczZJh3VGitJoO/VCIDdty
+RPXOAqVDd7EpP1UUehZlKVWZ0OZMEfRgKbRCel5abQKBgQDwgwrIQ5+BiZv6a0VT
+ETeXB+hRbvBinRykNo/RvLc3j1enRh9/zO/ShadZIXgOAiM1Jnr5Gp8KkNGca6K1
+myhtad7xYPODYzNXXp6T1OPgZxHZLIYzVUj6ypXeV64Te5ZiDaJ1D49czsq+PqsQ
+XRcgBJSNpFtDFiXWpjXWfx8PxwKBgQDhAnLY5Sl2eeQo+ud0MvjwftB/mN2qCzJY
+5AlQpRI4ThWxJgGPuHTR29zVa5iWNYuA5LWrC1y/wx+t5HKUwq+5kxvs+npYpDJD
+ZX/w0Glc6s0Jc/mFySkbw9B2LePedL7lRF5OiAyC6D106Sc9V2jlL4IflmOzt4CD
+ZTNbLtC6hwKBgHfIzBXxl/9sCcMuqdg1Ovp9dbcZCaATn7ApfHd5BccmHQGyav27
+k7XF2xMJGEHhzqcqAxUNrSgV+E9vTBomrHvRvrd5Ec7eGTPqbBA0d0nMC5eeFTh7
+wV0miH20LX6Gjt9G6yJiHYSbeV5G1+vOcTYBEft5X/qJjU7aePXbWh0BAoGBAJlV
+5tgCCuhvFloK6fHYzqZtdT6O+PfpW20SMXrgkvMF22h2YvgDFrDwqKRUB47NfHzg
+3yBpxNH1ccA5/w97QO8w3gX3h6qicpJVOAPusu6cIBACFZfjRv1hyszOZwvw+Soa
+Fj5kHkqTY1YpkREPYS9V2dIW1Wjic1SXgZDw7VM/AoGAP/cZ3ZHTSCDTFlItqy5C
+rIy2AiY0WJsx+K0qcvtosPOOwtnGjWHb1gdaVdfX/IRkSsX4PAOdnsyidNC5/l/m
+y8oa+5WEeGFclWFhr4dnTA766o8HrM2UjIgWWYBF2VKdptGnHxFeJWFUmeQC/xeW
+w37pCS7ykL+7gp7V0WShYsw=
+-----END PRIVATE KEY-----

examples/websockets-http2/src/main.rs

@@ -0,0 +1,97 @@
+//! Run with
+//!
+//! ```not_rust
+//! cargo run -p example-websockets-http2
+//! ```
+
+use axum::{
+    extract::{
+        ws::{self, WebSocketUpgrade},
+        State,
+    },
+    http::Version,
+    routing::any,
+    Router,
+};
+use axum_server::tls_rustls::RustlsConfig;
+use std::{net::SocketAddr, path::PathBuf};
+use tokio::sync::broadcast;
+use tower_http::services::ServeDir;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+#[tokio::main]
+async fn main() {
+    tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::EnvFilter::try_from_default_env()
+                .unwrap_or_else(|_| format!("{}=debug", env!("CARGO_CRATE_NAME")).into()),
+        )
+        .with(tracing_subscriber::fmt::layer())
+        .init();
+
+    let assets_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("assets");
+
+    // configure certificate and private key used by https
+    let config = RustlsConfig::from_pem_file(
+        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+            .join("self_signed_certs")
+            .join("cert.pem"),
+        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+            .join("self_signed_certs")
+            .join("key.pem"),
+    )
+    .await
+    .unwrap();
+
+    // build our application with some routes and a broadcast channel
+    let app = Router::new()
+        .fallback_service(ServeDir::new(assets_dir).append_index_html_on_directories(true))
+        .route("/ws", any(ws_handler))
+        .with_state(broadcast::channel::<String>(16).0);
+
+    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
+    tracing::debug!("listening on {}", addr);
+
+    let mut server = axum_server::bind_rustls(addr, config);
+
+    // IMPORTANT: This is required to advertise our support for HTTP/2 websockets to the client.
+    // If you use axum::serve, it is enabled by default.
+    server.http_builder().http2().enable_connect_protocol();
+
+    server.serve(app.into_make_service()).await.unwrap();
+}
+
+async fn ws_handler(
+    ws: WebSocketUpgrade,
+    version: Version,
+    State(sender): State<broadcast::Sender<String>>,
+) -> axum::response::Response {
+    tracing::debug!("accepted a WebSocket using {version:?}");
+    let mut receiver = sender.subscribe();
+    ws.on_upgrade(|mut ws| async move {
+        loop {
+            tokio::select! {
+                // Since `ws` is a `Stream`, it is by nature cancel-safe.
+                res = ws.recv() => {
+                    match res {
+                        Some(Ok(ws::Message::Text(s))) => {
+                            let _ = sender.send(s);
+                        }
+                        Some(Ok(_)) => {}
+                        Some(Err(e)) => tracing::debug!("client disconnected abruptly: {e}"),
+                        None => break,
+                    }
+                }
+                // Tokio guarantees that `broadcast::Receiver::recv` is cancel-safe.
+                res = receiver.recv() => {
+                    match res {
+                        Ok(msg) => if let Err(e) = ws.send(ws::Message::Text(msg)).await {
+                            tracing::debug!("client disconnected abruptly: {e}");
+                        }
+                        Err(_) => continue,
+                    }
+                }
+            }
+        }
+    })
+}

examples/websockets/src/main.rs

@@ -19,7 +19,7 @@
 use axum::{
     extract::ws::{Message, WebSocket, WebSocketUpgrade},
     response::IntoResponse,
-    routing::get,
+    routing::any,
     Router,
 };
 use axum_extra::TypedHeader;
@@ -57,7 +57,7 @@ async fn main() {
     // build our application with some routes
     let app = Router::new()
         .fallback_service(ServeDir::new(assets_dir).append_index_html_on_directories(true))
-        .route("/ws", get(ws_handler))
+        .route("/ws", any(ws_handler))
         // logging so we can see whats going on
         .layer(
             TraceLayer::new_for_http()
@@ -77,7 +77,7 @@ async fn main() {
     .unwrap();
 }
 
-/// The handler for the HTTP request (this gets called when the HTTP GET lands at the start
+/// The handler for the HTTP request (this gets called when the HTTP request lands at the start
 /// of websocket negotiation). After this completes, the actual switching from HTTP to
 /// websocket protocol will occur.
 /// This is the last point where we can extract TCP/IP metadata such as IP address of the client