feat(Cacher): support optional interface{ Size() int64 } for Conten…

…t-Length header (#60) * Update response.go Signed-off-by: Luke Young <[email protected]> * Update response.go Signed-off-by: Luke Young <[email protected]> * Update cacher.go Signed-off-by: Luke Young <[email protected]> * Update cacher.go Signed-off-by: Aofei Sheng <[email protected]> * Update response.go Signed-off-by: Aofei Sheng <[email protected]> * Update response.go Signed-off-by: Aofei Sheng <[email protected]> * Update response.go Signed-off-by: Aofei Sheng <[email protected]> * test(Cacher): cover `interface{ Size() int64 }` Signed-off-by: Aofei Sheng <[email protected]> --------- Signed-off-by: Luke Young <[email protected]> Signed-off-by: Aofei Sheng <[email protected]> Co-authored-by: Aofei Sheng <[email protected]>
goproxy · Jun 23, 2024 · 546d218 · 546d218
1 parent 4754c19
commit 546d218
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 27 deletions.
diff --git a/cacher.go b/cacher.go
@@ -16,13 +16,15 @@ type Cacher interface {
 	// The returned [io.ReadCloser] may optionally implement the following
 	// interfaces:
 	//  1. [io.Seeker], mainly for the Range request header.
-	//  2. interface{ LastModified() time.Time }, mainly for the
+	//  2. interface{ Size() int64 }, mainly for the Content-Length response
+	//     header. It takes effect only if 1 is not implemented.
+	//  3. interface{ LastModified() time.Time }, mainly for the
 	//     Last-Modified response header. Also for the If-Unmodified-Since,
 	//     If-Modified-Since, and If-Range request headers when 1 is
 	//     implemented.
-	//  3. interface{ ModTime() time.Time }, same as 2 but with lower
+	//  4. interface{ ModTime() time.Time }, same as 2 but with lower
 	//     priority.
-	//  4. interface{ ETag() string }, mainly for the ETag response header.
+	//  5. interface{ ETag() string }, mainly for the ETag response header.
 	//     Also for the If-Match, If-None-Match, and If-Range request
 	//     headers when 1 is implemented. Note that the return value will be
 	//     assumed to have complied with RFC 7232, section 2.3, so it will

diff --git a/response.go b/response.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/fs"
 	"net/http"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -88,6 +89,12 @@ func responseSuccess(rw http.ResponseWriter, req *http.Request, content io.Reade
 		return
 	}
 
+	if s, ok := content.(interface{ Size() int64 }); ok {
+		if size := s.Size(); size > 0 {
+			rw.Header().Set("Content-Length", strconv.FormatInt(size, 10))
+		}
+	}
+
 	if !lastModified.IsZero() {
 		rw.Header().Set("Last-Modified", lastModified.UTC().Format(http.TimeFormat))
 	}

diff --git a/response_test.go b/response_test.go
@@ -147,6 +147,13 @@ func TestResponseInternalServerError(t *testing.T) {
 	}
 }
 
+type successResponseBody_Size struct {
+	io.Reader
+	size int64
+}
+
+func (srb successResponseBody_Size) Size() int64 { return srb.size }
+
 type successResponseBody_LastModified struct {
 	io.Reader
 	lastModified time.Time
@@ -170,66 +177,84 @@ func (srb successResponseBody_ETag) ETag() string { return srb.etag }
 
 func TestResponseSuccess(t *testing.T) {
 	for _, tt := range []struct {
-		n                int
-		method           string
-		content          io.Reader
-		wantLastModified string
-		wantETag         string
-		wantContent      string
+		n                 int
+		method            string
+		content           io.Reader
+		wantContentLength int64
+		wantLastModified  string
+		wantETag          string
+		wantContent       string
 	}{
 		{
-			n:           1,
-			content:     strings.NewReader("foobar"),
-			wantContent: "foobar",
+			n:                 1,
+			content:           strings.NewReader("foobar"),
+			wantContentLength: 6,
+			wantContent:       "foobar",
 		},
 		{
-			n:       2,
-			method:  http.MethodHead,
-			content: strings.NewReader("foobar"),
+			n:                 2,
+			method:            http.MethodHead,
+			wantContentLength: 6,
+			content:           strings.NewReader("foobar"),
 		},
 		{
 			n: 3,
+			content: successResponseBody_Size{
+				Reader: strings.NewReader("foobar"),
+				size:   6,
+			},
+			wantContentLength: 6,
+			wantContent:       "foobar",
+		},
+		{
+			n: 4,
 			content: successResponseBody_LastModified{
 				Reader:       strings.NewReader("foobar"),
 				lastModified: time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC),
 			},
-			wantLastModified: "Sat, 01 Jan 2000 00:00:00 GMT",
-			wantContent:      "foobar",
+			wantContentLength: -1,
+			wantLastModified:  "Sat, 01 Jan 2000 00:00:00 GMT",
+			wantContent:       "foobar",
 		},
 		{
-			n: 4,
+			n: 5,
 			content: successResponseBody_ModTime{
 				Reader:  strings.NewReader("foobar"),
 				modTime: time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC),
 			},
-			wantLastModified: "Sat, 01 Jan 2000 00:00:00 GMT",
-			wantContent:      "foobar",
+			wantContentLength: -1,
+			wantLastModified:  "Sat, 01 Jan 2000 00:00:00 GMT",
+			wantContent:       "foobar",
 		},
 		{
-			n: 5,
+			n: 6,
 			content: successResponseBody_ETag{
 				Reader: strings.NewReader("foobar"),
 				etag:   `"foobar"`,
 			},
-			wantETag:    `"foobar"`,
-			wantContent: "foobar",
+			wantContentLength: -1,
+			wantETag:          `"foobar"`,
+			wantContent:       "foobar",
 		},
 		{
-			n: 6,
+			n: 7,
 			content: struct {
 				io.Reader
+				successResponseBody_Size
 				successResponseBody_LastModified
 				successResponseBody_ModTime
 				successResponseBody_ETag
 			}{
 				strings.NewReader("foobar"),
+				successResponseBody_Size{size: 6},
 				successResponseBody_LastModified{lastModified: time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)},
 				successResponseBody_ModTime{modTime: time.Date(2000, 1, 2, 0, 0, 0, 0, time.UTC)},
 				successResponseBody_ETag{etag: `"foobar"`},
 			},
-			wantLastModified: "Sat, 01 Jan 2000 00:00:00 GMT",
-			wantETag:         `"foobar"`,
-			wantContent:      "foobar",
+			wantContentLength: 6,
+			wantLastModified:  "Sat, 01 Jan 2000 00:00:00 GMT",
+			wantETag:          `"foobar"`,
+			wantContent:       "foobar",
 		},
 	} {
 		rec := httptest.NewRecorder()
@@ -244,6 +269,9 @@ func TestResponseSuccess(t *testing.T) {
 		if got, want := recr.Header.Get("Cache-Control"), "public, max-age=60"; got != want {
 			t.Errorf("test(%d): got %q, want %q", tt.n, got, want)
 		}
+		if got, want := recr.ContentLength, tt.wantContentLength; got != want {
+			t.Errorf("test(%d): got %d, want %d", tt.n, got, want)
+		}
 		if got, want := recr.Header.Get("Last-Modified"), tt.wantLastModified; got != want {
 			t.Errorf("test(%d): got %q, want %q", tt.n, got, want)
 		}