FIX: Use theme screenshot names in theme fields (#31852)

Currently we allow for 2 theme screenshots to be specified,
with a lightweight spec to allow both a light and dark
version of the screenshot. However, we were not storing this
screenshot name anywhere, so we would not be able to use it
for light/dark switching.

This commit fixes that issue, and also does some general refactoring
around theme screenshots, and adds more tests.

app/models/remote_theme.rb

@@ -31,9 +31,6 @@ class RemoteTheme < ActiveRecord::Base
   MAX_ASSET_FILE_SIZE = 8.megabytes
   MAX_THEME_FILE_COUNT = 1024
   MAX_THEME_SIZE = 256.megabytes
-  MAX_THEME_SCREENSHOT_FILE_SIZE = 1.megabyte
-  MAX_THEME_SCREENSHOT_DIMENSIONS = [3840, 2160] # 4K resolution
-  THEME_SCREENSHOT_ALLOWED_FILE_TYPES = %w[.jpg .jpeg .gif .png].freeze
 
   has_one :theme, autosave: false
   scope :joined_remotes,
@@ -253,7 +250,7 @@ class RemoteTheme < ActiveRecord::Base
 
     theme_info["assets"]&.each do |name, relative_path|
       if path = importer.real_path(relative_path)
-        upload = create_upload(path, relative_path)
+        upload = RemoteTheme.create_upload(theme, path, relative_path)
         if !upload.errors.empty?
           raise ImportError,
                 I18n.t(
@@ -275,61 +272,15 @@ class RemoteTheme < ActiveRecord::Base
     # TODO (martin): Until we are ready to roll this out more
     # widely, let's avoid doing this work for most sites.
     if SiteSetting.theme_download_screenshots
-      theme_info["screenshots"] = Array.wrap(theme_info["screenshots"]).take(2)
+      begin
-      theme_info["screenshots"].each_with_index do |relative_path, idx|
+        updated_fields.concat(
-        if path = importer.real_path(relative_path)
+          ThemeScreenshotsHandler.new(theme).parse_screenshots_as_theme_fields!(
-          if !THEME_SCREENSHOT_ALLOWED_FILE_TYPES.include?(File.extname(path))
+            theme_info["screenshots"],
-            raise ImportError,
+            importer,
-                  I18n.t(
+          ),
-                    "themes.import_error.screenshot_invalid_type",
+        )
-                    file_name: File.basename(path),
+      rescue ThemeScreenshotsHandler::ThemeScreenshotError => err
-                    accepted_formats: THEME_SCREENSHOT_ALLOWED_FILE_TYPES.join(","),
+        raise ImportError, err.message
-                  )
-          end
-
-          if File.size(path) > MAX_THEME_SCREENSHOT_FILE_SIZE
-            raise ImportError,
-                  I18n.t(
-                    "themes.import_error.screenshot_invalid_size",
-                    file_name: File.basename(path),
-                    max_size:
-                      ActiveSupport::NumberHelper.number_to_human_size(
-                        MAX_THEME_SCREENSHOT_FILE_SIZE,
-                      ),
-                  )
-          end
-
-          screenshot_width, screenshot_height = FastImage.size(path)
-          if (screenshot_width.nil? || screenshot_height.nil?) ||
-               screenshot_width > MAX_THEME_SCREENSHOT_DIMENSIONS[0] ||
-               screenshot_height > MAX_THEME_SCREENSHOT_DIMENSIONS[1]
-            raise ImportError,
-                  I18n.t(
-                    "themes.import_error.screenshot_invalid_dimensions",
-                    file_name: File.basename(path),
-                    width: screenshot_width.to_i,
-                    height: screenshot_height.to_i,
-                    max_width: MAX_THEME_SCREENSHOT_DIMENSIONS[0],
-                    max_height: MAX_THEME_SCREENSHOT_DIMENSIONS[1],
-                  )
-          end
-
-          upload = create_upload(path, relative_path)
-          if !upload.errors.empty?
-            raise ImportError,
-                  I18n.t(
-                    "themes.import_error.screenshot",
-                    errors: upload.errors.full_messages.join(","),
-                  )
-          end
-
-          updated_fields << theme.set_field(
-            target: :common,
-            name: "screenshot_#{idx + 1}",
-            type: :theme_screenshot_upload_var,
-            upload_id: upload.id,
-          )
-        end
       end
     end
 
@@ -523,7 +474,7 @@ class RemoteTheme < ActiveRecord::Base
     remote_url.present?
   end
 
-  def create_upload(path, relative_path)
+  def self.create_upload(theme, path, relative_path)
     new_path = "#{File.dirname(path)}/#{SecureRandom.hex}#{File.extname(path)}"
 
     # OptimizedImage has strict file name restrictions, so rename temporarily

lib/compression/strategy.rb

@@ -38,18 +38,6 @@ module Compression
       Pathname.new(filename).realpath.to_s
     end
 
-    # https://guides.rubyonrails.org/security.html#file-uploads
-    def sanitize_filename(filename)
-      filename.strip.tap do |name|
-        # NOTE: File.basename doesn't work right with Windows paths on Unix
-        # get only the filename, not the whole path
-        name.sub! %r{\A.*(\\|/)}, ""
-        # Finally, replace all non alphanumeric, underscore
-        # or periods with underscore
-        name.gsub! /[^\w\.\-]/, "_"
-      end
-    end
-
     def calculate_available_size(max_size)
       1024**2 * (max_size / 1.049) # Mb to Mib
     end

lib/compression/tar.rb

@@ -9,7 +9,7 @@ module Compression
     end
 
     def compress(path, target_name)
-      tar_filename = sanitize_filename("#{target_name}.tar")
+      tar_filename = FileHelper.sanitize_filename("#{target_name}.tar")
       Discourse::Utils.execute_command(
         "tar",
         "--create",

lib/file_helper.rb

@@ -40,6 +40,20 @@ class FileHelper
     filename.match?(supported_playable_media_regexp)
   end
 
+  # https://guides.rubyonrails.org/security.html#file-uploads
+  def self.sanitize_filename(filename)
+    filename.strip.tap do |name|
+      # NOTE: File.basename doesn't work right with Windows paths on Unix
+      # get only the filename, not the whole path
+      name.sub! %r{\A.*(\\|/)}, ""
+      # Replace all non alphanumeric, underscore
+      # or periods with underscore
+      name.gsub! /[^\w\.\-]/, "_"
+      # Finally, replace all double underscores with a single one
+      name.gsub! /_+/, "_"
+    end
+  end
+
   class FakeIO
     attr_accessor :status
   end

lib/theme_screenshots_handler.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+class ThemeScreenshotsHandler
+  MAX_THEME_SCREENSHOT_FILE_SIZE = 1.megabyte
+  MAX_THEME_SCREENSHOT_DIMENSIONS = [3840, 2160] # 4K resolution
+  MAX_THEME_SCREENSHOT_COUNT = 2
+  THEME_SCREENSHOT_ALLOWED_FILE_TYPES = %w[.jpg .jpeg .gif .png].freeze
+
+  class ThemeScreenshotError < StandardError
+  end
+
+  def initialize(theme)
+    @theme = theme
+  end
+
+  # Screenshots here come from RemoteTheme.extract_theme_info, which
+  # in turn parses the theme about.json file, which is where screenshots
+  # are defined.
+  def parse_screenshots_as_theme_fields!(screenshots, theme_importer)
+    updated_theme_fields = []
+    screenshots = Array.wrap(screenshots).take(MAX_THEME_SCREENSHOT_COUNT)
+    screenshots.each do |relative_path|
+      path = theme_importer.real_path(relative_path)
+      next if !path.present?
+
+      screenshot_filename = File.basename(path)
+      screenshot_extension = File.extname(path)
+
+      if !THEME_SCREENSHOT_ALLOWED_FILE_TYPES.include?(screenshot_extension)
+        raise ThemeScreenshotError,
+              I18n.t(
+                "themes.import_error.screenshot_invalid_type",
+                file_name: screenshot_filename,
+                accepted_formats: THEME_SCREENSHOT_ALLOWED_FILE_TYPES.join(","),
+              )
+      end
+
+      if File.size(path) > MAX_THEME_SCREENSHOT_FILE_SIZE
+        raise ThemeScreenshotError,
+              I18n.t(
+                "themes.import_error.screenshot_invalid_size",
+                file_name: screenshot_filename,
+                max_size:
+                  ActiveSupport::NumberHelper.number_to_human_size(MAX_THEME_SCREENSHOT_FILE_SIZE),
+              )
+      end
+
+      screenshot_width, screenshot_height = FastImage.size(path)
+      if (screenshot_width.nil? || screenshot_height.nil?) ||
+           screenshot_width > MAX_THEME_SCREENSHOT_DIMENSIONS[0] ||
+           screenshot_height > MAX_THEME_SCREENSHOT_DIMENSIONS[1]
+        raise ThemeScreenshotError,
+              I18n.t(
+                "themes.import_error.screenshot_invalid_dimensions",
+                file_name: screenshot_filename,
+                width: screenshot_width.to_i,
+                height: screenshot_height.to_i,
+                max_width: MAX_THEME_SCREENSHOT_DIMENSIONS[0],
+                max_height: MAX_THEME_SCREENSHOT_DIMENSIONS[1],
+              )
+      end
+
+      upload = RemoteTheme.create_upload(@theme, path, relative_path)
+      if !upload.errors.empty?
+        raise ThemeScreenshotError,
+              I18n.t(
+                "themes.import_error.screenshot",
+                errors: upload.errors.full_messages.join(","),
+              )
+      end
+
+      screenshot_filename_clean =
+        FileHelper.sanitize_filename(screenshot_filename.gsub(screenshot_extension, ""))
+      updated_theme_fields << @theme.set_field(
+        target: :common,
+        name: "screenshot_#{screenshot_filename_clean}",
+        type: :theme_screenshot_upload_var,
+        upload_id: upload.id,
+      )
+    end
+    updated_theme_fields
+  end
+end

spec/lib/theme_screenshots_handler_spec.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+RSpec.describe ThemeScreenshotsHandler do
+  subject(:handler) { ThemeScreenshotsHandler.new(theme) }
+
+  fab!(:remote_theme) { RemoteTheme.create!(remote_url: "https://github.com/org/remote-theme1") }
+  fab!(:theme) { Fabricate(:theme, remote_theme: remote_theme) }
+
+  let(:importer) { ThemeStore::GitImporter.new("https://github.com/org/remote-theme1") }
+
+  after(:each) { FileUtils.rm_rf(importer.temp_folder) }
+
+  def write_temp_screenshots_for_importer(screenshots)
+    FileUtils.mkdir_p("#{importer.temp_folder}/screenshots")
+
+    # Random data is added to each file so the sha1 hash is different
+    # and distinct uploads are created
+    screenshots.each do |screenshot|
+      File.write(
+        "#{importer.temp_folder}/#{screenshot}",
+        File.read(file_from_fixtures("logo.jpg", "images")) + SecureRandom.hex,
+      )
+    end
+  end
+
+  it "sanitizes filenames for screenshots before saving them" do
+    screenshots = ["screenshots/light.jpeg", "screenshots/some Absolutely silly $%&* name --.jpeg"]
+    write_temp_screenshots_for_importer(screenshots)
+    handler.parse_screenshots_as_theme_fields!(screenshots, importer)
+    theme.save!
+    expect(theme.theme_fields.reload.map(&:name)).to match_array(
+      %w[screenshot_light screenshot_some_Absolutely_silly_name_--],
+    )
+  end
+
+  it "generates correct theme fields" do
+    screenshots = %w[screenshots/light.jpeg screenshots/dark.jpeg]
+    write_temp_screenshots_for_importer(screenshots)
+    handler.parse_screenshots_as_theme_fields!(screenshots, importer)
+    theme.save!
+    expect(theme.theme_fields.pluck(:type_id)).to eq(
+      [
+        ThemeField.types[:theme_screenshot_upload_var],
+        ThemeField.types[:theme_screenshot_upload_var],
+      ],
+    )
+    expect(theme.theme_fields.pluck(:name)).to match_array(%w[screenshot_light screenshot_dark])
+    expect(theme.theme_fields.pluck(:upload_id)).to match_array(Upload.last(2).pluck(:id))
+  end
+
+  it "raises an error if the screenshot is not an allowed file type" do
+    screenshots = ["screenshots/light.tga"]
+    write_temp_screenshots_for_importer(screenshots)
+    expect { handler.parse_screenshots_as_theme_fields!(screenshots, importer) }.to raise_error(
+      ThemeScreenshotsHandler::ThemeScreenshotError,
+      I18n.t(
+        "themes.import_error.screenshot_invalid_type",
+        file_name: "light.tga",
+        accepted_formats: ThemeScreenshotsHandler::THEME_SCREENSHOT_ALLOWED_FILE_TYPES.join(","),
+      ),
+    )
+  end
+
+  it "raises an error if the screenshot is too big" do
+    screenshots = ["screenshots/light.jpeg"]
+    write_temp_screenshots_for_importer(screenshots)
+    stub_const(ThemeScreenshotsHandler, "MAX_THEME_SCREENSHOT_FILE_SIZE", 1.byte) do
+      expect { handler.parse_screenshots_as_theme_fields!(screenshots, importer) }.to raise_error(
+        ThemeScreenshotsHandler::ThemeScreenshotError,
+        I18n.t(
+          "themes.import_error.screenshot_invalid_size",
+          file_name: "light.jpeg",
+          max_size: "1 Bytes",
+        ),
+      )
+    end
+  end
+
+  it "raises an error if the screenshot has invalid dimensions" do
+    screenshots = ["screenshots/light.jpeg"]
+    write_temp_screenshots_for_importer(screenshots)
+    stub_const(ThemeScreenshotsHandler, "MAX_THEME_SCREENSHOT_DIMENSIONS", [1, 1]) do
+      expect { handler.parse_screenshots_as_theme_fields!(screenshots, importer) }.to raise_error(
+        ThemeScreenshotsHandler::ThemeScreenshotError,
+        I18n.t(
+          "themes.import_error.screenshot_invalid_dimensions",
+          file_name: "light.jpeg",
+          width: 512,
+          height: 512,
+          max_width: 1,
+          max_height: 1,
+        ),
+      )
+    end
+  end
+end

spec/models/remote_theme_spec.rb

@@ -6,7 +6,8 @@ RSpec.describe RemoteTheme do
       love_color: "FAFAFA",
       tertiary_low_color: "FFFFFF",
       color_scheme_name: "Amazing",
-      about_url: "https://www.site.com/about"
+      about_url: "https://www.site.com/about",
+      screenshots: %w[screenshots/light.jpeg screenshots/dark.jpeg]
     )
       <<~JSON
         {
@@ -35,7 +36,7 @@ RSpec.describe RemoteTheme do
               "value": "list_setting"
             }
           },
-          "screenshots": ["screenshots/1.jpeg", "screenshots/2.jpeg"]
+          "screenshots": #{screenshots.to_json}
         }
       JSON
     end
