refactor: refactor DNS resolution

This commit also adds some of tests for the Domain model. It was during the writing of these tests that the DNS resolution refactoring requirement became apparent.

app/controllers/domains_controller.rb

@@ -71,7 +71,7 @@ def verify
     when "Email"
       if params[:code]
         if @domain.verification_token == params[:code].to_s.strip
-          @domain.verify
+          @domain.mark_as_verified
           redirect_to_with_json [:setup, organization, @server, @domain], notice: "#{@domain.name} has been verified successfully. You now need to configure your DNS records."
         else
           respond_to do |wants|

app/models/concerns/has_dns_checks.rb

@@ -43,8 +43,8 @@ def check_dns(source = :manual)
   #
 
   def check_spf_record
-    result = resolver.getresources(name, Resolv::DNS::Resource::IN::TXT)
-    spf_records = result.map(&:data).grep(/\Av=spf1/)
+    result = resolver.txt(name)
+    spf_records = result.grep(/\Av=spf1/)
     if spf_records.empty?
       self.spf_status = "Missing"
       self.spf_error = "No SPF record exists for this domain"
@@ -73,8 +73,7 @@ def check_spf_record!
 
   def check_dkim_record
     domain = "#{dkim_record_name}.#{name}"
-    result = resolver.getresources(domain, Resolv::DNS::Resource::IN::TXT)
-    records = result.map(&:data)
+    records = resolver.txt(domain)
     if records.empty?
       self.dkim_status = "Missing"
       self.dkim_error = "No TXT records were returned for #{domain}"
@@ -104,8 +103,7 @@ def check_dkim_record!
   #
 
   def check_mx_records
-    result = resolver.getresources(name, Resolv::DNS::Resource::IN::MX)
-    records = result.map(&:exchange)
+    records = resolver.mx(name).map(&:last)
     if records.empty?
       self.mx_status = "Missing"
       self.mx_error = "There are no MX records for #{name}"
@@ -134,8 +132,7 @@ def check_mx_records!
   #
 
   def check_return_path_record
-    result = resolver.getresources(return_path_domain, Resolv::DNS::Resource::IN::CNAME)
-    records = result.map { |r| r.name.to_s.downcase }
+    records = resolver.cname(return_path_domain)
     if records.empty?
       self.return_path_status = "Missing"
       self.return_path_error = "There is no return path record at #{return_path_domain}"

app/models/domain.rb

@@ -77,7 +77,9 @@ def verified?
     verified_at.present?
   end
 
-  def verify
+  def mark_as_verified
+    return false if verified?
+
     self.verified_at = Time.now
     save!
   end
@@ -94,6 +96,8 @@ def generate_dkim_key
   end
 
   def dkim_key
+    return nil unless dkim_private_key
+
     @dkim_key ||= OpenSSL::PKey::RSA.new(dkim_private_key)
   end
 
@@ -114,28 +118,37 @@ def spf_record
   end
 
   def dkim_record
+    return if dkim_key.nil?
+
     public_key = dkim_key.public_key.to_s.gsub(/-+[A-Z ]+-+\n/, "").gsub(/\n/, "")
     "v=DKIM1; t=s; h=sha256; p=#{public_key};"
   end
 
   def dkim_identifier
+    return nil unless dkim_identifier_string
+
     Postal.config.dns.dkim_identifier + "-#{dkim_identifier_string}"
   end
 
   def dkim_record_name
-    "#{dkim_identifier}._domainkey"
+    identifier = dkim_identifier
+    return if identifier.nil?
+
+    "#{identifier}._domainkey"
   end
 
   def return_path_domain
     "#{Postal.config.dns.custom_return_path_prefix}.#{name}"
   end
 
-  def nameservers
-    @nameservers ||= get_nameservers
-  end
-
+  # Returns a DNSResolver instance that can be used to perform DNS lookups needed for
+  # the verification and DNS checking for this domain.
+  #
+  # @return [DNSResolver]
   def resolver
-    @resolver ||= Postal.config.general.use_local_ns_for_domains? ? Resolv::DNS.new : Resolv::DNS.new(nameserver: nameservers)
+    return DNSResolver.local if Postal.config.general.use_local_ns_for_domains?
+
+    @resolver ||= DNSResolver.for_domain(name)
   end
 
   def dns_verification_string
@@ -145,32 +158,14 @@ def dns_verification_string
   def verify_with_dns
     return false unless verification_method == "DNS"
 
-    result = resolver.getresources(name, Resolv::DNS::Resource::IN::TXT)
-    if result.map { |d| d.data.to_s.strip }.include?(dns_verification_string)
-      self.verified_at = Time.now
-      save
-    else
-      false
-    end
-  end
-
-  private
+    result = resolver.txt(name)
 
-  def get_nameservers
-    local_resolver = Resolv::DNS.new
-    ns_records = []
-    parts = name.split(".")
-    (parts.size - 1).times do |n|
-      d = parts[n, parts.size - n + 1].join(".")
-      ns_records = local_resolver.getresources(d, Resolv::DNS::Resource::IN::NS).map { |s| s.name.to_s }
-      break if ns_records.present?
+    if result.include?(dns_verification_string)
+      self.verified_at = Time.now
+      return save
     end
