feat: allow Workload Identity Federation for Google connector

A new Google connector option, `useCloudIdentityApi`, has been introduced.
If the value it `true`, dex will use cloud identity api to fetch
groups. In particular, no user impersonation happens. The linked service
account needs to be a assigned to a custom admin role created in Google
workspace. This role requires group read rights.
assumed that Workload Identity Federation shall be used, and the

Moreover, if the JSON propvided for the service account is not of type
`service_account` it is assumed it's configued for Workload Identity
Federation. In that case, the linked service account needs to include
`Service Account Token Creator`.

Signed-off-by: Michael Dudzinski <michael.dudzinski@aetherize.com>

Author: midu-git

connector/google/google.go

@@ -17,6 +17,7 @@ import (
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
 	admin "google.golang.org/api/admin/directory/v1"
+	"google.golang.org/api/cloudidentity/v1"
 	"google.golang.org/api/impersonate"
 	"google.golang.org/api/option"
 
@@ -53,11 +54,14 @@ type Config struct {
 	// Deprecated: Use DomainToAdminEmail
 	AdminEmail string
 
-	// Required if ServiceAccountFilePath
+	// If ServiceAccountFilePath is set, this value is ignored if UseCloudIdentityAPI is set. Otherwise, it's required.
 	// The map workspace domain to email of a GSuite super user which the service account will impersonate
 	// when listing groups
 	DomainToAdminEmail map[string]string
 
+	// If set, Cloud Identity API is used to fetch groups for a user. Defaults to false.
+	UseCloudIdentityAPI bool `json:"useCloudIdentityAPI"`
+
 	// If this field is true, fetch direct group membership and transitive group membership
 	FetchTransitiveGroupMembership bool `json:"fetchTransitiveGroupMembership"`
 
@@ -66,6 +70,14 @@ type Config struct {
 	PromptType *string `json:"promptType"`
 }
 
+func validateConfigForCloudIdentity(c *Config, logger *slog.Logger) error {
+	if len(c.DomainToAdminEmail) > 0 || len(c.AdminEmail) > 0 {
+		logger.Warn("For cloud identity calls \"DomainToAdminEmail\" and \"AdminEmail\" are ignored. It's safe to remove both configuration options.")
+	}
+
+	return nil
+}
+
 // Open returns a connector which can be used to login users through Google.
 func (c *Config) Open(id string, logger *slog.Logger) (conn connector.Connector, err error) {
 	logger = logger.With(slog.Group("connector", "type", "google", "id", id))
@@ -94,22 +106,38 @@ func (c *Config) Open(id string, logger *slog.Logger) (conn connector.Connector,
 
 	adminSrv := make(map[string]*admin.Service)
 
-	// We know impersonation is required when using a service account credential
-	// TODO: or is it?
-	if len(c.DomainToAdminEmail) == 0 && c.ServiceAccountFilePath != "" {
-		cancel()
-		return nil, fmt.Errorf("directory service requires the domainToAdminEmail option to be configured")
-	}
+	var groupsMembershipsService *cloudidentity.GroupsMembershipsService
 
-	if (len(c.DomainToAdminEmail) > 0) || slices.Contains(scopes, "groups") {
-		for domain, adminEmail := range c.DomainToAdminEmail {
-			srv, err := createDirectoryService(c.ServiceAccountFilePath, adminEmail, logger)
-			if err != nil {
-				cancel()
-				return nil, fmt.Errorf("could not create directory service: %v", err)
-			}
+	if c.UseCloudIdentityAPI {
+		err = validateConfigForCloudIdentity(c, logger)
+		if err != nil {
+			cancel()
+			return nil, err
+		}
+
+		groupsMembershipsService, err = createGroupsMembershipsService(c.ServiceAccountFilePath, logger)
+		if err != nil {
+			cancel()
+			return nil, fmt.Errorf("could not create groups memebership service: %v", err)
+		}
+	} else {
+		// We know impersonation is required when using a service account credential
+		// TODO: or is it?
+		if len(c.DomainToAdminEmail) == 0 && c.ServiceAccountFilePath != "" {
+			cancel()
+			return nil, fmt.Errorf("directory service requires the domainToAdminEmail option to be configured")
+		}
+
+		if (len(c.DomainToAdminEmail) > 0) || slices.Contains(scopes, "groups") {
+			for domain, adminEmail := range c.DomainToAdminEmail {
+				srv, err := createDirectoryService(c.ServiceAccountFilePath, adminEmail, logger)
+				if err != nil {
+					cancel()
+					return nil, fmt.Errorf("could not create directory service: %v", err)
+				}
 
-			adminSrv[domain] = srv
+				adminSrv[domain] = srv
+			}
 		}
 	}
 
@@ -137,8 +165,10 @@ func (c *Config) Open(id string, logger *slog.Logger) (conn connector.Connector,
 		groups:                         c.Groups,
 		serviceAccountFilePath:         c.ServiceAccountFilePath,
 		domainToAdminEmail:             c.DomainToAdminEmail,
+		useCloudIdentityAPI:            c.UseCloudIdentityAPI,
 		fetchTransitiveGroupMembership: c.FetchTransitiveGroupMembership,
 		adminSrv:                       adminSrv,
+		groupsMembershipsService:       groupsMembershipsService,
 		promptType:                     promptType,
 	}, nil
 }
@@ -158,8 +188,10 @@ type googleConnector struct {
 	groups                         []string
 	serviceAccountFilePath         string
 	domainToAdminEmail             map[string]string
+	useCloudIdentityAPI            bool
 	fetchTransitiveGroupMembership bool
 	adminSrv                       map[string]*admin.Service
+	groupsMembershipsService       *cloudidentity.GroupsMembershipsService
 	promptType                     string
 }
 
@@ -262,14 +294,25 @@ func (c *googleConnector) createIdentity(ctx context.Context, identity connector
 	}
 
 	var groups []string
-	if s.Groups && len(c.adminSrv) > 0 {
+	if s.Groups {
 		checkedGroups := make(map[string]struct{})
-		groups, err = c.getGroups(claims.Email, c.fetchTransitiveGroupMembership, checkedGroups)
-		if err != nil {
-			return identity, fmt.Errorf("google: could not retrieve groups: %v", err)
+		if c.useCloudIdentityAPI {
+			if c.groupsMembershipsService != nil {
+				groups, err = c.getGroupsFromCloudIdentityAPI(claims.Email, c.fetchTransitiveGroupMembership, checkedGroups)
+				if err != nil {
+					return identity, fmt.Errorf("google: could not retrieve groups from Cloud Identity API: %v", err)
+				}
+			}
+		} else {
+			if len(c.adminSrv) > 0 {
+				groups, err = c.getGroupsFromAdminAPI(claims.Email, c.fetchTransitiveGroupMembership, checkedGroups)
+				if err != nil {
+					return identity, fmt.Errorf("google: could not retrieve groups form Admin API: %v", err)
+				}
+			}
 		}
 
-		if len(c.groups) > 0 {
+		if len(c.groups) > 0 && len(groups) > 0 {
 			groups = pkg_groups.Filter(groups, c.groups)
 			if len(groups) == 0 {
 				return identity, fmt.Errorf("google: user %q is not in any of the required groups", claims.Username)
@@ -288,9 +331,9 @@ func (c *googleConnector) createIdentity(ctx context.Context, identity connector
 	return identity, nil
 }
 
-// getGroups creates a connection to the admin directory service and lists
+// getGroupsFromAdminAPI creates a connection to the admin directory service and lists
 // all groups the user is a member of
-func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership bool, checkedGroups map[string]struct{}) ([]string, error) {
+func (c *googleConnector) getGroupsFromAdminAPI(email string, fetchTransitiveGroupMembership bool, checkedGroups map[string]struct{}) ([]string, error) {
 	var userGroups []string
 	var err error
 	groupsList := &admin.Groups{}
@@ -321,7 +364,54 @@ func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership
 			}
 
 			// getGroups takes a user's email/alias as well as a group's email/alias
-			transitiveGroups, err := c.getGroups(group.Email, fetchTransitiveGroupMembership, checkedGroups)
+			transitiveGroups, err := c.getGroupsFromAdminAPI(group.Email, fetchTransitiveGroupMembership, checkedGroups)
+			if err != nil {
+				return nil, fmt.Errorf("could not list transitive groups: %v", err)
+			}
+
+			userGroups = append(userGroups, transitiveGroups...)
+		}
+
+		if groupsList.NextPageToken == "" {
+			break
+		}
+	}
+
+	return userGroups, nil
+}
+
+// getGroupsFromCloudIdentityAPI creates a connection to the cloud identity service and lists
+// all groups the user is a member of
+func (c *googleConnector) getGroupsFromCloudIdentityAPI(email string, fetchTransitiveGroupMembership bool, checkedGroups map[string]struct{}) ([]string, error) {
+	var userGroups []string
+	var err error
+	groupsList := &cloudidentity.SearchDirectGroupsResponse{}
+	groupsMembershipService := c.groupsMembershipsService
+
+	for {
+		query := fmt.Sprintf("member_key_id=='%s'", email)
+		groupsList, err = groupsMembershipService.SearchDirectGroups("groups/-").
+			Query(query).PageToken(groupsList.NextPageToken).Do()
+		if err != nil {
+			return nil, fmt.Errorf("could not list groups: %v", err)
+		}
+
+		for _, membership := range groupsList.Memberships {
+			groupEmail := membership.GroupKey.Id
+			if _, exists := checkedGroups[groupEmail]; exists {
+				continue
+			}
+
+			checkedGroups[groupEmail] = struct{}{}
+			// TODO (joelspeed): Make desired group key configurable
+			userGroups = append(userGroups, groupEmail)
+
+			if !fetchTransitiveGroupMembership {
+				continue
+			}
+
+			// getGroups takes a user's email/alias as well as a group's email/alias
+			transitiveGroups, err := c.getGroupsFromCloudIdentityAPI(groupEmail, fetchTransitiveGroupMembership, checkedGroups)
 			if err != nil {
 				return nil, fmt.Errorf("could not list transitive groups: %v", err)
 			}
@@ -455,3 +545,45 @@ func createDirectoryService(serviceAccountFilePath, email string, logger *slog.L
 
 	return admin.NewService(ctx, option.WithHTTPClient(config.Client(ctx)))
 }
+
+func createGroupsMembershipsService(serviceAccountFilePath string, logger *slog.Logger) (service *cloudidentity.GroupsMembershipsService, err error) {
+	var jsonCredentials []byte
+
+	ctx := context.Background()
+	if serviceAccountFilePath == "" {
+		logger.Warn("the application default credential is used since the service account file path is not used")
+		credential, err := google.FindDefaultCredentials(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to fetch application default credentials: %w", err)
+		}
+
+		if credential.JSON == nil {
+			return nil, fmt.Errorf("application default credentials returned empty JSON")
+		}
+
+		jsonCredentials = credential.JSON
+	} else {
+		logger.Debug("Using credentials at", "sa_path", serviceAccountFilePath)
+		jsonCredentials, err = os.ReadFile(serviceAccountFilePath)
+		if err != nil {
+			return nil, fmt.Errorf("error reading credentials from file: %v", err)
+		}
+	}
+
+	var cloudIdentityService *cloudidentity.Service
+
+	config, err := google.JWTConfigFromJSON(jsonCredentials, cloudidentity.CloudIdentityGroupsReadonlyScope)
+
+	if err == nil {
+		cloudIdentityService, err = cloudidentity.NewService(ctx, option.WithHTTPClient(config.Client(ctx)))
+	} else {
+		// most probably type != service_account
+		cloudIdentityService, err = cloudidentity.NewService(ctx, option.WithCredentialsJSON(jsonCredentials), option.WithScopes(cloudidentity.CloudIdentityGroupsReadonlyScope))
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("error creating cloud identity service: %v", err)
+	}
+
+	return cloudidentity.NewGroupsMembershipsService(cloudIdentityService), nil
+}