test/jdk/com/sun/jndi/dns/lib/DNSServer.java

@@ -24,10 +24,12 @@
 import sun.security.util.HexDumpEncoder;
 
 import java.io.IOException;
 import java.net.DatagramPacket;
 import java.net.DatagramSocket;
+import java.net.InetAddress;
+import java.net.SocketException;
 import java.nio.ByteBuffer;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;

@@ -47,11 +49,11 @@
  * xxxx: 00 11 22 33 44 55 66 77   88 99 aa bb cc dd ee ff  ................
  *
  * Typically, DNS protocol exchange is generated by DNSTracer who captures
  * communication messages between DNS application program and real DNS server
  */
-public class DNSServer implements Runnable {
+public class DNSServer extends Thread implements Server {
 
     public class Pair<F, S> {
         private F first;
         private S second;
 

@@ -85,23 +87,25 @@
     private DatagramSocket socket;
     private String filename;
     private boolean loop;
     private final List<Pair<byte[], byte[]>> cache = new ArrayList<>();
     private ByteBuffer reqBuffer = ByteBuffer.allocate(DNS_PACKET_SIZE);
+    private volatile boolean isRunning;
 
-    public DNSServer(DatagramSocket socket, String filename) {
-        this(socket, filename, false);
+    public DNSServer(String filename) throws SocketException {
+        this(filename, false);
     }
 
-    public DNSServer(DatagramSocket socket, String filename, boolean loop) {
-        this.socket = socket;
+    public DNSServer(String filename, boolean loop) throws SocketException {
+        this.socket = new DatagramSocket(0, InetAddress.getLoopbackAddress());
         this.filename = filename;
         this.loop = loop;
     }
 
     public void run() {
         try {
+            isRunning = true;
             System.out.println(
                     "DNSServer: Loading DNS cache data from : " + filename);
             loadCaptureFile(filename);
 
             System.out.println(

@@ -110,47 +114,85 @@
             System.out.println("DNSServer: loop playback: " + loop);
 
             int playbackIndex = 0;
 
             while (playbackIndex < cache.size()) {
-                DatagramPacket reqPacket = new DatagramPacket(reqBuffer.array(),
-                        reqBuffer.array().length);
-                socket.receive(reqPacket);
-
-                System.out.println(
-                        "DNSServer: received query message from " + reqPacket
-                                .getSocketAddress());
+                DatagramPacket reqPacket = receiveQuery();
 
                 if (!verifyRequestMsg(reqPacket, playbackIndex)) {
+                    if (playbackIndex > 0 && verifyRequestMsg(reqPacket,
+                            playbackIndex - 1)) {
+                        System.out.println(
+                                "DNSServer: received retry query, resend");
+                        playbackIndex--;
+                    } else {
                     throw new RuntimeException(
                             "DNSServer: Error: Failed to verify DNS request. "
                                     + "Not identical request message : \n"
                                     + encoder.encodeBuffer(
                                     Arrays.copyOf(reqPacket.getData(),
                                             reqPacket.getLength())));
                 }
+                }
 
-                byte[] payload = generateResponsePayload(reqPacket,
-                        playbackIndex);
-                socket.send(new DatagramPacket(payload, payload.length,
-                        reqPacket.getSocketAddress()));
-                System.out.println(
-                        "DNSServer: send response message to " + reqPacket
-                                .getSocketAddress());
+                sendResponse(reqPacket, playbackIndex);
 
                 playbackIndex++;
                 if (loop && playbackIndex >= cache.size()) {
                     playbackIndex = 0;
                 }
             }
 
             System.out.println(
                     "DNSServer: Done for all cached messages playback");
+
+            System.out.println(
+                    "DNSServer: Still listening for possible retry query");
+            while (true) {
+                DatagramPacket reqPacket = receiveQuery();
+
+                // here we only handle the retry query for last one
+                if (!verifyRequestMsg(reqPacket, playbackIndex - 1)) {
+                    throw new RuntimeException(
+                            "DNSServer: Error: Failed to verify DNS request. "
+                                    + "Not identical request message : \n"
+                                    + encoder.encodeBuffer(
+                                    Arrays.copyOf(reqPacket.getData(),
+                                            reqPacket.getLength())));
+                }
+
+                sendResponse(reqPacket, playbackIndex - 1);
+            }
         } catch (Exception e) {
+            if (isRunning) {
             System.err.println("DNSServer: Error: " + e);
+                e.printStackTrace();
+            } else {
+                System.out.println("DNSServer: Exit");
         }
     }
+    }
+
+    private DatagramPacket receiveQuery() throws IOException {
+        DatagramPacket reqPacket = new DatagramPacket(reqBuffer.array(),
+                reqBuffer.array().length);
+        socket.receive(reqPacket);
+
+        System.out.println("DNSServer: received query message from " + reqPacket
+                .getSocketAddress());
+
+        return reqPacket;
+    }
+
+    private void sendResponse(DatagramPacket reqPacket, int playbackIndex)
+            throws IOException {
+        byte[] payload = generateResponsePayload(reqPacket, playbackIndex);
+        socket.send(new DatagramPacket(payload, payload.length,
+                reqPacket.getSocketAddress()));
+        System.out.println("DNSServer: send response message to " + reqPacket
+                .getSocketAddress());
+    }
 
     /*
      * Load a capture file containing an DNS protocol exchange in the
      * hexadecimal dump format emitted by sun.misc.HexDumpEncoder:
      *

@@ -297,6 +339,25 @@
         if ('a' <= ch && ch <= 'f') {
             return ch - 'a' + 10;
         }
         return -1;
     }
+
+    @Override public void stopServer() {
+        isRunning = false;
+        if (socket != null) {
+            try {
+                socket.close();
+            } catch (Exception e) {
+                // ignore
+            }
+        }
+    }
+
+    @Override public int getPort() {
+        if (socket != null) {
+            return socket.getLocalPort();
+        } else {
+            return -1;
+        }
+    }
 }