1 files changed, 48 insertions, 10 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,7 +1,7 @@
-// TODO: whitespacePingLayer_->setInactive();
+// TODO: Send out better errors
 
 #include "Swiften/Session/BasicSessionStream.h"
 
 #include <boost/bind.hpp>
 
 #include "Swiften/StreamStack/XMPPLayer.h"
@@ -10,57 +10,84 @@
 #include "Swiften/StreamStack/WhitespacePingLayer.h"
 #include "Swiften/StreamStack/TLSLayer.h"
 #include "Swiften/StreamStack/TLSLayerFactory.h"
 
 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() {
 	xmppLayer = boost::shared_ptr<XMPPLayer>(
 			new XMPPLayer(payloadParserFactories, payloadSerializers));
 	xmppLayer->onStreamStart.connect(boost::bind(&BasicSessionStream::handleStreamStartReceived, shared_from_this(), _1));
 	xmppLayer->onElement.connect(boost::bind(&BasicSessionStream::handleElementReceived, shared_from_this(), _1));
 	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() {
   delete streamStack;
 }
 
 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() {
   xmppLayer->resetParser();
 }
 
@@ -70,14 +97,25 @@ void BasicSessionStream::handleStreamStartReceived(const ProtocolHeader& header)
 
 void BasicSessionStream::handleElementReceived(boost::shared_ptr<Element> element) {
 	onElementReceived(element);
 }
 
 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()));
 }
 
 };