From 5b2f6ecd2f6439e8d575cf12a4a2871af6d8ae53 Mon Sep 17 00:00:00 2001
From: Bobby DeSimone <bobbydesimone@gmail.com>
Date: Fri, 12 Jul 2019 15:46:05 -0700
Subject: [PATCH] update tests

---
 internal/sessions/cookie_store.go      | 30 +++++++++++++-------------
 internal/sessions/cookie_store_test.go | 25 ++++++++-------------
 2 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/internal/sessions/cookie_store.go b/internal/sessions/cookie_store.go
index 4492312e4..3d237820f 100644
--- a/internal/sessions/cookie_store.go
+++ b/internal/sessions/cookie_store.go
@@ -130,22 +130,22 @@ func (s *CookieStore) makeCSRFCookie(req *http.Request, value string, expiration
 func (s *CookieStore) SetCookie(w http.ResponseWriter, cookie *http.Cookie) {
 	if len(cookie.String()) <= MaxChunkSize {
 		http.SetCookie(w, cookie)
-	} else {
-		chunks := chunk(cookie.Value, MaxChunkSize)
-		for i, c := range chunks {
-			// start with a copy of our original cookie
-			nc := *cookie
-			if i == 0 {
-				// if this is the first cookie, add our canary byte
-				nc.Value = fmt.Sprintf("%s%s", string(ChunkedCanaryByte), c)
-			} else {
-				// subsequent parts will be postfixed with their part number
-				nc.Name = fmt.Sprintf("%s_%d", cookie.Name, i)
-				nc.Value = fmt.Sprintf("%s", c)
-			}
-			log.Info().Interface("new cookie", nc).Msg("SetCookie: chunked")
-			http.SetCookie(w, &nc)
+		return
+	}
+	chunks := chunk(cookie.Value, MaxChunkSize)
+	for i, c := range chunks {
+		// start with a copy of our original cookie
+		nc := *cookie
+		if i == 0 {
+			// if this is the first cookie, add our canary byte
+			nc.Value = fmt.Sprintf("%s%s", string(ChunkedCanaryByte), c)
+		} else {
+			// subsequent parts will be postfixed with their part number
+			nc.Name = fmt.Sprintf("%s_%d", cookie.Name, i)
+			nc.Value = fmt.Sprintf("%s", c)
 		}
+		log.Info().Interface("new cookie", nc).Msg("SetCookie: chunked")
+		http.SetCookie(w, &nc)
 	}
 
 }
diff --git a/internal/sessions/cookie_store_test.go b/internal/sessions/cookie_store_test.go
index 377d34643..44b484400 100644
--- a/internal/sessions/cookie_store_test.go
+++ b/internal/sessions/cookie_store_test.go
@@ -1,7 +1,9 @@
 package sessions
 
 import (
+	"crypto/rand"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
@@ -204,6 +206,10 @@ func TestCookieStore_SaveSession(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	hugeString := make([]byte, 4097)
+	if _, err := rand.Read(hugeString); err != nil {
+		t.Fatal(err)
+	}
 	tests := []struct {
 		name         string
 		sessionState *SessionState
@@ -211,22 +217,9 @@ func TestCookieStore_SaveSession(t *testing.T) {
 		wantErr      bool
 		wantLoadErr  bool
 	}{
-		{"good",
-			&SessionState{
-				AccessToken:     "token1234",
-				RefreshToken:    "refresh4321",
-				RefreshDeadline: time.Now().Add(1 * time.Hour).Truncate(time.Second).UTC(),
-				Email:           "user@domain.com",
-				User:            "user",
-			}, cipher, false, false},
-		{"bad cipher",
-			&SessionState{
-				AccessToken:     "token1234",
-				RefreshToken:    "refresh4321",
-				RefreshDeadline: time.Now().Add(1 * time.Hour).Truncate(time.Second).UTC(),
-				Email:           "user@domain.com",
-				User:            "user",
-			}, mockCipher{}, true, true},
+		{"good", &SessionState{AccessToken: "token1234", RefreshToken: "refresh4321", RefreshDeadline: time.Now().Add(1 * time.Hour).Truncate(time.Second).UTC(), Email: "user@domain.com", User: "user"}, cipher, false, false},
+		{"bad cipher", &SessionState{AccessToken: "token1234", RefreshToken: "refresh4321", RefreshDeadline: time.Now().Add(1 * time.Hour).Truncate(time.Second).UTC(), Email: "user@domain.com", User: "user"}, mockCipher{}, true, true},
+		{"huge cookie", &SessionState{AccessToken: fmt.Sprintf("%x", hugeString), RefreshToken: "refresh4321", RefreshDeadline: time.Now().Add(1 * time.Hour).Truncate(time.Second).UTC(), Email: "user@domain.com", User: "user"}, cipher, false, false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {