Support to proxy request with Transfer-Encoding: chunked and proxy_buffering: off

tkan145 · tkan145 · commit addb342f5ae0 · 2023-10-10T16:17:14.000+10:00
diff --git a/gateway/src/apicast/http_proxy.lua b/gateway/src/apicast/http_proxy.lua
@@ -7,9 +7,22 @@ local round_robin = require 'resty.balancer.round_robin'
 local http_proxy = require 'resty.http.proxy'
 local file_reader = require("resty.file").file_reader
 local file_size = require("resty.file").file_size
+local get_client_body_reader = require('resty.http.raw_request').get_client_body_reader
+local send_response = require('resty.http.raw_response').send_response
+local http_ver = ngx.req.http_version
+local ngx_send_headers = ngx.send_headers
+local ngx_req = ngx.req
+local ngx_flush = ngx.flush
+local ngx_var = ngx.var
 
 local _M = { }
 
+local http_methods_with_body = {
+  POST = true,
+  PUT = true,
+  PATCH = true
+}
+
 function _M.reset()
     _M.balancer = round_robin.new()
     _M.resolver = resty_resolver
@@ -103,17 +116,70 @@ local function modify_chunked_request_headers(req, file)
     set_header(req.headers, "Content-Length", contentLength)
 end
 
+local function handle_expect()
+    local expect = ngx_var.http_expect
+    if type(expect) == "table" then
+        expect = expect[1]
+    end
 
-local function forward_https_request(proxy_uri, uri, skip_https_connect)
-    -- This is needed to call ngx.req.get_body_data() below.
-    ngx.req.read_body()
+    if expect and expect:lower() == "100-continue" then
+        ngx.status = 100
+        local ok, err = ngx_send_headers()
 
-    local request = {
-        uri = uri,
-        method = ngx.req.get_method(),
-        headers = ngx.req.get_headers(0, true),
-        path = format('%s%s%s', ngx.var.uri, ngx.var.is_args, ngx.var.query_string or ''),
+        if not ok then
+            return nil, "failed to send response header: " .. (err or "unknown")
+        end
+
+        ok, err = ngx_flush(true)
+        if not ok then
+            return nil, "failed to flush response header: " .. (err or "unknown")
+        end
+    end
+end
+
+local function get_request_body(req, opts)
+    local chunksize = 32 * 1024
+    local sock, err
+    local body = nil
+    local encoding = ngx_var.http_transfer_encoding
+
+    if not opts.request_buffering then
+        if http_ver() ~= 1.1 then
+            ngx.log(ngx.ERR, "bad http version")
+            ngx.exit(ngx.HTTP_BAD_REQUEST)
+        end
+
+        if type(encoding) == "table" then
+            encoding = encoding[1]
+        end
+
+        _, err = handle_expect()
+        if err then
+            ngx.log(ngx.ERR, err)
+            return ngx.exit(ngx.HTTP_BAD_GATEWAY)
+        end
 
+        if encoding and encoding:lower() == "chunked" then
+            -- The default ngx reader does not support chunked request
+            -- so we will need to get the raw request socket and manually
+            -- decode the chunked request
+            sock, err = ngx.req.socket(true)
+            body = get_client_body_reader(sock, chunksize, true)
+        else
+            sock, err = ngx.req.socket()
+            body = get_client_body_reader(sock, chunksize)
+        end
+
+        if not sock then
+            ngx.log(ngx.ERR, "unable to obtain request socket: ", err)
+            return ngx.exit(ngx.HTTP_BAD_GATEWAY)
+        end
+
+        req.body = body
+        req.sock = sock
+    else
+        -- This is needed to call ngx.req.get_body_data() below.
+        ngx.req.read_body()
         -- We cannot use resty.http's .get_client_body_reader().
         -- In POST requests with HTTPS, the result of that call is nil, and it
         -- results in a time-out.
@@ -123,37 +189,53 @@ local function forward_https_request(proxy_uri, uri, skip_https_connect)
         -- read and need to be cached in a local file. This request will return
         -- nil, so after this we need to read the temp file.
         -- https://github.com/openresty/lua-nginx-module#ngxreqget_body_data
-        body = ngx.req.get_body_data(),
-        proxy_uri = proxy_uri
-    }
+        req.body = ngx_req.get_body_data()
+
+        if not req.body then
+            local temp_file_path = ngx.req.get_body_file()
+            ngx.log(ngx.INFO, "HTTPS Proxy: Request body is bigger than client_body_buffer_size, read the content from path='", temp_file_path, "'")
+
+            if temp_file_path then
+              body, err = file_reader(temp_file_path)
+              if err then
+                ngx.log(ngx.ERR, "HTTPS proxy: Failed to read temp body file, err: ", err)
+                ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
+              end
+              if encoding == "chunked" then
+                -- If the body is smaller than "client_boby_buffer_size" the Content-Length header is
+                -- set based on the size of the buffer. However, when the body is rendered to a file,
+                -- we will need to calculate and manually set the Content-Length header based on the
+                -- file size
+                modify_chunked_request_headers(req, temp_file_path)
+              end
+              req.body = body
+            end
+        end
 
