feat: add max_header_len & capture_error_body 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.
+* `capture_error_body`
+
+    Specifies whether to read the response body when the WebSocket upgrade is refused with a status other than 101, and append it to the returned error message. The upgrade is rejected either way; this only controls whether the body is included, since reading it means waiting on the socket again. 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 capture_error_body
     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
+        capture_error_body = opts.capture_error_body
 
         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
+    capture_error_body = capture_error_body 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,
+        capture_error_body = capture_error_body,
     }, mt)
 end
 
@@ -265,9 +272,21 @@ function _M.connect(self, uri, opts)
         return nil, "failed to send the handshake request: " .. err
     end
 
+    -- read the response up to the end of the headers, optionally bounded
+    -- by max_header_len
+    local header
     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 header and #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
@@ -284,8 +303,19 @@ function _M.connect(self, uri, opts)
     -- RFC 6455 section 4.1: a status code other than 101 means the server
     -- has not accepted the upgrade, so the client must fail the connection
     if m[1] ~= "101" then
-        return nil, "failed websocket handshake: unexpected response status: "
-                    .. m[1], header
+        local msg = "failed websocket handshake: unexpected response status: "
+                    .. m[1]
+
+        -- the body usually explains why the upgrade was refused, but reading
+        -- it means waiting on the socket again, so it is opt-in
+        if self.capture_error_body then
+            local body = sock:receive("*a")
+            if body then
+                msg = msg .. ", body: " .. body
+            end
+        end
+
+        return nil, msg, header
     end
 
     return 1, nil, header

t/cs.t

@@ -2695,3 +2695,72 @@ received text frame: reused connection
 --- no_error_log
 [error]
 [warn]
+
+
+=== TEST 40: capture_error_body appends the body to the handshake error
+--- http_config eval: $::HttpConfig
+--- config
+    location = /c {
+        content_by_lua_block {
+            local client = require "resty.websocket.client"
+            local wb, err = client:new{ capture_error_body = 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: "failed websocket handshake: unexpected response status: 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]