Writing WebSockets server with few lines of Lua on nginx/OpenResty

In case if you've been hunting for a super-fast way of doing WebSockets on your OpenResty or nginx with Lua module, there is a way to do it directly as a location handler. In most cases, with WebSockets, it is all just about getting data from users and sending something back, but, besides that, we should have it to easily inject PUB/SUBSCRIBE services, like NATS driven ones or custom made ones.

location /ws{
    client_max_body_size 32k;

    lua_socket_log_errors off;
    lua_check_client_abort on;

    content_by_lua_block {
        local server = require('websockets.server')

        -- server got data
        server.on.message = function(data, send)
            print(json.encode(data))
            send('{"type": "debug", "data": "---"}')
        end

        server.run()
 
    }
}

As you can see from the code above, there are just a few lines that can do the whole job for you. You add 'on.message' handler that receives data and send() function that you can use to send your data back to the user.

And the code for the /websockets/server.lua is the following:

---
--- Websockets Server for NGINX/OpenResty
--- Created by skitsanos.
---
local server = require "resty.websocket.server"

local response = {
    info = function()
        return json.encode({
            type = 'info',
            data = {
                version = '1.0.0',
                id = ngx.var.request_id
            }
        })
    end,

    message = function(data)
        return json.encode({
            type = 'message',
            data ?= data
        })
    end
}

local m = {
    version = '1.0.0',

    debug = true,

    json = true,

    send = nil,

    on = {
        message = nil
    }
}

function m.log(...)
    if (m.debug) then
        ngx.log(...)
    end
end

function m.run()
    local wb, err = server:new {
        max_payload_len = 32768,
        timeout = 5000
    }

    if not wb then
        m.log(ngx.ERR, "failed to new websocket: ", err)
        return ngx.exit(444)
    end

    wb:send_text(response.info())

    m.send = function(data)
        wb:send_text(data)
    end

    while true do
        -- try to receive
        local data, typ, errReceive = wb:recv_frame()

        -- check socked timeout
        if not data then
            if not string.find(errReceive, "timeout", 1, true) then
                m.log(ngx.ERR, "failed to receive a frame: ", errReceive)
                return ngx.exit(444)
            end
        end

        local token = ngx.var.arg_token;
        if (token == nil) then
            token = ngx.var.request_id
        end

        if (typ ~= nil) then
            m.log(ngx.INFO, 'received a frame of type ', typ, ' and payload ', data)
        end

        local switch = {
            ['close'] = function()
                -- for typ "close", err contains the status code
                local code = errReceive

                -- send a close frame back:
                local bytes, errSendClose = wb:send_close(1000)
                if (not bytes) then
                    m.log(ngx.ERR, 'failed to send the close frame: ', errSendClose)
                end

                m.log(ngx.INFO, 'closing with status code ', code, ' and message ', data)
                return 1, errSendClose
            end,

            ['ping'] = function()
                local bytes, SendPong = wb:send_pong()
                if (not bytes) then
                    ngx.log(ngx.ERR, "failed to send frame: ", SendPong)
                    return 1, SendPong
                end
            end,

            ['text'] = function()
                if (not m.json) then
                    if (m.on.message ~= nil) then
                        m.on.message(data, m.send)
                    end

                    return 0, nil
                end

                -- decode payload
                local ok, payload = pcall(json.decode, data)
                if (ok ~= true) then
                    wb:send_close(1003, 'Incorrect payload format. Must be valid JSON')
                    return 1, ok
                end

                if (m.on.message ~= nil) then
                    m.on.message(payload, m.send)
                end

                return 0, nil
            end
        }

        if (typ == 'close') then
            break
        end

        if (typ ~= nil) then
            -- check for return code
            local _, errRet = switch[typ]();
            if (errRet) then
                break
            end
        end
    end

    wb:send_close()
end

return m

And now the client-side, another bit of code, this time with JavaScript:

const intervalReconnect = 3000;

let socket;

const reconnectWebSocket = () =>
{
    socket = new WebSocket(`ws://localhost/wsrelay`);

    socket.onopen = () =>
    {
        console.log('connected');
    };

    socket.onclose = () =>
    {
        setTimeout(() =>
        {
            reconnectWebSocket();
        }, intervalReconnect);
    };

    socket.onerror = e =>
    {
        if (e.currentTarget.readyState === 3)
        {
            console.log('-- connection closed');
        }
    };

    socket.onmessage = e =>
    {
        const {type, data} = JSON.parse(e.data);
        switch (type)
        {
            case 'info':
                console.log(data);
                socket.send(JSON.stringify({type: 'subscribe', subject: 'app.*'}));
                break;

            default:
                console.log(data);
                break;
        }
    };
};


reconnectWebSocket();

This is just a basic example of reconnecting a WebSocket client. In case if the connection dropped, it will reconnect in 3 seconds. You can extend your 'onmessage' handler in any possible way.

Hope it helps anyone.