Google: Add Support for Multiple Admin Emails to Retrieve Group Lists

connector/google/google.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
@@ -22,7 +23,8 @@ import (
 )
 
 const (
-	issuerURL = "https://accounts.google.com"
+	issuerURL                  = "https://accounts.google.com"
+	wildcardDomainToAdminEmail = "*"
 )
 
 // Config holds configuration options for Google logins.
@@ -46,17 +48,28 @@ type Config struct {
 	// check groups with the admin directory api
 	ServiceAccountFilePath string `json:"serviceAccountFilePath"`
 
+	// Deprecated: Use DomainToAdminEmail
+	AdminEmail string
+
 	// Required if ServiceAccountFilePath
-	// The email of a GSuite super user which the service account will impersonate
+	// The map workspace domain to email of a GSuite super user which the service account will impersonate
 	// when listing groups
-	AdminEmail string
+	DomainToAdminEmail map[string]string
 
 	// If this field is true, fetch direct group membership and transitive group membership
 	FetchTransitiveGroupMembership bool `json:"fetchTransitiveGroupMembership"`
 }
 
 // Open returns a connector which can be used to login users through Google.
 func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, err error) {
+	if c.AdminEmail != "" {
+		log.Deprecated(logger, `google: use "domainToAdminEmail.*: %s" option instead of "adminEmail: %s".`, c.AdminEmail, c.AdminEmail)
+		if c.DomainToAdminEmail == nil {
+			c.DomainToAdminEmail = make(map[string]string)
+		}
+
+		c.DomainToAdminEmail[wildcardDomainToAdminEmail] = c.AdminEmail
+	}
 	ctx, cancel := context.WithCancel(context.Background())
 
 	provider, err := oidc.NewProvider(ctx, issuerURL)
@@ -72,17 +85,26 @@ func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, e
 		scopes = append(scopes, "profile", "email")
 	}
 
-	var adminSrv *admin.Service
+	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")
+	}
 
 	// Fixing a regression caused by default config fallback: https://github.com/dexidp/dex/issues/2699
-	if (c.ServiceAccountFilePath != "" && c.AdminEmail != "") || slices.Contains(scopes, "groups") {
-		srv, err := createDirectoryService(c.ServiceAccountFilePath, c.AdminEmail, logger)
-		if err != nil {
-			cancel()
-			return nil, fmt.Errorf("could not create directory service: %v", err)
-		}
+	if (c.ServiceAccountFilePath != "" && 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 = srv
+			adminSrv[domain] = srv
+		}
 	}
 
 	clientID := c.ClientID
@@ -103,7 +125,7 @@ func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, e
 		hostedDomains:                  c.HostedDomains,
 		groups:                         c.Groups,
 		serviceAccountFilePath:         c.ServiceAccountFilePath,
-		adminEmail:                     c.AdminEmail,
+		domainToAdminEmail:             c.DomainToAdminEmail,
 		fetchTransitiveGroupMembership: c.FetchTransitiveGroupMembership,
 		adminSrv:                       adminSrv,
 	}, nil
