--[[ $%BEGINLICENSE%$
 Copyright (c) 2007, 2008, Oracle and/or its affiliates. All rights reserved.

 This program is free software; you can redistribute it and/or
 modify it under the terms of the GNU General Public License as
 published by the Free Software Foundation; version 2 of the
 License.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 02110-1301  USA

 $%ENDLICENSE%$ --]]


---
--[[ 
    (Not so) simple example of a run-time module loader for MySQL Proxy
    
    Usage: 
    1. load this module in the Proxy.
    
    2. From a client, run the command 
       PLOAD name_of_script
       and you can use the features implemented in the new script
       immediately.
    
    3. Successive calls to PLOAD will load a new one.
    
    4. Previous scripts will still be active, and the loader will
    use all of them in sequence to see if one handles the 
    request

    5. CAVEAT. If your script has a read_query function that 
    *always* returns a non null value (e.g. a logging feature)
    then other modules, although loaded, are not used.
    If you have any such modules, you must load them at the
    end of the sequence.

    6. to remove a module, use
    UNLOAD module_name

    IMPORTANT NOTICE:
    the proxy will try to load the file in the directory where
    the proxy (NOT THE CLIENT) was started.
    To use modules not in such directory, you need to
    specify an absolute path.
--]]

local VERSION = '0.1.3'

local DEBUG = os.getenv('DEBUG') or 0
DEBUG = DEBUG + 0

---
-- print_debug()
-- conditionally print a message
--
function proxy.global.print_debug (msg)
    if DEBUG > 0 then
        print(msg)
    end
end

---
-- Handlers for MySQL Proxy hooks
-- Initially, they are void.
-- If the loaded module has implemented any
-- functions, they are associated with these handlers
--
if proxy.global.loaded_handlers == nil then
    proxy.global.loaded_handlers = {
        rq   = {},
        rqr  = {},
        dc   = {},
        cs   = {},
        rh   = {},
        ra   = {},
        rar  = {},
    }
end

---
-- list of functions loaded from user modules
--
if proxy.global.handled_function == nil then
    proxy.global.handled_functions = {
        rq  = 'read_query',
        rqr = 'read_query_result',
        cs  = 'connect_server',
        rh  = 'read_handshake',
        ra  = 'read_auth',
        rar = 'read_auth_result',
        dc  = 'disconnect_client',
    }
end

if proxy.global.handler_status == nil then
    proxy.global.handler_status = {}
end

local funcs_to_create = {
    cs  = 0,
    ra  = 1,
    rar = 1,
    rh  = 1,
    dc  = 0,
}

---
-- creates the hooks for Proxy handled functions
--
for id, has_param in pairs(funcs_to_create) do
    local parameter = has_param == 1 and 'myarg' or ''
    local fstr = string.format(
[[
    function %s(%s)
    if #proxy.global.loaded_handlers['%s'] then
        for i, h in pairs(proxy.global.loaded_handlers['%s'])
        do
            if h then
                proxy.global.print_debug ( 'handling "%s" using ' .. h.name)
                local result = h.handler(%s)
                if result then
                    return result
                end
            end
        end
    end
    -- return
end
]],
    proxy.global.handled_functions[id], 
    parameter,
    id, id, proxy.global.handled_functions[id], parameter
)
    proxy.global.print_debug ('creating function ' .. proxy.global.handled_functions[id])
    assert(loadstring(fstr))()
end

---
-- an error string to describe the error occurring
--
local load_multi_error_str = ''

---
-- set_error()
--
-- sets the error string 
--
-- @param msg the message to be assigned to the error string
--
function set_error(msg)
    load_multi_error_str = msg
    proxy.global.print_debug (msg)
    return nil
end

---
-- file_exists()
--
-- checks if a file exists
--
-- @param fname the file name
--
function file_exists(fname)
    local fh=io.open( fname, 'r')
    if fh then 
        fh:close()
        return true
    else
        return false
    end
end

local module_counter = 0

