From d46690ba4729bd0be5fca0bb2d68931336db638c Mon Sep 17 00:00:00 2001 From: Jason Parraga Date: Fri, 8 Nov 2024 13:06:48 -0800 Subject: [PATCH] Translate access token instead of id token to gRPC backend Signed-off-by: Jason Parraga --- flyteadmin/auth/handlers.go | 19 ++++++++++++++----- flyteadmin/auth/handlers_test.go | 22 +++++++++++++--------- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/flyteadmin/auth/handlers.go b/flyteadmin/auth/handlers.go index d8bc626652..29ba10067e 100644 --- a/flyteadmin/auth/handlers.go +++ b/flyteadmin/auth/handlers.go @@ -353,8 +353,8 @@ func WithAuditFields(ctx context.Context, subject string, clientIds []string, to func GetHTTPRequestCookieToMetadataHandler(authCtx interfaces.AuthenticationContext) HTTPRequestToMetadataAnnotator { return func(ctx context.Context, request *http.Request) metadata.MD { // TODO: Improve error handling - idToken, _, _, _ := authCtx.CookieManager().RetrieveTokenValues(ctx, request) - if len(idToken) == 0 { + idToken, accessToken, _, _ := authCtx.CookieManager().RetrieveTokenValues(ctx, request) + if len(idToken) == 0 && len(accessToken) == 0 { // If no token was found in the cookies, look for an authorization header, starting with a potentially // custom header set in the Config object if len(authCtx.Options().HTTPAuthorizationHeader) > 0 { @@ -372,9 +372,18 @@ func GetHTTPRequestCookieToMetadataHandler(authCtx interfaces.AuthenticationCont return nil } - // IDtoken is injected into grpc authorization metadata - meta := metadata.MD{ - DefaultAuthorizationHeader: []string{fmt.Sprintf("%s %s", IDTokenScheme, idToken)}, + var meta metadata.MD + + if len(accessToken) > 0 { + // Access token is injected into grpc authorization metadata + meta = metadata.MD{ + DefaultAuthorizationHeader: []string{fmt.Sprintf("%s %s", BearerScheme, accessToken)}, + } + } else { + // IDtoken is injected into grpc authorization metadata + meta = metadata.MD{ + DefaultAuthorizationHeader: []string{fmt.Sprintf("%s %s", IDTokenScheme, idToken)}, + } } userInfo, err := authCtx.CookieManager().RetrieveUserInfo(ctx, request) diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go index ee106e92cb..8866f6d1c7 100644 --- a/flyteadmin/auth/handlers_test.go +++ b/flyteadmin/auth/handlers_test.go @@ -396,22 +396,26 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) { mockAuthCtx.OnCookieManager().Return(&cookieManager) mockAuthCtx.OnOptions().Return(&config.Config{}) handler := GetHTTPRequestCookieToMetadataHandler(&mockAuthCtx) - req, err := http.NewRequest("GET", "/api/v1/projects", nil) - assert.NoError(t, err) - accessTokenCookie, err := NewSecureCookie(accessTokenCookieNameSplitFirst, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) + accessTokenCookie1, err := NewSecureCookie(accessTokenCookieNameSplitFirst, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) assert.NoError(t, err) - req.AddCookie(&accessTokenCookie) - accessTokenCookieSplit, err := NewSecureCookie(accessTokenCookieNameSplitSecond, ".d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) + accessTokenCookie2, err := NewSecureCookie(accessTokenCookieNameSplitSecond, ".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) + idCookie, err := NewSecureCookie(idTokenCookieName, "x.y.z", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode) assert.NoError(t, err) - req.AddCookie(&idCookie) - assert.Equal(t, "IDToken a.b.c.d.e.f", handler(ctx, req)["authorization"][0]) + t.Run("access token and ID token cookies present", func(t *testing.T) { + req, err := http.NewRequest("GET", "/api/v1/projects", nil) + assert.NoError(t, err) + + req.AddCookie(&accessTokenCookie1) + req.AddCookie(&accessTokenCookie2) + req.AddCookie(&idCookie) + + assert.Equal(t, "Bearer a.b.c.d.e.f", handler(ctx, req)["authorization"][0]) + }) } func TestGetHTTPMetadataTaggingHandler(t *testing.T) {