dnsruby 1.74.0 → 2.0.0

data/lib/dnsruby/packet_sender.rb

@@ -15,6 +15,7 @@
 # ++
 require 'dnsruby/select_thread'
 require 'ipaddr'
+require 'securerandom'
 # require 'dnsruby/iana_ports'
 module Dnsruby
 
@@ -24,40 +25,41 @@ module Dnsruby
 
     @@authoritative_cache = Cache.new
     @@recursive_cache = Cache.new
+    PIPELINE_MUTEX_CREATION = Mutex.new
+    private_constant :PIPELINE_MUTEX_CREATION
 
 
-    def PacketSender.cache(query, response)
-      return if response.cached
-      #  ONLY cache the response if it is not an update response
+    def PacketSender.cache(query, response, scope: response.instance_variable_get(:@dnsruby_cache_scope))
+      return if response.cached || query.send_raw || scope == false || !query.do_caching || query.class == Update
       question = query.question()[0]
-      if (query.do_caching && (query.class != Update) &&
-          (question.qtype != Types::AXFR) && (question.qtype != Types::IXFR) &&
-          (response.rcode == RCode::NOERROR) &&(!response.tsig) &&
-          (query.class != Update) &&
-          (response.header.ancount > 0))
-        # # @TODO@ What about TSIG-signed responses?
-        #  Don't cache any packets with "*" in the query name! (RFC1034 sec 4.3.3)
-        if (!question.qname.to_s.include? "*")
-          #  Now cache response RRSets
-          if (query.header.rd)
-            PacketSender.cache_recursive(response);
-          else
-            PacketSender.cache_authoritative(response);
-          end
-        end
+      return unless question && question.qtype != Types::AXFR && question.qtype != Types::IXFR &&
+        response.rcode == RCode::NOERROR && response.header.ancount > 0
+      # TSIG authentication belongs to this transaction. Verification removes
+      # the response's TSIG record, so its absence does not make it cacheable.
+      return if query.tsig || query.signed? || response.signed? || response.tsig
+      # Alias-bearing NODATA belongs to the scoped negative cache, even
+      # though its CNAME records make the answer section nonempty.
+      return if Cache.negative_soa(response)
+      # Don't cache packets with "*" in the query name (RFC 1034 section 4.3.3).
+      return if question.qname.to_s.include?('*')
+      if query.header.rd
+        PacketSender.cache_recursive(response, scope: scope)
+      else
+        PacketSender.cache_authoritative(response, scope: scope)
       end
     end
 
-    def PacketSender.cache_authoritative(answer)
+    def PacketSender.cache_authoritative(answer, scope: nil)
       return if !answer.header.aa
-      @@authoritative_cache.add(answer)
+      @@authoritative_cache.add(answer, scope: scope)
     end
 
-    def PacketSender.cache_recursive(answer)
-      @@recursive_cache.add(answer)
+    def PacketSender.cache_recursive(answer, scope: nil)
+      @@recursive_cache.add(answer, scope: scope)
     end
 
     def PacketSender.clear_caches
+      Cache.invalidate_resolutions
       @@recursive_cache.clear
       @@authoritative_cache.clear
     end
@@ -86,6 +88,18 @@ module Dnsruby
     # Limit the number of queries per pipeline
     attr_accessor :tcp_pipelining_max_queries
 
+    # Maximum idle time for a reusable TCP connection, in seconds.
+    def tcp_idle_timeout
+      @tcp_idle_timeout || Resolver::DefaultTcpIdleTimeout
+    end
+
+    def tcp_idle_timeout=(seconds)
+      unless seconds.is_a?(Numeric) && seconds.finite? && seconds > 0
+        raise ArgumentError, 'TCP idle timeout must be a positive finite number'
+      end
+      @tcp_idle_timeout = seconds
+    end
+
     #  Use UDP only - don't use TCP
     #  For test/debug purposes only
     #  Defaults to false
@@ -109,16 +123,15 @@ module Dnsruby
     #  Defaults to true
     attr_accessor :recurse
 
-    #  The max UDP packet size
-    # 
-    #  Defaults to 512
+    # Maximum UDP DNS payload size. Defaults to 1232 bytes; the minimum is 512.
+    # A caller-supplied OPT record sets the receive size for that query.
     attr_reader :udp_size
 
     #  The address of the resolver to send queries to
     attr_reader :server
 
     #  Use DNSSEC for this PacketSender
-    #  dnssec defaults to ON
+    #  dnssec defaults to false
     attr_reader :dnssec
 
     #  Set the source address. If the arg is nil, do nothing
@@ -138,6 +151,8 @@ module Dnsruby
     # Sets the TSIG to sign outgoing messages with.
     # Pass in either a Dnsruby::RR::TSIG, or a key_name and key (or just a key)
     # Pass in nil to stop tsig signing.
+    # Defaults to hmac-sha256; the algorithm must match the server's key.
+    # You can select hmac-md5 for legacy servers, which logs a warning.
     # It is possible for client code to sign packets prior to sending - see
     # Dnsruby::RR::TSIG#apply and Dnsruby::Message#sign
     # Note that pre-signed packets will not be signed by PacketSender.
@@ -152,8 +167,8 @@ module Dnsruby
       @dnssec=on
       if (on)
         #  Set the UDP size (RFC 4035 section 4.1)
-        if (udp_packet_size < Resolver::MinDnssecUdpSize)
-          self.udp_size = Resolver::MinDnssecUdpSize
+        if (udp_packet_size < Resolver::AbsoluteMinDnssecUdpSize)
+          self.udp_size = Resolver::AbsoluteMinDnssecUdpSize
         end
       end
     end
@@ -176,6 +191,7 @@ module Dnsruby
     #  * :use_tcp
     #  * :tcp_pipelining
     #  * :tcp_pipelining_max_queries
+    #  * :tcp_idle_timeout
     #  * :no_tcp
     #  * :ignore_truncation
     #  * :src_address
@@ -202,7 +218,9 @@ module Dnsruby
       @recurse = true
       @tcp_pipelining = false
       @tcp_pipelining_max_queries = :infinite
+      @tcp_idle_timeout = Resolver::DefaultTcpIdleTimeout
       @use_counts = {}
+      @process_pid = Process.pid
 
       ignore_config_resolv_errors = false;
       if arg.nil?
@@ -223,6 +241,7 @@ module Dnsruby
               end
             end
           rescue Exception => e
+            raise if attr.to_s == 'tsig' || ValidationBudget.abort?(e)
             Dnsruby.log.error { "PacketSender : Argument #{attr}, #{arg[attr]} not valid : #{e}\n" }
           end
         end