@@ -73,8 +74,8 @@ RSpec.describe RemoteTheme do
         "settings.yaml" => settings,
         "locales/en.yml" => "sometranslations",
         "migrations/settings/0001-some-migration.js" => migration_js,
-        "screenshots/1.jpeg" => file_from_fixtures("logo.jpg", "images"),
+        "screenshots/light.jpeg" => file_from_fixtures("logo.jpg", "images"),
-        "screenshots/2.jpeg" => file_from_fixtures("logo.jpg", "images"),
+        "screenshots/dark.jpeg" => file_from_fixtures("logo.jpg", "images"),
       )
     end
 
@@ -392,12 +393,12 @@ RSpec.describe RemoteTheme do
       before { SiteSetting.theme_download_screenshots = true }
 
       it "fails if any of the provided screenshots is not an accepted file type" do
-        stub_const(RemoteTheme, "THEME_SCREENSHOT_ALLOWED_FILE_TYPES", [".bmp"]) do
+        stub_const(ThemeScreenshotsHandler, "THEME_SCREENSHOT_ALLOWED_FILE_TYPES", [".bmp"]) do
           expect { RemoteTheme.import_theme(initial_repo_url) }.to raise_error(
             RemoteTheme::ImportError,
             I18n.t(
               "themes.import_error.screenshot_invalid_type",
-              file_name: "1.jpeg",
+              file_name: "light.jpeg",
               accepted_formats: ".bmp",
             ),
           )
