/*
* author: kenny huang
* email: [email protected]
* brief: 向lua注册一简单的网络接口，在lua中实现简单的网络功能
*        网络部分代码采用了windows网络编程中的一个例子
*/
#include <winsock2.h>
#include <windows.h>
#include <stdio.h>
#include "lua.h"
#include "lauxlib.h"
#include "lualib.h"
void BindFunction(lua_State *lState);
__declspec(dllexport) void RegisterFunction(lua_State *L)
{
    BindFunction(L);
}
#define DATA_BUFSIZE 8192
typedef struct _SOCKET_INFORMATION {
   CHAR Buffer[DATA_BUFSIZE];
   WSABUF DataBuf;
   SOCKET Socket;
   OVERLAPPED Overlapped;
   DWORD BytesSEND;
   DWORD BytesRECV;
   int index;
} SOCKET_INFORMATION, * LPSOCKET_INFORMATION;
DWORD TotalSockets = 0;
LPSOCKET_INFORMATION SocketArray[FD_SETSIZE];
FD_SET WriteSet;
FD_SET ReadSet;
BOOL CreateSocketInformation(SOCKET s)
{
   LPSOCKET_INFORMATION SI;
      
   printf("Accepted socket number %d/n", s);
   if ((SI = (LPSOCKET_INFORMATION) GlobalAlloc(GPTR,
      sizeof(SOCKET_INFORMATION))) == NULL)
   {
      printf("GlobalAlloc() failed with error %d/n", GetLastError());
      return FALSE;
   }
   // Prepare SocketInfo structure for use.
   SI->Socket = s;
   SI->BytesSEND = 0;
   SI->BytesRECV = 0;
   SocketArray[TotalSockets] = SI;
   SI->index = TotalSockets;    
   TotalSockets++;
   return TRUE;
}
void FreeSocketInformation(DWORD Index)
{
   LPSOCKET_INFORMATION SI = SocketArray[Index];
   DWORD i;
   closesocket(SI->Socket);
   printf("Closing socket number %d/n", SI->Socket);
   GlobalFree(SI);
   // Squash the socket array
   for (i = Index; i < TotalSockets; i++)
   {
      SocketArray[i] = SocketArray[i + 1];
      if(SocketArray[i])
        SocketArray[i]->index = i;
   }
   TotalSockets--;
}
BOOL SetNonBlock(SOCKET sock)
{
   unsigned long NonBlock = 1;
   if (ioctlsocket(sock, FIONBIO, &NonBlock) == SOCKET_ERROR)
   {
      printf("ioctlsocket() failed with error %d/n", WSAGetLastError());
      return FALSE;
   }
   return TRUE;
}
SOCKET TCP_Listen(/*const char *addr,*/int port)
{
   SOCKET ListenSocket;
   SOCKADDR_IN InternetAddr;
   if ((ListenSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0,
      WSA_FLAG_OVERLAPPED)) == INVALID_SOCKET) 
   {
      printf("WSASocket() failed with error %d/n", WSAGetLastError());
      return 0;
   }
   InternetAddr.sin_family = AF_INET;
   InternetAddr.sin_addr.s_addr = htonl(INADDR_ANY);
   InternetAddr.sin_port = htons(port);
   if (bind(ListenSocket, (PSOCKADDR) &InternetAddr, sizeof(InternetAddr))
      == SOCKET_ERROR)
   {
      printf("bind() failed with error %d/n", WSAGetLastError());
      return 0;
   }
   if (listen(ListenSocket, 5))
   {
      printf("listen() failed with error %d/n", WSAGetLastError());
      return 0;
   }
   return ListenSocket;
}
int Select()
{
    return select(0, &ReadSet, &WriteSet, NULL, NULL);
}
BOOL Accept(SOCKET sock)
{
    SOCKET newSock;
    if ((newSock = accept(sock, NULL, NULL)) != INVALID_SOCKET)
    {
        if(!SetNonBlock(newSock))
            return FALSE;
        if (CreateSocketInformation(newSock) == FALSE)
            return FALSE;
        return TRUE;
    }
    else
    {        
      if (WSAGetLastError() != WSAEWOULDBLOCK)
      {
         printf("accept() failed with error %d/n", WSAGetLastError());
      }
      return FALSE;
    }
}

//下面都是注册到lua中的C函数
int lua_PrepareSend(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-2);
    int len = lua_tonumber(L,-1);
    sockInfo->DataBuf.buf = sockInfo->Buffer + sockInfo->BytesSEND;
    sockInfo->DataBuf.len = len;
    return 0;
}

