WebSocket для Lua

//

#include "ws.h"



int luaopen_ws( lua_State *L ) {
    luaL_newlib(L, __index);

    luaL_newlib(L, _server_index);
    lua_setfield(L, -2, "server");

    luaL_newlib(L, _client_index);
    lua_setfield(L, -2, "client");

    return 1;
}



static int lua_ws_xor( lua_State *L ) {
    size_t data_len, key_len;
    const char *data = luaL_checklstring(L, 1, &data_len);
    const char *key  = luaL_checklstring(L, 2, &key_len);
    char result[data_len];

    for ( size_t d=0, k=0; d<data_len; ) {
        result[d] = data[d] ^ key[k];
        k++; d++;
        if ( k >= key_len ) k = 0;
    }

    lua_pushlstring(L, result, data_len);
    return 1;
}



static int lua_ws_client_newmask( lua_State *L ) {
    char mask[WS_MASKLEN];
    struct timeval tv;

    for ( short i=0; i<WS_MASKLEN; i++ ) {
        gettimeofday(&tv, NULL);
        srand((unsigned int)tv.tv_usec);
        mask[i] = (char)rand();
    }

    lua_pushlstring(L, mask, WS_MASKLEN);
    return 1;
}



// пакет закрытия от клиента к серверу
static int lua_ws_client_close( lua_State *L ) {
    uchar str[] = { 0x88, 0x80, 0x11, 0x11, 0x11, 0x11 };
    lua_pushlstring(L, (char *)str, 6);
    return 1;
}



// пакет закрытия от сервера к клиенту
static int lua_ws_server_close( lua_State *L ) {
    uchar str[] = { 0x88, 0x00 };
    lua_pushlstring(L, (char *)str, 2);
    return 1;
}



// пакет от клиента к серверу
static int lua_ws_client_pack( lua_State *L ) {
    ulong len = luaL_checknumber(L, 1);
    ushort len2;
    uchar head[10];
    uchar head_len = 2;

    if ( len > 65535 ) { // long
        head_len += 8;
        head[1] = 127;
        len = (ulong)ntohll(len);
        memcpy(&head[2], &len, 8);
    } else if ( len > 125 ) { // medium
        head_len += 2;
        head[1] = 126;
        len2 = (ushort)ntohs((ushort)len);
        memcpy(&head[2], &len2, 2);
    } else { // short
        head[1] = (uchar)len;
    }

    head[0] = 0x81;
    head[1] |= 0x80;

    lua_pushlstring(L, (char *)head, head_len);
    return 1;
}



// пакет от сервера к клиенту
static int lua_ws_server_pack( lua_State *L ) {
    ulong len = luaL_checknumber(L, 1);
    ushort len2;
    uchar head[10];
    uchar head_len = 2;

    if ( len > 65535 ) { // long
        head_len += 8;
        head[1] = 127;
        len = (ulong)ntohll(len);
        memcpy(&head[2], &len, 8);
    } else if ( len > 125 ) { // medium
        head_len += 2;
        head[1] = 126;
        len2 = (ushort)ntohs((ushort)len);
        memcpy(&head[2], &len2, 2);
    } else { // short
        head[1] = (uchar)len;
    }

    head[0] = 0x81;

    lua_pushlstring(L, (char *)head, head_len);
    return 1;
}



/* errors:
    0 - no error
    1 - not enough data
    2 - connection closed
    3 - unknown opcode or not implemented
*/
static int lua_ws_server_unpack( lua_State *L ) {
    size_t str_len;
    const char *str = luaL_checklstring(L, 1, &str_len);
    uchar FIN, OPCODE, MASK, PLEN;
    ulong data_len; // real length of body
    ulong head_len; // real length of header
    char mask[WS_MASKLEN];

    head_len = 2; // min header len
    if ( str_len < head_len ) ws_err(1);

    FIN = ((str[0] & 0x80) > 0);
    OPCODE = (str[0] & 0x0F);
    MASK = ((str[1] & 0x80) > 0);
    PLEN = (str[1] & 0x7F);

    if ( OPCODE == 8 ) ws_err(2);
    if ( OPCODE != 1 ) ws_err(3);
    if ( FIN != 1 )    ws_err(3);
    if ( MASK != 1 )   ws_err(3);

    if ( PLEN == 127 ) {
        head_len += 8; // long (+8b)
        if ( str_len < head_len ) ws_err(1);
        data_len = (ulong)ntohll(*((ulong *)&str[2]));
    } else if ( PLEN == 126 ) {
        head_len += 2; // medium (+2b)
        if ( str_len < head_len ) ws_err(1);
        data_len = (ulong)ntohs(*((ushort *)&str[2]));
    } else { // short
        data_len = (ulong)PLEN;
    }

    head_len += WS_MASKLEN; // + mask
    if ( str_len < head_len ) ws_err(1);

    for ( size_t i=0; i<WS_MASKLEN; i++ ) {
        mask[i] = str[head_len-WS_MASKLEN+i];
    }

    if ( str_len < head_len+data_len ) ws_err(1);

    lua_pushnumber(L, 0); // success
    lua_pushlstring(L, str+head_len, data_len); // masked data
    lua_pushnumber(L, head_len+data_len); // packet full len
    lua_pushlstring(L, mask, WS_MASKLEN); // mask
    return 4;
}



static int lua_ws_client_unpack( lua_State *L ) {
    size_t str_len;
    const char *str = luaL_checklstring(L, 1, &str_len);
    uchar FIN, OPCODE, MASK, PLEN;
    ulong data_len; // real length of body
    ulong head_len; // real length of header

    head_len = 2; // min header len
    if ( str_len < head_len ) ws_err(1);

    FIN = ((str[0] & 0x80) > 0);
    OPCODE = (str[0] & 0x0F);
    MASK = ((str[1] & 0x80) > 0);
    PLEN = (str[1] & 0x7F);

    if ( OPCODE == 8 ) ws_err(2);
    if ( OPCODE != 1 ) ws_err(3);
    if ( FIN != 1 )    ws_err(3);
    if ( MASK != 0 )   ws_err(3);

    if ( PLEN == 127 ) {
        head_len += 8; // long (+8b)
        if ( str_len < head_len ) ws_err(1);
        data_len = (ulong)ntohll(*((ulong *)&str[2]));
    } else if ( PLEN == 126 ) {
        head_len += 2; // medium (+2b)
        if ( str_len < head_len ) ws_err(1);
        data_len = (ulong)ntohs(*((ushort *)&str[2]));
    } else { // short
        data_len = (ulong)PLEN;
    }

    if ( str_len < head_len+data_len ) ws_err(1);

    lua_pushnumber(L, 0); // success
    lua_pushlstring(L, str+head_len, data_len); // data
    lua_pushnumber(L, head_len+data_len); // packet full len
    return 3;
}