mirror of
https://github.com/tokio-rs/axum.git
synced 2026-08-25 00:00:23 +02:00
Change WebSocket API to use an extractor (#121)
Fixes https://github.com/tokio-rs/axum/issues/111 Example usage: ```rust use axum::{ prelude::*, extract::ws::{WebSocketUpgrade, WebSocket}, response::IntoResponse, }; let app = route("/ws", get(handler)); async fn handler(ws: WebSocketUpgrade) -> impl IntoResponse { ws.on_upgrade(handle_socket) } async fn handle_socket(mut socket: WebSocket) { while let Some(msg) = socket.recv().await { let msg = if let Ok(msg) = msg { msg } else { // client disconnected return; }; if socket.send(msg).await.is_err() { // client disconnected return; } } } ```
This commit is contained in:
+13
-7
@@ -7,11 +7,14 @@
|
||||
//! ```
|
||||
|
||||
use axum::{
|
||||
extract::TypedHeader,
|
||||
extract::{
|
||||
ws::{Message, WebSocket, WebSocketUpgrade},
|
||||
TypedHeader,
|
||||
},
|
||||
prelude::*,
|
||||
response::IntoResponse,
|
||||
routing::nest,
|
||||
service::ServiceExt,
|
||||
ws::{ws, Message, WebSocket},
|
||||
};
|
||||
use http::StatusCode;
|
||||
use std::net::SocketAddr;
|
||||
@@ -44,7 +47,7 @@ async fn main() {
|
||||
)
|
||||
// routes are matched from bottom to top, so we have to put `nest` at the
|
||||
// top since it matches all routes
|
||||
.route("/ws", ws(handle_socket))
|
||||
.route("/ws", get(ws_handler))
|
||||
// logging so we can see whats going on
|
||||
.layer(
|
||||
TraceLayer::new_for_http().make_span_with(DefaultMakeSpan::default().include_headers(true)),
|
||||
@@ -59,15 +62,18 @@ async fn main() {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn handle_socket(
|
||||
mut socket: WebSocket,
|
||||
// websocket handlers can also use extractors
|
||||
async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
user_agent: Option<TypedHeader<headers::UserAgent>>,
|
||||
) {
|
||||
) -> impl IntoResponse {
|
||||
if let Some(TypedHeader(user_agent)) = user_agent {
|
||||
println!("`{}` connected", user_agent.as_str());
|
||||
}
|
||||
|
||||
ws.on_upgrade(handle_socket)
|
||||
}
|
||||
|
||||
async fn handle_socket(mut socket: WebSocket) {
|
||||
if let Some(msg) = socket.recv().await {
|
||||
if let Ok(msg) = msg {
|
||||
println!("Client says: {:?}", msg);
|
||||
|
||||
Reference in New Issue
Block a user