3 files changed, 34 insertions, 20 deletions
diff --git a/Swiften/Client/ClientSession.cpp b/Swiften/Client/ClientSession.cpp
index 9400b56..9950a76 100644
--- a/Swiften/Client/ClientSession.cpp
+++ b/Swiften/Client/ClientSession.cpp
@@ -62,7 +62,7 @@ ClientSession::~ClientSession() {
 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->onClosed.connect(boost::bind(&ClientSession::handleStreamFinished, shared_from_this(), _1));
+	stream->onClosed.connect(boost::bind(&ClientSession::handleStreamClosed, shared_from_this(), _1));
 	stream->onTLSEncrypted.connect(boost::bind(&ClientSession::handleTLSEncrypted, shared_from_this()));
 
 	assert(state == Initial);
@@ -367,20 +367,10 @@ void ClientSession::continueAfterTLSEncrypted() {
 	sendStreamHeader();
 }
 
-void ClientSession::handleStreamFinished(boost::shared_ptr<Swift::Error> error) {
-	finishSession(error);
-}
-
-void ClientSession::finish() {
-	finishSession(boost::shared_ptr<Error>());
-}
-
-void ClientSession::finishSession(Error::Type error) {
-	finishSession(boost::shared_ptr<Swift::ClientSession::Error>(new Swift::ClientSession::Error(error)));
-}
-
-void ClientSession::finishSession(boost::shared_ptr<Swift::Error> error) {
+void ClientSession::handleStreamClosed(boost::shared_ptr<Swift::Error> streamError) {
+	State previousState = state;
 	state = Finished;
+
 	if (stanzaAckRequester_) {
 		stanzaAckRequester_->onRequestAck.disconnect(boost::bind(&ClientSession::requestAck, shared_from_this()));
 		stanzaAckRequester_->onStanzaAcked.disconnect(boost::bind(&ClientSession::handleStanzaAcked, shared_from_this(), _1));
@@ -393,14 +383,32 @@ void ClientSession::finishSession(boost::shared_ptr<Swift::Error> error) {
 	stream->setWhitespacePingEnabled(false);
 	stream->onStreamStartReceived.disconnect(boost::bind(&ClientSession::handleStreamStart, shared_from_this(), _1));
 	stream->onElementReceived.disconnect(boost::bind(&ClientSession::handleElement, shared_from_this(), _1));
-	stream->onClosed.disconnect(boost::bind(&ClientSession::handleStreamFinished, shared_from_this(), _1));
+	stream->onClosed.disconnect(boost::bind(&ClientSession::handleStreamClosed, shared_from_this(), _1));
 	stream->onTLSEncrypted.disconnect(boost::bind(&ClientSession::handleTLSEncrypted, shared_from_this()));
-	if (stream->isAvailable()) {
-		stream->writeFooter();
+
+	if (previousState == Finishing) {
+		onFinished(error_);
 	}
-	onFinished(error);
+	else {
+		onFinished(streamError);
+	}
+}
+
+void ClientSession::finish() {
+	finishSession(boost::shared_ptr<Error>());
 }
 
+void ClientSession::finishSession(Error::Type error) {
+	finishSession(boost::shared_ptr<Swift::ClientSession::Error>(new Swift::ClientSession::Error(error)));
+}
+
+void ClientSession::finishSession(boost::shared_ptr<Swift::Error> error) {
+	state = Finishing;
+	error_ = error;
+	assert(stream->isOpen());
+	stream->writeFooter();
+	stream->close();
+}
 
 void ClientSession::requestAck() {
 	stream->writeElement(boost::shared_ptr<StanzaAckRequest>(new StanzaAckRequest()));
diff --git a/Swiften/Client/ClientSession.h b/Swiften/Client/ClientSession.h
index f35c298..be0f89e 100644
--- a/Swiften/Client/ClientSession.h
+++ b/Swiften/Client/ClientSession.h
@@ -37,6 +37,7 @@ namespace Swift {
 				BindingResource,
 				StartingSession,
 				Initialized,
+				Finishing,
 				Finished
 			};
 
@@ -115,7 +116,7 @@ namespace Swift {
 
 			void handleElement(boost::shared_ptr<Element>);
 			void handleStreamStart(const ProtocolHeader&);
-			void handleStreamFinished(boost::shared_ptr<Swift::Error>);
+			void handleStreamClosed(boost::shared_ptr<Swift::Error>);
 
 			void handleTLSEncrypted();
 
@@ -139,6 +140,7 @@ namespace Swift {
 			ClientAuthenticator* authenticator;
 			boost::shared_ptr<StanzaAckRequester> stanzaAckRequester_;
 			boost::shared_ptr<StanzaAckResponder> stanzaAckResponder_;
+			boost::shared_ptr<Swift::Error> error_;
 			CertificateTrustChecker* certificateTrustChecker;
 	};
 }
diff --git a/Swiften/Client/UnitTest/ClientSessionTest.cpp b/Swiften/Client/UnitTest/ClientSessionTest.cpp
index 2b0241a..5d0e2aa 100644
--- a/Swiften/Client/UnitTest/ClientSessionTest.cpp
+++ b/Swiften/Client/UnitTest/ClientSessionTest.cpp
@@ -296,7 +296,11 @@ class ClientSessionTest : public CppUnit::TestFixture {
 				MockSessionStream() : available(true), canTLSEncrypt(true), tlsEncrypted(false), compressed(false), whitespacePingEnabled(false), resetCount(0) {
 				}
 
-				virtual bool isAvailable() {
+				virtual void close() {
+					onClosed(boost::shared_ptr<Error>());
+				}
+
+				virtual bool isOpen() {
 					return available;
 				}