dashboard/app: pre-gzip all responses

Author: tarasmadan

dashboard/app/api.go

@@ -111,16 +111,16 @@ func handleJSON(fn JSONHandler) http.Handler {
 			http.Error(w, err.Error(), status)
 			return
 		}
-		w.Header().Set("Content-Type", "application/json")
-		wJS := w.(io.Writer)
-		if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
-			w.Header().Set("Content-Encoding", "gzip")
-			gw := gzip.NewWriter(w)
-			defer gw.Close()
-			wJS = gw
-		}
+
+		wJS := newGzipResponseWriterCloser(w)
+		defer wJS.Close()
 		if err := json.NewEncoder(wJS).Encode(reply); err != nil {
 			log.Errorf(c, "failed to encode reply: %v", err)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		if err := wJS.writeResult(r); err != nil {
+			log.Errorf(c, "wJS.writeResult: %s", err.Error())
 		}
 	})
 }

dashboard/app/handler.go

@@ -5,11 +5,13 @@ package main
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"sort"
 	"strings"
@@ -43,37 +45,45 @@ func handleContext(fn contextHandler) http.Handler {
 			}
 			defer backpressureRobots(c, r)()
 		}
-		if err := fn(c, w, r); err != nil {
-			hdr := commonHeaderRaw(c, r)
-			data := &struct {
-				Header  *uiHeader
-				Error   string
-				TraceID string
-			}{
-				Header:  hdr,
-				Error:   err.Error(),
-				TraceID: strings.Join(r.Header["X-Cloud-Trace-Context"], " "),
-			}
-			if err == ErrAccess {
-				if hdr.LoginLink != "" {
-					http.Redirect(w, r, hdr.LoginLink, http.StatusTemporaryRedirect)
-					return
-				}
-				http.Error(w, "403 Forbidden", http.StatusForbidden)
+
+		w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		gzw := newGzipResponseWriterCloser(w)
+		defer gzw.Close()
+		err := fn(c, gzw, r)
+		if err == nil {
+			if err = gzw.writeResult(r); err == nil {
 				return
 			}
-			var redir *ErrRedirect
-			if errors.As(err, &redir) {
-				http.Redirect(w, r, redir.Error(), http.StatusFound)
+		}
+		hdr := commonHeaderRaw(c, r)
+		data := &struct {
+			Header  *uiHeader
+			Error   string
+			TraceID string
+		}{
+			Header:  hdr,
+			Error:   err.Error(),
+			TraceID: strings.Join(r.Header["X-Cloud-Trace-Context"], " "),
+		}
+		if err == ErrAccess {
+			if hdr.LoginLink != "" {
+				http.Redirect(w, r, hdr.LoginLink, http.StatusTemporaryRedirect)
 				return
 			}
+			http.Error(w, "403 Forbidden", http.StatusForbidden)
+			return
+		}
+		var redir *ErrRedirect
+		if errors.As(err, &redir) {
+			http.Redirect(w, r, redir.Error(), http.StatusFound)
+			return
+		}
 
-			status := logErrorPrepareStatus(c, err)
-			w.WriteHeader(status)
-			if err1 := templates.ExecuteTemplate(w, "error.html", data); err1 != nil {
-				combinedError := fmt.Sprintf("got err \"%v\" processing ExecuteTemplate() for err \"%v\"", err1, err)
-				http.Error(w, combinedError, http.StatusInternalServerError)
-			}
+		status := logErrorPrepareStatus(c, err)
+		w.WriteHeader(status)
+		if err1 := templates.ExecuteTemplate(w, "error.html", data); err1 != nil {
+			combinedError := fmt.Sprintf("got err \"%v\" processing ExecuteTemplate() for err \"%v\"", err1, err)
+			http.Error(w, combinedError, http.StatusInternalServerError)
 		}
 	})
 }
@@ -339,3 +349,64 @@ func encodeCookie(w http.ResponseWriter, cd *cookieData) {
 }
 
 var templates = html.CreateGlob("*.html")
