Skip to content

Commit

Permalink
Refactor to leverage client-body-reader instead of implementing our o…
Browse files Browse the repository at this point in the history
…wn chunked reader
  • Loading branch information
tkan145 committed Jul 18, 2023
1 parent c1cae29 commit 04a9b0c
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 99 deletions.
4 changes: 2 additions & 2 deletions gateway/src/apicast/http_proxy.lua
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
130 changes: 34 additions & 96 deletions 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
Expand All @@ -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

Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion t/http-proxy.t
Expand Up @@ -15,7 +15,7 @@ $ENV{'LARGE_BODY'} = large_body();

require("policies.pl");

repeat_each(1);
repeat_each(3);

run_tests();

Expand Down

0 comments on commit 04a9b0c

Please sign in to comment.