bandage.lua tcp-async client-server lib

local socket     = require('socket')
local tcp_master = socket.tcp4 or socket.tcp

--[[
    структура сообщения:
    [служебный символ][длина сообщения]\r\n
    [тело сообщения]

    (почти аналогично redis)

    Служебные символы (типы сообщений):
    @ - коннект
    # - дисконнект
    $ - обычное сообщение
    % - пинг
]]


local client = {}
client.__index = client

-- создание самостоятельного клиента
function client:new(host, port)
    local sock = tcp_master()
    assert(type(host) == 'string', 'arg#1: string host expected, got ' .. type(host))
    assert(type(port) == 'number', 'arg#2: number port expected, got ' .. type(host))

    local status, error = sock:connect(host, port)

    if not status then
        return nil, error
    end

    local self = self:_new(sock)

    self:onAccept()

    return self
end

-- заполнение полей отдельно
function client:_new(sock, server)
    local host, port = sock:getpeername()
    local self = setmetatable({}, self)

    self:__debug('is created' .. (server and ' by ' .. server or ''))

    self.sock = sock
    self.host = host
    self.port = port

    self.server  = server -- связь с родительским классом, если создаёт сервер

    self.time    = socket.gettime()
    self.dt      = self.time
    self.timeout = 60 -- ping
    self.timer   = 0
    self.delay   = 2

    self.msgtype = nil -- для состояния 'приём сообщения',
    self.length  = nil -- тип и длина принимаемого
    self.message = ""  -- контейнер для частей сообщения

    self.dbg = nil -- отладочный вывод

    -- мы не будем ждать пока он нам что-то присылает
    -- потому что в это время нам кто-то другой может что-то слать
    self.sock:settimeout(0)
    return self
end

function client:ping()
    local time = socket.gettime()
    local dt   = time - self.dt

    self.dt    = time
    self.timer = self.timer - dt

    if self.timer < 0 then
        self.timer = self.delay
        self:send('ping', '%')
    end

    if time - self.time > self.timeout then
        self:disconnect('disconnect by timeout')
    end
end

-- отправка сообщения соединения (передача инфы всякой, например)
function client:connect(message)
    message = message or "HELLO"
    self:send(message, '@')
end

function client:disconnect(message)
    message = message or 'disconnect'

    if self.server then
        local server = self.server
        self.server = nil
        return server:removeClient(self, message)
    end

    self:send(message, '#')
    self:onClosed()

    self.sock:close()
    self.sock = nil
end

function client:onAccept(msg)
    self:__debug('call onAccept')
end

function client:onClosed(msg)
    self:__debug('call onClosed')
end

function client:onConnect(msg)
    self:__debug('call onConnect', msg:sub(1, 10))
end

function client:onMessage(msg)
    self:__debug('call onMessage', msg:sub(1, 10))
end

function client:onDisconnect(msg)
    self:__debug('call onDisconnect', msg:sub(1, 10))
end

function client:onPing(message)
    if message == 'ping' then
        self:send('pong', '%')
    end
    if message == 'pong' then
        self.time = socket.gettime()
    end
end

