Implemented a custom DNS resolver

Author: maiph

src/main/java/org/java_websocket/DnsResolver.java

@@ -0,0 +1,30 @@
+packageorg.java_websocket;
+
+importjava.net.InetAddress;
+importjava.net.URI;
+importjava.net.UnknownHostException;
+
+/**
+ * Users may implement this interface to override the default DNS lookup offered
+ * by the OS.
+ *
+ * @since 1.4.1
+ */
+publicinterfaceDnsResolver{
+
+/**
+ * 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.
+ */
+InetAddressresolve(URIuri) throwsUnknownHostException;
+
+}

src/main/java/org/java_websocket/client/WebSocketClient.java

@@ -28,10 +28,12 @@
 importjava.io.IOException;
 importjava.io.InputStream;
 importjava.io.OutputStream;
+importjava.net.InetAddress;
 importjava.net.InetSocketAddress;
 importjava.net.Proxy;
 importjava.net.Socket;
 importjava.net.URI;
+importjava.net.UnknownHostException;
 importjava.nio.ByteBuffer;
 importjava.util.Collection;
 importjava.util.Collections;
@@ -46,6 +48,7 @@
 importjavax.net.ssl.SSLSocketFactory;
 
 importorg.java_websocket.AbstractWebSocket;
+importorg.java_websocket.DnsResolver;
 importorg.java_websocket.WebSocket;
 importorg.java_websocket.WebSocketImpl;
 importorg.java_websocket.drafts.Draft;
@@ -131,6 +134,14 @@ public abstract class WebSocketClient extends AbstractWebSocket implements Runna
  */
 privateintconnectTimeout = 0;
 
+/**
+ * DNS resolver that translates a URI to an InetAddress
+ *
+ * @see InetAddress
+ * @since 1.4.1
+ */
+privateDnsResolverdnsResolver = 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 = newDnsResolver(){
+@Override
+publicInetAddressresolve(URIuri) throwsUnknownHostException{
+returnInetAddress.getByName(uri.getHost());
+ }
+ };
 if(httpHeaders != null){
 headers = newTreeMap<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
+ */
+publicvoidsetDnsResolver(DnsResolverdnsResolver){
+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( newInetSocketAddress( uri.getHost(), getPort() ), connectTimeout );
+if (!socket.isBound()){
+InetSocketAddressaddr = newInetSocketAddress(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;
 intport = getPort();
-Stringhost = uri.getHost() + (
+Stringhost = uri.getHost() + (
  (port != WebSocketImpl.DEFAULT_PORT && port != WebSocketImpl.DEFAULT_WSS_PORT)
- ? ":" + port
+ ? ":" + port
  : "" );
 
 HandshakeImpl1Clienthandshake = newHandshakeImpl1Client();
@@ -844,7 +873,7 @@ public InetSocketAddress getLocalSocketAddress(){
 publicInetSocketAddressgetRemoteSocketAddress(){
 returnengine.getRemoteSocketAddress();
  }
-
+
 @Override
 publicStringgetResourceDescriptor(){
 returnuri.getPath();