From 8a309675429033a7d619f24bc891577d5a75250c Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Mon, 16 Sep 2024 09:59:03 -0700
Subject: [PATCH] Limit S3 Select JSON documents to 10MB (#20439)

Closes #20430

Limit allocations from badly formed documents.
---
 internal/s3select/json/reader.go | 6 +++++-
 internal/s3select/select.go      | 1 +
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/internal/s3select/json/reader.go b/internal/s3select/json/reader.go
index 4285c23fa..52eda1404 100644
--- a/internal/s3select/json/reader.go
+++ b/internal/s3select/json/reader.go
@@ -26,6 +26,10 @@ import (
 	"github.com/bcicen/jstream"
 )
 
+// Limit single document size to 10MiB, 10x the AWS limit:
+// https://docs.aws.amazon.com/AmazonS3/latest/userguide/selecting-content-from-objects.html
+const maxDocumentSize = 10 << 20
+
 // Reader - JSON record reader for S3Select.
 type Reader struct {
 	args       *ReaderArgs
@@ -80,7 +84,7 @@ func (r *Reader) Close() error {
 // NewReader - creates new JSON reader using readCloser.
 func NewReader(readCloser io.ReadCloser, args *ReaderArgs) *Reader {
 	readCloser = &syncReadCloser{rc: readCloser}
-	d := jstream.NewDecoder(readCloser, 0).ObjectAsKVS()
+	d := jstream.NewDecoder(io.LimitReader(readCloser, maxDocumentSize), 0).ObjectAsKVS()
 	return &Reader{
 		args:       args,
 		decoder:    d,
diff --git a/internal/s3select/select.go b/internal/s3select/select.go
index 6026c9b3c..74b8f43ed 100644
--- a/internal/s3select/select.go
+++ b/internal/s3select/select.go
@@ -442,6 +442,7 @@ func (s3Select *S3Select) Open(rsc io.ReadSeekCloser) error {
 				s3Select.recordReader = json.NewPReader(s3Select.progressReader, &s3Select.Input.JSONArgs)
 			}
 		} else {
+			// Document mode.
 			s3Select.recordReader = json.NewReader(s3Select.progressReader, &s3Select.Input.JSONArgs)
 		}