Rust:axum學習筆記(7) websocket

接上一篇繼續，今天來學習下如何用axum實現websocket，代碼如下：

Cargo.toml添加依賴項

[package]
name = "websocket"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
axum =  {version = "0.4.3", features = ["headers","ws"] }
tokio = { version = "1.0", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version="0.3", features = ["env-filter"] }
tower-http = { version = "0.2.0", features = ["fs", "trace"] }
headers = "0.3"

關鍵是axum的features中的ws，接下來是示例代碼main.rs

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        TypedHeader,
    },
    response::IntoResponse,
    routing::{get},
    Router,
};
use std::net::SocketAddr;
use tower_http::{
    trace::{DefaultMakeSpan, TraceLayer},
};

#[tokio::main]
async fn main() {
    if std::env::var_os("RUST_LOG").is_none() {
        std::env::set_var("RUST_LOG", "example_websockets=debug,tower_http=debug")
    }
    tracing_subscriber::fmt::init();
    

    let app = Router::new()
        .route("/", get(|| async { "Hello, World!" }))
        //綁定websocket路由
        .route("/ws", get(ws_handler))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::default().include_headers(true)),
        );

    
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    tracing::debug!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

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);
            //客戶端發什么，服務端就回什么（只是演示而已）
            if socket
                .send(Message::Text(format!("{:?}", msg)))
                .await
                .is_err()
            {
                println!("client disconnected");
                return;
            }
        } else {
            println!("client disconnected");
            return;
        }
    }
}

核心就是handle_socket這個function，這里我們只是簡單的將收到的內容，原封不動的發回瀏覽器。

運行一下：

先瀏覽http://localhost:3000/ 然后F12打開Console控制台，輸入下面幾行js代碼：

socket = new WebSocket('ws://localhost:3000/ws');

socket.addEventListener('message', function (event) {
    console.log('Message from server ', event.data);
});

socket.send('你好，RUST!');

就能看到服務端回過來的內容