sockmgr1.cpp

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

// sockmgr1.cpp: implementation of the CSockMgr class.
//
//////////////////////////////////////////////////////////////////////

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

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

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

CSockMgr::CSockMgr()
{
	flags = 0;
	sock = INVALID_SOCKET;
	sib.buffer = NULL;
	sob.buffer = NULL;
	packlink = NULL;
}

CSockMgr::~CSockMgr()
{
	SetSocket(INVALID_SOCKET);
}

SOCKET CSockMgr::SetSocket(SOCKET sock1)
{
	SOCKET s = sock;
	sock = sock1;

	while (packlink != NULL)
	{
		PPACK_LINK p = packlink;
		packlink = packlink->next;
		delete [] p;
	}

	FreeBuffer(&sib);
	FreeBuffer(&sob);

	return s;
}

bool CSockMgr::AllocBuffer(PSOCK_IO_BUFFER psb, int BufferSize)
{
	if (psb->buffer != NULL)
	{
		if (BufferSize > psb->maxlength)
			return false;

		psb->length = (BufferSize + 0xfff) & 0xfffff000;
		psb->buffer = (char *)VirtualAlloc(psb->buffer, psb->length, MEM_COMMIT, PAGE_READWRITE);
	}
	else
	{
		psb->maxlength = 0x200000;
		psb->pointer = 0;
		psb->datalength = 0;

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

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

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

void CSockMgr::FreeBuffer(PSOCK_IO_BUFFER psb)
{
	if (psb->buffer != NULL)
	{
		VirtualFree(psb->buffer, 0, MEM_RELEASE);
		psb->buffer = NULL;
	}

	psb->length = 0;
	psb->pointer = 0;
	psb->datalength = 0;
}

int CSockMgr::ReadBuffer(char *buffer, int length)
{
	if (flags & SOCKMGR_RECVING)
		return SOCKMGR_BUFFER_LOCKED;

	if (length > (int)(sib.datalength - (sib.pointer - sizeof(PACK_TYPE_1))))
		return SOCKMGR_SPACE_NOT_ENOUGH;

	memcpy(buffer, sib.buffer + sib.pointer, length);
	sib.pointer += length;

	if ((int)(sib.pointer - sizeof(PACK_TYPE_1)) >= sib.datalength)
		FreeBuffer(&sib);

	return SOCKMGR_SUCCESS;
}

int CSockMgr::WriteBuffer(const char *buffer, int length)
{
	if (flags & SOCKMGR_SENDING)
		return SOCKMGR_BUFFER_LOCKED;

	if (sob.buffer == NULL)
	{
		AllocBuffer(&sob);
	}

	if (sob.length - sob.pointer < length)
	{
		if (sob.maxlength - sob.pointer < length)
			return SOCKMGR_SPACE_NOT_ENOUGH;

		AllocBuffer(&sob, sob.pointer + length);
	}

	memcpy(sob.buffer + sob.pointer, buffer, length);
	sob.pointer += length;
	sob.datalength = sob.pointer;

	return SOCKMGR_SUCCESS;
}

int CSockMgr::SendPack()
{
	if (packlink == NULL)
		return 0;

	PPACK_LINK p = packlink;

	int ret = send(sock, ((char *)p->pack) + p->pointer, p->length - p->pointer, 0);

	if (ret == SOCKET_ERROR)
	{
		if (WSAGetLastError() == WSAEWOULDBLOCK)
			return SOCKMGR_PENDING;
		else
		{
			lastErrorCode = WSAGetLastError();
			packlink = packlink->next;
			delete [] p;
			return SOCKMGR_SOCK_FAILED;
		}
	}

	p->pointer += ret;

	if (p->pointer >= p->length)
	{
		packlink = packlink->next;
		delete [] p;
		return SOCKMGR_SUCCESS;
	}

	return SOCKMGR_PENDING;
}

int CSockMgr::RecvPack()
{
	if (!(flags & SOCKMGR_RECVING))
	{
		if (sib.buffer == NULL)
			AllocBuffer(&sib);
		AllocBuffer(&sib, sizeof(PACK_TYPE_1));
		sib.pointer = 0;
		flags |= SOCKMGR_RECVING;
	}

	if (sib.pointer < sizeof(PACK_TYPE_1))
	{
		int ret = recv(sock, (char *)sib.buffer + sib.pointer, sizeof(PACK_TYPE_1) - sib.pointer, 0);

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

		sib.pointer += ret;
	}

	if (sib.pointer < sizeof(PACK_TYPE_1))
		return SOCKMGR_PENDING;

	PPACK_TYPE_1 ppt1 = (PPACK_TYPE_1)sib.buffer;

	if ((int)ppt1->nPackSize > sib.length)
	{
		if ((int)ppt1->nPackSize > sib.maxlength)
		{
			FreeBuffer(&sib);
			return SOCKMGR_SPACE_NOT_ENOUGH;
		}

		AllocBuffer(&sib, ppt1->nPackSize);
	}

	int ret = recv(sock, sib.buffer + sib.pointer, ppt1->nPackSize - (sib.pointer - sizeof(PACK_TYPE_1)), 0);

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

	sib.pointer += ret;

	if (sib.pointer - sizeof(PACK_TYPE_1) >= ppt1->nPackSize)
	{
		sib.pointer = sizeof(PACK_TYPE_1);
		sib.datalength = ppt1->nPackSize;
		flags &= ~SOCKMGR_RECVING;
		return SOCKMGR_SUCCESS;
	}

	return SOCKMGR_PENDING;
}

int CSockMgr::GetBufferDataSize()
{
	return sib.datalength;
}

void CSockMgr::FlushBuffer()
{
	PPACK_LINK *pp = &packlink;

	if (sob.buffer == NULL)
		return;

	while (*pp != NULL)
		pp = &(*pp)->next;

	*pp = (PPACK_LINK)new char[sizeof(PACK_LINK) + sizeof(PACK_TYPE_1) + sob.datalength];

	if (*pp != NULL)
	{
		(*pp)->pack = (PPACK_TYPE_1)((char *)(*pp) + sizeof(PACK_LINK));
		(*pp)->length = sizeof(PACK_TYPE_1) + sob.datalength;
		(*pp)->pointer = 0;
		(*pp)->next = NULL;

		(*pp)->pack->dwPackType = 0;
		(*pp)->pack->nPackSize = sob.datalength;
		memcpy((*pp)->pack->bPackData, sob.buffer, sob.datalength);
	}

	FreeBuffer(&sob);
}

char *CSockMgr::GetErrorString()
{
	LPVOID lpMsgBuf;

	FormatMessage(	FORMAT_MESSAGE_ALLOCATE_BUFFER |
					FORMAT_MESSAGE_FROM_SYSTEM |
					FORMAT_MESSAGE_IGNORE_INSERTS,
					NULL,
					lastErrorCode,
					0,
					(LPTSTR)&lpMsgBuf,
					0,
					NULL
					);

	return (char *)lpMsgBuf;
}

int CSockMgr::SendPack(SOCKET sock, const char *buffer, int length, int flags)
{
	if (Inited == false)
	{
		InitializeCriticalSection(&SendCS);
		InitializeCriticalSection(&RecvCS);
		Inited = true;
	}

	PTEMP_BUFFER p = ptb;
	int result = SOCKMGR_PENDING;

	EnterCriticalSection(&SendCS);

	while (p != NULL)
	{
		if (p->sock == sock && p->method == 1)
			break;
		p = p->next;
	}

	if (p != NULL && buffer != NULL)
	{
		LeaveCriticalSection(&SendCS);
		return SOCKMGR_SOCK_LOCKED;
	}

	if (p == NULL)
	{
		if (buffer == NULL)
			return SOCKMGR_NOT_BUFFER;

		if (flags == 0)
		{
			p = (PTEMP_BUFFER)new char[sizeof(TEMP_BUFFER) + sizeof(PACK_TYPE_1) + length];

			if (p == NULL)
				return SOCKMGR_NOT_ENOUGH_MEMORY;

			p->sock = sock;
			p->buffer = (char *)p + sizeof(TEMP_BUFFER);
			p->length = sizeof(PACK_TYPE_1) + length;
			p->pointer = 0;
			p->method = 1;
			p->next = ptb;
			ptb = p;

			PPACK_TYPE_1 pack = (PPACK_TYPE_1)(p->buffer);
			pack->dwPackType = 0;
			pack->nPackSize = length;
			memcpy(pack->bPackData, buffer, length);
		}
		else
		{
			p = (PTEMP_BUFFER)new char[sizeof(TEMP_BUFFER) + sizeof(PACK_TYPE_2) + length + 0x1000];

			if (p == NULL)
				return SOCKMGR_NOT_ENOUGH_MEMORY;

			p->buffer = (char *)p + sizeof(TEMP_BUFFER);

			PPACK_TYPE_2 pack = (PPACK_TYPE_2)(p->buffer);

			if (flags & SOCKMGR_SEND_COMPRESS)
			{
				pack->bCompressed = TRUE;
				pack->nPackSize = Lz77Compress(p->buffer, (void *)buffer, length, 2);
			}
			else
			{
				pack->bCompressed = FALSE;
				pack->nPackSize = length;
			}

			if (flags & SOCKMGR_SEND_ENCRYPT)
			{
				rc5_encrypt(pack->bPackData, pack->bPackData, pack->nPackSize);
				pack->bEncrypted = TRUE;
			}
			else
			{
				pack->bEncrypted = FALSE;
			}

			pack->dwPackType = 1;
			pack->nOriginalSize = length;
			pack->dwCrc32 = crc32(0, pack->bPackData, pack->nPackSize);

			p->sock = sock;
			p->length = sizeof(PACK_TYPE_2) + pack->nPackSize;
			p->pointer = 0;
			p->method = 1;
			p->next = ptb;
			ptb = p;
		}
	}

	LeaveCriticalSection(&SendCS);

	int ret = send(p->sock, p->buffer + p->pointer, p->length - p->pointer, 0);

	if (ret == SOCKET_ERROR)
	{
		if (GetLastError() == WSAEWOULDBLOCK)
			return SOCKMGR_PENDING;
		else
		{
			result = SOCKMGR_SOCK_FAILED;
			goto cleanup;
		}
	}

	p->pointer += ret;

	if (p->pointer == p->length)
	{
		result = SOCKMGR_SUCCESS;
		goto cleanup;
	}

	return SOCKMGR_PENDING;

cleanup:
	EnterCriticalSection(&SendCS);

	for (PTEMP_BUFFER *prev = &ptb; *prev != NULL; prev = &((*prev)->next))
	{
		if ((*prev)->sock == sock && (*prev)->method == 1)
		{
			*prev = (*prev)->next;
			break;
		}
	}

	delete [] p;

	LeaveCriticalSection(&SendCS);

	return result;
}

int CSockMgr::RecvPack(SOCKET sock, IoCallBack funback, LPVOID funparam)
{
	if (Inited == false)
	{
		InitializeCriticalSection(&SendCS);
		InitializeCriticalSection(&RecvCS);
		Inited = true;
	}

	PTEMP_BUFFER p = ptb;
	int result = SOCKMGR_PENDING;
	int ret;

	EnterCriticalSection(&RecvCS);

	while (p != NULL)
	{
		if (p->sock == sock && p->method == 2)
			break;
		p = p->next;
	}

	if (p == NULL)
	{
		p = new TEMP_BUFFER;

		if (p == NULL)
		{
			LeaveCriticalSection(&RecvCS);
			return SOCKMGR_NOT_ENOUGH_MEMORY;
		}

		p->sock = sock;
		p->buffer = (char *)HeapAlloc(GetProcessHeap(), 0, sizeof(PACK_TYPE_1));

		if (p->buffer == NULL)
		{
			delete p;
			LeaveCriticalSection(&RecvCS);
			return SOCKMGR_NOT_ENOUGH_MEMORY;
		}

		p->length = sizeof(PACK_TYPE_1);
		p->pointer = 0;
		p->method = 2;
		p->funback = funback;
		p->funparam = funparam;
		p->next = ptb;
		ptb = p;
	}

	LeaveCriticalSection(&RecvCS);

	if (p->pointer < sizeof(PACK_TYPE_1))
	{
		int ret = recv(p->sock, p->buffer + p->pointer, sizeof(PACK_TYPE_1) - p->pointer, 0);

		if (ret == SOCKET_ERROR)
		{
			if (GetLastError() == WSAEWOULDBLOCK)
				return SOCKMGR_PENDING;
			else
			{
				result = SOCKMGR_SOCK_FAILED;
				goto cleanup;
			}
		}

		p->pointer += ret;

		if (p->pointer < sizeof(PACK_TYPE_1))
			return SOCKMGR_PENDING;

		PPACK_TYPE_1 pack = (PPACK_TYPE_1)(p->buffer);
		p->length = (pack->dwPackType == 0 ? sizeof(PACK_TYPE_1) : sizeof(PACK_TYPE_2)) + pack->nPackSize;
		p->buffer = (char *)HeapReAlloc(GetProcessHeap(), 0, p->buffer, p->length);

		if (p->buffer == NULL)
		{
			result = SOCKMGR_NOT_ENOUGH_MEMORY;
			goto cleanup;
		}
	}

	ret = recv(p->sock, p->buffer + p->pointer, p->length - p->pointer, 0);

	if (ret == SOCKET_ERROR)
	{
		if (GetLastError() == WSAEWOULDBLOCK)
			return SOCKMGR_PENDING;
		else
		{
			result = SOCKMGR_SOCK_FAILED;
			goto cleanup;
		}
	}

	p->pointer += ret;

	if (p->pointer == p->length)
	{
		if (((PPACK_TYPE_1)p->buffer)->dwPackType == 0)
		{
			PPACK_TYPE_1 pack = (PPACK_TYPE_1)(p->buffer);
			PACK_INFO pi;

			pi.bCompress = false;
			pi.bEncrypt = false;
			pi.pData = (char *)pack->bPackData;
			pi.nDataSize = pack->nPackSize;
			pi.nPackSize = pack->nPackSize;

			if (p->funback != NULL)
			{
				p->funback(&pi, p->funparam);
				result = SOCKMGR_SUCCESS;
				goto cleanup;
			}
		}
		else
		{
			PPACK_TYPE_2 pack = (PPACK_TYPE_2)(p->buffer);
			PACK_INFO pi;

			pi.bCompress = pack->bCompressed ? true : false;
			pi.bEncrypt = pack->bEncrypted ? true : false;
			pi.pData = (char *)pack->bPackData;
			pi.nDataSize = pack->nOriginalSize;
			pi.nPackSize = pack->nPackSize;

			if (crc32(0, (char *)pack->bPackData, pack->nPackSize) != pack->dwCrc32)
			{
				result = SOCKMGR_PACK_CRC_FAILED;
				goto cleanup;
			}

			if (pack->bEncrypted != FALSE)
				rc5_decrypt(pack->bPackData, pack->bPackData, pack->nPackSize);

			if (pack->bCompressed != FALSE)
			{
				char *t = (char *)HeapAlloc(GetProcessHeap(), 0, pack->nOriginalSize);

				if (t == NULL)
				{
					result = SOCKMGR_NOT_ENOUGH_MEMORY;
					goto cleanup;
				}

				Lz77Decompress(t, pack->bPackData, pack->nPackSize);
				p->length = pack->nOriginalSize;
				HeapFree(GetProcessHeap(), 0, p->buffer);
				p->buffer = t;
				pi.pData = t;
			}

			if (p->funback != NULL)
			{
				p->funback(&pi, p->funparam);
				result = SOCKMGR_SUCCESS;
				goto cleanup;
			}
		}
	}
	return SOCKMGR_PENDING;

cleanup:
	EnterCriticalSection(&RecvCS);

	for (PTEMP_BUFFER *prev = &ptb; *prev != NULL; prev = &((*prev)->next))
	{
		if ((*prev)->sock == sock && (*prev)->method == 2)
		{
			*prev = (*prev)->next;
			break;
		}
	}

	if (p->buffer != NULL)
		HeapFree(GetProcessHeap(), 0, p->buffer);

	delete p;
	LeaveCriticalSection(&RecvCS);

	return result;
}

void CSockMgr::InterruptSend(SOCKET sock)
{
	if (Inited == false)
		return;

	PTEMP_BUFFER p = NULL;

	EnterCriticalSection(&SendCS);

	for (PTEMP_BUFFER *prev = &ptb; *prev != NULL; prev = &((*prev)->next))
	{
		if ((*prev)->sock == sock && (*prev)->method == 1)
		{
			p = *prev;
			*prev = (*prev)->next;
			break;
		}
	}

	if (p != NULL)
		delete [] p;

	LeaveCriticalSection(&SendCS);
}

void CSockMgr::InterruptRecv(SOCKET sock)
{
	if (Inited == false)
		return;

	PTEMP_BUFFER p = NULL;

	EnterCriticalSection(&RecvCS);

	for (PTEMP_BUFFER *prev = &ptb; *prev != NULL; prev = &((*prev)->next))
	{
		if ((*prev)->sock == sock && (*prev)->method == 2)
		{
			p = *prev;
			*prev = (*prev)->next;
			break;
		}
	}

	if (p != NULL)
	{
		if (p->buffer != NULL)
			HeapFree(GetProcessHeap(), 0, p->buffer);

		delete p;
	}

	LeaveCriticalSection(&RecvCS);
}