From 3432b4fe32654f968ed0d83f68e046b530b3477d Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Thu, 8 Feb 2024 00:59:28 -0800 Subject: [PATCH] split access token into half and store --- 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=