From a29476f61e0475c8e66f68ae5e52f911361622e1 Mon Sep 17 00:00:00 2001
From: Caleb Doxsey <cdoxsey@pomerium.com>
Date: Tue, 31 Oct 2023 10:24:56 -0600
Subject: [PATCH] core/hpke: reduce memory usage from zstd (#4650)

* core/hpke: reduce memory usage from zstd

* use default compression, use default concurrency
---
 pkg/hpke/url.go      | 15 +++++++++++----
 pkg/hpke/url_test.go | 13 +++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/pkg/hpke/url.go b/pkg/hpke/url.go
index 12bbc272172..147ae65eed0 100644
--- a/pkg/hpke/url.go
+++ b/pkg/hpke/url.go
@@ -145,22 +145,29 @@ func withoutHPKEParams(values url.Values) url.Values {
 	return filtered
 }
 
-var zstdEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedBestCompression))
-
 func encodeQueryStringV1(values url.Values) []byte {
 	return []byte(values.Encode())
 }
 
+const zstdWindowSize = 8 << 10 // 8kiB
+
+var zstdEncoder, _ = zstd.NewWriter(nil,
+	zstd.WithEncoderLevel(zstd.SpeedDefault),
+	zstd.WithWindowSize(zstdWindowSize),
+)
+
 func encodeQueryStringV2(values url.Values) []byte {
 	return zstdEncoder.EncodeAll([]byte(values.Encode()), nil)
 }
 
-var zstdDecoder, _ = zstd.NewReader(nil)
-
 func decodeQueryStringV1(raw []byte) (url.Values, error) {
 	return url.ParseQuery(string(raw))
 }
 
+var zstdDecoder, _ = zstd.NewReader(nil,
+	zstd.WithDecoderLowmem(true),
+)
+
 func decodeQueryStringV2(raw []byte) (url.Values, error) {
 	bs, err := zstdDecoder.DecodeAll(raw, nil)
 	if err != nil {
diff --git a/pkg/hpke/url_test.go b/pkg/hpke/url_test.go
index 32f673b269c..8f14016c2d9 100644
--- a/pkg/hpke/url_test.go
+++ b/pkg/hpke/url_test.go
@@ -77,3 +77,16 @@ func TestEncryptURLValues(t *testing.T) {
 		assert.Less(t, len(encrypted.Encode()), 1024)
 	})
 }
+
+func BenchmarkZSTD(b *testing.B) {
+	payload := url.Values{
+		"a": {strings.Repeat("b", 128)},
+	}
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		bs := encodeQueryStringV2(payload)
+		_, _ = decodeQueryStringV2(bs)
+	}
+}