dnsruby 1.57.0 → 1.58.0

data/lib/dnsruby/resource/GPOS.rb

@@ -31,14 +31,20 @@ module Dnsruby
       #     latitude:   '20.0',
       #     altitude:   '30.0',
       #  }
-      def self.from_hash(gpos_params_hash)
+      #
+      # Since the type is assumed to be GPOS, it will be assigned
+      # automatially, and any other value will be overwritten.
+      # Therefore, having it present in the hash is not necessary.
+
+      def self.new_from_hash(gpos_params_hash)
+        gpos_params_hash[:type] = Types::GPOS
         RR.new_from_hash(gpos_params_hash)
       end
 
 
       # Create an instance from a string containing parameters, e.g.:
       # 'a.dnsruby.com.  10800  IN  GPOS  10.0  20.0  30.0'
-      def self.from_string(gpos_params_string)
+      def self.new_from_string(gpos_params_string)
         RR.new_from_string(gpos_params_string)
       end
 
@@ -49,7 +55,7 @@ module Dnsruby
       #   [EXAMPLE_HOSTNAME, Types::GPOS, Classes::IN, EXAMPLE_TTL, rdata.length, rdata, 0]
       # end
       # self.from_data(*EXAMPLE_GPOS_DATA)
-      def self.from_data(*gpos_params_data)
+      def self.new_from_data(*gpos_params_data)
         RR.new_from_data(*gpos_params_data)
       end
 

data/lib/dnsruby/resource/HIP.rb

@@ -126,7 +126,7 @@ module Dnsruby
         public_key = msg.get_bytes(pk_length)
         rsvs = []
         #  Load in the RSV names, if there are any
-        while (msg.has_remaining)
+        while (msg.has_remaining?)
           name = msg.get_name
           rsvs.push(name)
         end

data/lib/dnsruby/resource/IN.rb

@@ -52,7 +52,9 @@ module Dnsruby
       Types::SSHFP => SSHFP,
       Types::IPSECKEY => IPSECKEY,
       Types::HIP => HIP,
-      Types::DHCID => DHCID
+      Types::DHCID => DHCID,
+      Types::GPOS => GPOS,
+      Types::NXT => NXT
     } #:nodoc: all
 
     #  module IN contains ARPA Internet specific RRs

data/lib/dnsruby/resource/NAPTR.rb

@@ -30,7 +30,7 @@ module Dnsruby
       #  The NAPTR RR service field
       attr_accessor :service
       #  The NAPTR RR regexp field
-      attr_accessor :regexp
+      attr_reader :regexp
       #  The NAPTR RR replacement field
       attr_accessor :replacement
 
@@ -95,4 +95,4 @@ module Dnsruby
       end
     end
   end
-end
+end

data/lib/dnsruby/resource/NSEC3.rb

@@ -242,7 +242,7 @@ module Dnsruby
 
       def decode_next_hashed(input)
         @next_hashed = NSEC3.decode_next_hashed(input)
-        end
+      end
 
       def NSEC3.decode_next_hashed(input)
         return Base32.decode32hex(input)
@@ -329,4 +329,4 @@ module Dnsruby
       end
     end
   end
-end
+end

data/lib/dnsruby/resource/NXT.rb

