Refactor to leverage client-body-reader instead of implementing our own chunked reader

Author: tkan145

gateway/src/apicast/http_proxy.lua

@@ -6,7 +6,7 @@ local round_robin = require 'resty.balancer.round_robin'
 local http_proxy = require 'resty.http.proxy'
 local file_reader = require("resty.file").file_reader
 local chunked_reader = require('resty.http.chunked').chunked_reader
-local chunked_writer = require('resty.http.chunked').chunked_writer
+local write_response = require('resty.http.chunked').write_response
 local req_headers = ngx.req.get_headers
 local http_ver = ngx.req.http_version
 local ngx_send_headers = ngx.send_headers
@@ -189,7 +189,7 @@ local function forward_https_request(proxy_uri, uri, skip_https_connect)
     if res then
         -- if we are using raw socket we will need to send the response back with sock:send
         if sock then
-            chunked_writer(sock, res ,chunksize)
+            write_response(sock, res ,chunksize)
         else
             httpc:proxy_response(res)
         end

gateway/src/resty/http/chunked.lua

@@ -1,11 +1,9 @@
-local co_wrap_iter = require("resty.coroutines").co_wrap_iter
-local co_yield = coroutine._yield
+local httpc = require "resty.resolver.http"
 
 local _M = {
 }
 
 local cr_lf = "\r\n"
-local default_max_chunk_size = 32 * 1024 -- 32K
 
 local function send(socket, data)
     if not data or data == '' then
@@ -21,100 +19,9 @@ local function print_err(error_code, ...)
     return ngx.exit(error_code)
 end
 
--- This is a copy of lua-resty-http _chunked_body_reader function but with
--- extra bits to make chunked encoding work with raw socket
--- https://github.com/ledgetech/lua-resty-http/blob/v0.16.1/lib/resty/http.lua#L418
-
--- chunked_reader return a body reader that translates the data read from sock
--- out of HTTP "chunked" format before returning it
---
--- The chunked reader return nil when the final 0-length chunk is read
-function _M.chunked_reader(sock, max_chunk_size)
-    max_chunk_size = max_chunk_size or default_max_chunk_size
-
-    if not sock then
-        return nil, "chunked_reader: invalid sock"
-    end
-
-    return co_wrap_iter(function()
-        local eof = false
-        local remaining = 0
-        local size = 0
-        repeat
-            -- If we still have data on this chunk
-            if max_chunk_size and remaining > 0 then
-                if remaining > max_chunk_size then
-                    -- Consume up to max_chunk_size
-                    size = max_chunk_size
-                    remaining = remaining - max_chunk_size
-                else
-                    -- Consume all remaining
-                    size = remaining
-                    remaining = 0
-                end
-            else
-                --  read a line from socket
-                --  chunk-size CRLF
-                local line, err = sock:receive()
-                if not line then
-                    co_yield(nil, "chunked_reader: failed to receive chunk size, err: " .. (err or "unknown"))
-                end
-
-                size = tonumber(line, 16)
-                if not size then
-                    co_yield(nil, "chunked_reader: unable to read chunksize")
-                end
-
-                if max_chunk_size and size > max_chunk_size then
-                    -- Consume up to max_chunk_size
-                    remaining = size - max_chunk_size
-                    size = max_chunk_size
-                end
-            end
-
-
-            if size > 0 then
-                -- Receive the chunk
-                local chunk, err = sock:receive(size)
-                if not chunk then
-                    co_yield(nil, "chunked_reader: failed to receive chunk of size " .. size .. " err: " .. (err or "unknown"))
-                end
-
-                if remaining == 0 then
-                    -- We're at the end of a chunk, read the next  two bytes
-                    -- and verify they are  "\r\n"
-                    local data, err = sock:receive(2)
-                    if not data then
-                        co_yield(nil, "chunked_reader: failed to receive chunk terminator, err: " .. (err or "unknown"))
-                    end
-                end
-
-                chunk = string.format("%x\r\n", size) .. chunk .. cr_lf
-
-                co_yield(chunk)
-            else
-                -- we're at the end of a chunk, read the next two
-                -- bytes to verify they are "\r\n".
-                local chunk, err = sock:receive(2)
-                if not chunk then
-                    co_yield(nil, "chunked_reader: failed to receive chunk terminator, err: " .. (err or "unknown"))
-                end
-
-                if chunk ~= "\r\n" then
-                    co_yield(nil, "chunked_reader: bad chunk terminator")
-                end
-
-                eof = true
-                co_yield("0\r\n\r\n")
-                break
-            end
-        until eof
-    end)
-end
-
--- chunked_writer writes response body reader to sock in the HTTP/1.x server response format,
+-- write_response writes response body reader to sock in the HTTP/1.x server response format,
 -- including the status line, headers, body, and optional trailer.
-function _M.chunked_writer(sock, res, chunksize)
+function _M.write_response(sock, res, chunksize)
     local bytes, err
     chunksize = chunksize or 65536
 
@@ -163,4 +70,35 @@ function _M.chunked_writer(sock, res, chunksize)
     end
 end
 
+-- chunked_reader return a body reader that translates the data read from
+-- lua-resty-http client_body_reader to HTTP "chunked" format before returning it
+--
+-- The chunked reader return nil when the final 0-length chunk is read
+function _M.chunked_reader(sock, chunksize)
+    chunksize = chunksize or 65536
+    local eof = false
+    local reader = httpc:get_client_body_reader(chunksize, sock)
+    if not reader then
+        return nil
+    end
+
+    return function()
+        if eof then
+            return nil
+        end
+
+        local buffer, err = reader()
+        if err then
+            return nil, err
+        end
+        if buffer then
+            local chunk = string.format("%x\r\n", #buffer) .. buffer .. cr_lf
+            return chunk
+        else
+            eof = true
+            return "0\r\n\r\n"
+        end
+    end
+end
+
 return _M

t/http-proxy.t

@@ -15,7 +15,7 @@ $ENV{'LARGE_BODY'} = large_body();
 
 require("policies.pl");
 
-repeat_each(1);
+repeat_each(3);
 
 run_tests();