split access token into half and store

flyteadmin/auth/cookie.go

@@ -19,6 +19,8 @@ import (
 const (
 	// #nosec
 	accessTokenCookieName = "flyte_at"
+	// nosec
+	accessTokenCookieNameSplit = "flyte_at_1"
 	// #nosec
 	idTokenCookieName = "flyte_idt"
 	// #nosec

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

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":              "[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", 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) {

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

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=