@@ -405,12 +406,12 @@ RSpec.describe RemoteTheme do
       end
 
       it "fails if any of the provided screenshots is too big" do
-        stub_const(RemoteTheme, "MAX_THEME_SCREENSHOT_FILE_SIZE", 1.byte) do
+        stub_const(ThemeScreenshotsHandler, "MAX_THEME_SCREENSHOT_FILE_SIZE", 1.byte) do
           expect { RemoteTheme.import_theme(initial_repo_url) }.to raise_error(
             RemoteTheme::ImportError,
             I18n.t(
               "themes.import_error.screenshot_invalid_size",
-              file_name: "1.jpeg",
+              file_name: "light.jpeg",
               max_size: "1 Bytes",
             ),
           )
@@ -420,14 +421,14 @@ RSpec.describe RemoteTheme do
       it "fails if any of the provided screenshots has dimensions that are too big" do
         FastImage
           .expects(:size)
-          .with { |arg| arg.match(%r{/screenshots/1\.jpeg}) }
+          .with { |arg| arg.match(%r{/screenshots/light\.jpeg}) }
           .returns([512, 512])
-        stub_const(RemoteTheme, "MAX_THEME_SCREENSHOT_DIMENSIONS", [1, 1]) do
+        stub_const(ThemeScreenshotsHandler, "MAX_THEME_SCREENSHOT_DIMENSIONS", [1, 1]) do
           expect { RemoteTheme.import_theme(initial_repo_url) }.to raise_error(
             RemoteTheme::ImportError,
             I18n.t(
               "themes.import_error.screenshot_invalid_dimensions",
-              file_name: "1.jpeg",
+              file_name: "light.jpeg",
               width: 512,
               height: 512,
               max_width: 1,
@@ -440,17 +441,17 @@ RSpec.describe RemoteTheme do
       it "creates uploads and associated theme fields for all theme screenshots" do
         FastImage
           .stubs(:size)
-          .with { |arg| arg.match(%r{/screenshots/1\.jpeg}) }
+          .with { |arg| arg.match(%r{/screenshots/light\.jpeg}) }
           .returns([800, 600])
         FastImage
           .stubs(:size)
-          .with { |arg| arg.match(%r{/screenshots/2\.jpeg}) }
+          .with { |arg| arg.match(%r{/screenshots/dark\.jpeg}) }
           .returns([1024, 768])
 
         theme = RemoteTheme.import_theme(initial_repo_url)
 
-        screenshot_1 = theme.theme_fields.find_by(name: "screenshot_1")
+        screenshot_1 = theme.theme_fields.find_by(name: "screenshot_light")
-        screenshot_2 = theme.theme_fields.find_by(name: "screenshot_2")
+        screenshot_2 = theme.theme_fields.find_by(name: "screenshot_dark")
 
         expect(screenshot_1).to be_present
         expect(screenshot_1.type_id).to eq(ThemeField.types[:theme_screenshot_upload_var])
@@ -462,8 +463,8 @@ RSpec.describe RemoteTheme do
         expect(UploadReference.exists?(target: screenshot_1)).to eq(true)
         expect(UploadReference.exists?(target: screenshot_2)).to eq(true)
 
-        expect(screenshot_1.upload.original_filename).to eq("1.jpeg")
+        expect(screenshot_1.upload.original_filename).to eq("light.jpeg")
-        expect(screenshot_2.upload.original_filename).to eq("2.jpeg")
+        expect(screenshot_2.upload.original_filename).to eq("dark.jpeg")
       end
     end
   end