Procházet zdrojové kódy

refactor(response): unified response builder

Any string and []byte can be wrapped into io.Reader.
There's no need in target-specific code.
gudvinr před 3 týdny
rodič
revize
f4f655b350

+ 122 - 60
internal/http/response/builder.go

@@ -4,6 +4,7 @@
 package response // import "miniflux.app/v2/internal/http/response"
 
 import (
+	"bytes"
 	"compress/flate"
 	"compress/gzip"
 	"fmt"
@@ -17,7 +18,6 @@ import (
 	"time"
 
 	"github.com/andybalholm/brotli"
-	"github.com/andybalholm/brotli/matchfinder"
 )
 
 const compressionThreshold = 1024
@@ -45,6 +45,24 @@ var (
 	}
 )
 
+type writeResetCloser interface {
+	io.WriteCloser
+
+	Reset(io.Writer)
+}
+
+func pooledWriter(pool *sync.Pool, dst io.Writer) (io.WriteCloser, func()) {
+	w := pool.Get().(writeResetCloser)
+	w.Reset(dst)
+
+	cancel := func() {
+		w.Reset(io.Discard)
+		pool.Put(w)
+	}
+
+	return w, cancel
+}
+
 // Builder generates HTTP responses.
 type Builder struct {
 	w                 http.ResponseWriter
@@ -52,7 +70,7 @@ type Builder struct {
 	statusCode        int
 	headers           http.Header
 	enableCompression bool
-	body              any
+	body              io.Reader
 }
 
 // NewBuilder creates a new response builder.
@@ -74,13 +92,13 @@ func (b *Builder) WithHeader(key, value string) *Builder {
 
 // WithBodyAsBytes uses the given bytes to build the response.
 func (b *Builder) WithBodyAsBytes(body []byte) *Builder {
-	b.body = body
+	b.body = bytes.NewReader(body)
 	return b
 }
 
 // WithBodyAsString uses the given string to build the response.
 func (b *Builder) WithBodyAsString(body string) *Builder {
-	b.body = body
+	b.body = strings.NewReader(body)
 	return b
 }
 
@@ -126,6 +144,22 @@ func (b *Builder) WithCaching(etag string, duration time.Duration, callback func
 	}
 }
 
+type nopWriteCloser struct {
+	io.Writer
+}
+
+func (nopWriteCloser) Close() error { return nil }
+
+type nopWriteCloserReaderFrom struct {
+	io.Writer
+}
+
+func (nopWriteCloserReaderFrom) Close() error { return nil }
+
+func (c nopWriteCloserReaderFrom) ReadFrom(r io.Reader) (n int64, err error) {
+	return c.Writer.(io.ReaderFrom).ReadFrom(r)
+}
+
 // Write generates the HTTP response.
 func (b *Builder) Write() {
 	if b.body == nil {
@@ -133,19 +167,56 @@ func (b *Builder) Write() {
 		return
 	}
 
-	switch v := b.body.(type) {
-	case []byte:
-		b.compress(v)
-	case string:
-		b.compress([]byte(v))
-	case io.Reader:
-		// Compression not implemented in this case
-		b.writeHeaders()
-		_, err := io.Copy(b.w, v)
-		if err != nil {
-			slog.Error("Unable to write response body", slog.Any("error", err))
+	cancel := func() {}
+	defer func() { cancel() }()
+
+	var wc io.WriteCloser
+
+	switch b.acceptEncoding() {
+	case "":
+		// A hack to avoid infinite loop as [HTMLNotAcceptable] also calls [NewBuilder] from request.
+		cloned := b.r.Clone(b.r.Context())
+		cloned.Header.Del("Content-Encoding")
+
+		// [Builder.acceptEncoding] returns "" as per [HTTP Semantics]:
+		// 	As long as the identity;q=0 or *;q=0 directives do not explicitly forbid
+		// 	the identity value that means no encoding, the server must never return
+		// 	a 406 Not Acceptable error.
+		//
+		// [HTTP Semantics]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Accept-Encoding.
+		HTMLNotAcceptable(b.w, cloned)
+		return
+	case brotliEncoding:
+		b.headers.Set("Content-Encoding", brotliEncoding)
+		wc, cancel = pooledWriter(&brotliWriterPool, b.w)
+	case gzipEncoding:
+		b.headers.Set("Content-Encoding", gzipEncoding)
+		wc, cancel = pooledWriter(&gzipWriterPool, b.w)
+	case deflateEncoding:
+		b.headers.Set("Content-Encoding", deflateEncoding)
+		wc, cancel = pooledWriter(&flateWriterPool, b.w)
+	default:
+		// Fallback writer with no compression.
+
+		// Wraps [io.Writer] to make it implement [io.WriteCloser].
+		//
+		// That is straight-up copy from [io.NopCloser] for [io.Reader]
+		// as [io] doesn't have an alternative for [io.Writer].
+		wc = nopWriteCloser{b.w}
+		if _, ok := b.w.(io.ReaderFrom); ok {
+			wc = nopWriteCloserReaderFrom{b.w}
 		}
 	}
+
+	b.headers.Set("Vary", "Accept-Encoding")
+	b.writeHeaders()
+
+	_, err := io.Copy(wc, b.body)
+	if err != nil {
+		slog.Error("Unable to write response body", "error", err)
+	}
+
+	wc.Close()
 }
 
 func (b *Builder) writeHeaders() {
@@ -158,53 +229,44 @@ func (b *Builder) writeHeaders() {
 	b.w.WriteHeader(b.statusCode)
 }
 
-// values should be in sync with [Builder.compress] switch/case.
-var acceptEncoding = AcceptEncoding("br", "gzip", "deflate")
-
-func (b *Builder) compress(data []byte) {
-	if b.enableCompression && len(data) > compressionThreshold {
-		b.headers.Set("Vary", "Accept-Encoding")
-
-		encoding := acceptEncoding.Parse(b.r.Header.Get("Accept-Encoding"))
-		switch encoding {
-		case "br":
-			b.headers.Set("Content-Encoding", "br")
-			b.writeHeaders()
-
-			brotliWriter := brotliWriterPool.Get().(*matchfinder.Writer)
-			brotliWriter.Reset(b.w)
-			brotliWriter.Write(data)
-			brotliWriter.Close()
-			brotliWriter.Reset(io.Discard)
-			brotliWriterPool.Put(brotliWriter)
-			return
-		case "gzip":
-			b.headers.Set("Content-Encoding", "gzip")
-			b.writeHeaders()
-
-			gzipWriter := gzipWriterPool.Get().(*gzip.Writer)
-			gzipWriter.Reset(b.w)
-			gzipWriter.Write(data)
-			gzipWriter.Close()
-			gzipWriter.Reset(io.Discard)
-			gzipWriterPool.Put(gzipWriter)
-			return
-		case "deflate":
-			b.headers.Set("Content-Encoding", "deflate")
-			b.writeHeaders()
-
-			flateWriter := flateWriterPool.Get().(*flate.Writer)
-			flateWriter.Reset(b.w)
-			flateWriter.Write(data)
-			flateWriter.Close()
-			flateWriter.Reset(io.Discard)
-			flateWriterPool.Put(flateWriter)
-			return
-		}
+const (
+	brotliEncoding   = "br"
+	gzipEncoding     = "gzip"
+	deflateEncoding  = "deflate"
+	identityEncoding = "identity"
+)
+
+var acceptEncoding = AcceptEncoding(
+	brotliEncoding,
+	gzipEncoding,
+	deflateEncoding,
+)
+
+// acceptEncoding returns appropriate encoding for the response.
+//
+// Encoding will always be "identity" when compression explicitly disabled
+// or size of the body under [compressionThreshold] bytes.
+//
+// Otherwise, it is decided by "Accept-Encoding" header.
+func (b *Builder) acceptEncoding() string {
+	type sizer interface {
+		Size() int64
 	}
 
-	b.writeHeaders()
-	b.w.Write(data)
+	encoding := acceptEncoding.Parse(b.r.Header.Get("Accept-Encoding"))
+	if encoding == "" || encoding == identityEncoding {
+		return encoding
+	}
+
+	if !b.enableCompression {
+		return identityEncoding
+	}
+
+	if body, ok := b.body.(sizer); ok && body.Size() <= compressionThreshold {
+		return identityEncoding
+	}
+
+	return encoding
 }
 
 func normalizeETag(etag string) string {

+ 2 - 2
internal/http/response/builder_test.go

@@ -456,7 +456,7 @@ func TestBuildResponseWithCompressionDisabled(t *testing.T) {
 		t.Fatalf(`Unexpected header value, got %q instead of %q`, actual, expected)
 	}
 
-	expectedVary := ""
+	expectedVary := "Accept-Encoding"
 	actualVary := resp.Header.Get("Vary")
 	if actualVary != expectedVary {
 		t.Fatalf(`Unexpected vary header value, got %q instead of %q`, actualVary, expectedVary)
@@ -486,7 +486,7 @@ func TestBuildResponseWithDeflateCompressionAndSmallPayload(t *testing.T) {
 		t.Fatalf(`Unexpected header value, got %q instead of %q`, actual, expected)
 	}
 
-	expectedVary := ""
+	expectedVary := "Accept-Encoding"
 	actualVary := resp.Header.Get("Vary")
 	if actualVary != expectedVary {
 		t.Fatalf(`Unexpected vary header value, got %q instead of %q`, actualVary, expectedVary)

+ 23 - 0
internal/http/response/html.go

@@ -121,6 +121,29 @@ func HTMLNotFound(w http.ResponseWriter, r *http.Request) {
 		Write()
 }
 
+// HTMLNotAcceptable sends a not acceptable error to the client.
+func HTMLNotAcceptable(w http.ResponseWriter, r *http.Request) {
+	slog.Warn(http.StatusText(http.StatusNotAcceptable),
+		slog.String("client_ip", request.ClientIP(r)),
+		slog.Group("request",
+			slog.String("method", r.Method),
+			slog.String("uri", r.RequestURI),
+			slog.String("user_agent", r.UserAgent()),
+			slog.String("accept_encoding", r.Header.Get("Accept-Encoding")),
+		),
+		slog.Group("response",
+			slog.Int("status_code", http.StatusNotAcceptable),
+		),
+	)
+
+	NewBuilder(w, r).
+		WithStatus(http.StatusNotAcceptable).
+		WithHeader("Content-Type", "text/html; charset=utf-8").
+		WithHeader("Cache-Control", "no-cache, max-age=0, must-revalidate, no-store").
+		WithBodyAsString("Unknown Accept-Encoding").
+		Write()
+}
+
 // HTMLRedirect redirects the user to a relative path or an absolute http(s) URL.
 func HTMLRedirect(w http.ResponseWriter, r *http.Request, uri string) {
 	if !urllib.IsRelativePath(uri) && !urllib.IsAbsoluteURL(uri) {