From a639c0efba9053db88c6c80641524b0f57c62624 Mon Sep 17 00:00:00 2001
From: Matt Butcher <matt.butcher@fermyon.com>
Date: Tue, 1 Feb 2022 16:39:56 -0700
Subject: [PATCH] Added CACHE_CONTROL env vars to set cache control headers

Signed-off-by: Matt Butcher <matt.butcher@fermyon.com>
---
 README.md     | 12 ++++++++++++
 fileserver.gr | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index bd344cf..531a686 100644
--- a/README.md
+++ b/README.md
@@ -90,6 +90,18 @@ volumes = {"/" = "/path/to/fileserver"}
 
 The above configures Wagi to map the path `/static/...` to the `fileserver.gr.wasm` module. Then it serves all of the files in this project.
 
+### Environment Variables
+
+The following environment variables can be passed via Wagi's `-e` flag:
+
+- `CACHE_CONTROL`: The string value of a cache-control header. If not specified, this will set cache-control to `no-cache`. Google recommends setting this value to `CACHE_CONTROL="max-age=31536000"` (cache for up to 1 year).
+- Type-specific cache controls can be set using the following env vars (all of which default to `CACHE_CONTROL` if not set):
+    - `CSS_CACHE_CONTROL`
+    - `FONT_CACHE_CONTROL`
+    - `IMAGE_CACHE_CONTROL`
+    - `JS_CACHE_CONTROL`
+
+
 ### Testing the Static Fileserver with `curl`
 
 This step is the same whether you use Bindle or a `modules.toml`.
diff --git a/fileserver.gr b/fileserver.gr
index ff96a14..1de7b2b 100644
--- a/fileserver.gr
+++ b/fileserver.gr
@@ -45,6 +45,43 @@ let rec pipe = (in, out) => {
   }
 }
 
+// Determine cache control values.
+let cache_duration = (env, mtype) => {
+  let default_cache_control = match (Map.get("CACHE_CONTROL", env)) {
+    Some(val) => val,
+    None => "no-cache",
+  }
+
+  if (String.indexOf("image/", mtype) == Some(0)) {
+    match (Map.get("IMAGE_CACHE_CONTROL", env)) {
+      Some(val) => val,
+      None => default_cache_control,
+    }
+  } else if (String.indexOf("font/", mtype) == Some(0)) {
+    match (Map.get("FONT_CACHE_CONTROL", env)) {
+      Some(val) => val,
+      None => default_cache_control,
+    }
+  } else if (String.indexOf("text/css", mtype) == Some(0)) {
+    match (Map.get("CSS_CACHE_CONTROL", env)) {
+      Some(val) => val,
+      None => default_cache_control,
+    }
+  } else if (String.indexOf("text/javascript", mtype) == Some(0)) {
+    match (Map.get("JS_CACHE_CONTROL", env)) {
+      Some(val) => val,
+      None => default_cache_control,
+    }
+  } else {
+    default_cache_control
+  }
+}
+
+let headers = (env, path) => {
+  let mtype = Mediatype.guess(path);
+  "Content-Type: " ++ mtype ++ "\nCache-Control: " ++ cache_duration(env, mtype) ++ "\n\n"
+}
+
 let serve = (abs_path, env) => {
   // If PATH_PREFIX is set, then the path prefix is prepended onto the incoming path.
   // This allows you to map to a directory that does not match the directory name in the URL.
@@ -70,7 +107,7 @@ let serve = (abs_path, env) => {
         "Unexpected error when writing Content-Type",
         File.fdWrite(
           File.stdout,
-          "Content-Type: " ++ Mediatype.guess(path) ++ "\n\n",
+          headers(env, path),
         ),
       )
 
@@ -84,6 +121,8 @@ let serve = (abs_path, env) => {
   }
 }
 
+
+
 let guestpath = env => {
   // Backward compat for an older version of Wagi that had PATH_INFO wrong.
   // X_RELATIVE_PATH was removed before Wagi 0.4