[grid] Modifying ws upgrade to handle VNC client.

Author: diemol

java/server/src/org/openqa/selenium/grid/router/httpd/RouterServer.java

@@ -20,6 +20,7 @@
 import com.google.auto.service.AutoService;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+
 import org.openqa.selenium.BuildInfo;
 import org.openqa.selenium.UsernameAndPassword;
 import org.openqa.selenium.cli.CliCommand;
@@ -48,7 +49,6 @@
 import org.openqa.selenium.grid.web.GridUiRoute;
 import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.remote.http.HttpClient;
-import org.openqa.selenium.remote.http.HttpHandler;
 import org.openqa.selenium.remote.http.HttpResponse;
 import org.openqa.selenium.remote.http.Routable;
 import org.openqa.selenium.remote.http.Route;

java/server/src/org/openqa/selenium/netty/server/WebSocketUpgradeHandler.java

@@ -17,6 +17,9 @@
 
 package org.openqa.selenium.netty.server;
 
+import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.remote.http.Message;
+
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelFuture;
@@ -38,9 +41,6 @@
 import io.netty.handler.codec.http.websocketx.WebSocketServerHandshakerFactory;
 import io.netty.util.AttributeKey;
 
-import org.openqa.selenium.internal.Require;
-import org.openqa.selenium.remote.http.Message;
-
 import java.util.Optional;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -69,6 +69,27 @@ public WebSocketUpgradeHandler(
     this.factory = Require.nonNull("Factory", factory);
   }
 
+  private static void sendHttpResponse(
+    ChannelHandlerContext ctx, HttpRequest req, FullHttpResponse res) {
+    // Generate an error page if response status code is not OK (200).
+    if (res.status().code() != 200) {
+      ByteBuf buf = Unpooled.copiedBuffer(res.status().toString(), UTF_8);
+      res.content().writeBytes(buf);
+      buf.release();
+      setContentLength(res, res.content().readableBytes());
+    }
+
+    // Send the response and close the connection if necessary.
+    ChannelFuture f = ctx.channel().writeAndFlush(res);
+    if (!isKeepAlive(req) || res.status().code() != 200) {
+      f.addListener(ChannelFutureListener.CLOSE);
+    }
+  }
+
+  private static String getWebSocketLocation(HttpRequest req) {
+    return "ws://" + req.headers().get(HttpHeaderNames.HOST);
+  }
+
   @Override
   public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
     if (msg instanceof HttpRequest) {
@@ -100,26 +121,25 @@ private void handleHttpRequest(ChannelHandlerContext ctx, HttpRequest req) {
     }
 
     // Only handle the initial HTTP upgrade request
-    if (!(req.headers().contains("Connection", "upgrade", true) &&
-        req.headers().contains("Sec-WebSocket-Version"))) {
+    if (!(req.headers().containsValue("Connection", "upgrade", true) &&
+          req.headers().contains("Sec-WebSocket-Version"))) {
       ctx.fireChannelRead(req);
       return;
     }
 
     // Is this something we should try and handle?
     Optional<Consumer<Message>> maybeHandler = factory.apply(
       req.uri(),
-      msg -> {
-        ctx.channel().writeAndFlush(Require.nonNull("Message to send", msg));
-      });
+      msg -> ctx.channel().writeAndFlush(Require.nonNull("Message to send", msg)));
     if (!maybeHandler.isPresent()) {
-      sendHttpResponse(ctx, req, new DefaultFullHttpResponse(HTTP_1_1, BAD_REQUEST, ctx.alloc().buffer(0)));
+      sendHttpResponse(ctx, req,
+                       new DefaultFullHttpResponse(HTTP_1_1, BAD_REQUEST, ctx.alloc().buffer(0)));
       return;
     }
 
     // Handshake
     WebSocketServerHandshakerFactory wsFactory = new WebSocketServerHandshakerFactory(
-      getWebSocketLocation(req), null, false, Integer.MAX_VALUE);
+      getWebSocketLocation(req), null, true, Integer.MAX_VALUE);
     handshaker = wsFactory.newHandshaker(req);
     if (handshaker == null) {
       WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
@@ -155,29 +175,8 @@ private void handleWebSocketFrame(ChannelHandlerContext ctx, WebSocketFrame fram
     }
   }
 
-  private static void sendHttpResponse(
-    ChannelHandlerContext ctx, HttpRequest req, FullHttpResponse res) {
-    // Generate an error page if response status code is not OK (200).
-    if (res.status().code() != 200) {
-      ByteBuf buf = Unpooled.copiedBuffer(res.status().toString(), UTF_8);
-      res.content().writeBytes(buf);
-      buf.release();
-      setContentLength(res, res.content().readableBytes());
-    }
-
-    // Send the response and close the connection if necessary.
-    ChannelFuture f = ctx.channel().writeAndFlush(res);
-    if (!isKeepAlive(req) || res.status().code() != 200) {
-      f.addListener(ChannelFutureListener.CLOSE);
-    }
-  }
-
   @Override
   public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
     ctx.close();
   }
-
-  private static String getWebSocketLocation(HttpRequest req) {
-    return "ws://" + req.headers().get(HttpHeaderNames.HOST);
-  }
 }