From 597ffeb1217058afd71182e8cec1de670f8b6998 Mon Sep 17 00:00:00 2001
From: Ilia Pertsev <peppered@users.noreply.github.com>
Date: Tue, 5 Jan 2021 04:21:17 +0300
Subject: [PATCH] Fix joined cookie name for those containing underline in the
 suffix (#970)

* properly handle splitted cookies with names ending with _

* test update

* provide cookieName into joinCookies instead of processing the suffix

* changelog update

* test update
---
 CHANGELOG.md                              |  1 +
 pkg/sessions/cookie/session_store.go      |  7 ++--
 pkg/sessions/cookie/session_store_test.go | 51 ++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2846d28e7b..2272fb470a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -59,6 +59,7 @@
 
 ## Changes since v6.1.1
 
+- [#970](https://github.com/oauth2-proxy/oauth2-proxy/pull/970) Fix joined cookie name for those containing underline in the suffix (@peppered)
 - [#953](https://github.com/oauth2-proxy/oauth2-proxy/pull/953) Migrate Keycloak to EnrichSession & support multiple groups for authorization (@NickMeves)
 - [#957](https://github.com/oauth2-proxy/oauth2-proxy/pull/957) Use X-Forwarded-{Proto,Host,Uri} on redirect as last resort (@linuxgemini)
 - [#630](https://github.com/oauth2-proxy/oauth2-proxy/pull/630) Add support for Gitlab project based authentication (@factorysh)
diff --git a/pkg/sessions/cookie/session_store.go b/pkg/sessions/cookie/session_store.go
index 4a546cfce3..ce51ed07ef 100644
--- a/pkg/sessions/cookie/session_store.go
+++ b/pkg/sessions/cookie/session_store.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
-	"strings"
 	"time"
 
 	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/options"
@@ -220,12 +219,12 @@ func loadCookie(req *http.Request, cookieName string) (*http.Cookie, error) {
 	if len(cookies) == 0 {
 		return nil, fmt.Errorf("could not find cookie %s", cookieName)
 	}
-	return joinCookies(cookies)
+	return joinCookies(cookies, cookieName)
 }
 
 // joinCookies takes a slice of cookies from the request and reconstructs the
 // full session cookie
-func joinCookies(cookies []*http.Cookie) (*http.Cookie, error) {
+func joinCookies(cookies []*http.Cookie, cookieName string) (*http.Cookie, error) {
 	if len(cookies) == 0 {
 		return nil, fmt.Errorf("list of cookies must be > 0")
 	}
@@ -236,7 +235,7 @@ func joinCookies(cookies []*http.Cookie) (*http.Cookie, error) {
 	for i := 1; i < len(cookies); i++ {
 		c.Value += cookies[i].Value
 	}
-	c.Name = strings.TrimRight(c.Name, "_0")
+	c.Name = cookieName
 	return c, nil
 }
 
diff --git a/pkg/sessions/cookie/session_store_test.go b/pkg/sessions/cookie/session_store_test.go
index 5ef9eff1a1..bf9dda141e 100644
--- a/pkg/sessions/cookie/session_store_test.go
+++ b/pkg/sessions/cookie/session_store_test.go
@@ -154,9 +154,58 @@ func Test_splitCookie_joinCookies(t *testing.T) {
 				Value: value,
 			}
 			splitCookies := splitCookie(cookie)
-			joinedCookie, err := joinCookies(splitCookies)
+			joinedCookie, err := joinCookies(splitCookies, cookie.Name)
 			assert.NoError(t, err)
 			assert.Equal(t, *cookie, *joinedCookie)
 		})
 	}
 }
+
+func Test_joinCookies_withUnderlineSuffix(t *testing.T) {
+	testCases := map[string]struct {
+		CookieName string
+		SplitOrder []int
+	}{
+		"Ascending order split with \"_\" suffix": {
+			CookieName: "_cookie_name_",
+			SplitOrder: []int{0, 1, 2, 3, 4},
+		},
+		"Descending order split with \"_\" suffix": {
+			CookieName: "_cookie_name_",
+			SplitOrder: []int{4, 3, 2, 1, 0},
+		},
+		"Arbitrary order split with \"_\" suffix": {
+			CookieName: "_cookie_name_",
+			SplitOrder: []int{3, 1, 2, 0, 4},
+		},
+		"Arbitrary order split with \"_0\" suffix": {
+			CookieName: "_cookie_name_0",
+			SplitOrder: []int{1, 3, 0, 2, 4},
+		},
+		"Arbitrary order split with \"_1\" suffix": {
+			CookieName: "_cookie_name_1",
+			SplitOrder: []int{4, 1, 3, 0, 2},
+		},
+		"Arbitrary order split with \"__\" suffix": {
+			CookieName: "_cookie_name__",
+			SplitOrder: []int{1, 0, 4, 3, 2},
+		},
+	}
+
+	for testName, testCase := range testCases {
+		t.Run(testName, func(t *testing.T) {
+			cookieName := testCase.CookieName
+			var splitCookies []*http.Cookie
+			for _, splitSuffix := range testCase.SplitOrder {
+				cookie := &http.Cookie{
+					Name:  splitCookieName(cookieName, splitSuffix),
+					Value: strings.Repeat("v", 1000),
+				}
+				splitCookies = append(splitCookies, cookie)
+			}
+			joinedCookie, err := joinCookies(splitCookies, cookieName)
+			assert.NoError(t, err)
+			assert.Equal(t, cookieName, joinedCookie.Name)
+		})
+	}
+}