From c0e0f9bb83221f6e914e86abb5f61c773c00bbc0 Mon Sep 17 00:00:00 2001
From: Joakim Erdfelt <joakim.erdfelt@gmail.com>
Date: Mon, 5 May 2014 11:37:35 -0700
Subject: [PATCH] 431459 - Jetty WebSocket compression extensions fails to
 handle big messages properly

+ Setting up Assume for deregistered compression extension tests
---
 .../eclipse/jetty/websocket/server/ChromeTest.java |  3 +++
 .../server/PerMessageDeflateExtensionTest.java     |  6 ++++--
 .../websocket/server/SimpleServletServer.java      | 14 ++++++++++++++
 .../jetty/websocket/servlet/WebSocketServlet.java  |  2 ++
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ChromeTest.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ChromeTest.java
index 38ab89072c0..48b4664ca05 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ChromeTest.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ChromeTest.java
@@ -30,6 +30,7 @@ import org.eclipse.jetty.websocket.common.test.HttpResponse;
 import org.eclipse.jetty.websocket.server.examples.MyEchoServlet;
 import org.junit.AfterClass;
 import org.junit.Assert;
+import org.junit.Assume;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -53,6 +54,8 @@ public class ChromeTest
     @Test
     public void testUpgradeWithWebkitDeflateExtension() throws Exception
     {
+        Assume.assumeTrue("Server has x-webkit-deflate-frame registered",server.getWebSocketServletFactory().getExtensionFactory().isAvailable("x-webkit-deflate-frame"));
+        
         BlockheadClient client = new BlockheadClient(server.getServerUri());
         try
         {
diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/PerMessageDeflateExtensionTest.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/PerMessageDeflateExtensionTest.java
index 6f2a4fec1d2..9de6b938b91 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/PerMessageDeflateExtensionTest.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/PerMessageDeflateExtensionTest.java
@@ -57,6 +57,9 @@ public class PerMessageDeflateExtensionTest
     @Test
     public void testPerMessageDeflateDefault() throws Exception
     {
+        Assume.assumeTrue("Server has x-webkit-deflate-frame registered",
+                server.getWebSocketServletFactory().getExtensionFactory().isAvailable("permessage-deflate"));
+
         BlockheadClient client = new BlockheadClient(server.getServerUri());
         client.clearExtensions();
         client.addExtensions("permessage-deflate");
@@ -70,8 +73,7 @@ public class PerMessageDeflateExtensionTest
             client.sendStandardRequest();
             HttpResponse resp = client.expectUpgradeResponse();
 
-            // Stop test here if server doesn't have permessage-deflate enabled.
-            Assume.assumeThat("Response",resp.getExtensionsHeader(),containsString("permessage-deflate"));
+            Assert.assertThat("Response",resp.getExtensionsHeader(),containsString("permessage-deflate"));
 
             String msg = "Hello";
 
diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/SimpleServletServer.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/SimpleServletServer.java
index 6d6c730460b..4ec28fb229a 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/SimpleServletServer.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/SimpleServletServer.java
@@ -34,6 +34,7 @@ import org.eclipse.jetty.toolchain.test.MavenTestingUtils;
 import org.eclipse.jetty.util.log.Log;
 import org.eclipse.jetty.util.log.Logger;
 import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
 
 public class SimpleServletServer
 {
@@ -147,4 +148,17 @@ public class SimpleServletServer
             e.printStackTrace(System.err);
         }
     }
+
+    public WebSocketServletFactory getWebSocketServletFactory()
+    {
+        // Try filter approach first
+        WebSocketUpgradeFilter filter = (WebSocketUpgradeFilter)this.servlet.getServletContext().getAttribute(WebSocketUpgradeFilter.class.getName());
+        if (filter != null)
+        {
+            return filter.getFactory();
+        }
+
+        // Try servlet next
+        return (WebSocketServletFactory)this.servlet.getServletContext().getAttribute(WebSocketServletFactory.class.getName());
+    }
 }
diff --git a/jetty-websocket/websocket-servlet/src/main/java/org/eclipse/jetty/websocket/servlet/WebSocketServlet.java b/jetty-websocket/websocket-servlet/src/main/java/org/eclipse/jetty/websocket/servlet/WebSocketServlet.java
index f5a96865f73..85b44f64f7b 100644
--- a/jetty-websocket/websocket-servlet/src/main/java/org/eclipse/jetty/websocket/servlet/WebSocketServlet.java
+++ b/jetty-websocket/websocket-servlet/src/main/java/org/eclipse/jetty/websocket/servlet/WebSocketServlet.java
@@ -130,6 +130,8 @@ public abstract class WebSocketServlet extends HttpServlet
             configure(factory);
 
             factory.init();
+            
+            getServletContext().setAttribute(WebSocketServletFactory.class.getName(),factory);
         }
         catch (Exception x)
         {