@@ -251,7 +270,7 @@ module Dnsruby
         @ipv6=false
       rescue Exception
         begin
-          IPv6.create(@server)
+          Config.ipv6_server_address(@server)
           #           @src_address6 = '::'
           @ipv6=true
         rescue Exception
@@ -260,9 +279,61 @@ module Dnsruby
       end
     end
 
+    # A child must never write to a stream shared with its parent, including
+    # connections that have not yet been registered with SelectThread.
+    def ensure_current_process
+      return if @process_pid == Process.pid
+      selector = @selector
+      selector.send(:ensure_current_process) if selector
+      PIPELINE_MUTEX_CREATION.synchronize do
+        if @process_pid != Process.pid
+          ((@use_counts || {}).keys + [@pipeline_socket]).compact.uniq.each do |socket|
+            begin
+              socket.close unless socket.closed?
+            rescue IOError, SystemCallError
+            end
+          end
+          @use_counts = {}
+          @pipeline_socket = @pipeline_socket_configuration = @pipeline_socket_max_queries = nil
+          @pipeline_mutex = @pipeline_connect_context = @tcp_pipeline_local_port = nil
+          @selector = nil
+          @process_pid = Process.pid
+        end
+      end
+    end
+    private :ensure_current_process
+
     def close
-      #  @TODO@ What about closing?
-      #  Any queries to complete? Sockets to close?
+      ensure_current_process
+      # Idle streams close immediately; streams still serving another Resolver
+      # drain before closing. Do not wait for a writer blocked in connect.
+      @selector.retire_sender(self) if @selector
+    end
+
+    # Reserve ownership while Resolver still holds its query lock, then perform
+    # the synchronous public send after releasing that lock. close/timeout can
+    # cancel this reservation even before send_async has started.
+    def prepare_send_async(message, queue, id, retire_connection = false, control = QueryControl.current, transmission_budget = nil) #:nodoc:
+      ensure_current_process
+      selector = @selector = SelectThread.instance
+      pending = SelectThread::QuerySettings.new(nil, message, @ignore_truncation,
+        queue, id, nil, @server, @port, TransportClock.now + @packet_timeout, udp_packet_size, self)
+      pending.send_token = SelectThread::SendToken.new(control, transmission_budget)
+      pending.endtime = [pending.endtime, control&.deadline].compact.min
+      pending.retire_connection = retire_connection
+      return unless selector.reserve_send(pending)
+      -> do
+        next unless selector.query_settings_for_id(id).equal?(pending)
+        previous_context = Thread.current[:dnsruby_response_query]
+        begin
+          Thread.current[:dnsruby_response_query] = pending
+          send_async(message, queue, id)
+        rescue StandardError => error
+          selector.fail_query_if_current(id, queue, pending, error)
+        ensure
+          Thread.current[:dnsruby_response_query] = previous_context
+        end
+      end
     end
 
     # Asynchronously send a Message to the server. The send can be done using just
@@ -306,6 +377,7 @@ module Dnsruby
     # validation for this particular packet (assuming SingleResolver#dnssec == true)
     # Note that these options should not normally be used!
     def send_async(*args) # msg, client_queue, client_query_id, use_tcp=@use_tcp)
+      ensure_current_process
       #  @TODO@ Need to select a good Header ID here - see forgery-resilience RFC draft for details
       msg = args[0]
       client_query_id = nil
@@ -333,31 +405,36 @@ module Dnsruby
           end
         end
       end
-      #  Need to keep track of the request mac (if using tsig) so we can validate the response (RFC2845 4.1)
-      #       #Are we using EventMachine or native Dnsruby?
-      #       if (Resolver.eventmachine?)
-      #         return send_eventmachine(query_packet, msg, client_query_id, client_queue, use_tcp)
-      #       else
+      # Keep the request MAC for response verification (RFC 8945 section 4.3.1).
       if (!client_query_id)
         client_query_id = Time.now + rand(10000) # is this safe?!
       end
+      response_settings = response_query_settings(client_query_id)
+      disable_edns = args[4] || (response_settings && response_settings.disable_edns)
 
       begin
-        query_packet = make_query_packet(msg, use_tcp)
+        query_packet = make_query_packet(msg, use_tcp, disable_edns)
+        check_message_size(query_packet)
       rescue EncodeError => err
         Dnsruby.log.error { "#{err}" }
         st = SelectThread.instance
-        st.push_exception_to_select(client_query_id, client_queue, err, nil)
+        if response_settings
+          st.fail_query_if_current(client_query_id, client_queue, response_settings, err)
+        else
+          st.push_exception_to_select(client_query_id, client_queue, err, nil)
+        end
         return
       end
 
-      if (msg.do_caching && (msg.class != Update))
+      if (msg.do_caching && !msg.send_raw && (msg.class != Update) && !@tsig && !msg.tsig && !msg.signed?)
         #  Check the cache!!
         cachedanswer = nil
         if (msg.header.rd)
-          cachedanswer = @@recursive_cache.find(msg.question()[0].qname, msg.question()[0].type)
+          cachedanswer = @@recursive_cache.find(msg.question()[0].qname, msg.question()[0].type, msg.question()[0].qclass,
+            scope: positive_cache_scope(msg, query_packet))
         else
-          cachedanswer = @@authoritative_cache.find(msg.question()[0].qname, msg.question()[0].type)
+          cachedanswer = @@authoritative_cache.find(msg.question()[0].qname, msg.question()[0].type, msg.question()[0].qclass,
+            scope: positive_cache_scope(msg, query_packet))
         end
         if (cachedanswer)
           TheLog.debug("Sending cached answer to client\n")
@@ -366,7 +443,11 @@ module Dnsruby
           #  If we can find the answer, send it to the client straight away
           #  Post the result to the client using SelectThread
           st = SelectThread.instance
-          st.push_response_to_select(client_query_id, client_queue, cachedanswer, msg, self)
+          if response_settings
+            st.push_response_for_query(response_settings, cachedanswer, msg, self)
+          else
+            st.push_response_to_select(client_query_id, client_queue, cachedanswer, msg, self)
+          end
           return client_query_id
         end
       end
@@ -376,13 +457,18 @@ module Dnsruby
           #  Can't send the message - abort!
           err=IOError.new("Can't send message - too big for UDP and no_tcp=true")
           Dnsruby.log.error { "#{err}" }
-          st.push_exception_to_select(client_query_id, client_queue, err, nil)
+          st = SelectThread.instance
+          if response_settings
+            st.fail_query_if_current(client_query_id, client_queue, response_settings, err)
+          else
+            st.push_exception_to_select(client_query_id, client_queue, err, nil)
+          end
           return
         end
         Dnsruby.log.debug { "Query packet length exceeds max UDP packet size - using TCP" }
         use_tcp = true
       end