int lua_GetBuffer(lua_State *L)
{    
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-1);
    lua_pushstring(L,sockInfo->Buffer);
    return 1;
}
void lua_PrepareRecv(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-2);
    int len = lua_tonumber(L,-1);
    sockInfo->DataBuf.buf = sockInfo->Buffer;
    sockInfo->DataBuf.len = len;
    return 0;
}
int lua_FreeSocketInformation(lua_State *L)
{
    int index = lua_tonumber(L,-1);
    FreeSocketInformation(index);
    return 0;
}

int lua_SetNonBlock(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    lua_pushboolean(L,SetNonBlock(sock));
    return 1;
}
int lua_NetStart(lua_State *L)
{
   WSADATA wsaData;
   BOOL ret;
   if(WSAStartup(0x0202,&wsaData) != 0)
   {
      printf("WSAStartup() failed with error/n");
      WSACleanup();
      ret = FALSE;
   }
   ret = TRUE;
   lua_pushboolean(L,ret);
   return 1;
}
int lua_TCPListen(lua_State *L)
{
    int port = lua_tonumber(L,-1);
    SOCKET sock = TCP_Listen(port);
    lua_pushnumber(L,sock);
    return 1;
}
int lua_Add2Read(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    FD_SET(sock, &ReadSet);
    return 0;
    
}
int lua_Add2Write(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    FD_SET(sock, &WriteSet);
    return 0;
    
}
int lua_EmptyReadSet(lua_State *L)
{
    FD_ZERO(&ReadSet);
    return 0;
}
int lua_EmptyWriteSet(lua_State *L)
{
    FD_ZERO(&WriteSet);
    return 0;
}
int lua_GetSocketArray(lua_State *L)
{    
    int i = 0;
    int j = 1;
    for ( ; i < TotalSockets ; ++i)
    {
        char index[10];
        lua_getglobal(L,"add2SocketArray");
        lua_pushnumber(L,j);
        lua_pushlightuserdata(L,SocketArray[i]);
        lua_pcall(L,2,0,0);
        ++j;
    }
    return 0;
}
int lua_Select(lua_State *L)
{
    lua_pushnumber(L,Select());
    return 1;
}
int lua_Accept(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    BOOL ret = Accept(sock);
    lua_toboolean(L,ret);
    return 1;
}
int lua_CanRead(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    if(FD_ISSET(sock, &ReadSet))
        lua_pushboolean(L,TRUE);
    else
        lua_pushboolean(L,FALSE);
    return 1;
}
int lua_CanWrite(lua_State *L)
{
    SOCKET sock = lua_tonumber(L,-1);
    if(FD_ISSET(sock, &WriteSet))
        lua_pushboolean(L,TRUE);
    else
        lua_pushboolean(L,FALSE);
    return 1;
}
int lua_Recv(lua_State *L)
{
    DWORD RecvBytes = 0;
    DWORD Flags = 0;
    SOCKET_INFORMATION *SocketInfo = lua_touserdata(L,-1);
    if (WSARecv(SocketInfo->Socket, &(SocketInfo->DataBuf), 1, &RecvBytes,&Flags, NULL, NULL) == SOCKET_ERROR)
    {
        lua_pushnumber(L,-1);
    }
    else
        lua_pushnumber(L,(int)RecvBytes);
    return 1;
}

int lua_Send(lua_State *L)
{
    SOCKET_INFORMATION *SocketInfo = lua_touserdata(L,-1);
    DWORD SendBytes;
    if (WSASend(SocketInfo->Socket, &(SocketInfo->DataBuf), 1, &SendBytes, 0,NULL, NULL) == SOCKET_ERROR)
    {
        lua_pushnumber(L,-1);
    }
    else
        lua_pushnumber(L,SendBytes);
    return 1;
}
int lua_GetLastError(lua_State *L)
{
    int errorno = WSAGetLastError();
    lua_pushnumber(L,errorno);
    return 1;
}
int lua_GetByteSend(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-1);
    lua_pushnumber(L,sockInfo->BytesSEND);
    return 1;
}
int lua_SetByteSend(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-2);
    int bs = lua_tonumber(L,-1);
    sockInfo->BytesSEND = bs;
    return 0;
}
int lua_GetByteRecv(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-1);
    lua_pushnumber(L,sockInfo->BytesRECV);
    return 1;
}
int lua_SetByteRecv(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-2);
    int br = lua_tonumber(L,-1);
    sockInfo->BytesRECV = br;
    return 0;
}
int lua_GetSockIndex(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-1);
    lua_pushnumber(L,sockInfo->index);
    return 1;
}
int lua_GetSock(lua_State *L)
{
    SOCKET_INFORMATION *sockInfo = lua_touserdata(L,-1);
    lua_pushnumber(L,sockInfo->Socket);
    return 1;
}

