1 files changed, 31 insertions, 6 deletions
diff --git a/Swiften/FileTransfer/SOCKS5BytestreamProxyFinder.cpp b/Swiften/FileTransfer/SOCKS5BytestreamProxyFinder.cpp
index 9d7505b..78cf2e6 100644
--- a/Swiften/FileTransfer/SOCKS5BytestreamProxyFinder.cpp
+++ b/Swiften/FileTransfer/SOCKS5BytestreamProxyFinder.cpp
@@ -4,12 +4,19 @@
 * See Documentation/Licenses/BSD-simplified.txt for more information.
 */
+/*
+ * Copyright (c) 2015 Isode Limited.
+ * All rights reserved.
+ * See the COPYING file for more information.
+ */
 #include <Swiften/FileTransfer/SOCKS5BytestreamProxyFinder.h>
-#include <boost/smart_ptr/make_shared.hpp>
 #include <boost/bind.hpp>
+#include <boost/smart_ptr/make_shared.hpp>
 #include <Swiften/Base/Log.h>
+#include <Swiften/Base/foreach.h>
 #include <Swiften/Elements/S5BProxyRequest.h>
 #include <Swiften/Queries/GenericRequest.h>
 #include <Swiften/Queries/IQRouter.h>
@@ -25,19 +32,27 @@ SOCKS5BytestreamProxyFinder::~SOCKS5BytestreamProxyFinder() {
 void SOCKS5BytestreamProxyFinder::start() {
        serviceWalker = boost::make_shared<DiscoServiceWalker>(service, iqRouter);
        serviceWalker->onServiceFound.connect(boost::bind(&SOCKS5BytestreamProxyFinder::handleServiceFound, this, _1, _2));
+        serviceWalker->onWalkComplete.connect(boost::bind(&SOCKS5BytestreamProxyFinder::handleWalkEnded, this));
        serviceWalker->beginWalk();
 }
 void SOCKS5BytestreamProxyFinder::stop() {
+        typedef boost::shared_ptr<GenericRequest<S5BProxyRequest> > S5BProxyRequestGenericRequest;
+        foreach (S5BProxyRequestGenericRequest requester, pendingRequests) {
+                requester->onResponse.disconnect(boost::bind(&SOCKS5BytestreamProxyFinder::handleProxyResponse, this, requester, _1, _2));
+        }
        serviceWalker->endWalk();
        serviceWalker->onServiceFound.disconnect(boost::bind(&SOCKS5BytestreamProxyFinder::handleServiceFound, this, _1, _2));
+        serviceWalker->onWalkComplete.disconnect(boost::bind(&SOCKS5BytestreamProxyFinder::handleWalkEnded, this));
        serviceWalker.reset();
 }
 void SOCKS5BytestreamProxyFinder::sendBytestreamQuery(const JID& jid) {
        S5BProxyRequest::ref proxyRequest = boost::make_shared<S5BProxyRequest>();
        boost::shared_ptr<GenericRequest<S5BProxyRequest> > request = boost::make_shared<GenericRequest<S5BProxyRequest> >(IQ::Get, jid, proxyRequest, iqRouter);
-        request->onResponse.connect(boost::bind(&SOCKS5BytestreamProxyFinder::handleProxyResponse, this, _1, _2));
+        request->onResponse.connect(boost::bind(&SOCKS5BytestreamProxyFinder::handleProxyResponse, this, request, _1, _2));
+        pendingRequests.insert(request);
        request->send();
 }
@@ -47,16 +62,26 @@ void SOCKS5BytestreamProxyFinder::handleServiceFound(const JID& jid, boost::shar
        }
 }
-void SOCKS5BytestreamProxyFinder::handleProxyResponse(boost::shared_ptr<S5BProxyRequest> request, ErrorPayload::ref error) {
+void SOCKS5BytestreamProxyFinder::handleWalkEnded() {
+        if (pendingRequests.empty()) {
+                onProxiesFound(proxyHosts);
+        }
+}
+void SOCKS5BytestreamProxyFinder::handleProxyResponse(boost::shared_ptr<GenericRequest<S5BProxyRequest> > requester, boost::shared_ptr<S5BProxyRequest> request, ErrorPayload::ref error) {
+        requester->onResponse.disconnect(boost::bind(&SOCKS5BytestreamProxyFinder::handleProxyResponse, this, requester, _1, _2));
+        pendingRequests.erase(requester);
        if (error) {
                SWIFT_LOG(debug) << "ERROR" << std::endl;
        } else {
                if (request) {
-                        onProxyFound(request);
+                        SWIFT_LOG(debug) << "add request" << std::endl;
-                } else {
+                        proxyHosts.push_back(request);
-                        //assert(false);
                }
        }
+        if (pendingRequests.empty() && !serviceWalker->isActive()) {
+                onProxiesFound(proxyHosts);
+        }
 }
 }