Fix GZip handling for requests (#4165)

* Fix GZip handling for requests

* Flush GZipStream when compressing.

* Removing flush call - not needed for Test Proxy target platforms

* repair record mode operation. we were exhausting the request body stream.

* Update mock handler

Co-authored-by: scbedd <[email protected]>

Author: JoshLove-msft

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/PlaybackTests.cs

@@ -6,6 +6,7 @@
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
+using Azure.Sdk.Tools.TestProxy.Common;
 using Xunit;
 
 namespace Azure.Sdk.Tools.TestProxy.Tests
@@ -217,5 +218,35 @@ public async Task TestPlaybackSetsRetryAfterToZero()
             await testRecordingHandler.HandlePlaybackRequest(recordingId, request, response);
             Assert.False(response.Headers.ContainsKey("Retry-After"));
         }
+
+        [Fact]
+        public async Task TestPlaybackWithGZippedContentPlayback()
+        {
+            RecordingHandler testRecordingHandler = new RecordingHandler(Directory.GetCurrentDirectory());
+            var httpContext = new DefaultHttpContext();
+            var body = "{\"x-recording-file\":\"Test.RecordEntries/request_response_with_gzipped_content.json\"}";
+            httpContext.Request.Body = TestHelpers.GenerateStreamRequestBody(body);
+            httpContext.Request.ContentLength = body.Length;
+
+            var controller = new Playback(testRecordingHandler, new NullLoggerFactory())
+            {
+                ControllerContext = new ControllerContext()
+                {
+                    HttpContext = httpContext
+                }
+            };
+            await controller.Start();
+
+            var recordingId = httpContext.Response.Headers["x-recording-id"].ToString();
+            Assert.NotNull(recordingId);
+            Assert.True(testRecordingHandler.PlaybackSessions.ContainsKey(recordingId));
+            var entry = testRecordingHandler.PlaybackSessions[recordingId].Session.Entries[0];
+            HttpRequest request = TestHelpers.CreateRequestFromEntry(entry);
+
+            // compress the body to simulate what the request coming from the library will look like
+            request.Body = new MemoryStream(GZipUtilities.CompressBody(BinaryData.FromStream(request.Body).ToArray(), request.Headers));
+            HttpResponse response = new DefaultHttpContext().Response;
+            await testRecordingHandler.HandlePlaybackRequest(recordingId, request, response);
+        }
     }
 }

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs

@@ -12,6 +12,7 @@
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Net.Http.Headers;
 using System.Text;
 using System.Text.Json;
 using System.Threading;
@@ -696,6 +697,52 @@ public async Task TestRecordMaintainsUpstreamOverrideHostHeader(string upstreamH
             }
         }
 
+        [Fact]
+        public async Task TestRecordWithGZippedContent()
+        {
+            var httpContext = new DefaultHttpContext();
+            var bodyBytes = Encoding.UTF8.GetBytes("{\"hello\":\"world\"}");
+            var mockClient = new HttpClient(new MockHttpHandler(bodyBytes, "application/json", "gzip"));
+            var path = Directory.GetCurrentDirectory();
+            var recordingHandler = new RecordingHandler(path)
+            {
+                RedirectableClient = mockClient,
+                RedirectlessClient = mockClient
+            };
+
+            var relativePath = "recordings/gzip";
+            var fullPathToRecording = Path.Combine(path, relativePath) + ".json";
+
+            await recordingHandler.StartRecordingAsync(relativePath, httpContext.Response);
+
+            var recordingId = httpContext.Response.Headers["x-recording-id"].ToString();
+
+            httpContext.Request.ContentType = "application/json";
+            httpContext.Request.Headers["Content-Encoding"] = "gzip";
+            httpContext.Request.ContentLength = 0;
+            httpContext.Request.Headers["x-recording-id"] = recordingId;
+            httpContext.Request.Headers["x-recording-upstream-base-uri"] = "http://example.org";
+            httpContext.Request.Method = "GET";
+            httpContext.Request.Body = new MemoryStream(GZipUtilities.CompressBody(bodyBytes, httpContext.Request.Headers));
+
+            await recordingHandler.HandleRecordRequestAsync(recordingId, httpContext.Request, httpContext.Response);
+            recordingHandler.StopRecording(recordingId);
+
+            try
+            {
+                using var fileStream = File.Open(fullPathToRecording, FileMode.Open);
+                using var doc = JsonDocument.Parse(fileStream);
+                var record = RecordSession.Deserialize(doc.RootElement);
+                var entry = record.Entries.First();
+                Assert.Equal("{\"hello\":\"world\"}", Encoding.UTF8.GetString(entry.Request.Body));
+                Assert.Equal("{\"hello\":\"world\"}", Encoding.UTF8.GetString(entry.Response.Body));
+            }
+            finally
+            {
+                File.Delete(fullPathToRecording);
+            }
+        }
+
         #region SetRecordingOptions
         [Theory]
         [InlineData("{ \"HandleRedirects\": \"true\"}", true)]
