From 204ff0a78504b3c688d906bab4e03e87d8ea1b76 Mon Sep 17 00:00:00 2001
From: Ryan Nowak <nowakra@gmail.com>
Date: Sat, 22 Sep 2018 18:38:38 -0700
Subject: [PATCH] Set cache headers in health check middleware

---
 build/dependencies.props                      |  1 +
 .../HealthCheckMiddleware.cs                  | 12 +++-
 .../HealthCheckOptions.cs                     |  8 +++
 ...AspNetCore.Diagnostics.HealthChecks.csproj |  1 +
 .../HealthCheckMiddlewareTests.cs             | 58 ++++++++++++++++++-
 5 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/build/dependencies.props b/build/dependencies.props
index cc58ccf920..7a3fa9d10c 100644
--- a/build/dependencies.props
+++ b/build/dependencies.props
@@ -36,6 +36,7 @@
     <MicrosoftNETCoreApp21PackageVersion>2.1.3</MicrosoftNETCoreApp21PackageVersion>
     <MicrosoftNETCoreApp22PackageVersion>2.2.0-preview2-26905-02</MicrosoftNETCoreApp22PackageVersion>
     <MicrosoftNETTestSdkPackageVersion>15.6.1</MicrosoftNETTestSdkPackageVersion>
+    <MicrosoftNetHttpHeadersPackageVersion>2.2.0-preview3-35252</MicrosoftNetHttpHeadersPackageVersion>
     <MoqPackageVersion>4.7.49</MoqPackageVersion>
     <NETStandardLibrary20PackageVersion>2.0.3</NETStandardLibrary20PackageVersion>
     <NewtonsoftJsonPackageVersion>11.0.2</NewtonsoftJsonPackageVersion>
diff --git a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckMiddleware.cs b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckMiddleware.cs
index e06ef8509b..2e1897bc2a 100644
--- a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckMiddleware.cs
+++ b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckMiddleware.cs
@@ -8,6 +8,7 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Options;
+using Microsoft.Net.Http.Headers;
 
 namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
 {
@@ -70,6 +71,15 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
 
             httpContext.Response.StatusCode = statusCode;
 
+            if (!_healthCheckOptions.SuppressCacheHeaders)
+            {
+                // Similar to: https://github.com/aspnet/Security/blob/7b6c9cf0eeb149f2142dedd55a17430e7831ea99/src/Microsoft.AspNetCore.Authentication.Cookies/CookieAuthenticationHandler.cs#L377-L379
+                var headers = httpContext.Response.Headers;
+                headers[HeaderNames.CacheControl] = "no-store, no-cache";
+                headers[HeaderNames.Pragma] = "no-cache";
+                headers[HeaderNames.Expires] = "Thu, 01 Jan 1970 00:00:00 GMT";
+            }
+
             if (_healthCheckOptions.ResponseWriter != null)
             {
                 await _healthCheckOptions.ResponseWriter(httpContext, result);
@@ -103,7 +113,7 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
 
             if (notFound.Count > 0)
             {
-                var message = 
+                var message =
                     $"The following health checks were not found: '{string.Join(", ", notFound)}'. " +
                     $"Registered health checks: '{string.Join(", ", checks.Keys)}'.";
                 throw new InvalidOperationException(message);
diff --git a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckOptions.cs b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckOptions.cs
index 42996e3701..d66a59e124 100644
--- a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckOptions.cs
+++ b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/HealthCheckOptions.cs
@@ -46,5 +46,13 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
         /// of <see cref="HealthReport.Status"/> as a string.
         /// </remarks>
         public Func<HttpContext, HealthReport, Task> ResponseWriter { get; set; } = HealthCheckResponseWriters.WriteMinimalPlaintext;
+
+        /// <summary>
+        /// Gets or sets a value that controls whether the health check middleware will add HTTP headers to prevent
+        /// response caching. If the value is <c>false</c> the health check middleware will set or override the 
+        /// <c>Cache-Control</c>, <c>Expires</c>, and <c>Pragma</c> headers to prevent response caching. If the value 
+        /// is <c>true</c> the health check middleware will not modify the cache headers of the response.
+        /// </summary>
+        public bool SuppressCacheHeaders { get; set; }
     }
 }
diff --git a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/Microsoft.AspNetCore.Diagnostics.HealthChecks.csproj b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/Microsoft.AspNetCore.Diagnostics.HealthChecks.csproj
index b7dd85d4a6..15bbc62f78 100644
--- a/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/Microsoft.AspNetCore.Diagnostics.HealthChecks.csproj
+++ b/src/Microsoft.AspNetCore.Diagnostics.HealthChecks/Microsoft.AspNetCore.Diagnostics.HealthChecks.csproj
@@ -18,6 +18,7 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.AspNetCore.Http.Abstractions" Version="$(MicrosoftAspNetCoreHttpAbstractionsPackageVersion)" />
     <PackageReference Include="Microsoft.Extensions.Options" Version="$(MicrosoftExtensionsOptionsPackageVersion)" />
+    <PackageReference Include="Microsoft.Net.Http.Headers" Version="$(MicrosoftNetHttpHeadersPackageVersion)" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/test/Microsoft.AspNetCore.Diagnostics.HealthChecks.Tests/HealthCheckMiddlewareTests.cs b/test/Microsoft.AspNetCore.Diagnostics.HealthChecks.Tests/HealthCheckMiddlewareTests.cs
index 7a2435948c..4715a36da9 100644
--- a/test/Microsoft.AspNetCore.Diagnostics.HealthChecks.Tests/HealthCheckMiddlewareTests.cs
+++ b/test/Microsoft.AspNetCore.Diagnostics.HealthChecks.Tests/HealthCheckMiddlewareTests.cs
@@ -1,15 +1,16 @@
 // 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.Net;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.TestHost;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Net.Http.Headers;
 using Newtonsoft.Json;
+using System.Net;
+using System.Threading.Tasks;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
@@ -291,6 +292,57 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
             Assert.Equal("Healthy", await response.Content.ReadAsStringAsync());
         }
 
