From 57f99ca50fd368c9fdc4543e22d6e84b88571ae5 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Fri, 27 Mar 2020 10:47:22 +0800
Subject: [PATCH] Add set samesite in cookie. (#2306)

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
---
 context.go      | 13 +++++++++++--
 context_test.go |  6 ++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/context.go b/context.go
index 572dbb8..5c1fab2 100644
--- a/context.go
+++ b/context.go
@@ -71,6 +71,10 @@ type Context struct {
 	// formCache use url.ParseQuery cached PostForm contains the parsed form data from POST, PATCH,
 	// or PUT body parameters.
 	formCache url.Values
+
+	// SameSite allows a server to define a cookie attribute making it impossible for
+	// the browser to send this cookie along with cross-site requests.
+	sameSite http.SameSite
 }
 
 /************************************/
@@ -782,10 +786,15 @@ func (c *Context) GetRawData() ([]byte, error) {
 	return ioutil.ReadAll(c.Request.Body)
 }
 
+// SetSameSite with cookie
+func (c *Context) SetSameSite(samesite http.SameSite) {
+	c.sameSite = samesite
+}
+
 // SetCookie adds a Set-Cookie header to the ResponseWriter's headers.
 // The provided cookie must have a valid Name. Invalid cookies may be
 // silently dropped.
-func (c *Context) SetCookie(name, value string, maxAge int, path, domain string, sameSite http.SameSite, secure, httpOnly bool) {
+func (c *Context) SetCookie(name, value string, maxAge int, path, domain string, secure, httpOnly bool) {
 	if path == "" {
 		path = "/"
 	}
@@ -795,7 +804,7 @@ func (c *Context) SetCookie(name, value string, maxAge int, path, domain string,
 		MaxAge:   maxAge,
 		Path:     path,
 		Domain:   domain,
-		SameSite: sameSite,
+		SameSite: c.sameSite,
 		Secure:   secure,
 		HttpOnly: httpOnly,
 	})
diff --git a/context_test.go b/context_test.go
index 78b22c0..80c7b40 100644
--- a/context_test.go
+++ b/context_test.go
@@ -602,13 +602,15 @@ func TestContextPostFormMultipart(t *testing.T) {
 
 func TestContextSetCookie(t *testing.T) {
 	c, _ := CreateTestContext(httptest.NewRecorder())
-	c.SetCookie("user", "gin", 1, "/", "localhost", http.SameSiteLaxMode, true, true)
+	c.SetSameSite(http.SameSiteLaxMode)
+	c.SetCookie("user", "gin", 1, "/", "localhost", true, true)
 	assert.Equal(t, "user=gin; Path=/; Domain=localhost; Max-Age=1; HttpOnly; Secure; SameSite=Lax", c.Writer.Header().Get("Set-Cookie"))
 }
 
 func TestContextSetCookiePathEmpty(t *testing.T) {
 	c, _ := CreateTestContext(httptest.NewRecorder())
-	c.SetCookie("user", "gin", 1, "", "localhost", http.SameSiteLaxMode, true, true)
+	c.SetSameSite(http.SameSiteLaxMode)
+	c.SetCookie("user", "gin", 1, "", "localhost", true, true)
 	assert.Equal(t, "user=gin; Path=/; Domain=localhost; Max-Age=1; HttpOnly; Secure; SameSite=Lax", c.Writer.Header().Get("Set-Cookie"))
 }