From 4951314837af08eaee05d07e9b6bb0ea30446678 Mon Sep 17 00:00:00 2001
From: Jared Folkins <jfolkins@gmail.com>
Date: Thu, 13 Mar 2014 22:30:12 -0700
Subject: [PATCH] added FlashName,FlashSeperator, & Tests

---
 .gitignore     |  2 ++
 config.go      | 13 +++++++++++++
 config_test.go | 15 +++++++++++++++
 flash.go       | 17 ++++++-----------
 flash_test.go  | 40 ++++++++++++++++++++++++++++++++++++++++
 5 files changed, 76 insertions(+), 11 deletions(-)
 create mode 100644 config_test.go
 create mode 100644 flash_test.go

diff --git a/.gitignore b/.gitignore
index e43b0f98..482e34b1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,3 @@
 .DS_Store
+*.swp
+*.swo
diff --git a/config.go b/config.go
index 7c21d696..54de325c 100644
--- a/config.go
+++ b/config.go
@@ -58,6 +58,8 @@ var (
 	EnableAdmin            bool   // flag of enable admin module to log every request info.
 	AdminHttpAddr          string // http server configurations for admin module.
 	AdminHttpPort          int
+	FlashName              string // name of the flash variable found in response header and cookie
+	FlashSeperator         string // used to seperate flash key:value
 )
 
 func init() {
@@ -123,6 +125,9 @@ func init() {
 	AdminHttpAddr = "127.0.0.1"
 	AdminHttpPort = 8088
 
+	FlashName = "BEEGO_FLASH"
+	FlashSeperator = "BEEGOFLASH"
+
 	runtime.GOMAXPROCS(runtime.NumCPU())
 
 	// init BeeLogger
@@ -271,6 +276,14 @@ func ParseConfig() (err error) {
 			BeegoServerName = serverName
 		}
 
+		if flashname := AppConfig.String("FlashName"); flashname != "" {
+			FlashName = flashname
+		}
+
+		if flashseperator := AppConfig.String("FlashSeperator"); flashseperator != "" {
+			FlashSeperator = flashseperator
+		}
+
 		if sd := AppConfig.String("StaticDir"); sd != "" {
 			for k := range StaticDir {
 				delete(StaticDir, k)
diff --git a/config_test.go b/config_test.go
new file mode 100644
index 00000000..d400fd5d
--- /dev/null
+++ b/config_test.go
@@ -0,0 +1,15 @@
+package beego
+
+import (
+	"testing"
+)
+
+func TestDefaults(t *testing.T) {
+	if FlashName != "BEEGO_FLASH" {
+		t.Errorf("FlashName was not set to default.")
+	}
+
+	if FlashSeperator != "BEEGOFLASH" {
+		t.Errorf("FlashName was not set to default.")
+	}
+}
diff --git a/flash.go b/flash.go
index f435c155..3cb4a66a 100644
--- a/flash.go
+++ b/flash.go
@@ -6,9 +6,6 @@ import (
 	"strings"
 )
 
-// the separation string when encoding flash data.
-const BEEGO_FLASH_SEP = "#BEEGOFLASH#"
-
 // FlashData is a tools to maintain data when using across request.
 type FlashData struct {
 	Data map[string]string
@@ -54,29 +51,27 @@ func (fd *FlashData) Store(c *Controller) {
 	c.Data["flash"] = fd.Data
 	var flashValue string
 	for key, value := range fd.Data {
-		flashValue += "\x00" + key + BEEGO_FLASH_SEP + value + "\x00"
+		flashValue += "\x00" + key + "\x23" + FlashSeperator + "\x23" + value + "\x00"
 	}
-	c.Ctx.SetCookie("BEEGO_FLASH", url.QueryEscape(flashValue), 0, "/")
+	c.Ctx.SetCookie(FlashName, url.QueryEscape(flashValue), 0, "/")
 }
 
 // ReadFromRequest parsed flash data from encoded values in cookie.
 func ReadFromRequest(c *Controller) *FlashData {
-	flash := &FlashData{
-		Data: make(map[string]string),
-	}
-	if cookie, err := c.Ctx.Request.Cookie("BEEGO_FLASH"); err == nil {
+	flash := NewFlash()
+	if cookie, err := c.Ctx.Request.Cookie(FlashName); err == nil {
 		v, _ := url.QueryUnescape(cookie.Value)
 		vals := strings.Split(v, "\x00")
 		for _, v := range vals {
 			if len(v) > 0 {
-				kv := strings.Split(v, BEEGO_FLASH_SEP)
+				kv := strings.Split(v, FlashSeperator)
 				if len(kv) == 2 {
 					flash.Data[kv[0]] = kv[1]
 				}
 			}
 		}
 		//read one time then delete it
-		c.Ctx.SetCookie("BEEGO_FLASH", "", -1, "/")
+		c.Ctx.SetCookie(FlashName, "", -1, "/")
 	}
 	c.Data["flash"] = flash.Data
 	return flash
diff --git a/flash_test.go b/flash_test.go
new file mode 100644
index 00000000..d329bae4
--- /dev/null
+++ b/flash_test.go
@@ -0,0 +1,40 @@
+package beego
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+type TestFlashController struct {
+	Controller
+}
+
+func (this *TestFlashController) TestWriteFlash() {
+	flash := NewFlash()
+	flash.Notice("TestFlashString")
+	flash.Store(&this.Controller)
+	// we choose to serve json because we don't want to load a template html file
+	this.ServeJson(true)
+}
+
+func TestFlashHeader(t *testing.T) {
+	// create fake GET request
+	r, _ := http.NewRequest("GET", "/", nil)
+	w := httptest.NewRecorder()
+
+	// setup the handler
+	handler := NewControllerRegistor()
+	handler.Add("/", &TestFlashController{}, "get:TestWriteFlash")
+	handler.ServeHTTP(w, r)
+
+	// get the Set-Cookie value
+	sc := w.Header().Get("Set-Cookie")
+	// match for the expected header
+	res := strings.Contains(sc, "BEEGO_FLASH=%00notice%23BEEGOFLASH%23TestFlashString%00")
+	// validate the assertion
+	if res != true {
+		t.Errorf("TestFlashHeader() unable to validate flash message")
+	}
+}