From d2a16ff8f6d39c539d784a381adaf6ca309b4af8 Mon Sep 17 00:00:00 2001
From: astaxie <xiemengjun@gmail.com>
Date: Mon, 8 Jul 2013 16:17:08 +0800
Subject: [PATCH] fix #26 add xsrf function

---
 beego.go      |  2 ++
 config.go     |  3 +++
 context.go    |  8 +++++++
 controller.go | 59 +++++++++++++++++++++++++++++++++++++++++++++------
 4 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/beego.go b/beego.go
index 5a638d88..b6dda501 100644
--- a/beego.go
+++ b/beego.go
@@ -44,6 +44,7 @@ var (
 	EnbaleHotUpdate      bool //enable HotUpdate default is false
 	HttpServerTimeOut    int64
 	ErrorsShow           bool
+	XSRFKEY              string
 )
 
 func init() {
@@ -72,6 +73,7 @@ func init() {
 	AppConfigPath = path.Join(AppPath, "conf", "app.conf")
 	HttpServerTimeOut = 0
 	ErrorsShow = true
+	XSRFKEY = "beegoxsrf"
 	ParseConfig()
 }
 
diff --git a/config.go b/config.go
index 6aebbc2a..f652a38b 100644
--- a/config.go
+++ b/config.go
@@ -186,6 +186,9 @@ func ParseConfig() (err error) {
 		if errorsshow, err := AppConfig.Bool("errorsshow"); err == nil {
 			ErrorsShow = errorsshow
 		}
+		if xsrfkey := AppConfig.String("xsrfkey"); xsrfkey != "" {
+			XSRFKEY = xsrfkey
+		}
 	}
 	return nil
 }
diff --git a/context.go b/context.go
index 4d8ace94..01428c34 100644
--- a/context.go
+++ b/context.go
@@ -101,3 +101,11 @@ var cookieValueSanitizer = strings.NewReplacer("\n", " ", "\r", " ", ";", " ")
 func sanitizeValue(v string) string {
 	return cookieValueSanitizer.Replace(v)
 }
+
+func (ctx *Context) GetCookie(key string) string {
+	keycookie, err := ctx.Request.Cookie(key)
+	if err != nil {
+		return ""
+	}
+	return keycookie.Value
+}
diff --git a/controller.go b/controller.go
index c3caa2a4..c37109b8 100644
--- a/controller.go
+++ b/controller.go
@@ -4,9 +4,13 @@ import (
 	"bytes"
 	"compress/gzip"
 	"compress/zlib"
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base64"
 	"encoding/json"
 	"encoding/xml"
 	"errors"
+	"fmt"
 	"github.com/astaxie/beego/session"
 	"html/template"
 	"io"
@@ -18,16 +22,18 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"time"
 )
 
 type Controller struct {
-	Ctx        *Context
-	Data       map[interface{}]interface{}
-	ChildName  string
-	TplNames   string
-	Layout     string
-	TplExt     string
-	CruSession session.SessionStore
+	Ctx         *Context
+	Data        map[interface{}]interface{}
+	ChildName   string
+	TplNames    string
+	Layout      string
+	TplExt      string
+	_xsrf_token string
+	CruSession  session.SessionStore
 }
 
 type ControllerInterface interface {
@@ -331,3 +337,42 @@ func (c *Controller) DelSession(name interface{}) {
 func (c *Controller) IsAjax() bool {
 	return (c.Ctx.Request.Header.Get("HTTP_X_REQUESTED_WITH") == "XMLHttpRequest")
 }
+
+func (c *Controller) XsrfToken() string {
+	if c._xsrf_token == "" {
+		token := c.Ctx.GetCookie("_xsrf")
+		if token == "" {
+			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))
+			c.Ctx.SetCookie("_xsrf", token)
+		}
+		c._xsrf_token = token
+	}
+	return c._xsrf_token
+}
+
+func (c *Controller) CheckXsrfCookie() bool {
+	token := c.GetString("_xsrf")
+
+	if token == "" {
+		token = c.Ctx.Request.Header.Get("X-Xsrftoken")
+	}
+	if token == "" {
+		token = c.Ctx.Request.Header.Get("X-Csrftoken")
+	}
+	if token == "" {
+		c.Ctx.Abort(403, "'_xsrf' argument missing from POST")
+	}
+
+	if c._xsrf_token != token {
+		c.Ctx.Abort(403, "XSRF cookie does not match POST argument")
+	}
+	return true
+}
+
+func (c *Controller) XsrfFormHtml() string {
+	return "<input type=\"hidden\" name=\"_xsrf\" value=\"" +
+		c._xsrf_token + "\"/>"
+}