prosody/server.lua

--[[

        server.lua by blastbeat

        - this script contains the server loop of the program
        - other scripts can reg a server here

]]--

----------------------------------// DECLARATION //--

--// constants //--

local STAT_UNIT = 1 / ( 1024 * 1024 )    -- mb

--// lua functions //--

local function use( what ) return _G[ what ] end

local type = use "type"
local pairs = use "pairs"
local ipairs = use "ipairs"
local tostring = use "tostring"
local collectgarbage = use "collectgarbage"

--// lua libs //--

local table = use "table"
local coroutine = use "coroutine"

--// lua lib methods //--

local table_concat = table.concat
local table_remove = table.remove
local string_sub = use'string'.sub
local coroutine_wrap = coroutine.wrap
local coroutine_yield = coroutine.yield
local print = print;
local out_put = function () end --print;
local out_error = print;

--// extern libs //--

local luasec = require "ssl"
local luasocket = require "socket"

--// extern lib methods //--

local ssl_wrap = ( luasec and luasec.wrap )
local socket_bind = luasocket.bind
local socket_select = luasocket.select
local ssl_newcontext = ( luasec and luasec.newcontext )

--// functions //--

local loop
local stats
local addtimer
local closeall
local addserver
local firetimer
local closesocket
local removesocket
local wrapserver
local wraptcpclient
local wrapsslclient

--// tables //--

local listener
local readlist
local writelist
local socketlist
local timelistener

--// simple data types //--

local _
local readlen = 0    -- length of readlist
local writelen = 0    -- lenght of writelist

local sendstat= 0
local receivestat = 0

----------------------------------// DEFINITION //--

listener = { }    -- key = port, value = table
readlist = { }    -- array with sockets to read from
writelist = { }    -- arrary with sockets to write to
socketlist = { }    -- key = socket, value = wrapped socket
timelistener = { }

stats = function( )
    return receivestat, sendstat
end

wrapserver = function( listener, socket, ip, serverport, mode, sslctx )    -- this function wraps a server

    local dispatch, disconnect = listener.listener, listener.disconnect    -- dangerous

    local wrapclient, err

    if sslctx then
        if not ssl_newcontext then
            return nil, "luasec not found"
--        elseif not cfg_get "use_ssl" then
--            return nil, "ssl is deactivated"
        end
        if type( sslctx ) ~= "table" then
            out_error "server.lua: wrong server sslctx"
            return nil, "wrong server sslctx"
        end
        sslctx, err = ssl_newcontext( sslctx )
        if not sslctx then
            err = err or "wrong sslctx parameters"
            out_error( "server.lua: ", err )
            return nil, err
        end
        wrapclient = wrapsslclient
    else
        wrapclient = wraptcpclient
    end

    local accept = socket.accept
    local close = socket.close

    --// public methods of the object //--    

    local handler = { }

    handler.shutdown = function( ) end

    --[[handler.listener = function( data, err )
        return ondata( handler, data, err )
    end]]
    handler.ssl = function( )
        return sslctx and true or false
    end
    handler.close = function( closed )
        _ = not closed and close( socket )
        writelen = removesocket( writelist, socket, writelen )
        readlen = removesocket( readlist, socket, readlen )
        socketlist[ socket ] = nil
        handler = nil
    end
    handler.ip = function( )
        return ip
    end
    handler.serverport = function( )
        return serverport
    end
    handler.socket = function( )
        return socket
    end
    handler.receivedata = function( )
        local client, err = accept( socket )    -- try to accept
        if client then
            local ip, clientport = client:getpeername( )
            client:settimeout( 0 )
            local handler, client, err = wrapclient( listener, client, ip, serverport, clientport, mode, sslctx )    -- wrap new client socket
            if err then    -- error while wrapping ssl socket
                return false
            end
            out_put( "server.lua: accepted new client connection from ", ip, ":", clientport )
            return dispatch( handler )
        elseif err then    -- maybe timeout or something else
            out_put( "server.lua: error with new client connection: ", err )
            return false
        end
    end
    return handler
end

