Rust in Everything
基于 axum 0.8.9

WebSocketUpgrade

主线位置：请求主线的一条岔路。WebSocketUpgrade 是一个普通的 FromRequestParts 提取器，handler 用它返回 101（HTTP/2 上是 200）之后，这个请求在 axum 里就结束了，连接离开 hyper，交给 on_upgrade spawn 出来的任务。装配主线上只有一处相关：axum::serve 给每条连接的 HTTP/2 一侧打开 enable_connect_protocol。

sequenceDiagram
  participant C as 客户端
  participant H as hyper 连接
  participant X as 提取器
  participant F as handler
  participant T as 升级任务
  C->>H: GET，带 Upgrade 等头
  H->>X: 请求，extensions 里有 OnUpgrade
  X->>X: 查方法和头，取走 OnUpgrade
  X->>F: WebSocketUpgrade
  F->>T: on_upgrade：spawn
  F-->>H: 101 加 Sec-WebSocket-Accept
  H->>C: 写出 101
  H->>T: 交出 IO：Upgraded
  T->>T: 包成 WebSocketStream
  T->>C: 之后收发 WebSocket 帧

整个模块在一个文件里：axum/src/extract/ws.rs。它做三件事：提取时检查这是不是一个合法的 WebSocket 握手请求；on_upgrade 算出握手响应，同时安排一个任务去等 hyper 交出连接；拿到连接后包成 tokio-tungstenite 的 WebSocketStream。帧的编解码、Ping/Pong、关闭握手全是 tungstenite 做的，axum 只在外面套了一层自己的 Message 类型。

hyper 这一侧（什么请求会带 OnUpgrade、101 写完以后 IO 怎么交出来）在 hyper 卷的 协议升级 一章，这里不重复。

本页的例子用到的依赖：

axum = { version = "=0.8.9", features = ["ws"] }
tokio = { version = "=1.54.1", features = ["full"] }
tokio-tungstenite = "=0.29.0"
futures = "0.3"

tokio-tungstenite 写成 =0.29.0：axum 0.8.9 的 ws feature 依赖 tokio-tungstenite = "0.29.0"，Cargo 解析出来的就是 0.29.0（连带 tungstenite 0.29.0）。客户端用同一个版本，免得编出两份。平台都是 M2 Max、rustc 1.99.0、--release。

提取：查方法和四个头

WebSocketUpgrade 的 FromRequestParts 实现 按 HTTP 版本分两支。HTTP/1.x 这一支：

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        let sec_websocket_key = if parts.version <= Version::HTTP_11 {
            if parts.method != Method::GET {
                return Err(MethodNotGet.into());
            }

            if !header_contains(&parts.headers, header::CONNECTION, "upgrade") {
                return Err(InvalidConnectionHeader.into());
            }

            if !header_eq(&parts.headers, header::UPGRADE, "websocket") {
                return Err(InvalidUpgradeHeader.into());
            }

            Some(
                parts
                    .headers
                    .get(header::SEC_WEBSOCKET_KEY)
                    .ok_or(WebSocketKeyHeaderMissing)?
                    .clone(),
            )
        } else {
            // ... HTTP/2：见后面「HTTP/2：扩展 CONNECT」
        };

        if !header_eq(&parts.headers, header::SEC_WEBSOCKET_VERSION, "13") {
            return Err(InvalidWebSocketVersionHeader.into());
        }

        let on_upgrade = parts
            .extensions
            .remove::<hyper::upgrade::OnUpgrade>()
            .ok_or(ConnectionNotUpgradable)?;

        // ... 收集 Sec-WebSocket-Protocol

检查有先后，遇到第一个不满足的就返回。两个辅助函数的严格程度不一样：header_contains 把头转成小写后做子串查找（to_ascii_lowercase().contains(value)），所以 Connection: keep-alive, Upgrade 能过；header_eq 要求整个值忽略大小写后相等（eq_ignore_ascii_case），所以 Upgrade: WebSocket 能过，Upgrade: websocket, foo 不能过，尽管拒绝文本说的是「did not include」。两者都只看同名头的第一个值（headers.get）。Sec-WebSocket-Key 只查在不在，不查格式。

每个失败对应 rejection 模块里一个 define_rejection! 定义的类型，状态码和 body 文本写死在宏参数里：

检查拒绝类型状态码
HTTP/1 方法不是 GETMethodNotGet405
Connection 不含 upgradeInvalidConnectionHeader400
Upgrade 不等于 websocketInvalidUpgradeHeader400
没有 Sec-WebSocket-KeyWebSocketKeyHeaderMissing400
HTTP/2 方法不是 CONNECTMethodNotConnect405
HTTP/2 :protocol 不是 websocketInvalidProtocolPseudoheader400
Sec-WebSocket-Version 不等于 13InvalidWebSocketVersionHeader400
extensions 里没有 OnUpgradeConnectionNotUpgradable426

实验：手写握手，逐个去掉一个头

起一个真的 axum::serve，用裸 TcpStream 发手写的请求，把服务端回的字节原样打出来（date 行删掉了）。另外先用 oneshot 直接调一次 Router，不经过 hyper：

tower = { version = "=0.5.3", features = ["util"] }
http-body-util = "0.1"

use axum::{Router, body::Body, extract::ws::WebSocketUpgrade, http::Request, response::Response, routing::{any, get}};
use http_body_util::BodyExt;
use tower::ServiceExt;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

async fn handler(ws: WebSocketUpgrade) -> Response {
  ws.on_upgrade(|_socket| async {})
}

async fn proto(ws: WebSocketUpgrade) -> Response {
  // server preference order: superchat first
  ws.protocols(["superchat", "chat"]).on_upgrade(|_socket| async {})
}

// send raw bytes, print everything the server answers (until close or 300 ms of silence)
async fn send(addr: std::net::SocketAddr, label: &str, head: &str) {
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(head.as_bytes()).await.unwrap();
  let mut buf = Vec::new();
  let mut chunk = [0u8; 1024];
  loop {
    match tokio::time::timeout(Duration::from_millis(300), s.read(&mut chunk)).await {
      Ok(Ok(0)) | Err(_) => break,
      Ok(Ok(n)) => buf.extend_from_slice(&chunk[..n]),
      Ok(Err(e)) => { println!("read error {e}"); break; }
    }
  }
  let text = String::from_utf8_lossy(&buf);
  println!("--- {label}");
  for line in text.split("\r\n") {
    if line.starts_with("date:") { continue; }
    println!("  {line}");
  }
}

