From 946a42c021f3688ca411315b71fe989152df7f27 Mon Sep 17 00:00:00 2001
From: Chenrui <631807682@qq.com>
Date: Wed, 8 Jul 2020 17:14:52 +0800
Subject: [PATCH] fix: response http 413 when body size larger then MaxMemory.

---
 router.go      |  4 ++++
 router_test.go | 29 +++++++++++++++++++++++++----
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/router.go b/router.go
index e71366b4..9b257391 100644
--- a/router.go
+++ b/router.go
@@ -707,6 +707,10 @@ func (p *ControllerRegister) ServeHTTP(rw http.ResponseWriter, r *http.Request)
 
 	if r.Method != http.MethodGet && r.Method != http.MethodHead {
 		if BConfig.CopyRequestBody && !context.Input.IsUpload() {
+			if r.ContentLength > BConfig.MaxMemory {
+				exception("413", context)
+				goto Admin
+			}
 			context.Input.CopyBody(BConfig.MaxMemory)
 		}
 		context.Input.ParseFormOrMulitForm(BConfig.MaxMemory)
diff --git a/router_test.go b/router_test.go
index 2797b33a..e49f38db 100644
--- a/router_test.go
+++ b/router_test.go
@@ -15,6 +15,7 @@
 package beego
 
 import (
+	"bytes"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -71,7 +72,6 @@ func (tc *TestController) GetEmptyBody() {
 	tc.Ctx.Output.Body(res)
 }
 
-
 type JSONController struct {
 	Controller
 }
@@ -656,17 +656,14 @@ func beegoBeforeRouter1(ctx *context.Context) {
 	ctx.WriteString("|BeforeRouter1")
 }
 
-
 func beegoBeforeExec1(ctx *context.Context) {
 	ctx.WriteString("|BeforeExec1")
 }
 
-
 func beegoAfterExec1(ctx *context.Context) {
 	ctx.WriteString("|AfterExec1")
 }
 
-
 func beegoFinishRouter1(ctx *context.Context) {
 	ctx.WriteString("|FinishRouter1")
 }
@@ -709,3 +706,27 @@ func TestYAMLPrepare(t *testing.T) {
 		t.Errorf(w.Body.String())
 	}
 }
+
+func TestRouterEntityTooLargeCopyBody(t *testing.T) {
+	_MaxMemory := BConfig.MaxMemory
+	_CopyRequestBody := BConfig.CopyRequestBody
+	BConfig.CopyRequestBody = true
+	BConfig.MaxMemory = 20
+
+	b := bytes.NewBuffer([]byte("barbarbarbarbarbarbarbarbarbar"))
+	r, _ := http.NewRequest("POST", "/user/123", b)
+	w := httptest.NewRecorder()
+
+	handler := NewControllerRegister()
+	handler.Post("/user/:id", func(ctx *context.Context) {
+		ctx.Output.Body([]byte(ctx.Input.Param(":id")))
+	})
+	handler.ServeHTTP(w, r)
+
+	BConfig.CopyRequestBody = _CopyRequestBody
+	BConfig.MaxMemory = _MaxMemory
+
+	if w.Code != 413 {
+		t.Errorf("TestRouterRequestEntityTooLarge can't run")
+	}
+}