1 files changed, 10 insertions, 4 deletions
diff --git a/Swiften/StreamStack/StreamStack.h b/Swiften/StreamStack/StreamStack.h
index bd95811..263b1f5 100644
--- a/Swiften/StreamStack/StreamStack.h
+++ b/Swiften/StreamStack/StreamStack.h
@@ -19,30 +19,36 @@ namespace Swift {
     class LowLayer;
     class StreamLayer;
 
     class SWIFTEN_API StreamStack {
         public:
-            StreamStack(HighLayer* topLayer, LowLayer* bottomLayer);
+            StreamStack(std::unique_ptr<HighLayer> topLayer, std::unique_ptr<LowLayer> bottomLayer);
             ~StreamStack();
 
             void addLayer(std::unique_ptr<StreamLayer> /* streamLayer */);
 
             HighLayer* getTopLayer() const {
-                return topLayer_;
+                return topLayer_.get();
             }
 
             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:
-            HighLayer* topLayer_;
-            LowLayer* bottomLayer_;
+            std::unique_ptr<HighLayer> topLayer_;
+            std::unique_ptr<LowLayer> bottomLayer_;
             std::vector<std::unique_ptr<StreamLayer>> layers_;
     };
 }