From 393ab6a4f03b68d2a55336fc7ccf8c89f7c546ae Mon Sep 17 00:00:00 2001
From: BrennanConroy <brecon@microsoft.com>
Date: Mon, 11 Sep 2017 16:55:32 -0700
Subject: [PATCH] Add extension method to get HttpContext on
 HubConnectionContext (#849)

---
 .../HttpConnectionContextExtensions.cs        | 16 ++++++
 .../HubEndpointTests.cs                       | 53 +++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 src/Microsoft.AspNetCore.SignalR/HttpConnectionContextExtensions.cs

diff --git a/src/Microsoft.AspNetCore.SignalR/HttpConnectionContextExtensions.cs b/src/Microsoft.AspNetCore.SignalR/HttpConnectionContextExtensions.cs
new file mode 100644
index 0000000000..8220822413
--- /dev/null
+++ b/src/Microsoft.AspNetCore.SignalR/HttpConnectionContextExtensions.cs
@@ -0,0 +1,16 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Sockets.Http.Features;
+
+namespace Microsoft.AspNetCore.SignalR
+{
+    public static class HttpConnectionContextExtensions
+    {
+        public static HttpContext GetHttpContext(this HubConnectionContext connection)
+        {
+            return connection.Features.Get<IHttpContextFeature>()?.HttpContext;
+        }
+    }
+}
diff --git a/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs b/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
index 5a6602c69c..59956f2a46 100644
--- a/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
+++ b/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
@@ -8,8 +8,10 @@ using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Channels;
 using Microsoft.AspNetCore.Authorization;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.SignalR.Internal.Protocol;
 using Microsoft.AspNetCore.SignalR.Tests.Common;
+using Microsoft.AspNetCore.Sockets;
 using Microsoft.Extensions.DependencyInjection;
 using Moq;
 using Newtonsoft.Json;
@@ -852,6 +854,52 @@ namespace Microsoft.AspNetCore.SignalR.Tests
             }
         }
 
+        [Fact]
+        public async Task CanGetHttpContextFromHubConnectionContext()
+        {
+            var serviceProvider = CreateServiceProvider();
+
+            var endPoint = serviceProvider.GetService<HubEndPoint<MethodHub>>();
+
+            using (var client = new TestClient())
+            {
+                var httpContext = new DefaultHttpContext();
+                client.Connection.SetHttpContext(httpContext);
+                var endPointLifetime = endPoint.OnConnectedAsync(client.Connection);
+
+                await client.Connected.OrTimeout();
+
+                var result = (await client.InvokeAsync(nameof(MethodHub.HasHttpContext)).OrTimeout()).Result;
+                Assert.True((bool)result);
+
+                client.Dispose();
+
+                await endPointLifetime.OrTimeout();
+            }
+        }
+
+        [Fact]
+        public async Task GetHttpContextFromHubConnectionContextHandlesNull()
+        {
+            var serviceProvider = CreateServiceProvider();
+
+            var endPoint = serviceProvider.GetService<HubEndPoint<MethodHub>>();
+
+            using (var client = new TestClient())
+            {
+                var endPointLifetime = endPoint.OnConnectedAsync(client.Connection);
+
+                await client.Connected.OrTimeout();
+
+                var result = (await client.InvokeAsync(nameof(MethodHub.HasHttpContext)).OrTimeout()).Result;
+                Assert.False((bool)result);
+
+                client.Dispose();
+
+                await endPointLifetime.OrTimeout();
+            }
+        }
+
         private static void AssertHubMessage(HubMessage expected, HubMessage actual)
         {
             // We aren't testing InvocationIds here
@@ -1166,6 +1214,11 @@ namespace Microsoft.AspNetCore.SignalR.Tests
             {
                 return Clients.AllExcept(excludedIds).InvokeAsync("Send", message);
             }
+
+            public bool HasHttpContext()
+            {
+                return Context.Connection.GetHttpContext() != null;
+            }
         }
 
         private class InheritedHub : BaseHub