Handle TCP Requester #recv_reply incomplete data

Author: jsaubry

lib/resolv.rb

@@ -720,7 +720,8 @@ def request(sender, tout)
           begin
             reply, from = recv_reply(select_result[0])
           rescue Errno::ECONNREFUSED, # GNU/Linux, FreeBSD
-                 Errno::ECONNRESET # Windows
+                 Errno::ECONNRESET, # Windows
+                 EOFError
             # No name server running on the server?
             # Don't wait anymore.
             raise ResolvTimeout
@@ -930,10 +931,10 @@ def initialize(host, port=Port)
 
         def recv_reply(readable_socks)
           len_data = readable_socks[0].read(2)
-          raise Errno::ECONNRESET if len_data.nil?
+          raise EOFError if len_data.nil? || len_data.bytesize != 2
           len = len_data.unpack('n')[0]
           reply = @socks[0].read(len)
-          raise Errno::ECONNRESET if reply.nil?
+          raise EOFError if reply.nil? || reply.bytesize != len
           return reply, nil
         end
 

test/resolv/test_dns.rb

@@ -881,4 +881,65 @@ def test_tcp_connection_closed_after_length
       client_thread.join
     end
   end
+
+  def test_tcp_connection_closed_with_partial_length_prefix
+    with_tcp('127.0.0.1', 0) do |t|
+      _, server_port, _, server_address = t.addr
+
+      server_thread = Thread.new do
+        ct = t.accept
+        ct.recv(512)
+        ct.write "A" # 1 byte
+        ct.close
+      end
+
+      client_thread = Thread.new do
+        requester = Resolv::DNS::Requester::TCP.new(server_address, server_port)
+        begin
+          msg = Resolv::DNS::Message.new
+          msg.add_question('example.org', Resolv::DNS::Resource::IN::A)
+          sender = requester.sender(msg, msg)
+          assert_raise(Resolv::ResolvTimeout) do
+            requester.request(sender, 2)
+          end
+        ensure
+          requester.close
+        end
+      end
+
+      server_thread.join
+      client_thread.join
+    end
+  end
+
+  def test_tcp_connection_closed_with_partial_message_body
+    with_tcp('127.0.0.1', 0) do |t|
+      _, server_port, _, server_address = t.addr
+
+      server_thread = Thread.new do
+        ct = t.accept
+        ct.recv(512)
+        ct.write([10].pack('n')) # length 10
+        ct.write "12345" # 5 bytes (partial)
+        ct.close
+      end
+
+      client_thread = Thread.new do
+        requester = Resolv::DNS::Requester::TCP.new(server_address, server_port)
+        begin
+          msg = Resolv::DNS::Message.new
+          msg.add_question('example.org', Resolv::DNS::Resource::IN::A)
+          sender = requester.sender(msg, msg)
+          assert_raise(Resolv::ResolvTimeout) do
+            requester.request(sender, 2)
+          end
+        ensure
+          requester.close
+        end
+      end
+
+      server_thread.join
+      client_thread.join
+    end
+  end
 end