function client:send(message, msgtype)
    if not self.sock then
        return nil, 'Connection close'
    end
    message = message or ''
    msgtype = msgtype or '$' -- обычное сообщение по умолчанию

    -- [тип][длина сообщения]\r\n[тело сообщения]
    self.sock:send(msgtype .. #message .. '\r\n' .. message)
end

-- Приём проводится в два этапа:
-- приём заголовка (данных о сообщении), и приём самого тела сообщения.
function client:receive()
    if not self.sock then
        return nil, 'Connection closed'
    end

    -- Принимаем заголовок, если есть
    local _, status = self:receiveHeader()

        -- Принимаем тело, если есть
    local message, msgtype = self:receiveBody()

    if msgtype == '$' then
        msgtype = 'onMessage'
    end

    if msgtype == '@' then
        msgtype = 'onConnect'
    end

    if msgtype == '#' then
        msgtype = 'onDisconnect'
    end

    if msgtype == '%' then
        self:onPing(message)
        return
    end

    if self[msgtype] then
        self[msgtype](self, message)
    end

    -- если нас создал сервер и у сервера присутствует такой колбек
    if self.server and self.server[msgtype] then
        self.server[msgtype](self.server, message, self)
    end

    if self:ping() then
        msgtype = 'closed'
    end

    if msgtype == 'closed' then
        self:disconnect('timeout disconnect')
        return nil, msgtype
    end

    if status == 'closed' then
        return self:disconnect()
    end

    return message, msgtype
end

function client:receiveHeader()

    -- у нас уже может быть принятый заголовок
    if self.msglen then return end

    -- 1. Начало приёма сообщения
    -- принимаем 'строку' с \r\n
    local message, status = self.sock:receive('*l')

    -- все неправильные вещи - игнорируем
    if not message then return message, status end

    -- вырезаем первый символ и длину
    local msgtype, msglen = message:sub(1, 1), tonumber(message:sub(2))

    if not msglen then return nil, status end

    self.msglen  = msglen
    self.msgtype = msgtype
    self.message = ""
    return nil, status
end

function client:receiveBody()
    -- у нас может быть непринятый заголовок
    if not self.msglen then return end

    -- 2. Когда у нас уже есть длина сообщения:
    -- Мы должны асинхронно принять тело сообщения,
    -- в случае если на текущий момент идёт приём
    -- (мы приняли первую часть с типом и длиной)
    -- то теперь мы просто ждём пока сообщение не дойдёт окончательно.
    -- На малых данных, оба этих этапа выполняются подряд, в одно действие.

    local i = 1 -- чтобы сильно не забивали канал огромными сообщениями

    -- "стандатный" размер буфера приёма сообщений - 8192
    local bufsize = self:append("", 8192)


    local message, status, partial = self.sock:receive(bufsize)

    while message or partial and bufsize > 0 do
        i = i + 1
        message = message or partial
        bufsize = self:append(message or partial, bufsize)

        -- Прерываем на больших данных или если уже всё приняли
        if bufsize == 0 or i > 5 then break end

        message, status, partial = self.sock:receive(bufsize)
    end

    if bufsize == 0 then
        local message = self.message
        local msgtype = self.msgtype

        self.message = ""
        self.msgtype = nil
        self.msglen  = nil
        return message, msgtype
    end
end

-- Функция, которая цепляет кусок длинного сообщения,
-- и возвращает размер следующего блока для приёма

function client:append(chunk, bufsize)
    self.message = self.message .. chunk
    self.msglen = self.msglen - #chunk
    if self.msglen < bufsize then
        return self.msglen
    end
    return bufsize
end

function client:__tostring()
    if self.sock then
        local shost, sport = self.sock:getsockname()
        local phost, pport = self.sock:getpeername()
        return 'Client object [' .. shost .. ':' .. sport .. ']>[' .. phost .. ':' .. pport .. ']'
    end
    return 'Client object [unconnected]'
end

function client:__concat(other)
    return tostring(self) .. tostring(other)
end

function client:__debug(...)
    if self.dbg or self.server and self.server.dbg then
        print('debug ' .. self .. ' ' .. tostring(...), select(2, ...))
    end
end

setmetatable(client, {__call = client.new})


local server = {}
server.__index = server


function server:new(port)
    self = setmetatable({}, self)
    self.port = port or 7777
    self.sock = tcp_master()
    self.clients = {}

    local succ, res = self.sock:bind('*', self.port)
    if not succ then
        error('server:new(' .. port .. ') error: ' .. tostring(res), 2)
    end
    self.sock:listen()
    self.sock:settimeout(0)
    self.sock:setoption('keepalive',   true)
    self.sock:setoption('tcp-nodelay', true)

    self.dbg = nil

    return self
end

function server:onAccept(client)
    self:__debug('call onAccept ' .. client)
end

function server:onPing(message, client)
    client:onPing(message)
end

function server:onClosed(client)
    self:__debug('call onClosed ' .. client)
end

function server:onConnect(msg, client)
    self:__debug('call onConnect ' .. client, msg:sub(1, 10))
end

function server:onMessage(msg, client)
    self:__debug('call onMessage ' .. client, msg:sub(1, 10))
end

function server:onDisconnect(msg, client)
    self:__debug('call onDisconnect ' .. client, msg:sub(1, 10))
end

function server:newClient(sock)
    -- указываем заодно ссылку на себя
    local client  = client:_new(sock, self)

    self.clients[client.sock] = client
    return client
end

function server:removeClient(client, message)
    local sock = client.sock

    self:__debug('remove ' .. client .. ' start')
    local client = self.clients[sock]
    if not client then return end

    self:onClosed(client)
    client:disconnect(message)

    self.clients[sock] = nil
end

function server:getClientList()
    local l = {}
    for k, v in pairs(self.clients) do
        table.insert(l, v)
    end
    return l
end

function server:accept()
    local sock = self.sock:accept()
    while sock do
        socket.sleep(2)
        local client = self:newClient(sock)
        self:onAccept(client)
        sock = self.sock:accept()
    end
end

function server:update()
    self:accept()
    for sock, client in pairs(self.clients) do
        local message, status = client:receive()
        while message do
            -- Повторяем, пока есть что принять от данного клиента
            message, status = client:receive()
        end
    end
end

function server:send(client, message, type)
    client:send(message, type)
end

function server:shutdown(message)
    for sock, client in pairs(self.clients) do
        client:disconnect('server shutdown')
    end
    self.sock:close()
end

function server:__tostring()
    local shost, sport = self.sock:getsockname()

    return 'Server object [' .. shost .. ':' .. sport .. ']'
end

function server:__concat(other)
    return tostring(self) .. tostring(other)
end

function server:__debug(...)
    if self.dbg then
        print('debug ' .. self .. ' ' .. tostring(...), select(2, ...))
    end
end

setmetatable(server, {__call = server.new})

if ... then -- если цепляется как библиотека - возвращаем библиотеку
    return {
        server = server,
        client = client
    }
end
--[[ Схема работы:
    И клиенты и серверы могут иметь набор колбеков:
    onAccept     - момент стыковки,
    onConnect    - когда другая сторона прислала connect-сообщение [obj]:send(message, "@")
    onMessage    - другая сторона прислала message-сообщение [obj]:send(message)
    onDisconnect - другая сторона прислала disconnect-сообщение [obj]:send(message, "#")
    onClosed     - другая сторона закрыла соединение


    В случае сервера, если повесить колбеки на подсоединившихся клиентов,
    будут вызываться сразу оба колбека: соответствующий клиентский и соответствующий серверный.

    Уточнение: если одна из сторон вызвала client:disconnect(), то другая сторона вызовет только onClosed.
]]

-- Примеры

local srv = server:new(7777)

function srv:onAccept(client)
    print('Client connected! ' .. client)
    client:send("Hai, what's up?")
end

function srv:onConnect(message, client)
    print('Client sends us connect message! ' .. client, message)
end

function srv:onMessage(message, client)
    print('Client sends us message! ' .. client, message)
    -- второй аргумент - заставляет другую сторону вызвать колбек:
    -- @ -- onConnect, # - onDisconnect, $ - onMessage
    -- по умолчанию - onMessage
    -- Сервер может вызвать client:disconnect(),
    -- но тогда другая сторона сразу обрубится, с вызовом onClosed.
    client:send("Nope, get out here!", "#")
end

function srv:onDisconnect(message, client)
    print('Client disconnect from us! ' .. client, message)
end

function srv:onClosed(client)
    print(self .. ' client closes connection ' .. client)
end


local cli = client:new('localhost', 7777)

function cli:onMessage(message)
    print("Server sends us message: " .. message)
    self:send("Do you want to keep me?")
end

function cli:onDisconnect(message)
    print("Server disconnect us with message: " .. message)
    self:disconnect('(((')
end

function cli:onClosed()
    print(self .. ' server closes connection')
end


cli:connect("Hello! i'm novice!")

while true do
    srv:update()
    local a, b = cli:receive()
    socket.sleep(.1)
end