From 329ec6831a7dd2893fb108ccbaccf6aa56c8ada3 Mon Sep 17 00:00:00 2001
From: jub0bs <52325150+jub0bs@users.noreply.github.com>
Date: Tue, 11 Feb 2025 14:14:55 +0100
Subject: [PATCH] util/httputil: reduce heap allocations in
 newCompressedResponseWriter (#16001)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* util/httputil: Benchmark newCompressedResponseWriter

This benchmark illustrates that newCompressedResponseWriter incurs a
prohibitive amount of heap allocations when handling a request containing a
malicious Accept-Encoding header.¬

Signed-off-by: jub0bs <jcretel-infosec+github@protonmail.com>

* util/httputil: Improve newCompressedResponseWriter

This change dramatically reduces the heap allocations (in bytes)
incurred when handling a request containing a malicious Accept-Encoding header.

Below are some benchmark results; for conciseness, I've omitted the name of the
benchmark function (BenchmarkNewCompressionHandler_MaliciousAcceptEncoding):

```
goos: darwin
goarch: amd64
pkg: github.com/prometheus/prometheus/util/httputil
cpu: Intel(R) Core(TM) i7-6700HQ CPU @ 2.60GHz
│     old     │                 new                 │
│   sec/op    │   sec/op     vs base                │
  18.60m ± 2%   13.54m ± 3%  -27.17% (p=0.000 n=10)

│       old        │                 new                 │
│       B/op       │    B/op     vs base                 │
  16785442.50 ± 0%   32.00 ± 0%  -100.00% (p=0.000 n=10)

│    old     │                new                 │
│ allocs/op  │ allocs/op   vs base                │
  2.000 ± 0%   1.000 ± 0%  -50.00% (p=0.000 n=10)
```

Signed-off-by: jub0bs <jcretel-infosec+github@protonmail.com>

---------

Signed-off-by: jub0bs <jcretel-infosec+github@protonmail.com>
---
 util/httputil/compression.go      | 12 ++++++++++--
 util/httputil/compression_test.go | 12 ++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/util/httputil/compression.go b/util/httputil/compression.go
index 9a8a666453..d5bedb7fa9 100644
--- a/util/httputil/compression.go
+++ b/util/httputil/compression.go
@@ -56,8 +56,13 @@ func (c *compressedResponseWriter) Close() {
 
 // Constructs a new compressedResponseWriter based on client request headers.
 func newCompressedResponseWriter(writer http.ResponseWriter, req *http.Request) *compressedResponseWriter {
-	encodings := strings.Split(req.Header.Get(acceptEncodingHeader), ",")
-	for _, encoding := range encodings {
+	raw := req.Header.Get(acceptEncodingHeader)
+	var (
+		encoding   string
+		commaFound bool
+	)
+	for {
+		encoding, raw, commaFound = strings.Cut(raw, ",")
 		switch strings.TrimSpace(encoding) {
 		case gzipEncoding:
 			writer.Header().Set(contentEncodingHeader, gzipEncoding)
@@ -72,6 +77,9 @@ func newCompressedResponseWriter(writer http.ResponseWriter, req *http.Request)
 				writer:         zlib.NewWriter(writer),
 			}
 		}
+		if !commaFound {
+			break
+		}
 	}
 	return &compressedResponseWriter{
 		ResponseWriter: writer,
diff --git a/util/httputil/compression_test.go b/util/httputil/compression_test.go
index e166c7de79..fd3f1f66d4 100644
--- a/util/httputil/compression_test.go
+++ b/util/httputil/compression_test.go
@@ -18,6 +18,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/klauspost/compress/gzip"
@@ -72,6 +73,17 @@ func TestCompressionHandler_PlainText(t *testing.T) {
 	require.Equal(t, expected, actual, "expected response with content")
 }
 
+func BenchmarkNewCompressionHandler_MaliciousAcceptEncoding(b *testing.B) {
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/whatever", nil)
+	req.Header.Set("Accept-Encoding", strings.Repeat(",", http.DefaultMaxHeaderBytes))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		newCompressedResponseWriter(rec, req)
+	}
+}
+
 func TestCompressionHandler_Gzip(t *testing.T) {
 	tearDown := setup()
 	defer tearDown()