test_tlv_server.cpp

#undef UNICODE

#define WIN32_LEAN_AND_MEAN

#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>

#pragma comment (lib, "Ws2_32.lib")

#define SERVER_PORT "4444"


#pragma pack(push, 1)

typedef struct COMMON_PACKET_HEADER {
	uint32_t CommandId;
	uint32_t BodySize;
} COMMON_PACKET_HEADER;

typedef struct CHUNK_STRUCT {
	uint32_t ChunkId;
	int size;
	uint8_t *pBuf;
} CHUNK_STRUCT;

#pragma pack(pop)

#define CHUNK_LIST_SIZE 256
CHUNK_STRUCT** chunk_list = NULL;

void process_packet(char *packet, int packet_size) {
	COMMON_PACKET_HEADER* hdr = (COMMON_PACKET_HEADER*)packet;
							
	// this code will patch vuln
	/*if (hdr->BodySize < (packet_size - sizeof(COMMON_PACKET_HEADER))) {
		printf("[!] hdr->BodySize < (packet_size - sizeof(COMMON_PACKET_HEADER))\n");
		return;
	}*/

	uint8_t* body = (uint8_t*)(&packet[sizeof(COMMON_PACKET_HEADER)]);
	uint32_t pktIdx = 0;
	uint32_t ChunkId = 0;
	CHUNK_STRUCT* new_chunk = NULL;
	int available_slot = -1;

	printf("[+] CommandId = %d\n", hdr->CommandId);

	switch (hdr->CommandId) {
	case 1:
		// Allocate chunk
		
		for (int i = 0; i < CHUNK_LIST_SIZE; i++) {
			if (chunk_list[i] == NULL) {
				available_slot = i;
				break;
			}
		}

		if (available_slot == -1) {
			printf("[!] there's no available slot on chunk list.\n");
			return;
		}

		ChunkId = *((uint32_t *)(&body[pktIdx]));
		pktIdx += sizeof(uint32_t);

		new_chunk = (CHUNK_STRUCT*)calloc(1, sizeof(CHUNK_STRUCT));
		if (new_chunk == NULL) {
			printf("[!] calloc(%d) failed, GLE = 0x%08x\n", sizeof(CHUNK_STRUCT), GetLastError());
			return;
		}

		new_chunk->ChunkId = ChunkId;
		new_chunk->size = hdr->BodySize - sizeof(uint32_t);

		new_chunk->pBuf = (uint8_t*)calloc(1, new_chunk->size);

		if (new_chunk->pBuf == NULL) {
			printf("[!] calloc(%d) failed, GLE = 0x%08x\n", new_chunk->size, GetLastError());
			free(new_chunk);
			return;
		}

		memcpy(new_chunk->pBuf, &body[pktIdx], new_chunk->size);
		pktIdx += new_chunk->size;

		chunk_list[available_slot] = new_chunk;
		break;

	case 2:
		// Edit chunk
		ChunkId = *((uint32_t*)(&body[pktIdx]));
		pktIdx += sizeof(uint32_t);

		for (int i = 0; i < CHUNK_LIST_SIZE; i++) {
			if (chunk_list[i] == NULL) {
				continue;
			}

			if (chunk_list[i]->ChunkId == ChunkId) {

				// obviously, this one has integer underflow bug
				int new_buf_size = hdr->BodySize - sizeof(uint32_t);

				// this boundary check returns false if hdr->BodySize < sizeof(uint32_t)
				if (new_buf_size > chunk_list[i]->size) {
					uint8_t* new_buf = (uint8_t*)calloc(1, new_buf_size);
					if (new_buf == NULL) {
						printf("[!] calloc(%d) failed, GLE = 0x%08x\n", new_buf_size, GetLastError());
						return;
					}

					free(chunk_list[i]->pBuf);
					chunk_list[i]->pBuf = new_buf;
				}

				memcpy(chunk_list[i]->pBuf, &body[pktIdx], new_buf_size);
				pktIdx += new_buf_size;

				return;
			}
		}

		printf("[!] Couldn't find ChunkId 0x%x\n", ChunkId);
		break;

	case 3:
		// Delete Chunk
		ChunkId = *((uint32_t*)(&body[pktIdx]));
		pktIdx += sizeof(uint32_t);

		for (int i = 0; i < CHUNK_LIST_SIZE; i++) {
			if (chunk_list[i] == NULL) {
				continue;
			}

			if (chunk_list[i]->ChunkId == ChunkId) {

				free(chunk_list[i]->pBuf);
				free(chunk_list[i]);
				chunk_list[i] = NULL;

				printf("[+] ChunkId 0x%x deleted\n", ChunkId);

				return;
			}
		}

		printf("[!] Couldn't find ChunkId 0x%x\n", ChunkId);
		break;

	default:
		break;
	}
}


