mime/multipart: Allow ReadForm to process large non-file parts

Allow the memory limit passed into ReadForm to be used as the
memory limit for processing non-file form data as well as file
form data, rather than the existing behaviour of the memory limit
only applying to the file parts and the non-file parts being
arbitrarily limited to 10MB.

This ensures backwards compatibility while still providing the
user with control over the amount of non-file data that can be
processed instead of enforcing an arbitrary 10MB limit.

Change-Id: I53c09eae00147d3ff2d6bdfd4e50949267932c3d
Reviewed-on: https://go-review.googlesource.com/38195
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

src/mime/multipart/formdata.go

@@ -13,13 +13,20 @@ import (
 	"os"
 )
 
+// ErrMessageTooLarge is returned by ReadForm if the message form
+// data is too large to be processed.
+var ErrMessageTooLarge = errors.New("multipart: message too large")
+
 // TODO(adg,bradfitz): find a way to unify the DoS-prevention strategy here
 // with that of the http package's ParseForm.
 
 // ReadForm parses an entire multipart message whose parts have
 // a Content-Disposition of "form-data".
-// It stores up to maxMemory bytes of the file parts in memory
-// and the remainder on disk in temporary files.
+// It stores up to maxMemory bytes + 10MB (reserved for non-file parts)
+// in memory. File parts which can't be stored in memory will be stored on
+// disk in temporary files.
+// It returns ErrMessageTooLarge if all non-file parts can't be stored in
+// memory.
 func (r *Reader) ReadForm(maxMemory int64) (*Form, error) {
 	return r.readForm(maxMemory)
 }
@@ -32,7 +39,8 @@ func (r *Reader) readForm(maxMemory int64) (_ *Form, err error) {
 		}
 	}()
 
-	maxValueBytes := int64(10 << 20) // 10 MB is a lot of text.
+	// Reserve an additional 10 MB for non-file parts.
+	maxValueBytes := maxMemory + int64(10<<20)
 	for {
 		p, err := r.NextPart()
 		if err == io.EOF {
@@ -52,13 +60,13 @@ func (r *Reader) readForm(maxMemory int64) (_ *Form, err error) {
 
 		if filename == "" {
 			// value, store as string in memory
-			n, err := io.CopyN(&b, p, maxValueBytes)
+			n, err := io.CopyN(&b, p, maxValueBytes+1)
 			if err != nil && err != io.EOF {
 				return nil, err
 			}
 			maxValueBytes -= n
-			if maxValueBytes == 0 {
-				return nil, errors.New("multipart: message too large")
+			if maxValueBytes < 0 {
+				return nil, ErrMessageTooLarge
 			}
 			form.Value[name] = append(form.Value[name], b.String())
 			continue
@@ -93,6 +101,7 @@ func (r *Reader) readForm(maxMemory int64) (_ *Form, err error) {
 			fh.content = b.Bytes()
 			fh.Size = int64(len(fh.content))
 			maxMemory -= n
+			maxValueBytes -= n
 		}
 		form.File[name] = append(form.File[name], fh)
 	}

src/mime/multipart/formdata_test.go

@@ -8,14 +8,12 @@ import (
 	"bytes"
 	"io"
 	"os"
-	"regexp"
 	"strings"
 	"testing"
 )
 
 func TestReadForm(t *testing.T) {
-	testBody := regexp.MustCompile("\n").ReplaceAllString(message, "\r\n")
-	b := strings.NewReader(testBody)
+	b := strings.NewReader(strings.Replace(message, "\n", "\r\n", -1))
 	r := NewReader(b, boundary)
 	f, err := r.ReadForm(25)
 	if err != nil {
@@ -127,3 +125,44 @@ func (r *failOnReadAfterErrorReader) Read(p []byte) (n int, err error) {
 	r.sawErr = err
 	return
 }
+
+// TestReadForm_NonFileMaxMemory asserts that the ReadForm maxMemory limit is applied
+// while processing non-file form data as well as file form data.
+func TestReadForm_NonFileMaxMemory(t *testing.T) {
+	largeTextValue := strings.Repeat("1", (10<<20)+25)
+	message := `--MyBoundary
+Content-Disposition: form-data; name="largetext"
+
+` + largeTextValue + `
+--MyBoundary--
+`
+
+	testBody := strings.Replace(message, "\n", "\r\n", -1)
+	testCases := []struct {
+		name      string
+		maxMemory int64
+		err       error
+	}{
+		{"smaller", 50, nil},
+		{"exact-fit", 25, nil},
+		{"too-large", 0, ErrMessageTooLarge},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			b := strings.NewReader(testBody)
+			r := NewReader(b, boundary)
+			f, err := r.ReadForm(tc.maxMemory)
+			if err == nil {
+				defer f.RemoveAll()
+			}
+			if tc.err != err {
+				t.Fatalf("ReadForm error - got: %v; expected: %v", tc.err, err)
+			}
+			if err == nil {
+				if g := f.Value["largetext"][0]; g != largeTextValue {
+					t.Errorf("largetext mismatch: got size: %v, expected size: %v", len(g), len(largeTextValue))
+				}
+			}
+		})
+	}
+}