void BindFunction(lua_State *lState)
{
    lua_register(lState,"GetLastError",&lua_GetLastError);
    lua_register(lState,"Send",&lua_Send);
    lua_register(lState,"Recv",&lua_Recv);
    lua_register(lState,"CanWrite",&lua_CanWrite);
    lua_register(lState,"CanRead",&lua_CanRead);
    lua_register(lState,"Accept",&lua_Accept);
    lua_register(lState,"Select",&lua_Select);
    lua_register(lState,"GetSocketArray",&lua_GetSocketArray);
    lua_register(lState,"EmptyWriteSet",&lua_EmptyWriteSet);
    lua_register(lState,"EmptyReadSet",&lua_EmptyReadSet);
    lua_register(lState,"Add2Write",&lua_Add2Write);
    lua_register(lState,"Add2Read",&lua_Add2Read);
    lua_register(lState,"TCPListen",&lua_TCPListen);
    lua_register(lState,"NetStart",&lua_NetStart);
    lua_register(lState,"PrepareSend",&lua_PrepareSend);
    lua_register(lState,"PrepareRecv",&lua_PrepareRecv);
    lua_register(lState,"GetBuffer",&lua_GetBuffer);
    lua_register(lState,"FreeSocketInformation",&lua_FreeSocketInformation);
    lua_register(lState,"SetNonBlock",&lua_SetNonBlock);
    lua_register(lState,"GetByteSend",&lua_GetByteSend);
    lua_register(lState,"SetByteSend",&lua_SetByteSend);
    lua_register(lState,"GetByteRecv",&lua_GetByteRecv);
    lua_register(lState,"SetByteRecv",&lua_SetByteRecv);
    lua_register(lState,"GetSockIndex",&lua_GetSockIndex);
    lua_register(lState,"GetSock",&lua_GetSock);
}

SocketArray = {} --用于存放所有的客户连接
--下面导入dll,并调用里面的注册函数把所有网络接口都注册进来
local f = assert(package.loadlib("G:/luaNet/luaNet/Debug/luaNet.dll","RegisterFunction"))
f()
print("load ok")
function add2SocketArray(key,val)
    table.insert(SocketArray,key,val)
end
function netLoop()
   local ListenSocket   
   if NetStart() == false then
       print(" net start error ")
   end
   
   ListenSocket = TCPListen(5010)
   if ListenSocket == nil then
       print("listen error")
   end
   
   if SetNonBlock(ListenSocket) == false then
       print("error on setnonblock")
   end
   
   while true do    
      EmptyReadSet()    
      EmptyWriteSet()
      Add2Read(ListenSocket)
      
      SocketArray = {}
      GetSocketArray()
      local Size = table.getn(SocketArray)
      if Size > 0 then
        for k,v in ipairs(SocketArray) do
            if GetByteRecv(v) > GetByteSend(v) then
                Add2Write(GetSock(v))
            else
                Add2Read(GetSock(v))
            end
        end
      end
  
      local Total = Select()
      if Total == -1 then
              print(" error on select")
             return
      end
      if CanRead(ListenSocket) == true then
             Total = Total - 1
             if Accept(ListenSocket) == false then
                 print("error on accept")
                 return
             end
      end
             
      local index = 1       
      while Total > 0 do
        local continue = false
        local SocketInfo = SocketArray[index]
          if CanRead(GetSock(SocketInfo)) == true then
              Total = Total - 1
              PrepareRecv(SocketInfo,4096)
                local ret = Recv(SocketInfo)
                if ret == -1 then
                    if lua_GetLastError() ~= 10035 then
                          FreeSocketInformation(GetSockIndex(SocketInfo))
                    end
                    continue = true
                else
                    SetByteRecv(SocketInfo,ret)
                    if ret == 0 then
                         FreeSocketInformation(GetSockIndex(SocketInfo))
                         continue = true
                    else
                        print(GetBuffer(SocketInfo))
                    end                    
                    
                end
          end
          
          if continue == false then
              if CanWrite(GetSock(SocketInfo)) == true then
                    Total = Total - 1
                  PrepareSend(SocketInfo,GetByteRecv(SocketInfo) - GetByteSend(SocketInfo))
                  local ret = Send(SocketInfo)
                  if ret == -1 then
                 if    lua_GetLastError() ~= 10035 then
                     FreeSocketInformation(GetSockIndex(SocketInfo))
                 end
                  else
                 local tmp = GetByteSend(SocketInfo) + ret
                 SetByteSend(SocketInfo,tmp)
                           if GetByteSend(SocketInfo) == GetByteRecv(SocketInfo) then
                              SetByteSend(SocketInfo,0)
                              SetByteRecv(SocketInfo,0)
                          end
             end
              end
          end
          index = index + 1
      end     
                   
 end
 
end
netLoop()