Embrace https://github.com/kataras/iris/issues/250 - Flash messages available inside request life

context.go

@@ -54,6 +54,8 @@ const (
 
 	// stopExecutionPosition used inside the Context, is the number which shows us that the context's middleware manualy stop the execution
 	stopExecutionPosition = 255
+	// used inside GetFlash to store the lifetime request flash messages
+	flashMessagesStoreContextKey = "_iris_flash_messages_"
 )
 
 // this pool is used everywhere needed in the iris for example inside party-> Static
@@ -718,26 +720,48 @@ func (ctx *Context) RemoveCookie(name string) {
 }
 
 // GetFlash get a flash message by it's key
-// after this action the messages is removed
-// returns string, if the cookie doesn't exists the string is empty
-func (ctx *Context) GetFlash(key string) string {
-	val, err := ctx.GetFlashBytes(key)
-	if err != nil {
-		return ""
-	}
-	return string(val)
-}
+// after the request lifetime the value is removed
+// returns the value as string and an error
+//
+// if the cookie doesn't exists the string is empty and the error is filled
+func (ctx *Context) GetFlash(key string) (value string, err error) {
+
+	// first check if flash exists from this request's lifetime, if yes return that else continue to get the cookie
+	storeExists := false
+	if messages := ctx.Get(flashMessagesStoreContextKey); messages != nil {
+		m, isMap := messages.(map[string]string)
+		if !isMap {
+			return "", fmt.Errorf("Messages request's  store is not a map[string]string. This suppose will never happen, please report this bug.")
+		}
+		storeExists = true // in order to skip the check later
+		for k, v := range m {
+			if k == key && v != "" {
+				return v, nil
+			}
+		}
+	}
 
-// GetFlashBytes get a flash message by it's key
-// after this action the messages is removed
-// returns []byte along with an error if the cookie doesn't exists or decode fails
-func (ctx *Context) GetFlashBytes(key string) (value []byte, err error) {
 	cookieValue := string(ctx.RequestCtx.Request.Header.Cookie(key))
+
 	if cookieValue == "" {
 		err = errFlashNotFound.Return()
 	} else {
-		value, err = base64.URLEncoding.DecodeString(cookieValue)
-		//remove the message
+		v, e := base64.URLEncoding.DecodeString(cookieValue)
+		if e != nil {
+			return "", err
+		}
+		value = string(v)
+		// store this flash message to the lifetime request's local storage,
+		// I choose this method because no need to store it if not used at all
+		if storeExists {
+			ctx.Get(flashMessagesStoreContextKey).(map[string]string)[key] = value
+		} else {
+			flashStoreMap := make(map[string]string)
+			flashStoreMap[key] = value
+			ctx.Set(flashMessagesStoreContextKey, flashStoreMap)
+		}
+
+		//remove the real cookie, no need to have that, we stored it on lifetime request
 		ctx.RemoveCookie(key)
 		//it should'b be removed until the next reload, so we don't do that: ctx.Request.Header.SetCookie(key, "")
 	}
@@ -745,15 +769,11 @@ func (ctx *Context) GetFlashBytes(key string) (value []byte, err error) {
 }
 
 // SetFlash sets a flash message, accepts 2 parameters the key(string) and the value(string)
+// the value will be available on the NEXT request
 func (ctx *Context) SetFlash(key string, value string) {
-	ctx.SetFlashBytes(key, utils.StringToBytes(value))
-}
-
-// SetFlashBytes sets a flash message, accepts 2 parameters the key(string) and the value([]byte)
-func (ctx *Context) SetFlashBytes(key string, value []byte) {
 	c := fasthttp.AcquireCookie()
 	c.SetKey(key)
-	c.SetValue(base64.URLEncoding.EncodeToString(value))
+	c.SetValue(base64.URLEncoding.EncodeToString([]byte(value)))
 	c.SetPath("/")
 	c.SetHTTPOnly(true)
 	ctx.RequestCtx.Response.Header.SetCookie(c)

context/context.go

@@ -69,10 +69,8 @@ type (
 		SetCookie(*fasthttp.Cookie)
 		SetCookieKV(string, string)
 		RemoveCookie(string)
-		GetFlash(string) string
-		GetFlashBytes(string) ([]byte, error)
+		GetFlash(string) (string, error)
 		SetFlash(string, string)
-		SetFlashBytes(string, []byte)
 		Session() store.IStore
 		SessionDestroy()
 		Log(string, ...interface{})

test/render_test.go

@@ -60,27 +60,27 @@ func TestRenderRest(t *testing.T) {
 	})
 
 	e := tester(api, t)
-	dataT := e.GET("/data").Expect()
+	dataT := e.GET("/data").Expect().Status(iris.StatusOK)
 	dataT.Header("Content-Type").Equal("application/octet-stream")
 	dataT.Body().Equal(string(dataContents))
 
-	textT := e.GET("/text").Expect()
+	textT := e.GET("/text").Expect().Status(iris.StatusOK)
 	textT.Header("Content-Type").Equal("text/plain; charset=UTF-8")
 	textT.Body().Equal(textContents)
 
-	JSONPT := e.GET("/jsonp").Expect()
+	JSONPT := e.GET("/jsonp").Expect().Status(iris.StatusOK)
 	JSONPT.Header("Content-Type").Equal("application/javascript; charset=UTF-8")
 	JSONPT.Body().Equal(JSONPCallback + `({"hello":"jsonp"});`)
 
-	JSONT := e.GET("/json").Expect()
+	JSONT := e.GET("/json").Expect().Status(iris.StatusOK)
 	JSONT.Header("Content-Type").Equal("application/json; charset=UTF-8")
 	JSONT.JSON().Object().Equal(JSONXMLContents)
 
-	XMLT := e.GET("/xml").Expect()
+	XMLT := e.GET("/xml").Expect().Status(iris.StatusOK)
 	XMLT.Header("Content-Type").Equal("text/xml; charset=UTF-8")
 	XMLT.Body().Equal(`<` + JSONXMLContents.XMLName.Local + ` first="` + JSONXMLContents.FirstAttr + `" second="` + JSONXMLContents.SecondAttr + `"><name>` + JSONXMLContents.Name + `</name><birth>` + JSONXMLContents.Birth + `</birth><stars>` + strconv.Itoa(JSONXMLContents.Stars) + `</stars></info>`)
 
-	markdownT := e.GET("/markdown").Expect()
+	markdownT := e.GET("/markdown").Expect().Status(iris.StatusOK)
 	markdownT.Header("Content-Type").Equal("text/html; charset=UTF-8")
 	markdownT.Body().Equal("<h1>" + markdownContents[2:] + "</h1>\n")
 

test/sessions_test.go

@@ -72,3 +72,43 @@ func TestSessions(t *testing.T) {
 	e.POST("/set").WithJSON(values).Expect().Status(iris.StatusOK).Cookies().NotEmpty()
 	e.GET("/clear").Expect().Status(iris.StatusOK).JSON().Object().Empty()
 }
+
+func FlashMessagesTest(t *testing.T) {
+	api := iris.New()
+	values := map[string]string{"name": "kataras", "package": "iris"}
+
+	api.Put("/set", func(ctx *iris.Context) {
+		for k, v := range values {
+			ctx.SetFlash(k, v)
+		}
+	})
+
+	api.Get("/get", func(ctx *iris.Context) {
+		// one time one handler
+		kv := make(map[string]string)
+		for k := range values {
+			kv[k], _ = ctx.GetFlash(k)
+		}
+		//second time on the same handler
+		for k := range values {
+			kv[k], _ = ctx.GetFlash(k)
+		}
+
+	}, func(ctx *iris.Context) {
+		// third time on a next handler
+		// test the if next handler has access to them(must) because flash are request lifetime now.
+		kv := make(map[string]string)
+		for k := range values {
+			kv[k], _ = ctx.GetFlash(k)
+		}
+		// print them to the client for test the response also
+		ctx.JSON(iris.StatusOK, kv)
+	})
+
+	e := tester(api, t)
+	e.PUT("/set").Expect().Status(iris.StatusOK)
+	e.GET("/get").Expect().Status(iris.StatusOK).JSON().Object().Equal(values)
+	// secnd request lifetime ,the flash messages here should be not available
+	e.GET("/get").Expect().Status(iris.StatusOK).JSON().Object().Empty()
+
+}