int main() {
	// copied from https://docs.microsoft.com/ko-kr/windows/win32/winsock/complete-server-code

	WSADATA wsaData = { 0, };
	SOCKET ListenSocket = INVALID_SOCKET;
	SOCKET ClientSocket = INVALID_SOCKET;
	struct addrinfo hints = { 0, };
	struct addrinfo *pResult = NULL;

	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
		printf("[!] WSAStartup failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		return 1;
	}

	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
	hints.ai_flags = AI_PASSIVE;
	
	if (getaddrinfo(NULL, SERVER_PORT, &hints, &pResult) != 0) {
		printf("[!] getaddrinfo failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		WSACleanup();
		return 1;
	}

	ListenSocket = socket(pResult->ai_family, pResult->ai_socktype, pResult->ai_protocol);
	if (ListenSocket == INVALID_SOCKET) {
		printf("[!] socket failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		freeaddrinfo(pResult);
		WSACleanup();
		return 1;
	}

	if (bind(ListenSocket, pResult->ai_addr, pResult->ai_addrlen) == SOCKET_ERROR) {
		printf("[!] bind failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		closesocket(ListenSocket);
		freeaddrinfo(pResult);
		WSACleanup();
		return 1;
	}

	freeaddrinfo(pResult);

	printf("[+] Listening PORT %s...\n", SERVER_PORT);
	if (listen(ListenSocket, SOMAXCONN) == SOCKET_ERROR) {
		printf("[!] listen failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		closesocket(ListenSocket);
		WSACleanup();
		return 1;
	}

	ClientSocket = accept(ListenSocket, NULL, NULL); 
	if (ClientSocket == INVALID_SOCKET) {
		printf("[!] accept failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		closesocket(ListenSocket);
		WSACleanup();
		return 1;
	}

	printf("[+] accept done!\n");

	closesocket(ListenSocket);


	chunk_list = (CHUNK_STRUCT**)calloc(1, sizeof(CHUNK_STRUCT*) * CHUNK_LIST_SIZE);

	uint32_t buf_size = 0;
	char* buf = NULL;
	int recv_result = 0;
	while (1) {
		recv_result = recv(ClientSocket, (char*)&buf_size, sizeof(uint32_t), 0);
		
		if (recv_result == 0) {
			printf("[!] recv closed, bailout\n");
			break;
		}
		else if (recv_result == SOCKET_ERROR) {
			printf("[!] recv failed, WSA GLE = 0x%08x\n", WSAGetLastError());
			break;
		}

		if (buf_size >= 0x400) {
			printf("[!] buf_size( 0x%x ) >= 0x400, too big\n", buf_size);
			continue;
		}
		else if (buf_size == 0) {
			printf("[+] server exit\n");
			break;
		}

		buf = (char*)calloc(1, buf_size);
		if (buf == NULL) {
			printf("[!] calloc(%d) failed, GLE = 0x%08x\n", buf_size, GetLastError());
			break;
		}

		printf("[+] buffer size = %d\n", buf_size);

		recv_result = recv(ClientSocket, buf, buf_size, 0);

		if (recv_result == 0) {
			printf("recv closed, bailout\n");
			break;
		}
		else if (recv_result == SOCKET_ERROR) {
			printf("recv failed, WSA GLE = 0x%08x\n", WSAGetLastError());
			break;
		}

		process_packet(buf, buf_size);

		free(buf);
	}


	free(chunk_list);

	if(shutdown(ClientSocket, SD_SEND) == SOCKET_ERROR) {
		printf("shutdown failed, WSA GLE = 0x%08x\n", WSAGetLastError());
		closesocket(ClientSocket);
		WSACleanup();
		return 1;
	}

	closesocket(ClientSocket);
	WSACleanup();
	return 0;
}