Mesen/Utilities/UPnPPortMapper.cpp
2016-12-11 10:56:23 -05:00

156 lines
4.0 KiB
C++

#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