@@ -977,17 +1024,41 @@ public IgnoreOnLinuxFact()
     internal class MockHttpHandler : HttpMessageHandler
     {
         public const string DefaultResponse = "default response";
+        private readonly byte[] _responseContent;
+        private readonly string _contentType;
+        private readonly string _contentEncoding;
 
-        public MockHttpHandler()
+        public MockHttpHandler(byte[] responseContent = default, string contentType = default, string contentEncoding = default)
         {
+            _responseContent = responseContent ?? Encoding.UTF8.GetBytes(DefaultResponse);
+            _contentType = contentType ?? "application/json";
+            _contentEncoding = contentEncoding;
         }
-        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+
+        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
+            // should not throw as stream should not be disposed
+            var content = await request.Content.ReadAsStringAsync();
+            Assert.NotEmpty(content);
+            var response = new HttpResponseMessage(HttpStatusCode.OK);
+
+            // we need to set the content before the content headers as otherwise they will be cleared out after setting content.
+            if (_contentEncoding == "gzip")
+            {
+                response.Content = new ByteArrayContent(GZipUtilities.CompressBodyCore(_responseContent));
+            }
+            else
+            {
+                response.Content = new ByteArrayContent(_responseContent);
+            }
 
-            return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
+            response.Content.Headers.ContentType = new MediaTypeHeaderValue(_contentType);
+            if (_contentEncoding != null)
             {
-                Content = new StringContent(DefaultResponse)
-            });
+                response.Content.Headers.ContentEncoding.Add(_contentEncoding);
+            }
+
+            return response;
         }
     }
 }

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/Test.RecordEntries/request_response_with_gzipped_content.json

@@ -0,0 +1,48 @@
+{
+  "Entries": [
+    {
+      "RequestUri": "https://fakeazsdktestaccount.table.core.windows.net/Tables",
+      "RequestMethod": "POST",
+      "RequestHeaders": {
+        "Accept": "application/json;odata=minimalmetadata",
+        "Accept-Encoding": "gzip, deflate",
+        "Authorization": "Sanitized",
+        "Connection": "keep-alive",
+        "Content-Length": "34",
+        "Content-Type": "application/json",
+        "Content-Encoding": "gzip",
+        "DataServiceVersion": "3.0",
+        "Date": "Tue, 18 May 2021 23:27:42 GMT",
+        "User-Agent": "azsdk-python-data-tables/12.0.0b7 Python/3.8.6 (Windows-10-10.0.19041-SP0)",
+        "x-ms-client-request-id": "a4c24b7a-b830-11eb-a05e-10e7c6392c5a",
+        "x-ms-date": "Tue, 18 May 2021 23:27:42 GMT",
+        "x-ms-version": "2019-02-02"
+      },
+      "RequestBody": "{\u0022TableName\u0022:    \u0022listtable09bf2a3d\u0022}",
+      "StatusCode": 201,
+      "ResponseHeaders": {
+        "Cache-Control": "no-cache",
+        "Content-Type": "application/json",
+        "Content-Encoding": "gzip",
+        "Date": "Tue, 18 May 2021 23:27:43 GMT",
+        "Retry-After": "10",
+        "Location": "https://fakeazsdktestaccount.table.core.windows.net/Tables(\u0027listtable09bf2a3d\u0027)",
+        "Server": [
+          "Windows-Azure-Table/1.0",
+          "Microsoft-HTTPAPI/2.0"
+        ],
+        "Transfer-Encoding": "chunked",
+        "X-Content-Type-Options": "nosniff",
+        "x-ms-client-request-id": "a4c24b7a-b830-11eb-a05e-10e7c6392c5a",
+        "x-ms-request-id": "d2270777-c002-0072-313d-4ce19f000000",
+        "x-ms-version": "2019-02-02"
+      },
+      "ResponseBody": {
+        "odata.metadata": "https://fakeazsdktestaccount.table.core.windows.net/$metadata#Tables/@Element",
+        "TableName": "listtable09bf2a3d",
+        "connectionString":  null
+      }
+    }
+  ],
+  "Variables": {}
+}

tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/GZipUtilities.cs

@@ -0,0 +1,77 @@
+using System.Collections.Generic;
+using System.IO;
+using System.IO.Compression;
+using System.Linq;
+using System.Net.Http.Headers;
+using Microsoft.AspNetCore.Http;
+
+namespace Azure.Sdk.Tools.TestProxy.Common
+{
+    /// <summary>
+    /// Utility methods to compress and decompress content to/from GZip.
+    /// </summary>
+    public static class GZipUtilities
+    {
+        private const string Gzip = "gzip";
+        private const string ContentEncoding = "Content-Encoding";
+
+        public static byte[] CompressBody(byte[] incomingBody, IDictionary<string, string[]> headers)
+        {
+            if (headers.TryGetValue(ContentEncoding, out var values) && values.Contains(Gzip))
+            {
+                return CompressBodyCore(incomingBody);
+            }
+
+            return incomingBody;
+        }
+
+        public static byte[] CompressBody(byte[] incomingBody, IHeaderDictionary headers)
+        {
+            if (headers.TryGetValue(ContentEncoding, out var values) && values.Contains(Gzip))
+            {
+                return CompressBodyCore(incomingBody);
+            }
+
+            return incomingBody;
+        }
+
+        public static byte[] CompressBodyCore(byte[] body)
+        {
+            using var uncompressedStream = new MemoryStream(body);
+            using var resultStream = new MemoryStream();
+            using (var compressedStream = new GZipStream(resultStream, CompressionMode.Compress))
+            {
+                uncompressedStream.CopyTo(compressedStream);
+            }
+            return resultStream.ToArray();
+        }
+
+        public static byte[] DecompressBody(MemoryStream incomingBody, HttpContentHeaders headers)
+        {
+            if (headers.TryGetValues(ContentEncoding, out var values) && values.Contains(Gzip))
+            {
+                return DecompressBodyCore(incomingBody);
+            }
+
+            return incomingBody.ToArray();
+        }
+
+        public static byte[] DecompressBody(byte[] incomingBody, IHeaderDictionary headers)
+        {
+            if (headers.TryGetValue(ContentEncoding, out var values) && values.Contains(Gzip))
+            {
+                return DecompressBodyCore(new MemoryStream(incomingBody));
+            }
+
+            return incomingBody;
+        }
+
+        private static byte[] DecompressBodyCore(MemoryStream stream)
+        {
+            using var uncompressedStream = new GZipStream(stream, CompressionMode.Decompress);
+            using var resultStream = new MemoryStream();
+            uncompressedStream.CopyTo(resultStream);
+            return resultStream.ToArray();
+        }
+    }
+}

tools/test-proxy/Azure.Sdk.Tools.TestProxy/Properties/launchSettings.json

