From 3a245e49458b4799c105b34807ed8b2a4ebaf192 Mon Sep 17 00:00:00 2001
From: Douglas Adams <douglasryanadams@gmail.com>
Date: Mon, 10 Jul 2023 06:37:00 -0400
Subject: [PATCH] Auth: Fix US gov azure ad oauth URL parsing (#71254)

Updates regex for tenant ID parsing to support .us domains in addition
to .com domains for Azure AD.

Fixes #71252
---
 pkg/login/social/azuread_oauth.go      |  6 ++---
 pkg/login/social/azuread_oauth_test.go | 31 +++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/pkg/login/social/azuread_oauth.go b/pkg/login/social/azuread_oauth.go
index ace3c546abd..d05cb400340 100644
--- a/pkg/login/social/azuread_oauth.go
+++ b/pkg/login/social/azuread_oauth.go
@@ -334,7 +334,7 @@ func (s *SocialAzureAD) SupportBundleContent(bf *bytes.Buffer) error {
 
 func (s *SocialAzureAD) extractTenantID(authURL string) (string, error) {
 	if s.compiledTenantRegex == nil {
-		compiledTenantRegex, err := regexp.Compile(`https://login.microsoftonline.com/([^/]+)/oauth2`)
+		compiledTenantRegex, err := regexp.Compile(`https://login.microsoftonline.(com|us)/([^/]+)/oauth2`)
 		if err != nil {
 			return "", err
 		}
@@ -342,10 +342,10 @@ func (s *SocialAzureAD) extractTenantID(authURL string) (string, error) {
 	}
 
 	matches := s.compiledTenantRegex.FindStringSubmatch(authURL)
-	if len(matches) < 2 {
+	if len(matches) < 3 {
 		return "", fmt.Errorf("unable to extract tenant ID from URL")
 	}
-	return matches[1], nil
+	return matches[2], nil
 }
 
 func (s *SocialAzureAD) retrieveJWKS(client *http.Client) (*jose.JSONWebKeySet, error) {
diff --git a/pkg/login/social/azuread_oauth_test.go b/pkg/login/social/azuread_oauth_test.go
index 310282ffb99..490564573fa 100644
--- a/pkg/login/social/azuread_oauth_test.go
+++ b/pkg/login/social/azuread_oauth_test.go
@@ -36,6 +36,7 @@ func TestSocialAzureAD_UserInfo(t *testing.T) {
 		allowedGroups        []string
 		allowedOrganizations []string
 		forceUseGraphAPI     bool
+		usGovURL             bool
 	}
 	type args struct {
 		client *http.Client
@@ -89,6 +90,28 @@ func TestSocialAzureAD_UserInfo(t *testing.T) {
 			want:    nil,
 			wantErr: true,
 		},
+		{
+			name: "US Government domain",
+			claims: &azureClaims{
+				Email:             "me@example.com",
+				PreferredUsername: "",
+				Roles:             []string{},
+				Name:              "My Name",
+				ID:                "1234",
+			},
+			fields: fields{
+				SocialBase: newSocialBase("azuread", &oauth2.Config{}, &OAuthInfo{}, "Viewer", false, *featuremgmt.WithFeatures()),
+				usGovURL:   true,
+			},
+			want: &BasicUserInfo{
+				Id:     "1234",
+				Name:   "My Name",
+				Email:  "me@example.com",
+				Login:  "me@example.com",
+				Role:   "Viewer",
+				Groups: []string{},
+			},
+		},
 		{
 			name: "Email in preferred_username claim",
 			claims: &azureClaims{
@@ -476,6 +499,8 @@ func TestSocialAzureAD_UserInfo(t *testing.T) {
 	}
 
 	authURL := "https://login.microsoftonline.com/1234/oauth2/v2.0/authorize"
+	usGovAuthURL := "https://login.microsoftonline.us/1234/oauth2/v2.0/authorize"
+
 	cache := remotecache.NewFakeCacheStorage()
 	// put JWKS in cache
 	jwksDump, err := json.Marshal(jwks)
@@ -498,7 +523,11 @@ func TestSocialAzureAD_UserInfo(t *testing.T) {
 				s.SocialBase = newSocialBase("azuread", &oauth2.Config{}, &OAuthInfo{}, "", false, *featuremgmt.WithFeatures())
 			}
 
-			s.SocialBase.Endpoint.AuthURL = authURL
+			if tt.fields.usGovURL {
+				s.SocialBase.Endpoint.AuthURL = usGovAuthURL
+			} else {
+				s.SocialBase.Endpoint.AuthURL = authURL
+			}
 
 			cl := jwt.Claims{
 				Subject:   "subject",