1 files changed, 17 insertions, 11 deletions
diff --git a/Swiften/FileTransfer/ConnectivityManager.cpp b/Swiften/FileTransfer/ConnectivityManager.cpp
index de5eccb..f4e7e96 100644
--- a/Swiften/FileTransfer/ConnectivityManager.cpp
+++ b/Swiften/FileTransfer/ConnectivityManager.cpp
@@ -23,39 +23,44 @@ ConnectivityManager::ConnectivityManager(NATTraverser* worker) : natTraversalWor
 }
 
 ConnectivityManager::~ConnectivityManager() {
-	std::set<int> leftOpenPorts = ports;
-	foreach(int port, leftOpenPorts) {
-		removeListeningPort(port);
+	std::set<int> leftOpenUDPPorts = udpPorts;
+	foreach(int port, leftOpenUDPPorts) {
+		removeListeningPort(port, NATPortMapping::UDP);
+	}
+	std::set<int> leftOpenTCPPorts = tcpPorts;
+	foreach(int port, leftOpenTCPPorts) {
+		removeListeningPort(port, NATPortMapping::TCP);
 	}
 }
 
-void ConnectivityManager::addListeningPort(int port) {
-	ports.insert(port);
+void ConnectivityManager::addListeningPort(int port, NATPortMapping::Protocol protocol) {
+	udpPorts.insert(port);
 	boost::shared_ptr<NATTraversalGetPublicIPRequest> getIPRequest = natTraversalWorker->createGetPublicIPRequest();
 	if (getIPRequest) {
 		getIPRequest->onResult.connect(boost::bind(&ConnectivityManager::natTraversalGetPublicIPResult, this, _1));
 		getIPRequest->run();
 	}
 
-	boost::shared_ptr<NATTraversalForwardPortRequest> forwardPortRequest = natTraversalWorker->createForwardPortRequest(port, port);
+	boost::shared_ptr<NATTraversalForwardPortRequest> forwardPortRequest = natTraversalWorker->createForwardPortRequest(port, port, protocol);
 	if (forwardPortRequest) {
 		forwardPortRequest->onResult.connect(boost::bind(&ConnectivityManager::natTraversalForwardPortResult, this, _1));
 		forwardPortRequest->run();
 	}
 }
 
-void ConnectivityManager::removeListeningPort(int port) {
+void ConnectivityManager::removeListeningPort(int port, NATPortMapping::Protocol protocol) {
 	SWIFT_LOG(debug) << "remove listening port " << port << std::endl;
-	ports.erase(port);
-	boost::shared_ptr<NATTraversalRemovePortForwardingRequest> removePortForwardingRequest = natTraversalWorker->createRemovePortForwardingRequest(port, port);
+	udpPorts.erase(port);
+	boost::shared_ptr<NATTraversalRemovePortForwardingRequest> removePortForwardingRequest = natTraversalWorker->createRemovePortForwardingRequest(port, port, protocol);
 	if (removePortForwardingRequest) {
 		removePortForwardingRequest->run();
 	}
 }
 
-std::vector<HostAddressPort> ConnectivityManager::getHostAddressPorts() const {
+std::vector<HostAddressPort> ConnectivityManager::getHostAddressPorts(NATPortMapping::Protocol protocol) const {
 	PlatformNetworkEnvironment env;
 	std::vector<HostAddressPort> results;
+	const std::set<int>& ports = (protocol == NATPortMapping::TCP ? udpPorts : tcpPorts);
 
 	//std::vector<HostAddress> addresses;
 
@@ -71,10 +76,11 @@ std::vector<HostAddressPort> ConnectivityManager::getHostAddressPorts() const {
 	return results;
 }
 
-std::vector<HostAddressPort> ConnectivityManager::getAssistedHostAddressPorts() const {
+std::vector<HostAddressPort> ConnectivityManager::getAssistedHostAddressPorts(NATPortMapping::Protocol protocol) const {
 	std::vector<HostAddressPort> results;
 
 	if (publicAddress) {
+		const std::set<int>& ports = (protocol == NATPortMapping::TCP ? udpPorts : tcpPorts);
 		foreach (int port, ports) {
 			results.push_back(HostAddressPort(publicAddress.get(), port));
 		}