WW-5765 Re-store only the token validated in the request in execAndWait

core/src/main/java/org/apache/struts2/interceptor/ExecuteAndWaitInterceptor.java

@@ -281,10 +281,7 @@ protected String doIntercept(ActionInvocation actionInvocation) throws Exception
                 LOG.trace("BackgroundProcess instance is not done (wait processing) - executeAfterValidationPass: {}, secondTime: {}.", executeAfterValidationPass, secondTime);
                 actionInvocation.getStack().push(bp.getAction());
 
-                final String token = TokenHelper.getToken();
-                if (token != null) {
-                    TokenHelper.setSessionToken(TokenHelper.getTokenName(), token);
-                }
+                TokenHelper.restoreValidatedToken();
 
                 Map<String, ResultConfig> results = proxy.getConfig().getResults();
                 if (!results.containsKey(WAIT)) {

core/src/main/java/org/apache/struts2/util/TokenHelper.java

@@ -52,6 +52,7 @@ public class TokenHelper {
     public static final String TOKEN_NAME_FIELD = "struts.token.name";
     private static final Logger LOG = LogManager.getLogger(TokenHelper.class);
     private static final Random RANDOM = new SecureRandom();
+    private static final String VALIDATED_TOKEN_KEY = TokenHelper.class.getName() + ".validatedToken";
 
 
     /**
@@ -196,11 +197,33 @@ public static boolean validToken() {
 
         // remove the token so it won't be used again
         session.remove(tokenSessionName);
+        ActionContext.getContext().put(VALIDATED_TOKEN_KEY, new ValidatedToken(tokenName, token));
 
         return true;
     }
 
+    /**
+     * Puts back into the session the token accepted by {@link #validToken()} earlier in the current request, so the
+     * same token stays valid for a follow-up request. Does nothing if no token was accepted in the current request.
+     */
+    public static void restoreValidatedToken() {
+        ValidatedToken validated = (ValidatedToken) ActionContext.getContext().get(VALIDATED_TOKEN_KEY);
+        if (validated != null) {
+            setSessionToken(validated.name, validated.value);
+        }
+    }
+
     public static String generateGUID() {
         return new BigInteger(165, RANDOM).toString(36).toUpperCase();
     }
+
+    private static final class ValidatedToken {
+        private final String name;
+        private final String value;
+
+        private ValidatedToken(String name, String value) {
+            this.name = name;
+            this.value = value;
+        }
+    }
 }

core/src/test/java/org/apache/struts2/interceptor/ExecuteAndWaitInterceptorTest.java

@@ -36,6 +36,7 @@
 import com.opensymphony.xwork2.mock.MockResult;
 import com.opensymphony.xwork2.ognl.OgnlUtil;
 import com.opensymphony.xwork2.util.location.LocatableProperties;
+import org.apache.struts2.util.TokenHelper;
 import org.apache.struts2.StrutsInternalTestCase;
 import org.apache.struts2.dispatcher.HttpParameters;
 import org.apache.struts2.interceptor.exec.ExecutorProvider;
@@ -64,6 +65,8 @@
     private Map<String, Object> session;
     private ExecuteAndWaitInterceptor waitInterceptor;
     private ParametersInterceptor parametersInterceptor;
+    private TokenInterceptor tokenInterceptor;
+    private TokenSessionStoreInterceptor tokenSessionInterceptor;
 
     public void testOneWait() throws Exception {
         waitInterceptor.setDelay(0);
@@ -219,6 +222,63 @@
         assertEquals("success", result3);
     }
 
+    public void testWaitWithoutTokenCheckLeavesSessionTokensUnchanged() throws Exception {
+        waitInterceptor.setDelay(0);
+        waitInterceptor.setDelaySleepInterval(0);
+        String customTokenSessionName = TokenHelper.buildTokenSessionAttributeName("custom");
+        session.put(customTokenSessionName, "server-generated");
+
+        for (String tokenName : new String[]{TokenHelper.DEFAULT_TOKEN_NAME, "custom"}) {
+            Map<String, String[]> requestParams = new HashMap<>();
+            requestParams.put(TokenHelper.TOKEN_NAME_FIELD, new String[]{tokenName});
+            requestParams.put(TokenHelper.DEFAULT_TOKEN_NAME, new String[]{"request-supplied"});
+            requestParams.put("custom", new String[]{"request-supplied"});
+            withRequestParams(requestParams);
+
+            assertEquals("wait", buildProxy("plainWait").execute());
+
+            assertNull(session.get(TokenHelper.buildTokenSessionAttributeName(TokenHelper.DEFAULT_TOKEN_NAME)));
+            assertEquals("server-generated", session.get(customTokenSessionName));
+        }
+    }
+
+    public void testTokenAndWaitAcceptsWaitPageRefreshWithValidatedToken() throws Exception {
+        assertWaitPageRefreshAccepted("tokenWait");
+    }
+
+    public void testTokenSessionAndWaitAcceptsWaitPageRefreshWithValidatedToken() throws Exception {
+        assertWaitPageRefreshAccepted("tokenSessionWait");
+    }
+
+    private void assertWaitPageRefreshAccepted(String actionName) throws Exception {
+        waitInterceptor.setDelay(0);
+        waitInterceptor.setDelaySleepInterval(0);
+        String tokenSessionName = TokenHelper.buildTokenSessionAttributeName(TokenHelper.DEFAULT_TOKEN_NAME);
+        session.put(tokenSessionName, "server-generated");
+
+        Map<String, String[]> requestParams = new HashMap<>();
+        requestParams.put(TokenHelper.TOKEN_NAME_FIELD, new String[]{TokenHelper.DEFAULT_TOKEN_NAME});
+        requestParams.put(TokenHelper.DEFAULT_TOKEN_NAME, new String[]{"server-generated"});
+        withRequestParams(requestParams);
+
+        assertEquals("wait", buildProxy(actionName).execute());
+        assertEquals("server-generated", session.get(tokenSessionName));
+
+        Thread.sleep(300);
+
+        assertEquals("wait", buildProxy(actionName).execute());
+        assertEquals("server-generated", session.get(tokenSessionName));
+
+        Thread.sleep(500);
+
+        assertEquals("success", buildProxy(actionName).execute());
+        assertNull(session.get(tokenSessionName));
+    }
+
+    private void withRequestParams(Map<String, String[]> requestParams) {
+        context = ActionContext.of(context).withParameters(HttpParameters.create(requestParams).build()).getContextMap();
+    }
+
     protected ActionProxy buildProxy(String actionName) throws Exception {
         return actionProxyFactory.createActionProxy("", actionName, null, context);
     }
@@ -244,9 +304,13 @@
 
         container.inject(waitInterceptor);
         container.inject(parametersInterceptor);
+        container.inject(tokenInterceptor);
+        container.inject(tokenSessionInterceptor);
 
         waitInterceptor.init();
         parametersInterceptor.init();
+        tokenInterceptor.init();
+        tokenSessionInterceptor.init();
     }
 
     protected void tearDown() throws Exception {
@@ -273,17 +337,35 @@
             // interceptors
             waitInterceptor = new ExecuteAndWaitInterceptor();
             parametersInterceptor = new ParametersInterceptor();
+            tokenInterceptor = new TokenInterceptor();
+            tokenSessionInterceptor = new TokenSessionStoreInterceptor();
             PackageConfig wait = new PackageConfig.Builder("")
                 .addActionConfig("action1", new ActionConfig.Builder("", "action1", ExecuteAndWaitDelayAction.class.getName())
                     .addResultConfig(new ResultConfig.Builder(Action.SUCCESS, MockResult.class.getName()).build())
                     .addResultConfig(new ResultConfig.Builder(ExecuteAndWaitInterceptor.WAIT, MockResult.class.getName()).build())
                     .addInterceptor(new InterceptorMapping("params", parametersInterceptor))
                     .addInterceptor(new InterceptorMapping("execAndWait", waitInterceptor))
                     .build())
+                .addActionConfig("plainWait", waitActionConfig("plainWait"))
+                .addActionConfig("tokenWait", waitActionConfig("tokenWait", new InterceptorMapping("token", tokenInterceptor)))
+                .addActionConfig("tokenSessionWait", waitActionConfig("tokenSessionWait", new InterceptorMapping("tokenSession", tokenSessionInterceptor)))
                 .build();
             configuration.addPackageConfig("", wait);
         }
 
