[java] Update subPath for ProxyNodeWebsockets

Signed-off-by: Viet Nguyen Duc <[email protected]>
NDViet · Jan 8, 2024 · 7d864d4 · 7d864d4
1 parent b0f42c3
commit 7d864d4
Show file tree

Hide file tree

Showing 9 changed files with 91 additions and 10 deletions.
diff --git a/java/src/org/openqa/selenium/grid/commands/Standalone.java b/java/src/org/openqa/selenium/grid/commands/Standalone.java
@@ -222,7 +222,7 @@ protected Handlers createHandlers(Config config) {
     httpHandler = combine(httpHandler, Route.get("/readyz").to(() -> readinessCheck));
     Node node = createNode(config, bus, distributor, combinedHandler);
 
-    return new Handlers(httpHandler, new ProxyNodeWebsockets(clientFactory, node));
+    return new Handlers(httpHandler, new ProxyNodeWebsockets(clientFactory, node, subPath));
   }
 
   @Override

diff --git a/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java b/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java
@@ -56,18 +56,20 @@ public class ProxyNodeWebsockets
       ImmutableSet.of("goog:chromeOptions", "moz:debuggerAddress", "ms:edgeOptions");
   private final HttpClient.Factory clientFactory;
   private final Node node;
+  private final String subPath;
 
