FIX: Ensure used sockets are not closed

lib/discourse_antivirus/clamav.rb

@@ -25,7 +25,7 @@ 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") }
+            with_session(tcp_socket) { |socket| write_in_socket(socket, "zVERSION\0") }
 
           antivirus_version = clean_msg(antivirus_version).split("/")
 
@@ -41,7 +41,10 @@ def update_versions
     end
 
     def accepting_connections?
-      available = clamav_services_pool.all_tcp_sockets.any? { |socket| target_online?(socket) }
+      # 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 }
 
       update_status(!available)
 
@@ -64,7 +67,16 @@ def scan_upload(upload)
     end
 
     def scan_file(file)
-      scan_response = with_session { |socket| stream_file(socket, 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) }
+      rescue StandardError => e
+        e.message
+      end
 
       parse_response(scan_response)
     end
@@ -84,18 +96,19 @@ def update_status(unavailable)
     def target_online?(socket)
       return false if socket.nil?
 
-      ping_result = with_session(socket: socket) { |s| write_in_socket(s, "zPING\0") }
+      write_in_socket(socket, "zPING\0")
+
+      response = get_full_response_from(socket)
 
-      clean_msg(ping_result) == "PONG"
+      clean_msg(response) == "PONG"
     end
 
     def clean_msg(raw)
       raw.gsub("1: ", "").strip
     end
 
-    def with_session(socket: nil)
-      socket ||= clamav_services_pool.all_tcp_sockets.shuffle.find { |s| target_online?(s) }
-      raise "no online socket found" if !socket
+    def with_session(socket)
+      raise "ERROR: socket cannot be open" if !socket
 
       write_in_socket(socket, "zIDSESSION\0")
 

plugin.rb

@@ -67,7 +67,8 @@
       antivirus = DiscourseAntivirus::ClamAV.instance
 
       if antivirus.accepting_connections?
-        is_positive = antivirus.scan_file(file)[:found]
+        response = antivirus.scan_file(file)
+        is_positive = response[:found]
 
         upload.errors.add(:base, I18n.t("scan.virus_found")) if is_positive
       end

spec/lib/discourse_antivirus/clamav_spec.rb

@@ -104,7 +104,7 @@ def assert_version_was_requested(fake_socket)
   end
 
   def assert_file_was_sent_through(fake_socket, file)
-    expected = ["zIDSESSION\0", "zINSTREAM\0"]
+    expected = ["zPING\0", "zIDSESSION\0", "zINSTREAM\0"]
 
     file.rewind
     while data = file.read(2048)

spec/plugin_spec.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "rails_helper"
+require_relative "support/fake_pool"
 require_relative "support/fake_tcp_socket"
 
 describe DiscourseAntivirus do
@@ -15,15 +16,15 @@
     let(:file) { file_from_fixtures(filename, "pdf") }
 
     it "scans regular files and adds an error if the scan result is positive" do
-      mock_antivirus(FakeTCPSocket.positive(include_pong: true))
+      mock_antivirus(FakeTCPSocket.positive)
 
       scanned_upload = UploadCreator.new(file, filename).create_for(user.id)
 
       expect(scanned_upload.errors.to_a).to contain_exactly(I18n.t("scan.virus_found"))
     end
 
     it "scans regular files but does nothing if the scan result is negative" do
-      mock_antivirus(FakeTCPSocket.negative(include_pong: true))
+      mock_antivirus(FakeTCPSocket.negative)
 
       scanned_upload = UploadCreator.new(file, filename).create_for(user.id)
 
@@ -74,7 +75,7 @@
 
     it "scans the image if the live scan images setting is enabled" do
       SiteSetting.antivirus_live_scan_images = true
-      mock_antivirus(FakeTCPSocket.positive(include_pong: true))
+      mock_antivirus(FakeTCPSocket.positive)
 
       scanned_upload = UploadCreator.new(file, filename).create_for(user.id)
 
@@ -106,7 +107,7 @@
 
   def mock_antivirus(socket)
     IO.stubs(:select).returns(true)
-    pool = OpenStruct.new(tcp_socket: socket, all_tcp_sockets: [socket])
+    pool = FakePool.new([socket])
     antivirus = DiscourseAntivirus::ClamAV.new(Discourse.store, pool)
     DiscourseAntivirus::ClamAV.expects(:instance).returns(antivirus)
   end

spec/support/fake_pool.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+class FakePool
+  def initialize(sockets)
+    @sockets = sockets
+  end
+
+  def tcp_socket
+    @sockets.first.dup
+  end
+
+  def all_tcp_sockets
+    @sockets.map(&:dup)
+  end
+end

spec/support/fake_tcp_socket.rb

@@ -1,23 +1,20 @@
 # frozen_string_literal: true
 class FakeTCPSocket
-  def self.positive(include_pong: false)
-    responses = include_pong ? ["1: PONG\0"] : []
-    responses << "1: stream: Win.Test.EICAR_HDB-1 FOUND\0"
-    new(responses)
+  def self.positive
+    new(["1: PONG\0", "1: stream: Win.Test.EICAR_HDB-1 FOUND\0"])
   end
 
-  def self.negative(include_pong: false)
-    responses = include_pong ? ["1: PONG\0"] : []
-    responses << "1: stream: OK\0"
-    new(responses)
+  def self.negative
+    new(["1: PONG\0", "1: stream: OK\0"])
   end
 
   def self.error
-    new(["1: INSTREAM size limit exceeded. ERROR\0"])
+    new(["1: PONG\0", "1: INSTREAM size limit exceeded. ERROR\0"])
   end
 
   def initialize(canned_responses)
     @canned_responses = canned_responses
+    @closed = false
     @received_before_close = []
     @received = []
     @next_to_read = 0
@@ -28,10 +25,14 @@ def initialize(canned_responses)
   attr_accessor :canned_responses
 
   def sendmsg_nonblock(text, _, _)
+    raise if @closed
+
     received << text
   end
 
   def recvmsg_nonblock(maxlen)
+    raise if @closed
+
     interval_end = @next_to_read + maxlen
 
     response_chunk = canned_responses[@current_response][@next_to_read..interval_end]
@@ -47,6 +48,7 @@ def recvmsg_nonblock(maxlen)
   end
 
   def close
+    @closed = true
     @next_to_read = 0
     @received_before_close = @received
     @received = []