+        [Fact]
+        public async Task SetsCacheHeaders()
+        {
+            var builder = new WebHostBuilder()
+                .Configure(app =>
+                {
+                    app.UseHealthChecks("/health");
+                })
+                .ConfigureServices(services =>
+                {
+                    services.AddHealthChecks();
+                });
+            var server = new TestServer(builder);
+            var client = server.CreateClient();
+
+            var response = await client.GetAsync("/health");
+
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Healthy", await response.Content.ReadAsStringAsync());
+            Assert.Equal("no-store, no-cache", response.Headers.CacheControl.ToString());
+            Assert.Equal("no-cache", response.Headers.Pragma.ToString());
+            Assert.Equal(new string[] { "Thu, 01 Jan 1970 00:00:00 GMT" }, response.Content.Headers.GetValues(HeaderNames.Expires));
+        }
+
+        [Fact]
+        public async Task CanSuppressCacheHeaders()
+        {
+            var builder = new WebHostBuilder()
+                .Configure(app =>
+                {
+                    app.UseHealthChecks("/health", new HealthCheckOptions()
+                    {
+                        SuppressCacheHeaders = true,
+                    });
+                })
+                .ConfigureServices(services =>
+                {
+                    services.AddHealthChecks();
+                });
+            var server = new TestServer(builder);
+            var client = server.CreateClient();
+
+            var response = await client.GetAsync("/health");
+
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Healthy", await response.Content.ReadAsStringAsync());
+            Assert.Null(response.Headers.CacheControl);
+            Assert.Empty(response.Headers.Pragma.ToString());
+            Assert.False(response.Content.Headers.Contains(HeaderNames.Expires));
+        }
+
         [Fact]
         public async Task CanFilterChecks()
         {
@@ -363,7 +415,7 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
                 {
                     services.AddHealthChecks();
                 });
-            
+
             var server = new TestServer(builder);
             var client = server.CreateClient();