From d6759e7a059f4208f07aa781402841d7ddaaef96 Mon Sep 17 00:00:00 2001
From: Damien Neil <dneil@google.com>
Date: Fri, 10 Mar 2023 14:21:05 -0800
Subject: [PATCH] [release-branch.go1.19] net/textproto: avoid overpredicting
 the number of MIME header keys
 
Reviewed-on: https://team-review.git.corp.google.com/c/golang/go-private/+/1802452
Run-TryBot: Damien Neil <dneil@google.com>
Reviewed-by: Roland Shoemaker <bracewell@google.com>
Reviewed-by: Julie Qiu <julieqiu@google.com>
(cherry picked from commit f739f080a72fd5b06d35c8e244165159645e2ed6)
Reviewed-on: https://team-review.git.corp.google.com/c/golang/go-private/+/1802393
Reviewed-by: Damien Neil <dneil@google.com>
Run-TryBot: Roland Shoemaker <bracewell@google.com>
Change-Id: I675451438d619a9130360c56daf529559004903f
Reviewed-on: https://go-review.googlesource.com/c/go/+/481982
Run-TryBot: Michael Knyszek <mknyszek@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Matthew Dempsky <mdempsky@google.com>
Auto-Submit: Michael Knyszek <mknyszek@google.com>

Upstream-Status: Backport [https://github.com/golang/go/commit/d6759e7a059f4208f07aa781402841d7ddaaef96]
CVE: CVE-2023-24534
Signed-off-by: Vivek Kumbhar <vkumbhar@mvista.com>

---
 src/bytes/bytes.go               | 14 ++++++++
 src/net/textproto/reader.go      | 30 ++++++++++------
 src/net/textproto/reader_test.go | 59 ++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 11 deletions(-)

diff --git a/src/bytes/bytes.go b/src/bytes/bytes.go
index ce52649..95ff31c 100644
--- a/src/bytes/bytes.go
+++ b/src/bytes/bytes.go
@@ -1174,3 +1174,17 @@ func Index(s, sep []byte) int {
 	}
 	return -1
 }