---
--  get_module()
--  
--  This function scans an existing lua file, and turns it into
--  a closure exporting two function handlers, one for
--  read_query() and one for read_query_result().
--  If the input script does not contain the above functions,
--  get_module fails.
--
--  @param module_name the name of the existing script
--
function get_module(module_name)
    -- 
    -- assumes success
    --
    load_multi_error_str = ''

    --
    -- the module is copied to a temporary file
    -- on a given directory
    --
    module_counter = module_counter + 1
    local new_module = 'tmpmodule' .. module_counter
    local tmp_dir   = '/tmp/'
    local new_filename = tmp_dir .. new_module .. '.lua'
    local source_script = module_name
    if not source_script:match('.lua$') then
        source_script = source_script .. '.lua'
    end
    -- 
    -- if the new module does not exist
    -- an error is returned
    --
    if not file_exists(source_script) then
        set_error('file not found ' .. source_script)
        return
    end
    --
    -- if the module directory is not on the search path,
    -- we need to add it
    --
    if not package.path:match(tmp_dir) then
        package.path = tmp_dir .. '?.lua;' .. package.path
    end
    --
    -- Make sure that the module is not loaded.
    -- If we don't remove it from the list of loaded modules,
    -- subsequent load attempts will silently fail
    --
    package.loaded[new_module] = nil 
    local ofh = io.open(new_filename, 'w')
    --
    -- Writing to the new package, starting with the
    -- header and the wrapper function
    --
    ofh:write( string.format(
                 "module('%s', package.seeall)\n"
              .. "function make_funcs()\n" , new_module)
    )
    local found_funcs = {}
    --
    -- Copying contents from the original script
    -- to the new module, and checking for the existence
    -- of the handler functions
    --
    for line in io.lines(source_script) do
        ofh:write(line .. "\n")
        for i,v in pairs(proxy.global.handled_functions) do
            if line:match('^%s*function%s+' ..v .. '%s*%(') then
                found_funcs[i] = v
                    break
            end
        end
    end
    --
    -- closing the wrapper on the new module
    --
    local return_value = ''
    for i,v in pairs(found_funcs) do
        return_value = return_value .. i .. ' = ' .. v .. ', '
    end
    ofh:write(
           'return { ' ..  return_value ..  '}\n' .. 
        'end\n'
    )
    ofh:close()
    --
    -- Final check. If the handlers were not found, the load fails
    --
    --
    if (found_one == false ) then
        set_error('script ' .. source_script .. 
            ' does not contain a proxy handled function')
        return
    end
    --
    -- All set. The new module is loaded
    -- with a new function that will return the handlers
    --
    local result = require(new_module)
    os.remove(new_filename)
    return result
end

---
-- simple_dataset()
--
-- returns a dataset made of a header and a message
--
-- @param header the column name
-- @param message the contents of the column
function proxy.global.simple_dataset (header, message) 
    proxy.response.type = proxy.MYSQLD_PACKET_OK
    proxy.response.resultset = {
        fields = {{type = proxy.MYSQL_TYPE_STRING, name = header}},
        rows = { { message} }
    }
    return proxy.PROXY_SEND_RESULT
end

---
-- make_regexp_from_command()
--
-- creates a regular expression for fast scanning of the command
-- 
-- @param cmd the command to be converted to regexp
--
function proxy.global.make_regexp_from_command(cmd, options)
    local regexp= '^%s*';
    for ch in cmd:gmatch('(.)') do
        regexp = regexp .. '[' .. ch:upper() .. ch:lower() .. ']' 
    end
    if options and options.capture then
        regexp = regexp  .. '%s+(%S+)'
    end
    return regexp
end

-- 
-- The default command for loading a new module is PLOAD
-- You may change it through an environment variable
--
local proxy_load_command    = os.getenv('PLOAD')  or 'pload'
local proxy_unload_command  = os.getenv('PUNLOAD')  or 'punload'
local proxy_help_command    = os.getenv('PLOAD_HELP')  or 'pload_help'
local currently_using       = 0

local pload_regexp          = proxy.global.make_regexp_from_command(proxy_load_command, {capture = 1})
local punload_regexp        = proxy.global.make_regexp_from_command(proxy_unload_command,{ capture = 1} )
local pload_help_regexp     = proxy.global.make_regexp_from_command(proxy_help_command,{ capture = nil} )
local pload_help_dataset    = {
    {proxy_load_command   .. ' module_name',   'loads a given module'},
    {proxy_unload_command .. 'module_name',    'unloads and existing module'},
    {proxy_help_command,    'shows this help'},
}

