HTTP Multipart Encoded Body reader for ngx_lua

multipart.lua

-- Copyright (C) 2012 Zhang "agentzh" Yichun (章亦春)
-- Copyright (C) 2012 Matthieu Tourne

module("multipart", package.seeall)

local STATE_BEGIN = 1
local STATE_READING_HEADER = 2
local STATE_READING_BODY = 3
local STATE_ERROR = 4

local mt = { __index = multipart }

local req_socket = ngx.req.socket
local state_handlers

function new(self, chunk_size)
    local boundary = get_boundary()
    if not boundary then
        return nil, "no boundary defined in Content-Type"
    end

    boundary = '--' .. boundary

    local sock, err = req_socket()
    if not sock then
        return nil, err
    end

    local read2boundary, err = sock:receiveuntil("\r\n" .. boundary,
                                                 { inclusive = false })
    if not read2boundary then
        return nil, err
    end

    local read2data, err = sock:receiveuntil("\r\n\r\n",
                                             { inclusive = true })
    if not read2data then
        return nil, err
    end

    local read_line, err = sock:receiveuntil("\r\n",
                                             { inclusive = true })
    if not read_line then
        return nil, err
    end

    return setmetatable({
        sock = sock,
        size = chunk_size or 4096,
        eof = false,
        read2boundary = read2boundary,
        read_line = read_line,
        read2data = read2data,
        boundary = boundary,
        state = STATE_BEGIN,
        output_buffer = {}
    }, mt)
end


function set_timeout(self, timeout)
    local sock = self.sock
    if not sock then
        return nil, "not initialized"
    end

    return sock:settimeout(timeout)
end


function read(self)
   if self.eof then
      if #self.output_buffer > 0 then
         return "err", self:output()
      else
         return "eof", nil, nil
      end
   end

   local state = self.state
   local handler = state_handlers[state]
   if handler then
      return handler(self)
   end

   return nil, nil, "bad state: " .. state
end


function read_preamble(self)
    local sock = self.sock
    if not sock then
        return nil, nil, "not initialized"
    end

    local size = self.size

    -- discard the preamble data chunk
    local preamble, err = receive_wrap(self, sock.receive,
                                       { sock, #self.boundary })
    if err then
       return nil, nil, err
    end

    ngx.log(ngx.DEBUG, 'preamble: ', preamble)

    table.insert(self.output_buffer, preamble)

    if preamble ~= self.boundary then
       self.state = STATE_ERROR
       return "err", self:output()
    end

    local ok, err = discard_line(self)
    if err then
       return nil, nil, err
    end
    if not ok then
       return "err", self:output()
    end

    self.state = STATE_READING_HEADER
    return "preamble", self:output()
end


function discard_line(self)
   local read_line = self.read_line
   local size = self.size

   local line, err = self:receive_wrap(read_line, {size})
   if err then
      return nil, err
   end

   ngx.log(ngx.DEBUG, 'line: ', line)

   table.insert(self.output_buffer, line)

   local dummy, err = self:receive_wrap(read_line, {1})
   if err then
      return nil, err
   end

   ngx.log(ngx.DEBUG, 'dummy: ', dummy)

   if dummy then
      ngx.log(ngx.ERR, 'discard_line() line too long')
      table.insert(self.output_buffer, dummy)
      self.state = STATE_ERROR
      return nil, nil
   end

   return 1, nil
end


function read_header(self)
   local read2data = self.read2data

   local header, err = self:receive_wrap(read2data, {self.size})
   if err then
      return nil, nil, err
   end

   table.insert(self.output_buffer, header)

   ngx.log(ngx.DEBUG, 'header: ', header)

   local dummy, err = self:receive_wrap(read2data, {1})
   if err then
      return nil, nil, err
   end

   if dummy then
      ngx.log(ngx.ERR, 'multipart header line too long')
      self.state = STATE_ERROR
      table.insert(self.output_buffer, dummy)
      return "err", self:output()
   end

   self.state = STATE_READING_BODY
   if header then
      return "header", self:output()
   end
end


function read_body_part(self)
   local read2boundary = self.read2boundary

   local chunk, err = self:receive_wrap(read2boundary, {self.size})
   if err then
      return nil, nil, err
   end

   if not chunk then
      local sock = self.sock

      table.insert(self.output_buffer, '\r\n')
      table.insert(self.output_buffer, self.boundary)

      local data, err = self:receive_wrap(sock.receive, {sock, 2})
      if err then
         return nil, nil, err
      end

      table.insert(self.output_buffer, data)

      if data == "--" then
         -- set to error state, what could be after the last boundary
         self.state = STATE_ERROR
         return "end", self:output()
      end

      if data ~= "\r\n" then
         ok, err = discard_line(self)
         if err then
            return nil, nil, err
         end

         if not ok then
            self.state = STATE_ERROR
            return "err", self:output()
         end
      end

      self.state = STATE_READING_HEADER
      return "boundary", self:output()
   end

   table.insert(self.output_buffer, chunk)

   return "body", self:output()
end

function read_error(self)
   local sock = self.sock

   local data, err = self:receive_wrap(sock.receive, {sock, self.size})
   if err then
      return nil, nil, err
   end
   table.insert(self.output_buffer, data)
   return "err", self:output()
end

function receive_wrap(self, read_callback, params)
   if self.eof then
      return nil, nil
   end

   local data, err, partial = read_callback(unpack(params))

   if err == 'closed' then
      -- eof
      self.eof = true
      return partial, nil
   end
   if err then
      return nil, err
   end
   return data
end

function output(self)
   local data = table.concat(self.output_buffer)
   self.output_buffer = {}
   return data
end

function get_boundary()
    local header = ngx.var.content_type
    if not header then
        return nil
    end

    return string.match(header, ";%s+boundary=(%S+)")
end


state_handlers = {
    read_preamble,
    read_header,
    read_body_part,
    read_error
}