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

jkotalik · web-flow · commit 0646bb98b0fe · 2018-12-05T10:11:36.000-08:00
diff --git a/src/IISIntegration/IISIntegration.NoV1.sln b/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
diff --git a/src/IISIntegration/src/AspNetCoreModuleV2/InProcessRequestHandler/managedexports.cpp b/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(
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.IO.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.IO.cs
@@ -158,7 +158,7 @@ private async Task WriteBody(bool flush = false)
 
                         if (flush)
                         {
-                            await AsyncIO.FlushAsync();
+                            await AsyncIO.FlushAsync(moreData: true);
                             flush = false;
                         }
                     }
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.cs
@@ -243,12 +243,12 @@ private async Task ProduceStart(bool flushHeaders)
             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 @@ private async Task ProduceStart(bool flushHeaders)
                 }
             }
 
-            _writeBodyTask = WriteBody(!flushHeaders);
+            if (!canHaveNonEmptyBody)
+            {
+                _bodyOutput.Dispose();
+            }
+            else
+            {
+                _writeBodyTask = WriteBody(!flushHeaders);
+            }
+        }
+
+        private bool StatusCodeCanHaveBody()
+        {
+            return StatusCode != 204
+                && StatusCode != 304;
         }
 
         private void InitializeRequestIO()
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.Flush.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.Flush.cs
@@ -12,21 +12,23 @@ internal class AsyncFlushOperation : AsyncIOOperation
             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 @@ protected override void ResetOperation()
                 base.ResetOperation();
 
                 _requestHandler = default;
+                _moreData = false;
                 _engine.ReturnOperation(this);
             }
         }
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/AsyncIOEngine.cs
@@ -90,10 +90,10 @@ private void Run(AsyncIOOperation ioOperation)
         }
 
 
-        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);
         }
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/IAsyncIOEngine.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/IAsyncIOEngine.cs
@@ -11,7 +11,7 @@ internal interface IAsyncIOEngine: IDisposable
     {
         ValueTask<int> ReadAsync(Memory<byte> memory);
         ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);
-        ValueTask FlushAsync();
+        ValueTask FlushAsync(bool moreData);
         void NotifyCompletion(int hr, int bytes);
     }
 }
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.Initialize.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.Initialize.cs
@@ -25,7 +25,7 @@ public void Initialize(IntPtr requestHandler)
 
             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;
             }
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/Core/IO/WebSocketsAsyncIOEngine.cs
@@ -56,7 +56,7 @@ public ValueTask<int> WriteAsync(ReadOnlySequence<byte> data)
             }
         }
 
