Add cache control headers to 304 response

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 ### Change
 - `dpr` processing option doesn't enlarge image unless `enlarge` is true.
+- `304 Not Modified` responses includes `Cache-Control`, `Expires`, and `Vary` headers.
 
 ### Fix
 - Fix Client Hints behavior. `Width` is physical size, so we should divide it by `DPR` value.

ierrors/errors.go

@@ -62,13 +62,13 @@ func Wrap(err error, skip int) *Error {
 	return NewUnexpected(err.Error(), skip+1)
 }
 
-func WrapWithMessage(err error, skip int, msg string) *Error {
+func WrapWithPrefix(err error, skip int, prefix string) *Error {
 	if ierr, ok := err.(*Error); ok {
 		newErr := *ierr
-		ierr.Message = msg
+		newErr.Message = fmt.Sprintf("%s: %s", prefix, ierr.Message)
 		return &newErr
 	}
-	return NewUnexpected(err.Error(), skip+1)
+	return NewUnexpected(fmt.Sprintf("%s: %s", prefix, err), skip+1)
 }
 
 func callers(skip int) []uintptr {
@@ -87,10 +87,3 @@ func formatStack(stack []uintptr) string {
 
 	return strings.Join(lines, "\n")
 }
-
-func StatusCode(err error) int {
-	if ierr, ok := err.(*Error); ok {
-		return ierr.StatusCode
-	}
-	return 0
-}

imagedata/download.go

@@ -29,12 +29,19 @@ var (
 
 	// For tests
 	redirectAllRequestsTo string
-
-	ErrNotModified = ierrors.New(http.StatusNotModified, "Not Modified", "Not Modified")
 )
 
 const msgSourceImageIsUnreachable = "Source image is unreachable"
 
+type ErrorNotModified struct {
+	Message string
+	Headers map[string]string
+}
+
+func (e *ErrorNotModified) Error() string {
+	return e.Message
+}
+
 func initDownloading() error {
 	transport := &http.Transport{
 		Proxy:               http.ProxyFromEnvironment,
@@ -84,6 +91,18 @@ func initDownloading() error {
 	return nil
 }
 
+func headersToStore(res *http.Response) map[string]string {
+	m := make(map[string]string)
+
+	for _, h := range imageHeadersToStore {
+		if val := res.Header.Get(h); len(val) != 0 {
+			m[h] = val
+		}
+	}
+
+	return m
+}
+
 func requestImage(imageURL string, header http.Header) (*http.Response, error) {
 	req, err := http.NewRequest("GET", imageURL, nil)
 	if err != nil {
@@ -104,7 +123,7 @@ func requestImage(imageURL string, header http.Header) (*http.Response, error) {
 	}
 
 	if res.StatusCode == http.StatusNotModified {
-		return nil, ErrNotModified
+		return nil, &ErrorNotModified{Message: "Not Modified", Headers: headersToStore(res)}
 	}
 
 	if res.StatusCode != 200 {
@@ -152,12 +171,7 @@ func download(imageURL string, header http.Header) (*ImageData, error) {
 		return nil, err
 	}
 
-	imgdata.Headers = make(map[string]string)
+	imgdata.Headers = headersToStore(res)
-	for _, h := range imageHeadersToStore {
-		if val := res.Header.Get(h); len(val) != 0 {
-			imgdata.Headers[h] = val
-		}
-	}
 
 	return imgdata, nil
 }

imagedata/image_data.go

@@ -130,7 +130,11 @@ func FromFile(path, desc string) (*ImageData, error) {
 func Download(imageURL, desc string, header http.Header) (*ImageData, error) {
 	imgdata, err := download(imageURL, header)
 	if err != nil {
-		return nil, ierrors.WrapWithMessage(err, 1, fmt.Sprintf("Can't download %s: %s", desc, err))
+		if nmErr, ok := err.(*ErrorNotModified); ok {
+			nmErr.Message = fmt.Sprintf("Can't download %s: %s", desc, nmErr.Message)
+			return nil, nmErr
+		}
+		return nil, ierrors.WrapWithPrefix(err, 1, fmt.Sprintf("Can't download %s", desc))
 	}
 
 	return imgdata, nil

processing_handler.go

@@ -48,6 +48,37 @@ func initProcessingHandler() {
 	headerVaryValue = strings.Join(vary, ", ")
 }
 
+func setCacheControl(rw http.ResponseWriter, originHeaders map[string]string) {
+	var cacheControl, expires string
+
+	if config.CacheControlPassthrough && originHeaders != nil {
+		if val, ok := originHeaders["Cache-Control"]; ok {
+			cacheControl = val
+		}
+		if val, ok := originHeaders["Expires"]; ok {
+			expires = val
+		}
+	}
+
+	if len(cacheControl) == 0 && len(expires) == 0 {
+		cacheControl = fmt.Sprintf("max-age=%d, public", config.TTL)
+		expires = time.Now().Add(time.Second * time.Duration(config.TTL)).Format(http.TimeFormat)
+	}
+
+	if len(cacheControl) > 0 {
+		rw.Header().Set("Cache-Control", cacheControl)
+	}
+	if len(expires) > 0 {
+		rw.Header().Set("Expires", expires)
+	}
+}
+
+func setVary(rw http.ResponseWriter) {
+	if len(headerVaryValue) > 0 {
+		rw.Header().Set("Vary", headerVaryValue)
+	}
+}
+
 func respondWithImage(reqID string, r *http.Request, rw http.ResponseWriter, resultData *imagedata.ImageData, po *options.ProcessingOptions, originURL string, originData *imagedata.ImageData) {
 	var contentDisposition string
 	if len(po.Filename) > 0 {
@@ -70,32 +101,8 @@ func respondWithImage(reqID string, r *http.Request, rw http.ResponseWriter, res
 		}
 	}
 
-	var cacheControl, expires string
+	setCacheControl(rw, originData.Headers)
-
+	setVary(rw)
-	if config.CacheControlPassthrough && originData.Headers != nil {
-		if val, ok := originData.Headers["Cache-Control"]; ok {
-			cacheControl = val
-		}
-		if val, ok := originData.Headers["Expires"]; ok {
-			expires = val
-		}
-	}
-
-	if len(cacheControl) == 0 && len(expires) == 0 {
-		cacheControl = fmt.Sprintf("max-age=%d, public", config.TTL)
-		expires = time.Now().Add(time.Second * time.Duration(config.TTL)).Format(http.TimeFormat)
-	}
-
-	if len(cacheControl) > 0 {
-		rw.Header().Set("Cache-Control", cacheControl)
-	}
-	if len(expires) > 0 {
-		rw.Header().Set("Expires", expires)
-	}
-
-	if len(headerVaryValue) > 0 {
-		rw.Header().Set("Vary", headerVaryValue)
-	}
 
 	if config.EnableDebugHeaders {
 		rw.Header().Set("X-Origin-Content-Length", strconv.Itoa(len(originData.Data)))
@@ -120,7 +127,10 @@ func respondWithImage(reqID string, r *http.Request, rw http.ResponseWriter, res
 	)
 }
 
-func respondWithNotModified(reqID string, r *http.Request, rw http.ResponseWriter, po *options.ProcessingOptions, originURL string) {
+func respondWithNotModified(reqID string, r *http.Request, rw http.ResponseWriter, po *options.ProcessingOptions, originURL string, originHeaders map[string]string) {
+	setCacheControl(rw, originHeaders)
+	setVary(rw)
+
 	rw.WriteHeader(304)
 	router.LogResponse(
 		reqID, r, 304, nil,
@@ -209,14 +219,14 @@ func handleProcessing(reqID string, rw http.ResponseWriter, r *http.Request) {
 		defer metrics.StartDownloadingSegment(ctx)()
 		return imagedata.Download(imageURL, "source image", imgRequestHeader)
 	}()
-	switch {
+
-	case err == nil:
+	if err == nil {
 		defer originData.Close()
-	case ierrors.StatusCode(err) == http.StatusNotModified:
+	} else if nmErr, ok := err.(*imagedata.ErrorNotModified); ok && config.ETagEnabled {
 		rw.Header().Set("ETag", etagHandler.GenerateExpectedETag())
-		respondWithNotModified(reqID, r, rw, po, imageURL)
+		respondWithNotModified(reqID, r, rw, po, imageURL, nmErr.Headers)
 		return
-	default:
+	} else {
 		if ierr, ok := err.(*ierrors.Error); !ok || ierr.Unexpected {
 			errorreport.Report(err, r)
 		}
@@ -240,7 +250,7 @@ func handleProcessing(reqID string, rw http.ResponseWriter, r *http.Request) {
 		rw.Header().Set("ETag", etagHandler.GenerateActualETag())
 
 		if imgDataMatch && etagHandler.ProcessingOptionsMatch() {
-			respondWithNotModified(reqID, r, rw, po, imageURL)
+			respondWithNotModified(reqID, r, rw, po, imageURL, originData.Headers)
 			return
 		}
 	}