www.pudn.com > ntshell.rar > 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); 
}