3 files changed, 78 insertions, 13 deletions
diff --git a/Swiften/Session/BasicSessionStream.cpp b/Swiften/Session/BasicSessionStream.cpp
index 73eaf5b..115dc7c 100644
--- a/Swiften/Session/BasicSessionStream.cpp
+++ b/Swiften/Session/BasicSessionStream.cpp
@@ -1,4 +1,4 @@
-// TODO: whitespacePingLayer_->setInactive();
+// TODO: Send out better errors
 
 #include "Swiften/Session/BasicSessionStream.h"
 
@@ -13,7 +13,7 @@
 
 namespace Swift {
 
-BasicSessionStream::BasicSessionStream(boost::shared_ptr<Connection> connection, PayloadParserFactoryCollection* payloadParserFactories, PayloadSerializerCollection* payloadSerializers, TLSLayerFactory* tlsLayerFactory) : connection(connection), payloadParserFactories(payloadParserFactories), payloadSerializers(payloadSerializers), tlsLayerFactory(tlsLayerFactory) {
+BasicSessionStream::BasicSessionStream(boost::shared_ptr<Connection> connection, PayloadParserFactoryCollection* payloadParserFactories, PayloadSerializerCollection* payloadSerializers, TLSLayerFactory* tlsLayerFactory) : available(false), connection(connection), payloadParserFactories(payloadParserFactories), payloadSerializers(payloadSerializers), tlsLayerFactory(tlsLayerFactory) {
 }
 
 void BasicSessionStream::initialize() {
@@ -24,10 +24,13 @@ void BasicSessionStream::initialize() {
 	xmppLayer->onError.connect(boost::bind(
       &BasicSessionStream::handleXMPPError, shared_from_this()));
 
+	connection->onDisconnected.connect(boost::bind(&BasicSessionStream::handleConnectionError, shared_from_this(), _1));
 	connectionLayer = boost::shared_ptr<ConnectionLayer>(
       new ConnectionLayer(connection));
 
 	streamStack = new StreamStack(xmppLayer, connectionLayer);
+
+	available = true;
 }
 
 BasicSessionStream::~BasicSessionStream() {
@@ -35,29 +38,53 @@ BasicSessionStream::~BasicSessionStream() {
 }
 
 void BasicSessionStream::writeHeader(const ProtocolHeader& header) {
+	assert(available);
 	xmppLayer->writeHeader(header);
 }
 
 void BasicSessionStream::writeElement(boost::shared_ptr<Element> element) {
+	assert(available);
 	xmppLayer->writeElement(element);
 }
 
+void BasicSessionStream::writeFooter() {
+	assert(available);
+	xmppLayer->writeFooter();
+}
+
+bool BasicSessionStream::isAvailable() {
+	return available;
+}
+
 bool BasicSessionStream::supportsTLSEncryption() {
   return tlsLayerFactory && tlsLayerFactory->canCreate();
 }
 
 void BasicSessionStream::addTLSEncryption() {
+	assert(available);
 	tlsLayer = tlsLayerFactory->createTLSLayer();
-  streamStack->addLayer(tlsLayer);
-  // TODO: Add tls layer certificate if needed
-  tlsLayer->onError.connect(boost::bind(&BasicSessionStream::handleTLSError, shared_from_this()));
-  tlsLayer->connect();
+	if (hasTLSCertificate() && !tlsLayer->setClientCertificate(getTLSCertificate())) {
+		onError(boost::shared_ptr<Error>(new Error()));
+	}
+	else {
+		streamStack->addLayer(tlsLayer);
+		tlsLayer->onError.connect(boost::bind(&BasicSessionStream::handleTLSError, shared_from_this()));
+		tlsLayer->onConnected.connect(boost::bind(&BasicSessionStream::handleTLSConnected, shared_from_this()));
+		tlsLayer->connect();
+	}
 }
 
-void BasicSessionStream::addWhitespacePing() {
-  whitespacePingLayer = boost::shared_ptr<WhitespacePingLayer>(new WhitespacePingLayer());
-  streamStack->addLayer(whitespacePingLayer);
-  whitespacePingLayer->setActive();
+void BasicSessionStream::setWhitespacePingEnabled(bool enabled) {
+	if (enabled && !whitespacePingLayer) {
+		whitespacePingLayer = boost::shared_ptr<WhitespacePingLayer>(new WhitespacePingLayer());
+		streamStack->addLayer(whitespacePingLayer);
+	}
+	if (enabled) {
+		whitespacePingLayer->setActive();
+	}
+	else {
+		whitespacePingLayer->setInactive();
+	}
 }
 
 void BasicSessionStream::resetXMPPParser() {
@@ -73,10 +100,21 @@ void BasicSessionStream::handleElementReceived(boost::shared_ptr<Element> elemen
 }
 
 void BasicSessionStream::handleXMPPError() {
+	available = false;
 	onError(boost::shared_ptr<Error>(new Error()));
 }
 
+void BasicSessionStream::handleTLSConnected() {
+	onTLSEncrypted();
+}
+
 void BasicSessionStream::handleTLSError() {
+	available = false;
+	onError(boost::shared_ptr<Error>(new Error()));
+}
+
+void BasicSessionStream::handleConnectionError(const boost::optional<Connection::Error>&) {
+	available = false;
 	onError(boost::shared_ptr<Error>(new Error()));
 }
 
diff --git a/Swiften/Session/BasicSessionStream.h b/Swiften/Session/BasicSessionStream.h
index d248ebc..5fe0b4c 100644
--- a/Swiften/Session/BasicSessionStream.h
+++ b/Swiften/Session/BasicSessionStream.h
@@ -30,23 +30,29 @@ namespace Swift {
 
 			void initialize();
 
+			virtual bool isAvailable();
+
 			virtual void writeHeader(const ProtocolHeader& header);
 			virtual void writeElement(boost::shared_ptr<Element>);
+			virtual void writeFooter();
 
 			virtual bool supportsTLSEncryption();
 			virtual void addTLSEncryption();
 
-			virtual void addWhitespacePing();
+			virtual void setWhitespacePingEnabled(bool);
 
 			virtual void resetXMPPParser();
 
     private:
+			void handleConnectionError(const boost::optional<Connection::Error>& error);
       void handleXMPPError();
+			void handleTLSConnected();
       void handleTLSError();
 			void handleStreamStartReceived(const ProtocolHeader&);
 			void handleElementReceived(boost::shared_ptr<Element>);
 
     private:
+			bool available;
 			boost::shared_ptr<Connection> connection;
 			PayloadParserFactoryCollection* payloadParserFactories;
 			PayloadSerializerCollection* payloadSerializers;
diff --git a/Swiften/Session/SessionStream.h b/Swiften/Session/SessionStream.h
index 44a1980..b2444f5 100644
--- a/Swiften/Session/SessionStream.h
+++ b/Swiften/Session/SessionStream.h
@@ -6,6 +6,7 @@
 #include "Swiften/Elements/ProtocolHeader.h"
 #include "Swiften/Elements/Element.h"
 #include "Swiften/Base/Error.h"
+#include "Swiften/TLS/PKCS12Certificate.h"
 
 namespace Swift {
 	class SessionStream {
@@ -17,18 +18,38 @@ namespace Swift {
 
 			virtual ~SessionStream();
 
+			virtual bool isAvailable() = 0;
+
 			virtual void writeHeader(const ProtocolHeader& header) = 0;
+			virtual void writeFooter() = 0;
 			virtual void writeElement(boost::shared_ptr<Element>) = 0;
 
 			virtual bool supportsTLSEncryption() = 0;
 			virtual void addTLSEncryption() = 0;
-
-			virtual void addWhitespacePing() = 0;
+			virtual void setWhitespacePingEnabled(bool enabled) = 0;
 
 			virtual void resetXMPPParser() = 0;
 
+			void setTLSCertificate(const PKCS12Certificate& cert) {
+				certificate = cert;
+			}
+
+			virtual bool hasTLSCertificate() {
+				return !certificate.isNull();
+			}
+
+
 			boost::signal<void (const ProtocolHeader&)> onStreamStartReceived;
 			boost::signal<void (boost::shared_ptr<Element>)> onElementReceived;
 			boost::signal<void (boost::shared_ptr<Error>)> onError;
+			boost::signal<void ()> onTLSEncrypted;
+
+		protected:
+			const PKCS12Certificate& getTLSCertificate() const {
+				return certificate;
+			}
+
+		private:
+			PKCS12Certificate certificate;
 	};
 }