From 177458d36e8e2fb330df5a4b5b65f321f4dae546 Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Thu, 19 Dec 2024 14:46:42 -0300
Subject: [PATCH] mime/multipart: add field Reader.MaxMIMEHeaderSize

If the field is set, it is used instead of maxMIMEHeaderSize constant, allowing
to further constraint memory usage when parsing multipart streams.

Fixes https://github.com/golang/go/issues/68889
---
 src/mime/multipart/multipart.go      | 20 ++++++++++++++--
 src/mime/multipart/multipart_test.go | 35 ++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/src/mime/multipart/multipart.go b/src/mime/multipart/multipart.go
index 17088bc30e1b43..6379dfbe973763 100644
--- a/src/mime/multipart/multipart.go
+++ b/src/mime/multipart/multipart.go
@@ -331,6 +331,11 @@ func (p *Part) Close() error {
 // Reader's underlying parser consumes its input as needed. Seeking
 // isn't supported.
 type Reader struct {
+	// MaxMIMEHeaderSize is the maximum size of a MIME header we will parse,
+	// including header keys, values, and map overhead. If not set, then the
+	// default value of 10 << 20 is used.
+	MaxMIMEHeaderSize int64
+
 	bufReader *bufio.Reader
 	tempDir   string // used in tests
 
@@ -362,6 +367,17 @@ func maxMIMEHeaders() int64 {
 	return 10000
 }
 
+// maxMIMEHeaderSize returns the maximum size of a MIME header we will parse.
+// It uses the value of Reader.MaxMIMEHeaderSize if it is not 0, otherwise the
+// value of maxMIMEHeaderSize constant is used.
+func (r *Reader) maxMIMEHeaderSize() int64 {
+	if r.MaxMIMEHeaderSize != 0 {
+		return r.MaxMIMEHeaderSize
+	} else {
+		return maxMIMEHeaderSize
+	}
+}
+
 // NextPart returns the next part in the multipart or an error.
 // When there are no more parts, the error [io.EOF] is returned.
 //
@@ -369,7 +385,7 @@ func maxMIMEHeaders() int64 {
 // has a value of "quoted-printable", that header is instead
 // hidden and the body is transparently decoded during Read calls.
 func (r *Reader) NextPart() (*Part, error) {
-	return r.nextPart(false, maxMIMEHeaderSize, maxMIMEHeaders())
+	return r.nextPart(false, r.maxMIMEHeaderSize(), maxMIMEHeaders())
 }
 
 // NextRawPart returns the next part in the multipart or an error.
@@ -378,7 +394,7 @@ func (r *Reader) NextPart() (*Part, error) {
 // Unlike [Reader.NextPart], it does not have special handling for
 // "Content-Transfer-Encoding: quoted-printable".
 func (r *Reader) NextRawPart() (*Part, error) {
-	return r.nextPart(true, maxMIMEHeaderSize, maxMIMEHeaders())
+	return r.nextPart(true, r.maxMIMEHeaderSize(), maxMIMEHeaders())
 }
 
 func (r *Reader) nextPart(rawPart bool, maxMIMEHeaderSize, maxMIMEHeaders int64) (*Part, error) {
diff --git a/src/mime/multipart/multipart_test.go b/src/mime/multipart/multipart_test.go
index e0cb768c6967a5..b7bbf5ae320cab 100644
--- a/src/mime/multipart/multipart_test.go
+++ b/src/mime/multipart/multipart_test.go
@@ -931,6 +931,41 @@ Cases:
 	}
 }
 
+func TestParseMaxMIMEHeaderSize(t *testing.T) {
+	sep := "MyBoundary"
+	in := `This is a multi-part message.  This line is ignored.
+--MyBoundary
+Header1: value1
+Header2: value2
+Header3: value3
+Header4: value4
+Header5: value5
+Header6: value6
+Header7: value7
+Header8: value8
+
+My value
+The end.
+--MyBoundary--`
+
+	in = strings.Replace(in, "\n", "\r\n", -1)
+
+	// Control.
+	r := NewReader(strings.NewReader(in), sep)
+	_, err := r.NextPart()
+	if err != nil {
+		t.Fatalf("control failed: %v", err)
+	}
+
+	// Test MaxMIMEHeaderSize.
+	r = NewReader(strings.NewReader(in), sep)
+	r.MaxMIMEHeaderSize = 100
+	_, err = r.NextPart()
+	if err != ErrMessageTooLarge {
+		t.Fatalf("test failed: %v", err)
+	}
+}
+
 func partsFromReader(r *Reader) ([]headerBody, error) {
 	got := []headerBody{}
 	for {