Merge pull request #2 from pepabo/fix_notmodified

pepabo · Jan 31, 2023 · 3d95a50 · 3d95a50
2 parents edb1319 + 468c849
commit 3d95a50
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 7 deletions.
diff --git a/main.go b/main.go
@@ -86,6 +86,12 @@ func proxy(w http.ResponseWriter, r *http.Request) {
 	}
 	defer orgRes.Body.Close()
 
+	if orgRes.Header.Get("Last-Modified") != "" {
+		w.Header().Set("Last-Modified", orgRes.Header.Get("Last-Modified"))
+	} else {
+		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
+	}
+
 	if orgRes.StatusCode == http.StatusNotModified {
 		w.WriteHeader(http.StatusNotModified)
 		return
@@ -127,11 +133,6 @@ func proxy(w http.ResponseWriter, r *http.Request) {
 
 	w.Header().Set("Content-Type", "image/jpeg")
 	w.Header().Set("Content-Length", strconv.Itoa(buf.Len()))
-	if orgRes.Header.Get("Last-Modified") != "" {
-		w.Header().Set("Last-Modified", orgRes.Header.Get("Last-Modified"))
-	} else {
-		w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
-	}
 
 	if _, err := io.Copy(w, buf); err != nil {
 		// ignore already close client.

diff --git a/main_test.go b/main_test.go
@@ -79,25 +79,39 @@ func TestProxyJPEG(t *testing.T) {
 		t.Errorf("value of Content-Length header %d is larger than origin's one %d", res.ContentLength, orgRes.ContentLength)
 	}
 }
-func TestProxyNotModified(t *testing.T) {
+
+func TestOriginNotModifiedJPEG(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(proxy))
 
 	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Last-Modified", "2023-01-01T00:00:00")
 		w.WriteHeader(http.StatusNotModified)
+		http.ServeFile(w, r, "./testdata/oyaki.jpg")
 	}))
+
 	orgSrvURL = origin.URL
 
-	url := ts.URL + "/corn.png"
+	url := ts.URL + "/oyaki.jpg"
 
 	res, err := http.Get(url)
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	_, err = http.Get(orgSrvURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	if res.StatusCode != http.StatusNotModified {
 		t.Errorf("HTTP status is %d, want %d", res.StatusCode, http.StatusNotModified)
 	}
+
+	if res.ContentLength < 0 {
+		t.Errorf("Content-Length header does not exist")
+	}
 }
+
 func TestProxyPNG(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(proxy))
 
@@ -152,6 +166,41 @@ func TestOriginNotExist(t *testing.T) {
 	}
 }
 
+func TestOriginNotModifiedPNG(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(proxy))
+
+	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Last-Modified", "2023-01-01T00:00:00")
+		w.WriteHeader(http.StatusNotModified)
+		http.ServeFile(w, r, "./testdata/corn.png")
+	}))
+
+	orgSrvURL = origin.URL
+	url := ts.URL + "/corn.png"
+
+	res, err := http.Get(url)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = http.Get(orgSrvURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if res.StatusCode != http.StatusNotModified {
+		t.Errorf("HTTP status is %d, want %d", res.StatusCode, http.StatusNotModified)
+	}
+
+	if res.Header.Get("Last-Modified") == "" {
+		t.Errorf("Not set header")
+	}
+
+	if res.ContentLength < 0 {
+		t.Errorf("Content-Length header does not exist")
+	}
+}
+
 func BenchmarkProxyJpeg(b *testing.B) {
 	b.ResetTimer()
 	ts := httptest.NewServer(http.HandlerFunc(proxy))