util+server: Fix bug around chunked request handling. (#6906)

This commit fixes a request handling bug introduced in #6868, which
caused OPA to treat all incoming chunked requests as if they had
zero-length request bodies.

The fix detects cases where the request body size is unknown in the
DecodingLimits handler, and propagates a request context key down to
the `util.ReadMaybeCompressedBody` function, allowing it to correctly
select between using the original `io.ReadAll` style for chunked
requests, or the newer preallocated buffers approach (for requests of
known size).

This change has a small, but barely visible performance impact for large
requests (<5% increase in GC pauses for a 1GB request JSON blob), and
minimal, if any, effect on RPS under load.

Fixes: #6904

Signed-off-by: Philip Conrad <philip@chariot-chaser.net>

util/decoding/context.go

@@ -11,10 +11,20 @@ const (
 	reqCtxKeyGzipMaxLen = requestContextKey("server-decoding-plugin-context-gzip-max-length")
 )
 
+func AddServerDecodingMaxLen(ctx context.Context, maxLen int64) context.Context {
+	return context.WithValue(ctx, reqCtxKeyMaxLen, maxLen)
+}
+
 func AddServerDecodingGzipMaxLen(ctx context.Context, maxLen int64) context.Context {
 	return context.WithValue(ctx, reqCtxKeyGzipMaxLen, maxLen)
 }
 
+// Used for enforcing max body content limits when dealing with chunked requests.
+func GetServerDecodingMaxLen(ctx context.Context) (int64, bool) {
+	maxLength, ok := ctx.Value(reqCtxKeyMaxLen).(int64)
+	return maxLength, ok
+}
+
 func GetServerDecodingGzipMaxLen(ctx context.Context) (int64, bool) {
 	gzipMaxLength, ok := ctx.Value(reqCtxKeyGzipMaxLen).(int64)
 	return gzipMaxLength, ok

util/read_gzip_body.go

@@ -27,13 +27,24 @@ var gzipReaderPool = sync.Pool{
 // payload size, but not an unbounded amount of memory, as was potentially
 // possible before.
 func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
-	if r.ContentLength <= 0 {
-		return []byte{}, nil
-	}
-	// Read content from the request body into a buffer of known size.
-	content := bytes.NewBuffer(make([]byte, 0, r.ContentLength))
-	if _, err := io.CopyN(content, r.Body, r.ContentLength); err != nil {
-		return content.Bytes(), err
+	var content *bytes.Buffer
+	// Note(philipc): If the request body is of unknown length (such as what
+	// happens when 'Transfer-Encoding: chunked' is set), we have to do an
+	// incremental read of the body. In this case, we can't be too clever, we
+	// just do the best we can with whatever is streamed over to us.
+	// Fetch gzip payload size limit from request context.
+	if maxLength, ok := decoding.GetServerDecodingMaxLen(r.Context()); ok {
+		bs, err := io.ReadAll(io.LimitReader(r.Body, maxLength))
+		if err != nil {
+			return bs, err
+		}
+		content = bytes.NewBuffer(bs)
+	} else {
+		// Read content from the request body into a buffer of known size.
+		content = bytes.NewBuffer(make([]byte, 0, r.ContentLength))
+		if _, err := io.CopyN(content, r.Body, r.ContentLength); err != nil {
+			return content.Bytes(), err
+		}
 	}
 
 	// Decompress gzip content by reading from the buffer.