replace getHopByHop with getEndToEnd that does the right thing with headers on not-modified responses

Author: gregjones

httpcache.go

@@ -176,12 +176,9 @@ func (t *Transport) RoundTrip(req *http.Request) (resp *http.Response, err error
 		resp, err = transport.RoundTrip(req)
 		if err == nil && req.Method == "GET" && resp.StatusCode == http.StatusNotModified {
 			// Replace the 304 response with the one from cache, but update with some new headers
-			headersToMerge := getHopByHopHeaders(resp)
-			for _, headerKey := range headersToMerge {
-				key := http.CanonicalHeaderKey(headerKey)
-				if v, ok := resp.Header[key]; ok {
-					cachedResp.Header[key] = v
-				}
+			endToEndHeaders := getEndToEndHeaders(resp.Header)
+			for _, header := range endToEndHeaders {
+				cachedResp.Header[header] = resp.Header[header]
 			}
 			cachedResp.Status = fmt.Sprintf("%d %s", http.StatusOK, http.StatusText(http.StatusOK))
 			cachedResp.StatusCode = http.StatusOK
@@ -345,17 +342,32 @@ func getFreshness(respHeaders, reqHeaders http.Header) (freshness int) {
 	return stale
 }
 
-func getHopByHopHeaders(resp *http.Response) []string {
+func containsHeader(headers []string, header string) bool {
+	for _, v := range headers {
+		if http.CanonicalHeaderKey(v) == http.CanonicalHeaderKey(header) {
+			return true
+		}
+	}
+	return false
+}
+
+func getEndToEndHeaders(respHeaders http.Header) []string {
 	// These headers are always hop-by-hop
-	headers := []string{"connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailers", "transfer-encoding", "upgrade"}
+	hopByHopHeaders := []string{"connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailers", "transfer-encoding", "upgrade"}
 
-	for _, extra := range strings.Split(resp.Header.Get("connection"), ",") {
+	for _, extra := range strings.Split(respHeaders.Get("connection"), ",") {
 		// any header listed in connection, if present, is also considered hop-by-hop
 		if strings.Trim(extra, " ") != "" {
-			headers = append(headers, extra)
+			hopByHopHeaders = append(hopByHopHeaders, extra)
+		}
+	}
+	endToEndHeaders := []string{}
+	for respHeader, _ := range respHeaders {
+		if !containsHeader(hopByHopHeaders, respHeader) {
+			endToEndHeaders = append(endToEndHeaders, respHeader)
 		}
 	}
-	return headers
+	return endToEndHeaders
 }
 
 func canStore(reqCacheControl, respCacheControl cacheControl) (canStore bool) {

httpcache_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"testing"
 	"time"
 
@@ -89,6 +90,18 @@ func (s *S) SetUpSuite(c *C) {
 		w.Header().Set("Vary", "X-Madeup-Header")
 		w.Write([]byte("Some text content"))
 	}))
+
+	updateFieldsCounter := 0
+	mux.HandleFunc("/updatefields", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("X-Counter", strconv.Itoa(updateFieldsCounter))
+		w.Header().Set("Etag", `"e"`)
+		updateFieldsCounter++
+		if r.Header.Get("if-none-match") != "" {
+			w.WriteHeader(http.StatusNotModified)
+		} else {
+			w.Write([]byte("Some text content"))
+		}
+	}))
 }
 
 func (s *S) TearDownSuite(c *C) {
@@ -302,6 +315,22 @@ func (s *S) TestGetVaryUnused(c *C) {
 	c.Assert(resp2.Header.Get(XFromCache), Equals, "1")
 }
 
+func (s *S) TestUpdateFields(c *C) {
+	req, err := http.NewRequest("GET", s.server.URL+"/updatefields", nil)
+	resp, err := s.client.Do(req)
+	defer resp.Body.Close()
+	c.Assert(err, IsNil)
+	counter := resp.Header.Get("x-counter")
+
+	resp2, err2 := s.client.Do(req)
+	defer resp2.Body.Close()
+	c.Assert(err2, IsNil)
+	c.Assert(resp2.Header.Get(XFromCache), Equals, "1")
+	counter2 := resp2.Header.Get("x-counter")
+
+	c.Assert(counter, Not(Equals), counter2)
+}
+
 func (s *S) TestParseCacheControl(c *C) {
 	h := http.Header{}
 	for _ = range parseCacheControl(h) {
@@ -459,3 +488,54 @@ func (s *S) TestMaxStaleValue(c *C) {
 
 	c.Assert(getFreshness(respHeaders, reqHeaders), Equals, stale)
 }
+
+type containsHeaderChecker struct {
+	*CheckerInfo
+}
+
+func (c *containsHeaderChecker) Check(params []interface{}, names []string) (bool, string) {
+	items, ok := params[0].([]string)
+	if !ok {
+		return false, "Expected first param to be []string"
+	}
+	value, ok := params[1].(string)
+	if !ok {
+		return false, "Expected 2nd param to be string"
+	}
+	return containsHeader(items, value), ""
+}
+
+var ContainsHeader Checker = &containsHeaderChecker{&CheckerInfo{Name: "Contains", Params: []string{"Container", "expected to contain"}}}
+
+func (s *S) TestGetEndToEndHeaders(c *C) {
+	var (
+		headers http.Header
+		end2end []string
+	)
+
+	headers = http.Header{}
+	headers.Set("content-type", "text/html")
+	headers.Set("te", "deflate")
+
+	end2end = getEndToEndHeaders(headers)
+	c.Check(end2end, ContainsHeader, "content-type")
+	c.Check(end2end, Not(ContainsHeader), "te")
+
+	headers = http.Header{}
+	headers.Set("connection", "content-type")
+	headers.Set("content-type", "text/csv")
+	headers.Set("te", "deflate")
+	end2end = getEndToEndHeaders(headers)
+	c.Check(end2end, Not(ContainsHeader), "connection")
+	c.Check(end2end, Not(ContainsHeader), "content-type")
+	c.Check(end2end, Not(ContainsHeader), "te")
+
+	headers = http.Header{}
+	end2end = getEndToEndHeaders(headers)
+	c.Check(end2end, HasLen, 0)
+
+	headers = http.Header{}
+	headers.Set("connection", "content-type")
+	end2end = getEndToEndHeaders(headers)
+	c.Check(end2end, HasLen, 0)
+}