Mesen-SX/Utilities/UPnPPortMapper.cpp

#include "stdafx.h"
#include "UPnPPortMapper.h"

#ifdef _WIN32
#include <winsock2.h>
#include <natupnp.h>
#include <ws2tcpip.h>

bool UPnPPortMapper::AddNATPortMapping(uint16_t internalPort, uint16_t externalPort, IPProtocol protocol)
{
	bool result = false;

	CoInitializeEx(nullptr, COINIT_MULTITHREADED);

	IUPnPNAT* nat = nullptr;
	HRESULT hResult = CoCreateInstance(__uuidof(UPnPNAT), nullptr, CLSCTX_ALL, __uuidof(IUPnPNAT), (void**)&nat);

	BSTR proto = SysAllocString((protocol == IPProtocol::TCP) ? L"TCP" : L"UDP");

	if (SUCCEEDED(hResult) && nat)
	{
		IStaticPortMappingCollection* spmc = nullptr;
		hResult = nat->get_StaticPortMappingCollection(&spmc);
		if (SUCCEEDED(hResult) && spmc)
		{
			IStaticPortMapping* spm = nullptr;
			hResult = spmc->get_Item(externalPort, proto, &spm);
			if (spm != nullptr)
			{
				//An identical mapping already exists, remove it
				if (RemoveNATPortMapping(externalPort, protocol))
				{
					std::cout << "Removed existing UPnP mapping." << std::endl;
					spm->Release();
					spm = nullptr;
				}
			}

			if (!SUCCEEDED(hResult) || spm == nullptr)
			{
				std::cout << "Attempting to automatically forward port via UPnP..." << std::endl;

				vector<wstring> localIPs = GetLocalIPs();
				BSTR desc = SysAllocString(L"Mesen NetPlay");
				spm = nullptr;

				for (size_t i = 0, len = localIPs.size(); i < len; i++)
				{
					BSTR clientStr = SysAllocString(localIPs[i].c_str());
					hResult = spmc->Add(externalPort, proto, internalPort, clientStr, true, desc, &spm);
					SysFreeString(clientStr);
					SysFreeString(desc);

					if (SUCCEEDED(hResult) && spm)
					{
						//Successfully added a new port mapping
						std::cout << std::dec << "Forwarded port " << externalPort << " to IP " <<
							utf8::utf8::encode(localIPs[i]) << std::endl;
						result = true;
					}
					else
					{
						std::cout << "Unable to add UPnP port mapping. IP: " << utf8::utf8::encode(localIPs[i]) <<
							" HRESULT: 0x" << std::hex << hResult << std::endl;
					}

					if (spm)
					{
						spm->Release();
					}
				}
			}
			else
			{
				std::cout << "Unable to add UPnP port mapping." << std::endl;
			}
			spmc->Release();
		}
		nat->Release();
	}

	SysFreeString(proto);

	CoUninitialize();

	return result;
}

bool UPnPPortMapper::RemoveNATPortMapping(uint16_t externalPort, IPProtocol protocol)
{
	IUPnPNAT* nat = nullptr;
	IStaticPortMappingCollection* spmc;
	bool result = false;

	CoInitializeEx(nullptr, COINIT_MULTITHREADED);

	HRESULT hResult = ::CoCreateInstance(__uuidof(UPnPNAT), nullptr, CLSCTX_ALL, __uuidof(IUPnPNAT), (void**)&nat);

	BSTR proto = SysAllocString((protocol == IPProtocol::TCP) ? L"TCP" : L"UDP");

	if (SUCCEEDED(hResult) && nat)
	{
		hResult = nat->get_StaticPortMappingCollection(&spmc);
		if (SUCCEEDED(hResult) && spmc)
		{
			spmc->Remove(externalPort, proto);
			spmc->Release();
			result = true;
		}
		nat->Release();
	}

	SysFreeString(proto);

	CoUninitialize();

	return result;
}

vector<wstring> UPnPPortMapper::GetLocalIPs()
{
	vector<wstring> localIPs;
	ADDRINFOW* result = nullptr;
	ADDRINFOW* current = nullptr;
	ADDRINFOW hints;

	ZeroMemory(&hints, sizeof(hints));
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;

	wchar_t hostName[255];
	DWORD hostSize = 255;
	GetComputerName(hostName, &hostSize);

	if (GetAddrInfoW(hostName, nullptr, &hints, &result) == 0)
	{
		current = result;
		while (current != nullptr)
		{
			wchar_t ipAddr[255];
			DWORD ipSize = 255;

			if (WSAAddressToString(current->ai_addr, (DWORD)current->ai_addrlen, nullptr, ipAddr, &ipSize) == 0)
			{
				if (std::find(localIPs.begin(), localIPs.end(), ipAddr) == localIPs.end())
				{
					localIPs.push_back(ipAddr);
				}
			}
			current = current->ai_next;
		}
		FreeAddrInfoW(result);
	}

	return localIPs;
}

#else
	
bool UPnPPortMapper::AddNATPortMapping(uint16_t internalPort, uint16_t externalPort, IPProtocol protocol) 
{
	return false;
}

bool UPnPPortMapper::RemoveNATPortMapping(uint16_t externalPort, IPProtocol protocol) 
{
	return false;
}

vector<wstring> UPnPPortMapper::GetLocalIPs()
{
	return vector<wstring>();
}	
	
#endif