Split access token into half and store to avoid "securecookie: the va…

…lue is too long" error (#4863) * split access token into half and store Signed-off-by: Yubo Wang <[email protected]> * lint and fix unit test Signed-off-by: Yubo Wang <[email protected]> * fix lint again Signed-off-by: Yubo Wang <[email protected]> * keep old name and store new splitted names Signed-off-by: Yubo Wang <[email protected]> * update comment Signed-off-by: Yubo Wang <[email protected]> * correct spelling... Signed-off-by: Yubo Wang <[email protected]> --------- Signed-off-by: Yubo Wang <[email protected]> Co-authored-by: Yubo Wang <[email protected]>
flyteorg · Mar 29, 2024 · 9d8cf00 · 9d8cf00
1 parent a9847ef
commit 9d8cf00
Show file tree

Hide file tree

Showing 6 changed files with 154 additions and 19 deletions.
diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go
@@ -20,6 +20,10 @@ const (
 	// #nosec
 	accessTokenCookieName = "flyte_at"
 	// #nosec
+	accessTokenCookieNameSplitFirst = "flyte_at_1"
+	// #nosec
+	accessTokenCookieNameSplitSecond = "flyte_at_2"
+	// #nosec
 	idTokenCookieName = "flyte_idt"
 	// #nosec
 	refreshTokenCookieName = "flyte_rt"

diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go
@@ -52,6 +52,24 @@ func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded strin
 	}, nil
 }
 
+func (c CookieManager) RetrieveAccessToken(ctx context.Context, request *http.Request) (string, error) {
+	// If there is an old access token, we will retrieve it
+	oldAccessToken, err := retrieveSecureCookie(ctx, request, accessTokenCookieName, c.hashKey, c.blockKey)
+	if err == nil && oldAccessToken != "" {
+		return oldAccessToken, nil
+	}
+	// If there is no old access token, we will retrieve the new split access token
+	accessTokenFirstHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplitFirst, c.hashKey, c.blockKey)
+	if err != nil {
+		return "", err
+	}
+	accessTokenSecondHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplitSecond, c.hashKey, c.blockKey)
+	if err != nil {
+		return "", err
+	}
+	return accessTokenFirstHalf + accessTokenSecondHalf, nil
+}
+
 // TODO: Separate refresh token from access token, remove named returns, and use stdlib errors.
 // RetrieveTokenValues retrieves id, access and refresh tokens from cookies if they exist. The existence of a refresh token
 // in a cookie is optional and hence failure to find or read that cookie is tolerated. An error is returned in case of failure
@@ -64,7 +82,7 @@ func (c CookieManager) RetrieveTokenValues(ctx context.Context, request *http.Re
 		return "", "", "", err
 	}
 
-	accessToken, err = retrieveSecureCookie(ctx, request, accessTokenCookieName, c.hashKey, c.blockKey)
+	accessToken, err = c.RetrieveAccessToken(ctx, request)
 	if err != nil {
 		return "", "", "", err
 	}
@@ -135,20 +153,38 @@ func (c CookieManager) SetAuthCodeCookie(ctx context.Context, writer http.Respon
 	return nil
 }
 
