Setup the ServletRequestContext as part of the thread setup actions

undertow-io · Mar 27, 2015 · dcf6a7b · dcf6a7b
1 parent 4fa596b
commit dcf6a7b
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 35 deletions.
diff --git a/servlet/src/main/java/io/undertow/servlet/core/DeploymentManagerImpl.java b/servlet/src/main/java/io/undertow/servlet/core/DeploymentManagerImpl.java
@@ -158,6 +158,7 @@ public void deploy() {
         }
 
         final List<ThreadSetupAction> setup = new ArrayList<>();
+        setup.add(ServletRequestContextThreadSetupAction.INSTANCE);
         setup.add(new ContextClassLoaderSetupAction(deploymentInfo.getClassLoader()));
         setup.addAll(deploymentInfo.getThreadSetupActions());
         final CompositeThreadSetupAction threadSetupAction = new CompositeThreadSetupAction(setup);

diff --git a/servlet/src/main/java/io/undertow/servlet/core/SecurityActions.java b/servlet/src/main/java/io/undertow/servlet/core/SecurityActions.java
@@ -125,6 +125,45 @@ public ServletRequestContext run() {
         }
     }
 
+    static void setCurrentRequestContext(final ServletRequestContext servletRequestContext) {
+        if (System.getSecurityManager() == null) {
+            ServletRequestContext.setCurrentRequestContext(servletRequestContext);
+        } else {
+            AccessController.doPrivileged(new PrivilegedAction<Object>() {
+                @Override
+                public Object run() {
+                    ServletRequestContext.setCurrentRequestContext(servletRequestContext);
+                    return null;
+                }
+            });
+        }
+    }
+
+    static void clearCurrentServletAttachments() {
+        if (System.getSecurityManager() == null) {
+            ServletRequestContext.clearCurrentServletAttachments();
+        } else {
+            AccessController.doPrivileged(new PrivilegedAction<Object>() {
+                @Override
+                public Object run() {
+                    ServletRequestContext.clearCurrentServletAttachments();
+                    return null;
+                }
+            });
+        }
+    }
+    static ServletRequestContext requireCurrentServletRequestContext() {
+        if (System.getSecurityManager() == null) {
+            return ServletRequestContext.requireCurrent();
+        } else {
+            return AccessController.doPrivileged(new PrivilegedAction<ServletRequestContext>() {
+                @Override
+                public ServletRequestContext run() {
+                    return ServletRequestContext.requireCurrent();
+                }
+            });
+        }
+    }
     static ServletInitialHandler createServletInitialHandler(final ServletPathMatches paths, final HttpHandler next, final CompositeThreadSetupAction setupAction, final ServletContextImpl servletContext) {
         if (System.getSecurityManager() == null) {
             return new ServletInitialHandler(paths, next, setupAction, servletContext);

diff --git a/servlet/src/main/java/io/undertow/servlet/core/ServletRequestContextThreadSetupAction.java b/servlet/src/main/java/io/undertow/servlet/core/ServletRequestContextThreadSetupAction.java
@@ -0,0 +1,52 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2014 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package io.undertow.servlet.core;
+
+import io.undertow.server.HttpServerExchange;
+import io.undertow.servlet.api.ThreadSetupAction;
+import io.undertow.servlet.handlers.ServletRequestContext;
+
+/**
+ * @author Stuart Douglas
+ */
+class ServletRequestContextThreadSetupAction implements ThreadSetupAction {
+
+    static final ServletRequestContextThreadSetupAction INSTANCE = new ServletRequestContextThreadSetupAction();
+
+    private ServletRequestContextThreadSetupAction() {
+
+    }
+
+    private static final Handle HANDLE = new Handle() {
+        @Override
+        public void tearDown() {
+            SecurityActions.clearCurrentServletAttachments();
+        }
+    };
+
+    @Override
+    public Handle setup(HttpServerExchange exchange) {
+        if(exchange == null) {
+            return null;
+        }
+        ServletRequestContext servletRequestContext = exchange.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
+        SecurityActions.setCurrentRequestContext(servletRequestContext);
+        return HANDLE;
+    }
+}
diff --git a/servlet/src/main/java/io/undertow/servlet/handlers/SecurityActions.java b/servlet/src/main/java/io/undertow/servlet/handlers/SecurityActions.java
@@ -39,20 +39,6 @@ public HttpSessionImpl run() {
         }
     }
 
-    static void setCurrentRequestContext(final ServletRequestContext servletRequestContext) {
-        if (System.getSecurityManager() == null) {
-            ServletRequestContext.setCurrentRequestContext(servletRequestContext);
-        } else {
-            AccessController.doPrivileged(new PrivilegedAction<Object>() {
-                @Override
-                public Object run() {
-                    ServletRequestContext.setCurrentRequestContext(servletRequestContext);
-                    return null;
-                }
-            });
-        }
-    }
-
     static ServletRequestContext requireCurrentServletRequestContext() {
         if (System.getSecurityManager() == null) {
             return ServletRequestContext.requireCurrent();
@@ -65,18 +51,4 @@ public ServletRequestContext run() {
             });
         }
     }
-
-    static void clearCurrentServletAttachments() {
-            if (System.getSecurityManager() == null) {
-                ServletRequestContext.clearCurrentServletAttachments();
-            } else {
-                AccessController.doPrivileged(new PrivilegedAction<Object>() {
-                    @Override
-                    public Object run() {
-                        ServletRequestContext.clearCurrentServletAttachments();
-                        return null;
-                    }
-                });
-            }
-        }
 }
diff --git a/servlet/src/main/java/io/undertow/servlet/handlers/ServletInitialHandler.java b/servlet/src/main/java/io/undertow/servlet/handlers/ServletInitialHandler.java
@@ -270,8 +270,6 @@ public void handleFirstRequest(final HttpServerExchange exchange, final ServletC
                     request.setAttribute(entry.getKey(), entry.getValue());
                 }
             }
-
-            SecurityActions.setCurrentRequestContext(servletRequestContext);
             servletRequestContext.setRunningInsideHandler(true);
             try {
                 listeners.requestInitialized(request);
@@ -331,11 +329,7 @@ public void handleFirstRequest(final HttpServerExchange exchange, final ServletC
                 }
             }
         } finally {
-            try {
-                handle.tearDown();
-            } finally {
-                SecurityActions.clearCurrentServletAttachments();
-            }
+            handle.tearDown();
         }
     }