From 2cc019f9210d1174827ce9366b0b1da84f5c4805 Mon Sep 17 00:00:00 2001
From: Manabu Niseki <manabu.niseki@gmail.com>
Date: Wed, 29 Jul 2020 09:33:56 +0900
Subject: [PATCH] refactor: use async-dns

Use async-dns instead of Parallel
---
 lib/rogue_one/detector.rb | 52 +++++++++++++++++++++++++++++++++------
 lib/rogue_one/ping.rb     | 17 ++++++++++---
 rogue_one.gemspec         |  1 +
 spec/detector_spec.rb     |  1 -
 4 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/lib/rogue_one/detector.rb b/lib/rogue_one/detector.rb
index 9561b33..11f70d8 100644
--- a/lib/rogue_one/detector.rb
+++ b/lib/rogue_one/detector.rb
@@ -1,7 +1,13 @@
 # frozen_string_literal: true
 
+require "async"
+require "async/barrier"
+require "async/dns"
+require "async/reactor"
+require "async/semaphore"
+require "resolv"
 require "yaml"
-require "parallel"
+require "etc"
 
 module RogueOne
   class Detector
@@ -9,6 +15,7 @@ class Detector
     attr_reader :default_list
     attr_reader :custom_list
     attr_reader :verbose
+    attr_reader :max_concurrency
 
     GOOGLE_PUBLIC_DNS = "8.8.8.8"
 
@@ -19,6 +26,7 @@ def initialize(target:, default_list: "alexa", custom_list: nil, threshold: nil,
       @threshold = threshold
       @verbose = verbose
 
+      @max_concurrency = Etc.nprocessors * 2
       @memo = {}
       @verbose_memo = nil
     end
@@ -80,11 +88,16 @@ def occurrences
     def inspect
       return unless @memo.empty?
 
-      results = Parallel.map(domains) do |domain|
-        normal_results = normal_resolver.get_resources(domain, "A")
-        target_result = target_resolver.get_resource(domain, "A")
+      # read domains outside of the async blocks
+      domains
 
-        [domain, target_result] if target_result && !normal_results.include?(target_result)
+      normal = bulk_resolve(normal_resolver, domains)
+      resolutions = bulk_resolve(target_resolver, domains)
+
+      results = resolutions.map do |domain, addresses|
+        normal_addresses = normal.dig(domain) || []
+        address = (addresses || []).first
+        [domain, address] if address && !normal_addresses.include?(address)
       end.compact.to_h
 
       @memo = results.values.group_by(&:itself).map { |k, v| [k, v.length] }.to_h
@@ -116,12 +129,37 @@ def read_domains(path)
       raise ArgumentError, "Inputted an invalid list. Please input a list as an YAML file." unless list.valid_format?
     end
 
+    def bulk_resolve(resolver, domains)
+      results = []
+      Async do
+        barrier = Async::Barrier.new
+        semaphore = Async::Semaphore.new(max_concurrency, parent: barrier)
+
+        domains.each do |domain|
+          semaphore.async do
+            records = resolver.query(domain, Resolv::DNS::Resource::IN::A).answer.flatten
+
+            a_records = records.select do |record|
+              record.is_a? Resolv::DNS::Resource::IN::A
+            end
+
+            addresses = a_records.map do |record|
+              record.respond_to?(:address) ? record.address.to_s : nil
+            end.compact
+
+            results << [domain, addresses]
+          end
+        end
+      end
+      results.to_h.compact
+    end
+
     def normal_resolver
-      @normal_resolver ||= Resolver.new(nameserver: GOOGLE_PUBLIC_DNS)
+      Async::DNS::Resolver.new([[:udp, GOOGLE_PUBLIC_DNS, 53], [:tcp, GOOGLE_PUBLIC_DNS, 53]])
     end
 
     def target_resolver
-      @target_resolver ||= Resolver.new(nameserver: target)
+      Async::DNS::Resolver.new([[:udp, target, 53], [:tcp, target, 53]])
     end
   end
 end
diff --git a/lib/rogue_one/ping.rb b/lib/rogue_one/ping.rb
index 83e366d..b59da67 100644
--- a/lib/rogue_one/ping.rb
+++ b/lib/rogue_one/ping.rb
@@ -1,16 +1,27 @@
 # frozen_string_literal: true
 
+require "resolv"
+
 module RogueOne
   class Ping
     attr_reader :resolver
+    attr_reader :nameserver
 
     def initialize(nameserver)
-      @resolver = Resolver.new(nameserver: nameserver)
+      @nameserver = nameserver
+      @resolver = Resolv::DNS.new(nameserver: [nameserver])
+      @resolver.timeouts = 5
+    end
+
+    def get_a_record
+      resolver.getresource("example.com", Resolv::DNS::Resource::IN::A)
+    rescue Resolv::ResolvError => _e
+      nil
     end
 
     def pong?
-      result = resolver.get_resource("example.com", "A")
-      raise Error, "DNS resolve error: there is no resopnse from #{resolver.nameserver}" unless result
+      result = get_a_record
+      raise Error, "DNS resolve error: there is no resopnse from #{nameserver}" unless result
 
       true
     end
diff --git a/rogue_one.gemspec b/rogue_one.gemspec
index 2e4cb84..fedefbc 100644
--- a/rogue_one.gemspec
+++ b/rogue_one.gemspec
@@ -29,6 +29,7 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rake", "~> 13.0"
   spec.add_development_dependency "rspec", "~> 3.9"
 
+  spec.add_dependency "async-dns", "~> 1.2"
   spec.add_dependency "parallel", "~> 1.19"
   spec.add_dependency "thor", "~> 1.0"
 end
diff --git a/spec/detector_spec.rb b/spec/detector_spec.rb
index fcf57e2..a8cdad6 100644
--- a/spec/detector_spec.rb
+++ b/spec/detector_spec.rb
@@ -6,7 +6,6 @@
   describe "#report" do
     before do
       allow(subject).to receive(:top_100_domains).and_return(%w(google.com))
-      allow(Parallel).to receive(:processor_count).and_return(0)
     end
 
     let(:report) { subject.report }