[release/10.0.1xx] [BrowserRefresh] Handle gzip-compressed responses in ResponseStreamWrapper (#50535)

Co-authored-by: Mackinnon Buck <[email protected]>
Co-authored-by: javiercn <[email protected]>
Co-authored-by: Javier Calvarro Nelson <[email protected]>

Author: 3 people

src/BuiltInTools/BrowserRefresh/BrowserRefreshMiddleware.cs

@@ -46,6 +46,11 @@ public async Task InvokeAsync(HttpContext context)
                 try
                 {
                     await _next(context);
+
+                    // We complete the wrapper stream to ensure that any intermediate buffers
+                    // get fully flushed to the response stream. This is also required to
+                    // reliably determine whether script injection was performed.
+                    await responseStreamWrapper.CompleteAsync();
                 }
                 finally
                 {
@@ -219,7 +224,7 @@ internal static class Log
                 LogLevel.Warning,
                 new EventId(3, "FailedToConfiguredForRefreshes"),
                 "Unable to configure browser refresh script injection on the response. " +
-                $"Consider manually adding '{WebSocketScriptInjection.InjectedScript}' to the body of the page.");
+                $"Consider manually adding '{ScriptInjectingStream.InjectedScript}' to the body of the page.");
 
             private static readonly Action<ILogger, StringValues, Exception?> _responseCompressionDetected = LoggerMessage.Define<StringValues>(
                 LogLevel.Warning,

src/BuiltInTools/BrowserRefresh/BrowserScriptMiddleware.cs

@@ -44,7 +44,7 @@ internal static ReadOnlyMemory<byte> GetBlazorHotReloadJS()
         {
             var jsFileName = "Microsoft.AspNetCore.Watch.BrowserRefresh.BlazorHotReload.js";
             using var stream = new MemoryStream();
-            var manifestStream = typeof(WebSocketScriptInjection).Assembly.GetManifestResourceStream(jsFileName)!;
+            var manifestStream = typeof(BrowserScriptMiddleware).Assembly.GetManifestResourceStream(jsFileName)!;
             manifestStream.CopyTo(stream);
 
             return stream.ToArray();
@@ -61,7 +61,7 @@ internal static ReadOnlyMemory<byte> GetBrowserRefreshJS()
         internal static ReadOnlyMemory<byte> GetWebSocketClientJavaScript(string hostString, string serverKey)
         {
             var jsFileName = "Microsoft.AspNetCore.Watch.BrowserRefresh.WebSocketScriptInjection.js";
-            using var reader = new StreamReader(typeof(WebSocketScriptInjection).Assembly.GetManifestResourceStream(jsFileName)!);
+            using var reader = new StreamReader(typeof(BrowserScriptMiddleware).Assembly.GetManifestResourceStream(jsFileName)!);
             var script = reader.ReadToEnd()
                 .Replace("{{hostString}}", hostString)
                 .Replace("{{ServerKey}}", serverKey);

src/BuiltInTools/BrowserRefresh/ResponseStreamWrapper.cs

@@ -1,8 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Based on https://github.com/RickStrahl/Westwind.AspnetCore.LiveReload/blob/128b5f524e86954e997f2c453e7e5c1dcc3db746/Westwind.AspnetCore.LiveReload/ResponseStreamWrapper.cs
-
+using System.Diagnostics;
+using System.IO.Compression;
+using System.IO.Pipelines;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
 using Microsoft.Net.Http.Headers;
@@ -15,12 +16,18 @@ namespace Microsoft.AspNetCore.Watch.BrowserRefresh
     /// </summary>
     public class ResponseStreamWrapper : Stream
     {
-        private static readonly MediaTypeHeaderValue _textHtmlMediaType = new("text/html");
-        private readonly Stream _baseStream;
+        private static readonly MediaTypeHeaderValue s_textHtmlMediaType = new("text/html");
+
         private readonly HttpContext _context;
         private readonly ILogger _logger;
         private bool? _isHtmlResponse;
 
+        private Stream _baseStream;
+        private ScriptInjectingStream? _scriptInjectingStream;
+        private Pipe? _pipe;
+        private Task? _gzipCopyTask;
+        private bool _disposed;
+
         public ResponseStreamWrapper(HttpContext context, ILogger logger)
         {
             _context = context;
@@ -33,33 +40,25 @@ public ResponseStreamWrapper(HttpContext context, ILogger logger)
         public override bool CanWrite => true;
         public override long Length { get; }
         public override long Position { get; set; }
-        public bool ScriptInjectionPerformed { get; private set; }
-
-        public bool IsHtmlResponse => _isHtmlResponse ?? false;
+        public bool ScriptInjectionPerformed => _scriptInjectingStream?.ScriptInjectionPerformed == true;
+        public bool IsHtmlResponse => _isHtmlResponse == true;
 
         public override void Flush()
         {
             OnWrite();
             _baseStream.Flush();
         }
 
-        public override Task FlushAsync(CancellationToken cancellationToken)
+        public override async Task FlushAsync(CancellationToken cancellationToken)
         {
             OnWrite();
-            return _baseStream.FlushAsync(cancellationToken);
+            await _baseStream.FlushAsync(cancellationToken);
         }
 
         public override void Write(ReadOnlySpan<byte> buffer)
         {
             OnWrite();
-            if (IsHtmlResponse && !ScriptInjectionPerformed)
-            {
-                ScriptInjectionPerformed = WebSocketScriptInjection.TryInjectLiveReloadScript(_baseStream, buffer);
-            }
-            else
-            {
-                _baseStream.Write(buffer);
-            }
+            _baseStream.Write(buffer);
         }
 
         public override void WriteByte(byte value)
@@ -71,43 +70,19 @@ public override void WriteByte(byte value)
         public override void Write(byte[] buffer, int offset, int count)
         {
             OnWrite();
-
-            if (IsHtmlResponse && !ScriptInjectionPerformed)
-            {
-                ScriptInjectionPerformed = WebSocketScriptInjection.TryInjectLiveReloadScript(_baseStream, buffer.AsSpan(offset, count));
-            }
-            else
-            {
-                _baseStream.Write(buffer, offset, count);
-            }
+            _baseStream.Write(buffer.AsSpan(offset, count));
         }
 
         public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
             OnWrite();
-
-            if (IsHtmlResponse && !ScriptInjectionPerformed)
-            {
-                ScriptInjectionPerformed = await WebSocketScriptInjection.TryInjectLiveReloadScriptAsync(_baseStream, buffer.AsMemory(offset, count), cancellationToken);
-            }
-            else
-            {
-                await _baseStream.WriteAsync(buffer, offset, count, cancellationToken);
-            }
+            await _baseStream.WriteAsync(buffer.AsMemory(offset, count), cancellationToken);
         }
 
         public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
         {
             OnWrite();
-
-            if (IsHtmlResponse && !ScriptInjectionPerformed)
-            {
-                ScriptInjectionPerformed = await WebSocketScriptInjection.TryInjectLiveReloadScriptAsync(_baseStream, buffer, cancellationToken);
-            }
-            else
-            {
-                await _baseStream.WriteAsync(buffer, cancellationToken);
-            }
+            await _baseStream.WriteAsync(buffer, cancellationToken);
         }
 
         private void OnWrite()
@@ -122,15 +97,38 @@ private void OnWrite()
             _isHtmlResponse =
                 (response.StatusCode == StatusCodes.Status200OK || response.StatusCode == StatusCodes.Status500InternalServerError) &&
                 MediaTypeHeaderValue.TryParse(response.ContentType, out var mediaType) &&
-                mediaType.IsSubsetOf(_textHtmlMediaType) &&
+                mediaType.IsSubsetOf(s_textHtmlMediaType) &&
                 (!mediaType.Charset.HasValue || mediaType.Charset.Equals("utf-8", StringComparison.OrdinalIgnoreCase));
 
             if (_isHtmlResponse.Value)
             {
                 BrowserRefreshMiddleware.Log.SetupResponseForBrowserRefresh(_logger);
-
                 // Since we're changing the markup content, reset the content-length
                 response.Headers.ContentLength = null;
+
+                _scriptInjectingStream = new ScriptInjectingStream(_baseStream);
+
+                // By default, write directly to the script injection stream.
+                // We may change the base stream below if we detect that the response
+                // is compressed.
+                _baseStream = _scriptInjectingStream;
+
+                // Check if the response has gzip Content-Encoding
+                if (response.Headers.TryGetValue(HeaderNames.ContentEncoding, out var contentEncodingValues))
+                {
+                    var contentEncoding = contentEncodingValues.FirstOrDefault();
+                    if (string.Equals(contentEncoding, "gzip", StringComparison.OrdinalIgnoreCase))
+                    {
+                        // Remove the Content-Encoding header since we'll be serving uncompressed content
+                        response.Headers.Remove(HeaderNames.ContentEncoding);
+
+                        _pipe = new Pipe();
+                        var gzipStream = new GZipStream(_pipe.Reader.AsStream(leaveOpen: true), CompressionMode.Decompress, leaveOpen: true);
+
+                        _gzipCopyTask = gzipStream.CopyToAsync(_scriptInjectingStream);
+                        _baseStream = _pipe.Writer.AsStream(leaveOpen: true);
+                    }
+                }
             }
         }
 
@@ -145,5 +143,45 @@ public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken
              => throw new NotSupportedException();
 
         public override void SetLength(long value) => throw new NotSupportedException();
+
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                DisposeAsync().AsTask().GetAwaiter().GetResult();
+            }
+        }
+
+        public ValueTask CompleteAsync() => DisposeAsync();
+
+        public override async ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            _disposed = true;
+
+            if (_pipe is not null)
+            {
+                await _pipe.Writer.CompleteAsync();
+            }
+
+            if (_gzipCopyTask is not null)
+            {
+                await _gzipCopyTask;
+            }
+
+            if (_scriptInjectingStream is not null)
+            {
+                await _scriptInjectingStream.CompleteAsync();
+            }
+            else
+            {
+                Debug.Assert(_isHtmlResponse != true);
+                await _baseStream.FlushAsync();
+            }
+        }
     }
 }