从编解码层面理解 WebSocket ，手写一 个WebSocket（Rust）

asyncrustacean2023-10-26 16:54

本文旨在彻底讲清楚 websocket 是如何在 tcp 上面进行握手和数据传输的，并且用代码实现一个 ws server，希望读者看完这篇文章后彻底理解 websocket

websocket rfc

有兴趣的同学可以去看下完整的 websocket rfc

这里我只对重要的部分介绍

握手阶段

当在浏览器里执行 new Websocket("ws://localhost:8080") 的时候

浏览器首先会和服务器建立一个 tcp 连接，在 tcp 上传输 http 报文进行握手，具体的请求报文内容如下图所示

其中有几个特殊的 header 字段，分别是

  1. Upgrade: websocket
  2. Connection: Upgrade
  3. Sec-WebSocket-Protocol
  4. Sec-WebSocket-Key
  5. Sec-WebSocket-Version

其中 1，2 是固定的

Sec-WebSocket-Protocol 是客户端列举的一些子协议，服务端可以选择其中的一个或多个返回

Sec-WebSocket-Key 是客户端随机生成的( 先随机生成 16 字节的数据, 然后 base64 )，服务端需要将该字段的值拼接上一个固定的值258EAFA5-E914-47DA-95CA-C5AB0DC85B11, 然后用 sha-1 hash算法得到 hash 结果，再进行 base64 编码，结果放在响应头 Sec-WebSocket-Accept 返回

比如 Sec-WebSocket-Key 的值是 dGhlIHNhbXBsZSBub25jZQ== ，拼接得到 dGhlIHNhbXBsZSBub25jZQ==258EAFA5-E914-47DA-95CA-C5AB0DC85B11，将该值 sha1 hash 再 base64 得到 s3pPLMBiTxaQ9kYGzzhZRbK+xOo=，就是 Sec-WebSocket-Accept 的值

除了这些以外, 还有 Sec-WebSocket-Extensions, 表示一些扩展协议, 本文不讨论扩展协议

响应的头信息如下图

对应的字段都解释过了，需要注意的是响应状态码是 101, 所以整个握手的流程如下所示

sequenceDiagram Client->Server: 建立 tcp connection Client->>Server: http request 请求报文 Server->>Client: http response (Protocol, Accept) Client->Server: tcp 传输数据

传输数据

ws 通过在 tcp 传输一个个 frame 来传输消息, 具体 frame 的定义如下图

frame的格式实际上非常简单, 具体字段我一个个来解释

fin

第一个 bit 是 fin, 用来标识该 frame 是否是一个消息体里面的最后一个, 如果该消息体只有一个 frame, 那么 fin 也是 1

rsv1,rsv2,rsv3

rsv1,rsv2,rsv3, 这三个一般都是 0, 各 1bit 大小, 除非有扩展协议指定他们的值

opcode

opcode 指定了 payload data 的类型, 一共 4 bit

  • 0 表示这是一组连续的frame 其中的一个,和前面的 frame 合并
  • 1 表示这是文本消息(text)
  • 2 表示这是二进制消息
  • 3-7 是保留的类型
  • 8 关闭消息, 通知对端即将关闭
  • 9 ping 消息, 用来测活
  • 0xa pong 消息, 用来响应 ping 消息
  • 0xb-0xf 是保留的类型

mask

mask 表明 payload data 是否被掩码, 大小 1bit, 什么是掩码后面会提到

  1. 如果是客户端发往服务端的消息, 必须要进行掩码
  2. 如果是服务端发往客户端的消息, 必须不掩码

如果违反了上面的规则, 那么连接必须要断开

payload length

首先读取 7bit 的数字, 如果在 0-125 之间,那么就是 payload length

如果是 126, 那么读取后面 2 个字节的数字作为 payload length

如果是 127, 读取后面 8 个字节的数字作为 payload length

masking key

如果 mask 是 1, 那么就有 4 字节的 masking key, 否则该字段不存在

payload data 和掩码

如果 mask 是 1, 那么 payload data 就是被掩码的, 具体掩码的规则是:

假设 i 表示是 payload data 第 i 个字节

