From 41d9245d9a9f12eaddab699495a363deed3f1b3c Mon Sep 17 00:00:00 2001
From: Joakim Erdfelt <joakim.erdfelt@gmail.com>
Date: Thu, 15 Jun 2017 13:35:40 -0700
Subject: [PATCH] Testing / UntrustedWSServer improvements

---
 .../websocket/tests/UntrustedWSServer.java    | 58 ++++++++++++-------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/jetty-websocket/websocket-tests/src/main/java/org/eclipse/jetty/websocket/tests/UntrustedWSServer.java b/jetty-websocket/websocket-tests/src/main/java/org/eclipse/jetty/websocket/tests/UntrustedWSServer.java
index e321d55dc45..8d33b0be96f 100644
--- a/jetty-websocket/websocket-tests/src/main/java/org/eclipse/jetty/websocket/tests/UntrustedWSServer.java
+++ b/jetty-websocket/websocket-tests/src/main/java/org/eclipse/jetty/websocket/tests/UntrustedWSServer.java
@@ -23,6 +23,7 @@ import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -55,10 +56,11 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
     private URI wsUri;
     private boolean ssl = false;
     private SslContextFactory sslContextFactory;
-    
+    private Consumer<Server> serverConsumer;
+
     private Map<URI, CompletableFuture<UntrustedWSSession>> onOpenFutures = new ConcurrentHashMap<>();
     private final ServletContextHandler context = new ServletContextHandler();
-    
+
     @Override
     protected void doStart() throws Exception
     {
@@ -66,7 +68,7 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
         String name = "qtp-untrustedWSServer-" + hashCode();
         threadPool.setName(name);
         threadPool.setDaemon(true);
-        
+
         // Configure Server
         server = new Server(threadPool);
         if (ssl)
@@ -80,7 +82,7 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
             http_config.setResponseHeaderSize(8192);
             http_config.setSendServerVersion(true);
             http_config.setSendDateHeader(false);
-            
+
             sslContextFactory = new SslContextFactory();
             sslContextFactory.setKeyStorePath(MavenTestingUtils.getTestResourceFile("keystore").getAbsolutePath());
             sslContextFactory.setKeyStorePassword("storepwd");
@@ -88,11 +90,11 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
             sslContextFactory.setExcludeCipherSuites("SSL_RSA_WITH_DES_CBC_SHA", "SSL_DHE_RSA_WITH_DES_CBC_SHA", "SSL_DHE_DSS_WITH_DES_CBC_SHA",
                     "SSL_RSA_EXPORT_WITH_RC4_40_MD5", "SSL_RSA_EXPORT_WITH_DES40_CBC_SHA", "SSL_DHE_RSA_EXPORT_WITH_DES40_CBC_SHA",
                     "SSL_DHE_DSS_EXPORT_WITH_DES40_CBC_SHA");
-            
+
             // SSL HTTP Configuration
             HttpConfiguration https_config = new HttpConfiguration(http_config);
             https_config.addCustomizer(new SecureRequestCustomizer());
-            
+
             // SSL Connector
             connector = new ServerConnector(server, new SslConnectionFactory(sslContextFactory, HttpVersion.HTTP_1_1.asString()), new HttpConnectionFactory(https_config));
             connector.setPort(0);
@@ -104,58 +106,74 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
             connector.setPort(0);
         }
         server.addConnector(connector);
-        
+
         context.setContextPath("/");
         server.setHandler(context);
-        
+
         // Serve untrusted endpoint
         context.addServlet(UntrustedWSServlet.class, "/untrusted/*").setInitOrder(1);
-        
+
+        // Allow for server customization
+        if (serverConsumer != null)
+        {
+            serverConsumer.accept(server);
+        }
+
         // Start Server
         addBean(server);
-        
+
         super.doStart();
-        
+
         // Wireup Context related things
         UntrustedWSSessionFactory sessionFactory = (UntrustedWSSessionFactory) context.getServletContext().getAttribute(UntrustedWSSessionFactory.class.getName());
         sessionFactory.addListener(this);
-        
+
         // Establish the Server URI
         URI serverUri = server.getURI();
         wsUri = WSURI.toWebsocket(serverUri).resolve("/");
-        
+
         // Some debugging
         if (LOG.isDebugEnabled())
         {
             LOG.debug("WebSocket Server URI: " + wsUri.toASCIIString());
             server.dump();
         }
-        
+
         super.doStart();
     }
-    
+
+    public void setServerCustomizer(Consumer<Server> customizer)
+    {
+        this.serverConsumer = customizer;
+    }
+
+    public void join() throws InterruptedException
+    {
+        server.join();
+    }
+
     public URI getWsUri()
     {
         return wsUri;
     }
-    
+
     public URI getUntrustedWsUri(Class<?> clazz, TestName testname)
     {
         return wsUri.resolve("/untrusted/" + clazz.getSimpleName() + "/" + testname.getMethodName());
     }
-    
+
     public void registerHttpService(String urlPattern, BiConsumer<HttpServletRequest, HttpServletResponse> serviceConsumer)
     {
         ServletHolder holder = new ServletHolder(new BiConsumerServiceServlet(serviceConsumer));
         context.addServlet(holder, urlPattern);
     }
-    
+
     public void registerWebSocket(String urlPattern, WebSocketCreator creator)
     {
         ServletHolder holder = new ServletHolder(new UntrustedWSServlet(creator));
         context.addServlet(holder, urlPattern);
     }
-    
+
     @Override
     public void onSessionCreate(UntrustedWSSession session, URI requestURI)
     {
@@ -167,7 +185,7 @@ public class UntrustedWSServer extends ContainerLifeCycle implements UntrustedWS
             this.onOpenFutures.put(requestURI, sessionFuture);
         }
     }
-    
+
     public void registerOnOpenFuture(URI uri, CompletableFuture<UntrustedWSSession> sessionFuture)
     {
         this.onOpenFutures.put(uri, sessionFuture);