WFLY-18650 - Security roles lost following failover

Author: pedro-hos

auth/realm/base/src/main/java/org/wildfly/security/auth/realm/JaasSecurityRealm.java

@@ -168,10 +168,16 @@ public Principal getRealmIdentityPrincipal() {
             return principal;
         }
 
+        @Override
         public Subject getSubject() {
             return subject;
         }
 
+        @Override
+        public void setSubject(Subject subject) {
+            this.subject = subject;
+        }
+
         @Override
         public SupportLevel getCredentialAcquireSupport(final Class<? extends Credential> credentialType, final String algorithmName, final AlgorithmParameterSpec parameterSpec) throws RealmUnavailableException {
             return JaasSecurityRealm.this.getCredentialAcquireSupport(credentialType, algorithmName, parameterSpec);

auth/server/base/src/main/java/org/wildfly/security/auth/callback/CachedIdentityAuthorizeCallback.java

@@ -18,18 +18,22 @@
 
 package org.wildfly.security.auth.callback;
 
+import static org.wildfly.common.Assert.checkNotNullParam;
+
+import java.security.Principal;
+import java.util.Set;
+import java.util.function.Function;
+
+import javax.security.auth.Subject;
+
 import org.wildfly.common.Assert;
 import org.wildfly.security.auth.principal.NamePrincipal;
+import org.wildfly.security.auth.server.RealmIdentity;
 import org.wildfly.security.auth.server.SecurityDomain;
 import org.wildfly.security.auth.server.SecurityIdentity;
 import org.wildfly.security.cache.CachedIdentity;
 import org.wildfly.security.cache.IdentityCache;
 
-import java.security.Principal;
-import java.util.function.Function;
-
-import static org.wildfly.common.Assert.checkNotNullParam;
-
 /**
  * <p>A callback that is capable of perform authorization based on the identities managed by an {@link IdentityCache}.
  *
@@ -139,6 +143,25 @@ public CachedIdentityAuthorizeCallback(Principal principal, Function<SecurityDom
         this.localCache = localCache;
     }
 
+    /**
+     * Set the Roles present on {@link CachedIdentity} into the {@link RealmIdentity#setSubject(Subject)} in order to get authenticate on all HA nodes;
+     * @param realmIdentity
+     */
+    public void setSubject(RealmIdentity realmIdentity) {
+        checkNotNullParam("realmIdentity", realmIdentity);
+        Subject subject = realmIdentity.getSubject();
+        if (subject == null) {
+            CachedIdentity cachedIdentity = createDomainCache().get();
+            if (cachedIdentity != null) {
+                subject = new Subject();
+                Set<Principal> principals = subject.getPrincipals();
+                principals.add(realmIdentity.getRealmIdentityPrincipal());
+                cachedIdentity.getRoles().forEach(role -> principals.add(new Roles(role)));
+                realmIdentity.setSubject(subject);
+            }
+        }
+    }
+
     /**
      * Indicates if a cached identity was successfully authorized.
      *
@@ -229,4 +252,19 @@ public boolean needsInformation() {
     private IdentityCache createDomainCache() {
         return this.identityCache.apply(securityDomain);
     }
+
+    private static class Roles implements Principal {
+
+        private final String name;
+
+        Roles(final String name) {
+            this.name = name;
+        }
+
+        @Override
+        public String getName() {
+            return this.name;
+        }
+    }
+
 }

auth/server/base/src/main/java/org/wildfly/security/auth/server/RealmIdentity.java

@@ -24,6 +24,8 @@
 import java.security.spec.AlgorithmParameterSpec;
 import java.util.function.Function;
 
+import javax.security.auth.Subject;
+
 import org.wildfly.common.Assert;
 import org.wildfly.security.auth.SupportLevel;
 import org.wildfly.security.auth.principal.AnonymousPrincipal;
@@ -55,6 +57,12 @@ public interface RealmIdentity {
      */
     Principal getRealmIdentityPrincipal();
 
+    default Subject getSubject() {
+        return null;
+    }
+
+    default void setSubject(Subject subject) {}
+
     /**
      * @deprecated Transition method; remove before GA.
      */

auth/server/base/src/main/java/org/wildfly/security/auth/server/ServerAuthenticationContext.java

@@ -61,10 +61,10 @@
 import org.wildfly.security.auth.callback.EvidenceVerifyCallback;
 import org.wildfly.security.auth.callback.ExclusiveNameCallback;
 import org.wildfly.security.auth.callback.FastUnsupportedCallbackException;
-import org.wildfly.security.auth.callback.PrincipalAuthorizeCallback;
-import org.wildfly.security.auth.callback.MechanismInformationCallback;
 import org.wildfly.security.auth.callback.IdentityCredentialCallback;
+import org.wildfly.security.auth.callback.MechanismInformationCallback;
 import org.wildfly.security.auth.callback.PeerPrincipalCallback;
+import org.wildfly.security.auth.callback.PrincipalAuthorizeCallback;
 import org.wildfly.security.auth.callback.RequestInformationCallback;
 import org.wildfly.security.auth.callback.SSLCallback;
 import org.wildfly.security.auth.callback.SecurityIdentityCallback;
@@ -74,6 +74,7 @@
 import org.wildfly.security.auth.permission.RunAsPrincipalPermission;
 import org.wildfly.security.auth.principal.AnonymousPrincipal;
 import org.wildfly.security.auth.principal.NamePrincipal;
+import org.wildfly.security.auth.server._private.ElytronMessages;
 import org.wildfly.security.auth.server.event.RealmFailedAuthenticationEvent;
 import org.wildfly.security.auth.server.event.RealmIdentityFailedAuthorizationEvent;
 import org.wildfly.security.auth.server.event.RealmIdentitySuccessfulAuthorizationEvent;
@@ -98,7 +99,6 @@
 import org.wildfly.security.password.spec.ClearPasswordSpec;
 import org.wildfly.security.ssl.SSLConnection;
 import org.wildfly.security.x500.X500;
-import org.wildfly.security.auth.server._private.ElytronMessages;
 
 /**
  * Server-side authentication context.  Instances of this class are used to perform all authentication and re-authorization
@@ -1133,6 +1133,7 @@ private void handleOne(final Callback[] callbacks, final int idx) throws IOExcep
                         }
                         if (principal != null) {
                             setAuthenticationPrincipal(principal);
+                            authorizeCallback.setSubject(stateRef.get().getRealmIdentity());
                             if (authorize()) {
                                 authorizedIdentity = getAuthorizedIdentity();
                             }