From b01705161b86a75650241193bd52105bcfabe166 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Sun, 17 Jul 2022 03:26:46 -0700
Subject: [PATCH] zstd: Allow single segments up to "max decoded size" (#643)

Also set single-segment when less than allowed window size.
---
 zstd/decoder.go         |  3 +++
 zstd/encoder.go         |  4 ++--
 zstd/encoder_options.go |  2 +-
 zstd/framedec.go        | 19 +++++++++++++------
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/zstd/decoder.go b/zstd/decoder.go
index 286c8f9d71..d212f4737f 100644
--- a/zstd/decoder.go
+++ b/zstd/decoder.go
@@ -348,6 +348,9 @@ func (d *Decoder) DecodeAll(input, dst []byte) ([]byte, error) {
 			frame.history.setDict(&dict)
 		}
 		if frame.WindowSize > d.o.maxWindowSize {
+			if debugDecoder {
+				println("window size exceeded:", frame.WindowSize, ">", d.o.maxWindowSize)
+			}
 			return dst, ErrWindowSizeExceeded
 		}
 		if frame.FrameContentSize != fcsUnknown {
diff --git a/zstd/encoder.go b/zstd/encoder.go
index e6b1d01cf6..7aaaedb23e 100644
--- a/zstd/encoder.go
+++ b/zstd/encoder.go
@@ -528,8 +528,8 @@ func (e *Encoder) EncodeAll(src, dst []byte) []byte {
 		// If a non-single block is needed the encoder will reset again.
 		e.encoders <- enc
 	}()
-	// Use single segments when above minimum window and below 1MB.
-	single := len(src) < 1<<20 && len(src) > MinWindowSize
+	// Use single segments when above minimum window and below window size.
+	single := len(src) <= e.o.windowSize && len(src) > MinWindowSize
 	if e.o.single != nil {
 		single = *e.o.single
 	}
diff --git a/zstd/encoder_options.go b/zstd/encoder_options.go
index 44d8dbd199..a7c5e1aac4 100644
--- a/zstd/encoder_options.go
+++ b/zstd/encoder_options.go
@@ -283,7 +283,7 @@ func WithNoEntropyCompression(b bool) EOption {
 // a decoder is allowed to reject a compressed frame which requests a memory size beyond decoder's authorized range.
 // For broader compatibility, decoders are recommended to support memory sizes of at least 8 MB.
 // This is only a recommendation, each decoder is free to support higher or lower limits, depending on local limitations.
-// If this is not specified, block encodes will automatically choose this based on the input size.
+// If this is not specified, block encodes will automatically choose this based on the input size and the window size.
 // This setting has no effect on streamed encodes.
 func WithSingleSegment(b bool) EOption {
 	return func(o *encoderOptions) error {
diff --git a/zstd/framedec.go b/zstd/framedec.go
index fa0a633f38..9311ef51f5 100644
--- a/zstd/framedec.go
+++ b/zstd/framedec.go
@@ -231,20 +231,27 @@ func (d *frameDec) reset(br byteBuffer) error {
 		d.crc.Reset()
 	}
 
+	if d.WindowSize > d.o.maxWindowSize {
+		if debugDecoder {
+			printf("window size %d > max %d\n", d.WindowSize, d.o.maxWindowSize)
+		}
+		return ErrWindowSizeExceeded
+	}
+
 	if d.WindowSize == 0 && d.SingleSegment {
 		// We may not need window in this case.
 		d.WindowSize = d.FrameContentSize
 		if d.WindowSize < MinWindowSize {
 			d.WindowSize = MinWindowSize
 		}
-	}
-
-	if d.WindowSize > uint64(d.o.maxWindowSize) {
-		if debugDecoder {
-			printf("window size %d > max %d\n", d.WindowSize, d.o.maxWindowSize)
+		if d.WindowSize > d.o.maxDecodedSize {
+			if debugDecoder {
+				printf("window size %d > max %d\n", d.WindowSize, d.o.maxWindowSize)
+			}
+			return ErrDecoderSizeExceeded
 		}
-		return ErrWindowSizeExceeded
 	}
+
 	// The minimum Window_Size is 1 KB.
 	if d.WindowSize < MinWindowSize {
 		if debugDecoder {