feat(client): new method get_resp_headers() (#16)

KAG-4291

lib/resty/websocket/client.lua

@@ -14,8 +14,10 @@ local _send_frame = wbproto.send_frame
 local new_tab = wbproto.new_tab
 local tcp = ngx.socket.tcp
 local re_match = ngx.re.match
+local re_gmatch = ngx.re.gmatch
 local encode_base64 = ngx.encode_base64
 local concat = table.concat
+local insert = table.insert
 local char = string.char
 local str_find = string.find
 local str_sub = string.sub
@@ -27,6 +29,7 @@ local type = type
 local debug = ngx.config.debug
 local ngx_log = ngx.log
 local ngx_DEBUG = ngx.DEBUG
+local tostring = tostring
 local assert = assert
 local ssl_support = true
 
@@ -339,6 +342,8 @@ function _M.connect(self, uri, opts)
         return nil, "unexpected HTTP response code: " .. m[1], header
     end
 
+    self.resp_header = header
+
     return 1, nil, header
 end
 
@@ -480,4 +485,50 @@ function _M.set_keepalive(self, ...)
 end
 
 
+function _M.get_resp_headers(self)
+    if self.resp_headers then
+        return self.resp_headers
+    end
+
+    local iter, err = re_gmatch(self.resp_header .. "\r\n", "([^:\\s]+):\\s*(.*?)\r\n", "jo")
+    if err then
+        return nil, "failed to parse response header: " .. err
+    end
+
+    -- gather all response headers
+
+    local resp_headers = {}
+
+    while true do
+        local m, err = iter()
+        if err then
+            return nil, "failed to parse response header: " .. err
+        end
+
+        if not m then
+            -- no match found (any more)
+            break
+        end
+
+        local key = m[1]:lower():gsub("-", "_")
+        local val = m[2]
+
+        if resp_headers[key] then
+            if type(resp_headers[key]) ~= "table" then
+                resp_headers[key] = { resp_headers[key] }
+            end
+
+            insert(resp_headers[key], tostring(val))
+
+        else
+            resp_headers[key] = tostring(val)
+        end
+    end
+
+    self.resp_headers = resp_headers
+
+    return resp_headers
+end
+
+
 return _M

t/count.t

@@ -63,7 +63,7 @@ size: 5
 --- request
 GET /t
 --- response_body
-size: 13
+size: 14
 --- no_error_log
 [error]
 

t/cs.t

@@ -2881,3 +2881,65 @@ received: hello (text)
 [warn]
 
 
+=== TEST 43: client:get_resp_headers
+--- http_config eval: $::HttpConfig
+--- config
+    location = /c {
+        content_by_lua '
+            local client = require "resty.websocket.client"
+            local wb, err = client:new()
+            local uri = "ws://127.0.0.1:" .. ngx.var.server_port .. "/s"
+            -- ngx.say("uri: ", uri)
+            local ok, err = wb:connect(uri)
+            if not ok then
+                ngx.say("failed to connect: " .. err)
+                return
+            end
+
+            local data, typ, err = wb:recv_frame()
+            if not data then
+                ngx.say("failed to receive 1st frame: ", err)
+                return
+            end
+
+            ngx.say("1: received: ", data, " (", typ, ")")
+
+            local resp_headers = wb:get_resp_headers()
+
+            ngx.say(resp_headers.upgrade)
+            ngx.say(resp_headers.connection)
+            ngx.say(resp_headers.x_foo)
+        ';
+    }
+
+    location = /s {
+        content_by_lua '
+            local server = require "resty.websocket.server"
+
+            ngx.header["x-foo"] = "bar"
+
+            local wb, err = server:new()
+            if not wb then
+                ngx.log(ngx.ERR, "failed to new websocket: ", err)
+                return ngx.exit(444)
+            end
+
+            local bytes, err = wb:send_text("你好, WebSocket!")
+            if not bytes then
+                ngx.log(ngx.ERR, "failed to send the 1st text: ", err)
+                return ngx.exit(444)
+            end
+        ';
+    }
+--- request
+GET /c
+--- response_body
+1: received: 你好, WebSocket! (text)
+websocket
+upgrade
+bar
+--- no_error_log
+[error]
+[warn]
+
+