feat: add max_header_len & validate_handshake options

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)
 

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

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: <html>.*"
+--- 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]