sockmgr.cpp

远程控制系统,可以实现局域网内进行远程控制计算机,很方便,并且是学习用的非常好的资料.

// sockmgr.cpp: implementation of the CSockManager class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "common.h"
#include "ntshell.h"
#include "sockmgr.h"

bool CSockManager::Inited = false;
CSockManager::PTEMP_BUFFER CSockManager::ptb = NULL;
CRITICAL_SECTION CSockManager::SendCS;
CRITICAL_SECTION CSockManager::RecvCS;

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CSockManager::CSockManager()
{
	simCount = 0;
}

CSockManager::~CSockManager()
{
	while (simCount)
	{
		FreeBuffer(&ppsim[--simCount]->sib);
		delete ppsim[simCount];
	}
}

int CSockManager::GetIoStateCode(SOCKET sock)
{
	int index = GetSockIndex(sock);
	return index == SOCKMGR_SOCK_NOT_FOUND ? index : ppsim[index]->status;
}

int CSockManager::GetSockIndex(SOCKET sock)
{
	for (int i = 0; i < simCount; i++)
	{
		if (ppsim[i]->sock == sock)
			return i;
	}

	return SOCKMGR_SOCK_NOT_FOUND;
}

PVOID CSockManager::GetSim(SOCKET sock)
{
	int index = GetSockIndex(sock);

	if (index == SOCKMGR_SOCK_NOT_FOUND)
		return NULL;

	return ppsim[index];
}

int CSockManager::InsertSock(SOCKET sock)
{
	ppsim[simCount] = new SOCK_IO_MANAGER;
	ppsim[simCount]->packsize = 0;
	ppsim[simCount]->sock = sock;
	ppsim[simCount]->status = SOCKMGR_IO_READY;
	ppsim[simCount]->siblock = BUFFER_LOCKED_NONE;

	if (ppsim[simCount] != NULL)
	{
		if (InitBuffer(&ppsim[simCount]->sib) == false)
		{
			delete ppsim[simCount];
			return -1;
		}

		simCount++;
		return simCount - 1;
	}

	return SOCKMGR_SOCK_NOT_FOUND;
}

void CSockManager::DeleteSock(SOCKET sock)
{
	int index = GetSockIndex(sock);

	if (index == SOCKMGR_SOCK_NOT_FOUND)
		return;

	simCount--;
	FreeBuffer(&ppsim[index]->sib);
	delete ppsim[index];
	ppsim[index] = ppsim[simCount];
}

bool CSockManager::InitBuffer(PSOCK_IO_BUFFER psib, int BufferSize)
{
	psib->maxlength = 0x200000;
	psib->pointer = 0;
	psib->datalength = 0;

	if (BufferSize > psib->maxlength)
		return false;

	psib->length = (BufferSize + 0xfff) & 0xfffff000;
	psib->buffer = (char *)VirtualAlloc(NULL, psib->maxlength, MEM_RESERVE, PAGE_READWRITE);

	return psib->buffer == NULL ? false : true;
}

int CSockManager::ReadBuffer(SOCKET sock, char *buffer, int length)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	PSOCK_IO_BUFFER psib = &psim->sib;

	if (length > psib->datalength - psib->pointer)
		return SOCKMGR_SPACE_NOT_ENOUGH;

	memcpy(buffer, psib->buffer + psib->pointer, length);
	psib->pointer += length;

	return SOCKMGR_SUCCESS;
}

int CSockManager::WriteBuffer(SOCKET sock, char *buffer, int length)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	PSOCK_IO_BUFFER psib = &psim->sib;
	int used = psib->datalength;
	int need = 0;

	if (psib->buffer == NULL)
	{
		InitBuffer(psib);
	}

	if (psib->length - used < length)
	{
		if (psib->maxlength - used < length)
			return SOCKMGR_SPACE_NOT_ENOUGH;

		psib->length = (used + length + 0xfff) & 0xfffff000;
		psib->buffer = (char *)VirtualAlloc(psib->buffer, psib->length, MEM_COMMIT, PAGE_READWRITE);
	}

	memcpy(psib->buffer + used, buffer, length);
	psib->pointer += length;
	psib->datalength += length;

	return SOCKMGR_SUCCESS;
}

