-
Notifications
You must be signed in to change notification settings - Fork 170
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Initial support to proxy request with Transfer-Encoding: chunked
- Loading branch information
Showing
3 changed files
with
258 additions
and
24 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,166 @@ | ||
local co_wrap_iter = require("resty.coroutines").co_wrap_iter | ||
local co_yield = coroutine._yield | ||
|
||
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 | ||
ngx.log(ngx.DEBUG, 'skipping sending nil') | ||
return | ||
end | ||
|
||
return socket:send(data) | ||
end | ||
|
||
local function print_err(error_code, ...) | ||
ngx.log(ngx.ERR, ...) | ||
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, | ||
-- including the status line, headers, body, and optional trailer. | ||
function _M.chunked_writer(sock, res, chunksize) | ||
local bytes, err | ||
chunksize = chunksize or 65536 | ||
|
||
-- Status line | ||
-- FIXME: should get protocol version from res? | ||
local status = "HTTP/1.1 " .. res.status .. " " .. res.reason .. cr_lf | ||
bytes, err = send(sock, status) | ||
if not bytes then | ||
print_err(503, "chunked_writer: failed to send status line, err: " .. (err or "unknown")) | ||
end | ||
|
||
-- Rest of header | ||
for k, v in pairs(res.headers) do | ||
local header = k .. ": " .. v .. cr_lf | ||
bytes, err = send(sock, header) | ||
if not bytes then | ||
print_err(503, "chunked_writer: failed to send header, err: " .. (err or "unknown")) | ||
end | ||
end | ||
|
||
-- End-of-header | ||
bytes, err = send(sock, cr_lf) | ||
if not bytes then | ||
print_err(503, "chunked_writer: failed to send end of header, err: " .. (err or "unknown")) | ||
end | ||
|
||
-- Write body and trailer | ||
-- TODO: handle trailer | ||
if res.has_body then | ||
local reader = res.body_reader | ||
repeat | ||
local chunk, read_err | ||
|
||
chunk, read_err = reader(chunksize) | ||
if read_err then | ||
print_err(503, "chunked_writer: failed to read body, err: " .. (err or "unknown")) | ||
end | ||
|
||
if chunk then | ||
bytes, err = send(sock, chunk) | ||
if not bytes then | ||
print_err(503, "chunked_writer: failed to send body, err: " .. (err or "unknown")) | ||
end | ||
end | ||
until not chunk | ||
end | ||
end | ||
|
||
return _M |