4 files changed, 29 insertions, 34 deletions
diff --git a/Swiften/Session/BasicSessionStream.cpp b/Swiften/Session/BasicSessionStream.cpp
index 3e65640..c44961d 100644
--- a/Swiften/Session/BasicSessionStream.cpp
+++ b/Swiften/Session/BasicSessionStream.cpp
@@ -37,3 +37,3 @@ BasicSessionStream::BasicSessionStream(
             tlsOptions_(tlsOptions) {
-    xmppLayer = new XMPPLayer(payloadParserFactories, payloadSerializers, xmlParserFactory, streamType);
+    auto xmppLayer = std::make_unique<XMPPLayer>(payloadParserFactories, payloadSerializers, xmlParserFactory, streamType);
     xmppLayer->onStreamStart.connect(boost::bind(&BasicSessionStream::handleStreamStartReceived, this, _1));
@@ -46,6 +46,4 @@ BasicSessionStream::BasicSessionStream(
     connection->onDisconnected.connect(boost::bind(&BasicSessionStream::handleConnectionFinished, this, _1));
-    connectionLayer = new ConnectionLayer(connection);
-
-    streamStack = new StreamStack(xmppLayer, connectionLayer);
 
+    streamStack = std::make_unique<StreamStack>(std::move(xmppLayer), std::unique_ptr<ConnectionLayer>(new ConnectionLayer(connection)));
     available = true;
@@ -59,7 +57,6 @@ BasicSessionStream::~BasicSessionStream() {
     }
-    delete streamStack;
 
     connection->onDisconnected.disconnect(boost::bind(&BasicSessionStream::handleConnectionFinished, this, _1));
-    delete connectionLayer;
 
+    auto xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->onStreamStart.disconnect(boost::bind(&BasicSessionStream::handleStreamStartReceived, this, _1));
@@ -70,3 +67,2 @@ BasicSessionStream::~BasicSessionStream() {
     xmppLayer->onWriteData.disconnect(boost::bind(&BasicSessionStream::handleDataWritten, this, _1));
-    delete xmppLayer;
 }
@@ -75,2 +71,3 @@ void BasicSessionStream::writeHeader(const ProtocolHeader& header) {
     assert(available);
+    auto* xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->writeHeader(header);
@@ -80,2 +77,3 @@ void BasicSessionStream::writeElement(std::shared_ptr<ToplevelElement> element)
     assert(available);
+    auto* xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->writeElement(element);
@@ -85,2 +83,3 @@ void BasicSessionStream::writeFooter() {
     assert(available);
+    auto* xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->writeFooter();
@@ -90,2 +89,3 @@ void BasicSessionStream::writeData(const std::string& data) {
     assert(available);
+    auto* xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->writeData(data);
@@ -164,2 +164,3 @@ void BasicSessionStream::setWhitespacePingEnabled(bool enabled) {
 void BasicSessionStream::resetXMPPParser() {
+    auto* xmppLayer = streamStack->getLayer<XMPPLayer>();
     xmppLayer->resetParser();
diff --git a/Swiften/Session/BasicSessionStream.h b/Swiften/Session/BasicSessionStream.h
index 472b5cc..30a7e3b 100644
--- a/Swiften/Session/BasicSessionStream.h
+++ b/Swiften/Session/BasicSessionStream.h
@@ -85,5 +85,3 @@ namespace Swift {
             TimerFactory* timerFactory;
-            XMPPLayer* xmppLayer;
-            ConnectionLayer* connectionLayer;
-            StreamStack* streamStack;
+            std::unique_ptr<StreamStack> streamStack;
             TLSOptions tlsOptions_;
diff --git a/Swiften/Session/Session.cpp b/Swiften/Session/Session.cpp
index ebdb5d1..b1525b8 100644
--- a/Swiften/Session/Session.cpp
+++ b/Swiften/Session/Session.cpp
@@ -1,3 +1,3 @@
 /*
- * Copyright (c) 2010-2016 Isode Limited.
+ * Copyright (c) 2010-2018 Isode Limited.
  * All rights reserved.
@@ -24,5 +24,2 @@ Session::Session(
             xmlParserFactory(xmlParserFactory),
-            xmppLayer(nullptr),
-            connectionLayer(nullptr),
-            streamStack(nullptr),
             finishing(false) {
@@ -31,5 +28,2 @@ Session::Session(
 Session::~Session() {
-    delete streamStack;
-    delete connectionLayer;
-    delete xmppLayer;
 }
@@ -46,3 +40,3 @@ void Session::finishSession() {
     finishing = true;
-    if (xmppLayer) {
+    if (auto xmppLayer = getXMPPLayer()) {
         xmppLayer->writeFooter();
@@ -57,3 +51,3 @@ void Session::finishSession(const SessionError& /*error*/) {
     finishing = true;
-    if (xmppLayer) {
+    if (auto xmppLayer = getXMPPLayer()) {
         xmppLayer->writeFooter();
@@ -64,3 +58,3 @@ void Session::finishSession(const SessionError& /*error*/) {
 void Session::initializeStreamStack() {
-    xmppLayer = new XMPPLayer(payloadParserFactories, payloadSerializers, xmlParserFactory, ClientStreamType);
+    auto xmppLayer = std::unique_ptr<XMPPLayer>(new XMPPLayer(payloadParserFactories, payloadSerializers, xmlParserFactory, ClientStreamType));
     xmppLayer->onStreamStart.connect(
@@ -74,8 +68,16 @@ void Session::initializeStreamStack() {
             boost::bind(&Session::handleDisconnected, this, _1));
-    connectionLayer = new ConnectionLayer(connection);
-    streamStack = new StreamStack(xmppLayer, connectionLayer);
+    streamStack = std::unique_ptr<StreamStack>(new StreamStack(std::move(xmppLayer), std::unique_ptr<ConnectionLayer>(new ConnectionLayer(connection))));
 }
 
+XMPPLayer* Session::getXMPPLayer() const {
+    return dynamic_cast<XMPPLayer*>(streamStack->getTopLayer());
+}
+
+StreamStack* Session::getStreamStack() const {
+    return streamStack.get();
+}
+
+
 void Session::sendElement(std::shared_ptr<ToplevelElement> stanza) {
-    xmppLayer->writeElement(stanza);
+    getXMPPLayer()->writeElement(stanza);
 }
diff --git a/Swiften/Session/Session.h b/Swiften/Session/Session.h
index 04153ec..e6a0d53 100644
--- a/Swiften/Session/Session.h
+++ b/Swiften/Session/Session.h
@@ -1,3 +1,3 @@
 /*
- * Copyright (c) 2010-2017 Isode Limited.
+ * Copyright (c) 2010-2018 Isode Limited.
  * All rights reserved.
@@ -87,9 +87,4 @@ namespace Swift {
 
-            XMPPLayer* getXMPPLayer() const {
-                return xmppLayer;
-            }
-
-            StreamStack* getStreamStack() const {
-                return streamStack;
-            }
+            XMPPLayer* getXMPPLayer() const;
+            StreamStack* getStreamStack() const;
 
@@ -107,5 +102,4 @@ namespace Swift {
             XMLParserFactory* xmlParserFactory;
-            XMPPLayer* xmppLayer;
-            ConnectionLayer* connectionLayer;
-            StreamStack* streamStack;
+
+            std::unique_ptr<StreamStack> streamStack;
             bool finishing;