Only set Cache-Control header on status 200

ui/ui.go

@@ -72,10 +72,7 @@ func Build(config *Config, urls common.URLMap) (http.Handler, *common.Renderer,
 	})))
 
 	// Add cache-friendly headers.
-	h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		w.Header().Set("Cache-Control", "max-age=86400")
-		mux.ServeHTTP(w, req)
-	})
+	h := withCacheHeaders(mux)
 
 	return h, renderer, nil
 }
@@ -134,3 +131,21 @@ func toJSON(obj interface{}) string {
 	data, _ := json.Marshal(obj)
 	return string(data)
 }
+
+type cachedResponseWriter struct {
+	http.ResponseWriter
+}
+
+func (c *cachedResponseWriter) WriteHeader(statusCode int) {
+	if statusCode == http.StatusOK {
+		c.ResponseWriter.Header().Set("Cache-Control", "max-age=86400")
+	}
+	c.ResponseWriter.WriteHeader(statusCode)
+}
+
+func withCacheHeaders(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		cw := &cachedResponseWriter{ResponseWriter: w}
+		h.ServeHTTP(cw, req)
+	})
+}