1 files changed, 197 insertions, 136 deletions
diff --git a/Swiften/FileTransfer/SOCKS5BytestreamServerManager.cpp b/Swiften/FileTransfer/SOCKS5BytestreamServerManager.cpp
index 1d65d27..7fc1388 100644
--- a/Swiften/FileTransfer/SOCKS5BytestreamServerManager.cpp
+++ b/Swiften/FileTransfer/SOCKS5BytestreamServerManager.cpp
@@ -1,7 +1,7 @@
 /*
- * Copyright (c) 2012 Remko Tronçon
- * Licensed under the GNU General Public License v3.
- * See Documentation/Licenses/GPLv3.txt for more information.
+ * Copyright (c) 2012-2019 Isode Limited.
+ * All rights reserved.
+ * See the COPYING file for more information.
  */
 
 /*
@@ -12,19 +12,21 @@
 
 #include <Swiften/FileTransfer/SOCKS5BytestreamServerManager.h>
 
-#include <boost/smart_ptr/make_shared.hpp>
+#include <memory>
+
 #include <boost/bind.hpp>
 
-#include <Swiften/FileTransfer/SOCKS5BytestreamServerInitializeRequest.h>
-#include <Swiften/Base/foreach.h>
 #include <Swiften/Base/Log.h>
 #include <Swiften/FileTransfer/SOCKS5BytestreamServer.h>
+#include <Swiften/FileTransfer/SOCKS5BytestreamServerPortForwardingUser.h>
+#include <Swiften/FileTransfer/SOCKS5BytestreamServerResourceUser.h>
 #include <Swiften/Network/ConnectionServer.h>
 #include <Swiften/Network/ConnectionServerFactory.h>
-#include <Swiften/Network/NetworkEnvironment.h>
-#include <Swiften/Network/NATTraverser.h>
-#include <Swiften/Network/NATTraversalGetPublicIPRequest.h>
 #include <Swiften/Network/NATTraversalForwardPortRequest.h>
+#include <Swiften/Network/NATTraversalGetPublicIPRequest.h>
+#include <Swiften/Network/NATTraversalRemovePortForwardingRequest.h>
+#include <Swiften/Network/NATTraverser.h>
+#include <Swiften/Network/NetworkEnvironment.h>
 
 using namespace Swift;
 
@@ -32,164 +34,223 @@ static const int LISTEN_PORTS_BEGIN = 10000;
 static const int LISTEN_PORTS_END = 11000;
 
 SOCKS5BytestreamServerManager::SOCKS5BytestreamServerManager(
-		SOCKS5BytestreamRegistry* bytestreamRegistry,
-		ConnectionServerFactory* connectionServerFactory, 
-		NetworkEnvironment* networkEnvironment,
-		NATTraverser* natTraverser) :
-			bytestreamRegistry(bytestreamRegistry),
-			connectionServerFactory(connectionServerFactory),
-			networkEnvironment(networkEnvironment),
-			natTraverser(natTraverser),
-			state(Start),
-			server(NULL) {
+        SOCKS5BytestreamRegistry* bytestreamRegistry,
+        ConnectionServerFactory* connectionServerFactory,
+        NetworkEnvironment* networkEnvironment,
+        NATTraverser* natTraverser) :
+            bytestreamRegistry(bytestreamRegistry),
+            connectionServerFactory(connectionServerFactory),
+            networkEnvironment(networkEnvironment),
+            natTraverser(natTraverser),
+            state(Start),
+            server(nullptr),
+            attemptedPortMapping_(false) {
 }
 
 SOCKS5BytestreamServerManager::~SOCKS5BytestreamServerManager() {
-	SWIFT_LOG_ASSERT(!connectionServer, warning) << std::endl;
-	SWIFT_LOG_ASSERT(!getPublicIPRequest, warning) << std::endl;
-	SWIFT_LOG_ASSERT(!forwardPortRequest, warning) << std::endl;
-	SWIFT_LOG_ASSERT(state == Start, warning) << std::endl;
+    SWIFT_LOG_ASSERT(!connectionServer, warning);
+    SWIFT_LOG_ASSERT(!getPublicIPRequest, warning);
+    SWIFT_LOG_ASSERT(!forwardPortRequest, warning);
+    SWIFT_LOG_ASSERT(state == Start, warning);
+    if (portMapping && !unforwardPortRequest) {
+        SWIFT_LOG(warning) << "Port forwarding still alive. Trying to remove it now.";
+        unforwardPortRequest = natTraverser->createRemovePortForwardingRequest(portMapping.get().getLocalPort(), portMapping.get().getPublicPort());
+        unforwardPortRequest->start();
+    }
 }
 
-
-boost::shared_ptr<SOCKS5BytestreamServerInitializeRequest> SOCKS5BytestreamServerManager::createInitializeRequest() {
-	return boost::make_shared<SOCKS5BytestreamServerInitializeRequest>(this);
+std::shared_ptr<SOCKS5BytestreamServerResourceUser> SOCKS5BytestreamServerManager::aquireResourceUser() {
+    std::shared_ptr<SOCKS5BytestreamServerResourceUser> resourceUser;
+    if (s5bServerResourceUser_.expired()) {
+        resourceUser = std::make_shared<SOCKS5BytestreamServerResourceUser>(this);
+        s5bServerResourceUser_ = resourceUser;
+    }
+    else {
+        resourceUser = s5bServerResourceUser_.lock();
+    }
+    return resourceUser;
 }
 
-void SOCKS5BytestreamServerManager::stop() {
-	if (getPublicIPRequest) {
-		getPublicIPRequest->stop();
-		getPublicIPRequest.reset();
-	}
-	if (forwardPortRequest) {
-		forwardPortRequest->stop();
-		forwardPortRequest.reset();
-	}
-	if (connectionServer) {
-		connectionServer->stop();
-		connectionServer.reset();
-	}
-	// TODO: Remove port forwards
-	
-	state = Start;
+std::shared_ptr<SOCKS5BytestreamServerPortForwardingUser> SOCKS5BytestreamServerManager::aquirePortForwardingUser() {
+    std::shared_ptr<SOCKS5BytestreamServerPortForwardingUser> portForwardingUser;
+    if (s5bServerPortForwardingUser_.expired()) {
+        portForwardingUser = std::make_shared<SOCKS5BytestreamServerPortForwardingUser>(this);
+        s5bServerPortForwardingUser_ = portForwardingUser;
+    }
+    else {
+        portForwardingUser = s5bServerPortForwardingUser_.lock();
+    }
+    return portForwardingUser;
 }
 
 std::vector<HostAddressPort> SOCKS5BytestreamServerManager::getHostAddressPorts() const {
-	std::vector<HostAddressPort> result;
-	if (connectionServer) {
-		std::vector<NetworkInterface> networkInterfaces = networkEnvironment->getNetworkInterfaces();
-		foreach (const NetworkInterface& networkInterface, networkInterfaces) {
-			foreach (const HostAddress& address, networkInterface.getAddresses()) {
-				result.push_back(HostAddressPort(address, connectionServerPort));
-			}
-		}
-	}
-	return result;
+    std::vector<HostAddressPort> result;
+    if (connectionServer) {
+        std::vector<NetworkInterface> networkInterfaces = networkEnvironment->getNetworkInterfaces();
+        for (const auto& networkInterface : networkInterfaces) {
+            for (const auto& address : networkInterface.getAddresses()) {
+                result.push_back(HostAddressPort(address, connectionServerPort));
+            }
+        }
+    }
+    return result;
 }
 
 std::vector<HostAddressPort> SOCKS5BytestreamServerManager::getAssistedHostAddressPorts() const {
-	std::vector<HostAddressPort> result;
-	if (publicAddress && portMapping) {
-		result.push_back(HostAddressPort(*publicAddress, portMapping->getPublicPort()));
-	}
-	return result;
+    std::vector<HostAddressPort> result;
+    if (publicAddress && portMapping) {
+        result.push_back(HostAddressPort(*publicAddress, portMapping->getPublicPort()));
+    }
+    return result;
 }
 
 bool SOCKS5BytestreamServerManager::isInitialized() const {
-	return state == Initialized;
+    return state == Initialized;
 }
 
 void SOCKS5BytestreamServerManager::initialize() {
-	if (state == Start) {
-		state = Initializing;
-
-		// Find a port to listen on
-		assert(!connectionServer);
-		int port;
-		for (port = LISTEN_PORTS_BEGIN; port < LISTEN_PORTS_END; ++port) {
-			SWIFT_LOG(debug) << "Trying to start server on port " << port << std::endl;
-			connectionServer = connectionServerFactory->createConnectionServer(HostAddress("0.0.0.0"), port);
-			boost::optional<ConnectionServer::Error> error = connectionServer->tryStart();
-			if (!error) {
-				break;
-			}
-			else if (*error != ConnectionServer::Conflict) {
-				SWIFT_LOG(debug) << "Error starting server" << std::endl;
-				onInitialized(false);
-				return;
-			}
-			connectionServer.reset();
-		}
-		if (!connectionServer) {
-			SWIFT_LOG(debug) << "Unable to find an open port" << std::endl;
-			onInitialized(false);
-			return;
-		}
-		SWIFT_LOG(debug) << "Server started succesfully" << std::endl;
-		connectionServerPort = port;
-
-		// Start bytestream server. Should actually happen before the connectionserver is started
-		// but that doesn't really matter here.
-		assert(!server);
-		server = new SOCKS5BytestreamServer(connectionServer, bytestreamRegistry);
-		server->start();
-
-		// Retrieve public addresses
-		assert(!getPublicIPRequest);
-		publicAddress = boost::optional<HostAddress>();
-		if ((getPublicIPRequest = natTraverser->createGetPublicIPRequest())) {
-			getPublicIPRequest->onResult.connect(
-					boost::bind(&SOCKS5BytestreamServerManager::handleGetPublicIPResult, this, _1));
-			getPublicIPRequest->start();
-		}
-
-		// Forward ports
-		assert(!forwardPortRequest);
-		portMapping = boost::optional<NATPortMapping>();
-		if ((forwardPortRequest = natTraverser->createForwardPortRequest(port, port))) {
-			forwardPortRequest->onResult.connect(
-					boost::bind(&SOCKS5BytestreamServerManager::handleForwardPortResult, this, _1));
-			forwardPortRequest->start();
-		}
-	}
+    if (state == Start) {
+        state = Initializing;
+
+        // Find a port to listen on
+        assert(!connectionServer);
+        unsigned short port;
+        for (port = LISTEN_PORTS_BEGIN; port < LISTEN_PORTS_END; ++port) {
+            SWIFT_LOG(debug) << "Trying to start server on port " << port;
+            connectionServer = connectionServerFactory->createConnectionServer(HostAddress::fromString("::").get(), port);
+            boost::optional<ConnectionServer::Error> error = connectionServer->tryStart();
+            if (!error) {
+                break;
+            }
+            else if (*error != ConnectionServer::Conflict) {
+                SWIFT_LOG(debug) << "Error starting server";
+                onInitialized(false);
+                return;
+            }
+            connectionServer.reset();
+        }
+        if (!connectionServer) {
+            SWIFT_LOG(debug) << "Unable to find an open port";
+            onInitialized(false);
+            return;
+        }
+        SWIFT_LOG(debug) << "Server started succesfully";
+        connectionServerPort = port;
+
+        // Start bytestream server. Should actually happen before the connectionserver is started
+        // but that doesn't really matter here.
+        assert(!server);
+        server = new SOCKS5BytestreamServer(connectionServer, bytestreamRegistry);
+        server->start();
+        checkInitializeFinished();
+    }
 }
 
-void SOCKS5BytestreamServerManager::handleGetPublicIPResult(boost::optional<HostAddress> address) {
-	if (address) {
-		SWIFT_LOG(debug) << "Public IP discovered as " << address.get().toString() << "." << std::endl;
-	} 
-	else {
-		SWIFT_LOG(debug) << "No public IP discoverable." << std::endl;
-	}
+bool SOCKS5BytestreamServerManager::isPortForwardingReady() const {
+    return attemptedPortMapping_ && !getPublicIPRequest && !forwardPortRequest;
+}
 
-	publicAddress = address;
+void SOCKS5BytestreamServerManager::setupPortForwarding() {
+    assert(server);
+    attemptedPortMapping_ = true;
 
-	getPublicIPRequest->stop();
-	getPublicIPRequest.reset();
+    // Retrieve public addresses
+    assert(!getPublicIPRequest);
+    publicAddress = boost::optional<HostAddress>();
+    if ((getPublicIPRequest = natTraverser->createGetPublicIPRequest())) {
+        getPublicIPRequest->onResult.connect(
+                boost::bind(&SOCKS5BytestreamServerManager::handleGetPublicIPResult, this, _1));
+        getPublicIPRequest->start();
+    }
 
-	checkInitializeFinished();
+    // Forward ports
+    auto port = server->getAddressPort().getPort();
+    assert(!forwardPortRequest);
+    portMapping = boost::optional<NATPortMapping>();
+    if ((forwardPortRequest = natTraverser->createForwardPortRequest(port, port))) {
+        forwardPortRequest->onResult.connect(
+                boost::bind(&SOCKS5BytestreamServerManager::handleForwardPortResult, this, _1));
+        forwardPortRequest->start();
+    }
+}
+
+void SOCKS5BytestreamServerManager::removePortForwarding() {
+    // remove port forwards
+    if (portMapping) {
+        unforwardPortRequest = natTraverser->createRemovePortForwardingRequest(portMapping.get().getLocalPort(), portMapping.get().getPublicPort());
+        unforwardPortRequest->onResult.connect(boost::bind(&SOCKS5BytestreamServerManager::handleUnforwardPortResult, this, _1));
+        unforwardPortRequest->start();
+    }
+}
+
+void SOCKS5BytestreamServerManager::stop() {
+    if (getPublicIPRequest) {
+        getPublicIPRequest->stop();
+        getPublicIPRequest.reset();
+    }
+    if (forwardPortRequest) {
+        forwardPortRequest->stop();
+        forwardPortRequest.reset();
+    }
+    if (unforwardPortRequest) {
+        unforwardPortRequest->stop();
+        unforwardPortRequest.reset();
+    }
+    if (server) {
+        server->stop();
+        delete server;
+        server = nullptr;
+    }
+    if (connectionServer) {
+        connectionServer->stop();
+        connectionServer.reset();
+    }
+
+    state = Start;
+}
+
+void SOCKS5BytestreamServerManager::handleGetPublicIPResult(boost::optional<HostAddress> address) {
+    if (address) {
+        SWIFT_LOG(debug) << "Public IP discovered as " << address.get().toString() << ".";
+    }
+    else {
+        SWIFT_LOG(debug) << "No public IP discoverable.";
+    }
+
+    publicAddress = address;
+
+    getPublicIPRequest->stop();
+    getPublicIPRequest.reset();
 }
 
 void SOCKS5BytestreamServerManager::handleForwardPortResult(boost::optional<NATPortMapping> mapping) {
-	if (mapping) {
-		SWIFT_LOG(debug) << "Mapping port was successful." << std::endl;
-	} 
-	else {
-		SWIFT_LOG(debug) << "Mapping port has failed." << std::endl;
-	}
+    if (mapping) {
+        SWIFT_LOG(debug) << "Mapping port was successful.";
+    }
+    else {
+        SWIFT_LOG(debug) << "Mapping port has failed.";
+    }
 
-	portMapping = mapping;
+    portMapping = mapping;
+    onPortForwardingSetup(mapping.is_initialized());
 
-	forwardPortRequest->stop();
-	forwardPortRequest.reset();
+    forwardPortRequest->stop();
+    forwardPortRequest.reset();
+}
 
-	checkInitializeFinished();
+void SOCKS5BytestreamServerManager::handleUnforwardPortResult(boost::optional<bool> result) {
+    if (result.is_initialized() && result.get()) {
+        portMapping.reset();
+    }
+    else {
+        SWIFT_LOG(warning) << "Failed to remove port forwarding.";
+    }
+    attemptedPortMapping_ = false;
+    unforwardPortRequest.reset();
 }
 
 void SOCKS5BytestreamServerManager::checkInitializeFinished() {
-	assert(state == Initializing);
-	if (!getPublicIPRequest && !forwardPortRequest) {
-		state = Initialized;
-		onInitialized(true);
-	}
+    assert(state == Initializing);
+    state = Initialized;
+    onInitialized(true);
 }