@@ -0,0 +1,302 @@
+require_relative = ->(*args) do
+  this_file_dir = File.expand_path(File.dirname(__FILE__))
+  args.each { |arg| require(File.join(this_file_dir, arg)) }
+end
+
+require_relative.('../bitmap', '../bit_mapping', 'RR')
+
+module Dnsruby
+class RR
+
+# Class for NXT resource records.
+#
+# NXT-specific data types, present in RDATA, are:
+#   next_domain: the next domain name, as a Name instance
+#   types: array of record types as numbers
+#
+# RFC 2535 (https://www.ietf.org/rfc/rfc2535.txt)
+#
+# The RFC mentions that a low bit of zero in the type RDATA
+# indicates that the highest type code does not exceed 127,
+# and that a low bit of 1 indicates that some mechanism
+# other than a bitmap is being used.  This class does not
+# support such non-bitmap mechanisms, and assumes there
+# will always be a bitmap.
+class NXT < RR
+
+  ClassHash[[TypeValue = Types::NXT, Classes::IN]] = self #:nodoc: all
+
+  attr_accessor :next_domain, :types
+
+  REQUIRED_KEYS = [:next_domain, :types]
+
+  def from_hash(params_hash)
+    unless REQUIRED_KEYS.all? { |key| params_hash[key] }
+      raise ArgumentError.new("NXT hash must contain all of: #{REQUIRED_KEYS.join(', ')}.")
+    end
+    @next_domain = Name.create(params_hash[:next_domain]) unless @next_domain.is_a?(Name)
+    @types       = params_hash[:types]
+  end
+
+  def from_data(data)
+    next_domain, types = data
+    from_hash(next_domain: next_domain, types: types)
+  end
+
+  def from_string(string)
+    next_domain, *type_names = string.split  # type names are all but first
+    types = NxtTypes::names_to_codes(type_names)
+    from_hash(next_domain: next_domain, types: types)
+  end
+
+  # As with all resource record subclasses of RR, this class cannot be
+  # directly instantiated, but instead must be instantiated via use of
+  # one of the RR class methods.  These NXT class methods are wrappers
+  # around those RR methods, so that there is an interface on the NXT
+  # class for creating NXT instances.
+
+  # Create an instance from a hash of parameters, e.g.:
+  #
+  # rr = RR::NXT.new_from_hash(
+  #     name: 'b.dnsruby.com.',
+  #     ttl: 10800,
+  #     klass: Classes::IN,
+  #     next_domain: 'a.dnsruby.com.',
+  #     types: [Types::SOA, Types::NXT])
+  #
+  # Since the type is assumed to be NXT, it will be assigned
+  # automatically, and any other value will be overwritten.
+  # Therefore, having it present in the hash is not necessary.
+  def self.new_from_hash(params_hash)
+    params_hash[:type] = Types::NXT
+    RR.new_from_hash(params_hash)
+  end
+
+  # Create an instance from a string containing parameters, e.g.:
+  # b.dnsruby.com.  10800  IN  NXT  A.dnsruby.com.  SOA  NXT
+  def self.new_from_string(params_string)
+    RR.new_from_string(params_string)
+  end
+
+  # Create an instance from an ordered parameter list, e.g.:
+  # rdata = RR::NXT.build_rdata('a.dnsruby.com.', [Types::SOA, Types::NXT])
+  #
+  # rr = RR::NXT.new_from_data('b.dnsruby.com.', Types::NXT,
+  #     Classes::IN, 10800, rdata.size, rdata, 0)
+  def self.new_from_data(*params_data)
+    RR.new_from_data(*params_data)
+  end
+
+  # Builds rdata from the provided information.
+  # @param next_domain either a string or a Name
+  # @param types an array of types (where each type is the numeric type code)
+  #        or a TypeBitmap
+  def self.build_rdata(next_domain, types)
+    next_domain = Name.create(next_domain) if next_domain.is_a?(String)
+    types = TypeBitmap.from_type_codes(types) if types.is_a?(Array)
+
+    binary_string = ''.force_encoding('ASCII-8BIT')
+    binary_string << next_domain.canonical
+    binary_string << BitMapping.reverse_binary_string_bits(types.to_binary_string)
+    binary_string
+  end
+
+  # From the RFC:
+  # NXT has the following format:
+  # foo.nil.    NXT big.foo.nil NS KEY SOA NXT
+  # <owner> NXT <next_domain> <record types>
+  #
+  # We handle the rdata, the RR superclass does the rest.
+  def rdata_to_string
+    "#{next_domain} #{NxtTypes.codes_to_names(types).join(' ')}"
+  end
+
+  def encode_rdata(message_encoder, _canonical)
+    message_encoder.put_bytes(build_rdata)
+  end
+
+  def build_rdata
+    self.class.build_rdata(next_domain, types)
+  end
+
+  def self.decode_rdata(message_decoder)
+
+    start_index = message_decoder.index
+
+    rdata_len = -> do
+      rdata_length_str = message_decoder.data[start_index - 2, 2]
+      rdata_length_str.unpack('n').first
+    end
+
+    next_domain_and_bitmap = -> do
+      next_domain = message_decoder.get_name
+      bitmap_start_index = message_decoder.index
+
+      # If we're being called from new_from_data, the MessageDecoder
+      # contains only the rdata, not the entire message, and there will
+      # be no encoded length for us to read.
+      called_from_new_from_data = (start_index == 0)
+      bitmap_length = called_from_new_from_data \
+          ? message_decoder.data.size \
+          : rdata_len.() - (bitmap_start_index - start_index)
+
+      bitmap = message_decoder.get_bytes(bitmap_length)
+      bitmap = BitMapping.reverse_binary_string_bits(bitmap)
+      [next_domain, bitmap]
+    end
+
+    next_domain, type_bitmap = next_domain_and_bitmap.()
+    types = TypeBitmap.from_binary_string(type_bitmap).to_type_array
+    new(next_domain: next_domain, types: types)
+  end
+
+  # 'name' is used in the RR superclass, but 'owner' is the term referred to
+  # in the RFC, so we'll make owner an alias for name.
+  alias_method(:owner,  :name)
+  alias_method(:owner=, :name=)
+
+
+  # Methods used to manipulate the storage and representation of
+  # record types as stored in NXT record bitmaps.
+  module NxtTypes
+
+    module_function
+
+    # Maximum bitmap size is 128 bytes; since it's zero offset
+    # values are 0..(2 ** 128 - 1).  However, the least
+    # significant bit must not be set, so the maximum is 1 less than that.
+    MAX_BITMAP_NUMBER_VALUE = (2 ** 128) - 1 - 1
+
+    # Convert a numeric type code to its corresponding name (e.g. "A" => 1).
+    # Unknown types are named "TYPE#{number}".
+    def code_to_name(number)
+      Types.to_string(number) || "TYPE#{number}"
+    end
+
+    # Convert a type name to its corresponding numeric type code.
+    # Names matching /^TYPE(\d+)$/ are assumed to have a code
+    # corresponding to the numeric value of the substring following 'TYPE'.
+    def name_to_code(name)
+      code = Types.to_code(name)
+      if code.nil?
+        matches = /^TYPE(\d+)$/.match(name)
+        code = matches[1].to_i if matches
+      end
+      code
+    end
+
+    # For a given array of type names, return an array of codes.
+    def names_to_codes(names)
+      names.map { |s| name_to_code(s) }
+    end
+
+    # For the specified string containing names (e.g. 'A NS'),
+    # return an array containing the corresponding codes.
+    def names_string_to_codes(name_string)
+      names_to_codes(name_string.split(' '))
+    end
+
+    # For the given array of type codes, return an array of their
+    # corresponding names.
+    def codes_to_names(codes)
+      codes.map { |code| code_to_name(code) }
+    end
+
+    # Generate a string containing the names corresponding to the
+    # numeric type codes.  Sort it by the numeric type code, ascending.
+    def codes_to_string(codes)
+      codes.sort.map { |code| code_to_name(code) }.join(' ')
+    end
+
+    # From a binary string of type code bits, return an array
+    # of type codes.
+    def binary_string_to_codes(binary_string)
+      bitmap_number = BitMapping.binary_string_to_number(binary_string)
+      assert_legal_bitmap_value(bitmap_number)
+      BitMapping.number_to_set_bit_positions_array(bitmap_number)
+    end
+
+    # From a binary string of type code bits, return an array
+    # of type names.
+    def binary_string_to_names(binary_string)
+      codes = binary_string_to_codes(binary_string)
+      codes_to_names(codes)
+    end
+
+    # From an array of type codes, return a binary string.
+    def codes_to_binary_string(codes)
+      codes = codes.sort
+      unless legal_code_value?(codes.first) && legal_code_value?(codes.last)
+        raise ArgumentError.new("All codes must be between 1 and 127: #{codes.inspect}.")
+      end
+      bitmap_number = BitMapping.set_bit_position_array_to_number(codes)
+      BitMapping.number_to_binary_string(bitmap_number)
+    end
+
+    # Assert that the specified number is a legal value with which to
+    # instantiate a NXT type bitmap.  Raise on error, do nothing on success.
+    def assert_legal_bitmap_value(number)
+      max_value = NxtTypes::MAX_BITMAP_NUMBER_VALUE
+      if number > max_value
+        raise ArgumentError.new("Bitmap maximum value is #{max_value} (0x#{max_value.to_s(16)}).")
+      end
+      if number & 1 == 1
+        raise ArgumentError.new("Bitmap number must not have low bit set.")
+      end
+    end
+
+    def legal_code_value?(code)
+      (1..127).include?(code)
+    end
+  end
+
+
+  class TypeBitmap
+
+    attr_accessor :bitmap
+
+    # Create an instance from a string containing type names separated by spaces
+    # e.g. "A TXT NXT"
+    def self.from_names_string(names_string)
+      type_codes = BitMapping.names_string_to_codes(names_string)
+      from_type_codes(type_codes)
+    end
+
+    # Create an instance from type numeric codes (e.g. 30 for NXT).
+    def self.from_type_codes(type_codes)
+      new(BitMapping.set_bit_position_array_to_number(type_codes))
+    end
+
+    # Create an instance from a binary string, e.g. from a NXT record RDATA:
+    def self.from_binary_string(binary_string)
+      new(BitMapping.binary_string_to_number(binary_string))
+    end
+
+    # The constructor is made private so that the name of the method called
+    # to create the instance reveals to the reader the type of the initial data.
+    private_class_method :new
+    def initialize(bitmap_number)
+      NxtTypes.assert_legal_bitmap_value(bitmap_number)
+      @bitmap = Bitmap.from_number(bitmap_number)
+    end
+
+    # Returns a binary string representing this data, in as few bytes as possible
+    # (i.e. no leading zero bytes).
+    def to_binary_string
+      bitmap.to_binary_string
+    end
+
+    # Returns the instance's data as an array of type codes.
+    def to_type_array
+      bitmap.to_set_bit_position_array
+    end
+
+    # Output types in dig format, e.g. "A AAAA NXT"
+    def to_s
+      type_codes = bitmap.to_set_bit_position_array
+      NxtTypes.codes_to_string(type_codes)
+    end
+  end
+end
+end
+end

