Repository: valyala/fasthttp
Pull request #488 (2018-12-12 21:46:37)

Support SameSite cookie attribute
This is an implementation of the `SameSite` attribute for cookies. It is inspired by go v1.11's `net/http` implementation of the SameSite attribute, seen here: https://golang.org/src/net/http/cookie.go?s=4244:4276#L214

This PR addresses issue #489.


Patch:
diff --git a/cookie.go b/cookie.go
index d5abef2..1d9861d 100644
--- a/cookie.go
+++ b/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
diff --git a/server.go b/server.go
index bfa7efc..1f87d72 100644
--- a/server.go
+++ b/server.go
@@ -302,6 +302,11 @@ type Server struct {
 	//     * cONTENT-lenGTH -> Content-Length
 	DisableHeaderNamesNormalizing bool
 
+	// SleepWhenConcurrencyLimitsExceeded is a duration to be slept of if
+	// the concurrency limit in exceeded (default [when is 0]: don't sleep
+	// and accept new connections immidiatelly).
+	SleepWhenConcurrencyLimitsExceeded time.Duration
+
 	// NoDefaultServerHeader, when set to true, causes the default Server header
 	// to be excluded from the Response.
 	//
@@ -1566,7 +1571,11 @@ func (s *Server) Serve(ln net.Listener) error {
 			//
 			// There is a hope other servers didn't reach their
 			// concurrency limits yet :)
-			time.Sleep(100 * time.Millisecond)
+			//
+			// See also: https://github.com/valyala/fasthttp/pull/485#discussion_r239994990
+			if s.SleepWhenConcurrencyLimitsExceeded > 0 {
+				time.Sleep(s.SleepWhenConcurrencyLimitsExceeded)
+			}
 		}
 		c = nil
 	}
@@ -1741,6 +1750,13 @@ func (s *Server) GetCurrentConcurrency() uint32 {
 	return atomic.LoadUint32(&s.concurrency)
 }
 
+// GetOpenConnectionsCount returns a number of opened connections.
+//
+// This function is intended be used by monitoring systems
+func (s *Server) GetOpenConnectionsCount() int32 {
+	return atomic.LoadInt32(&s.open) - 1
+}
+
 func (s *Server) getConcurrency() int {
 	n := s.Concurrency
 	if n <= 0 {
diff --git a/strings.go b/strings.go
index afd70b4..0f05af5 100644
--- a/strings.go
+++ b/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")


Test patch:
diff --git a/cookie_test.go b/cookie_test.go
index cec8edd..c11a942 100644
--- a/cookie_test.go
+++ b/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