From 2db0327dc1cc11feb3b4135f723e9b6cda704c80 Mon Sep 17 00:00:00 2001
From: Richard <richard.scothern@gmail.com>
Date: Tue, 12 May 2015 17:49:18 -0700
Subject: [PATCH] Set cache headers for layers.

     - Set an Etag header
     - Check If-None-Match and respond appropriately
     - Set a Cache-Control header with a default of 1 week

Signed-off-by: Richard Scothern <richard.scothern@gmail.com>
---
 docs/handlers/api_test.go   | 37 +++++++++++++++++++++++++++++++++++++
 docs/storage/layerreader.go | 25 +++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/docs/handlers/api_test.go b/docs/handlers/api_test.go
index 1e31477f..6dc7a422 100644
--- a/docs/handlers/api_test.go
+++ b/docs/handlers/api_test.go
@@ -263,6 +263,43 @@ func TestLayerAPI(t *testing.T) {
 
 	checkResponse(t, "fetching layer bad digest", resp, http.StatusBadRequest)
 
+	// Cache headers
+	resp, err = http.Get(layerURL)
+	if err != nil {
+		t.Fatalf("unexpected error fetching layer: %v", err)
+	}
+
+	checkResponse(t, "fetching layer", resp, http.StatusOK)
+	checkHeaders(t, resp, http.Header{
+		"Content-Length":        []string{fmt.Sprint(layerLength)},
+		"Docker-Content-Digest": []string{layerDigest.String()},
+		"ETag":                  []string{layerDigest.String()},
+		"Cache-Control":         []string{"max-age=86400"},
+	})
+
+	// Matching etag, gives 304
+	etag := resp.Header.Get("Etag")
+	req, err = http.NewRequest("GET", layerURL, nil)
+	if err != nil {
+		t.Fatalf("Error constructing request: %s", err)
+	}
+	req.Header.Set("If-None-Match", etag)
+	resp, err = http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("Error constructing request: %s", err)
+	}
+
+	checkResponse(t, "fetching layer with etag", resp, http.StatusNotModified)
+
+	// Non-matching etag, gives 200
+	req, err = http.NewRequest("GET", layerURL, nil)
+	if err != nil {
+		t.Fatalf("Error constructing request: %s", err)
+	}
+	req.Header.Set("If-None-Match", "")
+	resp, err = http.DefaultClient.Do(req)
+	checkResponse(t, "fetching layer with invalid etag", resp, http.StatusOK)
+
 	// Missing tests:
 	// 	- Upload the same tarsum file under and different repository and
 	//       ensure the content remains uncorrupted.
diff --git a/docs/storage/layerreader.go b/docs/storage/layerreader.go
index ddca9741..044dab09 100644
--- a/docs/storage/layerreader.go
+++ b/docs/storage/layerreader.go
@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"fmt"
 	"net/http"
 	"time"
 
@@ -73,7 +74,31 @@ func (lr *layerReader) Handler(r *http.Request) (h http.Handler, err error) {
 	}
 
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// If the registry is serving this content itself, check
+		// the If-None-Match header and return 304 on match.  Redirected
+		// storage implementations do the same.
+
+		if etagMatch(r, lr.digest.String()) {
+			w.WriteHeader(http.StatusNotModified)
+			return
+		}
+		setCacheHeaders(w, 86400, lr.digest.String())
 		w.Header().Set("Docker-Content-Digest", lr.digest.String())
 		handlerFunc.ServeHTTP(w, r)
 	}), nil
 }
+
+func etagMatch(r *http.Request, etag string) bool {
+	for _, headerVal := range r.Header["If-None-Match"] {
+		if headerVal == etag {
+			return true
+		}
+	}
+	return false
+}
+
+func setCacheHeaders(w http.ResponseWriter, cacheAge int, etag string) {
+	w.Header().Set("ETag", etag)
+	w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", cacheAge))
+
+}