data/lib/dnsruby/resource/OPT.rb

@@ -248,9 +248,9 @@ module Dnsruby
       end
 
       def self.decode_rdata(msg)#:nodoc: all
-        if (msg.has_remaining)
+        if (msg.has_remaining?)
           options = []
-          while (msg.has_remaining) do
+          while (msg.has_remaining?) do
             code  = msg.get_unpack('n')[0]
             len = msg.get_unpack('n')[0]
             data = msg.get_bytes(len)

data/lib/dnsruby/resource/TXT.rb

@@ -15,7 +15,7 @@
 # ++
 begin
   require 'jcode'
-rescue LoadError => e
+rescue LoadError => _e
 end
 module Dnsruby
   class RR
@@ -189,4 +189,4 @@ module Dnsruby
       end
     end
   end
-end
+end

data/lib/dnsruby/resource/generic.rb

@@ -164,3 +164,4 @@ require 'dnsruby/resource/HIP'
 require 'dnsruby/resource/KX'
 require 'dnsruby/resource/DHCID'
 require 'dnsruby/resource/GPOS'
+require 'dnsruby/resource/NXT'

data/lib/dnsruby/select_thread.rb

@@ -20,6 +20,7 @@ begin
 rescue LoadError
   require 'thread'
 end
+require 'set'
 require 'singleton'
 require 'dnsruby/validator_thread.rb'
 module Dnsruby
