1 files changed, 63 insertions, 2 deletions
diff --git a/MediaBrowser.Server.Implementations/Session/SessionWebSocketListener.cs b/MediaBrowser.Server.Implementations/Session/SessionWebSocketListener.cs
index d1bb4b837..b581f9144 100644
--- a/MediaBrowser.Server.Implementations/Session/SessionWebSocketListener.cs
+++ b/MediaBrowser.Server.Implementations/Session/SessionWebSocketListener.cs
@@ -1,9 +1,12 @@
 using MediaBrowser.Controller.Net;
+using MediaBrowser.Controller.Security;
 using MediaBrowser.Controller.Session;
+using MediaBrowser.Model.Events;
 using MediaBrowser.Model.Logging;
 using MediaBrowser.Model.Serialization;
 using MediaBrowser.Model.Session;
 using System;
+using System.Collections.Specialized;
 using System.Globalization;
 using System.Linq;
 using System.Threading.Tasks;
@@ -13,7 +16,7 @@ namespace MediaBrowser.Server.Implementations.Session
     /// <summary>
     /// Class SessionWebSocketListener
     /// </summary>
-    public class SessionWebSocketListener : IWebSocketListener
+    public class SessionWebSocketListener : IWebSocketListener, IDisposable
     {
         /// <summary>
         /// The _true task result
@@ -35,17 +38,75 @@ namespace MediaBrowser.Server.Implementations.Session
         /// </summary>
         private readonly IJsonSerializer _json;
 
+        private readonly IHttpServer _httpServer;
+        private readonly IServerManager _serverManager;
+
+
         /// <summary>
         /// Initializes a new instance of the <see cref="SessionWebSocketListener" /> class.
         /// </summary>
         /// <param name="sessionManager">The session manager.</param>
         /// <param name="logManager">The log manager.</param>
         /// <param name="json">The json.</param>
-        public SessionWebSocketListener(ISessionManager sessionManager, ILogManager logManager, IJsonSerializer json)
+        /// <param name="httpServer">The HTTP server.</param>
+        /// <param name="serverManager">The server manager.</param>
+        public SessionWebSocketListener(ISessionManager sessionManager, ILogManager logManager, IJsonSerializer json, IHttpServer httpServer, IServerManager serverManager)
         {
             _sessionManager = sessionManager;
             _logger = logManager.GetLogger(GetType().Name);
             _json = json;
+            _httpServer = httpServer;
+            _serverManager = serverManager;
+            httpServer.WebSocketConnecting += _httpServer_WebSocketConnecting;
+            serverManager.WebSocketConnected += _serverManager_WebSocketConnected;
+        }
+
+        void _serverManager_WebSocketConnected(object sender, GenericEventArgs<IWebSocketConnection> e)
+        {
+            var session = GetSession(e.Argument.QueryString);
+
+            if (session != null)
+            {
+                var controller = session.SessionController as WebSocketController;
+
+                if (controller == null)
+                {
+                    controller = new WebSocketController(session, _logger, _sessionManager);
+                }
+
+                controller.AddWebSocket(e.Argument);
+
+                session.SessionController = controller;
+            }
+            else
+            {
+                _logger.Warn("Unable to determine session based on url: {0}", e.Argument.Url);
+            }
+        }
+
+        void _httpServer_WebSocketConnecting(object sender, WebSocketConnectingEventArgs e)
+        {
+            if (e.QueryString.AllKeys.Contains("api_key", StringComparer.OrdinalIgnoreCase))
+            {
+                var session = GetSession(e.QueryString);
+
+                if (session == null)
+                {
+                    e.AllowConnection = false;
+                }
+            }
+        }
+
+        private SessionInfo GetSession(NameValueCollection queryString)
+        {
+            var token = queryString["api_key"];
+            return _sessionManager.GetSessionByAuthenticationToken(token);
+        }
+
+        public void Dispose()
+        {
+            _httpServer.WebSocketConnecting -= _httpServer_WebSocketConnecting;
+            _serverManager.WebSocketConnected -= _serverManager_WebSocketConnected;
         }
 
         /// <summary>