optional lf line break compatibility

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
 
 
@@ -46,40 +49,64 @@ local function get_boundary()
 end
 
 
-function _M.new(self, chunk_size, max_line_size)
+function _M.new(self, chunk_size, max_line_size, restore_body_buffer, lf_line_break)
     local boundary = get_boundary()
 
     -- print("boundary: ", boundary)
 
     if not boundary then
         return nil, "no boundary defined in Content-Type"
     end
-
     -- print('boundary: "', boundary, '"')
 
     local sock, err = req_socket()
     if not sock then
         return nil, err
     end
 
-    local read2boundary, err = sock:receiveuntil("--" .. boundary)
+    if restore_body_buffer then
+        ngx_init_body(chunk_size)
+    end
+
+    boundary = "--" .. boundary
+    local read2boundary, err = sock:receiveuntil(boundary)
     if not read2boundary then
+        if restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, err
     end
 
-    local read_line, err = sock:receiveuntil("\r\n")
+    -- note that it matters when restore_body_buffer
+    -- because we shuold not change body length
+    local line_break
+    if lf_line_break then
+        line_break = "\n"
+    else
+        line_break = "\r\n"
+    end
+
+    local read_line, err = sock:receiveuntil(line_break)
     if not read_line then
+
+        if restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, err
     end
 
     return setmetatable({
         sock = sock,
         size = chunk_size or CHUNK_SIZE,
         line_size = max_line_size or MAX_LINE_SIZE,
+        restore_body_buffer = restore_body_buffer,
         read2boundary = read2boundary,
         read_line = read_line,
         boundary = boundary,
-        state = STATE_BEGIN
+        state = STATE_BEGIN,
+        line_break = line_break
     }, mt)
 end
 
@@ -99,15 +126,31 @@ local function discard_line(self)
 
     local line, err = read_line(self.line_size)
     if not line then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, err
     end
 
+    if self.restore_body_buffer then
+        ngx_append_body(line)
+        ngx_append_body(self.line_break)
+    end
+
     local dummy, err = read_line(1)
     if dummy then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, "line too long: " .. line .. dummy .. "..."
     end
 
     if err then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
         return nil, err
     end
 
@@ -122,9 +165,17 @@ local function discard_rest(self)
     while true do
         local dummy, err = sock:receive(size)
         if err and err ~= 'closed' then
+            if self.restore_body_buffer then
+                ngx_finish_body()
+            end
+
             return nil, err
         end
 
+        if self.restore_body_buffer then
+            ngx_append_body(dummy)
+        end
+
         if not dummy then
             return 1
         end
@@ -137,35 +188,72 @@ local function read_body_part(self)
 
     local chunk, err = read2boundary(self.size)
     if err then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, nil, err
     end
 
+    -- everything OK we got another body chunk
+    if chunk then
+        if self.restore_body_buffer then
+            ngx_append_body(chunk)
+            ngx_append_body(self.boundary)
+        end
+
+        if self.boundary:sub(1,1) == '\n' and chunk:sub(-1,-1) == '\r' then
+            chunk = chunk:sub(1,-2)
+        end
+
+        return "body", chunk
+    end
+
+    -- boundary not found, maybe end of body
     if not chunk then
         local sock = self.sock
+        local data = sock:receive(1)
 
-        local data = sock:receive(2)
-        if data == "--" then
+        -- neither -- or line_break, something's wrong
+        if data ~= "-" and data ~= self.line_break:sub(1,1) then
+            if self.restore_body_buffer then
+                ngx_append_body(data)
+            end
+            local ok, err = discard_line(self)
+            if not ok then
+                return nil, nil, err
+            end
+        end
+
+        if data ~= self.line_break then
+            local next = sock:receive(1)
+            data = data .. next
+        end
+
+        if self.restore_body_buffer then
+            ngx_append_body(data)
+        end
+
+        if data == self.line_break then
+            self.state = STATE_READING_HEADER
+            return "part_end"
+        elseif data == "--" then
             local ok, err = discard_rest(self)
             if not ok then
                 return nil, nil, err
             end
 
             self.state = STATE_EOF
             return "part_end"
-        end
-
-        if data ~= "\r\n" then
+        else
+            -- something's wrong
             local ok, err = discard_line(self)
             if not ok then
                 return nil, nil, err
             end
         end
-
-        self.state = STATE_READING_HEADER
-        return "part_end"
     end
 
-    return "body", chunk
 end
 
 
@@ -174,18 +262,38 @@ local function read_header(self)
 
     local line, err = read_line(self.line_size)
     if err then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, nil, err
     end
 
+    if self.restore_body_buffer then
+        ngx_append_body(line)
+        ngx_append_body(self.line_break)
+    end
+
     local dummy, err = read_line(1)
     if dummy then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, nil, "line too long: " .. line .. dummy .. "..."
     end
 
     if err then
+        if self.restore_body_buffer then
+            ngx_finish_body()
+        end
+
         return nil, nil, err
     end
 
+    if self.line_break == "\n" and line:sub(-1,-1) == '\r' then
+        line = line:sub(1,-2)
+    end
     -- print("read line: ", line)
 
     if line == "" then
@@ -203,7 +311,11 @@ local function read_header(self)
 end
 
 
-local function eof()
+local function eof(self)
+    if self.restore_body_buffer then
+        ngx_finish_body()
+    end
+
     return "eof", nil
 end
 
@@ -232,7 +344,15 @@ local function read_preamble(self)
     while true do
         local preamble = read2boundary(size)
         if not preamble then
+            if self.restore_body_buffer then
+                ngx_append_body(self.boundary)
+            end
+
             break
+
+        else if self.restore_body_buffer then
+                ngx_append_body(preamble)
+            end
         end
 
         -- discard the preamble data chunk
@@ -244,7 +364,8 @@ local function read_preamble(self)
         return nil, nil, err
     end
 
-    local read2boundary, err = sock:receiveuntil("\r\n--" .. self.boundary)
+    self.boundary = self.line_break .. self.boundary
+    local read2boundary, err = sock:receiveuntil(self.boundary)
     if not read2boundary then
         return nil, nil, err
     end