FIX: Scan files when a single server is online

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,11 +33,8 @@ def versions
 
     def update_versions
       antivirus_versions =
-        clamav_services_pool.all_tcp_sockets.map do |tcp_socket|
-          antivirus_version =
-            with_session(socket: tcp_socket) { |socket| write_in_socket(socket, "zVERSION\0") }
-
-          antivirus_version = clean_msg(antivirus_version).split("/")
+        clamav_services_pool.online_services.map do |service|
+          antivirus_version = service.version.split("/")
 
           {
             antivirus: antivirus_version[0],
@@ -41,38 +48,33 @@ def update_versions
     end
 
     def accepting_connections?
-      sockets = clamav_services_pool.all_tcp_sockets
-
-      if sockets.empty?
-        update_status(true)
-        return false
-      end
+      unavailable = clamav_services_pool.all_offline?
 
-      available = sockets.reduce(true) { |memo, socket| memo && target_online?(socket) }
+      PluginStore.set(PLUGIN_NAME, UNAVAILABLE, unavailable)
 
-      available.tap do |status|
-        unavailable = !status
-        update_status(unavailable)
-      end
+      !unavailable
     end
 
     def scan_upload(upload)
-      begin
-        file = get_uploaded_file(upload)
+      file = get_uploaded_file(upload)
 
-        return error_response(DOWNLOAD_FAILED) if file.nil?
+      return error_response(DOWNLOAD_FAILED) if file.nil?
 
-        scan_file(file)
-      rescue OpenURI::HTTPError
-        error_response(DOWNLOAD_FAILED)
-      rescue StandardError => e
-        Rails.logger.error("Could not scan upload #{upload.id}. Error: #{e.message}")
-        error_response(e.message)
-      end
+      scan_file(file)
+    rescue OpenURI::HTTPError
+      error_response(DOWNLOAD_FAILED)
+    rescue StandardError => e
+      Rails.logger.error("Could not scan upload #{upload.id}. Error: #{e.message}")
+      error_response(e.message)
     end
 
     def scan_file(file)
-      scan_response = with_session { |socket| stream_file(socket, file) }
+      online_service = clamav_services_pool.find_online_service
+
+      # We open one connection to check if the service is online and another
+      # to scan the file.
+      scan_response = online_service&.scan_file(file)
+      return error_response(UNAVAILABLE) unless scan_response
 
       parse_response(scan_response)
     end
@@ -85,54 +87,14 @@ def error_response(error_message)
       { error: true, found: false, message: error_message }
     end
 
-    def update_status(unavailable)
-      PluginStore.set(PLUGIN_NAME, UNAVAILABLE, unavailable)
-    end
-
-    def target_online?(socket)
-      return false if socket.nil?
-
-      ping_result = with_session(socket: socket) { |s| write_in_socket(s, "zPING\0") }
-
-      clean_msg(ping_result) == "PONG"
-    end
-
-    def clean_msg(raw)
-      raw.gsub("1: ", "").strip
-    end
-
-    def with_session(socket: clamav_services_pool.tcp_socket)
-      write_in_socket(socket, "zIDSESSION\0")
-
-      yield(socket)
-
-      write_in_socket(socket, "zEND\0")
-
-      response = get_full_response_from(socket)
-      socket.close
-      response
-    end
-
     def parse_response(scan_response)
       {
-        message: scan_response.gsub("1: stream:", "").gsub("\0", ""),
+        message: scan_response.gsub("stream:", "").gsub("\0", ""),
         found: scan_response.include?("FOUND"),
         error: scan_response.include?("ERROR"),
       }
     end
 
-    def stream_file(socket, file)
-      write_in_socket(socket, "zINSTREAM\0")
-
-      while data = file.read(2048)
-        write_in_socket(socket, [data.length].pack("N"))
-        write_in_socket(socket, data)
-      end
-
-      write_in_socket(socket, [0].pack("N"))
-      write_in_socket(socket, "")
-    end
-
     def get_uploaded_file(upload)
       if store.external?
         # Upload#filesize could be approximate.
@@ -143,31 +105,5 @@ def get_uploaded_file(upload)
         File.open(store.path_for(upload))
       end
     end
-
-    # ClamAV wants us to read/write in a non-blocking manner to prevent deadlocks.
-    # Read more about this [here](https://manpages.debian.org/testing/clamav-daemon/clamd.8.en.html#IDSESSION,)
-    #
-    # We need to peek into the socket buffer to make sure we can write/read from it,
-    # or we risk ClamAV abruptly closing the connection.
-    # For that, we use [IO#select](https://www.rubydoc.info/stdlib/core/IO.select)
-    def write_in_socket(socket, msg)
-      IO.select(nil, [socket])
-      socket.sendmsg_nonblock(msg, 0, nil)
-    end
-
-    def read_from_socket(socket)
-      IO.select([socket])
-
-      # Returns an array with the chunk as the first element
-      socket.recvmsg_nonblock(25).to_a.first.to_s
-    end
-
-    def get_full_response_from(socket)
-      buffer = ""
-
-      buffer += read_from_socket(socket) until buffer.ends_with?("\0")
-
-      buffer
-    end
   end
 end

lib/discourse_antivirus/clamav_service.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module DiscourseAntivirus
+  class ClamAVService
+    def initialize(connection_factory, hostname, port)
+      @connection_factory = connection_factory
+      @hostname = hostname
+      @port = port
+    end
+
+    def version
+      with_session { |s| write_in_socket(s, "zVERSION\0") }
+    end
+
+    def online?
+      ping_result = with_session { |s| write_in_socket(s, "zPING\0") }
+
+      ping_result == "PONG"
+    end
+
+    def scan_file(file)
+      with_session do |socket|
+        write_in_socket(socket, "zINSTREAM\0")
+
+        while data = file.read(2048)
+          write_in_socket(socket, [data.length].pack("N"))
+          write_in_socket(socket, data)
+        end
+
+        write_in_socket(socket, [0].pack("N"))
+        write_in_socket(socket, "")
+      end
+    end
+
+    private
+
+    attr_reader :connection_factory, :hostname, :port, :connection
+
+    def with_session
+      socket = connection_factory.call(hostname, port)
+      return if socket.nil?
+
+      write_in_socket(socket, "zIDSESSION\0")
+
+      yield(socket)
+
+      write_in_socket(socket, "zEND\0")
+
+      response = get_full_response_from(socket)
+      socket.close
+      response
+    end
+
+    # ClamAV wants us to read/write in a non-blocking manner to prevent deadlocks.
+    # Read more about this [here](https://manpages.debian.org/testing/clamav-daemon/clamd.8.en.html#IDSESSION,)
+    #
+    # We need to peek into the socket buffer to make sure we can write/read from it,
+    # or we risk ClamAV abruptly closing the connection.
+    # For that, we use [IO#select](https://www.rubydoc.info/stdlib/core/IO.select)
+    def write_in_socket(socket, msg)
+      IO.select(nil, [socket])
+      socket.sendmsg_nonblock(msg, 0, nil)
+    end
+
+    def get_full_response_from(socket)
+      buffer = +""
+
+      until buffer.ends_with?("\0")
+        IO.select([socket])
+
+        # Returns an array with the chunk as the first element
+        buffer << socket.recvmsg_nonblock(25).to_a.first.to_s
+      end
+
+      buffer.gsub("1: ", "").strip
+    end
+  end
+end

lib/discourse_antivirus/clamav_services_pool.rb

@@ -2,53 +2,46 @@
 
 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
+    def online_services
+      instances.select(&:online?)
     end
 
-    def tcp_socket
-      build_socket(service_instance.targets.first)
+    def all_offline?
+      instances.none?(&:online?)
     end
 
-    def all_tcp_sockets
-      service_instance.targets.map { |target| build_socket(target) }
+    def find_online_service
+      instances.find(&:online?)
     end
 
     private
 
-    def build_socket(target)
-      return if target.nil?
-      return if target.hostname.blank?
-      return if target.port.blank?
+    def connection_factory
+      @factory ||=
+        Proc.new do |hostname, port|
+          begin
+            TCPSocket.new(@hostname, @port, connect_timeout: 3)
+          rescue StandardError
+            nil
+          end
+        end
+    end
 
-      begin
-        TCPSocket.new(target.hostname, target.port)
-      rescue StandardError
-        nil
-      end
+    def instances
+      @instances ||=
+        servers
+          .filter { |server| server&.hostname.present? && server&.port.present? }
+          .map { |server| ClamAVService.new(connection_factory, server.hostname, server.port) }
     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