redo code changes in SimpleResolver and the IoClients

dnsjava · tw-datascientist · Oct 2, 2024 · Oct 3, 2024 · Oct 18, 2024 · Dec 5, 2024
commit 2a1192c822c90d3c7a4809a6daf766f2b952317d
diff --git a/src/main/java/org/xbill/DNS/DefaultIoClient.java b/src/main/java/org/xbill/DNS/DefaultIoClient.java
@@ -33,17 +33,6 @@ 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,
@@ -54,16 +43,4 @@ public CompletableFuture<byte[]> sendAndReceiveUdp(
       Duration timeout) {
     return udpIoClient.sendAndReceiveUdp(local, remote, query, data, max, timeout);
   }
-
-  @Override
-  public CompletableFuture<byte[]> sendAndReceiveUdp(
-    InetSocketAddress local,
-    InetSocketAddress remote,
-    Socks5Proxy proxy,
-    Message query,
-    byte[] data,
-    int max,
-    Duration timeout) {
-    return udpIoClient.sendAndReceiveUdp(local, remote, proxy, query, data, max, timeout);
-  }
 }
diff --git a/src/main/java/org/xbill/DNS/NioTcpClient.java b/src/main/java/org/xbill/DNS/NioTcpClient.java
@@ -44,9 +44,6 @@ 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) {
@@ -288,21 +285,10 @@ 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) {
@@ -323,14 +309,7 @@ public CompletableFuture<byte[]> sendAndReceiveTcp(
                     c.bind(local);
                   }
 
-                  if (proxy != null) {
-                    c.configureBlocking(true);
-                    c.connect(proxy.getProxyAddress());
-                    proxy.socks5TcpHandshake(c, remote);
-                  } else {
-                    c.connect(remote);
-                  }
-                  c.configureBlocking(false);
+                  c.connect(remote);
                   return new ChannelState(c);
                 } catch (IOException e) {
                   if (c != null) {

diff --git a/src/main/java/org/xbill/DNS/NioUdpClient.java b/src/main/java/org/xbill/DNS/NioUdpClient.java
@@ -10,7 +10,6 @@
 import java.nio.channels.DatagramChannel;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.Selector;
-import java.nio.channels.SocketChannel;
 import java.security.SecureRandom;
 import java.time.Duration;
 import java.util.Iterator;
@@ -64,7 +63,7 @@ private void processPendingRegistrations() {
 
       try {
         log.trace("Registering OP_READ for transaction with id {}", t.id);
-        t.udpChannel.register(selector(), SelectionKey.OP_READ, t);
+        t.channel.register(selector(), SelectionKey.OP_READ, t);
         t.send();
       } catch (IOException e) {
         t.completeExceptionally(e);
@@ -88,19 +87,17 @@ private class Transaction implements KeyProcessor {
     private final byte[] data;
     private final int max;
     private final long endTime;
-    private final DatagramChannel udpChannel;
-    private final SocketChannel tcpChannel;
-    private final Socks5Proxy proxy;
+    private final DatagramChannel channel;
     private final CompletableFuture<byte[]> f;
 
     void send() throws IOException {
       ByteBuffer buffer = ByteBuffer.wrap(data);
       verboseLog(
           "UDP write: transaction id=" + id,
-          udpChannel.socket().getLocalSocketAddress(),
-          udpChannel.socket().getRemoteSocketAddress(),
+          channel.socket().getLocalSocketAddress(),
+          channel.socket().getRemoteSocketAddress(),
           data);
-      int n = udpChannel.send(buffer, udpChannel.socket().getRemoteSocketAddress());
+      int n = channel.send(buffer, channel.socket().getRemoteSocketAddress());
       if (n == 0) {
         throw new EOFException(
             "Insufficient room for the datagram in the underlying output buffer for transaction "
@@ -141,14 +138,6 @@ public void processReadyKey(SelectionKey key) {
           keyChannel.socket().getRemoteSocketAddress(),
           resultingData);
       silentDisconnectAndCloseChannel();
-      if (proxy != null && tcpChannel != null) {
-        resultingData = proxy.removeUdpHeader(resultingData);
-        try {
-          tcpChannel.close();
-        } catch (IOException e) {
-          // ignore, we either already have everything we need or can't do anything
-        }
-      }
       f.complete(resultingData);
       pendingTransactions.remove(this);
     }
@@ -160,64 +149,32 @@ private void completeExceptionally(Exception e) {
 
     private void silentDisconnectAndCloseChannel() {
       try {
-        udpChannel.disconnect();
+        channel.disconnect();
       } catch (IOException e) {
         // ignore, we either already have everything we need or can't do anything
       } finally {
-        NioUdpClient.silentCloseChannel(udpChannel);
+        NioUdpClient.silentCloseChannel(channel);
       }
     }
   }
 
-  @Override
-  public CompletableFuture<byte[]> sendAndReceiveUdp(
-    InetSocketAddress local,
-    InetSocketAddress remote,
-    Message query,
-    byte[] data,
-    int max,
-    Duration timeout) {
-    return sendAndReceiveUdp(local, remote, null, query, data, max, timeout);
-  }
-
   @Override
   public CompletableFuture<byte[]> sendAndReceiveUdp(
       InetSocketAddress local,
       InetSocketAddress remote,
-      Socks5Proxy proxy,
       Message query,
       byte[] data,
       int max,
       Duration timeout) {
     long endTime = System.nanoTime() + timeout.toNanos();
     CompletableFuture<byte[]> f = new CompletableFuture<>();
-    DatagramChannel udpChannel = null;
-    SocketChannel tcpChannel = null;
+    DatagramChannel channel = null;
     try {
       final Selector selector = selector();
+      channel = DatagramChannel.open();
+      channel.configureBlocking(false);
 
-      // SOCKS5 handshake to set up the UDP association
-      if (proxy != null) {
-        data = proxy.addUdpHeader(data, remote);
-        try {
-          tcpChannel = SocketChannel.open();
-          if (local != null) {
-            tcpChannel.bind(local);
-          }
-          tcpChannel.connect(proxy.getProxyAddress());
-          remote = proxy.socks5UdpAssociateHandshake(tcpChannel);
-        } catch (IOException e) {
-          return new CompletableFuture<>().thenComposeAsync(in -> {
-            f.completeExceptionally(new WireParseException("Error in Udp Associate SOCKS5 handshake", e));
-            return f;
-          });
-        }
-      }
-
-      udpChannel = DatagramChannel.open();
-      udpChannel.configureBlocking(false);
-
-      Transaction t = new Transaction(query.getHeader().getID(), data, max, endTime, udpChannel, tcpChannel, proxy, f);
+      Transaction t = new Transaction(query.getHeader().getID(), data, max, endTime, channel, f);
       if (local == null || local.getPort() == 0) {
         boolean bound = false;
         for (int i = 0; i < 1024; i++) {
@@ -236,7 +193,7 @@ public CompletableFuture<byte[]> sendAndReceiveUdp(
               addr = new InetSocketAddress(local.getAddress(), port);
             }
 
-            udpChannel.bind(addr);
+            channel.bind(addr);
             bound = true;
             break;
           } catch (SocketException e) {
@@ -250,16 +207,16 @@ public CompletableFuture<byte[]> sendAndReceiveUdp(
         }
       }
 
-      udpChannel.connect(remote);
+      channel.connect(remote);
       pendingTransactions.add(t);
       registrationQueue.add(t);
       selector.wakeup();
     } catch (IOException e) {
-      silentCloseChannel(udpChannel);
+      silentCloseChannel(channel);
       f.completeExceptionally(e);
     } catch (Throwable e) {
       // Make sure to close the channel, no matter what, but only handle the declared IOException
-      silentCloseChannel(udpChannel);
+      silentCloseChannel(channel);
       throw e;
     }
 

diff --git a/src/main/java/org/xbill/DNS/SimpleResolver.java b/src/main/java/org/xbill/DNS/SimpleResolver.java
@@ -7,7 +7,6 @@
 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;
@@ -20,8 +19,6 @@
 import lombok.extern.slf4j.Slf4j;
 import org.xbill.DNS.io.DefaultIoClientFactory;
 import org.xbill.DNS.io.IoClientFactory;
-import org.xbill.DNS.io.TcpIoClient;
-import org.xbill.DNS.io.UdpIoClient;
 
 /**
  * An implementation of Resolver that sends one query to one server. SimpleResolver handles TCP
@@ -43,7 +40,6 @@ 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);
@@ -103,11 +99,6 @@ 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");
@@ -396,25 +387,16 @@ CompletableFuture<Message> sendAsync(Message query, boolean forceTcp, Executor e
     }
 
     CompletableFuture<byte[]> result;
-    SocketChannel c = null;
     if (tcp) {
-      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);
-      }
+      result =
+        ioClientFactory
+          .createOrGetTcpClient()
+          .sendAndReceiveTcp(localAddress, address, query, out, timeoutValue);
     } else {
-      UdpIoClient udpClient = ioClientFactory.createOrGetUdpClient();
-      if (proxy != null) {
-        localAddress = proxy.getLocalAddress();
-        address = proxy.getRemoteAddress();
-        result = udpClient.sendAndReceiveUdp(localAddress, address, proxy, query, out, udpSize, timeoutValue);
-      } else {
-        result = udpClient.sendAndReceiveUdp(localAddress, address, query, out, udpSize, timeoutValue);
-      }
+      result =
+        ioClientFactory
+          .createOrGetUdpClient()
+          .sendAndReceiveUdp(localAddress, address, query, out, udpSize, timeoutValue);
     }
 
     return result.thenComposeAsync(

diff --git a/src/main/java/org/xbill/DNS/io/TcpIoClient.java b/src/main/java/org/xbill/DNS/io/TcpIoClient.java
@@ -6,7 +6,6 @@
 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
@@ -33,12 +32,4 @@ CompletableFuture<byte[]> sendAndReceiveTcp(
       Message query,
       byte[] data,
       Duration timeout);
-
-  CompletableFuture<byte[]> sendAndReceiveTcp(
-    InetSocketAddress local,
-    InetSocketAddress remote,
-    Socks5Proxy proxy,
-    Message query,
-    byte[] data,
-    Duration timeout);
 }
diff --git a/src/main/java/org/xbill/DNS/io/UdpIoClient.java b/src/main/java/org/xbill/DNS/io/UdpIoClient.java
@@ -6,7 +6,6 @@
 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
@@ -35,13 +34,4 @@ CompletableFuture<byte[]> sendAndReceiveUdp(
       byte[] data,
       int max,
       Duration timeout);
-
-  CompletableFuture<byte[]> sendAndReceiveUdp(
-    InetSocketAddress local,
-    InetSocketAddress remote,
-    Socks5Proxy proxy,
-    Message query,
-    byte[] data,
-    int max,
-    Duration timeout);
 }