From efd7f2dde77245d6eee5edf12102da4f440d0489 Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Thu, 8 Feb 2024 00:59:28 -0800 Subject: [PATCH 1/6] split access token into half and store Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie.go | 2 + flyteadmin/auth/cookie_manager.go | 42 +++++++++--- flyteadmin/auth/cookie_manager_test.go | 95 ++++++++++++++++++++++++-- flyteadmin/go.mod | 1 + flyteadmin/go.sum | 2 + 5 files changed, 127 insertions(+), 15 deletions(-) diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go index bf80c1f920..d5ccb163f3 100644 --- a/flyteadmin/auth/cookie.go +++ b/flyteadmin/auth/cookie.go @@ -19,6 +19,8 @@ import ( const ( // #nosec accessTokenCookieName = "flyte_at" + // nosec + accessTokenCookieNameSplit = "flyte_at_1" // #nosec idTokenCookieName = "flyte_idt" // #nosec diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 9bc64b88cf..359219d3ab 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -52,6 +52,18 @@ func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded strin }, nil } +func (c CookieManager) RetrieveAccessToken(ctx context.Context, request *http.Request) (string, error) { + accessTokenFirstHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieName, c.hashKey, c.blockKey) + if err != nil { + return "", err + } + accessTokenSecondHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplit, 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 +76,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,19 +147,32 @@ 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:] + atCookie, err := NewSecureCookie(accessTokenCookieName, 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, &atCookie) + atCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplit, 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, &atCookieSplit) + 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()) - if err != nil { - logger.Errorf(ctx, "Error generating encrypted accesstoken cookie %s", err) - return err - } - - http.SetCookie(writer, &atCookie) + c.StoreAccessToken(ctx, token.AccessToken, writer) if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted { idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy()) @@ -188,6 +213,7 @@ 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(accessTokenCookieNameSplit)) 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..edfdcd4628 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" @@ -58,8 +59,9 @@ func TestCookieManager(t *testing.T) { 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[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", c[0].Name) + assert.Equal(t, "flyte_at_1", 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,22 @@ func TestCookieManager(t *testing.T) { manager.DeleteCookies(ctx, w) cookies := w.Result().Cookies() - require.Equal(t, 3, len(cookies)) + require.Equal(t, 4, 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, accessTokenCookieNameSplit, cookies[1].Name) + + assert.True(t, time.Now().After(cookies[2].Expires)) + assert.Equal(t, cookieSetting.Domain, cookies[2].Domain) + assert.Equal(t, refreshTokenCookieName, cookies[2].Name) + + assert.True(t, time.Now().After(cookies[3].Expires)) + assert.Equal(t, cookieSetting.Domain, cookies[3].Domain) + assert.Equal(t, idTokenCookieName, cookies[3].Name) }) t.Run("get_http_same_site_policy", func(t *testing.T) { diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod index badc2a3c88..f1dbf89525 100644 --- a/flyteadmin/go.mod +++ b/flyteadmin/go.mod @@ -20,6 +20,7 @@ require ( github.com/flyteorg/stow v0.3.8 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.1.2 github.com/golang/protobuf v1.5.3 diff --git a/flyteadmin/go.sum b/flyteadmin/go.sum index eb2eb48329..ca22700e39 100644 --- a/flyteadmin/go.sum +++ b/flyteadmin/go.sum @@ -608,6 +608,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= From 8f66dfb06d812e2e91f400a30f6089285fa94a6b Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Tue, 26 Mar 2024 08:41:56 +0800 Subject: [PATCH 2/6] lint and fix unit test Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie.go | 2 +- flyteadmin/auth/cookie_manager.go | 6 +++++- flyteadmin/auth/handlers_test.go | 14 +++++++++----- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go index d5ccb163f3..3e2c69d57e 100644 --- a/flyteadmin/auth/cookie.go +++ b/flyteadmin/auth/cookie.go @@ -19,7 +19,7 @@ import ( const ( // #nosec accessTokenCookieName = "flyte_at" - // nosec + // #nosec accessTokenCookieNameSplit = "flyte_at_1" // #nosec idTokenCookieName = "flyte_idt" diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 359219d3ab..554006ce6c 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -172,7 +172,11 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.Response return errors.Errorf(ErrTokenNil, "Attempting to set cookies with nil token") } - c.StoreAccessToken(ctx, token.AccessToken, writer) + err = c.StoreAccessToken(ctx, token.AccessToken, writer) + + if err != nil { + return logger.Errorf(ctx, "Error storing access token %s", err) + } if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted { idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy()) diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go index 452f797d9f..50f4761670 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(), 4) 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(), 4) }) 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(), 4) authCtx.AssertExpectations(t) hook.AssertExpectations(t) }) @@ -403,11 +403,15 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) { assert.NoError(t, err) req.AddCookie(&accessTokenCookie) - idCookie, err := NewSecureCookie(idTokenCookieName, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) + accessTokenCookieSplit, err = NewSecureCookie(accessTokenCookieNameSplit, ".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) { From b5cda52dcfd70cce5792fc63ad50c81cefcd5e3e Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Tue, 26 Mar 2024 23:50:18 +0800 Subject: [PATCH 3/6] fix lint again Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie_manager.go | 5 +++-- flyteadmin/auth/handlers_test.go | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 554006ce6c..6afad24d3d 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -172,10 +172,11 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.Response return errors.Errorf(ErrTokenNil, "Attempting to set cookies with nil token") } - err = c.StoreAccessToken(ctx, token.AccessToken, writer) + err := c.StoreAccessToken(ctx, token.AccessToken, writer) if err != nil { - return logger.Errorf(ctx, "Error storing access token %s", err) + logger.Errorf(ctx, "Error storing access token %s", err) + return err } if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted { diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go index 50f4761670..e636b188c1 100644 --- a/flyteadmin/auth/handlers_test.go +++ b/flyteadmin/auth/handlers_test.go @@ -403,7 +403,7 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) { assert.NoError(t, err) req.AddCookie(&accessTokenCookie) - accessTokenCookieSplit, err = NewSecureCookie(accessTokenCookieNameSplit, ".d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) + accessTokenCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplit, ".d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) assert.NoError(t, err) req.AddCookie(&accessTokenCookieSplit) From 9f953b96f1cf0284e8f1cfa249160bc5c48e201d Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Thu, 28 Mar 2024 11:37:37 -0700 Subject: [PATCH 4/6] keep old name and store new splitted names Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie.go | 4 +++- flyteadmin/auth/cookie_manager.go | 21 ++++++++++++++------- flyteadmin/auth/cookie_manager_test.go | 21 +++++++++++++-------- flyteadmin/auth/handlers_test.go | 10 +++++----- 4 files changed, 35 insertions(+), 21 deletions(-) diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go index 3e2c69d57e..2470220d24 100644 --- a/flyteadmin/auth/cookie.go +++ b/flyteadmin/auth/cookie.go @@ -20,7 +20,9 @@ const ( // #nosec accessTokenCookieName = "flyte_at" // #nosec - accessTokenCookieNameSplit = "flyte_at_1" + accessTokenCookieNameSplitFirst = "flyte_at_1" + // #nosec + accessTokenCookieNameSplitSecond = "flyte_at_2" // #nosec idTokenCookieName = "flyte_idt" // #nosec diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 6afad24d3d..66758275b8 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -53,11 +53,17 @@ func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded strin } func (c CookieManager) RetrieveAccessToken(ctx context.Context, request *http.Request) (string, error) { - accessTokenFirstHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieName, c.hashKey, c.blockKey) + // 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 access token + accessTokenFirstHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplitFirst, c.hashKey, c.blockKey) if err != nil { return "", err } - accessTokenSecondHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplit, c.hashKey, c.blockKey) + accessTokenSecondHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplitSecond, c.hashKey, c.blockKey) if err != nil { return "", err } @@ -151,18 +157,18 @@ func (c CookieManager) StoreAccessToken(ctx context.Context, accessToken string, midpoint := len(accessToken) / 2 firstHalf := accessToken[:midpoint] secondHalf := accessToken[midpoint:] - atCookie, err := NewSecureCookie(accessTokenCookieName, firstHalf, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy()) + 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, &atCookie) - atCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplit, secondHalf, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy()) + 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, &atCookieSplit) + http.SetCookie(writer, &atCookieSecond) return nil } @@ -218,7 +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(accessTokenCookieNameSplit)) + 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 edfdcd4628..09d8468e83 100644 --- a/flyteadmin/auth/cookie_manager_test.go +++ b/flyteadmin/auth/cookie_manager_test.go @@ -58,8 +58,8 @@ 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_at_1", c[1].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) }) @@ -139,8 +139,8 @@ 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_at_1", c[1].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) }) @@ -205,22 +205,27 @@ func TestCookieManager(t *testing.T) { manager.DeleteCookies(ctx, w) cookies := w.Result().Cookies() - require.Equal(t, 4, 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, accessTokenCookieNameSplit, cookies[1].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, refreshTokenCookieName, cookies[2].Name) + 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, idTokenCookieName, cookies[3].Name) + 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 e636b188c1..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(), 4) + 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(), 4) + 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(), 4) + require.Len(t, w.Result().Cookies(), 5) authCtx.AssertExpectations(t) hook.AssertExpectations(t) }) @@ -399,11 +399,11 @@ 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) - accessTokenCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplit, ".d.e.f", 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) From 3a9e4a18c2f9e48ee7b0590cb936fe0cf068968d Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Thu, 28 Mar 2024 11:39:31 -0700 Subject: [PATCH 5/6] update comment Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie_manager.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 66758275b8..000c9f8302 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -58,7 +58,7 @@ func (c CookieManager) RetrieveAccessToken(ctx context.Context, request *http.Re if err == nil && oldAccessToken != "" { return oldAccessToken, nil } - // If there is no old access token, we will retrieve the new access token + // If there is no old access token, we will retrieve the new splitted access token accessTokenFirstHalf, err := retrieveSecureCookie(ctx, request, accessTokenCookieNameSplitFirst, c.hashKey, c.blockKey) if err != nil { return "", err From 3d1103538d0a72024233b693fa5ed27d3e9f6ebd Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Thu, 28 Mar 2024 23:59:08 -0700 Subject: [PATCH 6/6] correct spelling... Signed-off-by: Yubo Wang --- flyteadmin/auth/cookie_manager.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go index 000c9f8302..ce360c9d3a 100644 --- a/flyteadmin/auth/cookie_manager.go +++ b/flyteadmin/auth/cookie_manager.go @@ -58,7 +58,7 @@ func (c CookieManager) RetrieveAccessToken(ctx context.Context, request *http.Re if err == nil && oldAccessToken != "" { return oldAccessToken, nil } - // If there is no old access token, we will retrieve the new splitted access token + // 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