Refactor ClamAV service pool

lib/discourse_antivirus/clamav.rb

@@ -12,6 +12,16 @@ def self.instance
       new(Discourse.store, DiscourseAntivirus::ClamAVServicesPool.new)
     end
 
+    def self.correctly_configured?
+      return true if Rails.env.test?
+
+      if Rails.env.production?
+        SiteSetting.antivirus_srv_record.present?
+      else
+        GlobalSetting.respond_to?(:clamav_hostname) && GlobalSetting.respond_to?(:clamav_port)
+      end
+    end
+
     def initialize(store, clamav_services_pool)
       @store = store
       @clamav_services_pool = clamav_services_pool
@@ -23,10 +33,8 @@ def versions
 
     def update_versions
       antivirus_versions =
-        clamav_services_pool.all_tcp_sockets.map do |tcp_socket|
-          antivirus_version =
-            with_session(tcp_socket) { |socket| write_in_socket(socket, "zVERSION\0") }
-
+        clamav_services_pool.instances.map do |instance|
+          antivirus_version = with_session(instance) { |s| write_in_socket(s, "zVERSION\0") }
           antivirus_version = clean_msg(antivirus_version).split("/")
 
           {
@@ -41,10 +49,7 @@ def update_versions
     end
 
     def accepting_connections?
-      # At least a server must be online
-      sockets = clamav_services_pool.all_tcp_sockets
-      available = sockets.any? { |s| target_online?(s) }
-      sockets.each { |s| s&.close }
+      available = online_target.present?
 
       update_status(!available)
 
@@ -67,14 +72,9 @@ def scan_upload(upload)
     end
 
     def scan_file(file)
-      # Find a random online server and close sockets to the other ones
-      sockets = clamav_services_pool.all_tcp_sockets
-      socket = sockets.shuffle.find { |s| target_online?(s) }
-      sockets.each { |s| s&.close if socket != s }
-
       scan_response =
         begin
-          with_session(socket) { |s| stream_file(s, file) }
+          with_session { |socket| stream_file(socket, file) }
         rescue StandardError => e
           e.message
         end
@@ -94,21 +94,21 @@ def update_status(unavailable)
       PluginStore.set(PLUGIN_NAME, UNAVAILABLE, unavailable)
     end
 
-    def target_online?(socket)
-      return false if socket.nil?
-
-      write_in_socket(socket, "zPING\0")
-
-      response = get_full_response_from(socket)
+    def online_target
+      clamav_services_pool.instances.find do |instance|
+        ping_result = with_session(instance) { |s| write_in_socket(s, "zPING\0") }
 
-      clean_msg(response) == "PONG"
+        clean_msg(ping_result) == "PONG"
+      end
     end
 
     def clean_msg(raw)
       raw.gsub("1: ", "").strip
     end
 
-    def with_session(socket)
+    def with_session(instance = online_target)
+      socket = instance&.connect!
+
       raise "ERROR: socket cannot be open" if !socket
 
       write_in_socket(socket, "zIDSESSION\0")

lib/discourse_antivirus/clamav_service.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module DiscourseAntivirus
+  class ClamAVService
+    def initialize(hostname, port)
+      @hostname = hostname
+      @port = port
+    end
+
+    def connect!
+      begin
+        TCPSocket.new(@hostname, @port, connect_timeout: 3)
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end

lib/discourse_antivirus/clamav_services_pool.rb

@@ -2,49 +2,23 @@
 
 module DiscourseAntivirus
   class ClamAVServicesPool
-    UNAVAILABLE = "unavailable"
-
-    def self.correctly_configured?
-      return true if Rails.env.test?
-
-      if Rails.env.production?
-        SiteSetting.antivirus_srv_record.present?
-      else
-        GlobalSetting.respond_to?(:clamav_hostname) && GlobalSetting.respond_to?(:clamav_port)
-      end
-    end
-
-    def all_tcp_sockets
-      service_instance.targets.map { |target| build_socket(target) }
+    def instances
+      @instances ||=
+        servers
+          .filter { |server| server&.hostname.present? && server&.port.present? }
+          .map { |server| ClamAVService.new(server.hostname, server.port) }
     end
 
     private
 
-    def build_socket(target)
-      return if target.nil?
-      return if target.hostname.blank?
-      return if target.port.blank?
-
-      begin
-        TCPSocket.new(target.hostname, target.port, connect_timeout: 3)
-      rescue StandardError
-        nil
-      end
-    end
-
-    def service_instance
-      @instance ||=
+    def servers
+      @servers ||=
         if Rails.env.production?
-          DNSSD::ServiceInstance.new(Resolv::DNS::Name.create(SiteSetting.antivirus_srv_record))
+          DNSSD::ServiceInstance.new(
+            Resolv::DNS::Name.create(SiteSetting.antivirus_srv_record),
+          ).targets
         else
-          OpenStruct.new(
-            targets: [
-              OpenStruct.new(
-                hostname: GlobalSetting.clamav_hostname,
-                port: GlobalSetting.clamav_port,
-              ),
-            ],
-          )
+          [OpenStruct.new(hostname: GlobalSetting.clamav_hostname, port: GlobalSetting.clamav_port)]
         end
     end
   end

lib/validators/enable_discourse_antivirus_validator.rb

@@ -8,7 +8,7 @@ def initialize(opts = {})
   def valid_value?(val)
     return true if val == "f"
 
-    DiscourseAntivirus::ClamAVServicesPool.correctly_configured?
+    DiscourseAntivirus::ClamAV.correctly_configured?
   end
 
   def error_message

plugin.rb

@@ -19,27 +19,19 @@
 add_admin_route "antivirus.title", "antivirus"
 
 after_initialize do
-  require_dependency File.expand_path(
-                       "../app/controllers/discourse_antivirus/antivirus_controller.rb",
-                       __FILE__,
-                     )
-  require_dependency File.expand_path(
-                       "../lib/discourse_antivirus/clamav_services_pool.rb",
-                       __FILE__,
-                     )
-  require_dependency File.expand_path("../lib/discourse_antivirus/clamav.rb", __FILE__)
-  require_dependency File.expand_path("../lib/discourse_antivirus/background_scan.rb", __FILE__)
-  require_dependency File.expand_path("../models/scanned_upload.rb", __FILE__)
-  require_dependency File.expand_path("../models/reviewable_upload.rb", __FILE__)
-  require_dependency File.expand_path("../serializers/reviewable_upload_serializer.rb", __FILE__)
-  require_dependency File.expand_path("../jobs/scheduled/scan_batch.rb", __FILE__)
-  require_dependency File.expand_path("../jobs/scheduled/create_scanned_uploads.rb", __FILE__)
-  require_dependency File.expand_path("../jobs/scheduled/fetch_antivirus_version.rb", __FILE__)
-  require_dependency File.expand_path(
-                       "../jobs/scheduled/remove_orphaned_scanned_uploads.rb",
-                       __FILE__,
-                     )
-  require_dependency File.expand_path("../jobs/scheduled/flag_quarantined_uploads.rb", __FILE__)
+  require_relative "app/controllers/discourse_antivirus/antivirus_controller.rb"
+  require_relative "lib/discourse_antivirus/clamav_services_pool.rb"
+  require_relative "lib/discourse_antivirus/clamav_service.rb"
+  require_relative "lib/discourse_antivirus/clamav.rb"
+  require_relative "lib/discourse_antivirus/background_scan.rb"
+  require_relative "models/scanned_upload.rb"
+  require_relative "models/reviewable_upload.rb"
+  require_relative "serializers/reviewable_upload_serializer.rb"
+  require_relative "jobs/scheduled/scan_batch.rb"
+  require_relative "jobs/scheduled/create_scanned_uploads.rb"
+  require_relative "jobs/scheduled/fetch_antivirus_version.rb"
+  require_relative "jobs/scheduled/remove_orphaned_scanned_uploads.rb"
+  require_relative "jobs/scheduled/flag_quarantined_uploads.rb"
 
   register_reviewable_type ReviewableUpload
 
@@ -76,10 +68,8 @@
   end
 
   if defined?(::DiscoursePrometheus)
-    require_dependency File.expand_path(
-                         "../lib/discourse_antivirus/clamav_health_metric.rb",
-                         __FILE__,
-                       )
+    require_relative "lib/discourse_antivirus/clamav_health_metric.rb"
+
     DiscoursePluginRegistry.register_global_collector(DiscourseAntivirus::ClamAVHealthMetric, self)
   end