fn req(method: &str, version: &str, headers: &[&str]) -> String {
  let mut s = format!("{method} /ws {version}\r\nhost: x\r\n");
  for h in headers { s.push_str(h); s.push_str("\r\n"); }
  s.push_str("\r\n");
  s
}

#[tokio::main]
async fn main() {
  let app = Router::new().route("/ws", any(handler)).route("/get", get(handler)).route("/proto", any(proto));

  // no hyper connection at all: the request has no OnUpgrade in its extensions
  let request = Request::get("/ws")
    .header("connection", "upgrade").header("upgrade", "websocket")
    .header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==").header("sec-websocket-version", "13")
    .body(Body::empty()).unwrap();
  let res = app.clone().oneshot(request).await.unwrap();
  let status = res.status();
  let body = res.into_body().collect().await.unwrap().to_bytes();
  println!("--- oneshot\n  {status} {}", String::from_utf8_lossy(&body));
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  let conn = "connection: upgrade";
  let upg = "upgrade: websocket";
  let key = "sec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==";
  let ver = "sec-websocket-version: 13";

  send(addr, "POST", &req("POST", "HTTP/1.1", &[conn, upg, key, ver])).await;
  send(addr, "POST to get(handler)", &req("POST", "HTTP/1.1", &[conn, upg, key, ver]).replace("/ws", "/get")).await;
  send(addr, "no Connection", &req("GET", "HTTP/1.1", &[upg, key, ver])).await;
  send(addr, "no Upgrade", &req("GET", "HTTP/1.1", &[conn, key, ver])).await;
  send(addr, "Upgrade: websocket, foo", &req("GET", "HTTP/1.1", &[conn, "upgrade: websocket, foo", key, ver])).await;
  send(addr, "no Sec-WebSocket-Key", &req("GET", "HTTP/1.1", &[conn, upg, ver])).await;
  send(addr, "Sec-WebSocket-Version: 8", &req("GET", "HTTP/1.1", &[conn, upg, key, "sec-websocket-version: 8"])).await;
  send(addr, "no Key, no Version", &req("GET", "HTTP/1.1", &[conn, upg])).await;
  send(addr, "HTTP/1.0", &req("GET", "HTTP/1.0", &[conn, upg, key, ver])).await;
  send(addr, "ok, Connection: keep-alive, Upgrade + Upgrade: WebSocket", &req("GET", "HTTP/1.1", &["connection: keep-alive, Upgrade", "upgrade: WebSocket", key, ver])).await;
  send(addr, "ok", &req("GET", "HTTP/1.1", &[conn, upg, key, ver])).await;
  send(addr, "subprotocols: client offers chat, superchat", &req("GET", "HTTP/1.1", &[conn, upg, key, ver, "sec-websocket-protocol: chat, superchat"]).replace("/ws", "/proto")).await;
  send(addr, "subprotocols: client offers v2", &req("GET", "HTTP/1.1", &[conn, upg, key, ver, "sec-websocket-protocol: v2"]).replace("/ws", "/proto")).await;
}

连跑 3 次，输出相同。前半部分：

--- oneshot
  426 Upgrade Required WebSocket request couldn't be upgraded since no upgrade state was present
--- POST
  HTTP/1.1 405 Method Not Allowed
  content-type: text/plain; charset=utf-8
  content-length: 28
  
  Request method must be `GET`
--- POST to get(handler)
  HTTP/1.1 405 Method Not Allowed
  allow: GET,HEAD
  content-length: 0
  
  
--- no Connection
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 43
  
  Connection header did not include 'upgrade'
--- no Upgrade
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 44
  
  `Upgrade` header did not include 'websocket'
--- Upgrade: websocket, foo
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 44
  
  `Upgrade` header did not include 'websocket'
--- no Sec-WebSocket-Key
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 34
  
  `Sec-WebSocket-Key` header missing
--- Sec-WebSocket-Version: 8
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 51
  
  `Sec-WebSocket-Version` header did not include '13'
--- no Key, no Version
  HTTP/1.1 400 Bad Request
  content-type: text/plain; charset=utf-8
  content-length: 34
  
  `Sec-WebSocket-Key` header missing
--- HTTP/1.0
  HTTP/1.0 426 Upgrade Required
  content-type: text/plain; charset=utf-8
  content-length: 73
  
  WebSocket request couldn't be upgraded since no upgrade state was present

几点：

  • any 和 get 的区别。 路由用 any(handler) 时，POST 走到了提取器，回的是提取器的 405 和文本；用 get(handler) 时，POST 在 MethodRouter 那一层就被拦下，是带 Allow 的空 405（MethodRouter），提取器根本没跑。模块文档的例子用的是 any，原因在后面 HTTP/2 一节。
  • 同时缺两个头，只报第一个。 没有 Key 也没有 Version，报的是 Key：Key 的检查在 HTTP/1 那一支里，Version 的检查在两支之后。
  • HTTP/1.0 是 426。 四个头都对，前面的检查全过，卡在最后一步：extensions 里没有 OnUpgrade。hyper 只在 HTTP/1.1 请求带 Upgrade 头（或者方法是 CONNECT）时才放这个值（协议升级：起点），HTTP/1.0 不算。
  • oneshot 也是 426。 不经过 hyper，就没人往 extensions 里放 OnUpgrade。不起服务器测路由的办法测不了 WebSocket 握手的成功路径，只能测到这里，见 不起服务器测路由。
📝NOTE
没有蠢问题

问： hyper 已经认出了这是升级请求，axum 为什么还要自己再查一遍 Connection 和 Upgrade？

答： 两边问的不是同一个问题。hyper 只决定「要不要准备交出 IO」，判断依据是 HTTP/1.1 加上有 Upgrade 头，不看 Connection，也不管升级成什么协议；它不认识 WebSocket。上面「no Connection」那一行就是证据：请求带着 Upgrade: websocket，hyper 放了 OnUpgrade，是 axum 的检查回了 400。反过来，「no Upgrade」的请求 hyper 根本没放 OnUpgrade，但 axum 先查头、最后才取 OnUpgrade，所以报的是缺哪个头，而不是笼统的 426。

握手响应：Sec-WebSocket-Accept

同一次运行的后半部分，三个成功的握手：

--- ok, Connection: keep-alive, Upgrade + Upgrade: WebSocket
  HTTP/1.1 101 Switching Protocols
  connection: upgrade
  upgrade: websocket
  sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
  
  
--- ok
  HTTP/1.1 101 Switching Protocols
  connection: upgrade
  upgrade: websocket
  sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
  
  
