Nginx WebSocket 源码优化详解

WebSocket 协议简介

WebSocket 是基于 TCP 的应用层协议，用于在 C/S 架构的应用中实现双向通信，RFC 文档为 RFC6455。

WebSocket 特点

1. 双向通信：会话通信实时性强
2. 持久连接：建立后可一直保持连接，期间可不断发送消息，避免了 HTTP 的非状态性
3. 头部开销小：相比 HTTP，数据交换时协议控制的数据包头部较小
4. 多格式支持：可互相发送 JSON、XML、HTML 或图片等任意格式数据
5. 二进制支持：更好的二进制支持和扩展能力

WebSocket 协议格式

0                   1                   2                   3
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-------+-+-------------+-------------------------------+
|F|R|R|R| opcode|M| Payload len |    Extended payload length    |
|I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
|N|V|V|V|       |S|             |   (if payload len==126/127)   |
| |1|2|3|       |K|             |                               |
+-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
|     Extended payload length continued, if payload len == 127  |
+ - - - - - - - - - - - - - - - +-------------------------------+
|                               |Masking-key, if MASK set to 1  |
+-------------------------------+-------------------------------+
| Masking-key (continued)       |          Payload Data         |
+-------------------------------- - - - - - - - - - - - - - - - +
:                     Payload Data continued ...                :
+ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
|                     Payload Data continued ...                |
+---------------------------------------------------------------+

Nginx WebSocket 基础配置

在 Nginx 中开启 WebSocket 支持的基本配置：

server {
    listen 80;
    server_name www.domain.com;
    
    location / {
        proxy_pass http://127.0.0.1:8080/;  # 回源地址
        proxy_http_version 1.1;
        proxy_read_timeout 600s;  # 超时设置
        
        # 启用支持 WebSocket
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
    }
}

这种配置会将 HTTP 连接升级为 WebSocket 连接，但存在一个问题：所有发往后端的流量都会带上 Upgrade 和 Connection 头。

Nginx WebSocket 源码优化

问题背景

当在 server{} 块中配置 WebSocket 时，所有发往后端 upstream 的流量都会添加 WebSocket 头。但在实际场景中，可能只有部分请求需要升级为 WebSocket，其他请求不需要。原生 Nginx 配置无法区分这两种情况。

解决方案

通过修改和优化 Nginx 源码，区分客户端的流量是否为 WebSocket 请求。

1. 定义 WebSocket 头

static ngx_keyval_t ngx_http_proxy_websocket_headers[] = {
    { ngx_string("Host"), ngx_string("$proxy_host") },
    { ngx_string("Connection"), ngx_string("Upgrade") },
    { ngx_string("Content-Length"), ngx_string("$proxy_internal_body_length") },
    { ngx_string("Transfer-Encoding"), ngx_string("$proxy_internal_chunked") },
    { ngx_string("TE"), ngx_string("Keep-Alive") },
    { ngx_string("Expect"), ngx_string("") },
    { ngx_string("Upgrade"), ngx_string("websocket") },
    { ngx_null_string, ngx_null_string }
};

2. 初始化 WebSocket 头

#ifdef NGX_WEBSOCKET_INNER
// 通过 headers 来 init
rc = ngx_http_proxy_init_headers(cf, conf, &conf->websocket_headers,
                                ngx_http_proxy_websocket_headers);
if (rc != NGX_OK) {
    return NGX_CONF_ERROR;
}
#endif

3. 处理连接头

static ngx_int_t
ngx_http_process_connection(ngx_http_request_t *r, ngx_table_elt_t *h,
                           ngx_uint_t offset)
{
    // 获取 upgrade 标志
    if (ngx_strcasestrn(h->value.data, "Upgrade", 7 - 1)) {
        r->is_websocket_request |= NGX_WEBSOCKET_HEADER_CONNECTION;
    }
}

4. 判断 WebSocket 请求

