Call Flush With fMoreData=false when response isn't allowed to have body (#4403)

src/IISIntegration/IISIntegration.NoV1.sln

@@ -38,7 +38,6 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "build", "build", "{7E80C58E
 		build\applicationhost.iis.config = build\applicationhost.iis.config
 		build\Build.Settings = build\Build.Settings
 		build\Config.Definitions.Props = build\Config.Definitions.Props
-		build\dependencies.props = build\dependencies.props
 		build\functional-test-assets.targets = build\functional-test-assets.targets
 		build\Key.snk = build\Key.snk
 		build\launchSettings.json = build\launchSettings.json
@@ -116,7 +115,7 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "IIS.BackwardsCompatibility.
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "IIS.ForwardsCompatibility.FunctionalTests", "test\IIS.ForwardsCompatibility.FunctionalTests\IIS.ForwardsCompatibility.FunctionalTests.csproj", "{D1EA5D99-28FD-4197-81DE-17098846B38B}"
 EndProject
-Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "InProcessForwardsCompatWebSite", "test\WebSites\InProcessForwardsCompatWebSite\InProcessWebSite.csproj", "{BBBC85B2-5D7A-4D09-90B1-8DBCC9059493}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "InProcessWebSite", "test\WebSites\InProcessForwardsCompatWebSite\InProcessWebSite.csproj", "{BBBC85B2-5D7A-4D09-90B1-8DBCC9059493}"
 EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution

src/IISIntegration/src/AspNetCoreModuleV2/InProcessRequestHandler/managedexports.cpp

@@ -288,13 +288,13 @@ EXTERN_C __MIDL_DECLSPEC_DLLEXPORT
 HRESULT
 http_flush_response_bytes(
     _In_ IN_PROCESS_HANDLER* pInProcessHandler,
+    _In_ BOOL fMoreData,
     _Out_ BOOL* pfCompletionExpected
 )
 {
     IHttpResponse *pHttpResponse = (IHttpResponse*)pInProcessHandler->QueryHttpContext()->GetResponse();
 
     BOOL fAsync = TRUE;
-    BOOL fMoreData = TRUE;
     DWORD dwBytesSent = 0;
 
     HRESULT hr = pHttpResponse->Flush(

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.IO.cs

@@ -158,7 +158,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
                         if (flush)
                         {
-                            await AsyncIO.FlushAsync();
+                            await AsyncIO.FlushAsync(moreData: true);
                             flush = false;
                         }
                     }

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.cs

@@ -243,12 +243,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             SetResponseHeaders();
 
             EnsureIOInitialized();
-
+            var canHaveNonEmptyBody = StatusCodeCanHaveBody();
             if (flushHeaders)
             {
                 try
                 {
-                    await AsyncIO.FlushAsync();
+                    await AsyncIO.FlushAsync(canHaveNonEmptyBody);
                 }
                 // Client might be disconnected at this point
                 // don't leak the exception
@@ -258,7 +258,20 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 }
             }
 
-            _writeBodyTask = WriteBody(!flushHeaders);
+            if (!canHaveNonEmptyBody)
+            {
+                _bodyOutput.Dispose();
+            }
+            else
+            {
+                _writeBodyTask = WriteBody(!flushHeaders);
+            }
+        }
+
+        private bool StatusCodeCanHaveBody()
+        {
+            return StatusCode != 204
+                && StatusCode != 304;
         }
 
         private void InitializeRequestIO()

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.Flush.cs

@@ -12,21 +12,23 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             private readonly AsyncIOEngine _engine;
 
             private IntPtr _requestHandler;
+            private bool _moreData;
 
             public AsyncFlushOperation(AsyncIOEngine engine)
             {
                 _engine = engine;
             }
 
-            public void Initialize(IntPtr requestHandler)
+            public void Initialize(IntPtr requestHandler, bool moreData)
             {
                 _requestHandler = requestHandler;
+                _moreData = moreData;
             }
 
             protected override bool InvokeOperation(out int hr, out int bytes)
             {
                 bytes = 0;
-                hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, out var fCompletionExpected);
+                hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, _moreData, out var fCompletionExpected);
 
                 return !fCompletionExpected;
             }
