From 8f9f6bcb6f4fc36f4794c8d11b7b9746cb9b09c4 Mon Sep 17 00:00:00 2001
From: BrennanConroy <brecon@microsoft.com>
Date: Fri, 24 Feb 2017 14:42:09 -0800
Subject: [PATCH] Adding Auth to endpoints

---
 .../HubEndPoint.cs                            |  1 -
 .../EndPointOptions.cs                        |  2 +-
 .../HttpConnectionDispatcher.cs               |  9 ++-
 .../Internal/AuthorizeHelper.cs               | 64 +++++++++++++++++++
 .../Microsoft.AspNetCore.Sockets.csproj       |  1 +
 5 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 src/Microsoft.AspNetCore.Sockets/Internal/AuthorizeHelper.cs

diff --git a/src/Microsoft.AspNetCore.SignalR/HubEndPoint.cs b/src/Microsoft.AspNetCore.SignalR/HubEndPoint.cs
index 6e34c7249c..de5edc16c3 100644
--- a/src/Microsoft.AspNetCore.SignalR/HubEndPoint.cs
+++ b/src/Microsoft.AspNetCore.SignalR/HubEndPoint.cs
@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
-using System.IO.Pipelines;
 using System.Linq;
 using System.Reflection;
 using System.Threading;
diff --git a/src/Microsoft.AspNetCore.Sockets/EndPointOptions.cs b/src/Microsoft.AspNetCore.Sockets/EndPointOptions.cs
index 6b1a4091a1..d00b14a1a4 100644
--- a/src/Microsoft.AspNetCore.Sockets/EndPointOptions.cs
+++ b/src/Microsoft.AspNetCore.Sockets/EndPointOptions.cs
@@ -11,4 +11,4 @@ namespace Microsoft.AspNetCore.Sockets
 
         public TransportType Transports { get; set; } = TransportType.All;
     }
-}
+}
\ No newline at end of file
diff --git a/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs b/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs
index 1fa6bf0dbf..0a58e615ff 100644
--- a/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs
+++ b/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs
@@ -35,9 +35,12 @@ namespace Microsoft.AspNetCore.Sockets
 
         public async Task ExecuteAsync<TEndPoint>(string path, HttpContext context) where TEndPoint : EndPoint
         {
-            // Get the end point mapped to this http connection
-            var endpoint = (EndPoint)context.RequestServices.GetRequiredService<TEndPoint>();
             var options = context.RequestServices.GetRequiredService<IOptions<EndPointOptions<TEndPoint>>>().Value;
+            // TODO: Authorize attribute on EndPoint
+            if (!await AuthorizeHelper.AuthorizeAsync(context, options.Policy))
+            {
+                return;
+            }
 
             if (context.Request.Path.StartsWithSegments(path + "/negotiate"))
             {
@@ -49,6 +52,8 @@ namespace Microsoft.AspNetCore.Sockets
             }
             else
             {
+                // Get the end point mapped to this http connection
+                var endpoint = (EndPoint)context.RequestServices.GetRequiredService<TEndPoint>();
                 await ExecuteEndpointAsync(path, context, endpoint, options);
             }
         }
diff --git a/src/Microsoft.AspNetCore.Sockets/Internal/AuthorizeHelper.cs b/src/Microsoft.AspNetCore.Sockets/Internal/AuthorizeHelper.cs
new file mode 100644
index 0000000000..34089bd5bb
--- /dev/null
+++ b/src/Microsoft.AspNetCore.Sockets/Internal/AuthorizeHelper.cs
@@ -0,0 +1,64 @@
+// 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 System.Security.Claims;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Authorization;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Internal;
+
+namespace Microsoft.AspNetCore.Sockets.Internal
+{
+    public static class AuthorizeHelper
+    {
+        public static async Task<bool> AuthorizeAsync(HttpContext context, AuthorizationPolicy policy)
+        {
+            if (policy != null)
+            {
+                if (policy.AuthenticationSchemes != null && policy.AuthenticationSchemes.Count > 0)
+                {
+                    ClaimsPrincipal newPrincipal = null;
+                    foreach (var scheme in policy.AuthenticationSchemes)
+                    {
+                        var result = await context.Authentication.AuthenticateAsync(scheme);
+                        if (result != null)
+                        {
+                            newPrincipal = SecurityHelper.MergeUserPrincipal(newPrincipal, result);
+                        }
+                    }
+
+                    if (newPrincipal == null)
+                    {
+                        newPrincipal = new ClaimsPrincipal(new ClaimsIdentity());
+                    }
+
+                    context.User = newPrincipal;
+                }
+
+                var authService = context.RequestServices.GetRequiredService<IAuthorizationService>();
+                if (await authService.AuthorizeAsync(context.User, context, policy))
+                {
+                    return true;
+                }
+
+                // Challenge
+                if (policy.AuthenticationSchemes != null && policy.AuthenticationSchemes.Count > 0)
+                {
+                    foreach (var scheme in policy.AuthenticationSchemes)
+                    {
+                        await context.Authentication.ChallengeAsync(scheme, properties: null);
+                    }
+                }
+                else
+                {
+                    await context.Authentication.ChallengeAsync(properties: null);
+                }
+
+                return false;
+            }
+
+            return true;
+        }
+    }
+}
diff --git a/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj b/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj
index ee57b2968a..47ed82dbbd 100644
--- a/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj
+++ b/src/Microsoft.AspNetCore.Sockets/Microsoft.AspNetCore.Sockets.csproj
@@ -17,6 +17,7 @@
     <PackageReference Include="Microsoft.AspNetCore.Authorization" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.AspNetCore.Hosting.Abstractions" Version="$(AspNetCoreVersion)" />
     <PackageReference Include="Microsoft.AspNetCore.Routing" Version="$(AspNetCoreVersion)" />
+    <PackageReference Include="Microsoft.Extensions.SecurityHelper.Sources" Version="$(AspNetCoreVersion)" PrivateAssets="All" />
     <PackageReference Include="Microsoft.Extensions.TaskCache.Sources" Version="$(AspNetCoreVersion)" PrivateAssets="All" />
     <PackageReference Include="System.Reflection.TypeExtensions" Version="$(CoreFxVersion)" />
     <PackageReference Include="System.Security.Claims" Version="$(CoreFxVersion)" />