2 files changed, 20 insertions, 7 deletions
diff --git a/Swiften/Server/ServerFromClientSession.cpp b/Swiften/Server/ServerFromClientSession.cpp
index fe4388e..9a3cf83 100644
--- a/Swiften/Server/ServerFromClientSession.cpp
+++ b/Swiften/Server/ServerFromClientSession.cpp
@@ -32,6 +32,15 @@ ServerFromClientSession::ServerFromClientSession(
 			authenticated_(false),
 			initialized_(false) {
 	xmppLayer_ = boost::shared_ptr<XMPPLayer>(new XMPPLayer(payloadParserFactories_, payloadSerializers_));
+	connectionLayer_ = boost::shared_ptr<ConnectionLayer>(new ConnectionLayer(connection_));
+	streamStack_ = new StreamStack(xmppLayer_, connectionLayer_);
+}
+
+ServerFromClientSession::~ServerFromClientSession() {
+	delete streamStack_;
+}
+
+void ServerFromClientSession::start() {
 	xmppLayer_->onStreamStart.connect(
 			boost::bind(&ServerFromClientSession::handleStreamStart, this, _2));
 	xmppLayer_->onElement.connect(
@@ -42,12 +51,7 @@ ServerFromClientSession::ServerFromClientSession(
 			boost::bind(boost::ref(onDataRead), _1));
 	xmppLayer_->onWriteData.connect(
 			boost::bind(boost::ref(onDataWritten), _1));
-	connectionLayer_ = boost::shared_ptr<ConnectionLayer>(new ConnectionLayer(connection_));
-	streamStack_ = new StreamStack(xmppLayer_, connectionLayer_);
-}
-
-ServerFromClientSession::~ServerFromClientSession() {
-	delete streamStack_;
+	connection_->onDisconnected.connect(boost::bind(&ServerFromClientSession::handleDisconnected, shared_from_this(), _1));
 }
 
 void ServerFromClientSession::handleElement(boost::shared_ptr<Element> element) {
@@ -114,5 +118,9 @@ void ServerFromClientSession::sendStanza(boost::shared_ptr<Stanza> stanza) {
 	xmppLayer_->writeElement(stanza);
 }
 
+void ServerFromClientSession::handleDisconnected(const boost::optional<Connection::Error>&) {
+	onSessionFinished();
+}
+
 
 }
diff --git a/Swiften/Server/ServerFromClientSession.h b/Swiften/Server/ServerFromClientSession.h
index 0110d5d..e7df99d 100644
--- a/Swiften/Server/ServerFromClientSession.h
+++ b/Swiften/Server/ServerFromClientSession.h
@@ -2,9 +2,11 @@
 
 #include <boost/shared_ptr.hpp>
 #include <boost/signal.hpp>
+#include <boost/enable_shared_from_this.hpp>
 
 #include "Swiften/Base/String.h"
 #include "Swiften/JID/JID.h"
+#include "Swiften/Network/Connection.h"
 
 namespace Swift {
 	class Element;
@@ -18,7 +20,7 @@ namespace Swift {
 	class Connection;
 	class ByteArray;
 
-	class ServerFromClientSession {
+	class ServerFromClientSession : public boost::enable_shared_from_this<ServerFromClientSession> {
 		public:
 			ServerFromClientSession(
 					const String& id,
@@ -28,6 +30,8 @@ namespace Swift {
 					UserRegistry* userRegistry);
 			~ServerFromClientSession();
 
+			void start();
+
 			void sendStanza(boost::shared_ptr<Stanza>);
 
 			const JID& getJID() const {
@@ -45,6 +49,7 @@ namespace Swift {
 			boost::signal<void (const ByteArray&)> onDataRead;
 
 		private:
+			void handleDisconnected(const boost::optional<Connection::Error>& error);
 			void handleElement(boost::shared_ptr<Element>);
 			void handleStreamStart(const String& domain);