refactor: get UserAuthentication directly from WithMockUserSecurityContextFactory

Author: nigel.zheng

src/main/java/com/github/ahunigel/test/security/WithClaims.java

@@ -1,21 +1,22 @@
 package com.github.ahunigel.test.security;
 
-import org.springframework.beans.BeanUtils;
-import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.security.authentication.AbstractAuthenticationToken;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.context.SecurityContext;
-import org.springframework.security.test.context.support.WithMockUser;
+import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.security.test.context.support.WithSecurityContext;
 import org.springframework.security.test.context.support.WithSecurityContextFactory;
 
 import java.lang.annotation.*;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.stream.Collectors;
 
 /**
  * Created by Nigel Zheng on 8/3/2018.
+ * <p>
+ * Attach claims as map to current authentication details
  */
 @Target({ElementType.METHOD, ElementType.TYPE})
 @Retention(RetentionPolicy.RUNTIME)
@@ -26,36 +27,41 @@
   /**
    * Return the contained {@link Claim} annotations.
    *
-   * @return the claim
+   * @return the claims
    */
   Claim[] value();
 
-  WithMockUser user() default @WithMockUser();
+  /**
+   * key-value paired string array, redundant value will be ignored
+   * <p>
+   * would merge with #value() if key is absent
+   *
+   * @return
+   */
+  String[] claims() default {};
 
   class WithClaimsSecurityContextFactory implements WithSecurityContextFactory<WithClaims> {
     @Override
     public SecurityContext createSecurityContext(WithClaims annotation) {
-      WithSecurityContext withSecurityContext = AnnotationUtils
-          .findAnnotation(annotation.user().getClass(), WithSecurityContext.class);
-      WithSecurityContextFactory factory = createFactory(withSecurityContext);
-      SecurityContext context = factory.createSecurityContext(annotation.user());
+      SecurityContext context = SecurityContextHolder.getContext();
       Authentication authentication = context.getAuthentication();
-      Map<String, String> claims = Arrays.stream(annotation.value()).collect(Collectors.toMap(Claim::name, Claim::value));
-      ((AbstractAuthenticationToken) authentication).setDetails(claims);
+      if (authentication != null && authentication instanceof AbstractAuthenticationToken) {
+        Map<String, String> claims = Arrays.stream(annotation.value()).collect(Collectors.toMap(Claim::name, Claim::value));
+        Map<String, String> stringMap = toMap(annotation.claims());
+        stringMap.entrySet().stream().forEach(entry -> claims.putIfAbsent(entry.getKey(), entry.getValue()));
+        if (!claims.isEmpty()) {
+          ((AbstractAuthenticationToken) authentication).setDetails(claims);
+        }
+      }
       return context;
     }
 
-    private WithSecurityContextFactory<? extends Annotation> createFactory(
-        WithSecurityContext withSecurityContext) {
-      Class<? extends WithSecurityContextFactory<? extends Annotation>> clazz = withSecurityContext
-          .factory();
-      try {
-        return BeanUtils.instantiateClass(clazz);
-      } catch (IllegalStateException e) {
-        return BeanUtils.instantiateClass(clazz);
-      } catch (Exception e) {
-        throw new RuntimeException(e);
+    private Map<String, String> toMap(String[] claims) {
+      final Map<String, String> map = new HashMap<>();
+      for (int i = 0; i + 1 < claims.length; i += 2) {
+        map.put(claims[i], claims[i + 1]);
       }
+      return map;
     }
   }
 

src/main/java/com/github/ahunigel/test/security/WithMockUserAndClaims.java

@@ -0,0 +1,56 @@
+package com.github.ahunigel.test.security;
+
+import com.github.ahunigel.test.security.util.MockUserUtils;
+import org.springframework.core.annotation.AliasFor;
+import org.springframework.security.authentication.AbstractAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.test.context.support.WithMockUser;
+import org.springframework.security.test.context.support.WithSecurityContext;
+import org.springframework.security.test.context.support.WithSecurityContextFactory;
+
+import java.lang.annotation.*;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * Created by Nigel Zheng on 8/3/2018.
+ * <p>
+ * Emulate running with a mocked user,
+ * attach claims as map to mocked authentication details
+ */
+@Target({ElementType.METHOD, ElementType.TYPE})
+@Retention(RetentionPolicy.RUNTIME)
+@Inherited
+@Documented
+@WithSecurityContext(factory = WithMockUserAndClaims.WithMockUserWithClaimsSecurityContextFactory.class)
+public @interface WithMockUserAndClaims {
+  /**
+   * Return the contained {@link Claim} annotations.
+   *
+   * @return the claims
+   */
+  @AliasFor("claims")
+  Claim[] value();
+
+  @AliasFor("value")
+  Claim[] claims();
+
+  WithMockUser user() default @WithMockUser();
+
+  class WithMockUserWithClaimsSecurityContextFactory implements WithSecurityContextFactory<WithMockUserAndClaims> {
+    @Override
+    public SecurityContext createSecurityContext(WithMockUserAndClaims annotation) {
+      SecurityContext context = MockUserUtils.getSecurityContext(annotation.user());
+      Authentication authentication = context.getAuthentication();
+      Map<String, String> claims = Arrays.stream(annotation.value()).collect(Collectors.toMap(Claim::name, Claim::value));
+      if (!claims.isEmpty() && authentication instanceof AbstractAuthenticationToken) {
+        ((AbstractAuthenticationToken) authentication).setDetails(claims);
+      }
+      return context;
+    }
+
+  }
+
+}

src/main/java/com/github/ahunigel/test/security/oauth2/WithMockOAuth2Client.java

@@ -17,6 +17,8 @@
 
 /**
  * Created by Nigel Zheng on 8/3/2018.
+ * <p>
+ * Emulate running with a mocked oauth2 client
  */
 @Retention(RetentionPolicy.RUNTIME)
 @WithSecurityContext(factory = WithMockOAuth2Client.WithMockOAuth2ClientSecurityContextFactory.class)

src/main/java/com/github/ahunigel/test/security/oauth2/WithMockOAuth2User.java

@@ -1,12 +1,9 @@
 package com.github.ahunigel.test.security.oauth2;
 
 import com.github.ahunigel.test.security.Claim;
-import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
-import org.springframework.security.core.GrantedAuthority;
-import org.springframework.security.core.authority.SimpleGrantedAuthority;
+import com.github.ahunigel.test.security.util.MockUserUtils;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
-import org.springframework.security.core.userdetails.User;
 import org.springframework.security.oauth2.provider.OAuth2Authentication;
 import org.springframework.security.test.context.support.WithMockUser;
 import org.springframework.security.test.context.support.WithSecurityContext;
@@ -15,14 +12,14 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.Map;
-import java.util.Set;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 /**
  * Created by Nigel Zheng on 8/3/2018.
+ * <p>
+ * Emulate running with a mocked oauth2 client on behalf of user,
+ * attach claims as map to mocked oauth2 authentication details
  */
 @Retention(RetentionPolicy.RUNTIME)
 @WithSecurityContext(factory = WithMockOAuth2User.WithMockOAuth2UserSecurityContextFactory.class)
@@ -31,40 +28,21 @@
 
   WithMockUser user() default @WithMockUser();
 
+  /**
+   * Return the contained {@link Claim} annotations.
+   *
+   * @return the claims
+   */
   Claim[] claims();
 
   class WithMockOAuth2UserSecurityContextFactory implements WithSecurityContextFactory<WithMockOAuth2User> {
 
-    /**
-     * Sadly, #WithMockUserSecurityContextFactory is not public,
-     * so re-implement mock user authentication creation
-     *
-     * @param user
-     * @return an Authentication with provided user details
-     */
-    public static UsernamePasswordAuthenticationToken getUserAuthentication(final WithMockUser user) {
-      final String principal = user.username().isEmpty() ? user.value() : user.username();
-
-      final Stream<String> grants = user.authorities().length == 0 ?
-          Stream.of(user.roles()).map(r -> "ROLE_" + r) :
-          Stream.of(user.authorities());
-
-      final Set<? extends GrantedAuthority> userAuthorities = grants
-          .map(auth -> new SimpleGrantedAuthority(auth))
-          .collect(Collectors.toSet());
-
-      return new UsernamePasswordAuthenticationToken(
-          new User(principal, user.password(), userAuthorities),
-          principal + ":" + user.password(),
-          userAuthorities);
-    }
-
     @Override
     public SecurityContext createSecurityContext(final WithMockOAuth2User annotation) {
       final SecurityContext ctx = SecurityContextHolder.createEmptyContext();
       OAuth2Authentication oAuth2Authentication = new OAuth2Authentication(
           WithMockOAuth2Client.WithMockOAuth2ClientSecurityContextFactory.getOAuth2Request(annotation.client()),
-          getUserAuthentication(annotation.user()));
+          MockUserUtils.getAuthentication(annotation.user()));
       Map<String, String> claims = Arrays.stream(annotation.claims()).collect(Collectors.toMap(Claim::name, Claim::value));
       if (!claims.isEmpty()) {
         oAuth2Authentication.setDetails(claims);
@@ -74,12 +52,5 @@ public SecurityContext createSecurityContext(final WithMockOAuth2User annotation
       return ctx;
     }
 
-    private Map<String, Object> toMap(String[] claims) {
-      final Map<String, Object> map = new HashMap<>();
-      for (int i = 0; i + 1 < claims.length; i += 2) {
-        map.put(claims[i], claims[i + 1]);
-      }
-      return map;
-    }
   }
 }

src/main/java/com/github/ahunigel/test/security/util/MockUserUtils.java

@@ -0,0 +1,68 @@
+package com.github.ahunigel.test.security.util;
+
+import org.springframework.beans.BeanUtils;
+import org.springframework.core.annotation.AnnotationUtils;
+import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.GrantedAuthority;
+import org.springframework.security.core.authority.SimpleGrantedAuthority;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.userdetails.User;
+import org.springframework.security.test.context.support.WithMockUser;
+import org.springframework.security.test.context.support.WithSecurityContext;
+import org.springframework.security.test.context.support.WithSecurityContextFactory;
+
+import java.lang.annotation.Annotation;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Created by Nigel Zheng on 8/3/2018.
+ */
+public class MockUserUtils {
+  public static SecurityContext getSecurityContext(WithMockUser withMockUser) {
+    WithSecurityContext withSecurityContext = AnnotationUtils
+        .findAnnotation(withMockUser.getClass(), WithSecurityContext.class);
+    WithSecurityContextFactory factory = createFactory(withSecurityContext);
+    return factory.createSecurityContext(withMockUser);
+  }
+
+  public static Authentication getAuthentication(WithMockUser withMockUser) {
+    return getSecurityContext(withMockUser).getAuthentication();
+  }
+
+  private static WithSecurityContextFactory<? extends Annotation> createFactory(
+      WithSecurityContext withSecurityContext) {
+    Class<? extends WithSecurityContextFactory<? extends Annotation>> clazz = withSecurityContext
+        .factory();
+    try {
+      return BeanUtils.instantiateClass(clazz);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /**
+   * User authentication creation
+   *
+   * @param user
+   * @return an Authentication with provided user details
+   */
+  public static UsernamePasswordAuthenticationToken createUserAuthentication(final WithMockUser user) {
+    final String principal = user.username().isEmpty() ? user.value() : user.username();
+
+    final Stream<String> grants = user.authorities().length == 0 ?
+        Stream.of(user.roles()).map(r -> "ROLE_" + r) :
+        Stream.of(user.authorities());
+
+    final Set<? extends GrantedAuthority> userAuthorities = grants
+        .map(auth -> new SimpleGrantedAuthority(auth))
+        .collect(Collectors.toSet());
+
+    return new UsernamePasswordAuthenticationToken(
+        new User(principal, user.password(), userAuthorities),
+        principal + ":" + user.password(),
+        userAuthorities);
+  }
+}