void CSockManager::FreeBuffer(PSOCK_IO_BUFFER psib)
{
	VirtualFree(psib->buffer, 0, MEM_RELEASE);

	psib->buffer = NULL;
	psib->length = 0;
	psib->pointer = 0;
	psib->datalength = 0;
}

int CSockManager::SendPack(SOCKET sock)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);
	int ret;

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	if (psim->siblock != BUFFER_LOCKED_NONE &&
		psim->siblock != BUFFER_LOCKED_SEND)
		return SOCKMGR_BUFFER_LOCKED;

	if (psim->status == SOCKMGR_IO_READY ||
		psim->status == SOCKMGR_IO_READ_FINISH ||
		psim->status == SOCKMGR_IO_WRITE_FINISH)
	{
		ret = send(psim->sock, (char *)&psim->sib.datalength, sizeof(psim->sib.datalength), 0);

		if (ret == SOCKET_ERROR)
		{
			lastErrorCode = WSAGetLastError();
			return SOCKMGR_SOCK_FAILED;
		}

		psim->status = SOCKMGR_IO_WRITING_PACKDATA;
		psim->sib.pointer = 0;
	}

	if (psim->status != SOCKMGR_IO_WRITING_PACKDATA)
	{
		return SOCKMGR_INVALID_PROCESS;
	}

	psim->siblock = BUFFER_LOCKED_SEND;

	ret = send(psim->sock, psim->sib.buffer + psim->sib.pointer, psim->sib.datalength - psim->sib.pointer, 0);

	if (ret == SOCKET_ERROR)
	{
		if (WSAGetLastError() == WSAEWOULDBLOCK)
			return SOCKMGR_PENDING;
		else
		{
			lastErrorCode = WSAGetLastError();
			psim->siblock = BUFFER_LOCKED_NONE;
			return SOCKMGR_SOCK_FAILED;
		}
	}

	psim->sib.pointer += ret;

	if (psim->sib.pointer == psim->sib.datalength)
	{
		psim->status = SOCKMGR_IO_WRITE_FINISH;
		psim->siblock = BUFFER_LOCKED_NONE;
		psim->sib.pointer = 0;
		psim->packsize = psim->sib.datalength;
		return SOCKMGR_SUCCESS;
	}

	return SOCKMGR_PENDING;
}