---
-- removes a module from the loaded list
--
function remove_module (module_name)
    local found_module = false
    local to_delete = { loaded = {}, status = {} }
    for i,lmodule in pairs(proxy.global.handler_status) do
        if i == module_name then
            found_module = true
            local counter = 0
            for j,h in pairs(lmodule) do
                -- proxy.global.print_debug('removing '.. module_name .. ' (' .. i ..') ' .. h.id .. ' -> ' .. h.ndx )
                to_delete['loaded'][h.id] = h.ndx
                counter = counter + 1
                to_delete['status'][i] = counter
            end
        end
    end
    for i,v in pairs (to_delete['loaded']) do
        table.remove(proxy.global.loaded_handlers[i], v)
    end
    for i,v in pairs (to_delete['status']) do
        table.remove(proxy.global.handler_status[i], v)
    end
    if found_module == false then
        return proxy.global.simple_dataset(module_name, 'NOT FOUND')
    end
    return proxy.global.simple_dataset(module_name, 'unloaded')
end

---
-- creates a dataset from a list of header names
-- and a list of rows
-- @param header a list of field names
-- @param dataset a list of row contents 
function proxy.global.make_dataset (header, dataset) 
    proxy.response.type = proxy.MYSQLD_PACKET_OK

    proxy.response.resultset = {
        fields = {},
        rows = {}
    }
    for i,v in pairs (header) do
        table.insert(proxy.response.resultset.fields, {type = proxy.MYSQL_TYPE_STRING, name = v})
    end
    for i,v in pairs (dataset) do
        table.insert(proxy.response.resultset.rows, v )
    end
    return proxy.PROXY_SEND_RESULT
end

--
-- This function is called at each query.
-- The request for loading a new script is handled here
--
function read_query (packet)
    currently_using = 0
    if packet:byte() ~= proxy.COM_QUERY then
        return
    end
    local query = packet:sub(2)
    -- Checks if a PLOAD command was issued.
    -- A regular expresion check is faster than 
    -- doing a full tokenization. (Especially on large queries)
    --
    if (query:match(pload_help_regexp)) then
        return proxy.global.make_dataset({'command','description'}, pload_help_dataset)
    end
    local unload_module = query:match(punload_regexp)
    if (unload_module) then
        return remove_module(unload_module)
    end
    local new_module = query:match(pload_regexp)
    if (new_module) then
        --[[
           If a request for loading is received, then
           we attempt to load the new module using the
           get_module() function
        --]]
        local new_tablespace = get_module(new_module)
        if (new_tablespace) then
            local handlers = new_tablespace.make_funcs()
            --
            -- The new module must have at least  handlers for read_query()
            -- or disconnect_client. read_query_result() is optional. 
            -- The loading function returns nil if no handlers were found
            --
    
            proxy.global.print_debug('')
            proxy.global.handler_status[new_module] = {}
            for i,v in pairs( proxy.global.handled_functions) do
                local handler_str = type(handlers[i]) 
                proxy.global.print_debug (i .. ' ' .. handler_str )
                if handlers[i] then
                    table.insert(proxy.global.loaded_handlers[i] , 
                { name = new_module, handler = handlers[i]} )
                table.insert(proxy.global.handler_status[new_module], 
                { func = proxy.global.handled_functions[i], id=i, ndx = #proxy.global.loaded_handlers[i] })
                end
            end
            if handlers['rqr'] and not handlers['rq'] then
                table.insert(proxy.global.loaded_handlers['rq'] , nil )
            end
            if handlers['rq'] and not handlers['rqr'] then
                table.insert(proxy.global.loaded_handlers['rqr'] , nil )
            end
            -- 
            -- Returns a confirmation that a new  module was loaded
            --
            return proxy.global.simple_dataset('info', 'module "' .. new_module .. '" loaded' )
        else
            --
            -- The load was not successful.
            -- Inform the user
            --
            return proxy.global.simple_dataset('ERROR', 'load of "' 
                .. new_module .. '" failed (' 
                .. load_multi_error_str .. ')' )
        end
    end
    --
    -- If a handler was installed from a new module, it is called
    -- now. 
    --
    if #proxy.global.loaded_handlers['rq'] then
        for i,rqh in pairs(proxy.global.loaded_handlers['rq'])
        do
            proxy.global.print_debug ( 'handling "read_query" using ' .. i .. ' -> ' .. rqh.name)
            local result = rqh.handler(packet)
            if (result) then
                currently_using = i
                return result
            end
        end
    end
end

--
-- this function is called every time a result set is
-- returned after an injected query
--
function read_query_result(inj)
    -- 
    -- If the dynamically loaded module had an handler for read_query_result()
    -- it is called now
    --
    local rqrh = proxy.global.loaded_handlers['rqr'][currently_using]
    if rqrh then
        proxy.global.print_debug ( 'handling "read_query_result" using ' .. currently_using .. ' -> ' .. rqrh.name)
        local result = rqrh.handler(inj)
        if result then
            return result
        end
    end
end
