native.rb in httpx-0.23.0

- old
+ new

@@ -31,10 +31,11 @@
 
     def initialize(_, options)
       super
       @ns_index = 0
       @resolver_options = DEFAULTS.merge(@options.resolver_options)
+      @socket_type = @resolver_options.fetch(:socket_type, :udp)
       @nameserver = Array(@resolver_options[:nameserver]) if @resolver_options[:nameserver]
       @ndots = @resolver_options[:ndots]
       @search = Array(@resolver_options[:search]).map { |srch| srch.scan(/[^.]+/) }
       @_timeouts = Array(@resolver_options[:timeouts])
       @timeouts = Hash.new { |timeouts, host| timeouts[host] = @_timeouts.dup }
@@ -154,11 +155,11 @@
         transition(:idle)
         resolve(connection)
       else
 
         @timeouts.delete(host)
-        @queries.delete(h)
+        reset_hostname(h, reset_candidates: false)
 
         return unless @queries.empty?
 
         @connections.delete(connection)
         # This loop_time passed to the exception is bogus. Ideally we would pass the total
@@ -167,81 +168,143 @@
       end
     end
 
     def dread(wsize = @resolver_options[:packet_size])
       loop do
+        wsize = @large_packet.capacity if @large_packet
+
         siz = @io.read(wsize, @read_buffer)
-        return unless siz && siz.positive?
 
-        parse(@read_buffer)
+        unless siz
+          ex = EOFError.new("descriptor closed")
+          ex.set_backtrace(caller)
+          raise ex
+        end
+
+        return unless siz.positive?
+
+        if @socket_type == :tcp
+          # packet may be incomplete, need to keep draining from the socket
+          if @large_packet
+            # large packet buffer already exists, continue pumping
+            @large_packet << @read_buffer
+
+            next unless @large_packet.full?
+
+            parse(@large_packet.to_s)
+
+            @socket_type = @resolver_options.fetch(:socket_type, :udp)
+            @large_packet = nil
+            transition(:closed)
+            return
+          else
+            size = @read_buffer[0, 2].unpack1("n")
+
+            if size > @read_buffer.bytesize
+              # only do buffer logic if it's worth it, and the whole packet isn't here already
+              @large_packet = Buffer.new(size)
+              @large_packet << @read_buffer.byteslice(2..-1)
+
+              next
+            else
+              parse(@read_buffer)
+            end
+          end
+        else # udp
+          parse(@read_buffer)
+        end
+
         return if @state == :closed
       end
     end
 
     def dwrite
       loop do
         return if @write_buffer.empty?
 
         siz = @io.write(@write_buffer)
-        return unless siz && siz.positive?
 
+        unless siz
+          ex = EOFError.new("descriptor closed")
+          ex.set_backtrace(caller)
+          raise ex
+        end
+
+        return unless siz.positive?
+
         return if @state == :closed
       end
     end
 
     def parse(buffer)
-      begin
-        addresses = Resolver.decode_dns_answer(buffer)
-      rescue Resolv::DNS::DecodeError => e
+      code, result = Resolver.decode_dns_answer(buffer)
+
+      case code
+      when :ok
+        parse_addresses(result)
+      when :no_domain_found
+        # Indicates no such domain was found.
         hostname, connection = @queries.first
-        @queries.delete(hostname)
-        @timeouts.delete(hostname)
+        reset_hostname(hostname)
+
         @connections.delete(connection)
-        ex = NativeResolveError.new(connection, connection.origin.host, e.message)
+        raise NativeResolveError.new(connection, connection.origin.host, "name or service not known (#{hostname})")
+      when :message_truncated
+        # TODO: what to do if it's already tcp??
+        return if @socket_type == :tcp
+
+        @socket_type = :tcp
+
+        hostname, _ = @queries.first
+        reset_hostname(hostname)
+        transition(:closed)
+      when :dns_error
+        hostname, connection = @queries.first
+        reset_hostname(hostname)
+        @connections.delete(connection)
+        ex = NativeResolveError.new(connection, connection.origin.host, "unknown DNS error (error code #{result})")
         ex.set_backtrace(e.backtrace)
         raise ex
+      when :decode_error
+        hostname, connection = @queries.first
+        reset_hostname(hostname)
+        @connections.delete(connection)
+        ex = NativeResolveError.new(connection, connection.origin.host, result.message)
+        ex.set_backtrace(result.backtrace)
+        raise ex
       end
+    end
 
-      if addresses.nil?
-        # Indicates no such domain was found.
-        hostname, connection = @queries.first
-        @queries.delete(hostname)
-        @timeouts.delete(hostname)
-
-        unless @queries.value?(connection)
-          @connections.delete(connection)
-          raise NativeResolveError.new(connection, connection.origin.host)
-        end
-      elsif addresses.empty?
+    def parse_addresses(addresses)
+      if addresses.empty?
         # no address found, eliminate candidates
         _, connection = @queries.first
