3 files changed, 20 insertions, 13 deletions
diff --git a/Swiften/Client/Client.cpp b/Swiften/Client/Client.cpp
index 9fb4ca0..9e38626 100644
--- a/Swiften/Client/Client.cpp
+++ b/Swiften/Client/Client.cpp
@@ -21,6 +21,9 @@ Client::Client(const JID& jid, const String& password) :
 }
 
 Client::~Client() {
+	if (session_ || connection_) {
+		std::cerr << "Warning: Client not disconnected properly" << std::endl;
+	}
 	delete tlsLayerFactory_;
 	delete connectionFactory_;
 }
@@ -52,24 +55,28 @@ void Client::handleConnectionConnectFinished(bool error) {
 		if (!certificate_.isEmpty()) {
 			sessionStream_->setTLSCertificate(PKCS12Certificate(certificate_, password_));
 		}
-		sessionStream_->onDataRead.connect(boost::bind(&Client::handleDataRead, shared_from_this(), _1));
-		sessionStream_->onDataWritten.connect(boost::bind(&Client::handleDataWritten, shared_from_this(), _1));
+		sessionStream_->onDataRead.connect(boost::bind(&Client::handleDataRead, this, _1));
+		sessionStream_->onDataWritten.connect(boost::bind(&Client::handleDataWritten, this, _1));
 		sessionStream_->initialize();
 
 		session_ = boost::shared_ptr<ClientSession>(new ClientSession(jid_, sessionStream_));
 		session_->onInitialized.connect(boost::bind(boost::ref(onConnected)));
-		session_->onFinished.connect(boost::bind(&Client::handleSessionFinished, shared_from_this(), _1));
-		session_->onNeedCredentials.connect(boost::bind(&Client::handleNeedCredentials, shared_from_this()));
-		session_->onElementReceived.connect(boost::bind(&Client::handleElement, shared_from_this(), _1));
+		session_->onFinished.connect(boost::bind(&Client::handleSessionFinished, this, _1));
+		session_->onNeedCredentials.connect(boost::bind(&Client::handleNeedCredentials, this));
+		session_->onElementReceived.connect(boost::bind(&Client::handleElement, this, _1));
 		session_->start();
 	}
 }
 
 void Client::disconnect() {
-	// TODO
-	//if (session_) {
-	//	session_->finishSession();
-	//}
+	if (session_) {
+		session_->finish();
+		session_.reset();
+	}
+	if (connection_) {
+		connection_->disconnect();
+		connection_.reset();
+	}
 }
 
 void Client::send(boost::shared_ptr<Stanza> stanza) {
diff --git a/Swiften/Client/Client.h b/Swiften/Client/Client.h
index 16127dd..5188789 100644
--- a/Swiften/Client/Client.h
+++ b/Swiften/Client/Client.h
@@ -3,7 +3,6 @@
 
 #include <boost/signals.hpp>
 #include <boost/shared_ptr.hpp>
-#include <boost/enable_shared_from_this.hpp>
 
 #include "Swiften/Base/Error.h"
 #include "Swiften/Client/ClientSession.h"
@@ -24,7 +23,7 @@ namespace Swift {
 	class ClientSession;
 	class BasicSessionStream;
 
-	class Client : public StanzaChannel, public IQRouter, public boost::enable_shared_from_this<Client> {
+	class Client : public StanzaChannel, public IQRouter, public boost::bsignals::trackable {
 		public:
 			Client(const JID& jid, const String& password);
 			~Client();
diff --git a/Swiften/Client/ClientSession.cpp b/Swiften/Client/ClientSession.cpp
index ed5d27d..a185ea0 100644
--- a/Swiften/Client/ClientSession.cpp
+++ b/Swiften/Client/ClientSession.cpp
@@ -25,13 +25,14 @@ ClientSession::ClientSession(
 			state(Initial), 
 			stream(stream),
 			needSessionStart(false) {
+}
+
+void ClientSession::start() {
 	stream->onStreamStartReceived.connect(boost::bind(&ClientSession::handleStreamStart, shared_from_this(), _1));
 	stream->onElementReceived.connect(boost::bind(&ClientSession::handleElement, shared_from_this(), _1));
 	stream->onError.connect(boost::bind(&ClientSession::handleStreamError, shared_from_this(), _1));
 	stream->onTLSEncrypted.connect(boost::bind(&ClientSession::handleTLSEncrypted, shared_from_this()));
-}
 
-void ClientSession::start() {
 	assert(state == Initial);
 	state = WaitingForStreamStart;
 	sendStreamHeader();