keep old name and store new splitted names

Signed-off-by: Yubo Wang <[email protected]>

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

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))
 }

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) {

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)