1 files changed, 20 insertions, 14 deletions
diff --git a/Swiften/Network/BoostConnection.cpp b/Swiften/Network/BoostConnection.cpp
index f7ff8c4..ac3b444 100644
--- a/Swiften/Network/BoostConnection.cpp
+++ b/Swiften/Network/BoostConnection.cpp
@@ -4,18 +4,23 @@
  * 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 <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 +31,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,7 +43,7 @@ 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_;
 };
 
@@ -75,18 +80,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));
 }
@@ -111,7 +116,8 @@ void BoostConnection::doRead() {
 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());
+		eventLoop->postEvent(boost::bind(boost::ref(onDataRead), createSafeByteArray(&readBuffer_[0], bytesTransferred)), shared_from_this());
+		std::fill(readBuffer_.begin(), readBuffer_.end(), 0);
 		doRead();
 	}
 	else if (/*error == boost::asio::error::eof ||*/ error == boost::asio::error::operation_aborted) {
@@ -135,7 +141,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 {