Merge pull request #11 from base-org/gzip

Add gzip compression on responses
base · Feb 27, 2024 · 539dd71 · 539dd71
2 parents c2304a5 + d87e315
commit 539dd71
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 28 deletions.
diff --git a/api/service/api.go b/api/service/api.go
@@ -39,8 +39,9 @@ func (e httpError) Error() string {
 }
 
 const (
-	sszAcceptType = "application/octet-stream"
-	serverTimeout = 60 * time.Second
+	jsonAcceptType = "application/json"
+	sszAcceptType  = "application/octet-stream"
+	serverTimeout  = 60 * time.Second
 )
 
 var (
@@ -97,6 +98,7 @@ func NewAPI(dataStoreClient storage.DataStoreReader, beaconClient client.BeaconB
 	r.Use(middleware.Timeout(serverTimeout))
 	r.Use(middleware.Recoverer)
 	r.Use(middleware.Heartbeat("/healthz"))
+	r.Use(middleware.Compress(5, jsonAcceptType, sszAcceptType))
 
 	recorder := opmetrics.NewPromHTTPRecorder(metrics.Registry(), m.MetricsNamespace)
 	r.Use(func(handler http.Handler) http.Handler {
@@ -187,6 +189,7 @@ func (a *API) blobSidecarHandler(w http.ResponseWriter, r *http.Request) {
 	responseType := r.Header.Get("Accept")
 
 	if responseType == sszAcceptType {
+		w.Header().Set("Content-Type", sszAcceptType)
 		res, err := blobSidecars.MarshalSSZ()
 		if err != nil {
 			a.logger.Error("unable to marshal blob sidecars to SSZ", "err", err)
@@ -202,6 +205,7 @@ func (a *API) blobSidecarHandler(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	} else {
+		w.Header().Set("Content-Type", jsonAcceptType)
 		err := json.NewEncoder(w).Encode(blobSidecars)
 		if err != nil {
 			a.logger.Error("unable to encode blob sidecars to JSON", "err", err)

diff --git a/api/service/api_test.go b/api/service/api_test.go
@@ -1,9 +1,11 @@
 package service
 
 import (
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http/httptest"
 	"os"
 	"testing"
@@ -228,38 +230,59 @@ func TestAPIService(t *testing.T) {
 
 	for _, test := range tests {
 		for _, rf := range responseFormat {
-			t.Run(fmt.Sprintf("%s-%s", test.name, rf), func(t *testing.T) {
-				request := httptest.NewRequest("GET", test.path, nil)
-				request.Header.Set("Accept", rf)
+			for _, compress := range []bool{true, false} {
+				testName := fmt.Sprintf("%s-%s", test.name, rf)
+				if compress {
+					testName = fmt.Sprintf("%s-%s", testName, "gzip")
+				}
 
-				response := httptest.NewRecorder()
+				t.Run(testName, func(t *testing.T) {
+					request := httptest.NewRequest("GET", test.path, nil)
+					request.Header.Set("Accept", rf)
 
-				a.router.ServeHTTP(response, request)
+					if compress {
+						request.Header.Set("Accept-Encoding", "gzip")
+					}
 
-				require.Equal(t, test.status, response.Code)
+					response := httptest.NewRecorder()
 
-				if test.status == 200 && test.expected != nil {
-					blobSidecars := storage.BlobSidecars{}
+					a.router.ServeHTTP(response, request)
 
-					var err error
-					if rf == "application/octet-stream" {
-						res := api.BlobSidecars{}
-						err = res.UnmarshalSSZ(response.Body.Bytes())
-						blobSidecars.Data = res.Sidecars
-					} else {
-						err = json.Unmarshal(response.Body.Bytes(), &blobSidecars)
-					}
+					require.Equal(t, test.status, response.Code)
 
-					require.NoError(t, err)
-					require.Equal(t, *test.expected, blobSidecars)
-				} else if test.status != 200 && rf == "application/json" && test.errMessage != "" {
-					var e httpError
-					err := json.Unmarshal(response.Body.Bytes(), &e)
-					require.NoError(t, err)
-					require.Equal(t, test.status, e.Code)
-					require.Equal(t, test.errMessage, e.Message)
-				}
-			})
+					if test.status == 200 && test.expected != nil {
+						var data []byte
+						if compress {
+							reader, err := gzip.NewReader(response.Body)
+							require.NoError(t, err)
+
+							data, err = io.ReadAll(reader)
+							require.NoError(t, err)
+						} else {
+							data = response.Body.Bytes()
+						}
+
+						blobSidecars := storage.BlobSidecars{}
+
+						if rf == "application/octet-stream" {
+							res := api.BlobSidecars{}
+							err = res.UnmarshalSSZ(data)
+							blobSidecars.Data = res.Sidecars
+						} else {
+							err = json.Unmarshal(data, &blobSidecars)
+						}
+
+						require.NoError(t, err)
+						require.Equal(t, *test.expected, blobSidecars)
+					} else if test.status != 200 && rf == "application/json" && test.errMessage != "" {
+						var e httpError
+						err := json.Unmarshal(response.Body.Bytes(), &e)
+						require.NoError(t, err)
+						require.Equal(t, test.status, e.Code)
+						require.Equal(t, test.errMessage, e.Message)
+					}
+				})
+			}
 		}
 	}
 }