From 9d8cf001db4df46b0767729d406cbf35d3d2aff5 Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Fri, 29 Mar 2024 14:01:18 -0700 Subject: [PATCH] Split access token into half and store to avoid "securecookie: the value is too long" error (#4863) * split access token into half and store Signed-off-by: Yubo Wang * lint and fix unit test Signed-off-by: Yubo Wang * fix lint again Signed-off-by: Yubo Wang * keep old name and store new splitted names Signed-off-by: Yubo Wang * update comment Signed-off-by: Yubo Wang * correct spelling... Signed-off-by: Yubo Wang --------- Signed-off-by: Yubo Wang Co-authored-by: Yubo Wang --- flyteadmin/auth/cookie.go | 4 + flyteadmin/auth/cookie_manager.go | 48 ++++++++++-- flyteadmin/auth/cookie_manager_test.go | 102 +++++++++++++++++++++++-- flyteadmin/auth/handlers_test.go | 16 ++-- flyteadmin/go.mod | 1 + flyteadmin/go.sum | 2 + 6 files changed, 154 insertions(+), 19 deletions(-) diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go index bf80c1f920..2470220d24 100644 --- 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 index 9bc64b88cf..ce360c9d3a 100644 --- 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 index 6dd67a0473..09d8468e83 100644 --- 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": "john.doe@example.com", + "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 index 452f797d9f..5428fb9b80 100644 --- 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 index 070c057741..39bde3ea23 100644 --- 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 index 9691e94d86..456493c095 100644 --- 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=