LuaYard/LuaNetwork.cpp

392 lines
6.7 KiB
C++

#include "LuaNetwork.h"
#include <WinSock2.h>
#include <ws2tcpip.h>
#include <Windows.h>
#include <vector>
#include <iostream>
using namespace std;
#pragma comment(lib, "ws2_32.lib")
void push_message(lua_State* L, int errcode)
{
char msgbuf[1024] = { 0 };
FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL, errcode, MAKELANGID(LANG_ENGLISH, SUBLANG_DEFAULT), msgbuf, 1024, NULL);
lua_pushstring(L, msgbuf);
}
void push_wserror(lua_State* L)
{
int errcode = WSAGetLastError();
lua_pushinteger(L, errcode);
push_message(L, errcode);
}
int socket_create(lua_State* L)
{
int mode = lua_tointeger(L, 1);
int option = lua_tointeger(L, 2);
int socket_type;
switch (mode)
{
case 1:
socket_type = SOCK_STREAM;
break;
case 2:
socket_type = SOCK_DGRAM;
break;
default:
socket_type = SOCK_STREAM;
}
int sfd = socket(AF_INET, socket_type, 0);
if (sfd < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
if (option & 0x1)
{
unsigned long op = 1;
int ret = ioctlsocket(sfd, FIONBIO, &op);
if (ret < 0)
{
closesocket(sfd);
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
}
lua_pushboolean(L, true);
lua_pushinteger(L, sfd);
return 2;
}
int socket_recv(lua_State* L)
{
int sfd = lua_tointeger(L, 1);
int nsz = lua_tointeger(L, 2);
static char buffer[1024 * 1024]; // 1M
int ret = recv(sfd, buffer, max(min(nsz, 1024 * 1024), 0), 0);
if (ret < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
else
{
lua_pushboolean(L, true);
lua_pushinteger(L, ret);
if (ret > 0)
{
lua_pushlstring(L, buffer, ret);
}
else
{
lua_pushstring(L, "");
}
return 3;
}
}
int socket_send(lua_State* L)
{
int sfd = lua_tointeger(L, 1);
size_t len;
const char* data = lua_tolstring(L, 2, &len);
int ret = send(sfd, data, len, 0);
if (ret < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
else
{
lua_pushboolean(L, true);
lua_pushinteger(L, ret);
lua_pushboolean(L, ret == len);
return 3;
}
}
int socket_select(lua_State* L)
{
timeval tval;
tval.tv_sec = lua_tointeger(L, 4);
tval.tv_usec = lua_tointeger(L, 5);
vector<int> readvec;
lua_pushnil(L);
while (lua_next(L, 1))
{
int fd = lua_tointeger(L, -1);
readvec.push_back(fd);
lua_pop(L, 1);
}
lua_pop(L, 1);
vector<int> writevec;
lua_pushnil(L);
while (lua_next(L, 2))
{
int fd = lua_tointeger(L, -1);
writevec.push_back(fd);
lua_pop(L, 1);
}
lua_pop(L, 1);
vector<int> errorvec;
lua_pushnil(L);
while (lua_next(L, 3))
{
int fd = lua_tointeger(L, -1);
errorvec.push_back(fd);
lua_pop(L, 1);
}
lua_pop(L, 1);
FD_SET readset, writeset, errorset, * preadset, * pwriteset, * perrorset;
FD_ZERO(&readset);
FD_ZERO(&writeset);
FD_ZERO(&errorset);
if (!readvec.empty())
{
for (const int& i : readvec)
{
FD_SET(i, &readset);
}
preadset = &readset;
}
else
{
preadset = NULL;
}
if (!writevec.empty())
{
for (const int& i : writevec)
{
FD_SET(i, &writeset);
}
pwriteset = &writeset;
}
else
{
pwriteset = NULL;
}
if (!errorvec.empty())
{
for (const int& i : errorvec)
{
FD_SET(i, &errorset);
}
perrorset = &errorset;
}
else
{
perrorset = NULL;
}
// cout << preadset << " " << pwriteset << " " << perrorset << " " << tval.tv_sec << " " << tval.tv_usec << endl;
int ret = select(0, preadset, pwriteset, perrorset, &tval);
// cout << "select() return: " << ret << endl;
// Sleep(1000);
if (ret < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
else
{
lua_pushboolean(L, true);
lua_pushinteger(L, ret);
if (ret > 0)
{
lua_newtable(L);
if (preadset)
{
int cnt = 1;
for (const int& fd : readvec)
{
if (FD_ISSET(fd, preadset))
{
lua_pushinteger(L, fd);
lua_seti(L, -2, cnt++);
}
}
}
lua_newtable(L);
if (pwriteset)
{
int cnt = 1;
for (const int& fd : writevec)
{
if (FD_ISSET(fd, pwriteset))
{
lua_pushinteger(L, fd);
lua_seti(L, -2, cnt++);
}
}
}
lua_newtable(L);
if (perrorset)
{
int cnt = 1;
for (const int& fd : errorvec)
{
if (FD_ISSET(fd, perrorset))
{
lua_pushinteger(L, fd);
lua_seti(L, -2, cnt++);
}
}
}
return 5;
}
else
{
return 2;
}
}
}
int socket_close(lua_State* L)
{
int sfd = lua_tointeger(L, 1);
closesocket(sfd);
return 0;
}
int socket_connect(lua_State* L)
{
int sfd = lua_tointeger(L, 1);
const char* ip = lua_tostring(L, 2);
int port = lua_tointeger(L, 3);
sockaddr_in saddr;
memset(&saddr, 0, sizeof(saddr));
saddr.sin_family = AF_INET;
saddr.sin_addr.s_addr = inet_addr(ip);
saddr.sin_port = htons(port);
int ret = connect(sfd, (const sockaddr*)&saddr, sizeof(saddr));
if (ret < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
else
{
lua_pushboolean(L, true);
return 1;
}
}
int socket_connect_finished(lua_State* L)
{
int sfd = lua_tointeger(L, 1);
int result;
int result_len = sizeof(result);
int ret = getsockopt(sfd, SOL_SOCKET, SO_ERROR, (char*)&result, &result_len);
cout << "getsockopt " << ret << " " << result << endl;
if (ret < 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
else
{
lua_pushboolean(L, true);
return 1;
}
}
int network_dnsresolve(lua_State* L)
{
const char* ip = lua_tostring(L, 1);
struct addrinfo hints;
memset(&hints, 0, sizeof(hints));
hints.ai_family = AF_INET; // Specified to IPv4
hints.ai_socktype = SOCK_STREAM;
hints.ai_protocol = IPPROTO_TCP;
struct addrinfo* result = NULL;
if (getaddrinfo(ip, NULL, &hints, &result) != 0)
{
lua_pushboolean(L, false);
push_wserror(L);
return 3;
}
struct addrinfo* p = result;
vector<string> vec;
while (p)
{
if (p->ai_family == AF_INET)
{
struct sockaddr_in* psaddr = (sockaddr_in*)(p->ai_addr);
vec.push_back(inet_ntoa(psaddr->sin_addr));
}
p = p->ai_next;
}
lua_pushboolean(L, true);
lua_newtable(L);
for (size_t i = 0; i < vec.size(); i++)
{
lua_pushstring(L, vec[i].c_str());
lua_seti(L, -2, (lua_Integer)i + 1);
}
return 2;
}
// Act like require(...), returns a table.
int InitLuaNetwork(lua_State* L)
{
WORD wVersionRequested = MAKEWORD(2, 2);
WSADATA wsaData;
int err = WSAStartup(wVersionRequested, &wsaData);
if (err != 0) {
throw exception("WSAStartup failed with error: %d\n", err);
}
luaL_Reg funcs[] =
{
{"socket", socket_create},
{"send", socket_send},
{"recv", socket_recv},
{"connect", socket_connect},
{"connect_finished", socket_connect_finished},
{"select", socket_select},
{"close", socket_close},
{"resolve", network_dnsresolve},
{NULL, NULL}
};
luaL_newlib(L, funcs);
lua_pushstring(L, "LuaNetwork v0.1");
lua_setfield(L, -2, "version");
return 1;
}