[grid] Checking for IPv6 addresses

SeleniumHQ · Mar 13, 2020 · d38dfb3 · d38dfb3
1 parent d777cf5
commit d38dfb3
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 7 deletions.
diff --git a/java/server/src/org/openqa/selenium/events/zeromq/BoundZmqEventBus.java b/java/server/src/org/openqa/selenium/events/zeromq/BoundZmqEventBus.java
@@ -25,9 +25,13 @@
 import org.zeromq.ZContext;
 import org.zeromq.ZMQ;
 
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.Consumer;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 
 class BoundZmqEventBus implements EventBus {
@@ -46,12 +50,12 @@ class BoundZmqEventBus implements EventBus {
     LOG.info(String.format("XPUB binding to %s, XSUB binding to %s", xpubAddr, xsubAddr));
 
     xpub = context.createSocket(SocketType.XPUB);
-    xpub.setIPv6(true);
+    xpub.setIPv6(xpubAddr.isIPv6);
     xpub.setImmediate(true);
     xpub.bind(xpubAddr.bindTo);
 
     xsub = context.createSocket(SocketType.XSUB);
-    xsub.setIPv6(true);
+    xsub.setIPv6(xsubAddr.isIPv6);
     xsub.setImmediate(true);
     xsub.bind(xsubAddr.bindTo);
 
@@ -89,7 +93,7 @@ public void close() {
 
   private Addresses deriveAddresses(String host, String connection) {
     if (connection.startsWith("inproc:")) {
-      return new Addresses(connection, connection);
+      return new Addresses(connection, connection, false);
     }
 
     if (!connection.startsWith("tcp://")) {
@@ -109,19 +113,35 @@ private Addresses deriveAddresses(String host, String connection) {
       host = hostName;
     }
 
+    boolean isAddressIPv6 = false;
+    try {
+      if (InetAddress.getByName(host) instanceof Inet6Address ) {
+        isAddressIPv6 = true;
+        if (!host.startsWith("[")) {
+          host = String.format("[%s]", host);
+        }
+      }
+    } catch (UnknownHostException e) {
+      LOG.log(Level.WARNING, "Could not determine if host address is IPv6 or IPv4", e);
+    }
+
     return new Addresses(
         connection,
-        String.format("tcp://%s:%d", host, port));
+        String.format("tcp://%s:%d", host, port),
+        isAddressIPv6
+    );
   }
 
   private static class Addresses {
-    Addresses(String bindTo, String advertise) {
+    Addresses(String bindTo, String advertise, boolean isIPv6) {
       this.bindTo = bindTo;
       this.advertise = advertise;
+      this.isIPv6 = isIPv6;
     }
 
     String bindTo;
     String advertise;
+    boolean isIPv6;
 
     @Override
     public String toString() {

diff --git a/java/server/src/org/openqa/selenium/events/zeromq/UnboundEventBus.java b/java/server/src/org/openqa/selenium/events/zeromq/UnboundEventBus.java
@@ -29,6 +29,11 @@
 import org.zeromq.ZContext;
 import org.zeromq.ZMQ;
 
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.net.UnknownHostException;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -40,6 +45,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 
 class UnboundZmqEventBus implements EventBus {
@@ -64,12 +70,12 @@ class UnboundZmqEventBus implements EventBus {
     LOG.info(String.format("Connecting to %s and %s", publishConnection, subscribeConnection));
 
     sub = context.createSocket(SocketType.SUB);
-    sub.setIPv6(true);
+    sub.setIPv6(isSubAddressIPv6(publishConnection));
     sub.connect(publishConnection);
     sub.subscribe(new byte[0]);
 
     pub = context.createSocket(SocketType.PUB);
-    pub.setIPv6(true);
+    pub.setIPv6(isSubAddressIPv6(subscribeConnection));
     pub.connect(subscribeConnection);
 
     ZMQ.Poller poller = context.createPoller(1);
@@ -131,6 +137,15 @@ class UnboundZmqEventBus implements EventBus {
     }
   }
 
+  private boolean isSubAddressIPv6(String connection) {
+    try {
+      return InetAddress.getByName(new URI(connection).getHost()) instanceof Inet6Address;
+    } catch (UnknownHostException | URISyntaxException e) {
+      LOG.log(Level.WARNING, String.format("Could not determine if the address %s is IPv6 or IPv4", connection), e);
+    }
+    return false;
+  }
+
   @Override
   public void addListener(Type type, Consumer<Event> onType) {
     Objects.requireNonNull(type, "Event type must be set.");