1 files changed, 27 insertions, 23 deletions
diff --git a/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp b/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
index 9951f7a..477af4b 100644
--- a/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
+++ b/Swiften/FileTransfer/SOCKS5BytestreamServerSession.cpp
@@ -4,13 +4,17 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swiften/FileTransfer/SOCKS5BytestreamServerSession.h"
+#include <Swiften/FileTransfer/SOCKS5BytestreamServerSession.h>
 
 #include <boost/bind.hpp>
+#include <iostream>
 
-#include "Swiften/Base/ByteArray.h"
-#include "Swiften/FileTransfer/SOCKS5BytestreamRegistry.h"
-#include "Swiften/FileTransfer/BytestreamException.h"
+#include <Swiften/Base/ByteArray.h>
+#include <Swiften/Base/SafeByteArray.h>
+#include <Swiften/Base/Algorithm.h>
+#include <Swiften/Base/Concat.h>
+#include <Swiften/FileTransfer/SOCKS5BytestreamRegistry.h>
+#include <Swiften/FileTransfer/BytestreamException.h>
 
 namespace Swift {
 
@@ -33,52 +37,52 @@ void SOCKS5BytestreamServerSession::stop() {
 	finish(false);
 }
 
-void SOCKS5BytestreamServerSession::handleDataRead(const ByteArray& data) {
-	unprocessedData += data;
+void SOCKS5BytestreamServerSession::handleDataRead(const SafeByteArray& data) {
+	append(unprocessedData, data);
 	process();
 }
 
 void SOCKS5BytestreamServerSession::process() {
 	if (state == WaitingForAuthentication) {
-		if (unprocessedData.getSize() >= 2) {
+		if (unprocessedData.size() >= 2) {
 			size_t authCount = unprocessedData[1];
 			size_t i = 2;
-			while (i < 2 + authCount && i < unprocessedData.getSize()) {
+			while (i < 2 + authCount && i < unprocessedData.size()) {
 				// Skip authentication mechanism
 				++i;
 			}
 			if (i == 2 + authCount) {
 				// Authentication message is complete
-				if (i != unprocessedData.getSize()) {
+				if (i != unprocessedData.size()) {
 					std::cerr << "SOCKS5BytestreamServerSession: Junk after authentication mechanism";
 				}
 				unprocessedData.clear();
-				connection->write(ByteArray("\x05\x00", 2));
+				connection->write(createSafeByteArray("\x05\x00", 2));
 				state = WaitingForRequest;
 			}
 		}
 	}
 	else if (state == WaitingForRequest) {
-		if (unprocessedData.getSize() >= 5) {
+		if (unprocessedData.size() >= 5) {
 			ByteArray requestID;
 			size_t i = 5;
 			size_t hostnameSize = unprocessedData[4];
-			while (i < 5 + hostnameSize && i < unprocessedData.getSize()) {
-				requestID += unprocessedData[i];
+			while (i < 5 + hostnameSize && i < unprocessedData.size()) {
+				requestID.push_back(unprocessedData[i]);
 				++i;
 			}
 			// Skip the port:
 			i += 2;
-			if (i >= unprocessedData.getSize()) {
-				if (i != unprocessedData.getSize()) {
+			if (i >= unprocessedData.size()) {
+				if (i != unprocessedData.size()) {
 					std::cerr << "SOCKS5BytestreamServerSession: Junk after authentication mechanism";
 				}
-				bytestream = bytestreams->getBytestream(requestID.toString());
-				ByteArray result("\x05", 1);
-				result += bytestream ? 0x0 : 0x4;
-				result += ByteArray("\x00\x03", 2);
-				result += static_cast<char>(requestID.getSize());
-				result += requestID + ByteArray("\x00\x00", 2);
+				bytestream = bytestreams->getBytestream(byteArrayToString(requestID));
+				SafeByteArray result = createSafeByteArray("\x05", 1);
+				result.push_back(bytestream ? 0x0 : 0x4);
+				append(result, createByteArray("\x00\x03", 2));
+				result.push_back(static_cast<char>(requestID.size()));
+				append(result, concat(requestID, createByteArray("\x00\x00", 2)));
 				if (!bytestream) {
 					connection->write(result);
 					finish(true);
@@ -96,9 +100,9 @@ void SOCKS5BytestreamServerSession::process() {
 void SOCKS5BytestreamServerSession::sendData() {
 	if (!bytestream->isFinished()) {
 		try {
-			connection->write(bytestream->read(chunkSize));
+			connection->write(createSafeByteArray(bytestream->read(chunkSize)));
 		}
-		catch (const BytestreamException& e) {
+		catch (const BytestreamException&) {
 			finish(true);
 		}
 	}