diff --git a/lib/resty/upload.lua b/lib/resty/upload.lua index 37da961..354a0d2 100644 --- a/lib/resty/upload.lua +++ b/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,7 +49,7 @@ 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) @@ -54,7 +57,6 @@ function _M.new(self, chunk_size, max_line_size) if not boundary then return nil, "no boundary defined in Content-Type" end - -- print('boundary: "', boundary, '"') local sock, err = req_socket() @@ -62,13 +64,36 @@ function _M.new(self, chunk_size, max_line_size) 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 @@ -76,10 +101,12 @@ function _M.new(self, chunk_size, max_line_size) 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,14 +188,56 @@ 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 @@ -152,20 +245,15 @@ local function read_body_part(self) 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