@@ -46,13 +47,15 @@ module Dnsruby
       @@mutex.synchronize {
         @@in_select=false
         #         @@notifier,@@notified=IO.pipe
-        @@sockets = [] # @@notified]
+        @@sockets = Set.new
         @@timeouts = Hash.new
         #     @@mutex.synchronize do
         @@query_hash = Hash.new
         @@socket_hash = Hash.new
+        @@socket_is_persistent = Hash.new
         @@observers = Hash.new
         @@tcp_buffers=Hash.new
+        @@socket_remaining_queries = Hash.new
         @@tick_observers = []
         @@queued_exceptions=[]
         @@queued_responses=[]
@@ -64,9 +67,8 @@ module Dnsruby
         BasicSocket.do_not_reverse_lookup = true
         #     end
         #  Now start the select thread
-        @@select_thread = Thread.new {
-          do_select
-        }
+        @@select_thread = Thread.new { do_select }
+
         #         # Start the validator thread
         #         @@validator = ValidatorThread.instance
       }
@@ -93,7 +95,7 @@ module Dnsruby
     class QuerySettings
       attr_accessor :query_bytes, :query, :ignore_truncation, :client_queue,
         :client_query_id, :socket, :dest_server, :dest_port, :endtime, :udp_packet_size,
-        :single_resolver
+        :single_resolver, :is_persistent_socket, :tcp_pipelining_max_queries
       #  new(query_bytes, query, ignore_truncation, client_queue, client_query_id,
       #      socket, dest_server, dest_port, endtime, , udp_packet_size, single_resolver)
       def initialize(*args)
