From d8b9db8d3ea655c708735bdda3fbbc96a9c4d7c4 Mon Sep 17 00:00:00 2001
From: slene <vslene@gmail.com>
Date: Sat, 22 Feb 2014 11:58:53 +0800
Subject: [PATCH] move SetSecureCookie / GetSecureCookie to *context.Context
 and alias in Controller

---
 context/context.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 controller.go      | 39 +++------------------------------------
 2 files changed, 48 insertions(+), 36 deletions(-)

diff --git a/context/context.go b/context/context.go
index ee308476..36d566b5 100644
--- a/context/context.go
+++ b/context/context.go
@@ -1,7 +1,14 @@
 package context
 
 import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base64"
+	"fmt"
 	"net/http"
+	"strconv"
+	"strings"
+	"time"
 
 	"github.com/astaxie/beego/middleware"
 )
@@ -59,3 +66,41 @@ func (ctx *Context) GetCookie(key string) string {
 func (ctx *Context) SetCookie(name string, value string, others ...interface{}) {
 	ctx.Output.Cookie(name, value, others...)
 }
+
+// Get secure cookie from request by a given key.
+func (ctx *Context) GetSecureCookie(Secret, key string) (string, bool) {
+	val := ctx.Input.Cookie(key)
+	if val == "" {
+		return "", false
+	}
+
+	parts := strings.SplitN(val, "|", 3)
+
+	if len(parts) != 3 {
+		return "", false
+	}
+
+	vs := parts[0]
+	timestamp := parts[1]
+	sig := parts[2]
+
+	h := hmac.New(sha1.New, []byte(Secret))
+	fmt.Fprintf(h, "%s%s", vs, timestamp)
+
+	if fmt.Sprintf("%02x", h.Sum(nil)) != sig {
+		return "", false
+	}
+	res, _ := base64.URLEncoding.DecodeString(vs)
+	return string(res), true
+}
+
+// Set Secure cookie for response.
+func (ctx *Context) SetSecureCookie(Secret, name, value string, others ...interface{}) {
+	vs := base64.URLEncoding.EncodeToString([]byte(value))
+	timestamp := strconv.FormatInt(time.Now().UnixNano(), 10)
+	h := hmac.New(sha1.New, []byte(Secret))
+	fmt.Fprintf(h, "%s%s", vs, timestamp)
+	sig := fmt.Sprintf("%02x", h.Sum(nil))
+	cookie := strings.Join([]string{vs, timestamp, sig}, "|")
+	ctx.Output.Cookie(name, cookie, others...)
+}
diff --git a/controller.go b/controller.go
index 9d783747..ecb21559 100644
--- a/controller.go
+++ b/controller.go
@@ -2,11 +2,7 @@ package beego
 
 import (
 	"bytes"
-	"crypto/hmac"
-	"crypto/sha1"
-	"encoding/base64"
 	"errors"
-	"fmt"
 	"html/template"
 	"io"
 	"io/ioutil"
@@ -17,7 +13,6 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
-	"time"
 
 	"github.com/astaxie/beego/context"
 	"github.com/astaxie/beego/session"
@@ -417,40 +412,12 @@ func (c *Controller) IsAjax() bool {
 
 // GetSecureCookie returns decoded cookie value from encoded browser cookie values.
 func (c *Controller) GetSecureCookie(Secret, key string) (string, bool) {
-	val := c.Ctx.GetCookie(key)
-	if val == "" {
-		return "", false
-	}
-
-	parts := strings.SplitN(val, "|", 3)
-
-	if len(parts) != 3 {
-		return "", false
-	}
-
-	vs := parts[0]
-	timestamp := parts[1]
-	sig := parts[2]
-
-	h := hmac.New(sha1.New, []byte(Secret))
-	fmt.Fprintf(h, "%s%s", vs, timestamp)
-
-	if fmt.Sprintf("%02x", h.Sum(nil)) != sig {
-		return "", false
-	}
-	res, _ := base64.URLEncoding.DecodeString(vs)
-	return string(res), true
+	return c.Ctx.GetSecureCookie(Secret, key)
 }
 
 // SetSecureCookie puts value into cookie after encoded the value.
-func (c *Controller) SetSecureCookie(Secret, name, val string, age int64) {
-	vs := base64.URLEncoding.EncodeToString([]byte(val))
-	timestamp := strconv.FormatInt(time.Now().UnixNano(), 10)
-	h := hmac.New(sha1.New, []byte(Secret))
-	fmt.Fprintf(h, "%s%s", vs, timestamp)
-	sig := fmt.Sprintf("%02x", h.Sum(nil))
-	cookie := strings.Join([]string{vs, timestamp, sig}, "|")
-	c.Ctx.SetCookie(name, cookie, age, "/")
+func (c *Controller) SetSecureCookie(Secret, name, value string, others ...interface{}) {
+	c.Ctx.SetSecureCookie(Secret, name, value, others...)
 }
 
 // XsrfToken creates a xsrf token string and returns.