From 7e8b679237328c6df9d6f52d61751b980c38b902 Mon Sep 17 00:00:00 2001
From: linoman <2051016+linoman@users.noreply.github.com>
Date: Wed, 21 Feb 2024 08:40:18 -0600
Subject: [PATCH] OAuth: Improve domain validation (#83110)

* enforce hd claim validation

* add tests
---
 pkg/login/social/connectors/google_oauth.go   | 20 +++++++++
 .../social/connectors/google_oauth_test.go    | 44 +++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/pkg/login/social/connectors/google_oauth.go b/pkg/login/social/connectors/google_oauth.go
index 6709d4e05cf..052691bb754 100644
--- a/pkg/login/social/connectors/google_oauth.go
+++ b/pkg/login/social/connectors/google_oauth.go
@@ -17,6 +17,7 @@ import (
 	ssoModels "github.com/grafana/grafana/pkg/services/ssosettings/models"
 	"github.com/grafana/grafana/pkg/services/ssosettings/validation"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/util/errutil"
 )
 
 const (
@@ -37,6 +38,7 @@ type googleUserData struct {
 	Email         string `json:"email"`
 	Name          string `json:"name"`
 	EmailVerified bool   `json:"email_verified"`
+	HD            string `json:"hd"`
 	rawJSON       []byte `json:"-"`
 }
 
@@ -115,6 +117,10 @@ func (s *SocialGoogle) UserInfo(ctx context.Context, client *http.Client, token
 		return nil, fmt.Errorf("user email is not verified")
 	}
 
+	if err := s.isHDAllowed(data.HD); err != nil {
+		return nil, err
+	}
+
 	groups, errPage := s.retrieveGroups(ctx, client, data)
 	if errPage != nil {
 		s.log.Warn("Error retrieving groups", "error", errPage)
@@ -290,3 +296,17 @@ func (s *SocialGoogle) getGroupsPage(ctx context.Context, client *http.Client, u
 
 	return &data, nil
 }
+
+func (s *SocialGoogle) isHDAllowed(hd string) error {
+	if len(s.info.AllowedDomains) == 0 {
+		return nil
+	}
+
+	for _, allowedDomain := range s.info.AllowedDomains {
+		if hd == allowedDomain {
+			return nil
+		}
+	}
+
+	return errutil.Forbidden("the hd claim found in the ID token is not present in the allowed domains", errutil.WithPublicMessage("Invalid domain"))
+}
diff --git a/pkg/login/social/connectors/google_oauth_test.go b/pkg/login/social/connectors/google_oauth_test.go
index 95345c94e53..7fe0d8b0244 100644
--- a/pkg/login/social/connectors/google_oauth_test.go
+++ b/pkg/login/social/connectors/google_oauth_test.go
@@ -890,3 +890,47 @@ func TestSocialGoogle_Reload(t *testing.T) {
 		})
 	}
 }
+
+func TestIsHDAllowed(t *testing.T) {
+	testCases := []struct {
+		name                 string
+		email                string
+		allowedDomains       []string
+		expectedErrorMessage string
+	}{
+		{
+			name:                 "should not fail if no allowed domains are set",
+			email:                "mycompany.com",
+			allowedDomains:       []string{},
+			expectedErrorMessage: "",
+		},
+		{
+			name:                 "should not fail if email is from allowed domain",
+			email:                "mycompany.com",
+			allowedDomains:       []string{"grafana.com", "mycompany.com", "example.com"},
+			expectedErrorMessage: "",
+		},
+		{
+			name:                 "should fail if email is not from allowed domain",
+			email:                "mycompany.com",
+			allowedDomains:       []string{"grafana.com", "example.com"},
+			expectedErrorMessage: "the hd claim found in the ID token is not present in the allowed domains",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			info := &social.OAuthInfo{}
+			info.AllowedDomains = tc.allowedDomains
+			s := NewGoogleProvider(info, &setting.Cfg{}, &ssosettingstests.MockService{}, featuremgmt.WithFeatures())
+			err := s.isHDAllowed(tc.email)
+
+			if tc.expectedErrorMessage != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.expectedErrorMessage)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}