Fixed problem with detecting session ID if response contains multiple Set-Cookie headers. This patch works with both Tomcat6 and 7.

Andrew McNair · Andrew McNair · commit 095b6fe7c17d · 2013-02-15T12:28:33.000-08:00
diff --git a/core/src/main/java/de/javakaffee/web/msm/RequestTrackingHostValve.java b/core/src/main/java/de/javakaffee/web/msm/RequestTrackingHostValve.java
@@ -17,6 +17,8 @@
 package de.javakaffee.web.msm;
 
 import java.io.IOException;
+import java.lang.reflect.Method;
+import java.util.Collection;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
 
@@ -63,6 +65,18 @@ public class RequestTrackingHostValve extends ValveBase {
 
     private static final String MSM_REQUEST_ID = "msm.requestId";
 
+	private static final boolean IS_TOMCAT_6;
+	static {
+		Method getHeaderValues = null;
+		try {
+			getHeaderValues = Response.class.getMethod("getHeaderValues", String.class);
+		}catch(NoSuchMethodException e) {
+			//Do nothing
+		}
+
+		IS_TOMCAT_6 = (getHeaderValues != null);
+	}
+
     /**
      * Creates a new instance with the given ignore pattern and
      * {@link SessionBackupService}.
@@ -239,24 +253,45 @@ private String getSessionId(final Request request, final Response response) {
         return sessionId != null ? sessionId : request.getRequestedSessionId();
     }
 
-    private String getSessionIdFromResponseSessionCookie( final Response response ) {
-        final String header = response.getHeader( "Set-Cookie" );
-        if ( header != null && header.contains( _sessionCookieName ) ) {
-            final String sessionIdPrefix = _sessionCookieName + "=";
-            final int idxNameStart = header.indexOf( sessionIdPrefix );
-            final int idxValueStart = idxNameStart + sessionIdPrefix.length();
-            int idxValueEnd = header.indexOf( ';', idxNameStart );
-            if ( idxValueEnd == -1 ) {
-                idxValueEnd = header.indexOf( ' ', idxValueStart );
-            }
-            if ( idxValueEnd == -1 ) {
-                idxValueEnd = header.length();
+    private String getSessionIdFromResponseSessionCookie(final Response response) {
+        String[] headers = getResponseSetCookieHeaders(response);
+        if (headers == null) {
+            return null;
+        }
+        for (String header : headers) {
+            if (header != null && header.contains(_sessionCookieName)) {
+                final String sessionIdPrefix = _sessionCookieName + "=";
+                final int idxNameStart = header.indexOf(sessionIdPrefix);
+                final int idxValueStart = idxNameStart + sessionIdPrefix.length();
+                int idxValueEnd = header.indexOf(';', idxNameStart);
+                if (idxValueEnd == -1) {
+                    idxValueEnd = header.indexOf(' ', idxValueStart);
+                }
+                if (idxValueEnd == -1) {
+                    idxValueEnd = header.length();
+                }
+                return header.substring(idxValueStart, idxValueEnd);
             }
-            return header.substring( idxValueStart, idxValueEnd );
         }
         return null;
     }
 
+    private String[] getResponseSetCookieHeaders(final Response response) {
+        try {
+            if (IS_TOMCAT_6) {
+                Method getHeaderValues = response.getClass().getMethod("getHeaderValues", String.class);
+                String[] result = (String[]) getHeaderValues.invoke(response, "Set-Cookie");
+                return result;
+            } else {
+                Method getHeaders = response.getClass().getMethod("getHeaders", String.class);
+                Collection<String> result = (Collection<String>) getHeaders.invoke(response, "Set-Cookie");
+                return result.toArray(new String[0]);
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private void logDebugResponseCookie( final Response response ) {
         final String header = response.getHeader("Set-Cookie");
         if ( header != null && header.contains( _sessionCookieName ) ) {
diff --git a/core/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveTest.java b/core/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveTest.java
@@ -43,7 +43,7 @@
  * @author <a href="mailto:martin.grotzke@javakaffee.de">Martin Grotzke</a>
  * @version $Id$
  */
-public class RequestTrackingHostValveTest {
+public abstract class RequestTrackingHostValveTest {
 
     protected MemcachedSessionService _service;
     private RequestTrackingHostValve _sessionTrackerValve;
@@ -109,23 +109,13 @@ public final void testBackupSessionNotInvokedWhenNoSessionIdPresent() throws IOE
     public final void testBackupSessionInvokedWhenResponseCookiePresent() throws IOException, ServletException {
         when( _request.getRequestedSessionId() ).thenReturn( null );
         final Cookie cookie = new Cookie( _sessionTrackerValve.getSessionCookieName(), "foo" );
-        when( _response.getHeader( eq( "Set-Cookie" ) ) ).thenReturn( generateCookieString( cookie ) );
+        setupGetResponseSetCookieHeadersExpectations(_response, new String[]{generateCookieString( cookie )});
         _sessionTrackerValve.invoke( _request, _response );
 
         verify( _service ).backupSession( eq( "foo" ), eq( false), anyString() );
 
     }
 
-    private String generateCookieString(final Cookie cookie) {
-        final StringBuffer sb = new StringBuffer();
-        ServerCookie.appendCookieValue
-        (sb, cookie.getVersion(), cookie.getName(), cookie.getValue(),
-             cookie.getPath(), cookie.getDomain(), cookie.getComment(),
-             cookie.getMaxAge(), cookie.getSecure(), true );
-        final String setSessionCookieHeader = sb.toString();
-        return setSessionCookieHeader;
-    }
-
     @Test
     public final void testChangeSessionIdForRelocatedSession() throws IOException, ServletException {
 
@@ -136,7 +126,7 @@ public final void testChangeSessionIdForRelocatedSession() throws IOException, S
         when( _request.getRequestedSessionId() ).thenReturn( sessionId );
 
         final Cookie cookie = new Cookie( _sessionTrackerValve.getSessionCookieName(), newSessionId );
-        when( _response.getHeader( eq( "Set-Cookie" ) ) ).thenReturn( generateCookieString( cookie ) );
+        setupGetResponseSetCookieHeadersExpectations(_response, new String[]{generateCookieString( cookie )});
 
         _sessionTrackerValve.invoke( _request, _response );
 
@@ -154,4 +144,16 @@ public final void testRequestFinishedShouldBeInvokedForIgnoredResources() throws
         verify( _service ).requestFinished( eq( "foo" ), anyString() );
     }
 
+    protected abstract void setupGetResponseSetCookieHeadersExpectations(Response response, String[] result);
+
+    private String generateCookieString(final Cookie cookie) {
+        final StringBuffer sb = new StringBuffer();
+        ServerCookie.appendCookieValue
+                (sb, cookie.getVersion(), cookie.getName(), cookie.getValue(),
+                        cookie.getPath(), cookie.getDomain(), cookie.getComment(),
+                        cookie.getMaxAge(), cookie.getSecure(), true);
+        final String setSessionCookieHeader = sb.toString();
+        return setSessionCookieHeader;
+    }
+
 }
diff --git a/tomcat6/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveT6Test.java b/tomcat6/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveT6Test.java
@@ -0,0 +1,16 @@
+package de.javakaffee.web.msm;
+
+import org.apache.catalina.connector.Response;
+import org.testng.annotations.Test;
+
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.when;
+
+@Test
+public class RequestTrackingHostValveT6Test extends RequestTrackingHostValveTest {
+
+    @Override
+    protected void setupGetResponseSetCookieHeadersExpectations(Response response, String[] result) {
+        when(response.getHeaderValues(eq("Set-Cookie"))).thenReturn(result);
+    }
+}
diff --git a/tomcat7/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveT7Test.java b/tomcat7/src/test/java/de/javakaffee/web/msm/RequestTrackingHostValveT7Test.java
@@ -0,0 +1,18 @@
+package de.javakaffee.web.msm;
+
+import org.apache.catalina.connector.Response;
+import org.testng.annotations.Test;
+
+import java.util.Arrays;
+
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.when;
+
+@Test
+public class RequestTrackingHostValveT7Test extends RequestTrackingHostValveTest {
+
+    @Override
+    protected void setupGetResponseSetCookieHeadersExpectations(Response response, String[] result) {
+        when(response.getHeaders(eq("Set-Cookie"))).thenReturn(Arrays.asList(result));
+    }
+}
