1 files changed, 15 insertions, 13 deletions
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,7 +1,7 @@
 /*
- * Copyright (c) 2010-2016 Isode Limited.
+ * Copyright (c) 2010-2018 Isode Limited.
  * All rights reserved.
  * See the COPYING file for more information.
  */
 
 #include <Swiften/Session/Session.h>
@@ -20,20 +20,14 @@ Session::Session(
         XMLParserFactory* xmlParserFactory) :
             connection(connection),
             payloadParserFactories(payloadParserFactories),
             payloadSerializers(payloadSerializers),
             xmlParserFactory(xmlParserFactory),
-            xmppLayer(nullptr),
-            connectionLayer(nullptr),
-            streamStack(nullptr),
             finishing(false) {
 }
 
 Session::~Session() {
-    delete streamStack;
-    delete connectionLayer;
-    delete xmppLayer;
 }
 
 void Session::startSession() {
     initializeStreamStack();
     handleSessionStarted();
@@ -42,44 +36,52 @@ void Session::startSession() {
 void Session::finishSession() {
     if (finishing) {
         return;
     }
     finishing = true;
-    if (xmppLayer) {
+    if (auto xmppLayer = getXMPPLayer()) {
         xmppLayer->writeFooter();
     }
     connection->disconnect();
 }
 
 void Session::finishSession(const SessionError& /*error*/) {
     if (finishing) {
         return;
     }
     finishing = true;
-    if (xmppLayer) {
+    if (auto xmppLayer = getXMPPLayer()) {
         xmppLayer->writeFooter();
     }
     connection->disconnect();
 }
 
 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(
             boost::bind(&Session::handleStreamStart, this, _1));
     xmppLayer->onElement.connect(boost::bind(&Session::handleElement, this, _1));
     xmppLayer->onError.connect(
             boost::bind(&Session::finishSession, this, XMLError));
     xmppLayer->onDataRead.connect(boost::bind(boost::ref(onDataRead), _1));
     xmppLayer->onWriteData.connect(boost::bind(boost::ref(onDataWritten), _1));
     connection->onDisconnected.connect(
             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);
 }
 
 void Session::handleDisconnected(const boost::optional<Connection::Error>& connectionError) {
     connection->onDisconnected.disconnect(
             boost::bind(&Session::handleDisconnected, this, _1));