From 38307c0b73b3b2fb70386f3c387818fd8ebf09cc Mon Sep 17 00:00:00 2001
From: Claudio Costa <cstcld91@gmail.com>
Date: Mon, 3 Aug 2020 08:25:47 +0200
Subject: [PATCH] [MM-27293] Remove heap pre-allocation for file uploads
 (#15143)

* Validate Content-Length

* Avoid pre-allocating memory for file uploads

* Add small pre-allocation to help mitigate small uploads performance
---
 api4/file.go | 14 ++++++++------
 app/file.go  | 38 ++++++++++++++++++--------------------
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/api4/file.go b/api4/file.go
index f7e48393b1..0b43c0ec1e 100644
--- a/api4/file.go
+++ b/api4/file.go
@@ -7,7 +7,6 @@ import (
 	"bytes"
 	"crypto/subtle"
 	"io"
-	"io/ioutil"
 	"mime"
 	"mime/multipart"
 	"net/http"
@@ -51,8 +50,7 @@ var MEDIA_CONTENT_TYPES = [...]string{
 	"audio/wav",
 }
 
-const maxUploadDrainBytes = (10 * 1024 * 1024) // 10Mb
-const maxMultipartFormDataBytes = 10 * 1024    // 10Kb
+const maxMultipartFormDataBytes = 10 * 1024 // 10Kb
 
 func (api *API) InitFile() {
 	api.BaseRoutes.Files.Handle("", api.ApiSessionRequired(uploadFileStream)).Methods("POST")
@@ -97,9 +95,6 @@ func multipartReader(req *http.Request, stream io.Reader) (*multipart.Reader, er
 }
 
 func uploadFileStream(c *Context, w http.ResponseWriter, r *http.Request) {
-	// Drain any remaining bytes in the request body, up to a limit
-	defer io.CopyN(ioutil.Discard, r.Body, maxUploadDrainBytes)
-
 	if !*c.App.Config().FileSettings.EnableFileAttachments {
 		c.Err = model.NewAppError("uploadFileStream",
 			"api.file.attachments.disabled.app_error",
@@ -120,6 +115,13 @@ func uploadFileStream(c *Context, w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	if r.ContentLength == 0 {
+		c.Err = model.NewAppError("uploadFileStream",
+			"api.file.upload_file.read_request.app_error",
+			nil, "Content-Length should not be 0", http.StatusBadRequest)
+		return
+	}
+
 	timestamp := time.Now()
 	var fileUploadResponse *model.FileUploadResponse
 
diff --git a/app/file.go b/app/file.go
index 4b615ddc3a..523af1121d 100644
--- a/app/file.go
+++ b/app/file.go
@@ -62,7 +62,7 @@ const (
 	ImageThumbnailRatio  = float64(ImageThumbnailHeight) / float64(ImageThumbnailWidth)
 	ImagePreviewWidth    = 1920
 
-	UploadFileInitialBufferSize = 2 * 1024 * 1024 // 2Mb
+	maxUploadInitialBufferSize = 1024 * 1024 // 1Mb
 
 	// Deprecated
 	IMAGE_THUMBNAIL_PIXEL_WIDTH  = 120
@@ -546,8 +546,17 @@ type UploadFileTask struct {
 
 func (t *UploadFileTask) init(a *App) {
 	t.buf = &bytes.Buffer{}
-	t.maxFileSize = *a.Config().FileSettings.MaxFileSize
-	t.limit = *a.Config().FileSettings.MaxFileSize
+	if t.ContentLength > 0 {
+		t.limit = t.ContentLength
+	} else {
+		t.limit = t.maxFileSize
+	}
+
+	if t.ContentLength > 0 && t.ContentLength < maxUploadInitialBufferSize {
+		t.buf.Grow(int(t.ContentLength))
+	} else {
+		t.buf.Grow(maxUploadInitialBufferSize)
+	}
 
 	t.fileinfo = model.NewInfo(filepath.Base(t.Name))
 	t.fileinfo.Id = model.NewId()
@@ -555,19 +564,6 @@ func (t *UploadFileTask) init(a *App) {
 	t.fileinfo.CreateAt = t.Timestamp.UnixNano() / int64(time.Millisecond)
 	t.fileinfo.Path = t.pathPrefix() + t.Name
 
-	// Prepare to read ContentLength if it is known, otherwise limit
-	// ourselves to MaxFileSize. Add an extra byte to check and fail if the
-	// client sent too many bytes.
-	if t.ContentLength > 0 {
-		t.limit = t.ContentLength
-		// Over-Grow the buffer to prevent bytes.ReadFrom from doing it
-		// at the very end.
-		t.buf.Grow(int(t.limit + 1 + bytes.MinRead))
-	} else {
-		// If we don't know the upload size, grow the buffer somewhat
-		// anyway to avoid extra reslicing.
-		t.buf.Grow(UploadFileInitialBufferSize)
-	}
 	t.limitedInput = &io.LimitedReader{
 		R: t.Input,
 		N: t.limit + 1,
@@ -588,14 +584,14 @@ func (a *App) UploadFileX(channelId, name string, input io.Reader,
 	opts ...func(*UploadFileTask)) (*model.FileInfo, *model.AppError) {
 
 	t := &UploadFileTask{
-		ChannelId: filepath.Base(channelId),
-		Name:      filepath.Base(name),
-		Input:     input,
+		ChannelId:   filepath.Base(channelId),
+		Name:        filepath.Base(name),
+		Input:       input,
+		maxFileSize: *a.Config().FileSettings.MaxFileSize,
 	}
 	for _, o := range opts {
 		o(t)
 	}
-	t.init(a)
 
 	if len(*a.Config().FileSettings.DriverName) == 0 {
 		return nil, t.newAppError("api.file.upload_file.storage.app_error",
@@ -606,6 +602,8 @@ func (a *App) UploadFileX(channelId, name string, input io.Reader,
 			"", http.StatusRequestEntityTooLarge, "Length", t.ContentLength, "Limit", t.maxFileSize)
 	}
 
+	t.init(a)
+
 	var aerr *model.AppError
 	if !t.Raw && t.fileinfo.IsImage() {
 		aerr = t.preprocessImage()