Issue #4349 - Add tests for client negotiation of internal extensions

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/core/server/internal/RFC6455Handshaker.java

@@ -45,7 +45,6 @@
 import org.eclipse.jetty.websocket.core.Behavior;
 import org.eclipse.jetty.websocket.core.ExtensionConfig;
 import org.eclipse.jetty.websocket.core.FrameHandler;
-import org.eclipse.jetty.websocket.core.WebSocketComponents;
 import org.eclipse.jetty.websocket.core.WebSocketConstants;
 import org.eclipse.jetty.websocket.core.WebSocketException;
 import org.eclipse.jetty.websocket.core.internal.ExtensionStack;
@@ -100,7 +99,7 @@ public boolean upgradeRequest(WebSocketNegotiator negotiator, HttpServletRequest
             baseRequest,
             request,
             response,
-            new WebSocketComponents());
+            negotiator.getWebSocketComponents());
         if (LOG.isDebugEnabled())
             LOG.debug("negotiation {}", negotiation);
 

jetty-websocket/websocket-core/src/main/resources/META-INF/services/org.eclipse.jetty.websocket.core.Extension

@@ -1,4 +1,5 @@
 org.eclipse.jetty.websocket.core.internal.IdentityExtension
 org.eclipse.jetty.websocket.core.internal.FragmentExtension
+org.eclipse.jetty.websocket.core.internal.PerMessageDeflateExtension
 org.eclipse.jetty.websocket.core.internal.ValidationExtension
-org.eclipse.jetty.websocket.core.internal.PerMessageDeflateExtension
+org.eclipse.jetty.websocket.core.internal.FrameCaptureExtension

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/core/WebSocketNegotiationTest.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
 
 import org.eclipse.jetty.client.HttpRequest;
 import org.eclipse.jetty.client.HttpResponse;
@@ -39,11 +40,15 @@
 import org.eclipse.jetty.websocket.core.client.ClientUpgradeRequest;
 import org.eclipse.jetty.websocket.core.client.UpgradeListener;
 import org.eclipse.jetty.websocket.core.client.WebSocketCoreClient;
+import org.eclipse.jetty.websocket.core.internal.WebSocketCoreSession;
 import org.eclipse.jetty.websocket.core.server.Negotiation;
 import org.eclipse.jetty.websocket.core.server.WebSocketNegotiator;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
@@ -69,11 +74,12 @@ public void onFrame(Frame frame)
 
     private WebSocketServer server;
     private WebSocketCoreClient client;
+    private WebSocketComponents components = new WebSocketComponents();
 
     @BeforeEach
     public void startup() throws Exception
     {
-        WebSocketNegotiator negotiator = new WebSocketNegotiator.AbstractNegotiator()
+        WebSocketNegotiator negotiator = new WebSocketNegotiator.AbstractNegotiator(components, null)
         {
             @Override
             public FrameHandler negotiate(Negotiation negotiation) throws IOException
@@ -120,7 +126,7 @@ public FrameHandler negotiate(Negotiation negotiation) throws IOException
         };
 
         server = new WebSocketServer(negotiator);
-        client = new WebSocketCoreClient();
+        client = new WebSocketCoreClient(null, components);
 
         server.start();
         client.start();
@@ -353,4 +359,80 @@ public void testInvalidUpgradeRequestNoKey() throws Exception
 
         assertThat(response, containsString("400 Bad Request"));
     }
+
+    public static Stream<Arguments> internalExtensionScenarios() throws Exception
+    {
+        return Stream.of(
+            Arguments.of("@int1, @int1", "@int1, @int1"),
+            Arguments.of("@int1, ext1, ext2", "@int1, ext1, ext2"),
+            Arguments.of("ext1, @int1, ext2", "ext1, @int1, ext2"),
+            Arguments.of("ext1, ext2, @int1", "ext1, ext2, @int1"),
+            Arguments.of("@int1, ext1, @int2, ext2, @int3", "@int1, ext1, @int2, ext2, @int3"),
+            Arguments.of("ext1, ext1, ext1, @int1, ext2", "ext1, @int1, ext2"),
+            Arguments.of("ext1, @int1, ext1, ext1, ext2", "ext1, @int1, ext2")
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("internalExtensionScenarios")
+    public void testClientRequestedInternalExtensions(String reqExts, String negExts) throws Exception
+    {
+        // Add some simple Extensions for to make test examples clearer.
+        WebSocketExtensionRegistry extRegistry = components.getExtensionRegistry();
+        extRegistry.register("ext1", EmptyExtension.class);
+        extRegistry.register("ext2", EmptyExtension.class);
+        extRegistry.register("ext3", EmptyExtension.class);
+        extRegistry.register("@int1", EmptyExtension.class);
+        extRegistry.register("@int2", EmptyExtension.class);
+        extRegistry.register("@int3", EmptyExtension.class);
+
+        TestFrameHandler clientHandler = new TestFrameHandler();
+        CompletableFuture<String> extensionHeader = new CompletableFuture<>();
+        ClientUpgradeRequest upgradeRequest = ClientUpgradeRequest.from(client, server.getUri(), clientHandler);
+        upgradeRequest.setSubProtocols("test");
+        upgradeRequest.addExtensions(reqExts.split(","));
+        upgradeRequest.addListener(new UpgradeListener()
+        {
+            @Override
+            public void onHandshakeResponse(HttpRequest request, HttpResponse response)
+            {
+                extensionHeader.complete(response.getHeaders().get(HttpHeader.SEC_WEBSOCKET_EXTENSIONS));
+            }
+        });
+
+        // Connect to the client then close the Session.
+        client.connect(upgradeRequest).get(5, TimeUnit.SECONDS);
+        clientHandler.sendClose();
+        assertTrue(clientHandler.closed.await(5, TimeUnit.SECONDS));
+        assertNull(clientHandler.getError());
+
+        // We had no internal Extensions in the response headers.
+        assertThat(extensionHeader.get(5, TimeUnit.SECONDS), not(containsString("@")));
+
+        // The list of Extensions on the client contains the internal Extensions.
+        StringBuilder negotiatedExtensions = new StringBuilder();
+        List<Extension> extensions = ((WebSocketCoreSession)clientHandler.coreSession).getExtensionStack().getExtensions();
+        for (int i = 0; i < extensions.size(); i++)
+        {
+            negotiatedExtensions.append(extensions.get(i).getConfig().getParameterizedName());
+            if (i != extensions.size() - 1)
+                negotiatedExtensions.append(", ");
+        }
+        assertThat(negotiatedExtensions.toString(), is(negExts));
+    }
+
+    public static class EmptyExtension extends AbstractExtension
+    {
+        @Override
+        public void onFrame(Frame frame, Callback callback)
+        {
+            nextIncomingFrame(frame, callback);
+        }
+
+        @Override
+        public void sendFrame(Frame frame, Callback callback, boolean batch)
+        {
+            nextOutgoingFrame(frame, callback, batch);
+        }
+    }
 }