Hot reload script injection improvements

src/BuiltInTools/BrowserRefresh/BrowserRefreshMiddleware.cs

@@ -3,6 +3,8 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
@@ -26,10 +28,10 @@ public async Task InvokeAsync(HttpContext context)
             // We only need to support this for requests that could be initiated by a browser.
             if (IsBrowserDocumentRequest(context))
             {
-                // Use a custom StreamWrapper to rewrite output on Write/WriteAsync
-                using var responseStreamWrapper = new ResponseStreamWrapper(context, _logger);
+                // Use a custom stream to buffer the response body for rewriting.
+                using var memoryStream = new MemoryStream();
                 var originalBodyFeature = context.Features.Get<IHttpResponseBodyFeature>();
-                context.Features.Set<IHttpResponseBodyFeature>(new StreamResponseBodyFeature(responseStreamWrapper));
+                context.Features.Set<IHttpResponseBodyFeature>(new StreamResponseBodyFeature(memoryStream));
 
                 try
                 {
@@ -40,15 +42,35 @@ public async Task InvokeAsync(HttpContext context)
                     context.Features.Set(originalBodyFeature);
                 }
 
-                if (responseStreamWrapper.IsHtmlResponse && _logger.IsEnabled(LogLevel.Debug))
+                if (memoryStream.TryGetBuffer(out var buffer) && buffer.Count > 0)
                 {
-                    if (responseStreamWrapper.ScriptInjectionPerformed)
+                    var response = context.Response;
+                    var baseStream = response.Body;
+
+                    if (IsHtmlResponse(response))
                     {
-                        Log.BrowserConfiguredForRefreshes(_logger);
+                        Log.SetupResponseForBrowserRefresh(_logger);
+
+                        // Since we're changing the markup content, reset the content-length
+                        response.Headers.ContentLength = null;
+
+                        var scriptInjectionPerformed = await WebSocketScriptInjection.TryInjectLiveReloadScriptAsync(baseStream, buffer);
+                        if (scriptInjectionPerformed)
+                        {
+                            Log.BrowserConfiguredForRefreshes(_logger);
+                        }
+                        else if (response.Headers.TryGetValue(HeaderNames.ContentEncoding, out var contentEncodings))
+                        {
+                            Log.ResponseCompressionDetected(_logger, contentEncodings);
+                        }
+                        else
+                        {
+                            Log.FailedToConfiguredForRefreshes(_logger);
+                        }
                     }
                     else
                     {
-                        Log.FailedToConfiguredForRefreshes(_logger);
+                        await baseStream.WriteAsync(buffer);
                     }
                 }
             }
@@ -92,26 +114,41 @@ internal static bool IsBrowserDocumentRequest(HttpContext context)
             return false;
         }
 
+        private bool IsHtmlResponse(HttpResponse response)
+            => (response.StatusCode == StatusCodes.Status200OK || response.StatusCode == StatusCodes.Status500InternalServerError) &&
+                MediaTypeHeaderValue.TryParse(response.ContentType, out var mediaType) &&
+                mediaType.IsSubsetOf(_textHtmlMediaType) &&
+                (!mediaType.Charset.HasValue || mediaType.Charset.Equals("utf-8", StringComparison.OrdinalIgnoreCase));
+
         internal static class Log
         {
             private static readonly Action<ILogger, Exception?> _setupResponseForBrowserRefresh = LoggerMessage.Define(
-               LogLevel.Debug,
+                LogLevel.Debug,
                 new EventId(1, "SetUpResponseForBrowserRefresh"),
-               "Response markup is scheduled to include browser refresh script injection.");
+                "Response markup is scheduled to include browser refresh script injection.");
 
             private static readonly Action<ILogger, Exception?> _browserConfiguredForRefreshes = LoggerMessage.Define(
-               LogLevel.Debug,
+                LogLevel.Debug,
                 new EventId(2, "BrowserConfiguredForRefreshes"),
-               "Response markup was updated to include browser refresh script injection.");
+                "Response markup was updated to include browser refresh script injection.");
 
             private static readonly Action<ILogger, Exception?> _failedToConfigureForRefreshes = LoggerMessage.Define(
-               LogLevel.Debug,
+                LogLevel.Warning,
                 new EventId(3, "FailedToConfiguredForRefreshes"),
-               "Unable to configure browser refresh script injection on the response.");
+                "Unable to configure browser refresh script injection on the response. " +
+                $"Consider manually adding '{WebSocketScriptInjection.InjectedScript}' to the body of the page.");
+
+            private static readonly Action<ILogger, StringValues, Exception?> _responseCompressionDetected = LoggerMessage.Define<StringValues>(
+                LogLevel.Warning,
+                new EventId(4, "ResponseCompressionDetected"),
+                "Unable to configure browser refresh script injection on the response. " +
+                $"This may have been caused by the response's {HeaderNames.ContentEncoding}: '{{encoding}}'. " +
+                "Consider disabling response compression.");
 
             public static void SetupResponseForBrowserRefresh(ILogger logger) => _setupResponseForBrowserRefresh(logger, null);
             public static void BrowserConfiguredForRefreshes(ILogger logger) => _browserConfiguredForRefreshes(logger, null);
             public static void FailedToConfiguredForRefreshes(ILogger logger) => _failedToConfigureForRefreshes(logger, null);
+            public static void ResponseCompressionDetected(ILogger logger, StringValues encoding) => _responseCompressionDetected(logger, encoding, null);
         }
     }
 }

src/BuiltInTools/BrowserRefresh/ReadOnlySpanOfByteExtensions.cs

@@ -0,0 +1,38 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace System;
+
+internal static class ReadOnlySpanOfByteExtensions
+{
+    public static int LastIndexOfNonWhiteSpace(this ReadOnlySpan<byte> buffer)
+    {
+        for (var i = buffer.Length - 1; i >= 0; i--)
+        {
+            if (!char.IsWhiteSpace(Convert.ToChar(buffer[i])))
+            {
+                return i;
+            }
+        }
+
+        return -1;
+    }
+
+    public static bool EndsWithIgnoreCase(this ReadOnlySpan<byte> buffer, ReadOnlySpan<byte> value)
+    {
+        if (buffer.Length < value.Length)
+        {
+            return false;
+        }
+
+        for (var i = 1; i <= value.Length; i++)
+        {
+            if (char.ToLowerInvariant(Convert.ToChar(value[^i])) != char.ToLowerInvariant(Convert.ToChar(buffer[^i])))
+            {
+                return false;
+            }
+        }
+
+        return true;
+    }
+}