@@ -36,6 +38,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 base.ResetOperation();
 
                 _requestHandler = default;
+                _moreData = false;
                 _engine.ReturnOperation(this);
             }
         }

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.cs

@@ -90,10 +90,10 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
         }
 
 
-        public ValueTask FlushAsync()
+        public ValueTask FlushAsync(bool moreData)
         {
             var flush = GetFlushOperation();
-            flush.Initialize(_handler);
+            flush.Initialize(_handler, moreData);
             Run(flush);
             return new ValueTask(flush, 0);
         }

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/IAsyncIOEngine.cs

@@ -11,7 +11,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
     {
         ValueTask<int> ReadAsync(Memory<byte> memory);
         ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);
-        ValueTask FlushAsync();
+        ValueTask FlushAsync(bool moreData);
         void NotifyCompletion(int hr, int bytes);
     }
 }

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.Initialize.cs

@@ -25,7 +25,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
 
             protected override bool InvokeOperation(out int hr, out int bytes)
             {
-                hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, out var completionExpected);
+                hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, fMoreData: true, out var completionExpected);
                 bytes = 0;
                 return !completionExpected;
             }

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.cs

@@ -56,7 +56,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             }
         }
 
-        public ValueTask FlushAsync()
+        public ValueTask FlushAsync(bool moreData)
         {
             lock (_contextLock)
             {

src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/NativeMethods.cs

@@ -67,7 +67,7 @@ namespace Microsoft.AspNetCore.Server.IIS
         private static extern unsafe int http_write_response_bytes(IntPtr pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_flush_response_bytes(IntPtr pInProcessHandler, out bool fCompletionExpected);
+        private static extern int http_flush_response_bytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern unsafe HttpApiTypes.HTTP_REQUEST_V2* http_get_raw_request(IntPtr pInProcessHandler);
@@ -171,9 +171,9 @@ namespace Microsoft.AspNetCore.Server.IIS
             return http_write_response_bytes(pInProcessHandler, pDataChunks, nChunks, out fCompletionExpected);
         }
 
-        public static int HttpFlushResponseBytes(IntPtr pInProcessHandler, out bool fCompletionExpected)
+        public static int HttpFlushResponseBytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected)
         {
-            return http_flush_response_bytes(pInProcessHandler, out fCompletionExpected);
+            return http_flush_response_bytes(pInProcessHandler, fMoreData, out fCompletionExpected);
         }
 
         public static unsafe HttpApiTypes.HTTP_REQUEST_V2* HttpGetRawRequest(IntPtr pInProcessHandler)

src/IISIntegration/test/Common.FunctionalTests/Inprocess/ResponseHeaderTests.cs

@@ -1,8 +1,10 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
 using System.Linq;
 using System.Net;
+using System.Net.Http;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Testing.xunit;
 using Microsoft.Net.Http.Headers;
@@ -75,6 +77,18 @@ namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
             Assert.Equal(body ?? string.Empty, await response.Content.ReadAsStringAsync());
         }
 
+        [ConditionalTheory]
+        [RequiresNewHandler]
+        [InlineData(204, "GET")]
+        [InlineData(304, "GET")]
+        public async Task TransferEncodingNotSetForStatusCodes(int code, string method)
+        {
+            var request = new HttpRequestMessage(new HttpMethod(method), _fixture.Client.BaseAddress + $"SetCustomErorCode?code={code}");
+            var response = await _fixture.Client.SendAsync(request);
+            Assert.Equal((HttpStatusCode)code, response.StatusCode);
+            Assert.DoesNotContain(response.Headers, h => h.Key.Equals("transfer-encoding", StringComparison.InvariantCultureIgnoreCase));
+        }
+
         [ConditionalFact]
         public async Task ServerHeaderIsOverriden()
         {