wrapsslclient = function( listener, socket, ip, serverport, clientport, mode, sslctx )    -- this function wraps a ssl cleint

    local dispatch, disconnect = listener.listener, listener.disconnect

    --// transform socket to ssl object //--

    local err
    socket, err = ssl_wrap( socket, sslctx )    -- wrap socket
    if err then
        out_put( "server.lua: ssl error: ", err )
        return nil, nil, err    -- fatal error
    end
    socket:settimeout( 0 )

    --// private closures of the object //--

    local writequeue = { }    -- buffer for messages to send

    local eol   -- end of buffer

    local sstat, rstat = 0, 0

    --// local import of socket methods //--

    local send = socket.send
    local receive = socket.receive
    local close = socket.close
    --local shutdown = socket.shutdown

    --// public methods of the object //--

    local handler = { }

    handler.getstats = function( )
        return rstat, sstat
    end

    handler.listener = function( data, err )
        return listener( handler, data, err )
    end
    handler.ssl = function( )
        return true
    end
    handler.send = function( _, data, i, j )
            return send( socket, data, i, j )
    end
    handler.receive = function( pattern, prefix )
            return receive( socket, pattern, prefix )
    end
    handler.shutdown = function( pattern )
        --return shutdown( socket, pattern )
    end
    handler.close = function( closed )
        close( socket )
        writelen = ( eol and removesocket( writelist, socket, writelen ) ) or writelen
        readlen = removesocket( readlist, socket, readlen )
        socketlist[ socket ] = nil
        out_put "server.lua: closed handler and removed socket from list"
    end
    handler.ip = function( )
        return ip
    end
    handler.serverport = function( )
        return serverport
    end
    handler.clientport = function( ) 
        return clientport
    end

    handler.write = function( data )
        if not eol then
            writelen = writelen + 1
            writelist[ writelen ] = socket
            eol = 0
        end
        eol = eol + 1
        writequeue[ eol ] = data
    end
    handler.writequeue = function( )
        return writequeue
    end
    handler.socket = function( )
        return socket
    end
    handler.mode = function( )
        return mode
    end
    handler._receivedata = function( )
        local data, err, part = receive( socket, mode )    -- receive data in "mode"
        if not err or ( err == "timeout" or err == "wantread" ) then    -- received something
            local data = data or part or ""
            local count = #data * STAT_UNIT
            rstat = rstat + count
            receivestat = receivestat + count
            out_put( "server.lua: read data '", data, "', error: ", err )
            return dispatch( handler, data, err )
        else    -- connections was closed or fatal error
            out_put( "server.lua: client ", ip, ":", clientport, " error: ", err )
            handler.close( )
            disconnect( handler, err )
            writequeue = nil
            handler = nil
            return false
        end
    end
    handler._dispatchdata = function( )    -- this function writes data to handlers
        local buffer = table_concat( writequeue, "", 1, eol )
        local succ, err, byte = send( socket, buffer )
        local count = ( succ or 0 ) * STAT_UNIT
        sstat = sstat + count
        sendstat = sendstat + count
        out_put( "server.lua: sended '", buffer, "', bytes: ", succ, ", error: ", err, ", part: ", byte, ", to: ", ip, ":", clientport )
        if succ then    -- sending succesful
            --writequeue = { }
            eol = nil
            writelen = removesocket( writelist, socket, writelen )    -- delete socket from writelist
            return true
        elseif byte and ( err == "timeout" or err == "wantwrite" ) then    -- want write
            buffer = string_sub( buffer, byte + 1, -1 )    -- new buffer
            writequeue[ 1 ] = buffer    -- insert new buffer in queue
            eol = 1
            return true
        else    -- connection was closed during sending or fatal error
            out_put( "server.lua: client ", ip, ":", clientport, " error: ", err )
            handler.close( )
            disconnect( handler, err )
            writequeue = nil
            handler = nil
            return false
        end
    end

    -- // COMPAT // --

    handler.getIp = handler.ip
    handler.getPort = handler.clientport

    --// handshake //--

    local wrote

    handler.handshake = coroutine_wrap( function( client )
            local err
            for i = 1, 10 do    -- 10 handshake attemps
                _, err = client:dohandshake( )
                if not err then
                    out_put( "server.lua: ssl handshake done" )
                    writelen = ( wrote and removesocket( writelist, socket, writelen ) ) or writelen
                    handler.receivedata = handler._receivedata    -- when handshake is done, replace the handshake function with regular functions
                    handler.dispatchdata = handler._dispatchdata
                    return dispatch( handler )
                else
                    out_put( "server.lua: error during ssl handshake: ", err )
                    if err == "wantwrite" then
                        if wrote == nil then
                            writelen = writelen + 1
                            writelist[ writelen ] = client
                            wrote = true
                        end
                    end
                    coroutine_yield( handler, nil, err )    -- handshake not finished
                end
            end
            _ = err ~= "closed" and close( socket )
            handler.close( )
            disconnect( handler, err )
            writequeue = nil
            handler = nil
            return false    -- handshake failed
        end
    )
    handler.receivedata = handler.handshake
    handler.dispatchdata = handler.handshake

    handler.handshake( socket )    -- do handshake

    socketlist[ socket ] = handler
    readlen = readlen + 1
    readlist[ readlen ] = socket

    return handler, socket
