[UNDERTOW-2418] Adjust properly session timeout also in case when FORM is combined with other mechanisms

Author: aogburn

core/src/main/java/io/undertow/security/impl/FormAuthenticationMechanism.java

@@ -172,10 +172,7 @@ public AuthenticationMechanismOutcome runFormAuth(final HttpServerExchange excha
     protected void handleRedirectBack(final HttpServerExchange exchange) {
         final Session session = Sessions.getSession(exchange);
         if (session != null) {
-            final Integer originalSessionTimeout = (Integer) session.removeAttribute(ORIGINAL_SESSION_TIMEOUT);
-            if (originalSessionTimeout != null) {
-                session.setMaxInactiveInterval(originalSessionTimeout);
-            }
+            restoreOriginalSessionTimeout(session);
             final String location = (String) session.removeAttribute(LOCATION_ATTRIBUTE);
             if(location != null) {
                 exchange.addDefaultResponseListener(new DefaultResponseListener() {
@@ -192,6 +189,20 @@ public boolean handleDefaultResponse(final HttpServerExchange exchange) {
         }
     }
 
+    protected void restoreOriginalSessionTimeout(final HttpServerExchange exchange) {
+        final Session session = Sessions.getSession(exchange);
+        restoreOriginalSessionTimeout(session);
+    }
+
+    protected void restoreOriginalSessionTimeout(final Session session) {
+        if (session != null) {
+            final Integer originalSessionTimeout = (Integer) session.removeAttribute(ORIGINAL_SESSION_TIMEOUT);
+            if (originalSessionTimeout != null) {
+                session.setMaxInactiveInterval(originalSessionTimeout);
+            }
+        }
+    }
+
     public ChallengeResult sendChallenge(final HttpServerExchange exchange, final SecurityContext securityContext) {
 
         // make sure a request to root context is handled with trailing slash. Otherwise call to j_security_check will not

core/src/main/java/io/undertow/security/impl/SecurityContextImpl.java

@@ -17,13 +17,17 @@
  */
 package io.undertow.security.impl;
 
+import static io.undertow.security.api.SecurityNotification.EventType.AUTHENTICATED;
+
 import io.undertow.UndertowLogger;
 import io.undertow.UndertowMessages;
 import io.undertow.security.api.AuthenticationMechanism;
 import io.undertow.security.api.AuthenticationMechanism.AuthenticationMechanismOutcome;
 import io.undertow.security.api.AuthenticationMechanism.ChallengeResult;
 import io.undertow.security.api.AuthenticationMechanismContext;
 import io.undertow.security.api.AuthenticationMode;
+import io.undertow.security.api.NotificationReceiver;
+import io.undertow.security.api.SecurityNotification;
 import io.undertow.security.idm.Account;
 import io.undertow.security.idm.IdentityManager;
 import io.undertow.security.idm.PasswordCredential;
@@ -168,6 +172,16 @@ public void addAuthenticationMechanism(final AuthenticationMechanism handler) {
             }
             cur.next = new Node<>(handler);
         }
+        if (handler instanceof FormAuthenticationMechanism) {
+            registerNotificationReceiver(new NotificationReceiver() {
+                @Override
+                public void handleNotification(final SecurityNotification notification) {
+                    if (notification.getEventType() == AUTHENTICATED) {
+                        ((FormAuthenticationMechanism) handler).restoreOriginalSessionTimeout(exchange);
+                    }
+                }
+            });
+        }
     }
 
     @Override

servlet/src/main/java/io/undertow/servlet/handlers/security/ServletFormAuthenticationMechanism.java

@@ -18,14 +18,10 @@
 
 package io.undertow.servlet.handlers.security;
 
-import static io.undertow.security.api.SecurityNotification.EventType.AUTHENTICATED;
 import static io.undertow.util.StatusCodes.OK;
 
 import io.undertow.security.api.AuthenticationMechanism;
 import io.undertow.security.api.AuthenticationMechanismFactory;
-import io.undertow.security.api.SecurityContext;
-import io.undertow.security.api.NotificationReceiver;
-import io.undertow.security.api.SecurityNotification;
 import io.undertow.security.idm.IdentityManager;
 import io.undertow.security.impl.FormAuthenticationMechanism;
 import io.undertow.server.HttpServerExchange;
@@ -157,19 +153,6 @@ public ServletFormAuthenticationMechanism(FormParserFactory formParserFactory, S
         this.overrideInitial = overrideInitial;
     }
 
-    @Override
-    public AuthenticationMechanismOutcome authenticate(final HttpServerExchange exchange, final SecurityContext securityContext) {
-        securityContext.registerNotificationReceiver(new NotificationReceiver() {
-            @Override
-            public void handleNotification(final SecurityNotification notification) {
-                if (notification.getEventType() == AUTHENTICATED) {
-                    getAndInitializeSession(exchange, false);
-                }
-            }
-        });
-        return super.authenticate(exchange, securityContext);
-    }
-
     @Override
     protected Integer servePage(final HttpServerExchange exchange, final String location) {
         final ServletRequestContext servletRequestContext = exchange.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
@@ -271,15 +254,17 @@ private Session getAndInitializeSession(final HttpServerExchange exchange, final
                 session.setMaxInactiveInterval(authenticationSessionTimeout);
             }
         } else {
-            final Integer originalSessionTimeout = (Integer) session.removeAttribute(ORIGINAL_SESSION_TIMEOUT);
-            if (originalSessionTimeout != null) {
-                session.setMaxInactiveInterval(originalSessionTimeout);
-            }
+            restoreOriginalSessionTimeout(session);
         }
 
         return session;
     }
 
+    @Override
+    protected void restoreOriginalSessionTimeout(final HttpServerExchange exchange) {
+        getAndInitializeSession(exchange, false);
+    }
+
     private static class FormResponseWrapper extends HttpServletResponseWrapper {
 
         private int status = OK;