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)
+    result = resolver.txt(name)
-    spf_records = result.map(&:data).grep(/\Av=spf1/)
+    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 = resolver.txt(domain)
-    records = result.map(&:data)
     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 = resolver.mx(name).map(&:last)
-    records = result.map(&:exchange)
     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 = resolver.cname(return_path_domain)
-    records = result.map { |r| r.name.to_s.downcase }
     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
+  # Returns a DNSResolver instance that can be used to perform DNS lookups needed for
-    @nameservers ||= get_nameservers
+  # the verification and DNS checking for this domain.
-  end
+  #
-
+  # @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)
+    result = resolver.txt(name)
-    if result.map { |d| d.data.to_s.strip }.include?(dns_verification_string)
+
+    if result.include?(dns_verification_string)
       self.verified_at = Time.now
-      save
+      return save
-    else
-      false
     end
-  end
 
-  private
+    false
-
-  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
   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 = domain.resolver.cname(full_name)
-    records = result.map { |r| r.name.to_s.downcase }
     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/mx_lookup.rb

@@ -1,36 +0,0 @@
-# frozen_string_literal: true
-
-module Postal
-  class MXLookup
-
-    class << self
-
-      def lookup(domain)
-        records = resolve(domain)
-        records = sort(records)
-        records.map { |m| m[1] }
-      end
-
-      private
-
-      def sort(records)
-        records.sort do |a, b|
-          if a[0] == b[0]
-            [-1, 1].sample
-          else
-            a[0] <=> b[0]
-          end
-        end
-      end
-
-      def resolve(domain)
-        Resolv::DNS.open do |dns|
-          dns.timeouts = [10, 5]
-          dns.getresources(domain, Resolv::DNS::Resource::IN::MX).map { |m| [m.preference.to_i, m.exchange.to_s] }
-        end
-      end
-
-    end
-
-  end
-end

lib/postal/received_header.rb

@@ -19,26 +19,13 @@ module Postal
         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

lib/postal/smtp_sender.rb

@@ -225,7 +225,7 @@ module Postal
 
     def servers
       @options[:servers] || self.class.relay_hosts || @servers ||= begin
-        mx_servers = MXLookup.lookup(@domain)
+        mx_servers = DNSResolver.local.mx(@domain).map(&:last)
         if mx_servers.empty?
           mx_servers = [@domain] # This will be resolved to an A or AAAA record later
         end
@@ -243,16 +243,13 @@ module Postal
 
     def lookup_ip_address(type, hostname)
       records = []
-      Resolv::DNS.open do |dns|
+      case type
-        dns.timeouts = [10, 5]
+      when :a
-        case type
+        records = DNSResolver.local.a(hostname)
-        when :a
+      when :aaaa
-          records = dns.getresources(hostname, Resolv::DNS::Resource::IN::A)
+        records = DNSResolver.local.aaaa(hostname)
-        when :aaaa
-          records = dns.getresources(hostname, Resolv::DNS::Resource::IN::AAAA)
-        end
       end
-      records.first&.address&.to_s&.downcase
+      records.first&.to_s&.downcase
     end
 
     class << self

spec/lib/postal/received_header_spec.rb

@@ -4,7 +4,7 @@ require "rails_helper"
 
 describe Postal::ReceivedHeader do
   before do
-    allow(Resolv::DNS).to receive(:open).and_return("hostname.com")
+    allow(DNSResolver.local).to receive(:ip_to_hostname).and_return("hostname.com")
   end
 
   describe ".generate" do

spec/models/domain_spec.rb

