1 files changed, 24 insertions, 16 deletions
diff --git a/Swiften/Network/BoostConnection.cpp b/Swiften/Network/BoostConnection.cpp
index f7ff8c4..043743e 100644
--- a/Swiften/Network/BoostConnection.cpp
+++ b/Swiften/Network/BoostConnection.cpp
@@ -4,18 +4,24 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swiften/Network/BoostConnection.h"
+#include <Swiften/Network/BoostConnection.h>
 
 #include <iostream>
+#include <string>
+#include <algorithm>
 #include <boost/bind.hpp>
 #include <boost/thread.hpp>
+#include <boost/asio/placeholders.hpp>
+#include <boost/asio/write.hpp>
+#include <boost/smart_ptr/make_shared.hpp>
 
 #include <Swiften/Base/Log.h>
-#include "Swiften/EventLoop/EventLoop.h"
-#include <string>
-#include "Swiften/Base/ByteArray.h"
-#include "Swiften/Network/HostAddressPort.h"
-#include "Swiften/Base/sleep.h"
+#include <Swiften/Base/Algorithm.h>
+#include <Swiften/EventLoop/EventLoop.h>
+#include <Swiften/Base/ByteArray.h>
+#include <Swiften/Network/HostAddressPort.h>
+#include <Swiften/Base/sleep.h>
+#include <Swiften/Base/SafeAllocator.h>
 
 namespace Swift {
 
@@ -26,8 +32,8 @@ static const size_t BUFFER_SIZE = 4096;
 // A reference-counted non-modifiable buffer class.
 class SharedBuffer {
 	public:
-		SharedBuffer(const ByteArray& data) : 
-				data_(new std::vector<char>(data.begin(), data.end())),
+		SharedBuffer(const SafeByteArray& data) :
+				data_(new std::vector<char, SafeAllocator<char> >(data.begin(), data.end())),
 				buffer_(boost::asio::buffer(*data_)) {
 		}
 
@@ -38,14 +44,14 @@ class SharedBuffer {
 		const boost::asio::const_buffer* end() const { return &buffer_ + 1; }
 
 	private:
-		boost::shared_ptr< std::vector<char> > data_;
+		boost::shared_ptr< std::vector<char, SafeAllocator<char> > > data_;
 		boost::asio::const_buffer buffer_;
 };
 
 // -----------------------------------------------------------------------------
 
 BoostConnection::BoostConnection(boost::shared_ptr<boost::asio::io_service> ioService, EventLoop* eventLoop) :
-		eventLoop(eventLoop), ioService(ioService), socket_(*ioService), readBuffer_(BUFFER_SIZE), writing_(false) {
+		eventLoop(eventLoop), ioService(ioService), socket_(*ioService), writing_(false) {
 }
 
 BoostConnection::~BoostConnection() {
@@ -75,18 +81,18 @@ void BoostConnection::disconnect() {
 	socket_.close();
 }
 
-void BoostConnection::write(const ByteArray& data) {
+void BoostConnection::write(const SafeByteArray& data) {
 	boost::lock_guard<boost::mutex> lock(writeMutex_);
 	if (!writing_) {
 		writing_ = true;
 		doWrite(data);
 	}
 	else {
-		writeQueue_ += data;
+		append(writeQueue_, data);
 	}
 }
 
-void BoostConnection::doWrite(const ByteArray& data) {
+void BoostConnection::doWrite(const SafeByteArray& data) {
 	boost::asio::async_write(socket_, SharedBuffer(data),
 			boost::bind(&BoostConnection::handleDataWritten, shared_from_this(), boost::asio::placeholders::error));
 }
@@ -103,15 +109,17 @@ void BoostConnection::handleConnectFinished(const boost::system::error_code& err
 }
 
 void BoostConnection::doRead() {
+	readBuffer_ = boost::make_shared<SafeByteArray>(BUFFER_SIZE);
 	socket_.async_read_some(
-			boost::asio::buffer(readBuffer_),
+			boost::asio::buffer(*readBuffer_),
 			boost::bind(&BoostConnection::handleSocketRead, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
 }
 
 void BoostConnection::handleSocketRead(const boost::system::error_code& error, size_t bytesTransferred) {
 	SWIFT_LOG(debug) << "Socket read " << error << std::endl;
 	if (!error) {
-		eventLoop->postEvent(boost::bind(boost::ref(onDataRead), ByteArray(&readBuffer_[0], bytesTransferred)), shared_from_this());
+		readBuffer_->resize(bytesTransferred);
+		eventLoop->postEvent(boost::bind(boost::ref(onDataRead), readBuffer_), shared_from_this());
 		doRead();
 	}
 	else if (/*error == boost::asio::error::eof ||*/ error == boost::asio::error::operation_aborted) {
@@ -135,7 +143,7 @@ void BoostConnection::handleDataWritten(const boost::system::error_code& error)
 	}
 	{
 		boost::lock_guard<boost::mutex> lock(writeMutex_);
-		if (writeQueue_.isEmpty()) {
+		if (writeQueue_.empty()) {
 			writing_ = false;
 		}
 		else {