From feb3a1564d23800f9c7aab08ac75d02b5e1d7e99 Mon Sep 17 00:00:00 2001
From: Ross Light <ross@zombiezen.com>
Date: Fri, 29 Sep 2017 09:23:29 -0700
Subject: [PATCH] add cache headers (#14)

Fixes #13
---
 README.md       |  6 ++++++
 handler.go      | 20 ++++++++++++++++----
 handler_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index d3c466b..ebb0ed6 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,7 @@ handler into larger Go servers.
 
 ```
 host: example.com
+cache_age: 3600
 paths:
   /foo:
     repo: https://github.com/example/foo
@@ -75,6 +76,11 @@ paths:
     </tr>
   </thead>
   <tbody>
+    <tr>
+      <th scope="row"><code>cache_age</code></th>
+      <td>optional</td>
+      <td>The amount of time to cache package pages in seconds.  Controls the <code>max-age</code> directive sent in the <a href="https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Cache-Control"><code>Cache-Control</code></a> HTTP header.</td>
+    </tr>
     <tr>
       <th scope="row"><code>host</code></th>
       <td>optional</td>
diff --git a/handler.go b/handler.go
index 6f5a3f9..897744c 100644
--- a/handler.go
+++ b/handler.go
@@ -16,6 +16,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"html/template"
 	"net/http"
@@ -26,8 +27,9 @@ import (
 )
 
 type handler struct {
-	host  string
-	paths pathConfigSet
+	host         string
+	cacheControl string
+	paths        pathConfigSet
 }
 
 type pathConfig struct {
@@ -39,8 +41,9 @@ type pathConfig struct {
 
 func newHandler(config []byte) (*handler, error) {
 	var parsed struct {
-		Host  string `yaml:"host,omitempty"`
-		Paths map[string]struct {
+		Host     string `yaml:"host,omitempty"`
+		CacheAge *int64 `yaml:"cache_max_age,omitempty"`
+		Paths    map[string]struct {
 			Repo    string `yaml:"repo,omitempty"`
 			Display string `yaml:"display,omitempty"`
 			VCS     string `yaml:"vcs,omitempty"`
@@ -50,6 +53,14 @@ func newHandler(config []byte) (*handler, error) {
 		return nil, err
 	}
 	h := &handler{host: parsed.Host}
+	cacheAge := int64(86400) // 24 hpurs (in seconds)
+	if parsed.CacheAge != nil {
+		cacheAge = *parsed.CacheAge
+		if cacheAge < 0 {
+			return nil, errors.New("cache_max_age is negative")
+		}
+	}
+	h.cacheControl = fmt.Sprintf("public, max-age=%d", cacheAge)
 	for path, e := range parsed.Paths {
 		pc := pathConfig{
 			path:    strings.TrimSuffix(path, "/"),
@@ -94,6 +105,7 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	w.Header().Set("Cache-Control", h.cacheControl)
 	if err := vanityTmpl.Execute(w, struct {
 		Import  string
 		Subpath string
diff --git a/handler_test.go b/handler_test.go
index 9b41c8a..9c41bf3 100644
--- a/handler_test.go
+++ b/handler_test.go
@@ -139,6 +139,10 @@ func TestBadConfigs(t *testing.T) {
 			"  /unknownvcs:\n" +
 			"    repo: https://bitbucket.org/zombiezen/gopdf\n" +
 			"    vcs: xyzzy\n",
+		"cache_max_age: -1\n" +
+			"paths:\n" +
+			"  /portmidi:\n" +
+			"    repo: https://github.com/rakyll/portmidi\n",
 	}
 	for _, config := range badConfigs {
 		_, err := newHandler([]byte(config))
@@ -227,3 +231,45 @@ func TestPathConfigSetFind(t *testing.T) {
 		}
 	}
 }
+
+func TestCacheHeader(t *testing.T) {
+	tests := []struct {
+		name         string
+		config       string
+		cacheControl string
+	}{
+		{
+			name:         "default",
+			cacheControl: "public, max-age=86400",
+		},
+		{
+			name:         "specify time",
+			config:       "cache_max_age: 60\n",
+			cacheControl: "public, max-age=60",
+		},
+		{
+			name:         "zero",
+			config:       "cache_max_age: 0\n",
+			cacheControl: "public, max-age=0",
+		},
+	}
+	for _, test := range tests {
+		h, err := newHandler([]byte("paths:\n  /portmidi:\n    repo: https://github.com/rakyll/portmidi\n" +
+			test.config))
+		if err != nil {
+			t.Errorf("%s: newHandler: %v", test.name, err)
+			continue
+		}
+		s := httptest.NewServer(h)
+		resp, err := http.Get(s.URL + "/portmidi")
+		if err != nil {
+			t.Errorf("%s: http.Get: %v", test.name, err)
+			continue
+		}
+		resp.Body.Close()
+		got := resp.Header.Get("Cache-Control")
+		if got != test.cacheControl {
+			t.Errorf("%s: Cache-Control header = %q; want %q", test.name, got, test.cacheControl)
+		}
+	}
+}