-      send_dnsruby(query_packet, msg, client_query_id, client_queue, use_tcp)
+      send_dnsruby(query_packet, msg, client_query_id, client_queue, use_tcp, disable_edns)
       return client_query_id
       #       end
     end
@@ -396,15 +482,26 @@ module Dnsruby
     #  In general, every subsequent call the function will either return the current tcp
     #  pipeline socket or a new connected socket if the current one was closed by the server
     def tcp_pipeline_socket(src_port)
+      ensure_current_process
       Dnsruby.log.debug("Using tcp_pipeline_socket")
-      sockaddr = Socket.sockaddr_in(@port, @server)
+      # A live connection retains its endpoint and source binding. Compare the
+      # configured source-port pool, not each randomly selected candidate port.
+      attempt = @pipeline_connect_context&.[](2) || transport_configuration
+      configuration = [attempt[:server], attempt[:port],
+        attempt[:ipv6] ? attempt[:src_address6] : attempt[:src_address], attempt[:src_port], attempt[:ipv6]]
+      sockaddr = Socket.sockaddr_in(configuration[1], configuration[0])
 
       reuse_pipeline_socket = -> do
         begin
-          max = @tcp_pipelining_max_queries
+          max = @pipeline_socket_max_queries
           use_count = @use_counts[@pipeline_socket]
-          if use_count && max != :infinite && use_count >= max
-             #we can't reuse the socket since max is reached
+          if !SelectThread.instance.reusable_socket?(@pipeline_socket) ||
+              max != attempt[:tcp_pipelining_max_queries] ||
+              @pipeline_socket_configuration != configuration ||
+              (use_count && max != :infinite && use_count >= max)
+            # Capacity and endpoint belong to a connection. Outstanding peers
+            # keep their old stream and drain before it closes.
+            SelectThread.instance.retire_socket(@pipeline_socket)
             @use_counts.delete(@pipeline_socket)
             @pipeline_socket = nil
             Dnsruby.log.debug("Max queries per connection attained - creating new socket")
@@ -414,6 +511,7 @@ module Dnsruby
         rescue Errno::EISCONN
           #already connected, do nothing and reuse!
         rescue IOError, Errno::ECONNRESET #close by remote host, reconnect
+          @use_counts.delete(@pipeline_socket)
           @pipeline_socket = nil
           Dnsruby.log.debug("Connection closed - recreating socket")
         end
@@ -421,14 +519,20 @@ module Dnsruby
 
       create_pipeline_socket = -> do
         @tcp_pipeline_local_port = src_port
-        src_address = @ipv6 ? @src_address6 : @src_address
+        src_address = configuration[2]
         begin
-          @pipeline_socket = Socket.new(Socket::AF_INET, Socket::SOCK_STREAM, 0)
-          @pipeline_socket.bind(Addrinfo.tcp(src_address, src_port))
-          @pipeline_socket.connect(sockaddr)
+          @pipeline_socket_max_queries = attempt[:tcp_pipelining_max_queries]
+          @pipeline_socket_configuration = configuration
+          deadline, token = @pipeline_connect_context ||
+            [TransportClock.now + attempt[:packet_timeout], SelectThread::SendToken.new]
+          @pipeline_socket = connect_tcp(configuration[0], configuration[1], src_address, src_port,
+            deadline, token)
           Dnsruby.log.debug("Creating socket #{src_address}:#{src_port}")
           @use_counts[@pipeline_socket] = 0
         rescue Exception => e
+          # A failed bind/connect never reaches the select thread, which would
+          # otherwise own closing the socket (for example, an unusable zone).
+          @pipeline_socket.close rescue nil
           @pipeline_socket = nil
           raise e
         end
@@ -446,35 +550,254 @@ module Dnsruby
     end
 
     #  This method sends the packet using the built-in pure Ruby event loop, with no dependencies.
