Fix bug: OAuth2AuthorizedClient not saved for ClientNeedConsentWhenLo…

…gin (#18715)
Azure · Jan 21, 2021 · 00e1c8e · 00e1c8e
1 parent 980e148
commit 00e1c8e
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/...g-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizedClientRepository.java b/...g-boot/src/main/java/com/azure/spring/aad/webapp/AADOAuth2AuthorizedClientRepository.java
@@ -4,13 +4,17 @@
 package com.azure.spring.aad.webapp;
 
 import com.azure.spring.aad.AADClientRegistrationRepository;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.oauth2.client.OAuth2AuthorizationContext;
 import org.springframework.security.oauth2.client.OAuth2AuthorizedClient;
 import org.springframework.security.oauth2.client.OAuth2AuthorizedClientProvider;
 import org.springframework.security.oauth2.client.RefreshTokenOAuth2AuthorizedClientProvider;
 import org.springframework.security.oauth2.client.web.OAuth2AuthorizedClientRepository;
 import org.springframework.security.oauth2.core.OAuth2AccessToken;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -24,6 +28,8 @@
  */
 public class AADOAuth2AuthorizedClientRepository implements OAuth2AuthorizedClientRepository {
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(AADOAuth2AuthorizedClientRepository.class);
+
     private final AADWebAppClientRegistrationRepository repo;
     private final OAuth2AuthorizedClientRepository delegate;
     private final OAuth2AuthorizedClientProvider provider;
@@ -73,7 +79,15 @@ public <T extends OAuth2AuthorizedClient> T loadAuthorizedClient(String id,
                 .principal(principal)
                 .attributes(getAttributesConsumer(scopes))
                 .build();
-            return (T) provider.authorize(context);
+            OAuth2AuthorizedClient clientGotByRefreshToken = provider.authorize(context);
+            try {
+                ServletRequestAttributes attributes =
+                    (ServletRequestAttributes) RequestContextHolder.currentRequestAttributes();
+                delegate.saveAuthorizedClient(clientGotByRefreshToken, principal, request, attributes.getResponse());
+            } catch (IllegalStateException exception) {
+                LOGGER.warn("Can not save OAuth2AuthorizedClient.", exception);
+            }
+            return (T) clientGotByRefreshToken;
         }
         return null;
     }