+
+// gzipResponseWriterCloser accumulates the gzipped result.
+// In case of error during the handler processing, we'll drop this gzipped data.
+// It allows to call http.Error in the middle of the response generation.
+//
+// For 200 Ok responses we return the compressed data or decompress it depending on the client Accept-Encoding header.
+type gzipResponseWriterCloser struct {
+	w                 *gzip.Writer
+	plainResponseSize int
+	buf               *bytes.Buffer
+	rw                http.ResponseWriter
+}
+
+func (g *gzipResponseWriterCloser) Write(p []byte) (n int, err error) {
+	g.plainResponseSize += len(p)
+	return g.w.Write(p)
+}
+
+func (g *gzipResponseWriterCloser) Close() {
+	if g.w != nil {
+		g.w.Close()
+	}
+}
+
+func (g *gzipResponseWriterCloser) Header() http.Header {
+	return g.rw.Header()
+}
+
+func (g *gzipResponseWriterCloser) WriteHeader(statusCode int) {
+	g.rw.WriteHeader(statusCode)
+}
+
+func (g *gzipResponseWriterCloser) writeResult(r *http.Request) error {
+	g.w.Close()
+	g.w = nil
+	clientSupportsGzip := strings.Contains(r.Header.Get("Accept-Encoding"), "gzip")
+	if clientSupportsGzip {
+		g.rw.Header().Set("Content-Encoding", "gzip")
+		_, err := g.rw.Write(g.buf.Bytes())
+		return err
+	}
+	if g.plainResponseSize > 31<<20 { // 32MB is the AppEngine hard limit for the response size.
+		return fmt.Errorf("len(response) > 31M, try to request gzipped: %w", ErrClientBadRequest)
+	}
+	gzr, err := gzip.NewReader(g.buf)
+	if err != nil {
+		return fmt.Errorf("gzip.NewReader: %w", err)
+	}
+	defer gzr.Close()
+	_, err = io.Copy(g.rw, gzr)
+	return err
+}
+
+func newGzipResponseWriterCloser(w http.ResponseWriter) *gzipResponseWriterCloser {
+	buf := &bytes.Buffer{}
+	return &gzipResponseWriterCloser{
+		w:   gzip.NewWriter(buf),
+		buf: buf,
+		rw:  w,
+	}
+}

dashboard/app/handler_test.go

@@ -0,0 +1,66 @@
+// Copyright 2025 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package main
+
+import (
+	"compress/gzip"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGzipResponseWriterCloser_no_compression(t *testing.T) {
+	res := httptest.NewRecorder()
+	gz := newGzipResponseWriterCloser(res)
+	gz.Write([]byte("test"))
+
+	err := gz.writeResult(httpRequestWithAcceptedEncoding(""))
+	assert.NoError(t, err)
+	assert.Equal(t, "test", res.Body.String())
+	assert.Equal(t, "", res.Header().Get("Content-Encoding"))
+}
+
+func TestGzipResponseWriterCloser_with_compression(t *testing.T) {
+	res := httptest.NewRecorder()
+	gz := newGzipResponseWriterCloser(res)
+	gz.Write([]byte("test"))
+
+	err := gz.writeResult(httpRequestWithAcceptedEncoding("gzip"))
+	assert.NoError(t, err)
+	assert.Equal(t, "gzip", res.Header().Get("Content-Encoding"))
+
+	gr, _ := gzip.NewReader(res.Body)
+	gotBytes := make([]byte, 28)
+	n, _ := gr.Read(gotBytes)
+	gotBytes = gotBytes[:n]
+	assert.Equal(t, "test", string(gotBytes))
+}
+
+func TestGzipResponseWriterCloser_headers(t *testing.T) {
+	res := httptest.NewRecorder()
+	gz := newGzipResponseWriterCloser(res)
+
+	gz.Header().Add("key", "val1")
+	gz.Header().Add("key", "val2")
+	err := gz.writeResult(httpRequestWithAcceptedEncoding(""))
+	assert.NoError(t, err)
+	assert.Equal(t, http.Header{
+		"Key": []string{"val1", "val2"},
+	}, res.Header())
+}
+
+func TestGzipResponseWriterCloser_status(t *testing.T) {
+	res := httptest.NewRecorder()
+	gz := newGzipResponseWriterCloser(res)
+
+	gz.WriteHeader(333)
+	gz.writeResult(httpRequestWithAcceptedEncoding("gzip"))
+	assert.Equal(t, 333, res.Code)
+}
+
+func httpRequestWithAcceptedEncoding(encoding string) *http.Request {
+	return &http.Request{Header: http.Header{"Accept-Encoding": []string{encoding}}}
+}