Rust基于tokio的websocket群聊

最新推荐文章于 2024-09-13 07:26:15 发布

昊天0

最新推荐文章于 2024-09-13 07:26:15 发布

阅读量311

点赞数 3

文章标签： rust websocket 开发语言

本文链接：https://blog.csdn.net/weixin_43343535/article/details/141891241

版权

原文:Rust基于tokio的websocket群聊

依赖

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

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

[dependencies]
tokio = { version = "1.23",features = ["full"] }
tokio-tungstenite = "*"
futures = "0.3"
serde={version = "1.0",features = ["derive"]}
serde_json = "1.0"
rand = "0.8.5"
tungstenite = "0.19.0"
futures-util = "0.3.30"

定义消息结构

use serde::{Deserialize, Serialize};

#[derive(Deserialize, Serialize, Debug)]
pub struct Message {
    pub username: String,
    pub message: String,
}

后台程序

mod message;

use crate::message::Message;
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::net::{TcpListener, TcpStream};
use tungstenite::protocol::Message as WsMessage;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

#[tokio::main]
async fn main() -> Result<()> {
    let addr = "127.0.0.1:8080";
    let listener = TcpListener::bind(&addr).await?;
    println!("服务端口 :{}", addr);

    //使用HashMap存储连接的客户端
    let clients: Arc<Mutex<HashMap<usize, mpsc::Sender<WsMessage>>>> =
        Arc::new(Mutex::new(HashMap::new()));
    while let Ok((stream, _)) = listener.accept().await {
        // 为每一个连接生成一个唯一的id
        let client_id = rand::random::<usize>();
        // 创建channels用于在服务器内部发送消息
        let (tx, rx) = mpsc::channel(0);
        clients.lock().unwrap().insert(client_id, tx);

        // 处理新的WebSocket连接
        let clients_clone = clients.clone();
        tokio::spawn(async move {
            if let Err(e) = handle_connection(client_id, stream, clients_clone, rx).await {
                println!("Error handling connection: {}", e);
            }
        });
    }

    Ok(())
}

async fn handle_connection(
    client_id: usize,
    stream: TcpStream,
    clients: Arc<Mutex<HashMap<usize, mpsc::Sender<WsMessage>>>>,
    mut rx: mpsc::Receiver<WsMessage>,
) -> Result<()> {
    let addr = stream.peer_addr()?;
    println!("新连接 :{}(ID:{})", addr, client_id);

    let ws_stream = tokio_tungstenite::accept_async(stream).await?;
    println!("WebSocket handshake has been successfully completed");

    //处理来自客户端的消息
    let (mut write, mut read) = ws_stream.split();

    loop {
        tokio::select! {
            msg = read.next()=>{
                match msg {
                    Some(Ok(message)) => {
                     if message.is_text(){
                            let message = message.to_text()?;
                            let message:Message = serde_json::from_str(message)?;
                            let broadcast_message = format!("{}:{}",message.username,message.message);
                            broadcast_message_to_all_clients(&clients,&broadcast_message).await?;
                        }
                    },
                    _ => {}
                }
            }
            // 处理来自服务器内部的消息
            msg = rx.next() =>{
                match msg{
                    Some(msg)=>{
                        if let Err(e) = write.send(msg).await{
                            println!("Error sending message: {}",e);
                            break;
                        }
                    }
                    None=>{
                        println!("Channel closed");
                        break;
                    }
                }
            }
        }
    }

    // 从客户端列表移除断开的客户端
    clients.lock().unwrap().remove(&client_id);
    Ok(())
}

/// 将消息广播给所有连接的客户端
async fn broadcast_message_to_all_clients(
    clients: &Arc<Mutex<HashMap<usize, mpsc::Sender<WsMessage>>>>,
    message: &str,
) -> Result<()> {
    let mut clients = clients.lock().unwrap();
    for (_client_id, tx) in clients.iter_mut() {
        let msg = WsMessage::Text(message.to_string());
        println!("{}", msg);
        tx.start_send(msg).unwrap()
    }
    Ok(())
}

客户端代码

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Title</title>
</head>
<body>
    <h1>Rust chat App</h1>
    <input type="text" id="username" placeholder="username">
    <input type="text" id="message" placeholder="Enter your message">
    <button onclick="sendMessage()">发送</button>
    <ul id="messages"></ul>
    <script>
        let websocket = new WebSocket('ws://localhost:8080')

        websocket.onopen = () =>{
            console.log("连接服务器")
        }

        websocket.onmessage = (event) =>{
            let messages = document.getElementById('messages')
            let li = document.createElement('li')
            li.textContent = event.data
            messages.appendChild(li)
        }

        function sendMessage(){
            let username = document.getElementById('username').value
            let message = document.getElementById('message').value
            let data = {
                username:username,
                message:message
            }
            websocket.send(JSON.stringify(data))
            document.getElementById('message').value = ''
        }
    </script>
</body>
</html>