Support SameSite cookie attribute (#488)

SameSite cookie attribute implementation.

cookie.go

@@ -18,6 +18,20 @@ var (
 	CookieExpireUnlimited = zeroTime
 )
 
+// CookieSameSite is an enum for the mode in which the SameSite flag should be set for the given cookie.
+// See https://tools.ietf.org/html/draft-ietf-httpbis-cookie-same-site-00 for details.
+type CookieSameSite int
+const (
+	// CookieSameSiteDisabled removes the SameSite flag
+	CookieSameSiteDisabled CookieSameSite = iota
+	// CookieSameSiteDefaultMode sets the SameSite flag
+	CookieSameSiteDefaultMode
+	// CookieSameSiteLaxMode sets the SameSite flag with the "Lax" parameter
+	CookieSameSiteLaxMode
+	// CookieSameSiteStrictMode sets the SameSite flag with the "Strict" parameter
+	CookieSameSiteStrictMode
+)
+
 // AcquireCookie returns an empty Cookie object from the pool.
 //
 // The returned object may be returned back to the pool with ReleaseCookie.
@@ -58,6 +72,7 @@ type Cookie struct {
 
 	httpOnly bool
 	secure   bool
+	sameSite CookieSameSite
 
 	bufKV argsKV
 	buf   []byte
@@ -74,6 +89,7 @@ func (c *Cookie) CopyTo(src *Cookie) {
 	c.path = append(c.path[:0], src.path...)
 	c.httpOnly = src.httpOnly
 	c.secure = src.secure
+	c.sameSite = src.sameSite
 }
 
 // HTTPOnly returns true if the cookie is http only.
@@ -96,6 +112,16 @@ func (c *Cookie) SetSecure(secure bool) {
 	c.secure = secure
 }
 
+// SameSite returns the SameSite mode.
+func (c *Cookie) SameSite() CookieSameSite {
+	return c.sameSite
+}
+
+// SetSameSite sets the cookie's SameSite flag to the given value.
+func (c *Cookie) SetSameSite(mode CookieSameSite) {
+	c.sameSite = mode
+}
+
 // Path returns cookie path.
 func (c *Cookie) Path() []byte {
 	return c.path
@@ -209,6 +235,7 @@ func (c *Cookie) Reset() {
 	c.path = c.path[:0]
 	c.httpOnly = false
 	c.secure = false
+	c.sameSite = CookieSameSiteDisabled
 }
 
 // AppendBytes appends cookie representation to dst and returns
@@ -246,6 +273,21 @@ func (c *Cookie) AppendBytes(dst []byte) []byte {
 		dst = append(dst, ';', ' ')
 		dst = append(dst, strCookieSecure...)
 	}
+	switch c.sameSite {
+	case CookieSameSiteDefaultMode:
+		dst = append(dst, ';', ' ')
+		dst = append(dst, strCookieSameSite...)
+	case CookieSameSiteLaxMode:
+		dst = append(dst, ';', ' ')
+		dst = append(dst, strCookieSameSite...)
+		dst = append(dst, '=')
+		dst = append(dst, strCookieSameSiteLax...)
+	case CookieSameSiteStrictMode:
+		dst = append(dst, ';', ' ')
+		dst = append(dst, strCookieSameSite...)
+		dst = append(dst, '=')
+		dst = append(dst, strCookieSameSiteStrict...)
+	}
 	return dst
 }
 
@@ -330,6 +372,21 @@ func (c *Cookie) ParseBytes(src []byte) error {
 				if caseInsensitiveCompare(strCookiePath, kv.key) {
 					c.path = append(c.path[:0], kv.value...)
 				}
+
+			case 's': // "samesite"
+				if caseInsensitiveCompare(strCookieSameSite, kv.key) {
+					// Case insensitive switch on first char
+					switch kv.value[0] | 0x20 {
+					case 'l': // "lax"
+						if caseInsensitiveCompare(strCookieSameSiteLax, kv.value) {
+							c.sameSite = CookieSameSiteLaxMode
+						}
+					case 's': // "strict"
+						if caseInsensitiveCompare(strCookieSameSiteStrict, kv.value) {
+							c.sameSite = CookieSameSiteStrictMode
+						}
+					}
+				}
 			}
 
 		} else if len(kv.value) != 0 {
@@ -343,6 +400,8 @@ func (c *Cookie) ParseBytes(src []byte) error {
 			case 's': // "secure"
 				if caseInsensitiveCompare(strCookieSecure, kv.value) {
 					c.secure = true
+				} else if caseInsensitiveCompare(strCookieSameSite, kv.value) {
+					c.sameSite = CookieSameSiteDefaultMode
 				}
 			}
 		} // else empty or no match

cookie_test.go

@@ -76,7 +76,7 @@ func TestCookieSecure(t *testing.T) {
 	if err := c.Parse("foo=bar"); err != nil {
 		t.Fatalf("unexpected error: %s", err)
 	}
-	if c.HTTPOnly() {
+	if c.Secure() {
 		t.Fatalf("Unexpected secure flag set")
 	}
 	s = c.String()
@@ -85,6 +85,54 @@ func TestCookieSecure(t *testing.T) {
 	}
 }
 
+func TestCookieSameSite(t *testing.T) {
+	var c Cookie
+
+	if err := c.Parse("foo=bar; samesite"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if c.SameSite() != CookieSameSiteDefaultMode {
+		t.Fatalf("SameSite must be set")
+	}
+	s := c.String()
+	if !strings.Contains(s, "; SameSite") {
+		t.Fatalf("missing SameSite flag in cookie %q", s)
+	}
+	
+	if err := c.Parse("foo=bar; samesite=lax"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if c.SameSite() != CookieSameSiteLaxMode {
+		t.Fatalf("SameSite Lax Mode must be set")
+	}
+	s = c.String()
+	if !strings.Contains(s, "; SameSite=Lax") {
+		t.Fatalf("missing SameSite flag in cookie %q", s)
+	}
+	
+	if err := c.Parse("foo=bar; samesite=strict"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if c.SameSite() != CookieSameSiteStrictMode {
+		t.Fatalf("SameSite Strict Mode must be set")
+	}
+	s = c.String()
+	if !strings.Contains(s, "; SameSite=Strict") {
+		t.Fatalf("missing SameSite flag in cookie %q", s)
+	}
+	
+	if err := c.Parse("foo=bar"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if c.SameSite() != CookieSameSiteDisabled {
+		t.Fatalf("Unexpected SameSite flag set")
+	}
+	s = c.String()
+	if strings.Contains(s, "SameSite") {
+		t.Fatalf("unexpected SameSite flag in cookie %q", s)
+	}
+}
+
 func TestCookieMaxAge(t *testing.T) {
 	var c Cookie
 

strings.go

@@ -53,13 +53,16 @@ var (
 	strRange            = []byte("Range")
 	strContentRange     = []byte("Content-Range")
 
-	strCookieExpires  = []byte("expires")
-	strCookieDomain   = []byte("domain")
-	strCookiePath     = []byte("path")
-	strCookieHTTPOnly = []byte("HttpOnly")
-	strCookieSecure   = []byte("secure")
-	strCookieMaxAge   = []byte("max-age")
-
+	strCookieExpires        = []byte("expires")
+	strCookieDomain         = []byte("domain")
+	strCookiePath           = []byte("path")
+	strCookieHTTPOnly       = []byte("HttpOnly")
+	strCookieSecure         = []byte("secure")
+	strCookieMaxAge         = []byte("max-age")
+	strCookieSameSite       = []byte("SameSite")
+	strCookieSameSiteLax    = []byte("Lax")
+	strCookieSameSiteStrict = []byte("Strict")
+	
 	strClose               = []byte("close")
 	strGzip                = []byte("gzip")
 	strDeflate             = []byte("deflate")