BSD socket in lua

lua_socket.c

#include <WinSock2.h>
#include <WS2tcpip.h>
#include <Windows.h>
#include <malloc.h>
#include <assert.h>
#include <lua.hpp>

#define SOCKET_HANDLE    "socket*."
#define check_socket(L)   (*(SOCKET*)luaL_checkudata(L, 1, SOCKET_HANDLE))


static const char* GetErrorMessage(DWORD dwErrorCode = WSAGetLastError())
{
    static __declspec(thread) char description[MAX_PATH];
    FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, dwErrorCode,
        0, description, MAX_PATH-1, NULL);
    return description;
}
static int socket_create(lua_State* L, SOCKET fd)
{
    void* udata = lua_newuserdata(L, sizeof(SOCKET));
    memcpy(udata, &fd, sizeof(fd));
    luaL_getmetatable(L, SOCKET_HANDLE);
    lua_setmetatable(L, -2);
    return 1;
}


static int socket_gc(lua_State* L)
{
    SOCKET* fd = (SOCKET*)luaL_checkudata(L, 1, SOCKET_HANDLE);
    closesocket(*fd);
    *fd = 0;
    return 0;
}

static int socket_tostring(lua_State* L)
{
    SOCKET fd = check_socket(L);
    lua_pushfstring(L, "socket* (%d)", fd);
    return 1;
}

static int socket_bind(lua_State* L)
{
    SOCKET fd = check_socket(L);
    const char* host = luaL_checkstring(L, 2);
    int port = luaL_checkint(L, 3);
    struct sockaddr_in addr = {};
    addr.sin_family = AF_INET;
    addr.sin_port = htons((u_short)port);
    if (inet_pton(AF_INET, host, &(addr.sin_addr.s_addr)) == SOCKET_ERROR)
    {
        luaL_error(L, "inet_pton() failed: %s", GetErrorMessage());
        return 0;
    }
    if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
    {
        luaL_error(L, "bind() failed: %s", GetErrorMessage());
    }
    return 0;
}

static int socket_setopt(lua_State* L)
{
    SOCKET fd = check_socket(L);
    const char* opt = luaL_checkstring(L, 2);
    if (strcmp(opt, "nonblock") == 0)
    {
        ULONG val = lua_toboolean(L, 3);
        if (ioctlsocket(fd, FIONBIO, &val) == SOCKET_ERROR)
        {
            luaL_error(L, "ioctlsocket() failed: %s", GetErrorMessage());
        }
    }
    return 0;
}

static int socket_connect(lua_State* L)
{
    SOCKET fd = check_socket(L);
    const char* host = luaL_checkstring(L, 2);
    int port = luaL_checkint(L, 3);
    struct sockaddr_in addr = {};
    addr.sin_family = AF_INET;
    addr.sin_port = htons((u_short)port);
    if (inet_pton(AF_INET, host, &(addr.sin_addr.s_addr)) == SOCKET_ERROR)
    {
        luaL_error(L, "inet_pton() failed: %s", GetErrorMessage());
        return 0;
    }
    if (connect(fd, (struct sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
    {
        luaL_error(L, "connect() failed: %s", GetErrorMessage());
        return 0;
    }
    return 0;
}

static int socket_listen(lua_State* L)
{
    SOCKET fd = check_socket(L);
    int max_conn = SOMAXCONN;
    if (lua_gettop(L) > 1)
    {
        max_conn = luaL_checkint(L, 2);
    }
    if (listen(fd, SOMAXCONN) == SOCKET_ERROR)
    {
        luaL_error(L, "listen() failed: %s", GetErrorMessage());
    }
    return 0;
}

static int socket_accept(lua_State* L)
{
    SOCKET acceptor = check_socket(L);
    SOCKET fd = accept(acceptor, NULL, NULL);
    if (fd == INVALID_SOCKET)
    {
        luaL_error(L, "accept() failed: %s", GetErrorMessage());
        return 0;
    }
    return socket_create(L, fd);
}

static int socket_recv(lua_State* L)
{
    SOCKET fd = check_socket(L);
    int bufsize = 1024;
    if (lua_gettop(L) > 1)
    {
        bufsize = luaL_checkint(L, 2);
    }
    char* buf = (char*)malloc(bufsize);
    if (buf == NULL)
    {
        luaL_error(L, "malloc() failed: %s", "no enough memory!");
        return 0;
    }
    int bytes = recv(fd, buf, bufsize, 0);
    if (bytes == SOCKET_ERROR)
    {
        lua_pushnil(L);
        lua_pushfstring(L, "recv() failed: %s", GetErrorMessage());
        return 2;
    }
    if (bytes == 0)
    {
        lua_pushnil(L);
        lua_pushfstring(L, "recv() failed: %s", "connection closed");
        return 2;
    }
    lua_pushlstring(L, buf, bytes);
    return 1;
}

static int socket_send(lua_State* L)
{
    SOCKET fd = check_socket(L);
    const char* data = luaL_checkstring(L, 2);
    size_t len = lua_rawlen(L, 2);
    int bytes = send(fd, data, len, 0);
    if (bytes == SOCKET_ERROR)
    {
        lua_pushnil(L);
        lua_pushfstring(L, "recv() failed: %s", GetErrorMessage());
        return 2;
    }
    if (bytes == 0)
    {
        lua_pushnil(L);
        lua_pushfstring(L, "recv() failed: %s", "connection closed");
        return 2;
    }
    return 0;
}

static int process_socket(lua_State* L)
{
    SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (fd == INVALID_SOCKET)
    {
        luaL_error(L, "socket() failed: %s", GetErrorMessage());
        return 0;
    }
    return socket_create(L, fd);
}


static int process_sleep(lua_State* L)
{
    int ms = luaL_checkinteger(L, 1);
    Sleep(ms);
    return 0;
}

static void make_meta(lua_State* L)
{
    static const luaL_Reg meta_lib[] =
    {
        { "__gc", socket_gc },
        { "__tostring", socket_tostring },
        { "bind", socket_bind },
        { "setopt", socket_setopt },
        { "connect", socket_connect },
        { "listen", socket_listen },
        { "accept", socket_accept },
        { "recv", socket_recv },
        { "send", socket_send },
        { NULL, NULL },
    };
    luaL_newmetatable(L, SOCKET_HANDLE);
    lua_pushvalue(L, -1);
    lua_setfield(L, -2, "__index");
    luaL_setfuncs(L, meta_lib, 0);
    lua_pop(L, 1);  /* pop new metatable */
}

extern "C" __declspec(dllexport)
int luaopen_process(lua_State* L)
{
    static const luaL_Reg lib[] =
    {
        { "sleep", process_sleep},
        { "socket", process_socket},
        { NULL, NULL },
    };
    luaL_newlib(L, lib);
    make_meta(L);
    return 1;
}

#ifdef _WIN32
struct WinsockInit
{
    WinsockInit()
    {
        WSADATA data;
        int r = WSAStartup(MAKEWORD(2, 2), &data);
        assert(r == 0);
    }
    ~WinsockInit()
    {
        WSACleanup();
    }
};

WinsockInit init;
#endif