j 为 i % 4 , 也就是对 4 取余

将masking key 第 j 位的数字和 payload 第 i 位的数字做异或, 就是掩码后 payload data 第 i 位的数字

掩码的具体内容后面代码会提到

总结

现在我们对整个 ws 的传输过程了解的差不多了, 下一步就是编码实现

实现一个 WebSocket Server

本文用 rust 实现, 如果你不懂 rust, 也不要担心, 不同语言之间的实现都是差不多的(你肯定看得懂), 如果你想学习 rust 可以去这里

代码除了用 ring 和 base64 这两个 crate, 就没有用其他 crate 了, 代码只是为了最快实现 ws 的传输过程, 不考虑性能等其他问题

实现握手环节

首先搭建一个 tcp 服务器

rust 复制代码 
use std::{
    error::Error,
    io::{BufRead, BufReader, BufWriter, Write},
    net::TcpListener,
};

fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("0.0.0.0:8080")?;

    while let Ok((stream, _)) = listener.accept() {
        let mut reader = BufReader::new(&stream);
        let mut writer = BufWriter::new(&stream);
        handshake(&mut reader, &mut writer)?;
    }

    Ok(())
}

fn handshake(reader: &mut impl BufRead, writer: &mut impl Write) -> Result<(), Box<dyn Error>> {
    // 在这里编写握手的过程
    Ok(())
}

在握手阶段写一个简单的 http parser 即可, 重点是获取 sec-websocket-key, 然后计算出 sec-websocket-accept 字段,并返回响应体, 完成握手

完整的代码如下所示

rust 复制代码 
use base64::{engine::general_purpose, Engine as _};
use ring::digest;
use std::{
    collections::BTreeMap,
    error::Error,
    io::{self, BufRead, BufReader, BufWriter, Write},
    net::TcpListener,
};

fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("0.0.0.0:8080")?;

    while let Ok((stream, _)) = listener.accept() {
        let mut reader = BufReader::new(&stream);
        let mut writer = BufWriter::new(&stream);
        handshake(&mut reader, &mut writer)?;
    }

    Ok(())
}

// 握手
fn handshake(reader: &mut impl BufRead, writer: &mut impl Write) -> Result<(), Box<dyn Error>> {
    let mut buffer = String::new();
    let size = reader.read_line(&mut buffer)?;
    // 读取 http 请求行
    let request_line: &str = &buffer[0..size];
    let _ = request_line;
    buffer.truncate(0);

    let mut headers = BTreeMap::<String, String>::new();

    loop {
        let size = reader.read_line(&mut buffer)?;
        // 读取每一个头信息
        let header_line: &str = &buffer[0..size];
        // 头信息完结
        if header_line == "\r\n" {
            break;
        }

        let header_line = &header_line[0..(size - 2)];

        if let Some((k, v)) = header_line.split_once(':') {
            headers.insert(k.to_lowercase(), v.trim_start().into());
        };

        buffer.truncate(0);
    }

    let sec_websocket_key = headers.get("sec-websocket-key").ok_or(io::Error::new(
        io::ErrorKind::ConnectionRefused,
        "no header Sec-Websocket-Key",
    ))?;

    const UUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    // sha1 加 base64
    let concat_str = [sec_websocket_key.as_bytes(), UUID].concat();
    let hash_result = digest::digest(&digest::SHA1_FOR_LEGACY_USE_ONLY, &concat_str);
    let sec_websocket_accept = general_purpose::STANDARD.encode(hash_result.as_ref());

    let response = format!(
        "HTTP/1.1 101 Switching Protocols\r\n\
        Upgrade: websocket\r\n\
        Connection: Upgrade\r\n\
        Sec-WebSocket-Accept: {}\r\n\r\n",
        sec_websocket_accept
    );

    writer.write_all(response.as_bytes())?;

    writer.flush()?;

    Ok(())
}

cargo run 将服务启动起来，然后在浏览器控制台输入

js 复制代码 
const ws = new WebSocket("ws://localhost:8080/")

ws.addEventListener("open", () => {
  console.log("open")
})

你会看到

说明 websocket 已经成功通过了握手的环节

实现消息的编码和解码