-    def send_dnsruby(query_bytes, query, client_query_id, client_queue, use_tcp) #:nodoc: all
-      endtime = Time.now + @packet_timeout
+    def send_dnsruby(query_bytes, query, client_query_id, client_queue, use_tcp, disable_edns = false) #:nodoc: all
+      ensure_current_process
+      st = @selector = SelectThread.instance
+      response_settings = response_query_settings(client_query_id)
+      # A configuration change applies to the next attempt, never half of this one.
+      configuration = transport_configuration
+      pipeline_mode = use_tcp && configuration[:tcp_pipelining]
+      begin
+        check_message_size(query_bytes)
+      rescue EncodeError => error
+        if response_settings
+          st.fail_query_if_current(client_query_id, client_queue, response_settings, error)
+        else
+          st.push_exception_to_select(client_query_id, client_queue, error, nil)
+        end
+        return
+      end
+      # A Resolver reservation already started this attempt's packet deadline.
+      # A protocol retry (which has bytes) receives a fresh packet deadline.
+      endtime = response_settings && response_settings.query_bytes.nil? ?
+        response_settings.endtime : TransportClock.now + configuration[:packet_timeout]
+      previous_settings = response_settings ||
+        (st.query_settings_for_id(client_query_id) if pipeline_mode)
+      pending = SelectThread::QuerySettings.new(query_bytes, query,
+        configuration[:ignore_truncation], client_queue, client_query_id, nil,
+        configuration[:server], configuration[:port],
+        endtime, udp_receive_size(query, configuration), self)
+      pending.transport_configuration = configuration
+      pending.send_token = previous_settings.send_token if previous_settings
+      # Carry the controlling lookup through selector retries and workers, whose
+      # thread-local control may differ from the thread which admitted it.
+      endtime = pending.endtime = [endtime, pending.send_token.control&.deadline].compact.min
+      pending.disable_edns = disable_edns || (previous_settings && previous_settings.disable_edns)
+      pending.cache_scope = positive_cache_scope(query, query_bytes, configuration)
+      # A protocol retry starts its own packet deadline at dispatch, including
+      # time queued for a worker. Retire the old socket membership atomically.
+      return unless st.reserve_send(pending, previous_settings)
+      previous_settings = pending
+      operation = -> do
+        send_dnsruby_with_socket(query_bytes, query, client_query_id, client_queue,
+          use_tcp, endtime, previous_settings, pipeline_mode)
+      end
+      send_operation = operation
+      if pipeline_mode
+        # Configuration-only senders remain Marshal-compatible. Allocate the
+        # writer mutex only on first use, including senders loaded from storage.
+        pipeline_mutex = PIPELINE_MUTEX_CREATION.synchronize { @pipeline_mutex ||= Mutex.new }
+        send_operation = -> do
+          locked = false
+          begin
+            until (locked = pipeline_mutex.try_lock)
+              break if previous_settings.send_token.cancelled? || TransportClock.now >= endtime ||
+                !st.query_settings_for_id(client_query_id).equal?(previous_settings)
+              sleep 0.005
+            end
+            if locked && !previous_settings.send_token.cancelled? && TransportClock.now < endtime
+              operation.call
+            else
+              st.fail_query_if_current(client_query_id, client_queue, previous_settings,
+                ResolvTimeout.new("Query timed out waiting for TCP pipeline"))
+            end
+          ensure
+            pipeline_mutex.unlock if locked
+          end
+        end
+      end
+      guarded_operation = -> do
+        send_operation.call
+      rescue StandardError => error
+        # Deferred callers do not read SendTask#value. Failures before socket
+        # registration still belong to this exact reserved attempt.
+        st.fail_query_if_current(client_query_id, client_queue, pending, error)
+      end
+      if use_tcp && st.selector_thread?
+        begin
+          st.defer_send(pipeline_mode ? self : nil, &guarded_operation)
+        rescue ThreadError => error
+          st.fail_query_if_current(client_query_id, client_queue, previous_settings,
+            ResolvTimeout.new("Cannot queue TCP send: #{error.message}"))
+        end
+      else
+        # Only the selector delegates I/O. Public caller sends remain synchronous.
+        guarded_operation.call
+      end
+    end
+
+    def check_message_size(query_bytes)
+      if query_bytes.bytesize > 65_535
+        raise EncodeError, "DNS message exceeds the 65535-byte wire limit (#{query_bytes.bytesize} bytes)"
+      end
+    end
+    private :check_message_size
+
+    class SendCancelled < StandardError; end
+    private_constant :SendCancelled
+
+    # Config has already resolved the server address. Check cancellation while
+    # connecting so an unresponsive peer cannot hold up a sender worker.
+    def connect_tcp(server, port, source, source_port, deadline, token)
+      socket = nil
+      connected = false
+      begin
+        raise SendCancelled if token.cancelled?
+        raise ResolvTimeout, 'Query timed out before TCP connect' if TransportClock.now >= deadline
+        socket = SelectThread.instance.track_socket(
+          Socket.new(server.include?(':') ? Socket::AF_INET6 : Socket::AF_INET, Socket::SOCK_STREAM, 0))
+        socket.bind(Socket.sockaddr_in(source_port, source))
+        result = socket.connect_nonblock(Socket.sockaddr_in(port, server), exception: false)
+        while result == :wait_writable || result == :wait_readable
+          raise SendCancelled if token.cancelled?
+          remaining = deadline - TransportClock.now
+          raise ResolvTimeout, 'Query timed out connecting TCP' if remaining <= 0
+          readers, writers = result == :wait_readable ? [[socket], nil] : [nil, [socket]]
+          next unless IO.select(readers, writers, nil, [remaining, 0.05].min)
+          error = socket.getsockopt(Socket::SOL_SOCKET, Socket::SO_ERROR).int
+          raise SystemCallError.new('TCP connect', error) unless error.zero?
+          result = 0
+        end
+        raise SendCancelled if token.cancelled?
+        raise ResolvTimeout, 'Query timed out connecting TCP' if TransportClock.now >= deadline
+        connected = true
+        socket
+      ensure
+        socket.close if socket && !connected && !socket.closed?
+      end
+    end
+    private :connect_tcp
+
+    # Socket#send may write only a prefix of a TCP frame. Nonblocking writes keep
+    # a stalled peer from holding the pipeline writer past the packet deadline.
+    def write_tcp_frame(socket, query_bytes, endtime, send_token)
+      check_message_size(query_bytes)
+      frame = [query_bytes.bytesize].pack('n') + query_bytes
+      offset = 0
+      while offset < frame.bytesize
+        remaining = endtime - TransportClock.now
+        begin
+          written = send_token.write do
+            raise ResolvTimeout, "Query timed out writing TCP frame" if remaining <= 0
+            socket.write_nonblock(frame.byteslice(offset, frame.bytesize - offset))
+          end
+          return offset if written == false # Cancelled between partial writes.
+          raise IOError, "TCP write made no progress" unless written && written > 0
+          offset += written
+        rescue IO::WaitWritable
+          # A canceled pipeline may still have peers and remain open. Bound the
+          # wait so cancellation releases the worker even if it never becomes
+          # writable; the next iteration retires any incomplete frame.
+          IO.select(nil, [socket], nil, [remaining, 0.05].min)
+        end
+      end
+      true
+    end
+
+    # A datagram is sent atomically. Order the actual write with cancellation,
+    # as for TCP, without holding the token while waiting for a writable socket.
+    def write_udp_packet(socket, query_bytes, endtime, send_token, wait_for_writable)
+      loop do
+        remaining = endtime - TransportClock.now
+        begin
+          written = send_token.write do
+            raise ResolvTimeout, "Query timed out writing UDP packet" if remaining <= 0
+            socket.write_nonblock(query_bytes)
+          end
+          return false if written == false
+          raise IOError, "Incomplete UDP datagram write" unless written == query_bytes.bytesize
+          return true
+        rescue IO::WaitWritable
+          raise unless wait_for_writable
+          IO.select(nil, [socket], nil, [remaining, 0.05].min)
+        end
+      end
+    end
+
+    def response_query_settings(client_query_id)
+      settings = Thread.current[:dnsruby_response_query]
+      settings if settings && settings.client_query_id == client_query_id &&
+        settings.single_resolver.equal?(self)
+    end
+
+    def send_dnsruby_with_socket(query_bytes, query, client_query_id, client_queue,
+        use_tcp, endtime, previous_settings, pipeline_mode = use_tcp && @tcp_pipelining)
       #  First send the query (synchronously)
       st = SelectThread.instance
+      report_error = ->(error) do
+        st.fail_query_if_current(client_query_id, client_queue, previous_settings, error)
+      end
+      return unless st.query_settings_for_id(client_query_id).equal?(previous_settings)
+      return if previous_settings.send_token.cancelled?
+      configuration = previous_settings.transport_configuration || transport_configuration
+      server, port = configuration.values_at(:server, :port)
+      unless previous_settings.send_token.transmission_budget.admit(server, port, use_tcp, query.question.first)
+        report_error.call(ResolvTimeout.new("DNS retry limit reached for server and transport"))
+        return
+      end
       socket = nil
       runnextportloop = true
       numtries = 0
-      src_address = @src_address
-      if (@ipv6)
-        src_address = @src_address6
-      end
+      source_port_configuration = configuration[:src_port]
+      source_ports = source_port_configuration.uniq
+      random_source_port = source_ports.first == 0
+      src_address = configuration[configuration[:ipv6] ? :src_address6 : :src_address]
       while (runnextportloop) do