@@ -10,7 +10,6 @@
   "profiles": {
     "Azure.Sdk.Tools.TestProxy": {
       "commandName": "Project",
-      "commandLineArgs": "--help",
       "environmentVariables": {
         "ASPNETCORE_ENVIRONMENT": "Development",
         "Logging__LogLevel__Microsoft": "Information"

tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs

@@ -200,7 +200,7 @@ public async Task HandleRecordRequestAsync(string recordingId, HttpRequest incom
 
             var entry = await CreateEntryAsync(incomingRequest).ConfigureAwait(false);
 
-            var upstreamRequest = CreateUpstreamRequest(incomingRequest, entry.Request.Body);
+            var upstreamRequest = CreateUpstreamRequest(incomingRequest, GZipUtilities.CompressBody(entry.Request.Body, entry.Request.Headers));
 
             HttpResponseMessage upstreamResponse = null;
 
@@ -218,7 +218,7 @@ public async Task HandleRecordRequestAsync(string recordingId, HttpRequest incom
             // HEAD requests do NOT have a body regardless of the value of the Content-Length header
             if (incomingRequest.Method.ToUpperInvariant() != "HEAD")
             {
-                body = DecompressBody((MemoryStream)await upstreamResponse.Content.ReadAsStreamAsync().ConfigureAwait(false), upstreamResponse.Content.Headers);
+                body = GZipUtilities.DecompressBody((MemoryStream)await upstreamResponse.Content.ReadAsStreamAsync().ConfigureAwait(false), upstreamResponse.Content.Headers);
             }
 
             entry.Response.Body = body.Length == 0 ? null : body;
@@ -250,7 +250,7 @@ public async Task HandleRecordRequestAsync(string recordingId, HttpRequest incom
 
             if (entry.Response.Body?.Length > 0)
             {
-                var bodyData = CompressBody(entry.Response.Body, entry.Response.Headers);
+                var bodyData = GZipUtilities.CompressBody(entry.Response.Body, entry.Response.Headers);
                 outgoingResponse.ContentLength = bodyData.Length;
                 await outgoingResponse.Body.WriteAsync(bodyData).ConfigureAwait(false);
             }
@@ -290,45 +290,6 @@ public static EntryRecordMode GetRecordMode(HttpRequest request)
             return mode;
         }
 
-        private byte[] CompressBody(byte[] incomingBody, SortedDictionary<string, string[]> headers)
-        {
-            if (headers.TryGetValue("Content-Encoding", out var values))
-            {
-                if (values.Contains("gzip"))
-                {
-                    using (var uncompressedStream = new MemoryStream(incomingBody))
-                    using (var resultStream = new MemoryStream())
-                    {
-                        using (var compressedStream = new GZipStream(resultStream, CompressionMode.Compress))
-                        {
-                            uncompressedStream.CopyTo(compressedStream);
-                        }
-                        return resultStream.ToArray();
-                    }
-                }
-            }
-
-            return incomingBody;
-        }
-
-        private byte[] DecompressBody(MemoryStream incomingBody, HttpContentHeaders headers)
-        {
-            if (headers.TryGetValues("Content-Encoding", out var values))
-            {
-                if (values.Contains("gzip"))
-                {
-                    using (var uncompressedStream = new GZipStream(incomingBody, CompressionMode.Decompress))
-                    using (var resultStream = new MemoryStream())
-                    {
-                        uncompressedStream.CopyTo(resultStream);
-                        return resultStream.ToArray();
-                    }
-                }
-            }
-
-            return incomingBody.ToArray();
-        }
-
         public HttpRequestMessage CreateUpstreamRequest(HttpRequest incomingRequest, byte[] incomingBody)
         {
             var upstreamRequest = new HttpRequestMessage();
@@ -484,7 +445,7 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
 
             if (match.Response.Body?.Length > 0)
             {
-                var bodyData = CompressBody(match.Response.Body, match.Response.Headers);
+                var bodyData = GZipUtilities.CompressBody(match.Response.Body, match.Response.Headers);
 
                 outgoingResponse.ContentLength = bodyData.Length;
 
@@ -506,7 +467,9 @@ public static async Task<RecordEntry> CreateEntryAsync(HttpRequest request)
                 }
             }
 
-            entry.Request.Body = await ReadAllBytes(request.Body).ConfigureAwait(false);
+            byte[] bytes = await ReadAllBytes(request.Body).ConfigureAwait(false);
+
+            entry.Request.Body = GZipUtilities.DecompressBody(bytes, request.Headers);
             return entry;
         }