From bc22394fdd4f26dddd576840b49d7c25802cda7d Mon Sep 17 00:00:00 2001
From: Adam Cooke <me@adamcooke.io>
Date: Tue, 27 Jul 2021 14:07:38 +0000
Subject: [PATCH] fix(smtp_server): refactor mx lookups to randomly order mx
 records with the same priority

closes #1408
closes #1405
---
 config/initializers/inflections.rb |  1 +
 lib/postal.rb                      |  3 ++-
 lib/postal/mx_lookup.rb            | 34 ++++++++++++++++++++++++++++++
 lib/postal/smtp_sender.rb          | 10 +++------
 4 files changed, 40 insertions(+), 8 deletions(-)
 create mode 100644 lib/postal/mx_lookup.rb

diff --git a/config/initializers/inflections.rb b/config/initializers/inflections.rb
index 8beea52..7cb810c 100644
--- a/config/initializers/inflections.rb
+++ b/config/initializers/inflections.rb
@@ -19,5 +19,6 @@ ActiveSupport::Inflector.inflections(:en) do |inflect|
   inflect.acronym 'UUID'
   inflect.acronym 'HTTP'
   inflect.acronym 'DB'
+  inflect.acronym 'MX'
   inflect.acronym 'DKIM'
 end
diff --git a/lib/postal.rb b/lib/postal.rb
index 2a0ecbf..06a04d8 100644
--- a/lib/postal.rb
+++ b/lib/postal.rb
@@ -19,12 +19,13 @@ module Postal
     autoload :MessageInspection
     autoload :MessageParser
     autoload :MessageRequeuer
+    autoload :MXLookup
     autoload :QueryString
     autoload :RabbitMQ
     autoload :ReplySeparator
     autoload :RspecHelpers
-    autoload :SendResult
     autoload :Sender
+    autoload :SendResult
     autoload :SMTPSender
     autoload :SMTPServer
     autoload :SpamCheck
diff --git a/lib/postal/mx_lookup.rb b/lib/postal/mx_lookup.rb
new file mode 100644
index 0000000..ed170ed
--- /dev/null
+++ b/lib/postal/mx_lookup.rb
@@ -0,0 +1,34 @@
+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
diff --git a/lib/postal/smtp_sender.rb b/lib/postal/smtp_sender.rb
index ede2b2b..0ea22bd 100644
--- a/lib/postal/smtp_sender.rb
+++ b/lib/postal/smtp_sender.rb
@@ -200,13 +200,9 @@ module Postal
 
     def servers
       @options[:servers] || self.class.relay_hosts || @servers ||= begin
-        mx_servers = []
-        Resolv::DNS.open do |dns|
-          dns.timeouts = [10,5]
-          mx_servers = dns.getresources(@domain, Resolv::DNS::Resource::IN::MX).map { |m| [m.preference.to_i, m.exchange.to_s] }.sort.map{ |m| m[1] }
-          if mx_servers.empty?
-            mx_servers = [@domain] # This will be resolved to an A or AAAA record later
-          end
+        mx_servers = MXLookup.lookup(@domain)
+        if mx_servers.empty?
+          mx_servers = [@domain] # This will be resolved to an A or AAAA record later
         end
         mx_servers
       end