end

wraptcpclient = function( listener, socket, ip, serverport, clientport, mode )    -- this function wraps a socket

    local dispatch, disconnect = listener.listener, listener.disconnect

    --// private closures of the object //--

    local writequeue = { }    -- list for messages to send

    local eol

    local rstat, sstat = 0, 0

    --// local import of socket methods //--

    local send = socket.send
    local receive = socket.receive
    local close = socket.close
    local shutdown = socket.shutdown

    --// public methods of the object //--

    local handler = { }

    handler.getstats = function( )
        return rstat, sstat
    end

    handler.listener = function( data, err )
        return listener( handler, data, err )
    end
    handler.ssl = function( )
        return false
    end
    handler.send = function( _, data, i, j )
            return send( socket, data, i, j )
    end
    handler.receive = function( pattern, prefix )
            return receive( socket, pattern, prefix )
    end
    handler.shutdown = function( pattern )
        return shutdown( socket, pattern )
    end
    handler.close = function( closed )
        _ = not closed and shutdown( socket )
        _ = not closed and close( socket )
        writelen = ( eol and removesocket( writelist, socket, writelen ) ) or writelen
        readlen = removesocket( readlist, socket, readlen )
        socketlist[ socket ] = nil
        out_put "server.lua: closed handler and removed socket from list"
    end
    handler.ip = function( )
        return ip
    end
    handler.serverport = function( )
        return serverport
    end
    handler.clientport = function( ) 
        return clientport
    end
    handler.write = function( data )
        if not eol then
            writelen = writelen + 1
            writelist[ writelen ] = socket
            eol = 0
        end
        eol = eol + 1
        writequeue[ eol ] = data
    end
    handler.writequeue = function( )
        return writequeue
    end
    handler.socket = function( )
        return socket
    end
    handler.mode = function( )
        return mode
    end
    handler.receivedata = function( )
        local data, err, part = receive( socket, mode )    -- receive data in "mode"
        if not err or ( err == "timeout" or err == "wantread" ) then    -- received something
            local data = data or part or ""
            local count = #data * STAT_UNIT
            rstat = rstat + count
            receivestat = receivestat + count
            out_put( "server.lua: read data '", data, "', error: ", err )
            return dispatch( handler, data, err )
        else    -- connections was closed or fatal error
            out_put( "server.lua: client ", ip, ":", clientport, " error: ", err )
            handler.close( )
            disconnect( handler, err )
            writequeue = nil
            handler = nil
            return false
        end
    end
    handler.dispatchdata = function( )    -- this function writes data to handlers
        local buffer = table_concat( writequeue, "", 1, eol )
        local succ, err, byte = send( socket, buffer )
        local count = ( succ or 0 ) * STAT_UNIT
        sstat = sstat + count
        sendstat = sendstat + count
        out_put( "server.lua: sended '", buffer, "', bytes: ", succ, ", error: ", err, ", part: ", byte, ", to: ", ip, ":", clientport )
        if succ then    -- sending succesful
            --writequeue = { }
            eol = nil
            writelen = removesocket( writelist, socket, writelen )    -- delete socket from writelist
            return true
        elseif byte and ( err == "timeout" or err == "wantwrite" ) then    -- want write
            buffer = string_sub( buffer, byte + 1, -1 )    -- new buffer
            writequeue[ 1 ] = buffer    -- insert new buffer in queue
            eol = 1
            return true
        else    -- connection was closed during sending or fatal error
            out_put( "server.lua: client ", ip, ":", clientport, " error: ", err )
            handler.close( )
            disconnect( handler, err )
            writequeue = nil
            handler = nil
            return false
        end
    end

    -- // COMPAT // --

    handler.getIp = handler.ip
    handler.getPort = handler.clientport

    socketlist[ socket ] = handler
    readlen = readlen + 1
    readlist[ readlen ] = socket

    return handler, socket