--- subprotocols: client offers chat, superchat
  HTTP/1.1 101 Switching Protocols
  connection: upgrade
  upgrade: websocket
  sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
  sec-websocket-protocol: superchat
  
  
--- subprotocols: client offers v2
  HTTP/1.1 101 Switching Protocols
  connection: upgrade
  upgrade: websocket
  sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

（101 响应同样带 date 头，被删掉了。）dGhlIHNhbXBsZSBub25jZQ== 是 RFC 6455 里的示例 key，s3pPLMBiTxaQ9kYGzzhZRbK+xOo= 是 RFC 给出的对应值。算法在 sign：

fn sign(key: &[u8]) -> HeaderValue {
    use base64::engine::Engine as _;

    let mut sha1 = Sha1::default();
    sha1.update(key);
    sha1.update(&b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"[..]);
    let b64 = Bytes::from(base64::engine::general_purpose::STANDARD.encode(sha1.finalize()));
    HeaderValue::from_maybe_shared(b64).expect("base64 is a valid value")
}

key 原样拼上 RFC 规定的 GUID，SHA-1，再 base64。这就是 ws feature 要额外拉 sha1 和 base64 两个依赖的原因（axum/Cargo.toml：ws = ["dep:hyper", "tokio", "dep:tokio-tungstenite", "dep:sha1", "dep:base64"]）。axum 自己做握手，tungstenite 只在握手之后接手，所以后面用的是 WebSocketStream::from_raw_socket，不是 tungstenite 的 accept。

子协议：提取时把所有 Sec-WebSocket-Protocol 头按逗号拆开、去空白，存进一个 BTreeSet。protocols 按服务端给的顺序找第一个客户端也提供了的，所以客户端写 chat, superchat、服务端写 ["superchat", "chat"]，选中的是 superchat。一个都对不上时，axum 照样回 101，只是不带 Sec-WebSocket-Protocol（最后一段输出）；要不要因此拒绝连接，由 handler 自己决定（requested_protocols 和 selected_protocol 可以查）。

on_upgrade：先 spawn，再返回响应

on_upgrade 是整个模块的中心：

    #[must_use = "to set up the WebSocket connection, this response must be returned"]
    pub fn on_upgrade<C, Fut>(self, callback: C) -> Response
    where
        C: FnOnce(WebSocket) -> Fut + Send + 'static,
        Fut: Future<Output = ()> + Send + 'static,
        F: OnFailedUpgrade,
    {
        let on_upgrade = self.on_upgrade;
        let config = self.config;
        let on_failed_upgrade = self.on_failed_upgrade;

        let protocol = self.protocol.clone();

        tokio::spawn(async move {
            let upgraded = match on_upgrade.await {
                Ok(upgraded) => upgraded,
                Err(err) => {
                    on_failed_upgrade.call(Error::new(err));
                    return;
                }
            };
            let upgraded = TokioIo::new(upgraded);

            let socket =
                WebSocketStream::from_raw_socket(upgraded, protocol::Role::Server, Some(config))
                    .await;
            let socket = WebSocket {
                inner: socket,
                protocol,
            };
            callback(socket).await;
        });

        let mut response = if let Some(sec_websocket_key) = &self.sec_websocket_key {
            // If `sec_websocket_key` was `Some`, we are using HTTP/1.1.
            // ...
            Response::builder()
                .status(StatusCode::SWITCHING_PROTOCOLS)
                .header(header::CONNECTION, UPGRADE)
                .header(header::UPGRADE, WEBSOCKET)
                .header(
                    header::SEC_WEBSOCKET_ACCEPT,
                    sign(sec_websocket_key.as_bytes()),
                )
                .body(Body::empty())
                .unwrap()
        } else {
            // Otherwise, we are HTTP/2+. As established in RFC 9113 section 8.5, we just respond
            // with a 2XX with an empty body:
            // <https://datatracker.ietf.org/doc/html/rfc9113#name-the-connect-method>.
            Response::new(Body::empty())
        };
        // ... 选中了子协议就加 Sec-WebSocket-Protocol
        response
    }

顺序是先 tokio::spawn，再造响应。原因在 hyper：OnUpgrade 要等 101 写到 socket 上、hyper 的连接 future 把 IO 交出来才会完成，而这一切发生在 handler 返回之后。handler 不能自己 await 它（会等一个要靠自己返回才能发生的事件），只能把等待交给另一个任务，自己立刻返回响应。#[must_use] 的提示也是这个意思：on_upgrade 返回的 Response 必须真的交回去，任务才等得到结果。

任务里三步：等 OnUpgrade；用 hyper-util 的 TokioIo 把 hyper 的 Read/Write 转成 tokio 的 AsyncRead/AsyncWrite（tokio-tungstenite 要的是后者，见 hyper 卷 运行时抽象）；包成 WebSocketStream，交给你的回调。WebSocketStream::from_raw_socket 不做握手，Role::Server 决定按服务端规则处理掩码（客户端发来的帧必须带掩码，accept_unmasked_frames 可以放宽），config 就是提取时的 Default::default() 加上你用 builder 方法改过的字段（后面「消息大小上限」）。

这个任务就是一个普通的 tokio 任务，axum 不保存它的 JoinHandle。它活多久完全由你的回调决定，这一点在停机时会显出来（本页最后一节）。

等不到的时候：on_failed_upgrade

OnUpgrade 可能以错误结束，任务就调 on_failed_upgrade。默认的 DefaultOnFailedUpgrade 什么也不做（fn call(self, _error: Error) {}）。on_failed_upgrade 换掉它，类型参数 F 也跟着换，OnFailedUpgrade 对所有 FnOnce(Error) + Send + 'static 都有实现，传一个闭包即可。

📝NOTE
冷知识：升级失败以前会让后台任务 panic

#1539（首次发布于 axum-v0.6.0-rc.5，正式版 axum-v0.6.0）加了 on_failed_upgrade。PR 描述：「Previously if upgrading a connection to a WebSocket connection failed the background task would panic. There was no way to customize that so users that might wanna report the error was out of luck.」还有一条理由：「Panicking also wasn't great because users might abort on panics which would bring down the server.」于是 spawn 的任务里那个 match on_upgrade.await 的 Err 分支从 panic 变成了调回调，默认回调什么都不做。

更早一些，提取器本身也会 panic：请求里没有 OnUpgrade 时（issue #1132 报的是经过反向代理的请求），#1135 改成返回上表最后一行的 ConnectionNotUpgradable，PR 原话「Rather than panicking we'll return a rejection」。CHANGELOG 把它记在 0.5.13。

什么时候会等不到？下面的程序试了四种「握手没走完」的情况。/refuse 的 handler 调了 on_upgrade（任务已经 spawn），却返回了一个 403；/accept 是正常的。

use axum::{
  Router,
  extract::ws::{WebSocket, WebSocketUpgrade},
  http::StatusCode,
  response::{IntoResponse, Response},
  routing::any,
};
use std::{error::Error as _, sync::OnceLock, time::{Duration, Instant}};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

static START: OnceLock<Instant> = OnceLock::new();
fn ms() -> u128 { START.get().unwrap().elapsed().as_millis() }

fn on_failed(e: axum::Error) {
  println!("{:>5} ms server: on_failed_upgrade: {e} / source: {:?}", ms(),
    e.source().and_then(|s| s.source()).map(|s| s.to_string()));
}

async fn on_socket(mut socket: WebSocket) {
  println!("{:>5} ms server: callback started", ms());
  match tokio::time::timeout(Duration::from_secs(1), socket.recv()).await {
    Err(_) => {
      println!("{:>5} ms server: recv still pending after 1 s", ms());
      let r = socket.recv().await;
      println!("{:>5} ms server: recv -> {r:?}", ms());
    }
    Ok(r) => println!("{:>5} ms server: recv -> {r:?}", ms()),
  }
}

// the handler builds the upgrade (which spawns the task) but returns something else
async fn refuse(ws: WebSocketUpgrade) -> Response {
  let _upgrade = ws.on_failed_upgrade(on_failed).on_upgrade(on_socket);
  println!("{:>5} ms server: handler returns 403 instead", ms());
  (StatusCode::FORBIDDEN, "no").into_response()
}

async fn accept(ws: WebSocketUpgrade) -> Response {
  ws.on_failed_upgrade(on_failed).on_upgrade(on_socket)
}

const HS: &str = "connection: upgrade\r\nupgrade: websocket\r\n\
sec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==\r\nsec-websocket-version: 13\r\n\r\n";

async fn read_some(s: &mut TcpStream) -> String {
  let mut buf = [0u8; 1024];
  let n = s.read(&mut buf).await.unwrap();
  String::from_utf8_lossy(&buf[..n]).lines().next().unwrap_or("").to_string()
}

#[tokio::main]
async fn main() {
  START.set(Instant::now()).unwrap();
  let app = Router::new().route("/refuse", any(refuse)).route("/accept", any(accept));
  let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  println!("=== A: handler returns 403, client keeps the connection 300 ms");
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(format!("GET /refuse HTTP/1.1\r\nhost: x\r\n{HS}").as_bytes()).await.unwrap();
  println!("{:>5} ms client: {}", ms(), read_some(&mut s).await);
  tokio::time::sleep(Duration::from_millis(300)).await;
  println!("{:>5} ms client: closes", ms());
  drop(s);
  tokio::time::sleep(Duration::from_millis(100)).await;

  println!("=== A2: after the 403, a second handshake on the same connection");
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(format!("GET /refuse HTTP/1.1\r\nhost: x\r\n{HS}").as_bytes()).await.unwrap();
  println!("{:>5} ms client: {}", ms(), read_some(&mut s).await);
  s.write_all(format!("GET /accept HTTP/1.1\r\nhost: x\r\n{HS}").as_bytes()).await.unwrap();
  println!("{:>5} ms client: {}", ms(), read_some(&mut s).await);
  tokio::time::sleep(Duration::from_millis(100)).await;
  println!("{:>5} ms client: closes", ms());
  drop(s);
  tokio::time::sleep(Duration::from_millis(100)).await;

  println!("=== B: client sends the handshake and closes at once");
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(format!("GET /accept HTTP/1.1\r\nhost: x\r\n{HS}").as_bytes()).await.unwrap();
  drop(s);
  println!("{:>5} ms client: closed", ms());
  tokio::time::sleep(Duration::from_millis(200)).await;

  println!("=== C: client reads 101, then sends nothing for 1.5 s");
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(format!("GET /accept HTTP/1.1\r\nhost: x\r\n{HS}").as_bytes()).await.unwrap();
  println!("{:>5} ms client: {}", ms(), read_some(&mut s).await);
  tokio::time::sleep(Duration::from_millis(1500)).await;
  println!("{:>5} ms client: closes", ms());
  drop(s);
  tokio::time::sleep(Duration::from_millis(100)).await;
}
❗IMPORTANT
动脑筋：情况 A 里，handler 调了 on_upgrade 但返回 403，客户端读到 403 后过 300 ms 才断开。on_failed_upgrade 和 on_socket 两个回调，哪一个会被调用？什么时候？

M2 Max、rustc 1.99.0、--release。答案在本页最后。

A2、B、C 三段的输出（连跑 8 次，下面是第 1 次）：

=== A2: after the 403, a second handshake on the same connection
  407 ms server: handler returns 403 instead
  407 ms client: HTTP/1.1 403 Forbidden
  407 ms server: on_failed_upgrade: operation was canceled / source: Some("upgrade expected but not completed")
  407 ms client: HTTP/1.1 101 Switching Protocols
  407 ms server: callback started
  508 ms client: closes
  508 ms server: recv -> Some(Err(Error { inner: Protocol(ResetWithoutClosingHandshake) }))
=== B: client sends the handshake and closes at once
  612 ms client: closed
  612 ms server: callback started
  612 ms server: recv -> Some(Err(Error { inner: Protocol(ResetWithoutClosingHandshake) }))
=== C: client reads 101, then sends nothing for 1.5 s
  815 ms server: callback started
  815 ms client: HTTP/1.1 101 Switching Protocols
 1816 ms server: recv still pending after 1 s
 2319 ms client: closes
 2319 ms server: recv -> Some(Err(Error { inner: Protocol(ResetWithoutClosingHandshake) }))

8 次运行里事件的先后都一样，只有两处会变：A2 的最后两行（「101」和「callback started」）先后互换过；B 的 recv 结果 6 次是 ResetWithoutClosingHandshake，2 次是 Io(Os { code: 54, kind: ConnectionReset, ... })（客户端关闭时还有没读的 101 留在接收缓冲区，内核发的是 RST）。

  • A2：同一条连接上又来了一个升级请求，第一个任务才收到错误。 错误的 Display 是 operation was canceled，来源是 upgrade expected but not completed，正是 hyper 卷 OnUpgrade 四种结局 表里「Pending 没送任何东西就被 drop」那一行。第二个请求解析时，hyper 的 State::prepare_upgrade 用新的 Pending 覆盖了旧的，旧的被 drop。
  • B：客户端发完握手就走，升级仍然「成功」。 101 写进了内核缓冲区，hyper 照常交出 IO，回调照常开始，第一次 recv 才发现对端没了。on_failed_upgrade 管的是「hyper 没交出 IO」，不是「对端没完成 WebSocket 会话」。
  • C：握手完成、对端不说话，recv 就一直挂着。 axum 和 tungstenite 都没有读超时，也不会主动发 Ping。要发现死连接，得在回调里自己定时发 Ping、或者给 recv 套 timeout。

所以在实践里，on_failed_upgrade 被调用的机会很少：axum 自己的 serve 总是用 serve_connection_with_upgrades，不会出现 hyper 卷里那种「low level API in use」的错误；剩下的主要是 A2 这类「Pending 被丢掉」的情况。更实际的教训是：要拒绝，就在调用 on_upgrade 之前拒绝，不要先 on_upgrade 再返回别的响应。

升级之后：WebSocket 与 Message

回调拿到的 WebSocket 只有两个字段：

pub struct WebSocket {
    inner: WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>,
    protocol: Option<HeaderValue>,
}

它实现了 Stream<Item = Result<Message, Error>> 和 Sink<Message>，recv 就是 self.next().await，send 就是 SinkExt::send。读写要并发时用 StreamExt::split（模块文档的「Read and write concurrently」一节）。四个 Sink 方法和 poll_next 都是转发给 inner，再把 tungstenite 的类型换成 axum 的。

inner 的 IO 类型写死成 TokioIo<Upgraded>，axum 不提供拿回 TcpStream 的办法。如果你绕开 WebSocketUpgrade、在 handler 里自己调 hyper::upgrade::on 并想 downcast 回原来的 IO，要知道 axum::serve 用的是 hyper-util 的 auto::Builder，交给 hyper 的 IO 外面多了一层私有的 Rewind，得用 auto::upgrade::downcast，见 auto：升级。

Message 和 tungstenite 的 Message 几乎一一对应，少了一个变体。from_tungstenite 的最后一行：

            // we can ignore `Frame` frames as recommended by the tungstenite maintainers
            // https://github.com/snapview/tungstenite-rs/issues/268
            ts::Message::Frame(_) => None,

返回 None 时 poll_next 的 loop 接着读下一条，所以 recv 永远看不到原始帧。Text 里装的是 axum 的 Utf8Bytes（包着 tungstenite 的 Utf8Bytes），Binary、Ping、Pong 是 Bytes，这是 0.8.0 跟着 tokio-tungstenite 0.26 改的（CHANGELOG：「axum::extract::ws::Message now uses Bytes in place of Vec<u8>, and a new Utf8Bytes type in place of String」，#3078）。

实验：用 tokio-tungstenite 当客户端

服务端是 echo，/drain 收到 Close 后继续读，/break 收到 Close 立刻跳出循环：

use axum::{
  Router,
  extract::ws::{Message, WebSocket, WebSocketUpgrade},
  response::Response,
  routing::any,
};
use futures::{SinkExt, StreamExt};
use tokio_tungstenite::{connect_async, tungstenite::Message as TMessage};

async fn break_on_close(ws: WebSocketUpgrade) -> Response {
  ws.on_upgrade(|socket| echo(socket, true))
}

async fn drain(ws: WebSocketUpgrade) -> Response {
  ws.on_upgrade(|socket| echo(socket, false))
}

async fn echo(mut socket: WebSocket, break_on_close: bool) {
  while let Some(msg) = socket.recv().await {
    let msg = match msg {
      Ok(msg) => msg,
      Err(e) => { println!("server: recv error: {e}"); return; }
    };
    println!("server: got {msg:?}");
    match msg {
      Message::Text(_) | Message::Binary(_) => socket.send(msg).await.unwrap(),
      Message::Ping(_) | Message::Pong(_) => {} // tungstenite answers pings itself
      Message::Close(_) => if break_on_close { break },
    }
  }
  println!("server: recv loop ended");
}

#[tokio::main]
async fn main() {
  let app = Router::new()
    .route("/break", any(break_on_close))
    .route("/drain", any(drain));
  let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  for path in ["/drain", "/break"] {
  println!("=== {path}");
  let (mut ws, res) = connect_async(format!("ws://{addr}{path}")).await.unwrap();
  println!("client: {:?} {}", res.version(), res.status());
  for (k, v) in res.headers() {
    if k != "date" { println!("client:   {k}: {}", v.to_str().unwrap()); }
  }

  for msg in [
    TMessage::text("hello"),
    TMessage::binary(vec![1u8, 2, 3]),
    TMessage::Ping(b"p".to_vec().into()),
  ] {
    ws.send(msg).await.unwrap();
    println!("client: got {:?}", ws.next().await.unwrap().unwrap());
  }
  ws.close(None).await.unwrap();
  while let Some(msg) = ws.next().await {
    println!("client: after close {msg:?}");
  }
  tokio::time::sleep(std::time::Duration::from_millis(50)).await;
  }
}

连跑 3 次，除了 sec-websocket-accept（connect_async 每次随机生成 key）以外完全相同：

=== /drain
client: HTTP/1.1 101 Switching Protocols
client:   connection: upgrade
client:   upgrade: websocket
client:   sec-websocket-accept: ok5Du/jhH8tJFjxJmV0RvWWFW98=
server: got Text(Utf8Bytes(Utf8Bytes(b"hello")))
client: got Text(Utf8Bytes(b"hello"))
server: got Binary(b"\x01\x02\x03")
client: got Binary(b"\x01\x02\x03")
server: got Ping(b"p")
client: got Pong(b"p")
server: got Close(None)
server: recv loop ended
client: after close Ok(Close(None))
=== /break
client: HTTP/1.1 101 Switching Protocols
client:   connection: upgrade
client:   upgrade: websocket
client:   sec-websocket-accept: u2AD3EiG3qlqSQtPCthlBit096w=
server: got Text(Utf8Bytes(Utf8Bytes(b"hello")))
client: got Text(Utf8Bytes(b"hello"))
server: got Binary(b"\x01\x02\x03")
client: got Binary(b"\x01\x02\x03")
server: got Ping(b"p")
client: got Pong(b"p")
server: got Close(None)
server: recv loop ended
client: after close Err(Protocol(ResetWithoutClosingHandshake))

服务端的 Text(Utf8Bytes(Utf8Bytes(...))) 两层就是上面说的 axum 包 tungstenite。Ping 会交给你的代码看，但 Pong 不用你回。

两条路由只差在收到 Close 之后。tungstenite 收到对端的 Close 时（do_close），只是把回复的 Close 帧放进 additional_send，并不立刻写；它在下一次读或 flush 时才写出去。/drain 再调一次 recv，回复被写出，tungstenite 随即以 ConnectionClosed 结束流，recv 返回 None，客户端收到 Close(None)，关闭握手完整。/break 直接跳出循环、drop 了 WebSocket，排队的回复没机会写，客户端看到的是 ResetWithoutClosingHandshake。自动回 Pong 是同一个机制：只有你继续读，Pong 才会被写出去。

Message::Close 的文档写的是「After receiving a close frame, axum will automatically respond with a close frame if necessary (you do not have to deal with this yourself). Since no further messages will be received, you may either do nothing or explicitly drop the connection.」实验说明「自动回复」有前提：收到 Close 之后要再读一次（或者 flush），回调立即返回是不够的。

📝NOTE
冷知识：WebSocket::close 在 0.8.0 被删掉了

0.7 的 WebSocket 有一个 close(self) 方法。#2974（首次发布于 axum-v0.8.0-rc.1，正式版 axum-v0.8.0）把它删了，PR 描述：「The current .close() method is mostly useless, because it does not allow sending a close frame and takes ownership of the WebSocket, which even forbids implementing the normal WebSocket close handshake.」同一个 PR 给 Message::Close 补了上面引的那段关闭握手文档。CHANGELOG 的说法是「Remove WebSocket::close. Users should explicitly send close messages themselves.」：要主动关闭，就 send(Message::Close(...))，然后继续 recv 直到 None，等对端的回复。

消息大小上限

提取时 config: Default::default()，axum 不改任何默认值，用的就是 tungstenite 0.29.0 的 WebSocketConfig::default()：

impl Default for WebSocketConfig {
    fn default() -> Self {
        Self {
            read_buffer_size: 128 * 1024,
            write_buffer_size: 128 * 1024,
            max_write_buffer_size: usize::MAX,
            max_message_size: Some(64 << 20),
            max_frame_size: Some(16 << 20),
            accept_unmasked_frames: false,
        }
    }
}

WebSocketUpgrade 为每个字段都有一个同名 builder 方法（max_message_size 等），在 on_upgrade 之前调用。max_message_size 和 max_frame_size 收的是 usize，内部写成 Some(max)，所以通过 axum 只能改上限，不能设成 tungstenite 的 None（不限）。

字段默认值管什么
max_message_size64 MiB一条消息（可能由多帧拼成）的总长度
max_frame_size16 MiB单帧 payload，读到帧头就检查，不等 payload
read_buffer_size128 KiB读缓冲，每条连接一开始就分配
write_buffer_size128 KiB攒够这么多才写 socket；flush 总会写完
max_write_buffer_sizeusize::MAX写缓冲上限，tungstenite 文档说默认不限
accept_unmasked_framesfalse是否接受客户端不带掩码的帧

WebSocketUpgrade 的 Debug 会打印 config，可以直接看。下面的程序有三条路由：默认配置、max_message_size(10)、max_frame_size(4)：

async fn report(mut socket: WebSocket) {
  while let Some(msg) = socket.recv().await {
    match msg {
      Ok(Message::Binary(b)) => println!("server: got Binary, {} bytes", b.len()),
      Ok(other) => println!("server: got {other:?}"),
      Err(e) => { println!("server: recv error: {e}"); return; }
    }
  }
}

async fn default_cfg(ws: WebSocketUpgrade) -> Response {
  println!("server: {ws:?}");
  ws.on_upgrade(report)
}

async fn small_message(ws: WebSocketUpgrade) -> Response {
  ws.max_message_size(10).on_upgrade(report)
}

async fn small_frame(ws: WebSocketUpgrade) -> Response {
  ws.max_frame_size(4).on_upgrade(report)
}

// main: same server setup as before, then for each route
//   connect_async, send one Binary of the given length, close(None), print what comes back
  for (path, len) in [("/default", 3 * 1024 * 1024), ("/small-message", 11), ("/small-frame", 5)] {
    println!("=== {path}, client sends {len} bytes");
    let (mut ws, _) = connect_async(format!("ws://{addr}{path}")).await.unwrap();
    ws.send(TMessage::binary(vec![b'x'; len])).await.unwrap();
    ws.close(None).await.ok();
    while let Some(msg) = ws.next().await {
      println!("client: {msg:?}");
    }
    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
  }

=== /default, client sends 3145728 bytes
server: WebSocketUpgrade { config: WebSocketConfig { read_buffer_size: 131072, write_buffer_size: 131072, max_write_buffer_size: 18446744073709551615, max_message_size: Some(67108864), max_frame_size: Some(16777216), accept_unmasked_frames: false }, protocol: None, sec_websocket_key: Some("nOqSlNiXZx6ntbtcnXsT3w=="), sec_websocket_protocol: {}, .. }
server: got Binary, 3145728 bytes
server: got Close(None)
client: Ok(Close(None))
=== /small-message, client sends 11 bytes
server: recv error: Space limit exceeded: Message too long: 11 > 10
client: Err(Io(Os { code: 54, kind: ConnectionReset, message: "Connection reset by peer" }))
=== /small-frame, client sends 5 bytes
server: recv error: Space limit exceeded: Message too long: 5 > 4
client: Err(Io(Os { code: 54, kind: ConnectionReset, message: "Connection reset by peer" }))

跑了 10 次，服务端的行除 key 外都相同；客户端那两行有 2 次 /small-message 是 Err(Protocol(ResetWithoutClosingHandshake))，其余都是 ConnectionReset。两种上限报的是同一个错误（tungstenite 的 CapacityError::MessageTooLong），文本里的 Message too long 不区分是帧还是整条消息。超限之后 tungstenite 不会自动回一个 1009（close_code::SIZE）Close 帧，回调返回、连接被 drop，客户端只看到连接断了。想让客户端知道原因，要在 Err 分支里自己 send(Message::Close(...))。

📝NOTE
没有蠢问题

问： axum 的 body 默认上限是 2 MB，上面却收下了一条 3 MiB 的消息。DefaultBodyLimit 对 WebSocket 不管用吗？

答： 不管用，它根本碰不到这些字节。DefaultBodyLimit 限的是请求 body（body 提取器），握手请求没有 body；101 之后的字节由 tungstenite 直接从 Upgraded 里读，不经过任何 http_body::Body，也不经过 tower 中间件。WebSocket 一侧唯一的上限是上表的 max_message_size / max_frame_size，默认 64 MiB / 16 MiB，比 HTTP 那边宽松得多。对外暴露的 WebSocket 端点通常应该按业务把它们调小。

HTTP/2：扩展 CONNECT

HTTP/2 没有 Upgrade 头，也没有 101。RFC 8441 的做法是：服务端在 SETTINGS 里宣布 SETTINGS_ENABLE_CONNECT_PROTOCOL，客户端发一个 CONNECT 请求，带伪头 :protocol: websocket，服务端回 2xx，这个流就成了 WebSocket 的通道。一条 HTTP/2 连接上可以同时有多个这样的流，也可以和普通请求混在一起。

三处代码配合。第一处在 axum::serve 的 handle_connection，每条连接建 Builder 时：

        let mut builder = Builder::new(TokioExecutor::new());
        // CONNECT protocol needed for HTTP/2 websockets
        #[cfg(feature = "http2")]
        builder.http2().enable_connect_protocol();

（hyper 的 enable_connect_protocol。）第二处在 hyper 的 HTTP/2 服务端：收到 CONNECT 时建一对 Pending / OnUpgrade，把 OnUpgrade 放进请求的 extensions，再把 h2 的 :protocol 换成 hyper::ext::Protocol 也放进去。第三处是提取器的 HTTP/2 分支：

        } else {
            if parts.method != Method::CONNECT {
                return Err(MethodNotConnect.into());
            }

            // if this feature flag is disabled, we won’t be receiving an HTTP/2 request to begin
            // with.
            #[cfg(feature = "http2")]
            if parts
                .extensions
                .get::<hyper::ext::Protocol>()
                .map_or(true, |p| p.as_str() != "websocket")
            {
                return Err(InvalidProtocolPseudoheader.into());
            }

            None
        };