@@ -108,9 +110,21 @@ module Dnsruby
         @endtime = args[8]
         @udp_packet_size = args[9]
         @single_resolver = args[10]
+        @is_persistent_socket = false
+        @tcp_pipelining_max_queries = nil
       end
     end
 
+    def tcp?(socket)
+      type = socket.getsockopt(Socket::SOL_SOCKET, Socket::SO_TYPE)
+      [Socket::SOCK_STREAM].pack("i") == type.data
+    end
+
+    def udp?(socket)
+      type = socket.getsockopt(Socket::SOL_SOCKET, Socket::SO_TYPE)
+      [Socket::SOCK_DGRAM].pack("i") == type.data
+    end
+
     def add_to_select(query_settings)
       #  Add the query to sockets, and then wake the select thread up
       @@mutex.synchronize {
@@ -118,9 +132,12 @@ module Dnsruby
         #  @TODO@ This assumes that all client_query_ids are unique!
         #  Would be a good idea at least to check this...
         @@query_hash[query_settings.client_query_id]=query_settings
-        @@socket_hash[query_settings.socket]=[query_settings.client_query_id] # @todo@ If we use persistent sockets then we need to update this array
+        @@socket_hash[query_settings.socket] ||= []
+        @@socket_hash[query_settings.socket] << query_settings.client_query_id
+        @@socket_remaining_queries[query_settings.socket] ||= query_settings.tcp_pipelining_max_queries if query_settings.tcp_pipelining_max_queries != :infinite
         @@timeouts[query_settings.client_query_id]=query_settings.endtime
-        @@sockets.push(query_settings.socket)
+        @@sockets << query_settings.socket
+        @@socket_is_persistent[query_settings.socket] = query_settings.is_persistent_socket
       }
       begin
         @@wakeup_sockets[0].send("wakeup!", 0)
@@ -158,14 +175,7 @@ module Dnsruby
         send_queued_responses
         send_queued_validation_responses
         timeout = tick_time = 0.1 # We provide a timer service to various Dnsruby classes
-        sockets=[]
-        timeouts=[]
-        has_observer = false
-        @@mutex.synchronize {
-          sockets = @@sockets
-          timeouts = @@timeouts.values
-          has_observer = !@@observers.empty?
-        }
+        sockets, timeouts, has_observer = @@mutex.synchronize { [@@sockets.to_a, @@timeouts.values, !@@observers.empty?] }
         if (timeouts.length > 0)
           timeouts.sort!
           timeout = timeouts[0] - Time.now
