1 files changed, 60 insertions, 35 deletions
diff --git a/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp b/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
index 2260fc20..cb34c58 100644
--- a/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
+++ b/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
@@ -1,4 +1,4 @@
 /*
- * Copyright (c) 2010 Remko Tronçon
+ * Copyright (c) 2010-2013 Remko Tronçon
  * Licensed under the GNU General Public License v3.
  * See Documentation/Licenses/GPLv3.txt for more information.
@@ -8,4 +8,5 @@
 
 #include <boost/bind.hpp>
+#include <boost/numeric/conversion/cast.hpp>
 #include <iostream>
 
@@ -15,4 +16,5 @@
 #include <Swiften/Base/Concat.h>
 #include <Swiften/Base/Log.h>
+#include <Swiften/Network/HostAddressPort.h>
 #include <Swiften/FileTransfer/SOCKS5BytestreamRegistry.h>
 #include <Swiften/FileTransfer/BytestreamException.h>
@@ -20,11 +22,18 @@
 namespace Swift {
 
-SOCKS5BytestreamServerSession::SOCKS5BytestreamServerSession(boost::shared_ptr<Connection> connection, SOCKS5BytestreamRegistry* bytestreams) : connection(connection), bytestreams(bytestreams), state(Initial), chunkSize(131072) {
-	connection->onDisconnected.connect(boost::bind(&SOCKS5BytestreamServerSession::handleDisconnected, this, _1));
+SOCKS5BytestreamServerSession::SOCKS5BytestreamServerSession(
+		boost::shared_ptr<Connection> connection, 
+		SOCKS5BytestreamRegistry* bytestreams) : 
+			connection(connection), 
+			bytestreams(bytestreams), 
+			state(Initial), 
+			chunkSize(131072), 
+			waitingForData(false) {
+	disconnectedConnection = connection->onDisconnected.connect(boost::bind(&SOCKS5BytestreamServerSession::handleDisconnected, this, _1));
 }
 
 SOCKS5BytestreamServerSession::~SOCKS5BytestreamServerSession() {
 	if (state != Finished && state != Initial) {
-		std::cerr << "Warning: SOCKS5BytestremServerSession unfinished" << std::endl;
+		std::cerr << "Warning: SOCKS5BytestreamServerSession unfinished" << std::endl;
 		finish(false);
 	}
@@ -33,32 +42,34 @@ SOCKS5BytestreamServerSession::~SOCKS5BytestreamServerSession() {
 void SOCKS5BytestreamServerSession::start() {
 	SWIFT_LOG(debug) << std::endl;
-	connection->onDataRead.connect(boost::bind(&SOCKS5BytestreamServerSession::handleDataRead, this, _1));
+	dataReadConnection = connection->onDataRead.connect(
+			boost::bind(&SOCKS5BytestreamServerSession::handleDataRead, this, _1));
 	state = WaitingForAuthentication;
 }
 
 void SOCKS5BytestreamServerSession::stop() {
-	connection->onDataWritten.disconnect(boost::bind(&SOCKS5BytestreamServerSession::sendData, this));
-	connection->onDataRead.disconnect(boost::bind(&SOCKS5BytestreamServerSession::handleDataRead, this, _1));
-	connection->disconnect();
-	state = Finished;
+	finish(false);
 }
 
-void SOCKS5BytestreamServerSession::startTransfer() {
-	if (state == ReadyForTransfer) {
-		if (readBytestream) {
+void SOCKS5BytestreamServerSession::startSending(boost::shared_ptr<ReadBytestream> stream) {
+	if (state != ReadyForTransfer) { SWIFT_LOG(debug) << "Not ready for transfer!" << std::endl; return; }
+
+	readBytestream = stream;
 	state = WritingData;
-			connection->onDataWritten.connect(boost::bind(&SOCKS5BytestreamServerSession::sendData, this));
+	dataAvailableConnection = readBytestream->onDataAvailable.connect(
+			boost::bind(&SOCKS5BytestreamServerSession::handleDataAvailable, this));
+	dataWrittenConnection = connection->onDataWritten.connect(
+			boost::bind(&SOCKS5BytestreamServerSession::sendData, this));
 	sendData();
 }
-		else if(writeBytestream) {
+
+void SOCKS5BytestreamServerSession::startReceiving(boost::shared_ptr<WriteBytestream> stream) {
+	if (state != ReadyForTransfer) { SWIFT_LOG(debug) << "Not ready for transfer!" << std::endl; return; }
+
+	writeBytestream = stream;
 	state = ReadingData;
 	writeBytestream->write(unprocessedData);
-			onBytesReceived(unprocessedData.size());
+	// onBytesReceived(unprocessedData.size());
 	unprocessedData.clear();
 }
-	} else {
-		SWIFT_LOG(debug) << "Not ready for transfer!" << std::endl;
-	}
-}
 
 HostAddressPort SOCKS5BytestreamServerSession::getAddressPort() const {
@@ -72,5 +83,11 @@ void SOCKS5BytestreamServerSession::handleDataRead(boost::shared_ptr<SafeByteArr
 	} else {
 		writeBytestream->write(createByteArray(vecptr(*data), data->size()));
-		onBytesReceived(data->size());
+		// onBytesReceived(data->size());
+	}
+}
+
+void SOCKS5BytestreamServerSession::handleDataAvailable() {
+	if (waitingForData) {
+		sendData();
 	}
 }
@@ -78,7 +95,5 @@ void SOCKS5BytestreamServerSession::handleDataRead(boost::shared_ptr<SafeByteArr
 void SOCKS5BytestreamServerSession::handleDisconnected(const boost::optional<Connection::Error>& error) {
 	SWIFT_LOG(debug) << (error ? (error == Connection::ReadError ? "Read Error" : "Write Error") : "No Error") << std::endl;
-	if (error) {
-		finish(true);
-	}
+	finish(error ? true : false);
 }
 
@@ -119,13 +134,12 @@ void SOCKS5BytestreamServerSession::process() {
 				}
 				unprocessedData.clear();
-				std::string streamID = byteArrayToString(requestID);
-				readBytestream = bytestreams->getReadBytestream(streamID);
-				writeBytestream = bytestreams->getWriteBytestream(streamID);
+				streamID = byteArrayToString(requestID);
+				bool hasBytestream = bytestreams->hasBytestream(streamID);
 				SafeByteArray result = createSafeByteArray("\x05", 1);
-				result.push_back((readBytestream || writeBytestream) ? 0x0 : 0x4);
+				result.push_back(hasBytestream ? 0x0 : 0x4);
 				append(result, createByteArray("\x00\x03", 2));
-				result.push_back(static_cast<char>(requestID.size()));
+				result.push_back(boost::numeric_cast<unsigned char>(requestID.size()));
 				append(result, concat(requestID, createByteArray("\x00\x00", 2)));
-				if (!readBytestream && !writeBytestream) {
+				if (!hasBytestream) {
 					SWIFT_LOG(debug) << "Readstream or Wrtiestream with ID " << streamID << " not found!" << std::endl;
 					connection->write(result);
@@ -133,7 +147,6 @@ void SOCKS5BytestreamServerSession::process() {
 				}
 				else {
-					SWIFT_LOG(debug) << "Found " << (readBytestream ? "Readstream" : "Writestream") << ". Sent OK." << std::endl;
+					SWIFT_LOG(debug) << "Found stream. Sent OK." << std::endl;
 					connection->write(result);
-					bytestreams->serverSessions[streamID] = this;
 					state = ReadyForTransfer;
 				}
@@ -146,7 +159,13 @@ void SOCKS5BytestreamServerSession::sendData() {
 	if (!readBytestream->isFinished()) {
 		try {
-			SafeByteArray dataToSend = createSafeByteArray(*readBytestream->read(chunkSize));
+			SafeByteArray dataToSend = createSafeByteArray(*readBytestream->read(boost::numeric_cast<size_t>(chunkSize)));
+			if (!dataToSend.empty()) {
 				connection->write(dataToSend);
 				onBytesSent(dataToSend.size());
+				waitingForData = false;
+			}
+			else {
+				waitingForData = true;
+			}
 		}
 		catch (const BytestreamException&) {
@@ -160,7 +179,13 @@ void SOCKS5BytestreamServerSession::sendData() {
 
 void SOCKS5BytestreamServerSession::finish(bool error) {
-	connection->onDataWritten.disconnect(boost::bind(&SOCKS5BytestreamServerSession::sendData, this));
-	connection->onDataRead.disconnect(boost::bind(&SOCKS5BytestreamServerSession::handleDataRead, this, _1));
-	connection->onDisconnected.disconnect(boost::bind(&SOCKS5BytestreamServerSession::handleDisconnected, this, _1));
+	SWIFT_LOG(debug) << error << " " << state << std::endl;
+	if (state == Finished) {
+		return;
+	}
+
+	disconnectedConnection.disconnect();
+	dataReadConnection.disconnect();
+	dataWrittenConnection.disconnect();
+	dataAvailableConnection.disconnect();
 	readBytestream.reset();
 	state = Finished;