Fix IOutputCache empty-body bug (#1702): bump feature revision in InvokeFeatures.Set (#2364)

Author: GarrettBeatty

.autover/changes/a0faddb0-c508-415c-b436-d99d68eafc56.json

@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.AspNetCoreServer",
+      "Type": "Patch",
+      "ChangelogMessages": [
+        "Fix InvokeFeatures.Set<TFeature> to bump the feature collection revision so middleware that wraps the response body (e.g. OutputCache, ResponseCompression) is properly visible to ASP.NET Core's FeatureReferences cache. Resolves https://github.com/aws/aws-lambda-dotnet/issues/1702 where IOutputCache stored empty response bodies."
+      ]
+    }
+  ]
+}

Libraries/src/Amazon.Lambda.AspNetCoreServer/Internal/InvokeFeatures.cs

@@ -117,10 +117,11 @@ public IEnumerator<KeyValuePair<Type, object>> GetEnumerator()
 
         public void Set<TFeature>(TFeature instance)
         {
-            if (instance == null)
-                return;
-
-            _features[typeof(TFeature)] = instance;
+            // Delegate to the indexer so _containerRevision is bumped, otherwise
+            // ASP.NET Core's FeatureReferences cache will return stale feature
+            // references after middleware (e.g. OutputCache, ResponseCompression)
+            // wraps the response body via HttpContext.Response.Body = wrapper.
+            this[typeof(TFeature)] = instance;
         }
 
         IEnumerator IEnumerable.GetEnumerator()

Libraries/test/Amazon.Lambda.AspNetCoreServer.Test/UtilitiesTest.cs

@@ -29,5 +29,57 @@ public void EnsureStatusCodeStartsAtIs200()
             var feature = new InvokeFeatures() as IHttpResponseFeature;
             Assert.Equal(200, feature.StatusCode);
         }
+
+        // Regression test for https://github.com/aws/aws-lambda-dotnet/issues/1702.
+        // ASP.NET Core's FeatureReferences cache uses Revision to detect when a
+        // feature has been swapped (e.g. OutputCache/ResponseCompression replacing
+        // IHttpResponseBodyFeature to wrap the response body). If Set<TFeature>
+        // does not bump the revision, cached references stay stale and writes
+        // bypass the wrapper.
+        [Fact]
+        public void SetFeatureBumpsRevision()
+        {
+            IFeatureCollection features = new InvokeFeatures();
+            var initialRevision = features.Revision;
+
+            features.Set<IHttpResponseBodyFeature>(new TestResponseBodyFeature());
+
+            Assert.NotEqual(initialRevision, features.Revision);
+        }
+
+        [Fact]
+        public void SetFeatureStoresAndRetrievesInstance()
+        {
+            IFeatureCollection features = new InvokeFeatures();
+            var replacement = new TestResponseBodyFeature();
+
+            features.Set<IHttpResponseBodyFeature>(replacement);
+
+            Assert.Same(replacement, features.Get<IHttpResponseBodyFeature>());
+        }
+
+        [Fact]
+        public void SetFeatureNullRemovesEntryAndBumpsRevision()
+        {
+            IFeatureCollection features = new InvokeFeatures();
+            // InvokeFeatures seeds itself as the IHttpResponseBodyFeature in its constructor.
+            Assert.NotNull(features.Get<IHttpResponseBodyFeature>());
+            var revisionBeforeRemove = features.Revision;
+
+            features.Set<IHttpResponseBodyFeature>(null);
+
+            Assert.Null(features.Get<IHttpResponseBodyFeature>());
+            Assert.NotEqual(revisionBeforeRemove, features.Revision);
+        }
+
+        private sealed class TestResponseBodyFeature : IHttpResponseBodyFeature
+        {
+            public System.IO.Stream Stream => System.IO.Stream.Null;
+            public System.IO.Pipelines.PipeWriter Writer => System.IO.Pipelines.PipeWriter.Create(System.IO.Stream.Null);
+            public System.Threading.Tasks.Task CompleteAsync() => System.Threading.Tasks.Task.CompletedTask;
+            public void DisableBuffering() { }
+            public System.Threading.Tasks.Task SendFileAsync(string path, long offset, long? count, System.Threading.CancellationToken cancellationToken) => System.Threading.Tasks.Task.CompletedTask;
+            public System.Threading.Tasks.Task StartAsync(System.Threading.CancellationToken cancellationToken) => System.Threading.Tasks.Task.CompletedTask;
+        }
     }
 }