From 2b7514b1a48147626e93346b0006cf589429b4aa Mon Sep 17 00:00:00 2001
From: gaohan <1135494872@qq.com>
Date: Tue, 8 Jun 2021 17:26:04 +0800
Subject: [PATCH] Add domain_hint in aad-stater. (#22134)

---
 .../CHANGELOG.md                                |  2 +-
 sdk/spring/azure-spring-boot/CHANGELOG.md       |  2 +-
 .../AADOAuth2AuthorizationRequestResolver.java  | 17 ++++++++++++-----
 .../webapp/AADWebSecurityConfigurerAdapter.java |  2 +-
 .../aad/AADAuthenticationProperties.java        | 13 +++++++++++++
 5 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/sdk/spring/azure-spring-boot-starter-active-directory/CHANGELOG.md b/sdk/spring/azure-spring-boot-starter-active-directory/CHANGELOG.md
index d5b665c1ef005..d7189f5a4bd80 100644
--- a/sdk/spring/azure-spring-boot-starter-active-directory/CHANGELOG.md
+++ b/sdk/spring/azure-spring-boot-starter-active-directory/CHANGELOG.md
@@ -1,7 +1,7 @@
 # Release History
 
 ## 3.6.0-beta.1 (Unreleased)
-
+- Support domain_hint in aad-starter.([#21517](https://github.com/Azure/azure-sdk-for-java/issues/21517))
 
 ## 3.5.0 (2021-05-24)
 ### New Features
diff --git a/sdk/spring/azure-spring-boot/CHANGELOG.md b/sdk/spring/azure-spring-boot/CHANGELOG.md
index c01299e98334d..c71ce3e525850 100644
--- a/sdk/spring/azure-spring-boot/CHANGELOG.md
+++ b/sdk/spring/azure-spring-boot/CHANGELOG.md
@@ -1,7 +1,7 @@
 # Release History
 
 ## 3.6.0-beta.1 (Unreleased)
-
+- Support domain_hint in aad-starter.([#21517](https://github.com/Azure/azure-sdk-for-java/issues/21517))
 
 ## 3.5.0 (2021-05-24)
 ### New Features
diff --git a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizationRequestResolver.java b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizationRequestResolver.java
index 2f3680065c058..a449ed1719309 100644
--- a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizationRequestResolver.java
+++ b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizationRequestResolver.java
@@ -3,6 +3,7 @@
 
 package com.azure.spring.aad.webapp;
 
+import com.azure.spring.autoconfigure.aad.AADAuthenticationProperties;
 import com.azure.spring.autoconfigure.aad.Constants;
 import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
 import org.springframework.security.oauth2.client.web.DefaultOAuth2AuthorizationRequestResolver;
@@ -21,11 +22,15 @@
 public class AADOAuth2AuthorizationRequestResolver implements OAuth2AuthorizationRequestResolver {
     private final OAuth2AuthorizationRequestResolver defaultResolver;
 
-    public AADOAuth2AuthorizationRequestResolver(ClientRegistrationRepository clientRegistrationRepository) {
+    private final AADAuthenticationProperties properties;
+
+    public AADOAuth2AuthorizationRequestResolver(ClientRegistrationRepository clientRegistrationRepository,
+                                                 AADAuthenticationProperties properties) {
         this.defaultResolver = new DefaultOAuth2AuthorizationRequestResolver(
             clientRegistrationRepository,
             OAuth2AuthorizationRequestRedirectFilter.DEFAULT_AUTHORIZATION_REQUEST_BASE_URI
         );
+        this.properties = properties;
     }
 
     @Override
@@ -56,11 +61,13 @@ private OAuth2AuthorizationRequest addClaims(HttpServletRequest httpServletReque
                         return claims;
                     })
                     .orElse(null);
-        if (conditionalAccessPolicyClaims == null) {
-            return oAuth2AuthorizationRequest;
-        }
         final Map<String, Object> additionalParameters = new HashMap<>();
-        additionalParameters.put(Constants.CLAIMS, conditionalAccessPolicyClaims);
+        if (conditionalAccessPolicyClaims != null) {
+            additionalParameters.put(Constants.CLAIMS, conditionalAccessPolicyClaims);
+        }
+        Optional.ofNullable(properties)
+                .map(AADAuthenticationProperties::getAuthenticateAdditionalParameters)
+                .ifPresent(additionalParameters::putAll);
         Optional.of(oAuth2AuthorizationRequest)
                 .map(OAuth2AuthorizationRequest::getAdditionalParameters)
                 .ifPresent(additionalParameters::putAll);
diff --git a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADWebSecurityConfigurerAdapter.java b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADWebSecurityConfigurerAdapter.java
index 0938f3d343904..5988440d5dd35 100644
--- a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADWebSecurityConfigurerAdapter.java
+++ b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/aad/webapp/AADWebSecurityConfigurerAdapter.java
@@ -71,6 +71,6 @@ protected OAuth2AccessTokenResponseClient<OAuth2AuthorizationCodeGrantRequest> a
     }
 
     protected OAuth2AuthorizationRequestResolver requestResolver() {
-        return new AADOAuth2AuthorizationRequestResolver(this.repo);
+        return new AADOAuth2AuthorizationRequestResolver(this.repo, properties);
     }
 }
diff --git a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/autoconfigure/aad/AADAuthenticationProperties.java b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/autoconfigure/aad/AADAuthenticationProperties.java
index b62ecbcf44932..c6472727377e6 100644
--- a/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/autoconfigure/aad/AADAuthenticationProperties.java
+++ b/sdk/spring/azure-spring-boot/src/main/java/com/azure/spring/autoconfigure/aad/AADAuthenticationProperties.java
@@ -64,6 +64,11 @@ public class AADAuthenticationProperties implements InitializingBean {
      */
     private String appIdUri;
 
+    /**
+     * Add additional parameters to the Authorization URL.
+     */
+    private Map<String, Object> authenticateAdditionalParameters;
+
     /**
      * Connection Timeout for the JWKSet Remote URL call.
      */
@@ -248,6 +253,14 @@ public void setAppIdUri(String appIdUri) {
         this.appIdUri = appIdUri;
     }
 
+    public Map<String, Object> getAuthenticateAdditionalParameters() {
+        return authenticateAdditionalParameters;
+    }
+
+    public void setAuthenticateAdditionalParameters(Map<String, Object> authenticateAdditionalParameters) {
+        this.authenticateAdditionalParameters = authenticateAdditionalParameters;
+    }
+
     public int getJwtConnectTimeout() {
         return jwtConnectTimeout;
     }