ELY-2921: Add principal tranformer support to TokenSecurityRealm

auth/realm/token/src/main/java/org/wildfly/security/auth/realm/token/TokenSecurityRealm.java

@@ -47,7 +47,8 @@ public final class TokenSecurityRealm implements SecurityRealm {
     private final String principalClaimName;
     /** A function that maps the set of token claims to a Principal. */
     private final Function<Attributes, Principal> claimToPrincipal;
-
+    /** A function that transofrms Principal  **/
+    private final Function<Principal, Principal> principalTransformer;
     /**
      * Returns a {@link Builder} instance that can be used to configure and create a {@link TokenSecurityRealm}.
      *
@@ -72,6 +73,8 @@ public static Builder builder() {
             this.claimToPrincipal = configuration.claimToPrincipal;
         }
 
+        this.principalTransformer = configuration.principalTransformer;
+
         this.strategy = Assert.checkNotNullParam("tokenValidationStrategy", configuration.strategy);
     }
 
@@ -134,12 +137,17 @@ final class TokenRealmIdentity implements RealmIdentity {
                 this.evidence = null;
             }
         }
+
         @Override
         public Principal getRealmIdentityPrincipal() {
             Principal principal = null;
             try {
                 if (exists()) {
                     principal = claimToPrincipal.apply(this.claims);
+
+                    if (principalTransformer != null) {
+                        principal = principalTransformer.apply(principal);
+                    }
                 }
             } catch (Exception e) {
                 throw ElytronMessages.log.tokenRealmFailedToObtainPrincipal(e);
@@ -223,7 +231,7 @@ public static class Builder {
         private String principalClaimName = "username";
         private Function<Attributes, Principal> claimToPrincipal;
         private TokenValidator strategy;
-
+        private Function<Principal, Principal> principalTransformer;
         /**
          * Construct a new instance.
          */
@@ -262,6 +270,15 @@ public Builder validator(TokenValidator strategy) {
             return this;
         }
 
+        /**
+         * Defines a {@link TokenValidator} that will be used to validate tokens.
+         *
+         * @return this instance
+         */
+        public Builder principalTransformer(Function<Principal, Principal> func) {
+            this.principalTransformer = func;
+            return this;
+        }
         /**
          * Creates a {@link TokenSecurityRealm} instance with all the given configuration.
          *

auth/realm/token/src/test/java/org/wildfly/security/auth/realm/token/OAuth2TokenSecurityRealmTest.java

@@ -24,6 +24,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.wildfly.security.auth.realm.token.validator.OAuth2IntrospectValidator;
+import org.wildfly.security.auth.server.NameRewriter;
 import org.wildfly.security.auth.server.RealmIdentity;
 import org.wildfly.security.auth.server.RealmUnavailableException;
 import org.wildfly.security.authz.Attributes;
@@ -42,6 +43,7 @@
 import java.nio.charset.StandardCharsets;
 import java.security.Principal;
 import java.util.Arrays;
+import java.util.Locale;
 import java.util.function.Function;
 
 import static org.junit.Assert.*;
@@ -103,6 +105,60 @@ public void testUserDefinedPrincipalClaimName() throws Exception {
         assertEquals("elytron@jboss.org", realmIdentityPrincipal.getName());
     }
 
+    @Test
+    public void testPrincipalTransformer() throws Exception {
+        configureReplayTokenIntrospection();
+
+        Function<Principal, Principal> principalTransformer = new CaseRewriter().asPrincipalRewriter();
+        TokenSecurityRealm securityRealm = TokenSecurityRealm.builder()
+                .validator(OAuth2IntrospectValidator.builder()
+                        .clientId("wildfly-elytron")
+                        .clientSecret("dont_tell_me")
+                        .tokenIntrospectionUrl(new URL("http://as.test.org/oauth2/token/introspect")).build())
+                        .principalTransformer(principalTransformer)
+                .build();
+
+        JsonObjectBuilder tokenBuilder = Json.createObjectBuilder();
+
+        tokenBuilder.add("active", true);
+        tokenBuilder.add("username", "elytron@jboss.org");
+
+        RealmIdentity realmIdentity = securityRealm.getRealmIdentity(new BearerTokenEvidence(tokenBuilder.build().toString()));
+
+        assertTrue(realmIdentity.exists());
+
+        Principal realmIdentityPrincipal = realmIdentity.getRealmIdentityPrincipal();
+
+        assertEquals("ELYTRON@JBOSS.ORG", realmIdentityPrincipal.getName());
+    }
+
+    @Test
+    public void testNullPrincipalTransformer() throws Exception {
+        configureReplayTokenIntrospection();
+
+        Function<Principal, Principal> principalTransformer = null;
+        TokenSecurityRealm securityRealm = TokenSecurityRealm.builder()
+                .validator(OAuth2IntrospectValidator.builder()
+                        .clientId("wildfly-elytron")
+                        .clientSecret("dont_tell_me")
+                        .tokenIntrospectionUrl(new URL("http://as.test.org/oauth2/token/introspect")).build())
+                        .principalTransformer(principalTransformer)
+                .build();
+
+        JsonObjectBuilder tokenBuilder = Json.createObjectBuilder();
+
+        tokenBuilder.add("active", true);
+        tokenBuilder.add("username", "elytron@jboss.org");
+
+        RealmIdentity realmIdentity = securityRealm.getRealmIdentity(new BearerTokenEvidence(tokenBuilder.build().toString()));
+
+        assertTrue(realmIdentity.exists());
+
+        Principal realmIdentityPrincipal = realmIdentity.getRealmIdentityPrincipal();
+
+        assertEquals("elytron@jboss.org", realmIdentityPrincipal.getName());
+    }
+
     @Test
     public void testNotActiveToken() throws Exception {
         configureReplayTokenIntrospection();
@@ -218,4 +274,13 @@ public JsonObject introspectAccessToken(URL tokenIntrospectionUrl, String client
             }
         };
     }
+
+    /*
+     * Function to convert string to all caps
+     */
+    private class CaseRewriter implements NameRewriter {
+        public String rewriteName(String original) {
+            return (original == null) ? null : original.toUpperCase(Locale.ROOT);
+        }
+    }
 }