先定义两种消息

rust 复制代码 
enum Message {
    Text(String), // 文本消息
    Binary(Vec<u8>), // 二进制消息
}

消息的编码

我们开始将消息编码成一个 frame, 如果你不知道消息的具体长度（比如流），你可以用多个 frame 将消息切片发送, 这里不考虑这个问题

开始定义 message 的编码方法, 因为服务端发往客户端的消息不掩码, 实现起来就很简单

rust 复制代码 
impl Message {
    fn as_bytes(&self) -> &[u8] {
        match &self {
            Message::Binary(data) => data,
            Message::Text(data) => data.as_bytes(),
        }
    }
    
    // opcode  表示当前的消息类型
    fn opcode(&self) -> u8 {
        match &self {
            Message::Binary(_) => 2,
            Message::Text(_) => 1,
        }
    }

    fn encode(&self) -> Vec<u8> {
        let payload_data = self.as_bytes();
        let payload_length = payload_data.len() as u64;

        // 初始的长度是 2个 字节 fin,rsv1...payload_length
        let mut total_frame_length = 2;

        if payload_length > 125 {
            // 扩展 payload_length
            if payload_length > u16::MAX as u64 {
                total_frame_length += 8;
            } else {
                total_frame_length += 2;
            }
        }

        total_frame_length += payload_length;

        let mut frame: Vec<u8> = Vec::with_capacity(total_frame_length as usize);

        let opcode = self.opcode();
        frame.push(0b1000_0000); // fin 是 1
        frame[0] |= opcode;

        if payload_length <= 125 {
            frame.push(payload_length as u8);
        } else if payload_length > u16::MAX as u64 {
            frame.push(127);
            frame.extend_from_slice(&payload_length.to_be_bytes());
        } else {
            frame.push(126);
            frame.extend_from_slice(&(payload_length as u16).to_be_bytes());
        }

        // 服务端不需要 mask, 直接拼接数据
        frame.extend_from_slice(payload_data);

        frame
    }
}

验证一下这个编码是否正确, 新增一个 handle_connection 函数, 分别向客户端发送 大,中,小 三种长度的文本消息,还有一个大的二进制消息, 在客户端验证这个效果

服务端的代码

rust 复制代码 
fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("0.0.0.0:8080")?;

    while let Ok((stream, _)) = listener.accept() {
        let mut reader = BufReader::new(&stream);
        let mut writer = BufWriter::new(&stream);
        handshake(&mut reader, &mut writer)?;
        handle_connection(&mut reader, &mut writer)?;
    }

    Ok(())
}

fn handle_connection(
    reader: &mut impl BufRead,
    writer: &mut impl Write,
) -> Result<(), Box<dyn Error>> {
    let message = Message::Text("hello from server side".into());
    writer.write_all(&message.encode())?;

    let middle_message = Message::Text(String::from_utf8_lossy(&[b'*'; 126]).to_string());
    writer.write_all(&middle_message.encode())?;

    let big_message =
        Message::Text(String::from_utf8_lossy(&[b'*'; (u16::MAX as usize + 1)]).to_string());
    writer.write_all(&big_message.encode())?;

    let big_binary_message = Message::Binary(vec![b'*'; u16::MAX as usize + 1]);

    writer.write_all(&big_binary_message.encode())?;

    writer.flush()?;
    Ok(())
}

客户端的代码

javascript 复制代码 
const ws = new WebSocket("ws://localhost:8080/")
ws.addEventListener("message", (ev) => {
  if (ev.data instanceof Blob) {
    // 二进制数据
    console.log({
      messageLength: ev.data.size,
      messageType: "binary",
      messageData: ev.data,
    })
  } else {
    // 文本消息
    console.log({
      messageLength: ev.data.length,
      messageType: "text",
      messageData: ev.data,
    })
  }
})

重新启动服务, 并在控制台运行这段代码, 不出意外的话, 你可以看到输出

说明我们的服务已经可以正确的编码并向客户端发送消息了

消息的解码

从客户端发往服务端的所有消息都是掩码的, ws 的掩码使用异或算法, 异或算法的特性是对同一个数字异或两次会还原数字, 服务端在接收消息的时候要还原数据, 代码如下

