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 @@ class DomainsController < ApplicationController
     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 @@ module HasDNSChecks
   #
 
   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 @@ module HasDNSChecks
 
   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 @@ module HasDNSChecks
   #
 
   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 @@ module HasDNSChecks
   #
 
   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 @@ class Domain < ApplicationRecord
     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 @@ class Domain < ApplicationRecord
   end
 
   def dkim_key
+    return nil unless dkim_private_key
+
     @dkim_key ||= OpenSSL::PKey::RSA.new(dkim_private_key)
   end
 
@@ -114,28 +118,37 @@ class Domain < ApplicationRecord
   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 @@ class Domain < ApplicationRecord
   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)
+    result = resolver.txt(name)
+
+    if result.include?(dns_verification_string)
       self.verified_at = Time.now
-      save
-    else
-      false
+      return save
     end
-  end
 
-  private
-
-  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?
-    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 @@ class TrackDomain < ApplicationRecord
   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