feature: preserve body after parse (#63)

Co-authored-by: Johnny Wang <[email protected]>

Author: suikabreaker

README.markdown

@@ -161,6 +161,19 @@ not want to save the data on local file systems.
 
 [Back to TOC](#table-of-contents)
 
+Usage
+=====
+
+```lua
+local upload = require "resty.upload"
+local form, err = upload:new(self, chunk_size, max_line_size, preserve_body)
+```
+`chunk_size` defaults to 4096. It is the size used to read data from the socket.
+
+`max_line_size` defaults to 512. It is the size limit to read the chunked body header.
+
+By Default, `lua-resty-upload` will consume the request body. For proxy mode this means upstream will not see the body. When `preserve_body` is set to true, the request body will be preserved. Note that this option is not free. When enabled, it will double the memory usage of `resty.upload`.
+
 Author
 ======
 

lib/resty/upload.lua

@@ -7,6 +7,9 @@ local match = string.match
 local setmetatable = setmetatable
 local type = type
 local ngx_var = ngx.var
+local ngx_init_body = ngx.req.init_body
+local ngx_finish_body = ngx.req.finish_body
+local ngx_append_body = ngx.req.append_body
 -- local print = print
 
 
@@ -21,11 +24,57 @@ local STATE_READING_HEADER = 2
 local STATE_READING_BODY = 3
 local STATE_EOF = 4
 
-
 local mt = { __index = _M }
 
 local state_handlers
 
+local function wrapped_receiveuntil(self, until_str)
+    local iter, err_outer = self:old_receiveuntil(until_str)
+    if iter == nil then
+        ngx_finish_body()
+    end
+
+    local function wrapped(size)
+        local ret, err = iter(size)
+        if ret then
+            ngx_append_body(ret)
+        end
+
+        -- non-nil ret for call with no size or successful size call and nil ret
+        if (not size and ret) or (size and not ret and not err) then
+            ngx_append_body(until_str)
+        end
+        return ret, err
+    end
+
+    return wrapped, err_outer
+end
+
+
+local function wrapped_receive(self, arg)
+    local ret, err, partial = self:old_receive(arg)
+    if ret then
+        ngx_append_body(ret)
+
+    elseif partial then
+        ngx_append_body(partial)
+    end
+
+    if ret == nil then
+        ngx_finish_body()
+    end
+
+    return ret, err
+end
+
+
+local function req_socket_body_collector(sock)
+    sock.old_receiveuntil = sock.receiveuntil
+    sock.old_receive = sock.receive
+    sock.receiveuntil = wrapped_receiveuntil
+    sock.receive = wrapped_receive
+end
+
 
 local function get_boundary()
     local header = ngx_var.content_type
@@ -46,7 +95,7 @@ local function get_boundary()
 end
 
 
-function _M.new(self, chunk_size, max_line_size)
+function _M.new(self, chunk_size, max_line_size, preserve_body)
     local boundary = get_boundary()
 
     -- print("boundary: ", boundary)
@@ -62,6 +111,11 @@ function _M.new(self, chunk_size, max_line_size)
         return nil, err
     end
 
+    if preserve_body then
+        ngx_init_body(chunk_size)
+        req_socket_body_collector(sock)
+    end
+
     local read2boundary, err = sock:receiveuntil("--" .. boundary)
     if not read2boundary then
         return nil, err
@@ -79,7 +133,8 @@ function _M.new(self, chunk_size, max_line_size)
         read2boundary = read2boundary,
         read_line = read_line,
         boundary = boundary,
-        state = STATE_BEGIN
+        state = STATE_BEGIN,
+        preserve_body = preserve_body
     }, mt)
 end
 
@@ -104,6 +159,10 @@ local function discard_line(self)
 
     local dummy, err = read_line(1)
     if dummy then
+        if self.preserve_body then
+            ngx_finish_body()
+        end
+
         return nil, "line too long: " .. line .. dummy .. "..."
     end
 
@@ -179,6 +238,10 @@ local function read_header(self)
 
     local dummy, err = read_line(1)
     if dummy then