@@ -185,8 +195,11 @@ module Dnsruby
         rescue SelectWakeup
           #  If SelectWakeup, then just restart this loop - the select call will be made with the new data
           next
-        rescue IOError => e# Don't worry if the socket was closed already
+        rescue IOError => e
           #           print "IO Error  =: #{e}\n"
+          exceptions = clean_up_closed_sockets
+          exceptions.each { |exception| send_exception_to_client(*exception) }
+
           next
         end
         if ready && ready.include?(@@wakeup_sockets[1])
@@ -201,72 +214,117 @@ module Dnsruby
           end
         end
         if (ready == nil)
-          #  proces the timeouts
+          #  process the timeouts
           process_timeouts
           unused_loop_count+=1
         else
           process_ready(ready)
           unused_loop_count=0
-          #                   process_error(errors)
+          # process_error(errors)
         end
-        @@mutex.synchronize{
+        @@mutex.synchronize do
           if (unused_loop_count > 10 && @@query_hash.empty? && @@observers.empty?)
-            Dnsruby.log.debug{"Stopping select loop"}
-            return
+            Dnsruby.log.debug("Try stop select loop")
+
+            non_persistent_sockets = @@sockets.select { |s| ! @@socket_is_persistent[s] }
+            non_persistent_sockets.each do |socket|
+              socket.close rescue nil
+              @@sockets.delete(socket)
+            end
+
+            Dnsruby.log.debug("Deleted #{non_persistent_sockets.size} non-persistent sockets," +
+                              " #{@@sockets.count} persistent sockets remain.")
+            @@socket_hash.clear
+
+            if @@sockets.empty?
+              Dnsruby.log.debug("Stopping select loop")
+              return
+            end
           end
-        }
+        end
         #               }
       end
     end
 
+    # Removes closed sockets from @@sockets, and returns an array containing 1
+    # exception for each closed socket contained in @@socket_hash.
+    def clean_up_closed_sockets
+      exceptions = @@mutex.synchronize do
+        closed_sockets_in_hash = @@sockets.select(&:closed?).select { |s| @@socket_hash[s] }
+        @@sockets.delete_if { | socket | socket.closed? }
+        closed_sockets_in_hash.each_with_object([]) do |socket, exceptions|
+          @@socket_hash[socket].each do | client_id |
+            exceptions << [SocketEofResolvError.new("TCP socket closed before all answers received"), socket, client_id]
+          end
+        end
+      end
+    end
+
     def process_error(errors)
       Dnsruby.log.debug{"Error! #{errors.inspect}"}
       #  @todo@ Process errors [can we do this in single socket environment?]
     end
 
+    def get_active_ids(queries, id)
+      queries.keys.select { |client_query_id| client_query_id[1].header.id == id }
+    end
+
     #         @@query_hash[query_settings.client_query_id]=query_settings
-    #         @@socket_hash[query_settings.socket]=[query_settings.client_query_id] # @todo@ If we use persistent sockets then we need to update this array
     def process_ready(ready)
-      ready.each do |socket|
-        query_settings = nil
-        @@mutex.synchronize{
-          #  Can do this if we have a query per socket, but not otherwise...
-          c_q_id = @@socket_hash[socket][0] # @todo@ If we use persistent sockets then this won't work
-          query_settings = @@query_hash[c_q_id]
-        }
+      persistent_sockets, nonpersistent_sockets = @@mutex.synchronize { ready.partition { |socket| persistent?(socket) } }
+
+      nonpersistent_sockets.each do |socket|
+        query_settings =  @@mutex.synchronize { @@query_hash[@@socket_hash[socket][0]] }
         next if !query_settings
+
         udp_packet_size = query_settings.udp_packet_size
         msg, bytes = get_incoming_data(socket, udp_packet_size)