+        # A pre-close reservation gets a one-use socket. Recheck on every
+        # admission retry without changing ownership or extending its deadline.
+        pipeline = pipeline_mode && !previous_settings.retire_connection
+        return unless st.query_settings_for_id(client_query_id).equal?(previous_settings)
+        return if previous_settings.send_token.cancelled?
         begin
+          if TransportClock.now >= endtime
+            report_error.call(ResolvTimeout.new("Query timed out before connect"))
+            return
+          end
           numtries += 1
-          src_port = get_next_src_port
+          src_port = if random_source_port
+            # Keep this attempt's mode even if a caller replaces the pool.
+            50_000 + SecureRandom.random_number(15_536)
+          else
+            source_ports[SecureRandom.random_number(source_ports.length)]
+          end
           if (use_tcp)
             begin
-              if (@tcp_pipelining)
-                socket, new_socket = tcp_pipeline_socket(src_port)
+              if pipeline
+                begin
+                  # RFC 7766 requires unique outstanding IDs on a connection.
+                  # Moving a collision to a fresh stream preserves caller IDs,
+                  # including raw messages and bytes covered by TSIG.
+                  if @pipeline_socket && st.query_id_in_use?(@pipeline_socket, query.header.id)
+                    st.retire_socket(@pipeline_socket)
+                    @use_counts.delete(@pipeline_socket)
+                    @pipeline_socket = nil
+                  end
+                  @pipeline_connect_context = [endtime, previous_settings.send_token, configuration]
+                  socket, new_socket = tcp_pipeline_socket(src_port)
+                ensure
+                  remove_instance_variable(:@pipeline_connect_context)
+                end
                 src_port = @tcp_pipeline_local_port
               else
-                socket = Socket.tcp(@server, @port, src_address, src_port, connect_timeout: @packet_timeout)
+                socket = connect_tcp(server, port, src_address, src_port,
+                  endtime, previous_settings.send_token)
                 new_socket = true
               end
             rescue Errno::EBADF, Errno::ENETUNREACH => e
               #  Can't create a connection
-              err=IOError.new("TCP connection error to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception = #{e.class}, #{e}")
+              err=IOError.new("TCP connection error to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception = #{e.class}, #{e}")
               Dnsruby.log.error { "#{err}" }
-              st.push_exception_to_select(client_query_id, client_queue, err, nil)
+              report_error.call(err)
               return
             end
           else
@@ -484,81 +807,150 @@ module Dnsruby
               socket = UDPSocket.new()
             else
               #               ipv6 = @src_address =~ /:/
-              socket = UDPSocket.new(@ipv6 ? Socket::AF_INET6 : Socket::AF_INET)
+              socket = UDPSocket.new(configuration[:ipv6] ? Socket::AF_INET6 : Socket::AF_INET)
             end
+            st.track_socket(socket)
             new_socket = true
             socket.bind(src_address, src_port)
-            socket.connect(@server, @port)
+            socket.connect(server, port)
           end
-          runnextportloop = false
+        rescue SendCancelled
+          return
         rescue Exception => e
           if (socket!=nil)
             begin
               #let the select thread close the socket if tcp_pipeli
-              socket.close unless @tcp_pipelining && !new_socket
+              socket.close unless pipeline && !new_socket
             rescue Exception
             end
           end
-          #  Try again if the error was EADDRINUSE and a random source port is used
-          #  Maybe try a max number of times?
-          if ((e.class != Errno::EADDRINUSE) || (numtries > 50) ||
-              ((e.class == Errno::EADDRINUSE) && (src_port == @src_port[0])))
-            err_msg = "dnsruby can't connect to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception = #{e.class}, #{e} #{e.backtrace}"
+          # Only an occupied bind removes a candidate. A pipeline admission
+          # retry has not found the port unusable and must retain the pool.
+          source_ports.delete(src_port) if e.is_a?(Errno::EADDRINUSE) && !random_source_port
+          # Random ports keep the bounded retry policy; configured pools do
+          # not revisit an occupied member during this attempt.
+          if !e.is_a?(Errno::EADDRINUSE) || (random_source_port && numtries >= 50) ||
+              (!random_source_port && source_ports.empty?)
+            err_msg = "dnsruby can't connect to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception = #{e.class}, #{e} #{e.backtrace}"
             err=IOError.new(err_msg)
             Dnsruby.log.error( "#{err}")
             Dnsruby.log.error(e.backtrace)
-            if @tcp_pipelining
-              st.push_exception_to_select(client_query_id, client_queue, SocketEofResolvError.new(err_msg), nil)
+            if e.kind_of?(Errno::ETIMEDOUT) || e.kind_of?(ResolvTimeout)
+              report_error.call(ResolvTimeout.new(err_msg))
+            elsif pipeline
+              report_error.call(SocketEofResolvError.new(err_msg))
             else
-              st.push_exception_to_select(client_query_id, client_queue, err, nil)
+              report_error.call(err)
             end
             return
           end
+          next # Retry a random source-port collision before registration.
         end
-      end
-      if (socket==nil)
-        err=IOError.new("dnsruby can't connect to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}")
-        Dnsruby.log.error { "#{err}" }
-        st.push_exception_to_select(client_query_id, client_queue, err, nil)
-        return
-      end
-      Dnsruby.log.debug { "Sending packet to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp} : #{query.question()[0].qname}, #{query.question()[0].qtype}" }
-      #             print "#{Time.now} : Sending packet to #{@server} : #{query.question()[0].qname}, #{query.question()[0].qtype}\n"
-      #  Listen for the response before we send the packet (to avoid any race conditions)
-      query_settings = SelectThread::QuerySettings.new(query_bytes, query, @ignore_truncation, client_queue, client_query_id, socket, @server, @port, endtime, udp_packet_size, self)
-      query_settings.is_persistent_socket = @tcp_pipelining if use_tcp
-      query_settings.tcp_pipelining_max_queries = @tcp_pipelining_max_queries if @tcp_pipelining
-      begin
-        if (use_tcp)
-          lenmsg = [query_bytes.length].pack('n')
-          socket.send(lenmsg, 0)
-        end
-        socket.send(query_bytes, 0)
-
-        #  The select thread will now wait for the response and send that or a
-        #  timeout back to the client_queue.
-        st.add_to_select(query_settings)
-      rescue Exception => e
-        err_msg = "Send failed to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception : #{e}"
-        err=IOError.new(err_msg)
-        Dnsruby.log.error { "#{err}" }
-        Dnsruby.log.error(e.backtrace)
-        if @tcp_pipelining
-          st.push_exception_to_select(client_query_id, client_queue, SocketEofResolvError.new(err_msg), nil) if new_socket
-        else
-          st.push_exception_to_select(client_query_id, client_queue, err, nil)
+        if (socket==nil)
+          err=IOError.new("dnsruby can't connect to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}")
+          Dnsruby.log.error { "#{err}" }
+          report_error.call(err)
+          return
         end