+        if self.preserve_body then
+            ngx_finish_body()
+        end
+ 
         return nil, nil, "line too long: " .. line .. dummy .. "..."
     end
 

t/sanity.t

@@ -648,3 +648,125 @@ read: ["eof"]
 --- no_error_log
 [error]
 
+
+
+=== TEST 11: body preserve off
+--- http_config eval: $::HttpConfig
+--- config
+    location /t {
+        content_by_lua '
+            local upload = require "resty.upload"
+            local ljson = require "ljson"
+
+            local form = upload:new(5)
+
+            form:set_timeout(1000) -- 1 sec
+
+            while true do
+                local typ, res, err = form:read()
+                if not typ then
+                    ngx.say("failed to read: ", err)
+                    return
+                end
+
+                ngx.say("read: ", ljson.encode({typ, res}))
+
+                if typ == "eof" then
+                    break
+                end
+            end
+
+            local typ, res, err = form:read()
+            ngx.say("read: ", ljson.encode({typ, res}))
+
+            ngx.say("remain body: ", ngx.req.get_body_data(), ",", ngx.req.get_body_file())
+        ';
+    }
+--- more_headers
+Content-Type: multipart/form-data; boundary=---------------------------820127721219505131303151179
+--- request eval
+qq{POST /t\n-----------------------------820127721219505131303151179\r
+Content-Disposition: form-data; name="file1"; filename="a.txt"\r
+Content-Type: text/plain\r
+\r
+Hello, world\r\n-----------------------------820127721219505131303151179\r
+Content-Disposition: form-data; name="test"\r
+\r
+value\r
+\r\n-----------------------------820127721219505131303151179--\r
+}
+--- response_body
+read: ["header",["Content-Disposition","form-data; name=\"file1\"; filename=\"a.txt\"","Content-Disposition: form-data; name=\"file1\"; filename=\"a.txt\""]]
+read: ["header",["Content-Type","text/plain","Content-Type: text/plain"]]
+read: ["body","Hello"]
+read: ["body",", wor"]
+read: ["body","ld"]
+read: ["part_end"]
+read: ["header",["Content-Disposition","form-data; name=\"test\"","Content-Disposition: form-data; name=\"test\""]]
+read: ["body","value"]
+read: ["body","\r\n"]
+read: ["part_end"]
+read: ["eof"]
+read: ["eof"]
+remain body: nil,nil
+--- no_error_log
+[error]
+
+
+
+=== TEST 12: body preserve on
+--- http_config eval: $::HttpConfig
+--- config
+    location /t {
+        content_by_lua '
+            local original_len = ngx.req.get_headers()["Content-Length"]
+            local upload = require "resty.upload"
+
+            local form = upload:new(5, nil, true)
+
+            form:set_timeout(1000) -- 1 sec
+
+            while true do
+                local typ, res, err = form:read()
+                if not typ then
+                    ngx.say("failed to read: ", err)
+                    return
+                end
+
+                if typ == "eof" then
+                    break
+                end
+            end
+
+            local typ, res, err = form:read()
+
+            local body = ngx.req.get_body_data()
+            local new_len
+            if body then
+                new_len = #body
+            else
+                new_len = io.open(ngx.req.get_body_file(), "r"):seek("end")
+            end
+            ngx.say("content_length: ", original_len)
+            ngx.say("remain body length: ", new_len)
+        ';
+    }
+--- more_headers
+Content-Type: multipart/form-data; boundary=---------------------------820127721219505131303151179
+--- request eval
+qq{POST /t\n-----------------------------820127721219505131303151179\r
+Content-Disposition: form-data; name="file1"; filename="a.txt"\r
+Content-Type: text/plain\r
+\r
+Hello, world\r\n-----------------------------820127721219505131303151179\r
+Content-Disposition: form-data; name="test"\r
+\r
+value\r
+\r\n-----------------------------820127721219505131303151179--\r
+}
+--- response_body
+content_length: 336
+remain body length: 336
+
+--- no_error_log
+[error]