DEV: Use new FileStore#download! method

discourse · Sep 19, 2024 · 631a010 · 631a010
1 parent e894058
commit 631a010
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/lib/discourse_antivirus/clam_av.rb b/lib/discourse_antivirus/clam_av.rb
@@ -61,7 +61,7 @@ def scan_upload(upload)
       return error_response(DOWNLOAD_FAILED) if file.nil?
 
       scan_file(file)
-    rescue OpenURI::HTTPError
+    rescue OpenURI::HTTPError, FileStore::DownloadError
       error_response(DOWNLOAD_FAILED)
     rescue StandardError => e
       Rails.logger.error("Could not scan upload #{upload.id}. Error: #{e.message}")
@@ -100,7 +100,7 @@ def get_uploaded_file(upload)
         # Upload#filesize could be approximate.
         # add two extra Mbs to make sure that we'll be able to download the upload.
         max_filesize = upload.filesize + 2.megabytes
-        store.download(upload, max_file_size_kb: max_filesize)
+        store.download!(upload, max_file_size_kb: max_filesize)
       else
         File.open(store.path_for(upload))
       end

diff --git a/spec/lib/discourse_antivirus/background_scan_spec.rb b/spec/lib/discourse_antivirus/background_scan_spec.rb
@@ -51,7 +51,7 @@
       store.stubs(:external?).returns(true)
       filesize = upload.filesize + 2.megabytes
       store
-        .expects(:download)
+        .expects(:download!)
         .with(upload, max_file_size_kb: filesize)
         .raises(OpenURI::HTTPError.new("forbidden", nil))
 
@@ -73,7 +73,7 @@
       store = Discourse.store
       store.stubs(:external?).returns(true)
       filesize = upload.filesize + 2.megabytes
-      store.expects(:download).with(upload, max_file_size_kb: filesize).returns(nil)
+      store.expects(:download!).with(upload, max_file_size_kb: filesize).returns(nil)
 
       antivirus = DiscourseAntivirus::ClamAv.new(store, build_fake_pool(socket))
       scanner = described_class.new(antivirus)

diff --git a/spec/lib/discourse_antivirus/clamav_spec.rb b/spec/lib/discourse_antivirus/clamav_spec.rb
@@ -48,6 +48,23 @@
       expect(scan_result[:error]).to eq(true)
       assert_file_was_sent_through(fake_socket, file)
     end
+
+    it "handles download errors" do
+      store = Discourse.store
+
+      store.stubs(:external?).returns(true)
+      store.stubs(:download!).raises(FileStore::DownloadError)
+
+      fake_socket = FakeTCPSocket.negative
+      pool = build_fake_pool(fake_socket)
+      antivirus = build_antivirus(pool, store)
+
+      scan_result = antivirus.scan_upload(upload)
+
+      expect(scan_result[:message]).to eq("Download failed")
+      expect(scan_result[:found]).to eq(false)
+      expect(scan_result[:error]).to eq(true)
+    end
   end
 
   describe "#version" do
@@ -124,7 +141,7 @@ def build_fake_pool(socket)
     FakePool.new([FakeTCPSocket.online, socket])
   end
 
-  def build_antivirus(pool)
-    described_class.new(Discourse.store, pool)
+  def build_antivirus(pool, store = Discourse.store)
+    described_class.new(store, pool)
   end
 end