linux007
9/13/2013 - 6:06 AM

lua-resty-multipart.lua

--[[
usage: 
    local form = multipart:new()
    -- set any opts
    form:parse()

    -- form.fields 
]]
local ngx          = require("ngx")
local upload       = require("resty.upload")
local say          = ngx.say
local sock         = ngx.req.sock
local setmetatable = setmetatable
local pcall        = pcall
local type         = type
local gsub         = string.gsub
local pairs        = pairs
local log          = say
local concat       = table.concat
local insert       = table.insert
local fopen        = io.open
local strlen       = string.len
local remove       = os.remove

module(...)

_M.__index = _M

local worker_pid = ngx.var.pid
local tmp_path = "/tmp/resty-multipart"
local header_handlers = {
    ['Content-Type'] = function(self, data, partial)
        self:get_cur_field().mime = data[2]
    end, 
    ['Content-Disposition'] = function(self, data, partial)
        local query = gsub(gsub(data[2], "; ", "&"), '"', '')
        query = ngx.decode_args(query)

        if type(query) == 'table' then
            --todo: parse sth like "model[name]" field
            if query.name then
                self.cur_field = query.name
                self.fields[query.name] = query

                local f = self:get_cur_field()
                f.value = {}
                if query.filename then
                    f.is_file = true
                    f.file_size = 0
                    f.tmpfile = self:gen_tmpfile()
                    local err
                    f.tfp, err = fopen(f.tmpfile, "w+")
                    if not f.tfp then
                        error(concat({f.tmpfile, ":", err}))
                    end
                end
            end
        end
    end,
}

function new(self, timeout, chunk_size)
    local form, err = upload:new(chunk_size)
    if not form then
        return nil, err
    end

    form:set_timeout(timeout or 1000)
    
    local obj = {
        form = form, 
        tmp_path = tmp_path,
        cur_field = nil, 
        fields = {},
        header_handlers = setmetatable({}, {__index=header_handlers}), 
    }

    return setmetatable(obj, self)
end

function parse(self)
    while true do
        local typ, data, partial = self.form:read()
        if typ then
            self:handle(typ, data, partial)
        end

        if typ=="eof" then break end
    end
end

function handle(self, typ, data, partial)
    local handler = "handler_"..typ
    if self[handler] then
        local status, err = pcall(self[handler], self, data, partial)
        if err then
            log(err)
        end
    end
end

function handler_header(self, data, partial)
    if type(data) == 'table' then
        local handler = self.header_handlers[data[1]]
        if handler then
            handler(self, data, partial)
        end
    elseif type(data) == 'string' then
        -- pass now 
    end
end

function handler_part_end(self)
    local f = self:get_cur_field()
    if f.is_file then
        if f.tfp then
            f.tfp:close()
        end
    else
        f.value = concat(f.value)
    end
    self.cur_field = nil
end

function handler_body(self, data, partial)
    local f = self:get_cur_field()
    if not f.is_file then
        insert(f.value, data)
    else
        -- write to tmpfile
        if f.tfp then
            f.tfp:write(data)
            f.file_size = f.file_size+strlen(data)
            ngx.sleep(0.001)
        end
    end
    
end

function clear_tmp_files(self)
    for k, v in pairs(self.fields) do
        if v.is_file then
            remove(v.tmpfile)   
        end
    end
end

function gen_tmpfile(self)
    ngx.update_time()
    local name_parts =  {self.tmp_path, '/', worker_pid,'-', ngx.now()}
    return concat(name_parts)
end

function get_cur_field(self)
    local f = self.fields[self.cur_field]
    if not f then
        error("cur_field not specified")
    end
    return f
end

local _class_mt = {
    __newindex = function(self, key, value) error"not acceptable" end   
}

return setmetatable(_M, _class_mt)