[grid] Adding vnc endpoints for live view in the UI

Author: diemol

java/server/src/org/openqa/selenium/grid/node/config/DriverServiceSessionFactory.java

@@ -67,7 +67,7 @@ public class DriverServiceSessionFactory implements SessionFactory {
   private final Predicate<Capabilities> predicate;
   private final DriverService.Builder<?, ?> builder;
   private final Capabilities stereotype;
-  private final BrowserOptionsMutator browserOptionsMutator;
+  private final SessionCapabilitiesMutator sessionCapabilitiesMutator;
 
   public DriverServiceSessionFactory(
       Tracer tracer,
@@ -80,7 +80,7 @@ public DriverServiceSessionFactory(
     this.stereotype = ImmutableCapabilities.copyOf(Require.nonNull("Stereotype", stereotype));
     this.predicate = Require.nonNull("Accepted capabilities predicate", predicate);
     this.builder = Require.nonNull("Driver service builder", builder);
-    this.browserOptionsMutator = new BrowserOptionsMutator(this.stereotype);
+    this.sessionCapabilitiesMutator = new SessionCapabilitiesMutator(this.stereotype);
   }
 
   @Override
@@ -101,7 +101,8 @@ public Either<WebDriverException, ActiveSession> apply(CreateSessionRequest sess
 
     try (Span span = tracer.getCurrentContext().createSpan("driver_service_factory.apply")) {
 
-      Capabilities capabilities = browserOptionsMutator.apply(sessionRequest.getDesiredCapabilities());
+      Capabilities capabilities = sessionCapabilitiesMutator
+        .apply(sessionRequest.getDesiredCapabilities());
 
       Optional<Platform> platformName = Optional.ofNullable(capabilities.getPlatformName());
       if (platformName.isPresent()) {
@@ -142,14 +143,13 @@ public Either<WebDriverException, ActiveSession> apply(CreateSessionRequest sess
         attributeMap.put(AttributeKey.DRIVER_RESPONSE.getKey(),
                          EventAttribute.setValue(response.toString()));
 
-        // TODO: This is a nasty hack. Try and make it elegant.
-
         Capabilities caps = new ImmutableCapabilities((Map<?, ?>) response.getValue());
         if (platformName.isPresent()) {
           caps = setInitialPlatform(caps, platformName.get());
         }
 
         caps = readDevToolsEndpointAndVersion(caps);
+        caps = readVncEndpoint(capabilities, caps);
 
         span.addEvent("Driver service created session", attributeMap);
         return Either.right(
@@ -214,6 +214,7 @@ public DevToolsInfo(URI cdpEndpoint, String version) {
       .filter(Optional::isPresent)
       .map(Optional::get)
       .findFirst();
+
     if (maybeInfo.isPresent()) {
       DevToolsInfo info = maybeInfo.get();
       return new PersistentCapabilities(caps)
@@ -223,6 +224,17 @@ public DevToolsInfo(URI cdpEndpoint, String version) {
     return caps;
   }
 
+  private Capabilities readVncEndpoint(Capabilities requestedCaps, Capabilities returnedCaps) {
+    String seVncEnabledCap = "se:vncEnabled";
+    String seVncEnabled = String.valueOf(requestedCaps.getCapability(seVncEnabledCap));
+    if (Boolean.parseBoolean(seVncEnabled)) {
+      returnedCaps = new PersistentCapabilities(returnedCaps)
+        .setCapability("se:vncLocalAddress", "ws://localhost:7900/websockify")
+        .setCapability(seVncEnabledCap, true);
+    }
+    return returnedCaps;
+  }
+
   // We set the platform to ANY before sending the caps to the driver because some drivers will
   // reject session requests when they cannot parse the platform.
   private Capabilities generalizePlatform(Capabilities caps) {

java/server/src/org/openqa/selenium/grid/node/config/BrowserOptionsMutator.java renamed to java/server/src/org/openqa/selenium/grid/node/config/SessionCapabilitiesMutator.java

@@ -21,27 +21,34 @@
 
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
+import org.openqa.selenium.PersistentCapabilities;
 
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Function;
 
-public class BrowserOptionsMutator implements Function<Capabilities, Capabilities> {
+public class SessionCapabilitiesMutator implements Function<Capabilities, Capabilities> {
 
   private static final ImmutableMap<String, String> BROWSER_OPTIONS = ImmutableMap.of(
     "chrome", "goog:chromeOptions",
     "firefox", "moz:firefoxOptions",
     "microsoftedge", "ms:edgeOptions");
+  private static final String SE_VNC_ENABLED = "se:vncEnabled";
   private final Capabilities slotStereotype;
 
 
-  public BrowserOptionsMutator(Capabilities slotStereotype) {
+  public SessionCapabilitiesMutator(Capabilities slotStereotype) {
     this.slotStereotype = slotStereotype;
   }
 
   @Override
   public Capabilities apply(Capabilities capabilities) {
+    if (slotStereotype.getCapability(SE_VNC_ENABLED) != null) {
+      capabilities = new PersistentCapabilities(capabilities)
+        .setCapability(SE_VNC_ENABLED, slotStereotype.getCapability(SE_VNC_ENABLED));
+    }
+
     if (!Objects.equals(slotStereotype.getBrowserName(), capabilities.getBrowserName())) {
       return capabilities;
     }

java/server/src/org/openqa/selenium/grid/node/local/LocalNode.java

@@ -81,7 +81,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
@@ -341,10 +340,9 @@ public Either<WebDriverException, CreateSessionResponse> newSession(CreateSessio
 
         // The session we return has to look like it came from the node, since we might be dealing
         // with a webdriver implementation that only accepts connections from localhost
-        Session externalSession = createExternalSession(
-          session,
-          externalUri,
-          slotToUse.isSupportingCdp() || caps.getCapability("se:cdp") != null);
+        boolean isSupportingCdp = slotToUse.isSupportingCdp() ||
+                                  caps.getCapability("se:cdp") != null;
+        Session externalSession = createExternalSession(session, externalUri, isSupportingCdp);
         return Either.right(new CreateSessionResponse(
           externalSession,
           getEncoder(session.getDownstreamDialect()).apply(externalSession)));
@@ -446,22 +444,31 @@ public void stop(SessionId id) throws NoSuchSessionException {
     tempFileSystems.invalidate(id);
   }
 
-  private Session createExternalSession(ActiveSession other, URI externalUri, boolean isSupportingCdp) {
+  private Session createExternalSession(ActiveSession other, URI externalUri,
+                                        boolean isSupportingCdp) {
     Capabilities toUse = ImmutableCapabilities.copyOf(other.getCapabilities());
 
-    // Rewrite the se:options if necessary to send the cdp url back
+    // Add se:cdp if necessary to send the cdp url back
     if (isSupportingCdp) {
       String cdpPath = String.format("/session/%s/se/cdp", other.getId());
       toUse = new PersistentCapabilities(toUse).setCapability("se:cdp", rewrite(cdpPath));
     }
 
+    // If enabled, set the VNC endpoint for live view
+    boolean isVncEnabled = toUse.getCapability("se:vncLocalAddress") != null;
+    if (isVncEnabled) {
+      String vncPath = String.format("/session/%s/se/vnc", other.getId());
+      toUse = new PersistentCapabilities(toUse).setCapability("se:vnc", rewrite(vncPath));
+    }
+
     return new Session(other.getId(), externalUri, other.getStereotype(), toUse, Instant.now());
   }
 
   private URI rewrite(String path) {
     try {
+      String scheme = "https".equals(gridUri.getScheme()) ? "wss" : "ws";
       return new URI(
-        "ws",
+        scheme,
         gridUri.getUserInfo(),
         gridUri.getHost(),
         gridUri.getPort(),