+        private ActionConfig waitActionConfig(String name, InterceptorMapping... tokenInterceptors) {
+            ActionConfig.Builder builder = new ActionConfig.Builder("", name, ExecuteAndWaitDelayAction.class.getName())
+                .addResultConfig(new ResultConfig.Builder(Action.SUCCESS, MockResult.class.getName()).build())
+                .addResultConfig(new ResultConfig.Builder(ExecuteAndWaitInterceptor.WAIT, MockResult.class.getName()).build());
+            for (InterceptorMapping tokenInterceptor : tokenInterceptors) {
+                builder.addInterceptor(tokenInterceptor);
+            }
+            return builder
+                .addInterceptor(new InterceptorMapping("params", parametersInterceptor))
+                .addInterceptor(new InterceptorMapping("execAndWait", waitInterceptor))
+                .build();
+        }
+
         public void register(ContainerBuilder builder, LocatableProperties props) throws ConfigurationException {
             builder.factory(ObjectFactory.class);
             builder.factory(ActionProxyFactory.class, DefaultActionProxyFactory.class);

core/src/test/java/org/apache/struts2/util/TokenHelperTest.java

@@ -76,6 +76,50 @@ public void testValidToken() {
         assertTrue(TokenHelper.validToken());
     }
 
+    public void testRestoreValidatedTokenPutsBackTheValidatedToken() {
+        String tokenName = "restoreTokenTest";
+        String token = TokenHelper.setToken(tokenName);
+        final String sessionTokenName = TokenHelper.buildTokenSessionAttributeName(tokenName);
+        withTokenParams(tokenName, token);
+
+        assertTrue(TokenHelper.validToken());
+        assertNull(session.get(sessionTokenName));
+
+        TokenHelper.restoreValidatedToken();
+
+        assertEquals(token, session.get(sessionTokenName));
+    }
+
+    public void testRestoreValidatedTokenAfterRejectedTokenLeavesSessionUnchanged() {
+        String tokenName = "rejectedTokenTest";
+        String token = TokenHelper.setToken(tokenName);
+        final String sessionTokenName = TokenHelper.buildTokenSessionAttributeName(tokenName);
+        Map<String, String[]> params = new HashMap<>();
+        params.put(TokenHelper.TOKEN_NAME_FIELD, new String[]{tokenName});
+        ActionContext.getContext().withParameters(HttpParameters.create(params).build());
+
+        assertFalse(TokenHelper.validToken());
+        TokenHelper.restoreValidatedToken();
+
+        assertEquals(token, session.get(sessionTokenName));
+    }
+
+    public void testRestoreValidatedTokenWithoutValidationLeavesSessionUnchanged() {
+        String tokenName = "unvalidatedTokenTest";
+        withTokenParams(tokenName, "request-supplied");
+
+        TokenHelper.restoreValidatedToken();
+
+        assertTrue(session.isEmpty());
+    }
+
+    private void withTokenParams(String tokenName, String token) {
+        Map<String, String[]> params = new HashMap<>();
+        params.put(TokenHelper.TOKEN_NAME_FIELD, new String[]{tokenName});
+        params.put(tokenName, new String[]{token});
+        ActionContext.getContext().withParameters(HttpParameters.create(params).build());
+    }
+
     public void testGetTokenDoesNotNpe() {
         String token = TokenHelper.getToken(null);
         assertNull(token);