From 1040e2e8823febef52a7f341b47f4288250d1da1 Mon Sep 17 00:00:00 2001
From: Mansour Rahimi <rahimi.mnr@gmail.com>
Date: Sun, 26 Apr 2020 21:04:23 +0200
Subject: [PATCH] Fix caching on JSON body read

When response is JSON and it doesn't have Content-length header
and Transfer-encoding header is not chunked, reader will use bufio
where we miss the EOF and it never set the refreshed cache.

This fix will check the read length besides EOF err.
---
 httpcache.go      |  2 +-
 httpcache_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/httpcache.go b/httpcache.go
index b41a63d..6f4991c 100644
--- a/httpcache.go
+++ b/httpcache.go
@@ -533,7 +533,7 @@ type cachingReadCloser struct {
 func (r *cachingReadCloser) Read(p []byte) (n int, err error) {
 	n, err = r.R.Read(p)
 	r.buf.Write(p[:n])
-	if err == io.EOF {
+	if err == io.EOF || n < len(p) {
 		r.OnEOF(bytes.NewReader(r.buf.Bytes()))
 	}
 	return n, err
diff --git a/httpcache_test.go b/httpcache_test.go
index a504641..6f66237 100644
--- a/httpcache_test.go
+++ b/httpcache_test.go
@@ -2,6 +2,7 @@ package httpcache
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"flag"
 	"io"
@@ -158,6 +159,15 @@ func setup() {
 			}
 		}
 	}))
+
+	mux.HandleFunc("/json", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Cache-Control", "max-age=3600")
+		w.Header().Set("Content-Type", "application/json")
+		// This will force using bufio.Read() instead of chunkedReader.Read()
+		// to miss the EOF.
+		w.Header().Set("Transfer-encoding", "identity")
+		json.NewEncoder(w).Encode(map[string]string{"k": "v"})
+	}))
 }
 
 func teardown() {
@@ -398,6 +408,43 @@ func TestCacheOnlyIfBodyRead(t *testing.T) {
 	}
 }
 
+func TestCacheOnJsonBodyRead(t *testing.T) {
+	resetTest()
+	{
+		req, err := http.NewRequest("GET", s.server.URL+"/json", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		resp, err := s.client.Do(req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer resp.Body.Close()
+		var r json.RawMessage
+		err = json.NewDecoder(resp.Body).Decode(&r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if resp.Header.Get(XFromCache) != "" {
+			t.Fatalf("XFromCache header isn't blank")
+		}
+	}
+	{
+		req, err := http.NewRequest("GET", s.server.URL+"/json", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		resp, err := s.client.Do(req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer resp.Body.Close()
+		if resp.Header.Get(XFromCache) != "1" {
+			t.Fatalf("XFromCache header isn't set")
+		}
+	}
+}
+
 func TestOnlyReadBodyOnDemand(t *testing.T) {
 	resetTest()