CVE-2017-7670

@@ -25,7 +25,11 @@
 import org.xbill.DNS.WireParseException;
 
 import java.net.InetAddress;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public abstract class AbstractProtocol implements Protocol {
     private static final Logger ACCESS = Logger.getLogger("com.comcast.cdn.traffic_control.traffic_router.core.access");
@@ -34,6 +38,7 @@ public abstract class AbstractProtocol implements Protocol {
     protected boolean shutdownRequested;
     private ExecutorService executorService;
     private NameServer nameServer;
+    private int taskTimeout = 5000; // default
 
     /**
      * Gets executorService.
@@ -138,9 +143,18 @@ protected byte[] query(final InetAddress client, final byte[] request) throws Wi
      * @param job
      *            the handler to be executed
      */
-    protected void submit(final Runnable job) {
-        executorService.submit(job);
-    }
+	protected void submit(final Runnable job) {
+		final Future<?> handler = executorService.submit(job);
+		executorService.submit(new Runnable() {
+			public void run() {
+				try {
+					handler.get(getTaskTimeout(), TimeUnit.MILLISECONDS);
+				} catch (InterruptedException | ExecutionException | TimeoutException e) {
+					handler.cancel(true);
+				}
+			}
+		});
+	}
 
     private Message createServerFail(final Message query) {
         final Message response = new Message();
@@ -155,4 +169,12 @@ private Message createServerFail(final Message query) {
         response.getHeader().setRcode(Rcode.SERVFAIL);
         return response;
     }
+
+	public int getTaskTimeout() {
+		return taskTimeout;
+	}
+
+	public void setTaskTimeout(final int taskTimeout) {
+		this.taskTimeout = taskTimeout;
+	}
 }

@@ -18,16 +18,19 @@
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.InetAddress;
 import java.net.ServerSocket;
 import java.net.Socket;
+import java.nio.channels.Channels;
 
 import org.apache.log4j.Logger;
 import org.xbill.DNS.Message;
 import org.xbill.DNS.WireParseException;
 
 public class TCP extends AbstractProtocol {
     private static final Logger LOGGER = Logger.getLogger(TCP.class);
+    private int readTimeout = 3000; // default
 
     private ServerSocket serverSocket;
 
@@ -87,11 +90,12 @@ class TCPSocketHandler implements Runnable {
         @SuppressWarnings("PMD.EmptyCatchBlock")
         public void run() {
             try {
+                socket.setSoTimeout(getReadTimeout());
                 final InetAddress client = socket.getInetAddress();
 
-                final DataInputStream is = new DataInputStream(socket.getInputStream());
+                final InputStream iis = Channels.newInputStream(Channels.newChannel(socket.getInputStream()));
+                final DataInputStream is = new DataInputStream(iis);
                 final DataOutputStream os = new DataOutputStream(socket.getOutputStream());
-
                 final int length = is.readUnsignedShort();
                 final byte[] request = new byte[length];
                 is.readFully(request);
@@ -121,5 +125,16 @@ public void shutdown() {
 		} catch (IOException e) {
 			LOGGER.warn("error on shutdown", e);
 		}
-    }
+	}
+
+
+	public int getReadTimeout() {
+		return readTimeout;
+	}
+
+	public void setReadTimeout(final int readTimeout) {
+		this.readTimeout = readTimeout;
+	}
+
+
 }

@@ -242,7 +242,12 @@
 		<property name="trafficRouterManager" ref="trafficRouterManager" />
 	</bean>
 
-	<bean id="RequestHandlerExecutorService" class="java.util.concurrent.Executors"
+	<bean id="UDPRequestHandlerExecutorService" class="java.util.concurrent.Executors"
+		factory-method="newFixedThreadPool">
+		<constructor-arg index="0" value="$[dns.max-threads:1000]" />
+	</bean>
+
+	<bean id="TCPRequestHandlerExecutorService" class="java.util.concurrent.Executors"
 		factory-method="newFixedThreadPool">
 		<constructor-arg index="0" value="$[dns.max-threads:1000]" />
 	</bean>
@@ -254,8 +259,10 @@
 
 	<bean id="TCP" class="com.comcast.cdn.traffic_control.traffic_router.core.dns.protocol.TCP">
 		<property name="serverSocket" ref="ServerSocket" />
-		<property name="executorService" ref="RequestHandlerExecutorService" />
+		<property name="executorService" ref="TCPRequestHandlerExecutorService" />
 		<property name="nameServer" ref="NameServer" />
+		<property name="readTimeout" value="$[dns.tcp.timeout.read:3000]" />
+		<property name="taskTimeout" value="$[dns.tcp.timeout.task:5000]" />
 	</bean>
 
 	<bean id="DatagramSocket" class="java.net.DatagramSocket">
@@ -264,8 +271,9 @@
 
 	<bean id="UDP" class="com.comcast.cdn.traffic_control.traffic_router.core.dns.protocol.UDP">
 		<property name="datagramSocket" ref="DatagramSocket" />
-		<property name="executorService" ref="RequestHandlerExecutorService" />
+		<property name="executorService" ref="UDPRequestHandlerExecutorService" />
 		<property name="nameServer" ref="NameServer" />
+		<property name="taskTimeout" value="$[dns.udp.timeout.task:5000]" />
 	</bean>
 
 	<bean id="ProtocolExecutorService" class="java.util.concurrent.Executors"

@@ -19,7 +19,6 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
-
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -130,11 +129,12 @@ public void testTCPSocketHandlerBadMessage() throws Exception {
         final DataOutputStream dos = new DataOutputStream(baos);
         dos.writeShort(wireRequest.length);
         dos.write(wireRequest);
+        in = new ByteArrayInputStream(baos.toByteArray());
+        when(socket.getInputStream()).thenReturn(in);
 
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
 
         when(socket.getOutputStream()).thenReturn(out);
-
         handler.run();
         assertThat(out.toByteArray().length, equalTo(0));
     }
@@ -147,6 +147,13 @@ public void testTCPSocketHandlerQueryFail() throws Exception {
         final Record question = Record.newRecord(name, Type.A, DClass.IN);
         final Message request = Message.newQuery(question);
 
+        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        final DataOutputStream dos = new DataOutputStream(baos);
+        dos.writeShort(request.toWire().length);
+        dos.write(request.toWire());
+        in = new ByteArrayInputStream(baos.toByteArray());
+        when(socket.getInputStream()).thenReturn(in);
+
         final Message response = new Message();
         response.setHeader(request.getHeader());