@@ -123,9 +145,9 @@ type googleConnector struct {
 	hostedDomains                  []string
 	groups                         []string
 	serviceAccountFilePath         string
-	adminEmail                     string
+	domainToAdminEmail             map[string]string
 	fetchTransitiveGroupMembership bool
-	adminSrv                       *admin.Service
+	adminSrv                       map[string]*admin.Service
 }
 
 func (c *googleConnector) Close() error {
@@ -226,7 +248,7 @@ func (c *googleConnector) createIdentity(ctx context.Context, identity connector
 	}
 
 	var groups []string
-	if s.Groups && c.adminSrv != nil {
+	if s.Groups && len(c.adminSrv) > 0 {
 		checkedGroups := make(map[string]struct{})
 		groups, err = c.getGroups(claims.Email, c.fetchTransitiveGroupMembership, checkedGroups)
 		if err != nil {
@@ -258,8 +280,14 @@ func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership
 	var userGroups []string
 	var err error
 	groupsList := &admin.Groups{}
+	domain := c.extractDomainFromEmail(email)
+	adminSrv, err := c.findAdminService(domain)
+	if err != nil {
+		return nil, err
+	}
+
 	for {
-		groupsList, err = c.adminSrv.Groups.List().
+		groupsList, err = adminSrv.Groups.List().
 			UserKey(email).PageToken(groupsList.NextPageToken).Do()
 		if err != nil {
 			return nil, fmt.Errorf("could not list groups: %v", err)
@@ -295,16 +323,37 @@ func (c *googleConnector) getGroups(email string, fetchTransitiveGroupMembership
 	return userGroups, nil
 }
 
+func (c *googleConnector) findAdminService(domain string) (*admin.Service, error) {
+	adminSrv, ok := c.adminSrv[domain]
+	if !ok {
+		adminSrv, ok = c.adminSrv[wildcardDomainToAdminEmail]
+		c.logger.Debugf("using wildcard (%s) admin email to fetch groups", c.domainToAdminEmail[wildcardDomainToAdminEmail])
+	}
+
+	if !ok {
+		return nil, fmt.Errorf("unable to find super admin email, domainToAdminEmail for domain: %s not set, %s is also empty", domain, wildcardDomainToAdminEmail)
+	}
+
+	return adminSrv, nil
+}
+
+// extracts the domain name from an email input. If the email is valid, it returns the domain name after the "@" symbol.
+// However, in the case of a broken or invalid email, it returns a wildcard symbol.
+func (c *googleConnector) extractDomainFromEmail(email string) string {
+	at := strings.LastIndex(email, "@")
+	if at >= 0 {
+		_, domain := email[:at], email[at+1:]
+
+		return domain
+	}
+
+	return wildcardDomainToAdminEmail
+}
+
 // createDirectoryService sets up super user impersonation and creates an admin client for calling
 // the google admin api. If no serviceAccountFilePath is defined, the application default credential
 // is used.
 func createDirectoryService(serviceAccountFilePath, email string, logger log.Logger) (*admin.Service, error) {
-	// We know impersonation is required when using a service account credential
-	// TODO: or is it?
-	if email == "" && serviceAccountFilePath != "" {
-		return nil, fmt.Errorf("directory service requires adminEmail")
-	}
-
 	var jsonCredentials []byte
 	var err error
 

connector/google/google_test.go

@@ -102,15 +102,15 @@ func TestOpen(t *testing.T) {
 				Scopes:                 []string{"openid", "groups"},
 				ServiceAccountFilePath: serviceAccountFilePath,
 			},
-			expectedErr: "requires adminEmail",
+			expectedErr: "requires the domainToAdminEmail",
 		},
 		"service_account_key_not_found": {
 			config: &Config{
 				ClientID:               "testClient",
 				ClientSecret:           "testSecret",
 				RedirectURI:            ts.URL + "/callback",
 				Scopes:                 []string{"openid", "groups"},
-				AdminEmail:             "[email protected]",
+				DomainToAdminEmail:     map[string]string{"*": "[email protected]"},
 				ServiceAccountFilePath: "not_found.json",
 			},
 			expectedErr: "error reading credentials",
@@ -121,18 +121,18 @@ func TestOpen(t *testing.T) {
 				ClientSecret:           "testSecret",
 				RedirectURI:            ts.URL + "/callback",
 				Scopes:                 []string{"openid", "groups"},
-				AdminEmail:             "[email protected]",
+				DomainToAdminEmail:     map[string]string{"bar.com": "[email protected]"},
 				ServiceAccountFilePath: serviceAccountFilePath,
 			},
 			expectedErr: "",
 		},
 		"adc": {
 			config: &Config{
-				ClientID:     "testClient",
-				ClientSecret: "testSecret",
-				RedirectURI:  ts.URL + "/callback",
-				Scopes:       []string{"openid", "groups"},
-				AdminEmail:   "[email protected]",
+				ClientID:           "testClient",
+				ClientSecret:       "testSecret",
+				RedirectURI:        ts.URL + "/callback",
+				Scopes:             []string{"openid", "groups"},
+				DomainToAdminEmail: map[string]string{"*": "[email protected]"},
 			},
 			adc:         serviceAccountFilePath,
 			expectedErr: "",
@@ -143,7 +143,7 @@ func TestOpen(t *testing.T) {
 				ClientSecret:           "testSecret",
 				RedirectURI:            ts.URL + "/callback",
 				Scopes:                 []string{"openid", "groups"},
-				AdminEmail:             "[email protected]",
+				DomainToAdminEmail:     map[string]string{"*": "[email protected]"},
 				ServiceAccountFilePath: serviceAccountFilePath,
 			},
 			adc:         "/dev/null",
@@ -176,15 +176,15 @@ func TestGetGroups(t *testing.T) {
 
 	os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", serviceAccountFilePath)
 	conn, err := newConnector(&Config{
-		ClientID:     "testClient",
-		ClientSecret: "testSecret",
-		RedirectURI:  ts.URL + "/callback",
-		Scopes:       []string{"openid", "groups"},
-		AdminEmail:   "[email protected]",
+		ClientID:           "testClient",
+		ClientSecret:       "testSecret",
+		RedirectURI:        ts.URL + "/callback",
+		Scopes:             []string{"openid", "groups"},
+		DomainToAdminEmail: map[string]string{"*": "[email protected]"},
 	})
 	assert.Nil(t, err)
 
-	conn.adminSrv, err = admin.NewService(context.Background(), option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
+	conn.adminSrv[wildcardDomainToAdminEmail], err = admin.NewService(context.Background(), option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
 	assert.Nil(t, err)
 	type testCase struct {
 		userKey                        string
@@ -236,3 +236,58 @@ func TestGetGroups(t *testing.T) {
 		})
 	}
 }
+
+func TestDomainToAdminEmailConfig(t *testing.T) {
+	ts := testSetup()
+	defer ts.Close()
+
+	serviceAccountFilePath, err := tempServiceAccountKey()
+	assert.Nil(t, err)
+
+	os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", serviceAccountFilePath)
+	conn, err := newConnector(&Config{
+		ClientID:           "testClient",
+		ClientSecret:       "testSecret",
+		RedirectURI:        ts.URL + "/callback",
+		Scopes:             []string{"openid", "groups"},
+		DomainToAdminEmail: map[string]string{"dexidp.com": "[email protected]"},
+	})
+	assert.Nil(t, err)
+
+	conn.adminSrv["dexidp.com"], err = admin.NewService(context.Background(), option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
+	assert.Nil(t, err)
+	type testCase struct {
+		userKey     string
+		expectedErr string
+	}
+
+	for name, testCase := range map[string]testCase{
+		"correct_user_request": {
+			userKey:     "[email protected]",
+			expectedErr: "",
+		},
+		"wrong_user_request": {
+			userKey:     "[email protected]",
+			expectedErr: "unable to find super admin email",
+		},
+		"wrong_connector_response": {
+			userKey:     "user_1_foo.bar",
+			expectedErr: "unable to find super admin email",
+		},
+	} {
+		testCase := testCase
+		callCounter = map[string]int{}
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			lookup := make(map[string]struct{})
+
+			_, err := conn.getGroups(testCase.userKey, true, lookup)
+			if testCase.expectedErr != "" {
+				assert.ErrorContains(err, testCase.expectedErr)
+			} else {
+				assert.Nil(err)
+			}
+			t.Logf("[%s] Amount of API calls per userKey: %+v\n", t.Name(), callCounter)
+		})
+	}
+}