From c961111d554fc888e3de7ce4c81a46c1d1cf066b Mon Sep 17 00:00:00 2001
From: Laura Hausmann <laura@hausmann.dev>
Date: Sun, 28 Jan 2024 03:33:20 +0100
Subject: [PATCH] Add authorization options to swagger-ui

---
 .../Core/Extensions/ServiceExtensions.cs      | 52 +++++++++++++++++++
 .../Extensions/WebApplicationExtensions.cs    | 14 +++++
 Iceshrimp.Backend/Startup.cs                  | 15 +-----
 3 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs b/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs
index 9521a647..61959abc 100644
--- a/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs
+++ b/Iceshrimp.Backend/Core/Extensions/ServiceExtensions.cs
@@ -1,3 +1,4 @@
+using System.Diagnostics.CodeAnalysis;
 using System.Threading.RateLimiting;
 using Iceshrimp.Backend.Controllers.Schemas;
 using Iceshrimp.Backend.Core.Configuration;
@@ -10,7 +11,9 @@ using Microsoft.AspNetCore.DataProtection;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption.ConfigurationModel;
 using Microsoft.AspNetCore.RateLimiting;
+using Microsoft.OpenApi.Models;
 using StackExchange.Redis;
+using Swashbuckle.AspNetCore.SwaggerGen;
 
 namespace Iceshrimp.Backend.Core.Extensions;
 
@@ -82,6 +85,55 @@ public static class ServiceExtensions {
 		});
 	}
 
+	public static void AddSwaggerGenWithOptions(this IServiceCollection services) {
+		services.AddSwaggerGen(options => {
+			options.SwaggerDoc("v1", new OpenApiInfo { Title = "Iceshrimp.NET", Version = "1.0" });
+			options.AddSecurityDefinition("user", new OpenApiSecurityScheme {
+				Name   = "Authorization token",
+				In     = ParameterLocation.Header,
+				Type   = SecuritySchemeType.Http,
+				Scheme = "bearer"
+			});
+			options.AddSecurityDefinition("admin", new OpenApiSecurityScheme {
+				Name   = "Authorization token",
+				In     = ParameterLocation.Header,
+				Type   = SecuritySchemeType.Http,
+				Scheme = "bearer"
+			});
+			options.OperationFilter<AuthorizeCheckOperationFilter>();
+		});
+	}
+
+	[SuppressMessage("ReSharper", "ClassNeverInstantiated.Local",
+	                 Justification = "SwaggerGenOptions.OperationFilter<T> instantiates this class at runtime")]
+	private class AuthorizeCheckOperationFilter : IOperationFilter {
+		public void Apply(OpenApiOperation operation, OperationFilterContext context) {
+			if (context.MethodInfo.DeclaringType is null)
+				return;
+
+			//TODO: separate admin & user authorize attributes
+			var hasAuthorize = context.MethodInfo.DeclaringType.GetCustomAttributes(true)
+			                          .OfType<AuthenticateAttribute>().Any() ||
+			                   context.MethodInfo.GetCustomAttributes(true)
+			                          .OfType<AuthenticateAttribute>().Any();
+
+			if (!hasAuthorize) return;
+			var schema = new OpenApiSecurityScheme {
+				Reference = new OpenApiReference {
+					Type = ReferenceType.SecurityScheme,
+					Id   = "user"
+				}
+			};
+
+			operation.Security = new List<OpenApiSecurityRequirement> {
+				new() {
+					[schema] = Array.Empty<string>()
+				}
+			};
+		}
+	}
+
+
 	public static void AddSlidingWindowRateLimiter(this IServiceCollection services) {
 		//TODO: separate limiter for authenticated users, partitioned by user id
 		//TODO: ipv6 /64 subnet buckets
diff --git a/Iceshrimp.Backend/Core/Extensions/WebApplicationExtensions.cs b/Iceshrimp.Backend/Core/Extensions/WebApplicationExtensions.cs
index 9c11180b..75b8fe93 100644
--- a/Iceshrimp.Backend/Core/Extensions/WebApplicationExtensions.cs
+++ b/Iceshrimp.Backend/Core/Extensions/WebApplicationExtensions.cs
@@ -16,6 +16,20 @@ public static class WebApplicationExtensions {
 		          .UseMiddleware<AuthorizedFetchMiddleware>();
 	}
 
+	public static IApplicationBuilder UseSwaggerWithOptions(this WebApplication app) {
+		app.UseSwagger();
+		app.UseSwaggerUI(options => {
+			options.DocumentTitle = "Iceshrimp API documentation";
+			options.SwaggerEndpoint("v1/swagger.json", "Iceshrimp.NET");
+			options.InjectStylesheet("/swagger/styles.css");
+			options.EnablePersistAuthorization();
+			options.EnableTryItOutByDefault();
+			options.DisplayRequestDuration();
+			options.DefaultModelsExpandDepth(-1); // Hide "Schemas" section
+		});
+		return app;
+	}
+
 	public static Config.InstanceSection Initialize(this WebApplication app, string[] args) {
 		var instanceConfig = app.Configuration.GetSection("Instance").Get<Config.InstanceSection>() ??
 		                     throw new Exception("Failed to read Instance config section");
diff --git a/Iceshrimp.Backend/Startup.cs b/Iceshrimp.Backend/Startup.cs
index 54b09db2..b1b22854 100644
--- a/Iceshrimp.Backend/Startup.cs
+++ b/Iceshrimp.Backend/Startup.cs
@@ -23,9 +23,7 @@ builder.Services.AddApiVersioning(options => {
 	options.UnsupportedApiVersionStatusCode = 501;
 });
 builder.Services.AddEndpointsApiExplorer();
-builder.Services.AddSwaggerGen(options => {
-	options.SwaggerDoc("v1", new OpenApiInfo { Title = "Iceshrimp.NET", Version = "1.0" });
-});
+builder.Services.AddSwaggerGenWithOptions();
 builder.Services.AddRazorPages();
 builder.Services.AddViteServices(options => {
 	options.PackageDirectory     = "../Iceshrimp.Frontend";
@@ -47,16 +45,7 @@ var config = app.Initialize(args);
 
 // This determines the order of middleware execution in the request pipeline
 app.UseRouting();
-app.UseSwagger();
-app.UseSwaggerUI(options => {
-	options.DocumentTitle = "Iceshrimp API documentation";
-	options.SwaggerEndpoint("v1/swagger.json", "Iceshrimp.NET");
-	options.InjectStylesheet("/swagger/styles.css");
-	options.EnablePersistAuthorization();
-	options.EnableTryItOutByDefault();
-	options.DisplayRequestDuration();
-	options.DefaultModelsExpandDepth(-1); // Hide "Schemas" section
-});
+app.UseSwaggerWithOptions();
 app.UseStaticFiles();
 app.UseRateLimiter();
 app.UseAuthorization();