-        begin
-          #we let the select_thread close the socket when doing tcp
-          #pipelining
-          socket.close unless @tcp_pipelining && !new_socket
-        rescue Exception
+        Dnsruby.log.debug { "Sending packet to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp} : #{query.question()[0].qname}, #{query.question()[0].qtype}" }
+        # A reused pipeline is already readable by the selector. Publish its new
+        # query before writing so an immediate answer always has an owner.
+        query_settings = SelectThread::QuerySettings.new(query_bytes, query, previous_settings.ignore_truncation,
+          client_queue, client_query_id, socket, server, port, endtime, previous_settings.udp_packet_size, self)
+        query_settings.transport_configuration = configuration
+        query_settings.send_token = previous_settings.send_token
+        query_settings.disable_edns = previous_settings.disable_edns
+        query_settings.cache_scope = previous_settings.cache_scope
+        query_settings.is_persistent_socket = pipeline
+        query_settings.tcp_pipelining_max_queries = (@pipeline_socket_max_queries || configuration[:tcp_pipelining_max_queries]) if pipeline
+        if use_tcp
+          registered = st.register_query_if_current(query_settings, previous_settings)
+          if registered == :retry_connection
+            # No bytes were written and the original reservation still owns the
+            # attempt. Peers on the retiring stream keep their registrations.
+            @use_counts.delete(socket)
+            @pipeline_socket = nil if @pipeline_socket == socket
+            next
+          end
+          unless registered
+            # Completion can win during connect. Release only this unused socket
+            # reservation; other queries may still own an existing pipeline.
+            @use_counts[socket] -= 1 if pipeline && @use_counts[socket]
+            socket.close if new_socket && !socket.closed?
+            if pipeline && socket.closed?
+              @use_counts.delete(socket)
+              @pipeline_socket = nil if @pipeline_socket == socket
+            end
+            return
+          end
+          previous_settings = query_settings
+        elsif !st.query_settings_for_id(client_query_id).equal?(previous_settings)
+          socket.close
+          return
         end
-        return
+        runnextportloop = false
       end
+      send_packet = -> do
+        begin
+          if use_tcp
+            completed = write_tcp_frame(socket, query_bytes, endtime, query_settings.send_token)
+            # Cancellation of a partially written frame retires its whole stream;
+            # otherwise another frame could be appended to a corrupt prefix.
+            unless completed == true
+              if pipeline && completed > 0
+                st.fail_tcp_send(query_settings, SocketEofResolvError.new("TCP send cancelled"))
+                @use_counts.delete(socket)
+                @pipeline_socket = nil if @pipeline_socket == socket
+              elsif !pipeline
+                socket.close rescue nil
+              end
+              return
+            end
+          else
+            unless write_udp_packet(socket, query_bytes, endtime, query_settings.send_token, !st.selector_thread?)
+              socket.close
+              return
+            end
+            socket.close unless st.register_query_if_current(query_settings, previous_settings)
+          end
+        rescue IO::WaitWritable
+          # Only a congested UDP send reaches this branch on the selector. Keep
+          # ordinary UDP traffic independent of the pool's blocked TCP writers.
+          return st.defer_send { send_packet.call }
+        rescue Exception => e
+          err_msg = "Send failed to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp}, exception : #{e}"
+          err=IOError.new(err_msg)
+          Dnsruby.log.error { "#{err}" }
+          Dnsruby.log.error(e.backtrace)
+          if pipeline
+            error = e.kind_of?(ResolvTimeout) ? e : SocketEofResolvError.new(err_msg)
+            st.fail_tcp_send(query_settings, error)
+            @use_counts.delete(socket)
+            @pipeline_socket = nil if @pipeline_socket == socket
+          elsif e.kind_of?(ResolvTimeout)
+            report_error.call(e)
+          else
+            report_error.call(err)
+          end
+          begin
+            #we let the select_thread close the socket when doing tcp
+            #pipelining
+            socket.close unless pipeline && !new_socket
+          rescue Exception
+          end
+          return
+        end
 
-      Dnsruby.log.debug { "Packet sent to #{@server}:#{@port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp} : #{query.question()[0].qname}, #{query.question()[0].qtype}" }
-      #       print "Packet sent to #{@server}:#{@port} from #{@src_address}:#{src_port}, use_tcp=#{use_tcp} : #{query.question()[0].qname}, #{query.question()[0].qtype}\n"
+        Dnsruby.log.debug { "Packet sent to #{server}:#{port} from #{src_address}:#{src_port}, use_tcp=#{use_tcp} : #{query.question()[0].qname}, #{query.question()[0].qtype}" }
+      end
+      send_packet.call
+    rescue StandardError => error
+      # Errors outside the expected connect/write failures must not disappear in
+      # a worker's result queue. Do not consume a replacement protocol attempt.
+      st.fail_query_if_current(client_query_id, client_queue, previous_settings, error)
+      # A failure between connect and registration has no selector-owned socket
+      # yet. Release that new descriptor here; reused streams retain their peers.
+      socket.close if socket && new_socket && !socket.closed?
     end
 
     #  The source port to send queries from
@@ -567,10 +959,8 @@ module Dnsruby
     # 
     #  Defaults to 0 - random port
     def src_port
-      if (@src_port.length == 1)
-        return @src_port[0]
-      end
-      return @src_port
+      ports = @src_port
+      ports.length == 1 ? ports[0] : ports
     end
 
     #  Can be a single Integer or a Range or an Array
@@ -582,8 +972,9 @@ module Dnsruby
     #         res.src_port=60015..60115
     # 
     def src_port=(p)
-      @src_port=[]
-      add_src_port(p)
+      ports = Resolver.get_ports_from(p)
+      raise ArgumentError, 'Source port pool must not be empty' if ports.empty?
+      @src_port = ports if Resolver.check_port(ports)
     end
 
     #  Can be a single Integer or a Range or an Array
@@ -599,57 +990,81 @@ module Dnsruby
     #         res.add_src_port(60015..60115)
     # 
     def add_src_port(p)