end

addtimer = function( listener )
    timelistener[ #timelistener + 1 ] = listener
end

firetimer = function( listener )
    for i, listener in ipairs( timelistener ) do
        listener( )
    end
end

addserver = function( listeners, port, addr, mode, sslctx )    -- this function provides a way for other scripts to reg a server
    local err
    if type( listeners ) ~= "table" then
        err = "invalid listener table"
    else
        for name, func in pairs( listeners ) do
            if type( func ) ~= "function" then
                err = "invalid listener function"
                break
            end
        end
    end
    if not type( port ) == "number" or not ( port >= 0 and port <= 65535 ) then
        err = "invalid port"
    elseif listener[ port ] then
        err=  "listeners on port '" .. port .. "' already exist"
    elseif sslctx and not luasec then
        err = "luasec not found"
    end
    if err then
        out_error( "server.lua: ", err )
        return nil, err
    end
    addr = addr or "*"
    local server, err = socket_bind( addr, port )
    if err then
        out_error( "server.lua: ", err )
        return nil, err
    end
    local handler, err = wrapserver( listeners, server, addr, port, mode, sslctx )    -- wrap new server socket
    if not handler then
        server:close( )
        return nil, err
    end
    server:settimeout( 0 )
    readlen = readlen + 1
    readlist[ readlen ] = server
    listener[ port ] = listeners
    socketlist[ server ] = handler
    out_put( "server.lua: new server listener on ", addr, ":", port )
    return true
end

removesocket = function( tbl, socket, len )    -- this function removes sockets from a list
    for i, target in ipairs( tbl ) do
        if target == socket then
            len = len - 1
            table_remove( tbl, i )
            return len
        end
    end
    return len
end

closeall = function( )
    for sock, handler in pairs( socketlist ) do
        handler.shutdown( )
        handler.close( )
        socketlist[ sock ] = nil
    end
    writelist, readlist, socketlist = { }, { }, { }
end

closesocket = function( socket )
    writelen = removesocket( writelist, socket, writelen )
    readlen = removesocket( readlist, socket, readlen )
    socketlist[ socket ] = nil
    socket:close( )
end

loop = function( )    -- this is the main loop of the program
    --signal_set( "hub", "run" )
    repeat
        local read, write, err = socket_select( readlist, writelist, 1 )    -- 1 sec timeout, nice for timers
        for i, socket in ipairs( write ) do    -- send data waiting in writequeues
            local handler = socketlist[ socket ]
            if handler then
                handler.dispatchdata( )
            else
                closesocket( socket )
                out_put "server.lua: found no handler and closed socket (writelist)"    -- this should not happen
            end
        end
        for i, socket in ipairs( read ) do    -- receive data
            local handler = socketlist[ socket ]
            if handler then
                handler.receivedata( )
            else
                closesocket( socket )
                out_put "server.lua: found no handler and closed socket (readlist)"    -- this can happen
            end
        end
        firetimer( )
        --collectgarbage "collect"
    until false --signal_get "hub" ~= "run"
    return --signal_get "hub"
end

----------------------------------// BEGIN //--

----------------------------------// PUBLIC INTERFACE //--

return {

    add = addserver,
    loop = loop,
    stats = stats,
    closeall = closeall,
    addtimer = addtimer,

}