-        public ValueTask FlushAsync()
+        public ValueTask FlushAsync(bool moreData)
         {
             lock (_contextLock)
             {
diff --git a/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/NativeMethods.cs b/src/IISIntegration/src/Microsoft.AspNetCore.Server.IIS/NativeMethods.cs
@@ -67,7 +67,7 @@ private static extern int register_callbacks(IntPtr pInProcessApplication,
         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 @@ public static unsafe int HttpWriteResponseBytes(IntPtr pInProcessHandler, HttpAp
             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)
diff --git a/src/IISIntegration/test/Common.FunctionalTests/Inprocess/ResponseHeaderTests.cs b/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 @@ public async Task CustomErrorCodeWorks(int code, string reason, string expectedR
             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()
         {

Original file line number	Diff line number	Diff line change
`@@ -288,13 +288,13 @@ EXTERN_C __MIDL_DECLSPEC_DLLEXPORT`
`288`	`288`	`HRESULT`
`289`	`289`	`http_flush_response_bytes(`
`290`	`290`	`_In_ IN_PROCESS_HANDLER* pInProcessHandler,`
	`291`	`+ _In_ BOOL fMoreData,`
`291`	`292`	`_Out_ BOOL* pfCompletionExpected`
`292`	`293`	`)`
`293`	`294`	`{`
`294`	`295`	`IHttpResponse pHttpResponse = (IHttpResponse)pInProcessHandler->QueryHttpContext()->GetResponse();`
`295`	`296`
`296`	`297`	`BOOL fAsync = TRUE;`
`297`		`- BOOL fMoreData = TRUE;`
`298`	`298`	`DWORD dwBytesSent = 0;`
`299`	`299`
`300`	`300`	`HRESULT hr = pHttpResponse->Flush(`
Original file line number	Diff line number	Diff line change
`@@ -158,7 +158,7 @@ private async Task WriteBody(bool flush = false)`
`158`	`158`
`159`	`159`	`if (flush)`
`160`	`160`	`{`
`161`		`- await AsyncIO.FlushAsync();`
	`161`	`+ await AsyncIO.FlushAsync(moreData: true);`
`162`	`162`	`flush = false;`
`163`	`163`	`}`
`164`	`164`	`}`
Original file line number	Diff line number	Diff line change
`@@ -243,12 +243,12 @@ private async Task ProduceStart(bool flushHeaders)`
`243`	`243`	`SetResponseHeaders();`
`244`	`244`
`245`	`245`	`EnsureIOInitialized();`
`246`		`-`
	`246`	`+ var canHaveNonEmptyBody = StatusCodeCanHaveBody();`
`247`	`247`	`if (flushHeaders)`
`248`	`248`	`{`
`249`	`249`	`try`
`250`	`250`	`{`
`251`		`- await AsyncIO.FlushAsync();`
	`251`	`+ await AsyncIO.FlushAsync(canHaveNonEmptyBody);`
`252`	`252`	`}`
`253`	`253`	`// Client might be disconnected at this point`
`254`	`254`	`// don't leak the exception`
`@@ -258,7 +258,20 @@ private async Task ProduceStart(bool flushHeaders)`
`258`	`258`	`}`
`259`	`259`	`}`
`260`	`260`
`261`		`- _writeBodyTask = WriteBody(!flushHeaders);`
	`261`	`+ if (!canHaveNonEmptyBody)`
	`262`	`+ {`
	`263`	`+ _bodyOutput.Dispose();`
	`264`	`+ }`
	`265`	`+ else`
	`266`	`+ {`
	`267`	`+ _writeBodyTask = WriteBody(!flushHeaders);`
	`268`	`+ }`
	`269`	`+ }`
	`270`	`+`
	`271`	`+ private bool StatusCodeCanHaveBody()`
	`272`	`+ {`
	`273`	`+ return StatusCode != 204`
	`274`	`+ && StatusCode != 304;`
`262`	`275`	`}`
`263`	`276`
`264`	`277`	`private void InitializeRequestIO()`
Original file line number	Diff line number	Diff line change
`@@ -12,21 +12,23 @@ internal class AsyncFlushOperation : AsyncIOOperation`
`12`	`12`	`private readonly AsyncIOEngine _engine;`
`13`	`13`
`14`	`14`	`private IntPtr _requestHandler;`
	`15`	`+ private bool _moreData;`
`15`	`16`
`16`	`17`	`public AsyncFlushOperation(AsyncIOEngine engine)`
`17`	`18`	`{`
`18`	`19`	`_engine = engine;`
`19`	`20`	`}`
`20`	`21`
`21`		`- public void Initialize(IntPtr requestHandler)`
	`22`	`+ public void Initialize(IntPtr requestHandler, bool moreData)`
`22`	`23`	`{`
`23`	`24`	`_requestHandler = requestHandler;`
	`25`	`+ _moreData = moreData;`
`24`	`26`	`}`
`25`	`27`
`26`	`28`	`protected override bool InvokeOperation(out int hr, out int bytes)`
`27`	`29`	`{`
`28`	`30`	`bytes = 0;`
`29`		`- hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, out var fCompletionExpected);`
	`31`	`+ hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, _moreData, out var fCompletionExpected);`
`30`	`32`
`31`	`33`	`return !fCompletionExpected;`
`32`	`34`	`}`
`@@ -36,6 +38,7 @@ protected override void ResetOperation()`
`36`	`38`	`base.ResetOperation();`
`37`	`39`
`38`	`40`	`_requestHandler = default;`
	`41`	`+ _moreData = false;`
`39`	`42`	`_engine.ReturnOperation(this);`
`40`	`43`	`}`
`41`	`44`	`}`
Original file line number	Diff line number	Diff line change
`@@ -90,10 +90,10 @@ private void Run(AsyncIOOperation ioOperation)`
`90`	`90`	`}`
`91`	`91`
`92`	`92`
`93`		`- public ValueTask FlushAsync()`
	`93`	`+ public ValueTask FlushAsync(bool moreData)`
`94`	`94`	`{`
`95`	`95`	`var flush = GetFlushOperation();`
`96`		`- flush.Initialize(_handler);`
	`96`	`+ flush.Initialize(_handler, moreData);`
`97`	`97`	`Run(flush);`
`98`	`98`	`return new ValueTask(flush, 0);`
`99`	`99`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@ internal interface IAsyncIOEngine: IDisposable`
`11`	`11`	`{`
`12`	`12`	`ValueTask<int> ReadAsync(Memory<byte> memory);`
`13`	`13`	`ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);`
`14`		`- ValueTask FlushAsync();`
	`14`	`+ ValueTask FlushAsync(bool moreData);`
`15`	`15`	`void NotifyCompletion(int hr, int bytes);`
`16`	`16`	`}`
`17`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ public void Initialize(IntPtr requestHandler)`
`25`	`25`
`26`	`26`	`protected override bool InvokeOperation(out int hr, out int bytes)`
`27`	`27`	`{`
`28`		`- hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, out var completionExpected);`
	`28`	`+ hr = NativeMethods.HttpFlushResponseBytes(_requestHandler, fMoreData: true, out var completionExpected);`
`29`	`29`	`bytes = 0;`
`30`	`30`	`return !completionExpected;`
`31`	`31`	`}`
Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ public ValueTask<int> WriteAsync(ReadOnlySequence<byte> data)`
`56`	`56`	`}`
`57`	`57`	`}`
`58`	`58`
`59`		`- public ValueTask FlushAsync()`
	`59`	`+ public ValueTask FlushAsync(bool moreData)`
`60`	`60`	`{`
`61`	`61`	`lock (_contextLock)`
`62`	`62`	`{`
Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ private static extern int register_callbacks(IntPtr pInProcessApplication,`
`67`	`67`	`private static extern unsafe int http_write_response_bytes(IntPtr pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected);`
`68`	`68`
`69`	`69`	`[DllImport(AspNetCoreModuleDll)]`
`70`		`- private static extern int http_flush_response_bytes(IntPtr pInProcessHandler, out bool fCompletionExpected);`
	`70`	`+ private static extern int http_flush_response_bytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected);`
`71`	`71`
`72`	`72`	`[DllImport(AspNetCoreModuleDll)]`
`73`	`73`	`private static extern unsafe HttpApiTypes.HTTP_REQUEST_V2* http_get_raw_request(IntPtr pInProcessHandler);`
`@@ -171,9 +171,9 @@ public static unsafe int HttpWriteResponseBytes(IntPtr pInProcessHandler, HttpAp`
`171`	`171`	`return http_write_response_bytes(pInProcessHandler, pDataChunks, nChunks, out fCompletionExpected);`
`172`	`172`	`}`
`173`	`173`
`174`		`- public static int HttpFlushResponseBytes(IntPtr pInProcessHandler, out bool fCompletionExpected)`
	`174`	`+ public static int HttpFlushResponseBytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected)`
`175`	`175`	`{`
`176`		`- return http_flush_response_bytes(pInProcessHandler, out fCompletionExpected);`
	`176`	`+ return http_flush_response_bytes(pInProcessHandler, fMoreData, out fCompletionExpected);`
`177`	`177`	`}`
`178`	`178`
`179`	`179`	`public static unsafe HttpApiTypes.HTTP_REQUEST_V2* HttpGetRawRequest(IntPtr pInProcessHandler)`