feat: Add Partitioned Parameter to SetCookie (#1048)

cloudwego · May 9, 2024 · b5b3f9d · b5b3f9d
1 parent 6f62cd9
commit b5b3f9d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/pkg/app/context.go b/pkg/app/context.go
@@ -1227,6 +1227,10 @@ func (ctx *RequestContext) Cookie(key string) []byte {
 //	4. ctx.SetCookie("user", "", 10, "/", "localhost",protocol.CookieSameSiteLaxMode, false, false)
 //	add response header --->  Set-Cookie: user=; max-age=10; domain=localhost; path=/; SameSite=Lax;
 func (ctx *RequestContext) SetCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly bool) {
+	ctx.setCookie(name, value, maxAge, path, domain, sameSite, secure, httpOnly, false)
+}
+
+func (ctx *RequestContext) setCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly, partitioned bool) {
 	if path == "" {
 		path = "/"
 	}
@@ -1240,9 +1244,20 @@ func (ctx *RequestContext) SetCookie(name, value string, maxAge int, path, domai
 	cookie.SetSecure(secure)
 	cookie.SetHTTPOnly(httpOnly)
 	cookie.SetSameSite(sameSite)
+	cookie.SetPartitioned(partitioned)
 	ctx.Response.Header.SetCookie(cookie)
 }
 
+// SetPartitionedCookie adds a partitioned cookie to the Response's headers.
+// Use protocol.CookieSameSiteNoneMode for cross-site cookies to work.
+//
+// Usage: ctx.SetPartitionedCookie("user", "name", 10, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+//
+// This adds the response header: Set-Cookie: user=name; Max-Age=10; Domain=localhost; Path=/; HttpOnly; Secure; SameSite=None; Partitioned
+func (ctx *RequestContext) SetPartitionedCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly bool) {
+	ctx.setCookie(name, value, maxAge, path, domain, sameSite, secure, httpOnly, true)
+}
+
 // UserAgent returns the value of the request user_agent.
 func (ctx *RequestContext) UserAgent() []byte {
 	return ctx.Request.Header.UserAgent()

diff --git a/pkg/app/context_test.go b/pkg/app/context_test.go
@@ -1615,8 +1615,14 @@ func TestSetBinder(t *testing.T) {
 
 func TestRequestContext_SetCookie(t *testing.T) {
 	c := NewContext(0)
-	c.SetCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteLaxMode, true, true)
-	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=Lax", c.Response.Header.Get("Set-Cookie"))
+	c.SetCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=None", c.Response.Header.Get("Set-Cookie"))
+}
+
+func TestRequestContext_SetPartitionedCookie(t *testing.T) {
+	c := NewContext(0)
+	c.SetPartitionedCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=None; Partitioned", c.Response.Header.Get("Set-Cookie"))
 }
 
 func TestRequestContext_SetCookiePathEmpty(t *testing.T) {