From c510926cb89d99fca4c22d967e93cbe14b82e948 Mon Sep 17 00:00:00 2001
From: Anker Jam <jianzhiyao020@qq.com>
Date: Sun, 18 Oct 2020 23:18:13 +0800
Subject: [PATCH] =?UTF-8?q?fix=204224=EF=BC=9Aform=20entity=20too=20large?=
 =?UTF-8?q?=20casue=20run=20out=20of=20memory?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 adapter/context/input.go    |  2 +-
 server/web/context/input.go |  5 +++--
 server/web/router.go        | 10 ++++++++--
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/adapter/context/input.go b/adapter/context/input.go
index 4d62d3c1..51bb9ea5 100644
--- a/adapter/context/input.go
+++ b/adapter/context/input.go
@@ -266,7 +266,7 @@ func (input *BeegoInput) SetData(key, val interface{}) {
 
 // ParseFormOrMulitForm parseForm or parseMultiForm based on Content-type
 func (input *BeegoInput) ParseFormOrMulitForm(maxMemory int64) error {
-	return (*context.BeegoInput)(input).ParseFormOrMulitForm(maxMemory)
+	return (*context.BeegoInput)(input).ParseFormOrMultiForm(maxMemory)
 }
 
 // Bind data from request.Form[key] to dest
diff --git a/server/web/context/input.go b/server/web/context/input.go
index 641e15cc..21cae0b0 100644
--- a/server/web/context/input.go
+++ b/server/web/context/input.go
@@ -420,10 +420,11 @@ func (input *BeegoInput) SetData(key, val interface{}) {
 	input.data[key] = val
 }
 
-// ParseFormOrMulitForm parseForm or parseMultiForm based on Content-type
-func (input *BeegoInput) ParseFormOrMulitForm(maxMemory int64) error {
+// ParseFormOrMultiForm parseForm or parseMultiForm based on Content-type
+func (input *BeegoInput) ParseFormOrMultiForm(maxMemory int64) error {
 	// Parse the body depending on the content type.
 	if strings.Contains(input.Header("Content-Type"), "multipart/form-data") {
+		input.Context.Request.Body = http.MaxBytesReader(input.Context.ResponseWriter, input.Context.Request.Body, maxMemory)
 		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 0f383f99..ca0918a0 100644
--- a/server/web/router.go
+++ b/server/web/router.go
@@ -666,6 +666,7 @@ func (p *ControllerRegister) ServeHTTP(rw http.ResponseWriter, r *http.Request)
 }
 
 func (p *ControllerRegister) serveHttp(ctx *beecontext.Context) {
+	var err error
 	startTime := time.Now()
 	r := ctx.Request
 	rw := ctx.ResponseWriter.ResponseWriter
@@ -718,12 +719,17 @@ func (p *ControllerRegister) serveHttp(ctx *beecontext.Context) {
 			}
 			ctx.Input.CopyBody(p.cfg.MaxMemory)
 		}
-		ctx.Input.ParseFormOrMulitForm(p.cfg.MaxMemory)
+
+		err = ctx.Input.ParseFormOrMultiForm(p.cfg.MaxMemory)
+		if err != nil {
+			logs.Error(errors.New("payload too large"))
+			exception("413", ctx)
+			goto Admin
+		}
 	}
 
 	// session init
 	if p.cfg.WebConfig.Session.SessionOn {
-		var err error
 		ctx.Input.CruSession, err = GlobalSessions.SessionStart(rw, r)
 		if err != nil {
 			logs.Error(err)