From f9a31ea00a43c689a1f88193c9d4952ce873efb3 Mon Sep 17 00:00:00 2001
From: astaxie <xiemengjun@gmail.com>
Date: Tue, 6 Aug 2013 23:21:52 +0800
Subject: [PATCH] EnableXSRF

---
 beego.go      |  3 ++-
 config.go     |  3 +++
 controller.go |  3 +--
 router.go     | 21 +++++++++++++++++++++
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/beego.go b/beego.go
index 66ebf6de..d4b1cacf 100644
--- a/beego.go
+++ b/beego.go
@@ -45,7 +45,8 @@ var (
 	HttpServerTimeOut    int64  //set httpserver timeout
 	ErrorsShow           bool   //set weather show errors
 	XSRFKEY              string //set XSRF
-	CopyRequestBody      bool   //When in raw application, You want to the reqeustbody
+	EnableXSRF           bool
+	CopyRequestBody      bool //When in raw application, You want to the reqeustbody
 )
 
 func init() {
diff --git a/config.go b/config.go
index af90bb87..bc6a6b2d 100644
--- a/config.go
+++ b/config.go
@@ -192,6 +192,9 @@ func ParseConfig() (err error) {
 		if xsrfkey := AppConfig.String("xsrfkey"); xsrfkey != "" {
 			XSRFKEY = xsrfkey
 		}
+		if enablexsrf, err := AppConfig.Bool("enablexsrf"); err == nil {
+			EnableXSRF = enablexsrf
+		}
 	}
 	return nil
 }
diff --git a/controller.go b/controller.go
index a563a56b..53a5497c 100644
--- a/controller.go
+++ b/controller.go
@@ -352,7 +352,7 @@ func (c *Controller) XsrfToken() string {
 			h := hmac.New(sha1.New, []byte(XSRFKEY))
 			fmt.Fprintf(h, "%s:%d", c.Ctx.Request.RemoteAddr, time.Now().UnixNano())
 			tok := fmt.Sprintf("%s:%d", h.Sum(nil), time.Now().UnixNano())
-			token := base64.URLEncoding.EncodeToString([]byte(tok))
+			token = base64.URLEncoding.EncodeToString([]byte(tok))
 			c.Ctx.SetCookie("_xsrf", token)
 		}
 		c._xsrf_token = token
@@ -362,7 +362,6 @@ func (c *Controller) XsrfToken() string {
 
 func (c *Controller) CheckXsrfCookie() bool {
 	token := c.GetString("_xsrf")
-
 	if token == "" {
 		token = c.Ctx.Request.Header.Get("X-Xsrftoken")
 	}
diff --git a/router.go b/router.go
index 67690a8a..d606e8e0 100644
--- a/router.go
+++ b/router.go
@@ -452,6 +452,17 @@ func (p *ControllerRegistor) ServeHTTP(rw http.ResponseWriter, r *http.Request)
 		method := vc.MethodByName("Prepare")
 		method.Call(in)
 
+		//if XSRF is Enable then check cookie where there has any cookie in the  request's cookie _csrf
+		if EnableXSRF {
+			method = vc.MethodByName("XsrfToken")
+			method.Call(in)
+			if r.Method == "POST" || r.Method == "DELETE" || r.Method == "PUT" ||
+				(r.Method == "POST" && (r.Form.Get("_method") == "delete" || r.Form.Get("_method") == "put")) {
+				method = vc.MethodByName("CheckXsrfCookie")
+				method.Call(in)
+			}
+		}
+
 		//if response has written,yes don't run next
 		if !w.started {
 			if r.Method == "GET" {
@@ -581,6 +592,16 @@ func (p *ControllerRegistor) ServeHTTP(rw http.ResponseWriter, r *http.Request)
 						method.Call(in)
 						method = vc.MethodByName(mName)
 						method.Call(in)
+						//if XSRF is Enable then check cookie where there has any cookie in the  request's cookie _csrf
+						if EnableXSRF {
+							method = vc.MethodByName("XsrfToken")
+							method.Call(in)
+							if r.Method == "POST" || r.Method == "DELETE" || r.Method == "PUT" ||
+								(r.Method == "POST" && (r.Form.Get("_method") == "delete" || r.Form.Get("_method") == "put")) {
+								method = vc.MethodByName("CheckXsrfCookie")
+								method.Call(in)
+							}
+						}
 						if !w.started {
 							if AutoRender {
 								method = vc.MethodByName("Render")