From cb9c6b8856820273c313618d4ed3b5c23f44d85b Mon Sep 17 00:00:00 2001
From: Yubo Wang <yubowang2019@gmail.com>
Date: Tue, 26 Mar 2024 08:45:30 +0800
Subject: [PATCH] lint and fix unit test

---
 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 d5ccb163f3d..3e2c69d57ef 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 359219d3ab2..554006ce6ca 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 452f797d9f6..50f4761670f 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) {