-        if (msg!=nil)
-          #  Check that the IP we received from was the IP we sent to!
-          answerip = msg.answerip.downcase
-          answerfrom = msg.answerfrom.downcase
-          dest_server = query_settings.dest_server
-          answeripaddr = IPAddr.new(answerip)
-          dest_server = IPAddr.new("0.0.0.0")
-          begin
-            destserveripaddr = IPAddr.new(dest_server)
-          rescue ArgumentError
-            #  Host name not IP address
-          end
-          if (dest_server && (dest_server != '0.0.0.0') &&
-                (answeripaddr != destserveripaddr) &&
-                (answerfrom != dest_server))
-            Dnsruby.log.warn("Unsolicited response received from #{answerip} instead of #{query_settings.dest_server}")
-          else
-            send_response_to_client(msg, bytes, socket)
-          end
-        end
+
+        process_message(msg, bytes, socket) if msg
+
+        ready.delete(socket)
+      end
+
+      persistent_sockets.each do |socket|
+        msg, bytes = get_incoming_data(socket, 0)
+        process_message(msg, bytes, socket) if msg
         ready.delete(socket)
       end
     end
 
+    def process_message(msg, bytes, socket)
+      @@mutex.synchronize do
+        ids = get_active_ids(@@query_hash, msg.header.id)
+        return if ids.empty? # should be only one
+        query_settings = @@query_hash[ids[0]].clone
+      end
+
+      answerip = msg.answerip.downcase
+      answerfrom = msg.answerfrom.downcase
+      answeripaddr = IPAddr.new(answerip)
+      dest_server = IPAddr.new("0.0.0.0")
+
+      begin
+        destserveripaddr = IPAddr.new(dest_server)
+      rescue ArgumentError
+        #  Host name not IP address
+      end
+
+      if (dest_server && (dest_server != '0.0.0.0') &&
+          (answeripaddr != destserveripaddr) &&
+          (answerfrom != dest_server))
+        Dnsruby.log.warn("Unsolicited response received from #{answerip} instead of #{query_settings.dest_server}")
+      else
+        send_response_to_client(msg, bytes, socket)
+      end
+    end
+
     def send_response_to_client(msg, bytes, socket)
       #  Figure out which client_ids we were expecting on this socket, then see if any header ids match up
       #  @TODO@ Can get rid of this, as we only have one query per socket.
       client_ids=[]
       @@mutex.synchronize{
-        client_ids = @@socket_hash[socket]
+        client_ids = @@socket_hash[socket].clone
       }
       #  get the queries associated with them
       client_ids.each do |id|
@@ -284,7 +342,7 @@ module Dnsruby
             res = @@query_hash[id].single_resolver
             query = @@query_hash[id].query
           }
-          tcp = (socket.class == TCPSocket)
+          tcp = tcp?(socket)
           #  At this point, we should check if the response is OK
           if (ret = res.check_response(msg, bytes, query, client_queue, id, tcp))
             remove_id(id)
@@ -307,30 +365,48 @@ module Dnsruby
       print("Stray packet - " + msg.question()[0].qname.to_s + " from " + msg.answerip.to_s + ", #{client_ids.length} client_ids\n")
     end
 
+    def persistent?(socket)
+      @@socket_is_persistent[socket]
+    end
+
     def remove_id(id)
-      socket=nil
-      @@mutex.synchronize{
+
+      @@mutex.synchronize do
         socket = @@query_hash[id].socket
         @@timeouts.delete(id)
         @@query_hash.delete(id)
-        @@socket_hash.delete(socket)
-        @@sockets.delete(socket) # @TODO@ Not if persistent!
-      }
-      Dnsruby.log.debug{"Closing socket #{socket}"}
-      begin
-        socket.close # @TODO@ Not if persistent!
-      rescue IOError # Don't worry if the socket was closed already
+        @@socket_hash[socket].delete(id)
+
+        decrement_remaining_queries(socket) if persistent?(socket)
+
+        if !persistent?(socket) || max_attained?(socket)
+          @@sockets.delete(socket)
+          @@socket_hash.delete(socket)
+          Dnsruby.log.debug("Closing socket #{socket}")
+          socket.close rescue nil
+        end
       end
     end
 