-    if not request.body then
-        local temp_file_path = ngx.req.get_body_file()
-        ngx.log(ngx.INFO, "HTTPS Proxy: Request body is bigger than client_body_buffer_size, read the content from path='", temp_file_path, "'")
-
-        if temp_file_path then
-          local body, err = file_reader(temp_file_path)
-          if err then
-            ngx.log(ngx.ERR, "HTTPS proxy: Failed to read temp body file, err: ", err)
-            ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
-          end
-          if ngx.var.http_transfer_encoding == "chunked" then
-            -- If the body is smaller than "client_boby_buffer_size" the Content-Length header is
-            -- set based on the size of the buffer. However, when the body is rendered to a file,
-            -- we will need to calculate and manually set the Content-Length header based on the
-            -- file size
-            modify_chunked_request_headers(request, temp_file_path)
-          end
-          request.body = body
+        -- The whole request is buffered in the memory so remove the Transfer-Encoding: chunked
+        if ngx.var.http_transfer_encoding == "chunked" then
+            set_header(req.headers, "Transfer-Encoding", nil)
         end
     end
+end
+
+local function forward_https_request(proxy_uri, uri, proxy_opts)
+    local method = ngx.req.get_method()
+    local chunksize = 32 * 1024
 
-    -- The whole request is buffered in the memory so remove the Transfer-Encoding: chunked
-    if ngx.var.http_transfer_encoding == "chunked" then
-        set_header(request.headers, "Transfer-Encoding", nil)
+    local request = {
+        uri = uri,
+        method = ngx.req.get_method(),
+        headers = ngx.req.get_headers(0, true),
+        path = format('%s%s%s', ngx.var.uri, ngx.var.is_args, ngx.var.query_string or ''),
+        proxy_uri = proxy_uri
+    }
+
+    if http_methods_with_body[method] then
+        get_request_body(request, proxy_opts)
     end
 
-    local httpc, err = http_proxy.new(request, skip_https_connect)
+    local httpc, err = http_proxy.new(request, proxy_opts.skip_https_connect)
 
     if not httpc then
         ngx.log(ngx.ERR, 'could not connect to proxy: ',  proxy_uri, ' err: ', err)
@@ -165,8 +247,16 @@ local function forward_https_request(proxy_uri, uri, skip_https_connect)
     res, err = httpc:request(request)
 
     if res then
-        httpc:proxy_response(res)
-        httpc:set_keepalive()
+        if not proxy_opts.request_buffering then
+            local bytes, err = send_response(request.sock, res, chunksize)
+            if not bytes then
+                ngx.log(ngx.ERR, "failed to send response: ", err)
+                return ngx.exit(ngx.HTTP_BAD_GATEWAY)
+            end
+        else
+            httpc:proxy_response(res)
+            httpc:set_keepalive()
+        end
     else
         ngx.log(ngx.ERR, 'failed to proxy request to: ', proxy_uri, ' err : ', err)
         return ngx.exit(ngx.HTTP_BAD_GATEWAY)
@@ -204,7 +294,11 @@ function _M.request(upstream, proxy_uri)
         return
     elseif uri.scheme == 'https' then
         upstream:rewrite_request()
-        forward_https_request(proxy_uri, uri, upstream.skip_https_connect)
+        local proxy_opts = {
+            skip_https_connect = upstream.skip_https_connect,
+            request_buffering = upstream.request_buffering
+        }
+        forward_https_request(proxy_uri, uri, proxy_opts)
         return ngx.exit(ngx.OK) -- terminate phase
     else
         ngx.log(ngx.ERR, 'could not connect to proxy: ',  proxy_uri, ' err: ', 'invalid request scheme')
diff --git a/gateway/src/apicast/upstream.lua b/gateway/src/apicast/upstream.lua
@@ -241,6 +241,8 @@ function _M:call(context)
           self:set_skip_https_connect_on_proxy();
         end
 
+        self.request_buffering = context.request_buffering
+
         http_proxy.request(self, proxy_uri)
     else
         local err = self:rewrite_request()
diff --git a/gateway/src/resty/http/raw_request.lua b/gateway/src/resty/http/raw_request.lua
@@ -0,0 +1,47 @@
+local httpc = require "resty.resolver.http"
+
+local _M = {
+}
+
+local cr_lf = "\r\n"
+
+-- 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
+local function 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
+
+function _M.get_client_body_reader(sock, chunksize, is_chunked)
+    if is_chunked then
+        return chunked_reader(sock, chunksize)
+    else
+        return httpc:get_client_body_reader(chunksize, sock)
+    end
+end
+
+return _M
diff --git a/gateway/src/resty/http/raw_response.lua b/gateway/src/resty/http/raw_response.lua
@@ -0,0 +1,74 @@
+local _M = {
+}
+
+local cr_lf = "\r\n"
+
+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
+
+-- 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.
+-- The connection is closed if send() fails or when returning a non-zero
+function _M.send_response(sock, res, chunksize)
+    local bytes, err
+    chunksize = chunksize or 65536
+
+    if ngx.headers_sent then
+      return nil, "headers have already been sent"
+    end
+
+    if not sock then
+        return nil, "socket not initialized yet"
+    end
+
+    -- 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
+        return nil, "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
+            return nil, "failed to send header, err: " .. (err or "unknown")
+        end
+    end
+
+    -- End-of-header
+    bytes, err = send(sock, cr_lf)
+    if not bytes then
+        return nil, "failed to send end of header, err: " .. (err or "unknown")
+    end
+
+    -- Write body
+    if res.has_body then
+        local reader = res.body_reader
+        repeat
+            local chunk, read_err
+
+            chunk, read_err = reader(chunksize)
+            if read_err then
+                return nil, "failed to read response body, err: " .. (err or "unknown")
+            end
+
+            if chunk then
+                bytes, err = send(sock, chunk)
+                if not bytes then
+                    return nil, "failed to send response body, err: " .. (err or "unknown")
+                end
+            end
+        until not chunk
+    end
+end
+
+return _M
