Implemented a custom DNS resolver

maiph · maiph · commit 5297299677a0 · 2019-06-13T11:43:47.000+01:00
diff --git a/src/main/java/org/java_websocket/DnsResolver.java b/src/main/java/org/java_websocket/DnsResolver.java
@@ -0,0 +1,30 @@
+package org.java_websocket;
+
+import java.net.InetAddress;
+import java.net.URI;
+import java.net.UnknownHostException;
+
+/**
+ * Users may implement this interface to override the default DNS lookup offered
+ * by the OS.
+ *
+ * @since 1.4.1
+ */
+public interface DnsResolver {
+
+  /**
+   * Resolves the IP address for the given URI.
+   *
+   * This method should never return null. If it's not able to resolve the IP
+   * address then it should throw an UnknownHostException
+   *
+   * @param uri The URI to be resolved
+   *
+   * @return The resolved IP address
+   *
+   * @throws UnknownHostException if no IP address for the <code>uri</code>
+   * could be found.
+   */
+  InetAddress resolve(URI uri) throws UnknownHostException;
+
+}
diff --git a/src/main/java/org/java_websocket/client/WebSocketClient.java b/src/main/java/org/java_websocket/client/WebSocketClient.java
@@ -28,10 +28,12 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.net.Socket;
 import java.net.URI;
+import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.util.Collection;
 import java.util.Collections;
@@ -46,6 +48,7 @@
 import javax.net.ssl.SSLSocketFactory;
 
 import org.java_websocket.AbstractWebSocket;
+import org.java_websocket.DnsResolver;
 import org.java_websocket.WebSocket;
 import org.java_websocket.WebSocketImpl;
 import org.java_websocket.drafts.Draft;
@@ -131,6 +134,14 @@ public abstract class WebSocketClient extends AbstractWebSocket implements Runna
 	 */
 	private int connectTimeout = 0;
 
+	/**
+	 * DNS resolver that translates a URI to an InetAddress
+	 *
+	 * @see InetAddress
+	 * @since 1.4.1
+	 */
+	private DnsResolver dnsResolver = null;
+
 	/**
 	 * Constructs a WebSocketClient instance and sets it to the connect to the
 	 * specified URI. The channel does not attampt to connect automatically. The connection
@@ -195,6 +206,12 @@ public WebSocketClient( URI serverUri , Draft protocolDraft , Map<String,String>
 		}
 		this.uri = serverUri;
 		this.draft = protocolDraft;
+		this.dnsResolver = new DnsResolver() {
+			@Override
+			public InetAddress resolve(URI uri) throws UnknownHostException {
+				return InetAddress.getByName(uri.getHost());
+			}
+		};
 		if(httpHeaders != null) {
 			headers = new TreeMap<String, String>(String.CASE_INSENSITIVE_ORDER);
 			headers.putAll(httpHeaders);
@@ -266,6 +283,17 @@ public void clearHeaders() {
 		headers = null;
 	}
 
+	/**
+	 * Sets a custom DNS resolver.
+	 *
+	 * @param dnsResolver The DnsResolver to use.
+	 *
+	 * @since 1.4.1
+	 */
+	public void setDnsResolver(DnsResolver dnsResolver) {
+		this.dnsResolver = dnsResolver;
+	}
+
 	/**
 	 * Reinitiates the websocket connection. This method does not block.
 	 * @since 1.3.8
@@ -431,8 +459,9 @@ public void run() {
 			socket.setTcpNoDelay( isTcpNoDelay() );
 			socket.setReuseAddress( isReuseAddr() );
 
-			if( !socket.isBound() ) {
-				socket.connect( new InetSocketAddress( uri.getHost(), getPort() ), connectTimeout );
+			if (!socket.isBound()) {
+				InetSocketAddress addr = new InetSocketAddress(dnsResolver.resolve(uri), uri.getPort());
+				socket.connect(addr, connectTimeout);
 			}
 
 			// if the socket is set by others we don't apply any TLS wrapper
@@ -509,9 +538,9 @@ private void sendHandshake() throws InvalidHandshakeException {
 		if( part2 != null )
 			path += '?' + part2;
 		int port = getPort();
-		String host = uri.getHost() + ( 
+		String host = uri.getHost() + (
 			(port != WebSocketImpl.DEFAULT_PORT && port != WebSocketImpl.DEFAULT_WSS_PORT)
-			? ":" + port 
+			? ":" + port
 			: "" );
 
 		HandshakeImpl1Client handshake = new HandshakeImpl1Client();
@@ -844,7 +873,7 @@ public InetSocketAddress getLocalSocketAddress() {
 	public InetSocketAddress getRemoteSocketAddress() {
 		return engine.getRemoteSocketAddress();
 	}
-	
+
 	@Override
 	public String getResourceDescriptor() {
 		return uri.getPath();
