| | | 1 | | using System.Management.Automation; |
| | | 2 | | using System.Management.Automation.Runspaces; |
| | | 3 | | using Kestrun.Hosting; |
| | | 4 | | using Kestrun.Hosting.Options; |
| | | 5 | | using Kestrun.Languages; |
| | | 6 | | using Kestrun.Models; |
| | | 7 | | using Microsoft.AspNetCore.Diagnostics; |
| | | 8 | | namespace Kestrun.Middleware; |
| | | 9 | | |
| | | 10 | | /// <summary> |
| | | 11 | | /// Extension methods for adding status code pages middleware. |
| | | 12 | | /// </summary> |
| | | 13 | | public static class StatusCodePageExtensions |
| | | 14 | | { |
| | | 15 | | /// <summary> |
| | | 16 | | /// Applies the configured status code pages middleware to the specified application builder. |
| | | 17 | | /// </summary> |
| | | 18 | | /// <param name="app">The application builder.</param> |
| | | 19 | | /// <param name="options">The status code options.</param> |
| | | 20 | | /// <returns>The application builder.</returns> |
| | | 21 | | public static IApplicationBuilder UseStatusCodePages(this IApplicationBuilder app, StatusCodeOptions options) |
| | | 22 | | { |
| | 10 | 23 | | ArgumentNullException.ThrowIfNull(options); |
| | | 24 | | // 1) direct delegate options |
| | 9 | 25 | | if (TryUseDirectOptions(app, options) is { } direct) |
| | | 26 | | { |
| | 2 | 27 | | return direct; |
| | | 28 | | } |
| | | 29 | | |
| | | 30 | | // 2) redirects |
| | 7 | 31 | | if (TryUseRedirects(app, options) is { } redirects) |
| | | 32 | | { |
| | 2 | 33 | | return redirects; |
| | | 34 | | } |
| | | 35 | | |
| | | 36 | | // 3) re-execute |
| | 5 | 37 | | if (TryUseReExecute(app, options) is { } reexec) |
| | | 38 | | { |
| | 2 | 39 | | return reexec; |
| | | 40 | | } |
| | | 41 | | |
| | | 42 | | // 4) static body |
| | 3 | 43 | | if (TryUseStaticBody(app, options) is { } staticBody) |
| | | 44 | | { |
| | 1 | 45 | | return staticBody; |
| | | 46 | | } |
| | | 47 | | |
| | | 48 | | // 5) custom script handler |
| | 2 | 49 | | if (TryUseScriptHandler(app, options) is { } script) |
| | | 50 | | { |
| | 1 | 51 | | return script; |
| | | 52 | | } |
| | | 53 | | |
| | | 54 | | // default to built-in behavior |
| | 1 | 55 | | return app.UseStatusCodePages(); |
| | | 56 | | } |
| | | 57 | | |
| | | 58 | | private static IApplicationBuilder? TryUseDirectOptions(IApplicationBuilder app, StatusCodeOptions options) |
| | 9 | 59 | | => options.Options is not null ? app.UseStatusCodePages(options.Options) : null; |
| | | 60 | | |
| | | 61 | | private static IApplicationBuilder? TryUseRedirects(IApplicationBuilder app, StatusCodeOptions options) |
| | 7 | 62 | | => HasValue(options.LocationFormat) ? app.UseStatusCodePagesWithRedirects(options.LocationFormat!) : null; |
| | | 63 | | |
| | | 64 | | private static IApplicationBuilder? TryUseReExecute(IApplicationBuilder app, StatusCodeOptions options) |
| | | 65 | | { |
| | 5 | 66 | | if (!HasValue(options.PathFormat)) |
| | | 67 | | { |
| | 3 | 68 | | return null; |
| | | 69 | | } |
| | 2 | 70 | | var query = NormalizeQuery(options.QueryFormat); |
| | 2 | 71 | | return app.UseStatusCodePagesWithReExecute(options.PathFormat!, query); |
| | | 72 | | } |
| | | 73 | | |
| | | 74 | | private static IApplicationBuilder? TryUseStaticBody(IApplicationBuilder app, StatusCodeOptions options) |
| | | 75 | | { |
| | 3 | 76 | | if (!HasValue(options.ContentType) || !HasValue(options.BodyFormat)) |
| | | 77 | | { |
| | 2 | 78 | | return null; |
| | | 79 | | } |
| | 1 | 80 | | var safeBody = EscapeTemplate(options.BodyFormat!); |
| | 1 | 81 | | return app.UseStatusCodePages(options.ContentType!, safeBody); |
| | | 82 | | } |
| | | 83 | | |
| | | 84 | | private static IApplicationBuilder? TryUseScriptHandler(IApplicationBuilder app, StatusCodeOptions options) |
| | | 85 | | { |
| | 2 | 86 | | if (options.LanguageOptions is null) |
| | | 87 | | { |
| | 1 | 88 | | return null; |
| | | 89 | | } |
| | 1 | 90 | | var compiled = options.Host.CompileScript(options.LanguageOptions); |
| | 1 | 91 | | var handler = BuildScriptHandler(options, compiled); |
| | 1 | 92 | | return app.UseStatusCodePages(handler); |
| | | 93 | | } |
| | 16 | 94 | | private static bool HasValue(string? s) => !string.IsNullOrWhiteSpace(s); |
| | | 95 | | |
| | | 96 | | /// <summary> |
| | | 97 | | /// Normalizes the query string to ensure it starts with '?' if not empty. |
| | | 98 | | /// </summary> |
| | | 99 | | /// <param name="query">The query string.</param> |
| | | 100 | | /// <returns>The normalized query string.</returns> |
| | | 101 | | private static string NormalizeQuery(string? query) |
| | | 102 | | { |
| | 2 | 103 | | if (query is null) |
| | | 104 | | { |
| | 1 | 105 | | return string.Empty; |
| | | 106 | | } |
| | | 107 | | |
| | 1 | 108 | | var q = query.Trim(); |
| | 1 | 109 | | return q.Length > 0 && q[0] == '?' |
| | 1 | 110 | | ? q |
| | 1 | 111 | | : "?" + q; |
| | | 112 | | } |
| | | 113 | | |
| | | 114 | | /// <summary> |
| | | 115 | | /// Escape curly braces for String.Format safety, but preserve the {0} placeholder used for status code. |
| | | 116 | | /// </summary> |
| | | 117 | | /// <param name="bodyFormat">The body format string to escape.</param> |
| | | 118 | | /// <returns>The escaped body format string.</returns> |
| | | 119 | | /// <remarks>If the template already contains escaped braces ({{ or }}), assume it is pre-escaped and skip.</remarks |
| | | 120 | | private static string EscapeTemplate(string bodyFormat) |
| | | 121 | | { |
| | 1 | 122 | | if (bodyFormat.Contains("{{", StringComparison.Ordinal) || bodyFormat.Contains("}}", StringComparison.Ordinal)) |
| | | 123 | | { |
| | 0 | 124 | | return bodyFormat; // already escaped |
| | | 125 | | } |
| | | 126 | | |
| | 1 | 127 | | var escaped = bodyFormat.Replace("{", "{{").Replace("}", "}}"); |
| | | 128 | | // restore the status-code placeholder |
| | 1 | 129 | | escaped = escaped.Replace("{{0}}", "{0}"); |
| | 1 | 130 | | return escaped; |
| | | 131 | | } |
| | | 132 | | |
| | | 133 | | /// <summary> |
| | | 134 | | /// Builds a status code handler that executes the compiled script delegate. |
| | | 135 | | /// </summary> |
| | | 136 | | /// <param name="options">The status code options.</param> |
| | | 137 | | /// <param name="compiled">The compiled request delegate.</param> |
| | | 138 | | /// <returns>A function that handles the status code context.</returns> |
| | | 139 | | private static Func<StatusCodeContext, Task> BuildScriptHandler(StatusCodeOptions options, RequestDelegate compiled) |
| | | 140 | | { |
| | 1 | 141 | | return async context => |
| | 1 | 142 | | { |
| | 0 | 143 | | var httpContext = context.HttpContext; |
| | 1 | 144 | | |
| | 1 | 145 | | // If running a PowerShell script but the runspace middleware did not execute |
| | 1 | 146 | | // (e.g., no matched endpoint so UseWhen predicate failed), bootstrap a temporary |
| | 1 | 147 | | // runspace and KestrunContext so the compiled delegate can run safely. |
| | 0 | 148 | | var needsBootstrap = options.LanguageOptions!.Language == Scripting.ScriptLanguage.PowerShell && |
| | 0 | 149 | | !httpContext.Items.ContainsKey(PowerShellDelegateBuilder.PS_INSTANCE_KEY); |
| | 1 | 150 | | |
| | 0 | 151 | | if (!needsBootstrap) |
| | 1 | 152 | | { |
| | 0 | 153 | | await compiled(httpContext); |
| | 0 | 154 | | return; |
| | 1 | 155 | | } |
| | 1 | 156 | | |
| | 0 | 157 | | var pool = options.Host.RunspacePool; // throws if not initialized |
| | 0 | 158 | | Runspace? runspace = null; |
| | 0 | 159 | | PowerShell? ps = null; |
| | 1 | 160 | | |
| | 1 | 161 | | try |
| | 1 | 162 | | { |
| | 0 | 163 | | runspace = await pool.AcquireAsync(httpContext.RequestAborted); |
| | 0 | 164 | | ps = PowerShell.Create(); |
| | 0 | 165 | | ps.Runspace = runspace; |
| | 1 | 166 | | |
| | 1 | 167 | | // Build Kestrun abstractions and inject into context for PS delegate to use |
| | 0 | 168 | | var kr = new KestrunContext(pool.Host, httpContext); |
| | 0 | 169 | | httpContext.Items[PowerShellDelegateBuilder.PS_INSTANCE_KEY] = ps; |
| | 0 | 170 | | httpContext.Items[PowerShellDelegateBuilder.KR_CONTEXT_KEY] = kr; |
| | 0 | 171 | | var ss = ps.Runspace.SessionStateProxy; |
| | 0 | 172 | | ss.SetVariable("Context", kr); |
| | 1 | 173 | | |
| | 0 | 174 | | await compiled(httpContext); |
| | 0 | 175 | | } |
| | 1 | 176 | | finally |
| | 1 | 177 | | { |
| | 0 | 178 | | ps?.Dispose(); |
| | 0 | 179 | | if (runspace is not null) |
| | 1 | 180 | | { |
| | 0 | 181 | | pool.Release(runspace); |
| | 1 | 182 | | } |
| | 0 | 183 | | _ = httpContext.Items.Remove(PowerShellDelegateBuilder.PS_INSTANCE_KEY); |
| | 0 | 184 | | _ = httpContext.Items.Remove(PowerShellDelegateBuilder.KR_CONTEXT_KEY); |
| | 1 | 185 | | } |
| | 1 | 186 | | }; |
| | | 187 | | } |
| | | 188 | | } |