From 62d65504f66e0361d178ea882c97bdcfbc426c25 Mon Sep 17 00:00:00 2001
From: Matt Jankowski <matt@jankowski.online>
Date: Thu, 14 Nov 2024 08:47:29 -0500
Subject: [PATCH] Add `DomainResource` class to wrap MX lookup/normalize
 (#32864)

---
 .../admin/email_domain_blocks_controller.rb   |  5 +----
 app/lib/domain_resource.rb                    | 22 +++++++++++++++++++
 app/models/user.rb                            |  8 +------
 .../admin/email_domain_blocks/new.html.haml   |  4 ++--
 lib/mastodon/cli/email_domain_blocks.rb       |  7 +-----
 spec/lib/domain_resource_spec.rb              | 19 ++++++++++++++++
 6 files changed, 46 insertions(+), 19 deletions(-)
 create mode 100644 app/lib/domain_resource.rb
 create mode 100644 spec/lib/domain_resource_spec.rb

diff --git a/app/controllers/admin/email_domain_blocks_controller.rb b/app/controllers/admin/email_domain_blocks_controller.rb
index fe822d8c99..9501ebd63a 100644
--- a/app/controllers/admin/email_domain_blocks_controller.rb
+++ b/app/controllers/admin/email_domain_blocks_controller.rb
@@ -58,10 +58,7 @@ module Admin
     private
 
     def set_resolved_records
-      Resolv::DNS.open do |dns|
-        dns.timeouts = 5
-        @resolved_records = dns.getresources(@email_domain_block.domain, Resolv::DNS::Resource::IN::MX).to_a
-      end
+      @resolved_records = DomainResource.new(@email_domain_block.domain).mx
     end
 
     def resource_params
diff --git a/app/lib/domain_resource.rb b/app/lib/domain_resource.rb
new file mode 100644
index 0000000000..59a29d8797
--- /dev/null
+++ b/app/lib/domain_resource.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+class DomainResource
+  attr_reader :domain
+
+  RESOLVE_TIMEOUT = 5
+
+  def initialize(domain)
+    @domain = domain
+  end
+
+  def mx
+    Resolv::DNS.open do |dns|
+      dns.timeouts = RESOLVE_TIMEOUT
+      dns
+        .getresources(domain, Resolv::DNS::Resource::IN::MX)
+        .to_a
+        .map { |mx| mx.exchange.to_s }
+        .compact_blank
+    end
+  end
+end
diff --git a/app/models/user.rb b/app/models/user.rb
index 9a215669b9..f717dcd860 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -457,13 +457,7 @@ class User < ApplicationRecord
 
     # Doing this conditionally is not very satisfying, but this is consistent
     # with the MX records validations we do and keeps the specs tractable.
-    unless self.class.skip_mx_check?
-      Resolv::DNS.open do |dns|
-        dns.timeouts = 5
-
-        records = dns.getresources(domain, Resolv::DNS::Resource::IN::MX).to_a.map { |e| e.exchange.to_s }.compact_blank
-      end
-    end
+    records = DomainResource.new(domain).mx unless self.class.skip_mx_check?
 
     EmailDomainBlock.requires_approval?(records + [domain], attempt_ip: sign_up_ip)
   end
diff --git a/app/views/admin/email_domain_blocks/new.html.haml b/app/views/admin/email_domain_blocks/new.html.haml
index 2dfdca9376..4db8fbe5e5 100644
--- a/app/views/admin/email_domain_blocks/new.html.haml
+++ b/app/views/admin/email_domain_blocks/new.html.haml
@@ -30,12 +30,12 @@
             %label.batch-table__row__select.batch-table__row__select--aligned.batch-checkbox
               = f.input_field :other_domains,
                               as: :boolean,
-                              checked_value: record.exchange.to_s,
+                              checked_value: record,
                               include_hidden: false,
                               multiple: true
             .batch-table__row__content.pending-account
               .pending-account__header
-                %samp= record.exchange.to_s
+                %samp= record
                 %br
                 = t('admin.email_domain_blocks.dns.types.mx')
 
diff --git a/lib/mastodon/cli/email_domain_blocks.rb b/lib/mastodon/cli/email_domain_blocks.rb
index a830ca3661..0cc9ccb705 100644
--- a/lib/mastodon/cli/email_domain_blocks.rb
+++ b/lib/mastodon/cli/email_domain_blocks.rb
@@ -45,12 +45,7 @@ module Mastodon::CLI
         end
 
         other_domains = []
-        if options[:with_dns_records]
-          Resolv::DNS.open do |dns|
-            dns.timeouts = 5
-            other_domains = dns.getresources(domain, Resolv::DNS::Resource::IN::MX).to_a.map { |e| e.exchange.to_s }.compact_blank
-          end
-        end
+        other_domains = DomainResource.new(domain).mx if options[:with_dns_records]
 
         email_domain_block = EmailDomainBlock.new(domain: domain, other_domains: other_domains)
         email_domain_block.save!
diff --git a/spec/lib/domain_resource_spec.rb b/spec/lib/domain_resource_spec.rb
new file mode 100644
index 0000000000..0d239fd9de
--- /dev/null
+++ b/spec/lib/domain_resource_spec.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+RSpec.describe DomainResource do
+  describe '#mx' do
+    subject { described_class.new(domain) }
+
+    let(:domain) { 'example.host' }
+    let(:exchange) { 'mx.host' }
+
+    before { configure_mx(domain: domain, exchange: exchange) }
+
+    it 'returns array of hostnames' do
+      expect(subject.mx)
+        .to eq([exchange])
+    end
+  end
+end