@@ -0,0 +1,317 @@
+# frozen_string_literal: true
+
+require "rails_helper"
+
+describe Domain do
+  subject(:domain) { build(:domain) }
+
+  describe "relationships" do
+    it { is_expected.to belong_to(:server).optional }
+    it { is_expected.to belong_to(:owner).optional }
+    it { is_expected.to have_many(:routes) }
+    it { is_expected.to have_many(:track_domains) }
+  end
+
+  describe "validations" do
+    it { is_expected.to validate_presence_of(:name) }
+    it { is_expected.to validate_uniqueness_of(:name).scoped_to([:owner_type, :owner_id]).case_insensitive.with_message("is already added") }
+    it { is_expected.to allow_value("example.com").for(:name) }
+    it { is_expected.to allow_value("example.co.uk").for(:name) }
+    it { is_expected.to_not allow_value("EXAMPLE.COM").for(:name) }
+    it { is_expected.to_not allow_value("example.com ").for(:name) }
+    it { is_expected.to_not allow_value("example com").for(:name) }
+    it { is_expected.to validate_inclusion_of(:verification_method).in_array(Domain::VERIFICATION_METHODS) }
+  end
+
+  describe "creation" do
+    it "creates a new dkim identifier string" do
+      expect { domain.save }.to change { domain.dkim_identifier_string }.from(nil).to(match(/\A[a-zA-Z0-9]{6}\z/))
+    end
+
+    it "generates a new dkim key" do
+      expect { domain.save }.to change { domain.dkim_private_key }.from(nil).to(match(/\A-+BEGIN RSA PRIVATE KEY-+/))
+    end
+
+    it "generates a UUID" do
+      expect { domain.save }.to change { domain.uuid }.from(nil).to(/[a-f0-9-]{36}/)
+    end
+  end
+
+  describe ".verified" do
+    it "returns verified domains only" do
+      verified_domain = create(:domain)
+      create(:domain, :unverified)
+      expect(described_class.verified).to eq [verified_domain]
+    end
+  end
+
+  context "when verification method changes" do
+    context "to DNS" do
+      let(:domain) { create(:domain, :unverified, verification_method: "Email") }
+
+      it "generates a DNS suitable verification token" do
+        domain.verification_method = "DNS"
+        expect { domain.save }.to change { domain.verification_token }.from(match(/\A\d{6}\z/)).to(match(/\A[A-Za-z0-9+]{32}\z/))
+      end
+    end
+
+    context "to Email" do
+      let(:domain) { create(:domain, :unverified, verification_method: "DNS") }
+
+      it "generates an email suitable verification token" do
+        domain.verification_method = "Email"
+        expect { domain.save }.to change { domain.verification_token }.from(match(/\A[A-Za-z0-9+]{32}\z/)).to(match(/\A\d{6}\z/))
+      end
+    end
+  end
+
+  describe "#verified?" do
+    context "when the domain is verified" do
+      it "returns true" do
+        expect(domain.verified?).to be true
+      end
+    end
+
+    context "when the domain is not verified" do
+      let(:domain) { build(:domain, :unverified) }
+
+      it "returns false" do
+        expect(domain.verified?).to be false
+      end
+    end
+  end
+
+  describe "#mark_as_verified" do
+    context "when already verified" do
+      it "returns false" do
+        expect(domain.mark_as_verified).to be false
+      end
+    end
+
+    context "when unverified" do
+      let(:domain) { create(:domain, :unverified) }
+
+      it "sets the verification time" do
+        expect { domain.mark_as_verified }.to change { domain.verified_at }.from(nil).to(kind_of(Time))
+      end
+    end
+  end
+
+  describe "#parent_domains" do
+    context "at level 1" do
+      let(:domain) { build(:domain, name: "example.com") }
+
+      it "returns the current domain only" do
+        expect(domain.parent_domains).to eq ["example.com"]
+      end
+    end
+
+    context "at level 2" do
+      let(:domain) { build(:domain, name: "test.example.com") }
+
+      it "returns the current domain plus its parent" do
+        expect(domain.parent_domains).to eq ["test.example.com", "example.com"]
+      end
+    end
+
+    context "at level 3 (and higher)" do
+      let(:domain) { build(:domain, name: "sub.test.example.com") }
+
+      it "returns the current domain plus its parents" do
+        expect(domain.parent_domains).to eq ["sub.test.example.com", "test.example.com", "example.com"]
+      end
+    end
+  end
+
+  describe "#generate_dkim_key" do
+    it "generates a new dkim key" do
+      expect { domain.generate_dkim_key }.to change { domain.dkim_private_key }.from(nil).to(match(/\A-+BEGIN RSA PRIVATE KEY-+/))
+    end
+  end
+
+  describe "#dkim_key" do
+    context "when the domain has a DKIM key" do
+      let(:domain) { create(:domain) }
+
+      it "returns the dkim key as a OpenSSL::PKey::RSA" do
+        expect(domain.dkim_key).to be_a OpenSSL::PKey::RSA
+        expect(domain.dkim_key.to_s).to eq domain.dkim_private_key
+      end
+    end
+
+    context "when the domain has no DKIM key" do
+      let(:domain) { build(:domain) }
+
+      it "returns nil" do
+        expect(domain.dkim_key).to be_nil
+      end
+    end
+  end
+
+  describe "#to_param" do
+    context "when the domain has not been saved" do
+      it "returns nil" do
+        expect(domain.to_param).to be_nil
+      end
+    end
+    context "when the domain has been saved" do
+      before do
+        domain.save
+      end
+
+      it "returns the UUID" do
+        expect(domain.to_param).to eq domain.uuid
+      end
+    end
+  end
+
+  describe "#verification_email_addresses" do
+    let(:domain) { build(:domain, name: "example.com") }
+
+    it "returns the verification email addresses" do
+      expect(domain.verification_email_addresses).to eq [
+        "webmaster@example.com",
+        "postmaster@example.com",
+        "admin@example.com",
+        "administrator@example.com",
+        "hostmaster@example.com"
+      ]
+    end
+  end
+
+  describe "#spf_record" do
+    it "returns the SPF record" do
+      expect(domain.spf_record).to eq "v=spf1 a mx include:#{Postal.config.dns.spf_include} ~all"
+    end
+  end
+
+  describe "#dkim_record" do
+    context "when the domain has no DKIM key" do
+      it "returns nil" do
+        expect(domain.dkim_record).to be_nil
+      end
+    end
+
+    context "when the domain has a DKIM key" do
+      before do
+        domain.save
+      end
+
+      it "returns the DKIM record" do
+        expect(domain.dkim_record).to match(/\Av=DKIM1; t=s; h=sha256; p=.*;\z/)
+      end
+    end
+  end
+
+  describe "#dkim_identifier" do
+    context "when the domain has no dkim identifier string" do
+      it "returns nil" do
+        expect(domain.dkim_identifier).to be_nil
+      end
+    end
+
+    context "when the domain has a dkim identifier string" do
+      before do
+        domain.save
+      end
+
+      it "returns the DKIM identifier" do
+        expect(domain.dkim_identifier).to eq "#{Postal.config.dns.dkim_identifier}-#{domain.dkim_identifier_string}"
+      end
+    end
+  end
+
+  describe "#dkim_record_name" do
+    context "when the domain has no dkim identifier string" do
+      it "returns nil" do
+        expect(domain.dkim_record_name).to be_nil
+      end
+    end
+
+    context "when the domain has a dkim identifier string" do
+      before do
+        domain.save
+      end
+
+      it "returns the DKIM identifier" do
+        expect(domain.dkim_record_name).to eq "#{Postal.config.dns.dkim_identifier}-#{domain.dkim_identifier_string}._domainkey"
+      end
+    end
+  end
+
+  describe "#return_path_domain" do
+    it "returns the return path domain" do
+      expect(domain.return_path_domain).to eq "#{Postal.config.dns.custom_return_path_prefix}.#{domain.name}"
+    end
+  end
+
+  describe "#dns_verification_string" do
+    let(:domain) { create(:domain, verification_method: "DNS") }
+
+    it "returns the DNS verification string" do
+      expect(domain.dns_verification_string).to eq "#{Postal.config.dns.domain_verify_prefix} #{domain.verification_token}"
+    end
+  end
+
+  describe "#resolver" do
+    context "when the local nameservers should be used" do
+      before do
+        allow(Postal.config.general).to receive(:use_local_ns_for_domains?).and_return(true)
+      end
+
+      it "uses the local DNS" do
+        expect(domain.resolver).to eq DNSResolver.local
+      end
+    end
+
+    context "when local nameservers should not be used" do
+      it "uses the a resolver for this domain" do
+        allow(DNSResolver).to receive(:for_domain).with(domain.name).and_return(DNSResolver.new(nameservers: ["1.2.3.4"]))
+        expect(domain.resolver).to be_a DNSResolver
+        expect(domain.resolver.nameservers).to eq ["1.2.3.4"]
+      end
+    end
+  end
+
+  describe "#verify_with_dns" do
+    context "when the verification method is not DNS" do
+      let(:domain) { build(:domain, verification_method: "Email") }
+
+      it "returns false" do
+        expect(domain.verify_with_dns).to be false
+      end
+    end
+
+    context "when a TXT record is found that matches" do
+      let(:domain) { create(:domain, :unverified) }
+
+      before do
+        allow(domain.resolver).to receive(:txt).with(domain.name).and_return([domain.dns_verification_string])
+      end
+
+      it "returns true" do
+        expect(domain.verify_with_dns).to be true
+      end
+
+      it "sets the verification time" do
+        expect { domain.verify_with_dns }.to change { domain.verified_at }.from(nil).to(kind_of(Time))
+      end
+    end
+
+    context "when no TXT record is found" do
+      let(:domain) { create(:domain, :unverified) }
+
+      before do
+        allow(domain.resolver).to receive(:txt).with(domain.name).and_return(["something", "something else"])
+      end
+
+      it "returns false" do
+        expect(domain.verify_with_dns).to be false
+      end
+
+      it "does not set the verification time" do
+        expect { domain.verify_with_dns }.to_not change { domain.verified_at } # rubocop:disable Lint/AmbiguousBlockAssociation
+      end
+    end
+  end
+end