From 48a2aaadc1a8dcb917870fcab60848c7dbede5dd Mon Sep 17 00:00:00 2001
From: Nick Meves <nick.meves@greenhouse.io>
Date: Fri, 3 Jul 2020 23:41:08 -0700
Subject: [PATCH] Count complete cookie content in byte splitting

---
 pkg/sessions/cookie/session_store.go      | 26 +++++---
 pkg/sessions/cookie/session_store_test.go | 76 +++++++++++------------
 2 files changed, 53 insertions(+), 49 deletions(-)

diff --git a/pkg/sessions/cookie/session_store.go b/pkg/sessions/cookie/session_store.go
index 32a71e56..833ed729 100644
--- a/pkg/sessions/cookie/session_store.go
+++ b/pkg/sessions/cookie/session_store.go
@@ -15,6 +15,13 @@ import (
 	"github.com/oauth2-proxy/oauth2-proxy/pkg/logger"
 )
 
+const (
+	// Cookies are limited to 4kb for all parts
+	// including the cookie name, value, attributes; IE (http.cookie).String()
+	// Most browsers' max is 4096 -- but we give ourselves some leeway
+	maxCookieLength = 4000
+)
+
 // Ensure CookieSessionStore implements the interface
 var _ sessions.SessionStore = &SessionStore{}
 
@@ -102,9 +109,7 @@ func (s *SessionStore) makeSessionCookie(req *http.Request, value string, now ti
 	}
 	c := s.makeCookie(req, s.CookieOptions.Name, value, s.CookieOptions.Expire, now)
 
-	// Cookies are limited to 4kb including the length of the cookie name,
-	// the cookie name can be up to 256 bytes
-	if len(c.Value) > 4096-len(s.CookieOptions.Name) {
+	if len(c.String()) > maxCookieLength {
 		return splitCookie(c)
 	}
 	return []*http.Cookie{c}
@@ -139,7 +144,7 @@ func NewCookieSessionStore(opts *options.SessionOptions, cookieOpts *options.Coo
 // it into a slice of cookies which fit within the 4kb cookie limit indexing
 // the cookies from 0
 func splitCookie(c *http.Cookie) []*http.Cookie {
-	if len(c.Value) < 4096-len(c.Name) {
+	if len(c.String()) < maxCookieLength {
 		return []*http.Cookie{c}
 	}
 
@@ -153,13 +158,16 @@ func splitCookie(c *http.Cookie) []*http.Cookie {
 		newCookie.Name = splitCookieName(c.Name, count)
 		count++
 
-		maxCookieLength := 4096 - len(newCookie.Name)
-		if len(valueBytes) < maxCookieLength {
-			newCookie.Value = string(valueBytes)
+		newCookie.Value = string(valueBytes)
+		cookieLength := len(newCookie.String())
+		if cookieLength <= maxCookieLength {
 			valueBytes = []byte{}
 		} else {
-			newValue := valueBytes[:maxCookieLength]
-			valueBytes = valueBytes[maxCookieLength:]
+			overflow := cookieLength - maxCookieLength
+			valueSize := len(valueBytes) - overflow
+
+			newValue := valueBytes[:valueSize]
+			valueBytes = valueBytes[valueSize:]
 			newCookie.Value = string(newValue)
 		}
 		cookies = append(cookies, newCookie)
diff --git a/pkg/sessions/cookie/session_store_test.go b/pkg/sessions/cookie/session_store_test.go
index f29edd16..c67aafbe 100644
--- a/pkg/sessions/cookie/session_store_test.go
+++ b/pkg/sessions/cookie/session_store_test.go
@@ -54,59 +54,55 @@ func Test_copyCookie(t *testing.T) {
 }
 
 func Test_splitCookie(t *testing.T) {
-	testCases := map[string]struct {
-		Cookie        *http.Cookie
-		ExpectedSizes []int
-	}{
+	testCases := map[string]*http.Cookie{
 		"Short cookie name": {
-			Cookie: &http.Cookie{
-				Name:  "short",
-				Value: strings.Repeat("v", 10000),
-			},
-			ExpectedSizes: []int{4089, 4089, 1822},
+			Name:  "short",
+			Value: strings.Repeat("v", 10000),
 		},
 		"Long cookie name": {
-			Cookie: &http.Cookie{
-				Name:  strings.Repeat("n", 251),
-				Value: strings.Repeat("a", 10000),
-			},
-			ExpectedSizes: []int{3843, 3843, 2314},
+			Name:  strings.Repeat("n", 251),
+			Value: strings.Repeat("a", 10000),
 		},
 		"Max cookie name": {
-			Cookie: &http.Cookie{
-				Name:  strings.Repeat("n", 256),
-				Value: strings.Repeat("a", 10000),
-			},
-			ExpectedSizes: []int{3840, 3840, 2320},
+			Name:  strings.Repeat("n", 256),
+			Value: strings.Repeat("a", 10000),
 		},
 		"Suffix overflow cookie name": {
-			Cookie: &http.Cookie{
-				Name:  strings.Repeat("n", 255),
-				Value: strings.Repeat("a", 10000),
-			},
-			ExpectedSizes: []int{3840, 3840, 2320},
+			Name:  strings.Repeat("n", 255),
+			Value: strings.Repeat("a", 10000),
 		},
-		"Double digit suffix cookie name results in different sizes": {
-			Cookie: &http.Cookie{
-				Name:  strings.Repeat("n", 253),
-				Value: strings.Repeat("a", 50000),
-			},
-			ExpectedSizes: []int{3841, 3841, 3841, 3841, 3841, 3841, 3841, 3841, 3841, 3841,
-				3840, 3840, 3840, 70},
+		"Double digit suffix cookie name overflow": {
+			Name:  strings.Repeat("n", 253),
+			Value: strings.Repeat("a", 50000),
 		},
-		"Double digit suffix overflow cookie name results in same sizes": {
-			Cookie: &http.Cookie{
-				Name:  strings.Repeat("n", 254),
-				Value: strings.Repeat("a", 50000),
-			},
-			ExpectedSizes: []int{3840, 3840, 3840, 3840, 3840, 3840, 3840, 3840, 3840, 3840,
-				3840, 3840, 3840, 80},
+		"With short name and attributes": {
+			Name:     "short",
+			Value:    strings.Repeat("v", 10000),
+			Path:     "/path",
+			Domain:   "x.y.z",
+			Secure:   true,
+			HttpOnly: true,
+			SameSite: http.SameSiteLaxMode,
+		},
+		"With max length name and attributes": {
+			Name:     strings.Repeat("n", 256),
+			Value:    strings.Repeat("v", 10000),
+			Path:     "/path",
+			Domain:   "x.y.z",
+			Secure:   true,
+			HttpOnly: true,
+			SameSite: http.SameSiteLaxMode,
 		},
 	}
 	for testName, tc := range testCases {
 		t.Run(testName, func(t *testing.T) {
-			for i, cookie := range splitCookie(tc.Cookie) {
-				assert.Equal(t, tc.ExpectedSizes[i], len(cookie.Value))
+			splitCookies := splitCookie(tc)
+			for i, cookie := range splitCookies {
+				if i < len(splitCookies)-1 {
+					assert.Equal(t, 4000, len(cookie.String()))
+				} else {
+					assert.GreaterOrEqual(t, 4000, len(cookie.String()))
+				}
 			}
 		})
 	}