3 files changed, 28 insertions, 15 deletions
diff --git a/Swiften/Component/ComponentSession.cpp b/Swiften/Component/ComponentSession.cpp
index 967e68d..c45f663 100644
--- a/Swiften/Component/ComponentSession.cpp
+++ b/Swiften/Component/ComponentSession.cpp
@@ -24,7 +24,7 @@ ComponentSession::~ComponentSession() {
 void ComponentSession::start() {
 	stream->onStreamStartReceived.connect(boost::bind(&ComponentSession::handleStreamStart, shared_from_this(), _1));
 	stream->onElementReceived.connect(boost::bind(&ComponentSession::handleElement, shared_from_this(), _1));
-	stream->onClosed.connect(boost::bind(&ComponentSession::handleStreamError, shared_from_this(), _1));
+	stream->onClosed.connect(boost::bind(&ComponentSession::handleStreamClosed, shared_from_this(), _1));
 
 	assert(state == Initial);
 	state = WaitingForStreamStart;
@@ -81,8 +81,19 @@ bool ComponentSession::checkState(State state) {
 	return true;
 }
 
-void ComponentSession::handleStreamError(boost::shared_ptr<Swift::Error> error) {
-	finishSession(error);
+void ComponentSession::handleStreamClosed(boost::shared_ptr<Swift::Error> streamError) {
+	State oldState = state;
+	state = Finished;
+	stream->setWhitespacePingEnabled(false);
+	stream->onStreamStartReceived.disconnect(boost::bind(&ComponentSession::handleStreamStart, shared_from_this(), _1));
+	stream->onElementReceived.disconnect(boost::bind(&ComponentSession::handleElement, shared_from_this(), _1));
+	stream->onClosed.disconnect(boost::bind(&ComponentSession::handleStreamClosed, shared_from_this(), _1));
+	if (oldState == Finishing) {
+		onFinished(error);
+	}
+	else {
+		onFinished(streamError);
+	}
 }
 
 void ComponentSession::finish() {
@@ -93,16 +104,12 @@ void ComponentSession::finishSession(Error::Type error) {
 	finishSession(boost::shared_ptr<Swift::ComponentSession::Error>(new Swift::ComponentSession::Error(error)));
 }
 
-void ComponentSession::finishSession(boost::shared_ptr<Swift::Error> error) {
-	state = Finished;
-	stream->setWhitespacePingEnabled(false);
-	stream->onStreamStartReceived.disconnect(boost::bind(&ComponentSession::handleStreamStart, shared_from_this(), _1));
-	stream->onElementReceived.disconnect(boost::bind(&ComponentSession::handleElement, shared_from_this(), _1));
-	stream->onClosed.disconnect(boost::bind(&ComponentSession::handleStreamError, shared_from_this(), _1));
-	if (stream->isAvailable()) {
-		stream->writeFooter();
-	}
-	onFinished(error);
+void ComponentSession::finishSession(boost::shared_ptr<Swift::Error> finishError) {
+	state = Finishing;
+	error = finishError;
+	assert(stream->isOpen());
+	stream->writeFooter();
+	stream->close();
 }
 
 }
diff --git a/Swiften/Component/ComponentSession.h b/Swiften/Component/ComponentSession.h
index cbfa227..dbe6e27 100644
--- a/Swiften/Component/ComponentSession.h
+++ b/Swiften/Component/ComponentSession.h
@@ -27,6 +27,7 @@ namespace Swift {
 				WaitingForStreamStart,
 				Authenticating,
 				Initialized,
+				Finishing,
 				Finished
 			};
 
@@ -68,7 +69,7 @@ namespace Swift {
 
 			void handleElement(boost::shared_ptr<Element>);
 			void handleStreamStart(const ProtocolHeader&);
-			void handleStreamError(boost::shared_ptr<Swift::Error>);
+			void handleStreamClosed(boost::shared_ptr<Swift::Error>);
 
 			bool checkState(State);
 
@@ -76,6 +77,7 @@ namespace Swift {
 			JID jid;
 			String secret;
 			boost::shared_ptr<SessionStream> stream;
+			boost::shared_ptr<Swift::Error> error;
 			State state;
 	};
 }
diff --git a/Swiften/Component/UnitTest/ComponentSessionTest.cpp b/Swiften/Component/UnitTest/ComponentSessionTest.cpp
index 3ad52f9..86776e8 100644
--- a/Swiften/Component/UnitTest/ComponentSessionTest.cpp
+++ b/Swiften/Component/UnitTest/ComponentSessionTest.cpp
@@ -95,7 +95,11 @@ class ComponentSessionTest : public CppUnit::TestFixture {
 				MockSessionStream() : available(true), whitespacePingEnabled(false), resetCount(0) {
 				}
 
-				virtual bool isAvailable() {
+				virtual void close() {
+					onClosed(boost::shared_ptr<Error>());
+				}
+
+				virtual bool isOpen() {
 					return available;
 				}