From 85ed1f4517090c6acb35ad2db8403c104d0ad3c4 Mon Sep 17 00:00:00 2001
From: Jacob Hoffman-Andrews <github@hoffman-andrews.com>
Date: Tue, 21 Jun 2022 23:52:48 -0700
Subject: [PATCH] Make Cache-Control configurable.

---
 src/config.rs      | 11 +++++++++++
 src/web/rustdoc.rs | 28 ++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index ecdfaf40e..631d6672b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -57,6 +57,12 @@ pub struct Config {
     // Content Security Policy
     pub(crate) csp_report_only: bool,
 
+    // Cache-Control header
+    // If both are absent, don't generate the header. If only one is present,
+    // generate just that directive. Values are in seconds.
+    pub(crate) cache_control_stale_while_revalidate: Option<u32>,
+    pub(crate) cache_control_max_age: Option<u32>,
+
     // Build params
     pub(crate) build_attempts: u16,
     pub(crate) rustwide_workspace: PathBuf,
@@ -130,6 +136,11 @@ impl Config {
 
             csp_report_only: env("DOCSRS_CSP_REPORT_ONLY", false)?,
 
+            cache_control_stale_while_revalidate: maybe_env(
+                "CACHE_CONTROL_STALE_WHILE_REVALIDATE",
+            )?,
+            cache_control_max_age: maybe_env("CACHE_CONTROL_MAX_AGE")?,
+
             local_archive_cache_path: env(
                 "DOCSRS_ARCHIVE_INDEX_CACHE_PATH",
                 prefix.join("archive_cache"),
diff --git a/src/web/rustdoc.rs b/src/web/rustdoc.rs
index 4006f9b88..12d824984 100644
--- a/src/web/rustdoc.rs
+++ b/src/web/rustdoc.rs
@@ -235,13 +235,13 @@ impl RustdocPage {
         let is_latest_url = self.is_latest_url;
         // Build the page of documentation
         let ctx = ctry!(req, tera::Context::from_serialize(self));
+        let config = extension!(req, Config);
         // Extract the head and body of the rustdoc file so that we can insert it into our own html
         // while logging OOM errors from html rewriting
         let html = match utils::rewrite_lol(rustdoc_html, max_parse_memory, ctx, templates) {
             Err(RewritingError::MemoryLimitExceeded(..)) => {
                 metrics.html_rewrite_ooms.inc();
 
-                let config = extension!(req, Config);
                 let err = anyhow!(
                     "Failed to serve the rustdoc file '{}' because rewriting it surpassed the memory limit of {} bytes",
                     file_path, config.max_parse_memory,
@@ -259,13 +259,21 @@ impl RustdocPage {
                 .headers
                 .set(CacheControl(vec![CacheDirective::MaxAge(0)]));
         } else {
-            response.headers.set(CacheControl(vec![
-                CacheDirective::Extension(
+            let mut directives = vec![];
+            if let Some(seconds) = config.cache_control_stale_while_revalidate {
+                directives.push(CacheDirective::Extension(
                     "stale-while-revalidate".to_string(),
-                    Some("2592000".to_string()), // sixty days
-                ),
-                CacheDirective::MaxAge(600u32), // ten minutes
-            ]));
+                    Some(format!("{}", seconds)),
+                ));
+            }
+
+            if let Some(seconds) = config.cache_control_max_age {
+                directives.push(CacheDirective::MaxAge(seconds));
+            }
+
+            if !directives.is_empty() {
+                response.headers.set(CacheControl(directives));
+            }
         }
         Ok(response)
     }
@@ -864,13 +872,17 @@ mod test {
     #[test]
     fn cache_headers() {
         wrapper(|env| {
+            env.override_config(|config| {
+                config.cache_control_max_age = Some(600);
+                config.cache_control_stale_while_revalidate = Some(2592000);
+            });
+
             env.fake_release()
                 .name("dummy")
                 .version("0.1.0")
                 .archive_storage(true)
                 .rustdoc_file("dummy/index.html")
                 .create()?;
-
             let resp = env.frontend().get("/dummy/latest/dummy/").send()?;
             assert_eq!(resp.headers().get("Cache-Control").unwrap(), &"max-age=0");