if (r->headers_in.upgrade == NULL) {
    goto not_websocket_request;
} else if ((r->is_websocket_request & NGX_WEBSOCKET_HEADER_CONNECTION) == 0) {
    goto not_websocket_request;
} else if (ngx_strcasestrn(r->headers_in.upgrade->value.data, "websocket", 9 - 1)) {
    r->is_websocket_request |= NGX_WEBSOCKET_HEADER_UPGRADE;
}

// 两种标志都有，r->websocket_request 标志位置位
if (r->is_websocket_request == (NGX_WEBSOCKET_HEADER_UPGRADE | NGX_WEBSOCKET_HEADER_CONNECTION)) {
    r->websocket_request = 1;
    r->http_version = NGX_HTTP_VERSION_11;
} else {
    r->websocket_request = 0;
}

5. 处理上游头

static ngx_int_t
ngx_http_proxy_process_header(ngx_http_request_t *r)
{
    if (u->headers_in.status_n == NGX_HTTP_SWITCHING_PROTOCOLS) {
        u->keepalive = 0;
        if (r->headers_in.upgrade) {
            u->upgrade = 1;
        }
    }
}

6. 发送响应

static void
ngx_http_upstream_send_response(ngx_http_request_t *r, ngx_http_upstream_t *u)
{
    if (u->upgrade) {
        ngx_http_upstream_upgrade(r, u);
        return;
    }
}

7. 升级连接

static void
ngx_http_upstream_upgrade(ngx_http_request_t *r, ngx_http_upstream_t *u)
{
    u->read_event_handler = ngx_http_upstream_upgraded_read_upstream;
    u->write_event_handler = ngx_http_upstream_upgraded_write_upstream;
    r->read_event_handler = ngx_http_upstream_upgraded_read_downstream;
    r->write_event_handler = ngx_http_upstream_upgraded_write_downstream;
}

8. 处理升级后的读写事件

static void
ngx_http_upstream_process_upgraded(ngx_http_request_t *r,
                                  ngx_uint_t from_upstream, ngx_uint_t do_write)
{
    // from_upstream 为 1，src 为 upstream(上游)，dst 为 downstream(下游)
    if (from_upstream) {
        src = upstream;
        dst = downstream;
        b = &u->buffer;
    }
    
    if (do_write) {
        // 写入处理
    }
    
    if (size && src->read->ready) {
        // src 为 upstream，用来读
        n = src->recv(src, b->last, size);
        
        // n >0 接收大于 0 的字节数，do_write 置为 1，continue 进行写入
        if (n > 0) {
            do_write = 1;
            b->last += n;
            if (from_upstream) {
                u->state->bytes_received += n;
            }
            continue;
        }
    }
    
    // 加入计时器
    if (upstream->write->active && !upstream->write->ready) {
        ngx_add_timer(upstream->write, u->conf->send_timeout);
    } else if (upstream->write->timer_set) {
        ngx_del_timer(upstream->write);
    }
    
    if (upstream->read->active && !upstream->read->ready) {
        ngx_add_timer(upstream->read, u->conf->read_timeout);
    } else if (upstream->read->timer_set) {
        ngx_del_timer(upstream->read);
    }
    
    if (downstream->write->active && !downstream->write->ready) {
        ngx_add_timer(downstream->write, clcf->send_timeout);
    } else if (downstream->write->timer_set) {
        ngx_del_timer(downstream->write);
    }
}

关键点总结

1. 协议识别：通过检查 Upgrade 和 Connection 头来识别 WebSocket 请求
2. 标志位设置：使用 NGX_WEBSOCKET_HEADER_UPGRADE 和 NGX_WEBSOCKET_HEADER_CONNECTION 标志位标记 WebSocket 请求
3. 连接升级：当识别为 WebSocket 请求时，将连接升级为 WebSocket 协议
4. 事件处理：为升级后的连接设置专门的读写事件处理函数
5. 计时器管理：为读写操作添加适当的超时控制

通过这种优化，Nginx 可以智能地区分普通 HTTP 请求和 WebSocket 请求，并只对需要升级的请求进行 WebSocket 处理，从而提高了代理的灵活性和效率。