Repository: markbates/goth
Pull request #325 (2020-03-17 16:56:32)

Added hosted domain (hd) parameter to google provider
Added possibility to add `hd` parameter to Auth URL, the option is quite handy when writing an enterprise app authenticated via Google Gsuite. 


Patch:
diff --git a/providers/google/google.go b/providers/google/google.go
index 33c7b86..bb5543f 100644
--- a/providers/google/google.go
+++ b/providers/google/google.go
@@ -32,13 +32,13 @@ func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
 
 // Provider is the implementation of `goth.Provider` for accessing Google.
 type Provider struct {
-	ClientKey    string
-	Secret       string
-	CallbackURL  string
-	HTTPClient   *http.Client
-	config       *oauth2.Config
-	prompt       oauth2.AuthCodeOption
-	providerName string
+	ClientKey       string
+	Secret          string
+	CallbackURL     string
+	HTTPClient      *http.Client
+	config          *oauth2.Config
+	authCodeOptions []oauth2.AuthCodeOption
+	providerName    string
 }
 
 // Name is the name used to retrieve this provider later.
@@ -61,11 +61,7 @@ func (p *Provider) Debug(debug bool) {}
 
 // BeginAuth asks Google for an authentication endpoint.
 func (p *Provider) BeginAuth(state string) (goth.Session, error) {
-	var opts []oauth2.AuthCodeOption
-	if p.prompt != nil {
-		opts = append(opts, p.prompt)
-	}
-	url := p.config.AuthCodeURL(state, opts...)
+	url := p.config.AuthCodeURL(state, p.authCodeOptions...)
 	session := &Session{
 		AuthURL: url,
 	}
@@ -152,12 +148,12 @@ func newConfig(provider *Provider, scopes []string) *oauth2.Config {
 	return c
 }
 
-//RefreshTokenAvailable refresh token is provided by auth provider or not
+// RefreshTokenAvailable refresh token is provided by auth provider or not
 func (p *Provider) RefreshTokenAvailable() bool {
 	return true
 }
 
-//RefreshToken get new access token based on the refresh token
+// RefreshToken get new access token based on the refresh token
 func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
 	token := &oauth2.Token{RefreshToken: refreshToken}
 	ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
@@ -176,5 +172,15 @@ func (p *Provider) SetPrompt(prompt ...string) {
 	if len(prompt) == 0 {
 		return
 	}
-	p.prompt = oauth2.SetAuthURLParam("prompt", strings.Join(prompt, " "))
+	p.authCodeOptions = append(p.authCodeOptions, oauth2.SetAuthURLParam("prompt", strings.Join(prompt, " ")))
+}
+
+// SetHostedDomain sets the hd parameter for google OAuth call.
+// Use this to force user to pick user from specific hosted domain.
+// See https://developers.google.com/identity/protocols/oauth2/openid-connect#hd-param
+func (p *Provider) SetHostedDomain(hd string) {
+	if hd == "" {
+		return
+	}
+	p.authCodeOptions = append(p.authCodeOptions, oauth2.SetAuthURLParam("hd", hd))
 }


Test patch:
diff --git a/providers/google/google_test.go b/providers/google/google_test.go
index c02460f..53f889e 100644
--- a/providers/google/google_test.go
+++ b/providers/google/google_test.go
@@ -53,6 +53,25 @@ func Test_BeginAuthWithPrompt(t *testing.T) {
 	a.Contains(s.AuthURL, "prompt=test+prompts")
 }
 
+func Test_BeginAuthWithHostedDomain(t *testing.T) {
+	// This exists because there was a panic caused by the oauth2 package when
+	// the AuthCodeOption passed was nil. This test uses it, Test_BeginAuth does
+	// not, to ensure both cases are covered.
+	t.Parallel()
+	a := assert.New(t)
+
+	provider := googleProvider()
+	provider.SetHostedDomain("example.com")
+	session, err := provider.BeginAuth("test_state")
+	s := session.(*google.Session)
+	a.NoError(err)
+	a.Contains(s.AuthURL, "accounts.google.com/o/oauth2/auth")
+	a.Contains(s.AuthURL, fmt.Sprintf("client_id=%s", os.Getenv("GOOGLE_KEY")))
+	a.Contains(s.AuthURL, "state=test_state")
+	a.Contains(s.AuthURL, "scope=email")
+	a.Contains(s.AuthURL, "hd=example.com")
+}
+
 func Test_Implements_Provider(t *testing.T) {
 	t.Parallel()
 	a := assert.New(t)