mirror of
https://github.com/tokio-rs/axum.git
synced 2026-08-27 00:00:24 +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:
+10
-6
@@ -14,9 +14,9 @@ use futures::{sink::SinkExt, stream::StreamExt};
|
||||
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::prelude::*;
|
||||
use axum::response::Html;
|
||||
use axum::ws::{ws, Message, WebSocket};
|
||||
use axum::response::{Html, IntoResponse};
|
||||
use axum::AddExtensionLayer;
|
||||
|
||||
// Our shared state
|
||||
@@ -33,7 +33,7 @@ async fn main() {
|
||||
let app_state = Arc::new(AppState { user_set, tx });
|
||||
|
||||
let app = route("/", get(index))
|
||||
.route("/websocket", ws(websocket))
|
||||
.route("/websocket", get(websocket_handler))
|
||||
.layer(AddExtensionLayer::new(app_state));
|
||||
|
||||
let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
|
||||
@@ -44,10 +44,14 @@ async fn main() {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn websocket(
|
||||
stream: WebSocket,
|
||||
async fn websocket_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
extract::Extension(state): extract::Extension<Arc<AppState>>,
|
||||
) {
|
||||
) -> impl IntoResponse {
|
||||
ws.on_upgrade(|socket| websocket(socket, state))
|
||||
}
|
||||
|
||||
async fn websocket(stream: WebSocket, state: Arc<AppState>) {
|
||||
// By splitting we can send and receive at the same time.
|
||||
let (mut sender, mut receiver) = stream.split();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user