Fix for #515 [websocket] Make sure FakeSession is shared amongst message

Atmosphere · Jul 16, 2012 · aad412f · aad412f
1 parent ecd7259
commit aad412f
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git a/integration-tests/cpr/src/test/java/org/atmosphere/tests/websocket/BaseTest.java b/integration-tests/cpr/src/test/java/org/atmosphere/tests/websocket/BaseTest.java
@@ -31,6 +31,7 @@
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
+import javax.servlet.http.HttpSession;
 import java.io.IOException;
 import java.net.ServerSocket;
 import java.util.concurrent.CountDownLatch;
@@ -133,6 +134,87 @@ public void onError(Throwable t) {
         c.close();
     }
 
+    @Test(timeOut = 60000, enabled = true)
+    public void sessionTest() {
+        logger.info("{}: running test: sessionTest", getClass().getSimpleName());
+
+        final AtomicReference<HttpSession> session = new AtomicReference<HttpSession>();
+        final AtomicReference<HttpSession> sessionOnPost = new AtomicReference<HttpSession>();
+
+        final CountDownLatch latch = new CountDownLatch(1);
+        atmoServlet.framework().addAtmosphereHandler(ROOT, new AbstractHttpAtmosphereHandler() {
+
+            boolean isSuspended = false;
+
+            public void onRequest(AtmosphereResource r) throws IOException {
+                if (!isSuspended) {
+                    session.set(r.suspend().session(true));
+                    isSuspended = true;
+                } else {
+                    String message = r.getRequest().getReader().readLine();
+                    r.getBroadcaster().broadcast(message);
+                }
+            }
+
+            public void onStateChange(AtmosphereResourceEvent event) throws IOException {
+                if (event.isSuspended()) {
+                    sessionOnPost.set(event.getResource().session());
+                    event.write(event.getMessage().toString().getBytes());
+                }
+            }
+        }, BroadcasterFactory.getDefault().get(DefaultBroadcaster.class, "suspend"));
+
+        AsyncHttpClient c = new AsyncHttpClient();
+        try {
+            final AtomicReference<String> response = new AtomicReference<String>(null);
+            WebSocket w = c.prepareGet(urlTarget).execute(new WebSocketUpgradeHandler.Builder().build())
+                    .get();
+            w.addWebSocketListener(new WebSocketTextListener() {
+
+                @Override
+                public void onMessage(String message) {
+                    response.set(message);
+                    latch.countDown();
+                }
+
+                @Override
+                public void onFragment(String fragment, boolean last) {
+                    fail();
+                }
+
+                @Override
+                public void onOpen(WebSocket websocket) {
+                }
+
+                @Override
+                public void onClose(WebSocket websocket) {
+                    latch.countDown();
+                }
+
+                @Override
+                public void onError(Throwable t) {
+                    t.printStackTrace();
+                    fail(t.getMessage());
+                }
+            }).sendTextMessage("echo");
+
+            try {
+                latch.await(20, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                fail(e.getMessage());
+            }
+
+            assertNotNull(response.get());
+            assertEquals(response.get(), "echo");
+            assertEquals(session.get().getId(), sessionOnPost.get().getId());
+
+        } catch (Exception e) {
+            logger.error("test failed", e);
+            fail(e.getMessage());
+        }
+        c.close();
+    }
+
     @Test(timeOut = 60000, enabled = true)
     public void echoTest() {
         logger.info("{}: running test: echoTest", getClass().getSimpleName());

diff --git a/modules/cpr/src/main/java/org/atmosphere/websocket/protocol/SimpleHttpProtocol.java b/modules/cpr/src/main/java/org/atmosphere/websocket/protocol/SimpleHttpProtocol.java
@@ -116,6 +116,7 @@ public List<AtmosphereRequest> onMessage(WebSocket webSocket, String d) {
                 .pathInfo(pathInfo)
                 .destroyable(destroyable)
                 .headers(resource.getRequest().headersMap())
+                .session(resource.session())
                 .build());
 
         return list;