-  public ProxyNodeWebsockets(HttpClient.Factory clientFactory, Node node) {
+  public ProxyNodeWebsockets(HttpClient.Factory clientFactory, Node node, String subPath) {
     this.clientFactory = Objects.requireNonNull(clientFactory);
     this.node = Objects.requireNonNull(node);
+    this.subPath = subPath;
   }
 
   @Override
   public Optional<Consumer<Message>> apply(String uri, Consumer<Message> downstream) {
-    UrlTemplate.Match fwdMatch = FWD_TEMPLATE.match(uri);
-    UrlTemplate.Match cdpMatch = CDP_TEMPLATE.match(uri);
-    UrlTemplate.Match bidiMatch = BIDI_TEMPLATE.match(uri);
-    UrlTemplate.Match vncMatch = VNC_TEMPLATE.match(uri);
+    UrlTemplate.Match fwdMatch = FWD_TEMPLATE.match(uri, this.subPath);
+    UrlTemplate.Match cdpMatch = CDP_TEMPLATE.match(uri, this.subPath);
+    UrlTemplate.Match bidiMatch = BIDI_TEMPLATE.match(uri, this.subPath);
+    UrlTemplate.Match vncMatch = VNC_TEMPLATE.match(uri, this.subPath);
 
     if (bidiMatch == null && cdpMatch == null && vncMatch == null && fwdMatch == null) {
       return Optional.empty();

diff --git a/java/src/org/openqa/selenium/grid/node/config/NodeFlags.java b/java/src/org/openqa/selenium/grid/node/config/NodeFlags.java
@@ -31,6 +31,7 @@
 import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_SESSION_TIMEOUT;
 import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_USE_SELENIUM_MANAGER;
 import static org.openqa.selenium.grid.node.config.NodeOptions.DEFAULT_VNC_ENV_VAR;
+import static org.openqa.selenium.grid.node.config.NodeOptions.NETWORK_SECTION;
 import static org.openqa.selenium.grid.node.config.NodeOptions.NODE_SECTION;
 import static org.openqa.selenium.grid.node.config.NodeOptions.OVERRIDE_MAX_SESSIONS;
 

diff --git a/java/src/org/openqa/selenium/grid/node/config/NodeOptions.java b/java/src/org/openqa/selenium/grid/node/config/NodeOptions.java
@@ -75,6 +75,7 @@ public class NodeOptions {
   public static final boolean DEFAULT_ENABLE_CDP = true;
   public static final boolean DEFAULT_ENABLE_BIDI = true;
   static final String NODE_SECTION = "node";
+  static final String NETWORK_SECTION = "network";
   static final boolean DEFAULT_DETECT_DRIVERS = true;
   static final boolean DEFAULT_USE_SELENIUM_MANAGER = false;
   static final boolean OVERRIDE_MAX_SESSIONS = false;
@@ -162,6 +163,25 @@ public boolean isManagedDownloadsEnabled() {
     return config.getBool(NODE_SECTION, "enable-managed-downloads").orElse(Boolean.FALSE);
   }
 
+  public String getSubPath() {
+    String configSubPath = normalizeSubPath(config.get(NETWORK_SECTION, "sub-path").orElse(""));
+    if (configSubPath.isEmpty()) {
+      configSubPath = normalizeSubPath(getPublicGridUri().map(URI::getPath).orElse(configSubPath));
+    }
+    return configSubPath;
+  }
+
+  public static String normalizeSubPath(String prefix) {
+    prefix = prefix.trim();
+    if (!prefix.startsWith("/")) {
+      prefix = "/" + prefix; // Prefix with a '/' if absent.
+    }
+    if (prefix.endsWith("/")) {
+      prefix = prefix.substring(0, prefix.length() - 1); // Remove the trailing '/' if present.
+    }
+    return prefix;
+  }
+
   public Node getNode() {
     return config.getClass(NODE_SECTION, "implementation", Node.class, DEFAULT_NODE_IMPLEMENTATION);
   }

diff --git a/java/src/org/openqa/selenium/grid/node/httpd/NodeServer.java b/java/src/org/openqa/selenium/grid/node/httpd/NodeServer.java
@@ -171,7 +171,8 @@ protected Handlers createHandlers(Config config) {
 
     Route httpHandler = Route.combine(node, get("/readyz").to(() -> readinessCheck));
 
-    return new Handlers(httpHandler, new ProxyNodeWebsockets(clientFactory, node));
+    return new Handlers(
+        httpHandler, new ProxyNodeWebsockets(clientFactory, node, nodeOptions.getSubPath()));
   }
 
   @Override

diff --git a/java/src/org/openqa/selenium/grid/node/local/LocalNode.java b/java/src/org/openqa/selenium/grid/node/local/LocalNode.java
@@ -127,6 +127,7 @@ public class LocalNode extends Node {
   private final int configuredSessionCount;
   private final boolean cdpEnabled;
   private final boolean managedDownloadsEnabled;
+  private final String subPath;
 
   private final boolean bidiEnabled;
   private final AtomicBoolean drainAfterSessions = new AtomicBoolean();
@@ -153,6 +154,7 @@ protected LocalNode(
       Duration heartbeatPeriod,
       List<SessionSlot> factories,
       Secret registrationSecret,
+      String subPath,
       boolean managedDownloadsEnabled) {
     super(tracer, new NodeId(UUID.randomUUID()), uri, registrationSecret);
 
@@ -171,6 +173,7 @@ protected LocalNode(
     this.cdpEnabled = cdpEnabled;
     this.bidiEnabled = bidiEnabled;
     this.managedDownloadsEnabled = managedDownloadsEnabled;
+    this.subPath = subPath;
 
     this.healthCheck =
         healthCheck == null
@@ -838,9 +841,7 @@ private Session createExternalSession(
   private URI rewrite(String path) {
     try {
       String scheme = "https".equals(gridUri.getScheme()) ? "wss" : "ws";
-      if (gridUri.getPath() != null && !gridUri.getPath().equals("/")) {
-        path = gridUri.getPath() + path;
-      }
+      path = subPath + path;
       return new URI(
           scheme, gridUri.getUserInfo(), gridUri.getHost(), gridUri.getPort(), path, null, null);
     } catch (URISyntaxException e) {
@@ -960,12 +961,14 @@ public static class Builder {
     private HealthCheck healthCheck;
     private Duration heartbeatPeriod = Duration.ofSeconds(NodeOptions.DEFAULT_HEARTBEAT_PERIOD);
     private boolean managedDownloadsEnabled = false;
+    private String subPath;
 
     private Builder(Tracer tracer, EventBus bus, URI uri, URI gridUri, Secret registrationSecret) {
       this.tracer = Require.nonNull("Tracer", tracer);
       this.bus = Require.nonNull("Event bus", bus);
       this.uri = Require.nonNull("Remote node URI", uri);
       this.gridUri = Require.nonNull("Grid URI", gridUri);
+      this.subPath = NodeOptions.normalizeSubPath(gridUri.getPath());
       this.registrationSecret = Require.nonNull("Registration secret", registrationSecret);
       this.factories = ImmutableList.builder();
     }
@@ -1014,6 +1017,13 @@ public Builder enableManagedDownloads(boolean enable) {
       return this;
     }
 
+    public Builder setSubPath(String subPath) {
+      if (!subPath.isEmpty()) {
+        this.subPath = NodeOptions.normalizeSubPath(subPath);
+      }
+      return this;
+    }
+
     public LocalNode build() {
       return new LocalNode(
           tracer,
@@ -1030,6 +1040,7 @@ public LocalNode build() {
           heartbeatPeriod,
           factories.build(),
           registrationSecret,
+          subPath,
           managedDownloadsEnabled);
     }
 

diff --git a/java/src/org/openqa/selenium/grid/node/local/LocalNodeFactory.java b/java/src/org/openqa/selenium/grid/node/local/LocalNodeFactory.java
@@ -70,6 +70,7 @@ public static Node create(Config config) {
             .enableCdp(nodeOptions.isCdpEnabled())
             .enableBiDi(nodeOptions.isBiDiEnabled())
             .enableManagedDownloads(nodeOptions.isManagedDownloadsEnabled())
+            .setSubPath(nodeOptions.getSubPath())
             .heartbeatPeriod(nodeOptions.getHeartbeatPeriod());
 
     List<DriverService.Builder<?, ?>> builders = new ArrayList<>();

diff --git a/java/src/org/openqa/selenium/remote/http/UrlTemplate.java b/java/src/org/openqa/selenium/remote/http/UrlTemplate.java
@@ -139,6 +139,18 @@ public UrlTemplate.Match match(String matchAgainst) {
     return compiled.apply(matchAgainst);
   }
 
+  /**
+   * @return A {@link Match} with all parameters filled if successful, null otherwise. Remove
+   *     subPath from matchAgainst before matching.
+   */
+  public UrlTemplate.Match match(String matchAgainst, String subPath) {
+    if (matchAgainst == null || subPath == null) {
+      return null;
+    }
+    matchAgainst = matchAgainst.replaceFirst(subPath, "");
+    return compiled.apply(matchAgainst);
+  }
+
   @SuppressWarnings("InnerClassMayBeStatic")
   public class Match {
     private final String url;

diff --git a/java/test/org/openqa/selenium/grid/node/config/NodeOptionsTest.java b/java/test/org/openqa/selenium/grid/node/config/NodeOptionsTest.java
@@ -673,6 +673,39 @@ void settingTheHubWithDefaultValueSetsTheGridUrlToTheNonLoopbackAddress() {
         .isEqualTo(Optional.of(URI.create(nonLoopbackAddressUrl)));
   }
 
+  @Test
+  void settingSubPathForNodeServerFromNetworkSection() {
+    String[] rawConfig =
+        new String[] {
+          "[network]", "sub-path = \"/mySubPath\"",
+        };
+    Config config = new TomlConfig(new StringReader(String.join("\n", rawConfig)));
+    NodeOptions nodeOptions = new NodeOptions(config);
+    assertThat(nodeOptions.getSubPath()).isEqualTo("/mySubPath");
+  }
+
+  @Test
+  void settingSubPathForNodeServerExtractFromGridUrl() {
+    String[] rawConfig =
+        new String[] {
+          "[node]", "grid-url = \"http://localhost:4444/mySubPath\"",
+        };
+    Config config = new TomlConfig(new StringReader(String.join("\n", rawConfig)));
+    NodeOptions nodeOptions = new NodeOptions(config);
+    assertThat(nodeOptions.getSubPath()).isEqualTo("/mySubPath");
+  }
+
+  @Test
+  void settingSubPathForNodeServerExtractFromHub() {
+    String[] rawConfig =
+        new String[] {
+          "[node]", "hub = \"http://0.0.0.0:4444/mySubPath\"",
+        };
+    Config config = new TomlConfig(new StringReader(String.join("\n", rawConfig)));
+    NodeOptions nodeOptions = new NodeOptions(config);
+    assertThat(nodeOptions.getSubPath()).isEqualTo("/mySubPath");
+  }
+
   @Test
   void notSettingSlotMatcherAvailable() {
     String[] rawConfig =