fix https://github.com/kataras/iris/issues/1485 by adding and using the new 'context.UpsertCookie' instead of 'context.SetCookie'

Former-commit-id: 31a50e580929616504b9bbbb1d602b0e9274a568

HISTORY.md

@@ -169,6 +169,7 @@ Other Improvements:
 
 New Context Methods:
 
+- `context.UpsertCookie(*http.Cookie, cookieOptions ...context.CookieOption)` upserts a cookie, fixes [#1485](https://github.com/kataras/iris/issues/1485) too
 - `context.StopWithStatus(int)` stops the handlers chain and writes the status code
 - `context.StopWithJSON(int, interface{})` stops the handlers chain, writes the status code and sends a JSON response
 - `context.StopWithProblem(int, iris.Problem)` stops the handlers, writes the status code and sends an `application/problem+json` response

context/context.go

@@ -933,6 +933,11 @@ type Context interface {
 	//
 	// Example: https://github.com/kataras/iris/tree/master/_examples/cookies/basic
 	SetCookie(cookie *http.Cookie, options ...CookieOption)
+	// UpsertCookie adds a cookie to the response like `SetCookie` does
+	// but it will also perform a replacement of the cookie
+	// if already set by a previous `SetCookie` call.
+	// It reports whether the cookie is new (true) or an existing one was updated (false).
+	UpsertCookie(cookie *http.Cookie, options ...CookieOption) bool
 	// SetSameSite sets a same-site rule for cookies to set.
 	// SameSite allows a server to define a cookie attribute making it impossible for
 	// the browser to send this cookie along with cross-site requests. The main
@@ -4651,7 +4656,7 @@ func CookieDecode(decode CookieDecoder) CookieOption {
 //
 // Example: https://github.com/kataras/iris/tree/master/_examples/cookies/basic
 func (ctx *context) SetCookie(cookie *http.Cookie, options ...CookieOption) {
-	cookie.SameSite = ctx.getSameSite()
+	cookie.SameSite = GetSameSite(ctx)
 
 	for _, opt := range options {
 		opt(cookie)
@@ -4660,6 +4665,37 @@ func (ctx *context) SetCookie(cookie *http.Cookie, options ...CookieOption) {
 	http.SetCookie(ctx.writer, cookie)
 }
 
+// UpsertCookie adds a cookie to the response like `SetCookie` does
+// but it will also perform a replacement of the cookie
+// if already set by a previous `SetCookie` call.
+// It reports whether the cookie is new (true) or an existing one was updated (false).
+func (ctx *context) UpsertCookie(cookie *http.Cookie, options ...CookieOption) bool {
+	cookie.SameSite = GetSameSite(ctx)
+
+	for _, opt := range options {
+		opt(cookie)
+	}
+
+	header := ctx.ResponseWriter().Header()
+
+	if cookies := header["Set-Cookie"]; len(cookies) > 0 {
+		s := cookie.Name + "=" // name=?value
+		for i, c := range cookies {
+			if strings.HasPrefix(c, s) {
+				// We need to update the Set-Cookie (to update the expiration or any other cookie's properties).
+				// Probably the cookie is set and then updated in the first session creation
+				// (e.g. UpdateExpiration, see https://github.com/kataras/iris/issues/1485).
+				cookies[i] = cookie.String()
+				header["Set-Cookie"] = cookies
+				return false
+			}
+		}
+	}
+
+	header.Add("Set-Cookie", cookie.String())
+	return true
+}
+
 const sameSiteContextKey = "iris.cookie_same_site"
 
 // SetSameSite sets a same-site rule for cookies to set.
@@ -4673,7 +4709,8 @@ func (ctx *context) SetSameSite(sameSite http.SameSite) {
 	ctx.Values().Set(sameSiteContextKey, sameSite)
 }
 
-func (ctx *context) getSameSite() http.SameSite {
+// GetSameSite returns the saved-to-context cookie http.SameSite option.
+func GetSameSite(ctx Context) http.SameSite {
 	if v := ctx.Values().Get(sameSiteContextKey); v != nil {
 		sameSite, ok := v.(http.SameSite)
 		if ok {

sessions/cookie.go

@@ -36,7 +36,8 @@ func AddCookie(ctx context.Context, cookie *http.Cookie, reclaim bool) {
 	if reclaim {
 		ctx.Request().AddCookie(cookie)
 	}
-	ctx.SetCookie(cookie)
+
+	ctx.UpsertCookie(cookie)
 }
 
 // RemoveCookie deletes a cookie by it's name/key

sessions/sessions_test.go

@@ -205,8 +205,9 @@ func TestSessionsUpdateExpiration(t *testing.T) {
 	cookieName := "mycustomsessionid"
 
 	sess := sessions.New(sessions.Config{
-		Cookie:  cookieName,
+		Cookie:       cookieName,
-		Expires: 30 * time.Minute,
+		Expires:      30 * time.Minute,
+		AllowReclaim: true,
 	})
 
 	app.Use(sess.Handler())
@@ -233,13 +234,17 @@ func TestSessionsUpdateExpiration(t *testing.T) {
 		writeResponse(ctx)
 	})
 
-	app.Get("/remember_me", func(ctx iris.Context) {
+	app.Post("/remember_me", func(ctx iris.Context) {
 		// re-sends the cookie with the new Expires and MaxAge fields,
 		// test checks that on same session id too.
 		sess.UpdateExpiration(ctx, 24*time.Hour)
 		writeResponse(ctx)
 	})
 
+	app.Get("/destroy", func(ctx iris.Context) {
+		sess.Destroy(ctx) // this will delete the cookie too.
+	})
+
 	e := httptest.New(t, app, httptest.URL("http://example.com"))
 
 	tt := e.GET("/set").Expect().Status(httptest.StatusOK)
@@ -250,7 +255,12 @@ func TestSessionsUpdateExpiration(t *testing.T) {
 	e.GET("/get").Expect().Status(httptest.StatusOK).
 		JSON().Equal(expectedResponse)
 
-	tt = e.GET("/remember_me").Expect().Status(httptest.StatusOK)
+	tt = e.POST("/remember_me").Expect().Status(httptest.StatusOK)
 	tt.Cookie(cookieName).MaxAge().Equal(24 * time.Hour)
 	tt.JSON().Equal(expectedResponse)
+
+	// Test call `UpdateExpiration` when cookie is firstly created.
+	e.GET("/destroy").Expect().Status(httptest.StatusOK)
+	e.POST("/remember_me").Expect().Status(httptest.StatusOK).
+		Cookie(cookieName).MaxAge().Equal(24 * time.Hour)
 }