+func (c CookieManager) StoreAccessToken(ctx context.Context, accessToken string, writer http.ResponseWriter) error {
+	midpoint := len(accessToken) / 2
+	firstHalf := accessToken[:midpoint]
+	secondHalf := accessToken[midpoint:]
+	atCookieFirst, err := NewSecureCookie(accessTokenCookieNameSplitFirst, firstHalf, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
+	if err != nil {
+		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie first half %s", err)
+		return err
+	}
+	http.SetCookie(writer, &atCookieFirst)
+	atCookieSecond, err := NewSecureCookie(accessTokenCookieNameSplitSecond, secondHalf, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
+	if err != nil {
+		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie second half %s", err)
+		return err
+	}
+	http.SetCookie(writer, &atCookieSecond)
+	return nil
+}
+
 func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) error {
 	if token == nil {
 		logger.Errorf(ctx, "Attempting to set cookies with nil token")
 		return errors.Errorf(ErrTokenNil, "Attempting to set cookies with nil token")
 	}
 
-	atCookie, err := NewSecureCookie(accessTokenCookieName, token.AccessToken, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
+	err := c.StoreAccessToken(ctx, token.AccessToken, writer)
+
 	if err != nil {
-		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie %s", err)
+		logger.Errorf(ctx, "Error storing access token %s", err)
 		return err
 	}
 
-	http.SetCookie(writer, &atCookie)
-
 	if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted {
 		idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
 		if err != nil {
@@ -188,6 +224,8 @@ func (c *CookieManager) getLogoutCookie(name string) *http.Cookie {
 
 func (c CookieManager) DeleteCookies(_ context.Context, writer http.ResponseWriter) {
 	http.SetCookie(writer, c.getLogoutCookie(accessTokenCookieName))
+	http.SetCookie(writer, c.getLogoutCookie(accessTokenCookieNameSplitFirst))
+	http.SetCookie(writer, c.getLogoutCookie(accessTokenCookieNameSplitSecond))
 	http.SetCookie(writer, c.getLogoutCookie(refreshTokenCookieName))
 	http.SetCookie(writer, c.getLogoutCookie(idTokenCookieName))
 }

diff --git a/flyteadmin/auth/cookie_manager_test.go b/flyteadmin/auth/cookie_manager_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang-jwt/jwt"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
@@ -57,9 +58,10 @@ func TestCookieManager(t *testing.T) {
 		assert.NoError(t, err)
 		fmt.Println(w.Header().Get("Set-Cookie"))
 		c := w.Result().Cookies()
-		assert.Equal(t, "flyte_at", c[0].Name)
-		assert.Equal(t, "flyte_idt", c[1].Name)
-		assert.Equal(t, "flyte_rt", c[2].Name)
+		assert.Equal(t, "flyte_at_1", c[0].Name)
+		assert.Equal(t, "flyte_at_2", c[1].Name)
+		assert.Equal(t, "flyte_idt", c[2].Name)
+		assert.Equal(t, "flyte_rt", c[3].Name)
 	})
 
 	t.Run("set_token_nil", func(t *testing.T) {
@@ -70,6 +72,79 @@ func TestCookieManager(t *testing.T) {
 		assert.EqualError(t, err, "[EMPTY_OAUTH_TOKEN] Attempting to set cookies with nil token")
 	})
 
+	t.Run("set_long_token_cookies", func(t *testing.T) {
+		ctx := context.Background()
+		// These were generated for unit testing only.
+		hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
+		blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
+		cookieSetting := config.CookieSettings{
+			SameSitePolicy: config.SameSiteDefaultMode,
+			Domain:         "default",
+		}
+		manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
+		assert.NoError(t, err)
+		longString := "asdfkjashdfkljqwpeuqwoieuiposdasdfasdfsdfcuzxcvjzvjlasfuo9qwuoiqwueoqoieulkszcjvhlkshcvlkasdhflkashfqoiwaskldfjhasdfljk" +
+			"aklsdjflkasdfhlkjasdhfklhfjkasdhfkasdhfjasdfhasldkfjhaskldfhaklsdfhlasjdfhalksjdfhlskdqoiweuqioweuqioweuqoiew" +
+			"aklsdjfqwoieuioqwerupweiruoqpweurpqweuropqweurpqweurpoqwuetopqweuropqwuerpoqwuerpoqweuropqweurpoqweyitoqpwety"
+		// These were generated for unit testing only.
+		tokenData := jwt.MapClaims{
+			"iss":                "https://login.microsoftonline.com/72f988bf-86f1-41af-91ab-2d7cd011db47/v2.0",
+			"aio":                "AXQAi/8UAAAAvfR1B135YmjOZGtNGTM/fgvUY0ugwwk2NWCjmNglWR9v+b5sI3cCSGXOp1Zw96qUNb1dm0jqBHHYDKQtc4UplZAtFULbitt3x2KYdigeS5tXl0yNIeMiYsA/1Dpd43xg9sXAtLU3+iZetXiDasdfkpsaldg==",
+			"sub":                "subject",
+			"aud":                "audience",
+			"exp":                1677642301,
+			"nbf":                1677641001,
+			"iat":                1677641001,
+			"jti":                "a-unique-identifier",
+			"user_id":            "123456",
+			"username":           "john_doe",
+			"preferred_username": "john_doe",
+			"given_name":         "John",
+			"family_name":        "Doe",
+			"email":              "[email protected]",
+			"scope":              "read write",
+			"role":               "user",
+			"is_active":          true,
+			"is_verified":        true,
+			"client_id":          "client123",
+			"custom_field1":      longString,
+			"custom_field2":      longString,
+			"custom_field3":      longString,
+			"custom_field4":      longString,
+			"custom_field5":      longString,
+			"custom_field6":      []string{longString, longString},
+			"additional_field1":  "extra_value1",
+			"additional_field2":  "extra_value2",
+		}
+		secretKey := []byte("tJL6Wr2JUsxLyNezPQh1J6zn6wSoDAhgRYSDkaMuEHy75VikiB8wg25WuR96gdMpookdlRvh7SnRvtjQN9b5m4zJCMpSRcJ5DuXl4mcd7Cg3Zp1C5")
+		rawToken := jwt.NewWithClaims(jwt.SigningMethodHS256, tokenData)
+		tokenString, err := rawToken.SignedString(secretKey)
+		if err != nil {
+			fmt.Println("Error:", err)
+			return
+		}
+		token := &oauth2.Token{
+			AccessToken:  tokenString,
+			RefreshToken: "refresh",
+		}
+
+		token = token.WithExtra(map[string]interface{}{
+			"id_token": "id token",
+		})
+
+		w := httptest.NewRecorder()
+		_, err = http.NewRequest("GET", "/api/v1/projects", nil)
+		assert.NoError(t, err)
+		err = manager.SetTokenCookies(ctx, w, token)
+		assert.NoError(t, err)
+		fmt.Println(w.Header().Get("Set-Cookie"))
+		c := w.Result().Cookies()
+		assert.Equal(t, "flyte_at_1", c[0].Name)
+		assert.Equal(t, "flyte_at_2", c[1].Name)
+		assert.Equal(t, "flyte_idt", c[2].Name)
+		assert.Equal(t, "flyte_rt", c[3].Name)
+	})
+
 	t.Run("set_token_cookies_wrong_key", func(t *testing.T) {
 		wrongKey := base64.RawStdEncoding.EncodeToString(bytes.Repeat([]byte("X"), 75))
 		wrongManager, err := NewCookieManager(ctx, wrongKey, wrongKey, cookieSetting)
@@ -130,16 +205,27 @@ func TestCookieManager(t *testing.T) {
 		manager.DeleteCookies(ctx, w)
 
 		cookies := w.Result().Cookies()
-		require.Equal(t, 3, len(cookies))
+		require.Equal(t, 5, len(cookies))
+
 		assert.True(t, time.Now().After(cookies[0].Expires))
 		assert.Equal(t, cookieSetting.Domain, cookies[0].Domain)
 		assert.Equal(t, accessTokenCookieName, cookies[0].Name)
+
 		assert.True(t, time.Now().After(cookies[1].Expires))
 		assert.Equal(t, cookieSetting.Domain, cookies[1].Domain)
-		assert.Equal(t, refreshTokenCookieName, cookies[1].Name)
-		assert.True(t, time.Now().After(cookies[1].Expires))
-		assert.Equal(t, cookieSetting.Domain, cookies[1].Domain)
-		assert.Equal(t, idTokenCookieName, cookies[2].Name)
+		assert.Equal(t, accessTokenCookieNameSplitFirst, cookies[1].Name)
+
+		assert.True(t, time.Now().After(cookies[2].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[2].Domain)
+		assert.Equal(t, accessTokenCookieNameSplitSecond, cookies[2].Name)
+
+		assert.True(t, time.Now().After(cookies[3].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[3].Domain)
+		assert.Equal(t, refreshTokenCookieName, cookies[3].Name)
+
+		assert.True(t, time.Now().After(cookies[4].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[4].Domain)
+		assert.Equal(t, idTokenCookieName, cookies[4].Name)
 	})
 
 	t.Run("get_http_same_site_policy", func(t *testing.T) {

diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go
@@ -305,7 +305,7 @@ func TestGetLogoutHandler(t *testing.T) {
 		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
 
 		assert.Equal(t, http.StatusOK, w.Code)
-		require.Len(t, w.Result().Cookies(), 3)
+		require.Len(t, w.Result().Cookies(), 5)
 		authCtx.AssertExpectations(t)
 	})
 
@@ -323,7 +323,7 @@ func TestGetLogoutHandler(t *testing.T) {
 
 		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
 		authCtx.AssertExpectations(t)
-		require.Len(t, w.Result().Cookies(), 3)
+		require.Len(t, w.Result().Cookies(), 5)
 	})
 
 	t.Run("with_hook_with_redirect", func(t *testing.T) {
@@ -349,7 +349,7 @@ func TestGetLogoutHandler(t *testing.T) {
 		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
 
 		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
-		require.Len(t, w.Result().Cookies(), 3)
+		require.Len(t, w.Result().Cookies(), 5)
 		authCtx.AssertExpectations(t)
 		hook.AssertExpectations(t)
 	})
@@ -399,15 +399,19 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) {
 	req, err := http.NewRequest("GET", "/api/v1/projects", nil)
 	assert.NoError(t, err)
 
-	accessTokenCookie, err := NewSecureCookie(accessTokenCookieName, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
+	accessTokenCookie, err := NewSecureCookie(accessTokenCookieNameSplitFirst, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
 	assert.NoError(t, err)
 	req.AddCookie(&accessTokenCookie)
 
-	idCookie, err := NewSecureCookie(idTokenCookieName, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
+	accessTokenCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplitSecond, ".d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
+	assert.NoError(t, err)
+	req.AddCookie(&accessTokenCookieSplit)
+
+	idCookie, err := NewSecureCookie(idTokenCookieName, "a.b.c.d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
 	assert.NoError(t, err)
 	req.AddCookie(&idCookie)
 
-	assert.Equal(t, "IDToken a.b.c", handler(ctx, req)["authorization"][0])
+	assert.Equal(t, "IDToken a.b.c.d.e.f", handler(ctx, req)["authorization"][0])
 }
 
 func TestGetHTTPMetadataTaggingHandler(t *testing.T) {

diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod
@@ -20,6 +20,7 @@ require (
 	github.com/flyteorg/stow v0.3.10
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-gormigrate/gormigrate/v2 v2.1.1
+	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/golang-jwt/jwt/v4 v4.5.0
 	github.com/golang/glog v1.2.0
 	github.com/golang/protobuf v1.5.3

diff --git a/flyteadmin/go.sum b/flyteadmin/go.sum
@@ -609,6 +609,8 @@ github.com/gogo/protobuf v1.2.0/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7a
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
+github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keLg81eXfW3O+oY=
+github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
 github.com/golang-jwt/jwt/v4 v4.5.0 h1:7cYmW1XlMY7h7ii7UhUyChSgS5wUJEnm9uZVTGqOWzg=
 github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang-jwt/jwt/v5 v5.0.0 h1:1n1XNM9hk7O9mnQoNBGolZvzebBQ7p93ULHRc28XJUE=