From e91afb1938eb7b085946e96b00a5b197c4db37a9 Mon Sep 17 00:00:00 2001
From: hemin <hemin.oak@gmail.com>
Date: Mon, 16 Oct 2017 13:49:05 +0800
Subject: [PATCH] add custom middleware options for beego.Run()

---
 app.go   | 10 +++++++++-
 beego.go | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/app.go b/app.go
index 25ea2a04..1cde9b4b 100644
--- a/app.go
+++ b/app.go
@@ -51,8 +51,10 @@ func NewApp() *App {
 	return app
 }
 
+type MiddleWare func(http.Handler) http.Handler
+
 // Run beego application.
-func (app *App) Run() {
+func (app *App) Run(mws ...MiddleWare) {
 	addr := BConfig.Listen.HTTPAddr
 
 	if BConfig.Listen.HTTPPort != 0 {
@@ -94,6 +96,12 @@ func (app *App) Run() {
 	}
 
 	app.Server.Handler = app.Handlers
+	for i:=len(mws)-1;i>=0;i-- {
+		if mws[i] == nil {
+			continue
+		}
+		app.Server.Handler = mws[i](app.Server.Handler)
+	}
 	app.Server.ReadTimeout = time.Duration(BConfig.Listen.ServerTimeOut) * time.Second
 	app.Server.WriteTimeout = time.Duration(BConfig.Listen.ServerTimeOut) * time.Second
 	app.Server.ErrorLog = logs.GetLogger("HTTP")
diff --git a/beego.go b/beego.go
index 0512dc60..92aeef1e 100644
--- a/beego.go
+++ b/beego.go
@@ -67,6 +67,20 @@ func Run(params ...string) {
 	BeeApp.Run()
 }
 
+func RunWithMiddleWares(addr string, mws ...MiddleWare) {
+	initBeforeHTTPRun()
+
+	strs := strings.Split(addr, ":")
+	if len(strs) > 0 && strs[0] != "" {
+		BConfig.Listen.HTTPAddr = strs[0]
+	}
+	if len(strs) > 1 && strs[1] != "" {
+		BConfig.Listen.HTTPPort, _ = strconv.Atoi(strs[1])
+	}
+
+	BeeApp.Run(mws...)
+}
+
 func initBeforeHTTPRun() {
 	//init hooks
 	AddAPPStartHook(