-        candidates = @queries.select { |_, conn| connection == conn }.keys
-        @queries.delete_if { |hs, _| candidates.include?(hs) }
-        @timeouts.delete_if { |hs, _| candidates.include?(hs) }
+        reset_hostname(hostname)
         @connections.delete(connection)
         raise NativeResolveError.new(connection, connection.origin.host)
       else
         address = addresses.first
         name = address["name"]
 
         connection = @queries.delete(name)
 
         unless connection
+          orig_name = name
           # absolute name
           name_labels = Resolv::DNS::Name.create(name).to_a
-          name = @queries.keys.first { |hname| name_labels == Resolv::DNS::Name.create(hname).to_a }
+          name = @queries.each_key.first { |hname| name_labels == Resolv::DNS::Name.create(hname).to_a }
 
           # probably a retried query for which there's an answer
-          return unless name
+          unless name
+            @timeouts.delete(orig_name)
+            return
+          end
 
           address["name"] = name
           connection = @queries.delete(name)
         end
 
-        # eliminate other candidates
-        @queries.delete_if { |_, conn| connection == conn }
-
         if address.key?("alias") # CNAME
           # clean up intermediate queries
           @timeouts.delete(name) unless connection.origin.host == name
 
           if catch(:coalesced) { early_resolve(connection, hostname: address["alias"]) }
@@ -263,10 +326,11 @@
       resolve
     end
 
     def resolve(connection = @connections.first, hostname = nil)
       raise Error, "no URI to resolve" unless connection
+
       return unless @write_buffer.empty?
 
       hostname ||= @queries.key(connection)
 
       if hostname.nil?
@@ -279,35 +343,49 @@
       else
         @queries[hostname] = connection
       end
       log { "resolver: query #{@record_type.name.split("::").last} for #{hostname}" }
       begin
-        @write_buffer << Resolver.encode_dns_query(hostname, type: @record_type)
+        @write_buffer << encode_dns_query(hostname)
       rescue Resolv::DNS::EncodeError => e
         emit_resolve_error(connection, hostname, e)
       end
     end
 
+    def encode_dns_query(hostname)
+      message_id = Resolver.generate_id
+      msg = Resolver.encode_dns_query(hostname, type: @record_type, message_id: message_id)
+      msg[0, 2] = [msg.size, message_id].pack("nn") if @socket_type == :tcp
+      msg
+    end
+
     def generate_candidates(name)
       return [name] if name.end_with?(".")
 
       candidates = []
       name_parts = name.scan(/[^.]+/)
       candidates = [name] if @ndots <= name_parts.size - 1
       candidates.concat(@search.map { |domain| [*name_parts, *domain].join(".") })
-      candidates << name unless candidates.include?(name)
+      fname = "#{name}."
+      candidates << fname unless candidates.include?(fname)
 
       candidates
     end
 
     def build_socket
-      return if @io
-
       ip, port = @nameserver[@ns_index]
       port ||= DNS_PORT
-      log { "resolver: server: #{ip}:#{port}..." }
-      @io = UDP.new(ip, port, @options)
+
+      case @socket_type
+      when :udp
+        log { "resolver: server: udp://#{ip}:#{port}..." }
+        UDP.new(ip, port, @options)
+      when :tcp
+        log { "resolver: server: tcp://#{ip}:#{port}..." }
+        origin = URI("tcp://#{ip}:#{port}")
+        TCP.new(origin, [ip], @options)
+      end
     end
 
     def transition(nextstate)
       case nextstate
       when :idle
@@ -317,11 +395,11 @@
         end
         @timeouts.clear
       when :open
         return unless @state == :idle
 
-        build_socket
+        @io ||= build_socket
 
         @io.connect
         return unless @io.connected?
 
         resolve if @queries.empty? && !@connections.empty?
@@ -343,8 +421,22 @@
       else
         @queries.each do |host, connection|
           emit_resolve_error(connection, host, error)
         end
       end
+    end
+
+    def reset_hostname(hostname, reset_candidates: true)
+      @timeouts.delete(hostname)
+      connection = @queries.delete(hostname)
+      @timeouts.delete(hostname)
+
+      return unless connection && reset_candidates
+
+      # eliminate other candidates
+      candidates = @queries.select { |_, conn| connection == conn }.keys
+      @queries.delete_if { |h, _| candidates.include?(h) }
+      # reset timeouts
+      @timeouts.delete_if { |h, _| candidates.include?(h) }
     end
   end
 end