简单ECHO服务

本章我们将使用 AXUM 和 Websocket 实现一个简单的 Echo 服务。

所谓 Echo 服务，就是无论客户端发送什么消息，服务端总是将该消息原样返回给客户端。

依赖

[dependencies]
tokio = { version = "1", features = ["full"] }
axum = { version = "0.7", features = ["ws"] }

为了使用 Websocket，axum 要启用 ws feature

源码解析

main() 函数

#[tokio::main]
async fn main() {
    // 创建TCP监听器
    let listener = TcpListener::bind("0.0.0.0:56789").await.unwrap();
    // 定义路由
    let app = Router::new().route("/ws", get(websocket_handler));
    // 打印日志
    println!("监听于 {}", listener.local_addr().unwrap());
    // 启动 axum 服务
    axum::serve(listener, app.into_make_service())
        .await
        .unwrap();
}

• 使用 tokio::io::TcpListener::bind() 来事先创建好一个TCP监听器
• 然后使用 axum::serve() 函数来启动 axum 服务

我们是这样定义路由的：

let app = Router::new().route("/ws", get(websocket_handler));

• URL 路径是 /ws
• 使用的是 GET 方法

我们来看看这个 websocket_handler 函数。

websocket_handler() 函数

async fn websocket_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_socket)
}

• 参数：ws: WebSocketUpgrade，WebSocketUpgrade是 axum 提供的一个 extract，用于将 HTTP 提升为 Websocket，建立 Webscoket 连接。注意，它只接受 GET 方法。
• 返回值：impl IntoResponse。我们在《漫游AXUM之各种响应》一文中讨论过这个 trait。就是说，这个函数的返回值可以是所有可以作为 axum 响应的数据类型。
• 函数体：很好，函数体只有一条语句：ws.on_upgrade(handle_socket)
  • on_upgrade()方法，用于完成协议的提升（HTTP 提升为 WEBSOCKET），并将Websocket消息流，传递给回调函数
  • 它的默认回调函数接收一个数据类型为WebSocket参数
    • 由于我们的回调函数 handle_socket的签名刚好和默认回调函数一致，所以可以简写成 ws.on_upgrade(handle_socket)
    • 完整的写法是：ws.on_upgrade(move |socket| handle_socket(socket))。当我们需要传递多个参数时，就需要用这种写法。后续章节有该写法的实例。

handle_socket() 函数

async fn handle_socket(mut socket: WebSocket) {
    while let Some(Ok(msg)) = socket.recv().await {
        match msg {
            Message::Close(_) => {
                println!("客户端断开连接");
                break;
            }
            Message::Text(text) => {
                println!("收到客户端文本消息：{}", text);
                // 向客户端原样发送收到的消息
                socket.send(Message::Text(text)).await.unwrap();
            }
            _ => println!("收到客户端消息：{:?}", msg),
        };
    }
}

流程如下：

• 通过 while let 来接收客户端发送过来的消息。
• Message枚举定义了多种消息，我们只处理其中的：
  • Close()：客户端断开连接，使用 break 退出循环
  • Text()：文本信息，我们接收到文本信息之后，再通过 send() 方法原封不动的将其发回给客户端 socket.send(Message::Text(text)).await.unwrap();
  • 其它消息：本案例不做处理，只是简单的打印一下
• 这个函数结束，意味着与客户端之间的连接断开

测试

你可以用任何语言写一个 websocket 客户端来与我们这个简单ECHO服务进行通信，更为简单的是使用现成的工具，比如 Postman、在线测试工具等。这里我们使用WebSocket King，测试结果截图在本文开篇已经给出。

本章代码在01/simple-echo分支。