added SOCKS5 support for DNS over tcp and udp (first working version)

dnsjava · tw-datascientist · Oct 2, 2024 · Oct 3, 2024 · Oct 18, 2024 · Dec 5, 2024
commit 5ff82ed97d4bfbb7a6ed95682ecb45cacb3ece28
diff --git a/src/main/java/org/xbill/DNS/DefaultIoClient.java b/src/main/java/org/xbill/DNS/DefaultIoClient.java
@@ -33,6 +33,17 @@ public CompletableFuture<byte[]> sendAndReceiveTcp(
     return tcpIoClient.sendAndReceiveTcp(local, remote, query, data, timeout);
   }
 
+  @Override
+  public CompletableFuture<byte[]> sendAndReceiveTcp(
+    InetSocketAddress local,
+    InetSocketAddress remote,
+    Socks5Proxy proxy,
+    Message query,
+    byte[] data,
+    Duration timeout) {
+    return tcpIoClient.sendAndReceiveTcp(local, remote, proxy, query, data, timeout);
+  }
+
   @Override
   public CompletableFuture<byte[]> sendAndReceiveUdp(
       InetSocketAddress local,

diff --git a/src/main/java/org/xbill/DNS/NioTcpClient.java b/src/main/java/org/xbill/DNS/NioTcpClient.java
@@ -44,6 +44,9 @@ private void processPendingRegistrations() {
         if (!state.channel.isConnected()) {
           state.channel.register(selector, SelectionKey.OP_CONNECT, state);
         } else {
+          if (state.channel.keyFor(selector) == null) {
+            state.channel.register(selector, SelectionKey.OP_CONNECT, state);
+          }
           state.channel.keyFor(selector).interestOps(SelectionKey.OP_WRITE);
         }
       } catch (IOException e) {
@@ -285,10 +288,21 @@ private static class ChannelKey {
     final InetSocketAddress remote;
   }
 
+  @Override
+  public CompletableFuture<byte[]> sendAndReceiveTcp(
+    InetSocketAddress local,
+    InetSocketAddress remote,
+    Message query,
+    byte[] data,
+    Duration timeout) {
+    return this.sendAndReceiveTcp(local, remote, null, query, data, timeout);
+  }
+
   @Override
   public CompletableFuture<byte[]> sendAndReceiveTcp(
       InetSocketAddress local,
       InetSocketAddress remote,
+      Socks5Proxy proxy,
       Message query,
       byte[] data,
       Duration timeout) {
@@ -309,7 +323,14 @@ public CompletableFuture<byte[]> sendAndReceiveTcp(
                     c.bind(local);
                   }
 
-                  c.connect(remote);
+                  if (proxy != null) {
+                    c.configureBlocking(true);
+                    c.connect(proxy.getProxyAddress());
+                    proxy.socks5TcpHandshake(c, remote);
+                  } else {
+                    c.connect(remote);
+                  }
+                  c.configureBlocking(false);
                   return new ChannelState(c);
                 } catch (IOException e) {
                   if (c != null) {

diff --git a/src/main/java/org/xbill/DNS/SimpleResolver.java b/src/main/java/org/xbill/DNS/SimpleResolver.java
@@ -7,6 +7,7 @@
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
+import java.nio.channels.SocketChannel;
 import java.time.Duration;
 import java.util.List;
 import java.util.Objects;
@@ -19,6 +20,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.xbill.DNS.io.DefaultIoClientFactory;
 import org.xbill.DNS.io.IoClientFactory;
+import org.xbill.DNS.io.TcpIoClient;
 
 /**
  * An implementation of Resolver that sends one query to one server. SimpleResolver handles TCP
@@ -40,6 +42,7 @@ public class SimpleResolver implements Resolver {
 
   private InetSocketAddress address;
   private InetSocketAddress localAddress;
+  private Socks5Proxy proxy;
   private boolean useTCP;
   private boolean ignoreTruncation;
   private OPTRecord queryOPT = new OPTRecord(DEFAULT_EDNS_PAYLOADSIZE, 0, 0, 0);
@@ -99,6 +102,11 @@ public SimpleResolver(InetSocketAddress host) {
     address = Objects.requireNonNull(host, "host must not be null");
   }
 
+  /** Creates a SimpleResolver that will query the specified host via the specified SOCKS5 proxy */
+  public SimpleResolver(Socks5Proxy socks5Proxy) {
+    proxy = Objects.requireNonNull(socks5Proxy, "proxy must not be null");
+  }
+
   /** Creates a SimpleResolver that will query the specified host */
   public SimpleResolver(InetAddress host) {
     Objects.requireNonNull(host, "host must not be null");
@@ -387,22 +395,68 @@ CompletableFuture<Message> sendAsync(Message query, boolean forceTcp, Executor e
     }
 
     CompletableFuture<byte[]> result;
+    SocketChannel c = null;
     if (tcp) {
-      result =
-          ioClientFactory
-              .createOrGetTcpClient()
-              .sendAndReceiveTcp(localAddress, address, query, out, timeoutValue);
+      TcpIoClient tcpClient = ioClientFactory.createOrGetTcpClient();
+      if (proxy != null) {
+        localAddress = proxy.getLocalAddress();
+        address = proxy.getRemoteAddress();
+        result = tcpClient.sendAndReceiveTcp(localAddress, address, proxy, query, out, timeoutValue);
+      } else {
+        result = tcpClient.sendAndReceiveTcp(localAddress, address, query, out, timeoutValue);
+      }
     } else {
+      if (proxy != null) {
+        try {
+          c = SocketChannel.open();
+          if (localAddress != null) {
+            c.bind(localAddress);
+          }
+          if (proxy != null) {
+            c.connect(proxy.getProxyAddress());
+            address = proxy.socks5UdpAssociateHandshake(c);
+          }
+        } catch (IOException e) {
+          if (c != null) {
+            try {
+              c.close();
+            } catch (IOException ee) {
+              // ignore
+            }
+          }
+          return new CompletableFuture<>().thenComposeAsync(in -> {
+            CompletableFuture<Message> f = new CompletableFuture<>();
+            f.completeExceptionally(new WireParseException("Error in Udp Associate SOCKS5 handshake", e));
+            return f;
+          }, executor);
+        }
+        out = proxy.addUdpHeader(out, proxy.getRemoteAddress());
+        localAddress = proxy.getLocalAddress();
+      }
       result =
           ioClientFactory
               .createOrGetUdpClient()
               .sendAndReceiveUdp(localAddress, address, query, out, udpSize, timeoutValue);
     }
 
+    SocketChannel finalC = c;
     return result.thenComposeAsync(
         in -> {
           CompletableFuture<Message> f = new CompletableFuture<>();
 
+          // finally close the tcp connection
+          // and remove SOCKS5 udp header from the response
+          if (proxy != null && !tcp) {
+            if (finalC != null) {
+              try {
+                finalC.close();
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            }
+            in = proxy.removeUdpHeader(in);
+          }
+
           // Check that the response is long enough.
           if (in.length < Header.LENGTH) {
             f.completeExceptionally(new WireParseException("invalid DNS header - too short"));

diff --git a/src/main/java/org/xbill/DNS/Socks5Proxy.java b/src/main/java/org/xbill/DNS/Socks5Proxy.java
@@ -0,0 +1,192 @@
+package org.xbill.DNS;
+
+import lombok.Getter;
+
+import java.net.InetSocketAddress;
+import java.net.UnknownHostException;
+import java.nio.ByteBuffer;
+import java.nio.channels.SocketChannel;
+
+@Getter
+public class Socks5Proxy {
+  private static final byte SOCKS5_VERSION = 0x05;
+
+  private static final byte SOCKS5_AUTH_NONE = 0x00;
+  private static final byte SOCKS5_AUTH_GSSAPI = 0x01;
+  private static final byte SOCKS5_AUTH_USER_PASS = 0x02;
+  private static final byte SOCKS5_AUTH_NO_ACCEPTABLE_METHODS = (byte) 0xFF;
+
+  private static final byte SOCKS5_CMD_CONNECT = 0x01;
+  private static final byte SOCKS5_CMD_BIND = 0x02;
+  private static final byte SOCKS5_CMD_UDP_ASSOCIATE = 0x03;
+
+  private static final byte SOCKS5_ATYP_IPV4 = 0x01;
+  private static final byte SOCKS5_ATYP_DOMAINNAME = 0x03;
+  private static final byte SOCKS5_ATYP_IPV6 = 0x04;
+
+  private static final byte SOCKS5_REP_SUCCEEDED = 0x00;
+  private static final byte SOCKS5_REP_GENERAL_FAILURE = 0x01;
+  private static final byte SOCKS5_REP_CONNECTION_NOT_ALLOWED = 0x02;
+  private static final byte SOCKS5_REP_NETWORK_UNREACHABLE = 0x03;
+  private static final byte SOCKS5_REP_HOST_UNREACHABLE = 0x04;
+  private static final byte SOCKS5_REP_CONNECTION_REFUSED = 0x05;
+  private static final byte SOCKS5_REP_TTL_EXPIRED = 0x06;
+  private static final byte SOCKS5_REP_COMMAND_NOT_SUPPORTED = 0x07;
+  private static final byte SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED = 0x08;
+
+  private static final byte SOCKS5_RESERVED = 0x00;
+
+  private final InetSocketAddress remoteAddress;
+  private final InetSocketAddress localAddress;
+  private final InetSocketAddress proxyAddress;
+
+
+  public Socks5Proxy(InetSocketAddress proxyAddress, InetSocketAddress remoteAddress, InetSocketAddress localAddress) {
+    this.remoteAddress = remoteAddress;
+    this.localAddress = localAddress;
+    this.proxyAddress = proxyAddress;
+  }
+
+  public void socks5MethodSelection(SocketChannel c) {
+    ByteBuffer buffer = ByteBuffer.allocate(3);
+    buffer.put(SOCKS5_VERSION);
+    buffer.put((byte) 1);
+    buffer.put(SOCKS5_AUTH_NONE);
+    buffer.flip();
+
+    try {
+      c.write(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to write to TCP channel", e);
+    }
+
+    buffer.clear();
+
+    try {
+      c.read(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to read from TCP channel", e);
+    }
+
+    buffer.flip();
+    if (buffer.get() != SOCKS5_VERSION) {
+      throw new IllegalArgumentException("Invalid version");
+    }
+
+    if (buffer.get() == SOCKS5_AUTH_NO_ACCEPTABLE_METHODS) {
+      throw new IllegalArgumentException("No acceptable methods");
+    }
+  }
+
+  public void socks5HeaderExchange(SocketChannel c, InetSocketAddress remote) {
+    ByteBuffer buffer = ByteBuffer.allocate(10);
+    buffer.put(SOCKS5_VERSION);
+    buffer.put(SOCKS5_CMD_CONNECT);
+    buffer.put(SOCKS5_RESERVED);
+    buffer.put(SOCKS5_ATYP_IPV4);
+    buffer.put(remote.getAddress().getAddress());
+    buffer.putShort((short) remote.getPort());
+    buffer.flip();
+
+    try {
+      c.write(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to write to TCP channel", e);
+    }
+    buffer.clear();
+
+    try {
+      c.read(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to read from TCP channel", e);
+    }
+    buffer.flip();
+
+    if (buffer.get() != SOCKS5_VERSION) {
+      throw new IllegalArgumentException("Invalid version");
+    }
+
+    byte reply = buffer.get();
+    if (reply != SOCKS5_REP_SUCCEEDED) {
+      throw new IllegalArgumentException("Failed to connect to remote server: " + reply);
+    }
+  }
+
+  public InetSocketAddress socks5UdpAssociateExchange(SocketChannel c) {
+    ByteBuffer buffer = ByteBuffer.allocate(10);
+    buffer.put(SOCKS5_VERSION);
+    buffer.put(SOCKS5_CMD_UDP_ASSOCIATE);
+    buffer.put(SOCKS5_RESERVED);
+    buffer.put(SOCKS5_ATYP_IPV4);
+    buffer.put(new byte[] {0, 0, 0, 0});
+    buffer.putShort((short) 0x00);
+    buffer.flip();
+
+    try {
+      c.write(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to write to TCP channel", e);
+    }
+    buffer.clear();
+
+    try {
+      c.read(buffer);
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to read from TCP channel", e);
+    }
+    buffer.flip();
+
+    if (buffer.get() != SOCKS5_VERSION) {
+      throw new IllegalArgumentException("Invalid version");
+    }
+
+    byte reply = buffer.get();
+    if (reply != SOCKS5_REP_SUCCEEDED) {
+      throw new IllegalArgumentException("Failed to connect to remote server: " + reply);
+    }
+
+    buffer.get(); // skip RSV byte
+
+    byte atyp = buffer.get();
+    if (atyp != SOCKS5_ATYP_IPV4) {
+      throw new IllegalArgumentException("Invalid address type");
+    }
+
+    byte[] addr = new byte[4];
+    buffer.get(addr);
+    int port = buffer.getShort() & 0xFFFF;
+    return new InetSocketAddress(this.getProxyAddress().getAddress(), port);
+  }
+
+  public byte[] addUdpHeader(byte[] in, InetSocketAddress to) {
+    ByteBuffer buffer = ByteBuffer.allocate(in.length + 10);
+    buffer.put(SOCKS5_VERSION);
+    buffer.put(SOCKS5_RESERVED);
+    buffer.put(SOCKS5_RESERVED);
+    buffer.put(SOCKS5_ATYP_IPV4);
+    buffer.put(to.getAddress().getAddress());
+    buffer.putShort((short) to.getPort());
+    buffer.put(in);
+
+    return buffer.array();
+  }
+
+  public byte[] removeUdpHeader(byte[] in) {
+    byte[] out = new byte[in.length - 10];
+    System.arraycopy(in, 10, out, 0, in.length - 10);
+    return out;
+  }
+
+  public void socks5TcpHandshake(
+    SocketChannel c, InetSocketAddress remote) {
+    this.socks5MethodSelection(c);
+    this.socks5HeaderExchange(c, remote);
+  }
+
+  public InetSocketAddress socks5UdpAssociateHandshake(
+    SocketChannel c
+    ) throws UnknownHostException {
+    this.socks5MethodSelection(c);
+    return this.socks5UdpAssociateExchange(c);
+  }
+}
diff --git a/src/main/java/org/xbill/DNS/io/TcpIoClient.java b/src/main/java/org/xbill/DNS/io/TcpIoClient.java
@@ -6,6 +6,7 @@
 import java.util.concurrent.CompletableFuture;
 import org.xbill.DNS.Message;
 import org.xbill.DNS.Resolver;
+import org.xbill.DNS.Socks5Proxy;
 
 /**
  * Serves as an interface from a {@link Resolver} to the underlying mechanism for sending bytes over
@@ -32,4 +33,12 @@ CompletableFuture<byte[]> sendAndReceiveTcp(
       Message query,
       byte[] data,
       Duration timeout);
+
+  CompletableFuture<byte[]> sendAndReceiveTcp(
+    InetSocketAddress local,
+    InetSocketAddress remote,
+    Socks5Proxy proxy,
+    Message query,
+    byte[] data,
+    Duration timeout);
 }