Merge pull request #1425 from lprimak/adaptable-ee-ini-env

[#1424] [Enhancement] Made Jakarta EE IniEnvironment more flexible
apache · Apr 19, 2024 · 10bb454 · 10bb454
2 parents f10d6a8 + 1ee2b75
commit 10bb454
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 17 deletions.
diff --git a/support/jakarta-ee/src/main/java/org/apache/shiro/ee/filters/FormResubmitSupport.java b/support/jakarta-ee/src/main/java/org/apache/shiro/ee/filters/FormResubmitSupport.java
@@ -32,8 +32,9 @@
 import java.net.URISyntaxException;
 import java.util.Collections;
 import org.apache.shiro.ee.filters.Forms.FallbackPredicate;
-import org.apache.shiro.ee.filters.ShiroFilter.WrappedSecurityManager;
 import static org.apache.shiro.ee.filters.FormResubmitSupportCookies.transformCookieHeader;
+import static org.apache.shiro.ee.filters.ShiroFilter.isSecurityManagerTypeOf;
+import static org.apache.shiro.ee.filters.ShiroFilter.unwrapSecurityManager;
 import static org.apache.shiro.ee.listeners.EnvironmentLoaderListener.isFormResubmitDisabled;
 import java.io.IOException;
 import java.net.CookieManager;
@@ -137,11 +138,11 @@ static class PartialAjaxResult {
     }
 
     static void savePostDataForResubmit(HttpServletRequest request, HttpServletResponse response, @NonNull String loginUrl) {
-        if (isPostRequest(request) && unwrapSecurityManager(SecurityUtils.getSecurityManager())
-                instanceof DefaultSecurityManager) {
+        if (isPostRequest(request) && isSecurityManagerTypeOf(SecurityUtils.getSecurityManager(),
+                DefaultSecurityManager.class)) {
             String postData = getPostData(request);
             var cacheKey = UUID.randomUUID();
-            var dsm = (DefaultSecurityManager) unwrapSecurityManager(SecurityUtils.getSecurityManager());
+            DefaultSecurityManager dsm = unwrapSecurityManager(SecurityUtils.getSecurityManager());
             if (dsm.getCacheManager() != null) {
                 var cache = dsm.getCacheManager().getCache(FORM_DATA_CACHE);
                 var rememberMeManager = (AbstractRememberMeManager) dsm.getRememberMeManager();
@@ -179,8 +180,8 @@ static String getPostData(ServletRequest request) {
 
     static String getSavedFormDataFromKey(@NonNull String savedFormDataKey) {
         String savedFormData = null;
-        if (unwrapSecurityManager(SecurityUtils.getSecurityManager()) instanceof DefaultSecurityManager) {
-            var dsm = (DefaultSecurityManager) unwrapSecurityManager(SecurityUtils.getSecurityManager());
+        if (isSecurityManagerTypeOf(SecurityUtils.getSecurityManager(), DefaultSecurityManager.class)) {
+            DefaultSecurityManager dsm = unwrapSecurityManager(SecurityUtils.getSecurityManager());
             if (dsm.getCacheManager() != null) {
                 var cache = dsm.getCacheManager().getCache(FORM_DATA_CACHE);
                 var cacheKey = UUID.fromString(savedFormDataKey);
@@ -543,15 +544,6 @@ public static DefaultWebSessionManager getNativeSessionManager(SecurityManager s
         return rv;
     }
 
-    private static org.apache.shiro.mgt.SecurityManager unwrapSecurityManager(SecurityManager securityManager) {
-        if (securityManager instanceof WrappedSecurityManager) {
-            WrappedSecurityManager wsm = (WrappedSecurityManager) securityManager;
-            return wsm.wrapped;
-        } else {
-            return securityManager;
-        }
-    }
-
     private static String getJSFNewViewState(URI savedRequest, HttpClient client, String savedFormData)
             throws IOException, InterruptedException {
         var getRequest = HttpRequest.newBuilder().uri(savedRequest).GET().build();

diff --git a/support/jakarta-ee/src/main/java/org/apache/shiro/ee/filters/ShiroFilter.java b/support/jakarta-ee/src/main/java/org/apache/shiro/ee/filters/ShiroFilter.java
@@ -45,6 +45,7 @@
 import lombok.extern.slf4j.Slf4j;
 import static org.apache.shiro.ee.listeners.EnvironmentLoaderListener.isServletNoPrincipal;
 import org.apache.shiro.mgt.DefaultSecurityManager;
+import org.apache.shiro.mgt.SecurityManager;
 import org.apache.shiro.session.Session;
 import org.apache.shiro.session.SessionException;
 import org.apache.shiro.subject.Subject;
@@ -177,6 +178,36 @@ public Subject createSubject(SubjectContext context) {
         }
     }
 
+    /**
+     * Determines if the specified security manager is of the specified type or a subclass of the specified type.
+     *
+     * @param securityManager
+     * @param type
+     * @return true if the security manager is of the specified type or a subclass of the specified type, false otherwise.
+     */
+    public static boolean isSecurityManagerTypeOf(SecurityManager securityManager,
+                                                  Class<? extends SecurityManager> type) {
+        return type.isAssignableFrom(unwrapSecurityManager(securityManager).getClass());
+    }
+
+    /**
+     * Unwraps the security manager. This method should only be used under limited circumstances,
+     * Because the returned value is no longer wrapped, it will not be able to create subjects correctly.
+     *
+     * @param <SM> the desired type of the security manager
+     * @param securityManager
+     * @return unwrapped security manager
+     */
+    @SuppressWarnings("unchecked")
+    public static <SM extends SecurityManager> SM unwrapSecurityManager(SecurityManager securityManager) {
+        if (securityManager instanceof WrappedSecurityManager) {
+            WrappedSecurityManager wsm = (WrappedSecurityManager) securityManager;
+            return (SM) wsm.wrapped;
+        } else {
+            return (SM) securityManager;
+        }
+    }
+
     @Override
     protected ServletRequest wrapServletRequest(HttpServletRequest request) {
         if (isShiroEEDisabled(request.getServletContext())) {

diff --git a/support/jakarta-ee/src/main/java/org/apache/shiro/ee/listeners/IniEnvironment.java b/support/jakarta-ee/src/main/java/org/apache/shiro/ee/listeners/IniEnvironment.java
@@ -15,11 +15,14 @@
 
 import java.nio.charset.StandardCharsets;
 
+import lombok.RequiredArgsConstructor;
 import org.apache.shiro.ee.filters.FormAuthenticationFilter;
 import org.apache.shiro.ee.filters.LogoutFilter;
 import org.apache.shiro.ee.filters.SslFilter;
 
 import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.servlet.Filter;
 
 import lombok.extern.slf4j.Slf4j;
@@ -44,7 +47,9 @@ public class IniEnvironment extends IniWebEnvironment {
     private String otherConfigLocation;
 
     @SuppressWarnings("deprecation")
+    @RequiredArgsConstructor
     private static final class SecurityManagerFactory extends WebIniSecurityManagerFactory {
+        private final Function<Supplier<byte[]>, SecurityManager> securityManagerSupplier;
         private final Lazy<AesCipherService> cipherService = new Lazy<>(AesCipherService::new);
 
         @Override
@@ -65,7 +70,7 @@ private static final class SecurityManagerFactory extends WebIniSecurityManagerF
 
         @Override
         protected SecurityManager createDefaultInstance() {
-            return new DefaultWebSecurityManager(this::generateCipherKey);
+            return securityManagerSupplier.apply(this::generateCipherKey);
         }
 
         private byte[] generateCipherKey() {
@@ -80,7 +85,11 @@ private byte[] generateCipherKey() {
     }
 
     public IniEnvironment() {
-        var securityManagerFactory = new SecurityManagerFactory();
+        this(DefaultWebSecurityManager::new);
+    }
+
+    public IniEnvironment(Function<Supplier<byte[]>, SecurityManager> securityManagerSupplier) {
+        var securityManagerFactory = new SecurityManagerFactory(securityManagerSupplier);
         securityManagerFactory.getReflectionBuilder().setAlternateObjectSupplier(Beans::getInstance);
         setSecurityManagerFactory(securityManagerFactory);
     }