From 7a729ea8d69b5cea16d7ba4a03748cd8a4caf871 Mon Sep 17 00:00:00 2001
From: crobibero <cody@robibe.ro>
Date: Wed, 2 Dec 2020 14:59:57 -0700
Subject: [PATCH] Move OpenApiSecurityScheme to OperationFilter

---
 .../ApiServiceCollectionExtensions.cs         | 13 +---
 .../SecurityRequirementsOperationFilter.cs    | 78 +++++++++++++++++++
 2 files changed, 79 insertions(+), 12 deletions(-)
 create mode 100644 Jellyfin.Server/Filters/SecurityRequirementsOperationFilter.cs

diff --git a/Jellyfin.Server/Extensions/ApiServiceCollectionExtensions.cs b/Jellyfin.Server/Extensions/ApiServiceCollectionExtensions.cs
index e5b9620f78..618a4e92b4 100644
--- a/Jellyfin.Server/Extensions/ApiServiceCollectionExtensions.cs
+++ b/Jellyfin.Server/Extensions/ApiServiceCollectionExtensions.cs
@@ -236,18 +236,6 @@ namespace Jellyfin.Server.Extensions
                     Description = "API key header parameter"
                 });
 
-                var securitySchemeRef = new OpenApiSecurityScheme
-                {
-                    Reference = new OpenApiReference { Type = ReferenceType.SecurityScheme, Id = AuthenticationSchemes.CustomAuthentication },
-                };
-
-                // TODO: Apply this with an operation filter instead of globally
-                // https://github.com/domaindrivendev/Swashbuckle.AspNetCore#add-security-definitions-and-requirements
-                c.AddSecurityRequirement(new OpenApiSecurityRequirement
-                {
-                    { securitySchemeRef, Array.Empty<string>() }
-                });
-
                 // Add all xml doc files to swagger generator.
                 var xmlFiles = Directory.GetFiles(
                     AppContext.BaseDirectory,
@@ -277,6 +265,7 @@ namespace Jellyfin.Server.Extensions
                 // TODO - remove when all types are supported in System.Text.Json
                 c.AddSwaggerTypeMappings();
 
+                c.OperationFilter<SecurityRequirementsOperationFilter>();
                 c.OperationFilter<FileResponseFilter>();
                 c.DocumentFilter<WebsocketModelFilter>();
             });
diff --git a/Jellyfin.Server/Filters/SecurityRequirementsOperationFilter.cs b/Jellyfin.Server/Filters/SecurityRequirementsOperationFilter.cs
new file mode 100644
index 0000000000..802662ce2f
--- /dev/null
+++ b/Jellyfin.Server/Filters/SecurityRequirementsOperationFilter.cs
@@ -0,0 +1,78 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Jellyfin.Api.Constants;
+using Microsoft.AspNetCore.Authorization;
+using Microsoft.OpenApi.Models;
+using Swashbuckle.AspNetCore.SwaggerGen;
+
+namespace Jellyfin.Server.Filters
+{
+    /// <summary>
+    /// Security requirement operation filter.
+    /// </summary>
+    public class SecurityRequirementsOperationFilter : IOperationFilter
+    {
+        /// <inheritdoc />
+        public void Apply(OpenApiOperation operation, OperationFilterContext context)
+        {
+            var requiredScopes = new List<string>();
+
+            // Add all method scopes.
+            foreach (var attribute in context.MethodInfo.GetCustomAttributes(true))
+            {
+                if (attribute is AuthorizeAttribute authorizeAttribute
+                    && authorizeAttribute.Policy != null
+                    && !requiredScopes.Contains(authorizeAttribute.Policy, StringComparer.Ordinal))
+                {
+                    requiredScopes.Add(authorizeAttribute.Policy);
+                }
+            }
+
+            // Add controller scopes if any.
+            var controllerAttributes = context.MethodInfo.DeclaringType?.GetCustomAttributes(true);
+            if (controllerAttributes != null)
+            {
+                foreach (var attribute in controllerAttributes)
+                {
+                    if (attribute is AuthorizeAttribute authorizeAttribute
+                        && authorizeAttribute.Policy != null
+                        && !requiredScopes.Contains(authorizeAttribute.Policy, StringComparer.Ordinal))
+                    {
+                        requiredScopes.Add(authorizeAttribute.Policy);
+                    }
+                }
+            }
+
+            if (requiredScopes.Count != 0)
+            {
+                if (!operation.Responses.ContainsKey("401"))
+                {
+                    operation.Responses.Add("401", new OpenApiResponse { Description = "Unauthorized" });
+                }
+
+                if (!operation.Responses.ContainsKey("403"))
+                {
+                    operation.Responses.Add("403", new OpenApiResponse { Description = "Forbidden" });
+                }
+
+                var scheme = new OpenApiSecurityScheme
+                {
+                    Reference = new OpenApiReference
+                    {
+                        Type = ReferenceType.SecurityScheme,
+                        Id = AuthenticationSchemes.CustomAuthentication
+                    }
+                };
+
+                operation.Security = new List<OpenApiSecurityRequirement>
+                {
+                    new OpenApiSecurityRequirement
+                    {
+                        [scheme] = requiredScopes
+                    }
+                };
+            }
+        }
+    }
+}
\ No newline at end of file