如何检测TcpStream传入协议？单端口适配WebSocket与普通Socket

1. 更优实现方案 #

优先建议拆分端口监听：给普通Socket和WebSocket分别分配独立端口（比如8081给普通Socket，8082给WebSocket），这种方式逻辑最清晰，完全避免协议解析冲突，也方便后续维护、扩展和流量隔离。

如果必须使用单端口，有两种可行方向：

• 基于协议特征自动检测：通过读取流的头部数据，判断是WebSocket握手请求还是自定义Socket协议
• 引入客户端前置标识：让自定义Socket客户端在连接后先发送一段固定标识，服务端据此区分协议类型

2. 标准的流类型检测方法 #

WebSocket是基于HTTP升级的协议，握手请求有明确的标准特征：

• 起始行为GET / HTTP/1.1（或兼容的HTTP版本）
• 必须携带Upgrade: websocket和Connection: Upgrade请求头

可以利用这些特征做检测，也可以借助tungstenite的握手逻辑通过错误处理区分：

方法一：利用tungstenite的握手错误区分 #

tungstenite::accept内部会自动校验WebSocket握手请求，若不符合则返回错误，可借此区分协议类型。注意要处理“已读取数据丢失”的问题，需用可回退的流包装：

use std::io::{self, Cursor, Chain};
use tungstenite::accept;

while let Ok((stream, _)) = server.accept() {
    // 先尝试WebSocket握手
    match accept(&mut stream.try_clone()?) {
        Ok(ws) => {
            // 处理WebSocket客户端逻辑
            handle_websocket(ws);
        }
        Err(_) => {
            // 握手失败，按普通Socket处理
            handle_plain_socket(stream);
        }
    }
}

方法二：手动检测WebSocket握手特征 #

通过预读流的头部数据，判断是否符合WebSocket握手格式，更灵活可控：

use std::io::{self, BufRead, BufReader, Peekable};

fn is_websocket_handshake(stream: &mut Peekable<BufReader<std::net::TcpStream>>) -> io::Result<bool> {
    // 读取请求起始行
    let mut line = String::new();
    stream.read_line(&mut line)?;
    if !line.starts_with("GET ") || !line.contains("HTTP/1.1") {
        return Ok(false);
    }

    // 遍历请求头，检查Upgrade和Connection字段
    loop {
        let mut header_line = String::new();
        let bytes_read = stream.read_line(&mut header_line)?;
        if bytes_read == 0 || header_line == "\r\n" {
            break;
        }
        let lower_header = header_line.to_lowercase();
        if lower_header.contains("upgrade: websocket") && lower_header.contains("connection: upgrade") {
            return Ok(true);
        }
    }
    Ok(false)
}

// 主循环中使用
while let Ok((stream, _)) = server.accept() {
    let mut peekable_stream = BufReader::new(stream).peekable();
    let is_ws = is_websocket_handshake(&mut peekable_stream).unwrap_or(false);
    
    if is_ws {
        let ws = accept(peekable_stream.into_inner())?;
        handle_websocket(ws);
    } else {
        handle_plain_socket(peekable_stream.into_inner());
    }
}

方法三：自定义客户端前置标识 #

如果自定义Socket协议的控制权在你手中，可让客户端连接后先发送固定标识（比如b"PLAIN_SOCKET\n"），服务端读取后判断：

use std::io::{self, Cursor, Chain};

while let Ok((mut stream, _)) = server.accept() {
    let mut buf = [0; 13]; // 对应"PLAIN_SOCKET\n"的长度
    let n = stream.read(&mut buf)?;
    
    if n == 13 && &buf == b"PLAIN_SOCKET\n" {
        handle_plain_socket(stream);
    } else {
        // 将已读取的字节放回流中，再处理WebSocket
        let cursor = Cursor::new(buf[..n].to_vec());
        let mut combined_stream = Chain::new(cursor, stream);
        let ws = accept(&mut combined_stream)?;
        handle_websocket(ws);
    }
}

总结 #

• 能拆分端口就优先拆分，是生产环境最稳妥的方案
• 单端口场景下，优先选择WebSocket握手特征检测（符合标准）或客户端前置标识（更可控）
• IP判断完全不适合生产环境，必须废弃

内容的提问来源于stack exchange，提问作者Raeisi