-      if (Resolver.check_port(p, @src_port))
-        a = Resolver.get_ports_from(p)
-        a.each do |x|
-          if ((@src_port.length > 0) && (x == 0))
-            raise ArgumentError.new("src_port of 0 only allowed as only src_port value (currently #{@src_port.length} values")
-          end
-          @src_port.push(x)
-        end
+      additions = Resolver.get_ports_from(p)
+      # Adding a concrete port starts a pool instead of retaining the default
+      # random-port sentinel, which would otherwise hide every added port.
+      ports = @src_port == [0] && !additions.empty? && !additions.include?(0) ? [] : @src_port
+      if Resolver.check_port(additions, ports)
+        @src_port = ports + additions
       end
     end
 
 
     def get_next_src_port
-      # Different OSes have different interpretations of "random port" here.
-      # Apparently, Linux will just give you the same port as last time, unless it is still
-      # open, in which case you get n+1.
-      # We need to determine an actual (random) number here, then ask the OS for it, and
-      # continue until we get one.
-      if (@src_port[0] == 0)
-        candidate = -1
-        #         # better to construct an array of all the ports we *can* use, and then just pick one at random!
-        #         candidate = Iana::UNRESERVED_PORTS[rand(Iana::UNRESERVED_PORTS.length())]
-        #         #        while (!(Resolver.port_in_range(candidate)))
-        #         #          candidate = (rand(65535-1024) + 1024)
-        #         #        end
-        #  @TODO@ Should probably construct a bitmap of the IANA ports...
-        candidate = 50000 + (rand(15535)) # pick one over 50000
-        return candidate
-      end
-      pos = rand(@src_port.length)
-      return @src_port[pos]
+      # Preserve the existing high-port range and configured pools, using
+      # unpredictable choices for the source-port contribution to DNS IDs.
+      ports = @src_port.dup
+      if (ports[0] == 0)
+        return 50_000 + SecureRandom.random_number(15_536)
+      end
+      ports[SecureRandom.random_number(ports.length)]
+    end
+
+    # Retain the request actually sent, including EDNS options, but not its
+    # transaction ID. Never share answers across upstream or source routes.
+    def positive_cache_scope(query, query_bytes = query.encode, configuration = nil) #:nodoc:
+      route = if configuration
+        configuration.values_at(:server, :port, :src_address, :src_address6, :src_port)
+      else
+        [@server.to_s.dup.freeze, @port, @src_address.to_s.dup.freeze,
+          @src_address6.to_s.dup.freeze, @src_port.dup.freeze]
+      end
+      (route + [query.send_raw, query_bytes.byteslice(2..-1).b.freeze]).freeze
+    end
+
+    # Keep the endpoint, binding and response policy together while connect or
+    # pipeline admission releases the caller to change resolver settings.
+    def transport_configuration
+      server = @server.to_s.dup.freeze
+      {server: server, port: @port, src_address: @src_address.to_s.dup.freeze,
+        src_address6: @src_address6.to_s.dup.freeze, src_port: @src_port.dup.freeze,
+        ipv6: server.include?(':'), packet_timeout: @packet_timeout, udp_size: udp_packet_size,
+        tcp_pipelining: @tcp_pipelining, tcp_pipelining_max_queries: @tcp_pipelining_max_queries,
+        tcp_idle_timeout: tcp_idle_timeout, ignore_truncation: @ignore_truncation,
+        no_tcp: @no_tcp, dnssec: @dnssec}.freeze
     end
+    private :transport_configuration
 
     def check_response(response, response_bytes, query, client_queue, client_query_id, tcp)
+      settings = response_query_settings(client_query_id)
+      configuration = settings&.transport_configuration
+      dnssec = configuration ? configuration[:dnssec] : @dnssec
+      ignore_truncation = settings ? settings.ignore_truncation : @ignore_truncation
+      no_tcp = settings ? settings.no_tcp : @no_tcp
       #  @TODO@ Should send_raw avoid this?
       if (!query.send_raw)
-        sig_value = check_tsig(query, response, response_bytes)
-        if (sig_value != :okay)
-          #  Should send error back up to Resolver here, and then NOT QUERY AGAIN!!!
-          return sig_value
+        # Reject unrelated packet headers before TSIG verification mutates the
+        # response or an error/truncation triggers another transport attempt.
+        unless response.header.qr && response.header.opcode == query.header.opcode &&
+            response.header.id == query.header.id
+          return false
         end
-        if ((response.header.get_header_rcode == RCode.FORMERR) &&
-            (query.header.arcount == 0))
-          # Raise an error
-          return true
+        checking_signature = query.tsig && response.tsig
+        sig_value = check_tsig(query, response, response_bytes)
+        authenticated_tsig_error = checking_signature && sig_value.kind_of?(TsigError) && response.verified? &&
+          response.header.get_header_rcode == RCode.NOTAUTH && [18, 22].include?(response.tsigerror)
+        if (sig_value != :okay) && !authenticated_tsig_error
+          # Keep waiting for a signed response without extending the deadline
+          # (RFC 8945 section 5.4). Save the error in case the query times out.
+          settings = response_query_settings(client_query_id)
+          settings.authentication_error = sig_value if settings
+          return false
         end
-        #  Should check that question section is same as question that was sent! RFC 5452
-        #  If it's not an update...
+        # Match the question when one is present. TCP error replies may omit
+        # it (RFC 7766 section 7); UDP replies must match (RFC 5452 section 9.1).
         if (query.class == Update)
-          #  @TODO@!!
-        else
+          # TODO: Validate the zone section for UPDATE replies.
+        elsif !(tcp && response.question.empty? && response.header.get_header_rcode != RCode.NOERROR)
           if ((response.question.size == 0) ||
               (response.question[0].qname.labels != query.question[0].qname.labels) ||
               (response.question[0].qtype != query.question[0].qtype) ||
@@ -660,19 +1075,29 @@ module Dnsruby
             return false
           end
         end
-        # end
-        #  IF WE GET FORMERR BACK HERE (and we have EDNS0 on) THEN
-        #  TRY AGAIN WITH NO OPT RECORDS! (rfc2671 section 5.3)
+        return sig_value if authenticated_tsig_error
+        # Retry a FORMERR without optional EDNS (RFC 6891 section 6.2.2).
+        # Queries requiring DNSSEC validation must keep EDNS and DO.
         if ((response.header.get_header_rcode == RCode.FORMERR) &&
-            (query.header.arcount > 0))
-          #  try resending the message with no OPT record
-          query.remove_additional
-          query.send_raw = true
-          send_async(query, client_queue, client_query_id, false)
+            query.get_opt && !(dnssec && query.do_validation))
+          # EDNS negotiation belongs to this attempt, not to the caller's raw
+          # mode or Message (which may also be in use by another nameserver).
+          retry_query = query.dup
+          retry_query.header = query.header.dup
+          retry_query.remove_additional
+          query.additional.each do |rr|
+            retry_query.add_additional(rr) unless rr.type == Types.OPT || rr.type == Types.TSIG
+          end
+          if query.tsig
+            return TsigError.new unless query.tsig.key
+            retry_query.tsigstate = :Unsigned
+            retry_query.set_tsig(query.tsig.dup)
+          end
+          send_async(retry_query, client_queue, client_query_id, tcp, true)
           return false
         end
