From 631a010a601b646c228db0a9bbf03b4a7cacd672 Mon Sep 17 00:00:00 2001
From: Ted Johansson <ted@discourse.org>
Date: Thu, 19 Sep 2024 10:06:00 +0800
Subject: [PATCH] DEV: Use new FileStore#download! method

---
 lib/discourse_antivirus/clam_av.rb            |  4 ++--
 .../background_scan_spec.rb                   |  4 ++--
 spec/lib/discourse_antivirus/clamav_spec.rb   | 21 +++++++++++++++++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/lib/discourse_antivirus/clam_av.rb b/lib/discourse_antivirus/clam_av.rb
index da7f392..0e33247 100644
--- 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
index 4cedec9..c5f4f82 100644
--- 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
index e426513..94b6d37 100644
--- 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