From d26683799a6b2de721530696dd6dfb6494ff1ea9 Mon Sep 17 00:00:00 2001
From: Ming Deng <mingflycash@gmail.com>
Date: Tue, 20 Oct 2020 22:06:24 +0800
Subject: [PATCH] add MaxUploadFile to provide more safety uploading controll

---
 server/web/config.go        | 19 ++++++++++++-------
 server/web/context/input.go |  3 +--
 server/web/router.go        | 11 ++++++++++-
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/server/web/config.go b/server/web/config.go
index 15a2dffe..10138e63 100644
--- a/server/web/config.go
+++ b/server/web/config.go
@@ -43,12 +43,16 @@ type Config struct {
 	RecoverFunc         func(*context.Context, *Config)
 	CopyRequestBody     bool
 	EnableGzip          bool
-	MaxMemory           int64
-	EnableErrorsShow    bool
-	EnableErrorsRender  bool
-	Listen              Listen
-	WebConfig           WebConfig
-	Log                 LogConfig
+	// MaxMemory and MaxUploadSize are used to limit the request body
+	// if the request is not uploading file, MaxMemory is the max size of request body
+	// if the request is uploading file, MaxUploadSize is the max size of request body
+	MaxMemory          int64
+	MaxUploadSize      int64
+	EnableErrorsShow   bool
+	EnableErrorsRender bool
+	Listen             Listen
+	WebConfig          WebConfig
+	Log                LogConfig
 }
 
 // Listen holds for http and https related config
@@ -215,6 +219,7 @@ func newBConfig() *Config {
 		CopyRequestBody:    false,
 		EnableGzip:         false,
 		MaxMemory:          1 << 26, // 64MB
+		MaxUploadSize:      1 << 30, // 1GB
 		EnableErrorsShow:   true,
 		EnableErrorsRender: true,
 		Listen: Listen{
@@ -302,7 +307,7 @@ func assignConfig(ac config.Configer) error {
 
 	err := ac.Unmarshaler("", BConfig)
 	if err != nil {
-		_, _ = fmt.Fprintln(os.Stderr, fmt.Sprintf("Unmarshaler config file to BConfig failed. " +
+		_, _ = fmt.Fprintln(os.Stderr, fmt.Sprintf("Unmarshaler config file to BConfig failed. "+
 			"And if you are working on v1.x config file, please ignore this, err: %s", err))
 		return err
 	}
diff --git a/server/web/context/input.go b/server/web/context/input.go
index 027ad527..504838a3 100644
--- a/server/web/context/input.go
+++ b/server/web/context/input.go
@@ -423,8 +423,7 @@ func (input *BeegoInput) SetData(key, val interface{}) {
 // ParseFormOrMultiForm parseForm or parseMultiForm based on Content-type
 func (input *BeegoInput) ParseFormOrMultiForm(maxMemory int64) error {
 	// Parse the body depending on the content type.
-	input.Context.Request.Body = http.MaxBytesReader(input.Context.ResponseWriter, input.Context.Request.Body, maxMemory)
-	if strings.Contains(input.Header("Content-Type"), "multipart/form-data") {
+	if input.IsUpload() {
 		if err := input.Context.Request.ParseMultipartForm(maxMemory); err != nil {
 			return errors.New("Error parsing request body:" + err.Error())
 		}
diff --git a/server/web/router.go b/server/web/router.go
index 515abcb0..7bb89d82 100644
--- a/server/web/router.go
+++ b/server/web/router.go
@@ -710,7 +710,12 @@ func (p *ControllerRegister) serveHttp(ctx *beecontext.Context) {
 	}
 
 	if r.Method != http.MethodGet && r.Method != http.MethodHead {
-		if p.cfg.CopyRequestBody && !ctx.Input.IsUpload() {
+
+		if ctx.Input.IsUpload() {
+			ctx.Input.Context.Request.Body = http.MaxBytesReader(ctx.Input.Context.ResponseWriter,
+				ctx.Input.Context.Request.Body,
+				p.cfg.MaxUploadSize)
+		} else if p.cfg.CopyRequestBody {
 			// connection will close if the incoming data are larger (RFC 7231, 6.5.11)
 			if r.ContentLength > p.cfg.MaxMemory {
 				logs.Error(errors.New("payload too large"))
@@ -718,6 +723,10 @@ func (p *ControllerRegister) serveHttp(ctx *beecontext.Context) {
 				goto Admin
 			}
 			ctx.Input.CopyBody(p.cfg.MaxMemory)
+		} else {
+			ctx.Input.Context.Request.Body = http.MaxBytesReader(ctx.Input.Context.ResponseWriter,
+				ctx.Input.Context.Request.Body,
+				p.cfg.MaxMemory)
 		}
 
 		err = ctx.Input.ParseFormOrMultiForm(p.cfg.MaxMemory)