int CSockManager::RecvPack(SOCKET sock)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);
	int ret, readed, packsize;

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	if (psim->siblock != BUFFER_LOCKED_NONE &&
		psim->siblock != BUFFER_LOCKED_RECV)
		return SOCKMGR_BUFFER_LOCKED;

	if (psim->status == SOCKMGR_IO_READY ||
		psim->status == SOCKMGR_IO_READ_FINISH ||
		psim->status == SOCKMGR_IO_WRITE_FINISH ||
		psim->status == SOCKMGR_IO_READING_PACKSIZE1 ||
		psim->status == SOCKMGR_IO_READING_PACKSIZE2 ||
		psim->status == SOCKMGR_IO_READING_PACKSIZE3)
	{
		switch (psim->status)
		{
		case SOCKMGR_IO_READY:
		case SOCKMGR_IO_READ_FINISH:
		case SOCKMGR_IO_WRITE_FINISH:
			readed = 0;
			break;
		case SOCKMGR_IO_READING_PACKSIZE1:
			readed = 1;
			break;
		case SOCKMGR_IO_READING_PACKSIZE2:
			readed = 2;
			break;
		case SOCKMGR_IO_READING_PACKSIZE3:
			readed = 3;
			break;
		}

		packsize = psim->packsize;
		ret = recv(psim->sock, (char *)&packsize + readed, sizeof(packsize) - readed, 0);

		if (ret == SOCKET_ERROR)
		{
			if (WSAGetLastError() == WSAEWOULDBLOCK)
				return SOCKMGR_PENDING;
			else
			{
				lastErrorCode = WSAGetLastError();
				return SOCKMGR_SOCK_FAILED;
			}
		}

		if (ret == 0)
			return SOCKMGR_SOCK_CLOSED;

		psim->packsize = packsize;

		switch (ret + readed)
		{
		case 1:
			psim->status = SOCKMGR_IO_READING_PACKSIZE1;
			return SOCKMGR_PENDING;
		case 2:
			psim->status = SOCKMGR_IO_READING_PACKSIZE2;
			return SOCKMGR_PENDING;
		case 3:
			psim->status = SOCKMGR_IO_READING_PACKSIZE3;
			return SOCKMGR_PENDING;
		case 4:
			psim->status = SOCKMGR_IO_READING_PACKDATA;
			psim->packsize >>= 2;
			psim->sib.pointer = 0;
			break;
		}
	}

	if (psim->status != SOCKMGR_IO_READING_PACKDATA)
		return SOCKMGR_INVALID_PROCESS;

	if (psim->packsize > psim->sib.length)
	{
		if (psim->packsize > psim->sib.maxlength)
			return SOCKMGR_SPACE_NOT_ENOUGH;

		psim->sib.length = (psim->packsize + 0xfff) & 0xfffff000;
		psim->sib.buffer = (char *)VirtualAlloc(psim->sib.buffer, psim->sib.length, MEM_COMMIT, PAGE_READWRITE);
	}

	psim->siblock = BUFFER_LOCKED_RECV;

	ret = recv(psim->sock, psim->sib.buffer + psim->sib.pointer, psim->packsize - psim->sib.pointer, 0);

	if (ret == SOCKET_ERROR)
	{
		if (WSAGetLastError() == WSAEWOULDBLOCK)
			return SOCKMGR_PENDING;
		else
		{
			lastErrorCode = WSAGetLastError();
			psim->siblock = BUFFER_LOCKED_NONE;
			return SOCKMGR_SOCK_FAILED;
		}
	}

	if (ret == 0)
	{
		psim->siblock = BUFFER_LOCKED_NONE;
		return SOCKMGR_SOCK_CLOSED;
	}

	psim->sib.pointer += ret;

	if (psim->sib.pointer >= psim->packsize)
	{
		psim->status = SOCKMGR_IO_READ_FINISH;
		psim->siblock = BUFFER_LOCKED_NONE;
		psim->sib.pointer = 0;
		psim->sib.datalength = psim->packsize;
		return SOCKMGR_SUCCESS;
	}

	return SOCKMGR_PENDING;
}

int CSockManager::OpenListenSock(char *bind_ip, int bind_port)
{
	SOCKET listen_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

	if (listen_sock == INVALID_SOCKET)
	{
		return -1;
	}

	struct sockaddr_in local;

	local.sin_family		= AF_INET;
	local.sin_port			= htons(bind_port);
	local.sin_addr.s_addr	= inet_addr(bind_ip);

	if (bind(listen_sock, (SOCKADDR *)&local, sizeof(local)) == SOCKET_ERROR)
	{
		closesocket(listen_sock);
		return -1;
	}

	if (listen(listen_sock, MAX_CONNECT) == SOCKET_ERROR) //最大连接数量
	{
		closesocket(listen_sock);
		return -1;
	}

	return 0;
}

int CSockManager::GetBufferDataSize(SOCKET sock)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim == NULL)
		return 0;

	return psim->packsize;
}

int CSockManager::GetBufferHandle(SOCKET sock, void **buffer, int *length)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	if (psim->siblock != BUFFER_LOCKED_NONE &&
		psim->siblock != BUFFER_LOCKED_USER)
		return SOCKMGR_BUFFER_LOCKED;

	psim->siblock = BUFFER_LOCKED_USER;

	if (length != NULL)
		*length = psim->sib.length;

	*buffer = psim->sib.buffer;

	return SOCKMGR_SUCCESS;
}

void CSockManager::ReleaseBufferHandle(SOCKET sock)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim != NULL && psim->siblock == BUFFER_LOCKED_USER)
		psim->siblock = BUFFER_LOCKED_NONE;
}

int CSockManager::MovePointer(SOCKET sock, int move, int method)
{
	PSOCK_IO_MANAGER psim = (PSOCK_IO_MANAGER)GetSim(sock);

	if (psim == NULL)
		return SOCKMGR_SOCK_NOT_FOUND;

	if (psim->siblock != BUFFER_LOCKED_NONE &&
		psim->siblock != BUFFER_LOCKED_USER)