1 files changed, 18 insertions, 12 deletions
diff --git a/Swiften/StreamStack/StreamStack.h b/Swiften/StreamStack/StreamStack.h
index b12a69f..263b1f5 100644
--- a/Swiften/StreamStack/StreamStack.h
+++ b/Swiften/StreamStack/StreamStack.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2010-2016 Isode Limited.
+ * Copyright (c) 2010-2018 Isode Limited.
  * All rights reserved.
  * See the COPYING file for more information.
  */
@@ -15,34 +15,40 @@
 #include <Swiften/Elements/Stanza.h>
 
 namespace Swift {
-    class XMPPLayer;
+    class HighLayer;
     class LowLayer;
     class StreamLayer;
 
     class SWIFTEN_API StreamStack {
         public:
-            StreamStack(XMPPLayer* xmppLayer, LowLayer* physicalLayer);
+            StreamStack(std::unique_ptr<HighLayer> topLayer, std::unique_ptr<LowLayer> bottomLayer);
             ~StreamStack();
 
-            void addLayer(StreamLayer*);
+            void addLayer(std::unique_ptr<StreamLayer> /* streamLayer */);
 
-            XMPPLayer* getXMPPLayer() const {
-                return xmppLayer_;
+            HighLayer* getTopLayer() const {
+                return topLayer_.get();
             }
 
-            template<typename T> T* getLayer() {
-                for (auto& i : layers_) {
-                    T* layer = dynamic_cast<T*>(i);
+            template<typename T> T* getLayer() const {
+                for (const auto& i : layers_) {
+                    T* layer = dynamic_cast<T*>(i.get());
                     if (layer) {
                         return layer;
                     }
                 }
+                if (T* layer = dynamic_cast<T*>(topLayer_.get())) {
+                    return layer;
+                }
+                if (T* layer = dynamic_cast<T*>(bottomLayer_.get())) {
+                    return layer;
+                }
                 return nullptr;
             }
 
         private:
-            XMPPLayer* xmppLayer_;
-            LowLayer* physicalLayer_;
-            std::vector<StreamLayer*> layers_;
+            std::unique_ptr<HighLayer> topLayer_;
+            std::unique_ptr<LowLayer> bottomLayer_;
+            std::vector<std::unique_ptr<StreamLayer>> layers_;
     };
 }