Translate access token instead of id token to gRPC backend

Signed-off-by: Jason Parraga <[email protected]>
flyteorg · Nov 8, 2024 · d46690b · d46690b
1 parent 25cfe16
commit d46690b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 14 deletions.
diff --git 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
@@ -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) {