这一支返回 None 作为 sec_websocket_key，on_upgrade 据此走 else 分支：回一个空 body 的 200，没有 Sec-WebSocket-Accept（RFC 8441 不用 key/accept 那一套）。Sec-WebSocket-Version: 13 两个版本都要查。还有一处小配合在路由层：RouteFuture::poll 对 CONNECT 的 2xx 响应不加 Content-Length，响应有 body 就换成空 body 并记一条 error 日志，因为 RFC 9110 不允许这两个头出现在 CONNECT 的成功响应里。

实验：hyper 的 HTTP/2 客户端发扩展 CONNECT

axum = { version = "=0.8.9", features = ["ws", "http2"] }
hyper = { version = "=1.12.0", features = ["client", "http2"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
http-body-util = "0.1"

use axum::{
  Router,
  body::Body,
  extract::ws::{Message, WebSocket, WebSocketUpgrade},
  http::{Method, Request},
  response::Response,
  routing::{any, get},
};
use futures::{SinkExt, StreamExt};
use http_body_util::BodyExt;
use hyper_util::rt::{TokioExecutor, TokioIo};
use tokio_tungstenite::{WebSocketStream, tungstenite::{Message as TMessage, protocol::Role}};

async fn handler(ws: WebSocketUpgrade) -> Response {
  ws.on_upgrade(|mut socket: WebSocket| async move {
    while let Some(Ok(msg)) = socket.recv().await {
      if let Message::Text(_) = msg {
        socket.send(msg).await.unwrap();
      }
    }
  })
}

#[tokio::main]
async fn main() {
  let app = Router::new().route("/ws", any(handler)).route("/get", get(handler));
  let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  let io = TokioIo::new(tokio::net::TcpStream::connect(addr).await.unwrap());
  let (mut sender, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new())
    .handshake(io)
    .await
    .unwrap();
  println!("right after handshake: extended CONNECT enabled = {}", conn.is_extended_connect_protocol_enabled());
  let mut conn = Box::pin(conn);
  // drive the connection until the server's SETTINGS frame has been processed
  for _ in 0..10 {
    tokio::select! {
      _ = &mut conn => unreachable!(),
      _ = tokio::task::yield_now() => {}
    }
    tokio::time::sleep(std::time::Duration::from_millis(5)).await;
  }
  println!("after server SETTINGS:  extended CONNECT enabled = {}", conn.is_extended_connect_protocol_enabled());
  tokio::spawn(conn);

  // 1) GET with the HTTP/1.1 handshake headers, over HTTP/2
  let req = Request::builder()
    .method(Method::GET)
    .uri(format!("http://{addr}/ws"))
    .header("connection", "upgrade")
    .header("upgrade", "websocket")
    .header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==")
    .header("sec-websocket-version", "13")
    .body(Body::empty())
    .unwrap();
  match sender.send_request(req).await {
    Ok(res) => {
      let status = res.status();
      let body = res.into_body().collect().await.unwrap().to_bytes();
      println!("GET over h2:          {status} {}", String::from_utf8_lossy(&body));
    }
    Err(e) => println!("GET over h2:          client error: {e}"),
  }

  // 2) CONNECT with :protocol = websocket (RFC 8441)
  let req = Request::builder()
    .method(Method::CONNECT)
    .uri(format!("http://{addr}/ws"))
    .extension(hyper::ext::Protocol::from_static("websocket"))
    .header("sec-websocket-version", "13")
    .body(Body::empty())
    .unwrap();
  let mut res = sender.send_request(req).await.unwrap();
  println!("CONNECT :protocol:    {:?} {}", res.version(), res.status());
  for (k, v) in res.headers() {
    if k != "date" { println!("  {k}: {}", v.to_str().unwrap()); }
  }
  let upgraded = hyper::upgrade::on(&mut res).await.unwrap();
  let mut ws = WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Client, None).await;
  ws.send(TMessage::text("over h2")).await.unwrap();
  println!("echo:                 {:?}", ws.next().await.unwrap().unwrap());

  // 3) the same connection still serves other streams
  let req = Request::builder()
    .method(Method::CONNECT)
    .uri(format!("http://{addr}/ws"))
    .extension(hyper::ext::Protocol::from_static("chat"))
    .header("sec-websocket-version", "13")
    .body(Body::empty())
    .unwrap();
  let res = sender.send_request(req).await.unwrap();
  let status = res.status();
  let body = res.into_body().collect().await.unwrap().to_bytes();
  println!("CONNECT :protocol=chat {status} {}", String::from_utf8_lossy(&body));

  // 4) the same extended CONNECT against a route registered with get()
  let req = Request::builder()
    .method(Method::CONNECT)
    .uri(format!("http://{addr}/get"))
    .extension(hyper::ext::Protocol::from_static("websocket"))
    .header("sec-websocket-version", "13")
    .body(Body::empty())
    .unwrap();
  let res = sender.send_request(req).await.unwrap();
  println!("CONNECT to get():     {} allow: {:?}", res.status(), res.headers().get("allow"));
}

连跑 3 次，输出相同：

right after handshake: extended CONNECT enabled = false
after server SETTINGS:  extended CONNECT enabled = true
GET over h2:          405 Method Not Allowed Request method must be `CONNECT`
CONNECT :protocol:    HTTP/2.0 200 OK
echo:                 Text(Utf8Bytes(b"over h2"))
CONNECT :protocol=chat 400 Bad Request `:protocol` pseudo-header did not include 'websocket'
CONNECT to get():     405 Method Not Allowed allow: Some("GET,HEAD")
  • 握手刚完成时客户端还没收到服务端的 SETTINGS，标志是 false；收到后变成 true，这就是 enable_connect_protocol 的效果。
  • 成功的响应是 200，除了 date 没有别的头。之后的 echo 用的是同一个 WebSocketStream，只是底下的 IO 换成了一个 HTTP/2 流。
  • 最后一行是 get 的代价：CONNECT 在 MethodRouter 就被 405 挡掉。用 get(ws_handler) 注册的端点只能走 HTTP/1.1。
📝NOTE
冷知识：HTTP/2 上的 WebSocket 是 0.8.0 才有的，代价是 get 要改成 any

#2894（首次发布于 axum-v0.8.0-rc.1，正式版 axum-v0.8.0）加的。PR 描述：「WebSockets over HTTP/2 are currently not supported.」做法只有两件事：「looking at the :protocol pseudo-header in the websocket extractor and advertising our support by enabling SETTINGS_ENABLE_CONNECT_PROTOCOL in HTTP/2」，也就是上面的提取器分支和 serve 里那一行。PR 还顺带修了一个问题：「Axum was incorrectly adding the Content-Length header to them」，即 RouteFuture 里对 CONNECT 2xx 的特殊处理。CHANGELOG 写明了用户要做的改动：「They can be enabled by changing get(ws_endpoint) handlers to any(ws_endpoint)」。这就是模块文档的例子都写 any(handler) 的原因，也是提取器里 HTTP/1 那一支还要自己查 GET 的原因：路由不再替它挡。

停机：升级后的连接不归 serve 管

101 写出之后，hyper 的连接 future 把 IO 交给 OnUpgrade 就以 Ok(()) 结束（协议升级：谁来交出 IO）。axum::serve 的停机靠的是每个连接任务手里的 watch 接收端，连接 future 一结束，这条连接就不在名单上了。on_upgrade spawn 的任务 axum 不跟踪，所以 with_graceful_shutdown 既不会通知 WebSocket 回调，也不会等它：停机「顺利」完成，进程退出时这些连接被直接切断，对端收不到 Close 帧。复现、原因和修法（让回调自己听停机信号、发 Close、用 TaskTracker 或 JoinSet 等它们结束）在 hyper 卷的 优雅停机一直不结束，那一页的「复现二」用的就是 axum 0.8.9。

动脑筋答案

情况 A 的输出（同一次运行，8 次里事件先后都相同，毫秒数差 1～2）：

=== A: handler returns 403, client keeps the connection 300 ms
    1 ms server: handler returns 403 instead
    0 ms client: HTTP/1.1 403 Forbidden
  303 ms client: closes
  304 ms server: callback started
  304 ms server: recv -> Some(Err(Error { inner: Protocol(ResetWithoutClosingHandshake) }))

被调用的是 on_socket，不是 on_failed_upgrade，时间是客户端断开的那一刻（303～304 ms）。 服务端从没发过 101，升级却「成功」了。

原因在 hyper。请求带着 Upgrade 头，hyper 解析时就调了 prepare_upgrade，Pending 留在连接的 State 里；响应是 403 还是 101，hyper 都不会清掉它（State.upgrade 只在 pending_upgrade 里被 take，或者被下一次 prepare_upgrade 覆盖）。403 写完后连接按 keep-alive 回到等待下一个请求的状态，on_socket 的任务一直挂在 OnUpgrade 上。客户端一断开，读写两端都关，poll_inner 看到 is_done()，又发现还有一个 Pending，就返回 Dispatched::Upgrade，UpgradeableConnection 把（已经被对端关掉的）IO 交了出去。于是回调开始，第一次 recv 就读到连接已断。

A2 和它的区别只在于：同一条连接上来了第二个升级请求，新的 Pending 覆盖了旧的，旧的被 drop，第一个任务才走到 on_failed_upgrade。两种结局都说明同一件事：调用 on_upgrade 就是承诺返回它的响应；要拒绝，在调用之前拒绝，此时 WebSocketUpgrade 连同里面的 OnUpgrade 一起被 drop，不会留下任何任务。

小结

  • WebSocketUpgrade 提取时按顺序查：HTTP/1 要 GET、Connection 含 upgrade（子串）、Upgrade 等于 websocket（整值）、有 Sec-WebSocket-Key；HTTP/2 要 CONNECT 加 :protocol: websocket；两者都要 Sec-WebSocket-Version: 13，最后取走 hyper 放的 OnUpgrade，没有就 426（HTTP/1.0、oneshot 都是这样）。
  • on_upgrade 先 tokio::spawn 一个等 OnUpgrade 的任务，再返回 101（Sec-WebSocket-Accept = base64(SHA-1(key + GUID))）或 HTTP/2 的空 200。任务拿到 Upgraded 后用 TokioIo 适配，from_raw_socket 包成 WebSocketStream。
  • on_failed_upgrade 只在 hyper 没交出 IO 时调用，默认什么都不做；客户端握手后立刻断开、或者一直不说话，都不算失败，要在回调里自己处理。先 on_upgrade 再返回别的响应，任务会在连接关闭时拿到一条死连接。
  • Message 去掉了 tungstenite 的 Frame 变体；Ping 自动回 Pong、对端 Close 自动回 Close，都要等下一次读才写出去，收到 Close 就 drop 会让对端看到 ResetWithoutClosingHandshake。
  • 默认 max_message_size 64 MiB、max_frame_size 16 MiB（tungstenite 0.29.0 的默认值，axum 不改），DefaultBodyLimit 不管 WebSocket；超限只是 recv 返回错误，不会自动发 1009。
  • HTTP/2 需要 http2 feature（serve 才会 enable_connect_protocol），并且路由要用 any。升级后的连接不在 serve 的停机名单上。

💬 关联讨论

发起讨论
加载中...