From 79b3032527950d29947da41bb8cf445fb44fa5ca Mon Sep 17 00:00:00 2001
From: drfish <drfish.me@gmail.com>
Date: Fri, 11 Oct 2019 11:20:26 +0800
Subject: [PATCH] Async granting group permission

---
 .../jenkins/azuread/AzureSecurityRealm.java   | 53 +++++++++++++------
 1 file changed, 38 insertions(+), 15 deletions(-)

diff --git a/src/main/java/com/microsoft/jenkins/azuread/AzureSecurityRealm.java b/src/main/java/com/microsoft/jenkins/azuread/AzureSecurityRealm.java
index 68239060..f103fbb9 100644
--- a/src/main/java/com/microsoft/jenkins/azuread/AzureSecurityRealm.java
+++ b/src/main/java/com/microsoft/jenkins/azuread/AzureSecurityRealm.java
@@ -76,6 +76,10 @@ public class AzureSecurityRealm extends SecurityRealm {
     private static final int NONCE_LENGTH = 10;
     public static final String CALLBACK_URL = "/securityRealm/finishLogin";
 
+    static {
+        SecurityContextHolder.setStrategyName(SecurityContextHolder.MODE_INHERITABLETHREADLOCAL);
+    }
+
     private Secret clientId;
     private Secret clientSecret;
     private Secret tenant;
@@ -202,20 +206,28 @@ public HttpResponse doFinishLogin(StaplerRequest request) throws InvalidJwtExcep
             }
             // validate the nonce to avoid CSRF
             final AzureAdUser userDetails = validateAndParseIdToken(expectedNonce, idToken);
-            final AzureAuthenticationToken auth = new AzureAuthenticationToken(userDetails);
-
-            // Enforce updating current identity
-            SecurityContextHolder.getContext().setAuthentication(auth);
-            User u = User.current();
-            if (u != null) {
-                String description = generateDescription(auth);
-                u.setDescription(description);
-                u.setFullName(auth.getAzureAdUser().getName());
-            }
-            SecurityListener.fireAuthenticated(userDetails);
+
+            refreshAuthentication(userDetails);
+
             AzureAdPlugin.sendLoginEvent(
                     AppInsightsUtils.hash(userDetails.getObjectID()),
                     AppInsightsUtils.hash(this.getTenant()));
+
+            Runnable runnable = () -> {
+                final Collection<ActiveDirectoryGroup> groups = AzureCachePool.get(getAzureClient())
+                        .getBelongingGroupsByOid(userDetails.getObjectID());
+                if (groups != null) {
+                    userDetails.setAuthorities(groups);
+                    refreshAuthentication(userDetails);
+                }
+            };
+
+            Thread thread = new Thread(runnable);
+            thread.setUncaughtExceptionHandler((th, ex) -> {
+                LOGGER.severe(String.format("Fail to grant group permission: %s", ex.toString()));
+            });
+            thread.start();
+
         } catch (Exception ex) {
             AzureAdPlugin.sendLoginFailEvent(this.getTenant(), ex.getMessage());
             throw ex;
@@ -234,17 +246,28 @@ public HttpResponse doFinishLogin(StaplerRequest request) throws InvalidJwtExcep
         }
     }
 
+    void refreshAuthentication(AzureAdUser userDetails) {
+        final AzureAuthenticationToken auth = new AzureAuthenticationToken(userDetails);
+
+        // Enforce updating current identity
+        SecurityContextHolder.getContext().setAuthentication(auth);
+        User u = User.current();
+        if (u != null) {
+            String description = generateDescription(auth);
+            u.setDescription(description);
+            u.setFullName(auth.getAzureAdUser().getName());
+        }
+        SecurityListener.fireAuthenticated(userDetails);
+    }
+
     AzureAdUser validateAndParseIdToken(String expectedNonce, String idToken)
-        throws InvalidJwtException, MalformedClaimException {
+            throws InvalidJwtException, MalformedClaimException {
         JwtClaims claims = getJwtConsumer().processToClaims(idToken);
         final String responseNonce = (String) claims.getClaimValue("nonce");
         if (StringUtils.isAnyEmpty(expectedNonce, responseNonce) || !expectedNonce.equals(responseNonce)) {
             throw new IllegalStateException("Invalid nonce in the response");
         }
         final AzureAdUser userDetails = AzureAdUser.createFromJwt(claims);
-        final Collection<ActiveDirectoryGroup> groups = AzureCachePool.get(getAzureClient())
-                .getBelongingGroupsByOid(userDetails.getObjectID());
-        userDetails.setAuthorities(groups);
         return userDetails;
     }