+
+// Cut slices s around the first instance of sep,
+// returning the text before and after sep.
+// The found result reports whether sep appears in s.
+// If sep does not appear in s, cut returns s, nil, false.
+//
+// Cut returns slices of the original slice s, not copies.
+func Cut(s, sep []byte) (before, after []byte, found bool) {
+	if i := Index(s, sep); i >= 0 {
+		return s[:i], s[i+len(sep):], true
+	}
+	return s, nil, false
+}
+
diff --git a/src/net/textproto/reader.go b/src/net/textproto/reader.go
index 6a680f4..fcbede8 100644
--- a/src/net/textproto/reader.go
+++ b/src/net/textproto/reader.go
@@ -493,8 +493,11 @@ func readMIMEHeader(r *Reader, lim int64) (MIMEHeader, error) {
 	// large one ahead of time which we'll cut up into smaller
 	// slices. If this isn't big enough later, we allocate small ones.
 	var strs []string
-	hint := r.upcomingHeaderNewlines()
+	hint := r.upcomingHeaderKeys()
 	if hint > 0 {
+		if hint > 1000 {
+			hint = 1000 // set a cap to avoid overallocation
+		}
 		strs = make([]string, hint)
 	}
 
@@ -589,9 +592,11 @@ func mustHaveFieldNameColon(line []byte) error {
 	return nil
 }
 
-// upcomingHeaderNewlines returns an approximation of the number of newlines
+var nl = []byte("\n")
+
+// upcomingHeaderKeys returns an approximation of the number of keys
 // that will be in this header. If it gets confused, it returns 0.
-func (r *Reader) upcomingHeaderNewlines() (n int) {
+func (r *Reader) upcomingHeaderKeys() (n int) {
 	// Try to determine the 'hint' size.
 	r.R.Peek(1) // force a buffer load if empty
 	s := r.R.Buffered()
@@ -599,17 +604,20 @@ func (r *Reader) upcomingHeaderNewlines() (n int) {
 		return
 	}
 	peek, _ := r.R.Peek(s)
-	for len(peek) > 0 {
-		i := bytes.IndexByte(peek, '\n')
-		if i < 3 {
-			// Not present (-1) or found within the next few bytes,
-			// implying we're at the end ("\r\n\r\n" or "\n\n")
-			return
+	for len(peek) > 0 && n < 1000 {
+		var line []byte
+		line, peek, _ = bytes.Cut(peek, nl)
+		if len(line) == 0 || (len(line) == 1 && line[0] == '\r') {
+			// Blank line separating headers from the body.
+			break
+		}
+		if line[0] == ' ' || line[0] == '\t' {
+			// Folded continuation of the previous line.
+			continue
 		}
 		n++
-		peek = peek[i+1:]
 	}
-	return
+	return n
 }
 
 // CanonicalMIMEHeaderKey returns the canonical format of the
diff --git a/src/net/textproto/reader_test.go b/src/net/textproto/reader_test.go
index 3124d43..3ae0de1 100644
--- a/src/net/textproto/reader_test.go
+++ b/src/net/textproto/reader_test.go
@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"io"
 	"reflect"
+	"runtime"
 	"strings"
 	"testing"
 )
@@ -127,6 +128,42 @@ func TestReadMIMEHeaderSingle(t *testing.T) {
 	}
 }
 
+// TestReaderUpcomingHeaderKeys is testing an internal function, but it's very
+// difficult to test well via the external API.
+func TestReaderUpcomingHeaderKeys(t *testing.T) {
+	for _, test := range []struct {
+		input string
+		want  int
+	}{{
+		input: "",
+		want:  0,
+	}, {
+		input: "A: v",
+		want:  1,
+	}, {
+		input: "A: v\r\nB: v\r\n",
+		want:  2,
+	}, {
+		input: "A: v\nB: v\n",
+		want:  2,
+	}, {
+		input: "A: v\r\n  continued\r\n  still continued\r\nB: v\r\n\r\n",
+		want:  2,
+	}, {
+		input: "A: v\r\n\r\nB: v\r\nC: v\r\n",
+		want:  1,
+	}, {
+		input: "A: v" + strings.Repeat("\n", 1000),
+		want:  1,
+	}} {
+		r := reader(test.input)
+		got := r.upcomingHeaderKeys()
+		if test.want != got {
+			t.Fatalf("upcomingHeaderKeys(%q): %v; want %v", test.input, got, test.want)
+		}
+	}
+}
+
 func TestReadMIMEHeaderNoKey(t *testing.T) {
 	r := reader(": bar\ntest-1: 1\n\n")
 	m, err := r.ReadMIMEHeader()
@@ -223,6 +260,28 @@ func TestReadMIMEHeaderTrimContinued(t *testing.T) {
 	}
 }
 
+// Test that reading a header doesn't overallocate. Issue 58975.
+func TestReadMIMEHeaderAllocations(t *testing.T) {
+	var totalAlloc uint64
+	const count = 200
+	for i := 0; i < count; i++ {
+		r := reader("A: b\r\n\r\n" + strings.Repeat("\n", 4096))
+		var m1, m2 runtime.MemStats
+		runtime.ReadMemStats(&m1)
+		_, err := r.ReadMIMEHeader()
+		if err != nil {
+			t.Fatalf("ReadMIMEHeader: %v", err)
+		}
+		runtime.ReadMemStats(&m2)
+		totalAlloc += m2.TotalAlloc - m1.TotalAlloc
+	}
+	// 32k is large and we actually allocate substantially less,
+	// but prior to the fix for #58975 we allocated ~400k in this case.
+	if got, want := totalAlloc/count, uint64(32768); got > want {
+		t.Fatalf("ReadMIMEHeader allocated %v bytes, want < %v", got, want)
+	}
+}
+
 type readResponseTest struct {
 	in       string
 	inCode   int
-- 
2.25.1

