从编解码层面理解 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 的传输过程, 谢谢大家

相关推荐
追逐时光者1 小时前
推荐 12 款开源美观、简单易用的 WPF UI 控件库，让 WPF 应用界面焕然一新！
后端·.net
Jagger_1 小时前
敏捷开发流程-精简版
前端·后端
苏打水com2 小时前
数据库进阶实战：从性能优化到分布式架构的核心突破
数据库·后端
间彧3 小时前
Spring Cloud Gateway与Kong或Nginx等API网关相比有哪些优劣势？
后端
间彧3 小时前
如何基于Spring Cloud Gateway实现灰度发布的具体配置示例？
后端
间彧3 小时前
在实际项目中如何设计一个高可用的Spring Cloud Gateway集群？
后端
间彧3 小时前
如何为Spring Cloud Gateway配置具体的负载均衡策略？
后端
间彧3 小时前
Spring Cloud Gateway详解与应用实战
后端
EnCi Zheng5 小时前
SpringBoot 配置文件完全指南-从入门到精通
java·spring boot·后端
烙印6015 小时前
Spring容器的心脏：深度解析refresh()方法（上）
java·后端·spring
热门推荐
01BongoCat - 跨平台键盘猫动画工具02两千字总结：Codex 国内如何安装和使用的教程，以及如何设置中文回答03GitHub 镜像站点04UV安装并设置国内源05智能库存管理的需求预测模型：从业务痛点到落地代码的完整实践0646个Nano-banana 精选提示词，持续更新中07Cursor Plan Mode：AI 终于知道先想后做了08Linux下V2Ray安装配置指南09GitLab 零基础入门指南：从安装到项目管理全流程10一文了解国产算子编程语言 TileLang，TileLang 对国产开源生态的影响与启示