rust 复制代码 
fn decode_message(reader: &mut impl BufRead) -> Result<Message, Box<dyn Error>> {
    let mut buffer = [0; 2];
    // 先获取前面两个字节
    reader.read_exact(&mut buffer)?;

    // 不考虑 fin 不为 1 的情况, 一次读取一个 frame 然后拼接成 message
    let opcode = buffer[0] & 0b1111;
    let mask = buffer[1] >> 7;
    if mask != 1 {
        // 客户端发来的消息必须是掩码的
        return Err(io::Error::new(io::ErrorKind::ConnectionRefused, "mask require").into());
    }

    let mut payload_length = (buffer[1] & 0b0111_1111) as u64;

    if payload_length == 126 {
        reader.read_exact(&mut buffer)?;
        payload_length = u16::from_be_bytes(buffer) as u64;
    } else if payload_length == 127 {
        let mut buffer = [0; 8];
        reader.read_exact(&mut buffer)?;
        payload_length = u64::from_be_bytes(buffer);
    }

    let mut mask_key = [0; 4];
    reader.read_exact(&mut mask_key)?;

    let mut payload_data: Vec<u8> = vec![0; payload_length as usize];
    reader.read_exact(&mut payload_data)?;

    // 还原原始的 payload_data
    (0..payload_data.len()).for_each(|i| {
        let j = i % 4;
        let cur_mask_key = mask_key[j];
        payload_data[i] ^= cur_mask_key;
    });

    Ok(if opcode == 1 {
        Message::Text(String::from_utf8_lossy(&payload_data).to_string())
    } else {
        Message::Binary(payload_data)
    })
}

测试一下, 更改 handle_connection 如下

rust 复制代码 
fn handle_connection(
    reader: &mut impl BufRead,
    writer: &mut impl Write,
) -> Result<(), Box<dyn Error>> {
    while let Ok(message) = decode_message(reader) {
        match message {
            Message::Text(text) => {
                println!("{text}");
            }
            Message::Binary(data) => {
                println!("{data:?}")
            }
        }
    }
    Ok(())
}

客户端代码如下, 控制台里面调用

javascript 复制代码 
const ws = new WebSocket("ws://localhost:8080/")
ws.onopen = () => {
  ws.send("hello from client side")
  ws.send("*".repeat(127))
  ws.send("*".repeat(65536))
}

服务端的输出结果如下

证明解码也没问题

总结

具体的代码放在 github 上面了, 希望对大家有帮助, 能更好的理解 websocket 的传输过程, 谢谢大家

相关推荐
一个小坑货几秒前
Cargo Rust 的包管理器
开发语言·后端·rust
bluebonnet274 分钟前
【Rust练习】22.HashMap
开发语言·后端·rust
uhakadotcom27 分钟前
如何实现一个基于CLI终端的AI 聊天机器人？
后端
Iced_Sheep1 小时前
干掉 if else 之策略模式
后端·设计模式
XINGTECODE1 小时前
海盗王集成网关和商城服务端功能golang版
开发语言·后端·golang
程序猿进阶1 小时前
堆外内存泄露排查经历
java·jvm·后端·面试·性能优化·oom·内存泄露
FIN技术铺2 小时前
Spring Boot框架Starter组件整理
java·spring boot·后端
凡人的AI工具箱2 小时前
15分钟学 Go 第 60 天 ：综合项目展示 - 构建微服务电商平台（完整示例25000字）
开发语言·后端·微服务·架构·golang
先天牛马圣体2 小时前
如何提升大型AI模型的智能水平
后端
java亮小白19972 小时前
Spring循环依赖如何解决的？
java·后端·spring
热门推荐
01(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明02PyTorch机器学习实现液态神经网络03【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总04玄机平台应急响应—webshell查杀05Coze扣子平台完整体验和实践（附国内和国际版对比）06Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO07【目标跟踪】相机运动补偿08Unity中PICO实现 隔空取物 和 接触抓取物体09RAG 实践- Ollama+RagFlow 部署本地知识库10量化方法怎么选？如何评估量化后的大模型LLM？