wltProxy/SProxy/CTool.cpp

#include "stdafx.h"
#include "CTool.h"

#include <WinSock.h>
#include <tcpmib.h>
#include <IPHlpApi.h>
#pragma comment(lib, "IPHlpApi.lib")
#pragma comment(lib, "WS2_32.lib")

CTool* SSingleton<CTool>::ms_Singleton = NULL;


CTool::CTool()
{
	WSADATA wsaData;
	WSAStartup(0x0201, &wsaData);
}

CTool::~CTool()
{
	WSACleanup();
}

std::vector<uint16_t> CTool::GetAllTcpConnectionsPort()
{
	std::vector<uint16_t> ret;
	ULONG size = 0;
	GetTcpTable(NULL, &size, TRUE);
	std::unique_ptr<char[]> buffer(new char[size]);

	PMIB_TCPTABLE tcpTable = reinterpret_cast<PMIB_TCPTABLE>(buffer.get());
	if (GetTcpTable(tcpTable, &size, FALSE) == NO_ERROR)
		for (size_t i = 0; i < tcpTable->dwNumEntries; i++)
			ret.push_back(ntohs((uint16_t)tcpTable->table[i].dwLocalPort));
	std::sort(std::begin(ret), std::end(ret));
	return ret;
}

std::vector<uint16_t> CTool::GetAllUdpConnectionsPort()
{
	std::vector<uint16_t> ret;
	ULONG size = 0;
	GetUdpTable(NULL, &size, TRUE);
	std::unique_ptr<char[]> buffer(new char[size]);

	PMIB_UDPTABLE udpTable = reinterpret_cast<PMIB_UDPTABLE>(buffer.get());
	if (GetUdpTable(udpTable, &size, FALSE) == NO_ERROR)
		for (size_t i = 0; i < udpTable->dwNumEntries; i++)
			ret.push_back(ntohs((uint16_t)udpTable->table[i].dwLocalPort));
	std::sort(std::begin(ret), std::end(ret));
	return ret;
}

uint16_t CTool::FindAvailableTcpPort(uint16_t begin, uint16_t end)
{
	auto vec = GetAllTcpConnectionsPort();
	for (uint16_t port = begin; port != end; ++port)
		if (!std::binary_search(std::begin(vec), std::end(vec), port))
			return port;
	return 0;
}

uint16_t CTool::FindAvailableUdpPort(uint16_t begin, uint16_t end)
{
	auto vecTcp = GetAllTcpConnectionsPort(),
		vecUdp = GetAllUdpConnectionsPort();
	for (uint16_t port = begin; port != end; ++port)
		if (!std::binary_search(std::begin(vecTcp), std::end(vecTcp), port) &&
			!std::binary_search(std::begin(vecUdp), std::end(vecUdp), port))
			return port;
	return 0;
}