-    return [] if ns_records.blank?
-
-    ns_records = ns_records.map { |r| local_resolver.getresources(r, Resolv::DNS::Resource::IN::A).map { |s| s.address.to_s } }.flatten
-    return [] if ns_records.blank?
 
-    ns_records
+    false
   end
 
 end

app/models/track_domain.rb

@@ -54,8 +54,7 @@ def dns_ok?
   end
 
   def check_dns
-    result = domain.resolver.getresources(full_name, Resolv::DNS::Resource::IN::CNAME)
-    records = result.map { |r| r.name.to_s.downcase }
+    records = domain.resolver.cname(full_name)
     if records.empty?
       self.dns_status = "Missing"
       self.dns_error = "There is no record at #{full_name}"

app/util/dns_resolver.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+class DNSResolver
+
+  attr_reader :nameservers
+  attr_reader :timeout
+
+  def initialize(nameservers: nil, timeout: 5)
+    @nameservers = nameservers
+    @timeout = timeout
+  end
+
+  # Return all A records for the given name
+  #
+  # @param [String] name
+  # @return [Array<String>]
+  def a(name)
+    dns do |dns|
+      dns.getresources(name, Resolv::DNS::Resource::IN::A).map do |s|
+        s.address.to_s
+      end
+    end
+  end
+
+  # Return all AAAA records for the given name
+  #
+  # @param [String] name
+  # @return [Array<String>]
+  def aaaa(name)
+    dns do |dns|
+      dns.getresources(name, Resolv::DNS::Resource::IN::AAAA).map do |s|
+        s.address.to_s
+      end
+    end
+  end
+
+  # Return all TXT records for the given name
+  #
+  # @param [String] name
+  # @return [Array<String>]
+  def txt(name)
+    dns do |dns|
+      dns.getresources(name, Resolv::DNS::Resource::IN::TXT).map do |s|
+        s.data.to_s.strip
+      end
+    end
+  end
+
+  # Return all CNAME records for the given name
+  #
+  # @param [String] name
+  # @return [Array<String>]
+  def cname(name)
+    dns do |dns|
+      dns.getresources(name, Resolv::DNS::Resource::IN::CNAME).map do |s|
+        s.name.to_s.downcase
+      end
+    end
+  end
+
+  # Return all MX records for the given name
+  #
+  # @param [String] name
+  # @return [Array<Array<Integer, String>>]
+  def mx(name)
+    dns do |dns|
+      records = dns.getresources(name, Resolv::DNS::Resource::IN::MX).map do |m|
+        [m.preference.to_i, m.exchange.to_s]
+      end
+      records.sort do |a, b|
+        if a[0] == b[0]
+          [-1, 1].sample
+        else
+          a[0] <=> b[0]
+        end
+      end
+    end
+  end
+
+  # Return the effective nameserver names for a given domain name.
+  #
+  # @param [String] name
+  # @return [Array<String>]
+  def effective_ns(name)
+    records = []
+    dns do |dns|
+      parts = name.split(".")
+      (parts.size - 1).times do |n|
+        d = parts[n, parts.size - n + 1].join(".")
+
+        records = dns.getresources(d, Resolv::DNS::Resource::IN::NS).map do |s|
+          s.name.to_s
+        end
+
+        break if records.present?
+      end
+    end
+
+    records
+  end
+
+  # Return the hostname for a given IP address.
+  # Returns the IP address itself if no hostname can be determined.
+  #
+  # @param [String] ip_address
+  # @return [String]
+  def ip_to_hostname(ip_address)
+    dns do |dns|
+      dns.getname(ip_address)&.to_s
+    end
+  rescue Resolv::ResolvError
+    ip_address
+  end
+
+  private
+
+  def dns
+    Resolv::DNS.open(nameserver: @nameservers || []) do |dns|
+      dns.timeouts = [@timeout, @timeout / 2]
+      yield dns
+    end
+  end
+
+  class << self
+
+    # Return a resolver which will use the nameservers for the given domain
+    #
+    # @param [String] name
+    # @return [DNSResolver]
+    def for_domain(name)
+      resolver = new
+      nameservers = resolver.effective_ns(name)
+      ips = nameservers.map do |ns|
+        resolver.a(ns)
+      end.flatten.uniq
+      new(nameservers: ips)
+    end
+
+    # Return a local resolver to use for lookups
+    #
+    # @return [DNSResolver]
+    def local
+      @local ||= new
+    end
+
+  end
+
+end

lib/postal/received_header.rb

@@ -19,26 +19,13 @@ def generate(server, helo, ip_address, method)
         header = "by #{our_hostname} with #{method.to_s.upcase}; #{Time.now.utc.rfc2822}"
 
         if server.nil? || server.privacy_mode == false
-          hostname = resolve_hostname(ip_address)
+          hostname = DNSResolver.local.ip_to_hostname(ip_address)
           header = "from #{helo} (#{hostname} [#{ip_address}]) #{header}"
         end
 
         header
       end
 
-      private
-
-      def resolve_hostname(ip_address)
-        Resolv::DNS.open do |dns|
-          dns.timeouts = [10, 5]
-          begin
-            dns.getname(ip_address)
-          rescue StandardError
-            ip_address
-          end
-        end
-      end
-
     end
 
   end