+    def decrement_remaining_queries(socket)
+      if @@socket_remaining_queries[socket]
+        @@socket_remaining_queries[socket] -= 1
+      end
+    end
+
+    def max_attained?(socket)
+      remaining = @@socket_remaining_queries[socket]
+      attained = persistent?(socket) && remaining && remaining <= 0
+      Dnsruby.log.debug("Max queries per conn attained") if attained
+      attained
+    end
+
     def process_timeouts
+      # NOTE: It's @@timeouts we need to protect; after the clone we're ok
+      timeouts = @@mutex.synchronize { @@timeouts.clone }
       time_now = Time.now
-      timeouts={}
-      @@mutex.synchronize {
-        timeouts = @@timeouts
-      }
       timeouts.each do |client_id, timeout|
-        if (timeout < time_now)
+        if timeout < time_now
           send_exception_to_client(ResolvTimeout.new("Query timed out"), nil, client_id)
         end
       end
@@ -354,8 +430,19 @@ module Dnsruby
         begin
           input, = socket.recv_nonblock(expected_length-buf.length)
           if (input=="")
-            TheLog.info("Bad response from server - no bytes read - ignoring")
-            #  @TODO@ Should we do anything about this?
+            Dnsruby.log.debug("EOF from server - no bytes read - closing socket")
+            socket.close #EOF closed by server, if we were interrupted we need to resend
+
+            exceptions = @@mutex.synchronize do
+              @@sockets.delete(socket) #remove ourselves from select, app will have to retry
+              #maybe fire an event
+              @@socket_hash[socket].map do | client_id |
+                [SocketEofResolvError.new("TCP socket closed before all answers received"), socket, client_id]
+              end
+            end
+
+            exceptions.each { |exception| send_exception_to_client(*exception) }
+
             return false
           end
           buf += input
@@ -391,21 +478,10 @@ module Dnsruby
     def get_incoming_data(socket, packet_size)
       answerfrom,answerip,answerport,answersize=nil
       ans,buf = nil
-      begin
-        if (socket.class == TCPSocket)
-          #  @todo@ Ruby Bug #9061 stops this working right
-          #  We'd like to do a socket.recvfrom, but that raises an Exception
-          #  on Windows for TCPSocket for Ruby 1.8.5 (and 1.8.6).
-          #  So, we need to do something different for TCP than UDP. *sigh*
-          #  @TODO@ This workaround will only work if there is exactly one socket per query
-          #     - *not* ideal TCP use!
-          @@mutex.synchronize{
-            client_id = @@socket_hash[socket][0]
-            answerfrom = @@query_hash[client_id].dest_server
-            answerip = answerfrom
-            answerport = @@query_hash[client_id].dest_port
-          }
+      is_tcp = tcp?(socket)
 
+      begin
+        if is_tcp
           #  Call TCP read here - that will take care of reading the 2 byte length,
           #  and then the full packet - without blocking select.
           buf = tcp_read(socket)
@@ -438,8 +514,23 @@ module Dnsruby
 
       begin
         ans = Message.decode(buf)
+
+        if is_tcp
+          @@mutex.synchronize do
+            ids = get_active_ids(@@query_hash, ans.header.id)
+            if ids.empty?
+              Dnsruby.log.error("Decode error from #{answerip} but can't determine packet id")
+              #todo add error event? The problem is we don't have a valid id so we don't
+              #know which client queue to send the exception to
+            end
+            answerfrom = @@query_hash[ids[0]].dest_server
+            answerip = answerfrom
+            answerport = @@query_hash[ids[0]].dest_port
+          end
+        end
+
       rescue Exception => e
-        Dnsruby.log.error{"Decode error! #{e.class}, #{e}\nfor msg (length=#{buf.length}) : #{buf}"}
+        Dnsruby.log.error("Decode error! #{e.class}, #{e}\nfor msg (length=#{buf.length}) : #{buf}")
         client_id=get_client_id_from_answerfrom(socket, answerip, answerport)
         if (client_id == nil)
           Dnsruby.log.error{"Decode error from #{answerip} but can't determine packet id"}