-        if (response.header.tc && !tcp && !@ignore_truncation)
-          if (@no_tcp)
+        if (response.header.tc && !tcp && !ignore_truncation)
+          if (no_tcp)
             Dnsruby.log.debug { "Truncated response - not resending over TCP as no_tcp==true" }
           else
             #  Try to resend over tcp
@@ -684,21 +1109,32 @@ module Dnsruby
           end
         end
       end
+      settings = response_query_settings(client_query_id)
+      if settings && settings.cache_scope && settings.cache_scope == positive_cache_scope(query, settings.query_bytes)
+        response.instance_variable_set(:@dnsruby_cache_scope, settings.cache_scope)
+      else
+        # Configuration changed during this exchange, or this response has no
+        # tracked transport. Do not seed a cache with an ambiguous route.
+        response.instance_variable_set(:@dnsruby_cache_scope, false)
+      end
       return true
     end
 
     def check_tsig(query, response, response_bytes)
       if (query.tsig)
-        if (response.tsig)
+        if (response_signature = response.tsig)
           if !query.tsig.verify(query, response, response_bytes)
+            authenticated_error = RR::TSIG.authenticated_error(response, response_signature)
+            return authenticated_error if authenticated_error
             #  Discard packet and wait for correctly signed response
             Dnsruby.log.error { "TSIG authentication failed!" }
             return TsigError.new
           end
         else
-          #  Treated as having format error and discarded (RFC2845, 4.6)
-          #  but return a different error code, because some servers fail at
-          #  this
+          # A missing signature is a format error (RFC 8945 section 5.4).
+          # Keep a distinct exception so callers can diagnose unsigned peers.
+          response.tsigstate = :Failed
+          response.tsigerror = RCode.FORMERR
           Dnsruby.log.error { "Expecting TSIG signed response, but got unsigned response - discarding" }
           return TsigNotSignedResponseError.new
         end
@@ -721,52 +1157,52 @@ module Dnsruby
     end
 
     #  Prepare the packet for sending
-    def make_query_packet(packet, use_tcp = @use_tcp) #:nodoc: all
-      if (!packet.send_raw) # Don't mess with this packet!
+    def make_query_packet(packet, use_tcp = @use_tcp, disable_edns = false) #:nodoc: all
+      packet.tsig.warn_if_legacy_algorithm if packet.tsig
+      # TSIG signs the header and every section, so leave signed messages intact.
+      # Decoded messages may have a TSIG without a signing state.
+      prepared = packet.signed? || packet.tsig
+      if !packet.send_raw && !prepared
         if (packet.header.opcode == OpCode.QUERY || @recurse)
           packet.header.rd=@recurse
         end
 
-        #  Only do this if the packet has not been prepared already!
-        if (@dnssec)
+        if @dnssec && !disable_edns
           prepare_for_dnssec(packet)
-        elsif ((udp_packet_size > Resolver::DefaultUDPSize) && !use_tcp)
-          #       if ((udp_packet_size > Resolver::DefaultUDPSize) && !use_tcp)
-          #  @TODO@ What if an existing OPT RR is not big enough? Should we replace it?
+        elsif !disable_edns && (udp_packet_size > 512) && !use_tcp
           add_opt_rr(packet)
         end
       end
 
-      if (@tsig && !packet.signed?)
-        @tsig.apply(packet)
+      # Prepare the whole message before encoding, which may sign it.
+      # Use the resolver's key if the Message has none.
+      encoded = packet.encode
+      if @tsig && !packet.signed? && !packet.tsig
+        packet.set_tsig(@tsig)
+        encoded = packet.encode
       end
-      return packet.encode
+      encoded
     end
 
     def add_opt_rr(packet)
       Dnsruby.log.debug { ";; Adding EDNS extension with UDP packetsize  #{udp_packet_size}.\n" }
-      #  RFC 3225
-      optrr = RR::OPT.new(udp_packet_size)
-
-      #  Only one OPT RR allowed per packet - do we already have one?
-      if (packet.additional.rrset(packet.question()[0].qname, Types::OPT).rrs.length == 0)
-        packet.add_additional(optrr)
-      end
+      # Find OPT by type: its class field stores the payload size. Reuse it and
+      # its options, since only one OPT is allowed (RFC 6891 section 6.1.1).
+      packet.add_additional(RR::OPT.new(udp_packet_size)) unless packet.get_opt
+      packet.get_opt
     end
 
     def prepare_for_dnssec(packet)
-      #  RFC 4035
+      # Request DNSSEC records with DO (RFC 4035 section 4.1).
       Dnsruby.log.debug { ";; Adding EDNS extension with UDP packetsize #{udp_packet_size} and DNS OK bit set\n" }
-      optrr = RR::OPT.new(udp_packet_size) # Decimal UDPpayload
+      optrr = add_opt_rr(packet)
       optrr.dnssec_ok=true
 
-      if (packet.additional.rrset(packet.question()[0].qname, Types::OPT).rrs.length == 0)
-        packet.add_additional(optrr)
-      end
-
-      packet.header.ad = false # RFC 4035 section 4.6
+      # We use DO to request DNSSEC records and leave AD clear
+      # (RFC 6840 section 5.7).
+      packet.header.ad = false
 
-      #  SHOULD SET CD HERE!!!
+      # Local validation needs the original data, even if the server rejects it.
       if (packet.do_validation)
         packet.header.cd = true
       end
@@ -778,10 +1214,16 @@ module Dnsruby
 
     #  Return the packet size to use for UDP
     def udp_packet_size
-      #  if @udp_size > DefaultUDPSize then we use EDNS and
-      #  @udp_size should be taken as the maximum packet_data length
-      ret = (@udp_size > Resolver::DefaultUDPSize ? @udp_size : Resolver::DefaultUDPSize)
-      return ret
+      # Values below 512 mean 512 in EDNS (RFC 6891 section 6.2.3).
+      [@udp_size, 512].max
+    end
+
+    # Match the receive buffer to this query's OPT so a valid reply isn't
+    # cut short when the caller advertises more than the resolver default.
+    def udp_receive_size(query, configuration = nil)
+      opt = query.get_opt
+      opt ? [opt.payloadsize, 512].max : (configuration ? configuration[:udp_size] : udp_packet_size)
     end
+    private :udp_receive_size
   end
 end