From 2db11aed17c8a461fd7e722c6256f5b6cf9e9a2d Mon Sep 17 00:00:00 2001 From: Taylor Jasko Date: Fri, 14 Mar 2025 21:25:40 -0500 Subject: [PATCH] feat: add `max_header_len` & `validate_handshake` options - Added `max_header_len` option to limit the maximum allowed header size during the WebSocket upgrade process. - Added `validate_handshake` option to enforce that the WebSocket handshake response must return HTTP 101. - Improved HTTP response parsing by checking the status line and extracting response headers properly. - Added new test cases --- README.markdown | 6 +++ lib/resty/websocket/client.lua | 44 +++++++++++++++++----- t/cs.t | 69 ++++++++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 9 deletions(-) diff --git a/README.markdown b/README.markdown index af3a1dc..477df29 100644 --- a/README.markdown +++ b/README.markdown @@ -185,12 +185,18 @@ An optional options table can be specified. The following options are as follows * `max_send_len` Specifies the maximal length of payload allowed when sending WebSocket frames. Defaults to the value of `max_payload_len`. +* `max_header_len` + + Specifies the maximal length of payload allowed when receiving headers during the WebSocket upgrade process. Defaults to `0`, disabling the check allowing unlimited length. * `send_masked` Specifies whether to send out masked WebSocket frames. When it is `true`, masked frames are always sent. Default to `false`. * `timeout` Specifies the network timeout threshold in milliseconds. You can change this setting later via the `set_timeout` method call. Note that this timeout setting does not affect the HTTP response header sending process for the websocket handshake; you need to configure the [send_timeout](http://nginx.org/en/docs/http/ngx_http_core_module.html#send_timeout) directive at the same time. +* `validate_handshake` + + Specifies whether to ensure the WebSocket upgrade returned an HTTP 101 status code. When the handshake fails, both the HTTP status code & response body will be captured in the returned error message. Default to `false`. [Back to TOC](#table-of-contents) diff --git a/lib/resty/websocket/client.lua b/lib/resty/websocket/client.lua index 55f77cb..57df078 100644 --- a/lib/resty/websocket/client.lua +++ b/lib/resty/websocket/client.lua @@ -51,11 +51,14 @@ function _M.new(self, opts) end local max_payload_len, send_unmasked, timeout - local max_recv_len, max_send_len + local max_recv_len, max_send_len, max_header_len + local validate_handshake if opts then max_payload_len = opts.max_payload_len max_recv_len = opts.max_recv_len max_send_len = opts.max_send_len + max_header_len = opts.max_header_len + validate_handshake = opts.validate_handshake send_unmasked = opts.send_unmasked timeout = opts.timeout @@ -68,12 +71,16 @@ function _M.new(self, opts) max_payload_len = max_payload_len or 65535 max_recv_len = max_recv_len or max_payload_len max_send_len = max_send_len or max_payload_len + max_header_len = max_header_len or 0 + validate_handshake = validate_handshake or false return setmetatable({ sock = sock, max_recv_len = max_recv_len, max_send_len = max_send_len, + max_header_len = max_header_len, send_unmasked = send_unmasked, + validate_handshake = validate_handshake, }, mt) end @@ -265,21 +272,40 @@ function _M.connect(self, uri, opts) return nil, "failed to send the handshake request: " .. err end + -- Parse request up to end of headers. + local header, err local header_reader = sock:receiveuntil("\r\n\r\n") - -- FIXME: check for too big response headers - local header, err, partial = header_reader() + if self.max_header_len > 0 then + header, err = header_reader(self.max_header_len + 1) + if string.len(header) > self.max_header_len then + return nil, "response headers too large (limit: " .. self.max_header_len .. " bytes)" + end + else + header, err = header_reader() + end if not header then return nil, "failed to receive response header: " .. err end - -- error("header: " .. header) - - -- FIXME: verify the response headers - - m, err = re_match(header, [[^\s*HTTP/1\.1\s+]], "jo") - if not m then + -- Validate HTTP status line. + local status_line_end = header:find("\r?\n") + local status_line + if not status_line_end then return nil, "bad HTTP response status line: " .. header end + status_line = header:sub(1, status_line_end - 1) + local status_code = status_line:match("^HTTP/1%.1 (%d+)") + if not status_code then + return nil, "bad HTTP response status code line: " .. header + end + + -- Ensure the status code is 101 (Switching Protocols) per RFC 6455. + -- This status code check is optional for backward compatibility. + if self.validate_handshake and status_code ~= "101" then + local body, body_err = sock:receive("*a") + body = body or "(no body received)" + return nil, "unexpected HTTP response, code: " .. status_code .. ", body: " .. body + end return 1, nil, header end diff --git a/t/cs.t b/t/cs.t index e1b7181..95f7d44 100644 --- a/t/cs.t +++ b/t/cs.t @@ -2695,3 +2695,72 @@ received text frame: reused connection --- no_error_log [error] [warn] + + +=== TEST 40: return full response body when handshake fails +--- http_config eval: $::HttpConfig +--- config + location = /c { + content_by_lua_block { + local client = require "resty.websocket.client" + local wb, err = client:new{ validate_handshake = true } + local uri = "ws://127.0.0.1:" .. ngx.var.server_port .. "/s" + local ok, err, res = wb:connect(uri) + if ok then + ngx.say("unexpected connection success") + return + end + + ngx.say("error: \"", err, "\"") + } + } + + location = /s { + return 400; + } +--- request +GET /c +--- response_body_like +^error: "unexpected HTTP response, code: 400, body: .*" +--- no_error_log +[error] +[warn] + + +=== TEST 41: response headers exceed max_header_len +--- http_config eval: $::HttpConfig +--- config + location = /c { + content_by_lua_block { + local client = require "resty.websocket.client" + local wb, err = client:new{ max_header_len = 1024 } + local uri = "ws://127.0.0.1:" .. ngx.var.server_port .. "/s" + local ok, err = wb:connect(uri) + if ok then + ngx.say("unexpected connection success") + return + end + + ngx.say("error: \"", err, "\"") + } + } + + location = /s { + content_by_lua_block { + ngx.header["X-Custom-1"] = string.rep("X", 5000) + + local server = require "resty.websocket.server" + local wb, err = server:new() + if not wb then + ngx.log(ngx.ERR, "failed to new websocket: ", err) + return ngx.exit(444) + end + } + } +--- request +GET /c +--- response_body_like +^error: "response headers too large \(limit: 1024 bytes\)" +--- no_error_log +[error] +[warn] \ No newline at end of file