Reject oversized projects

.env.example

@@ -67,7 +67,6 @@ SALESFORCE_CONNECT_USER=postgres
 
 SCRATCH_ASSET_CONFIG_BASE_URL=https://example.com/config/
 SCRATCH_ASSET_IMPORT_BASE_URL=https://example.com/assets/
-SCRATCH_ASSET_MAX_SIZE_MB=10
 
 HOSTNAME=localhost
 LEARNER_EXPERIENCE_TESTS_DISPATCH_TOKEN=changeme

app/controllers/api/scratch/assets_controller.rb

@@ -6,6 +6,7 @@ module Api
   module Scratch
     class AssetsController < ApiController
       include ActiveStorage::SetCurrent
+      include RemixSelection
 
       prepend_before_action :load_experience_cs_service_user, only: %i[create_global]
       before_action :authorize_user, except: %i[show]
@@ -50,7 +51,8 @@ def create_global
       def create_asset(reject_conflicting_content: false, **attributes)
         scratch_asset = ScratchAsset.find_or_initialize_by(attributes)
 
-        return reject_oversized_file if !scratch_asset.file.attached? && file_too_big?
+        reject_upload_over_limits(attributes[:project]) unless scratch_asset.file.attached?
+        return if performed?
 
         scratch_asset = save_new_asset(scratch_asset, attributes) if scratch_asset.new_record?
 
@@ -102,17 +104,43 @@ def file_matches?(scratch_asset)
         scratch_asset.file.blob.checksum == request_body_checksum
       end
 
+      def reject_upload_over_limits(project)
+        if file_too_big?
+          reject_oversized_file
+        elsif project && project_limit_exceeded?(project)
+          reject_project_over_limit
+        end
+      end
+
       def file_too_big?
-        request.body.size > max_file_size_mb.megabytes
+        request.body.size > ScratchAsset::MAX_SIZE_MB.megabytes
       end
 
       def reject_oversized_file
-        render json: { error: "The asset was larger than the limit of #{max_file_size_mb} MB" },
+        render json: { error: :asset_too_large, detail: "The asset was larger than the limit of #{ScratchAsset::MAX_SIZE_MB} MB" },
                status: :content_too_large
       end
 
-      def max_file_size_mb
-        ENV.fetch('SCRATCH_ASSET_MAX_SIZE_MB', 10).to_i
+      def project_limit_exceeded?(project)
+        current_total = ScratchAsset.project_byte_size(
+          project_ids: original_and_remix_ids(project),
+          uploaded_user_id: current_user.id
+        )
+        current_total + request.body.size > ScratchAsset::MAX_PROJECT_SIZE_MB.megabytes
+      end
+
+      def original_and_remix_ids(project)
+        original = users_remix?(project) ? project.parent : project
+        [original.id, project.id, remix_for_user(original, current_user)&.id].compact.uniq
+      end
+
+      def users_remix?(project)
+        project.parent.present? && project.user_id == current_user.id
+      end
+
+      def reject_project_over_limit
+        render json: { error: :project_too_large, detail: "The project would exceed the limit of #{ScratchAsset::MAX_PROJECT_SIZE_MB} MB" },
+               status: :unprocessable_content
       end
 
       def reject_conflicting_file(scratch_asset)

app/models/scratch_asset.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 class ScratchAsset < ApplicationRecord
+  MAX_SIZE_MB = 20
+  MAX_PROJECT_SIZE_MB = 100
+
   belongs_to :project, optional: true
   has_one_attached :file
 
@@ -25,6 +28,12 @@ def self.find_visible_to_project(project:, user:, filename:)
     global_assets.find_by(filename:)
   end
 
+  def self.project_byte_size(project_ids:, uploaded_user_id:)
+    joins(file_attachment: :blob)
+      .where(project_id: project_ids, uploaded_user_id:)
+      .sum('active_storage_blobs.byte_size')
+  end
+
   def global?
     project_id.nil?
   end

spec/features/scratch/creating_and_showing_a_scratch_asset_spec.rb

@@ -495,11 +495,8 @@ def make_request(body = upload)
       post request_path, headers: request_headers, params: body
     end
 
-    around do |example|
-      ClimateControl.modify(SCRATCH_ASSET_MAX_SIZE_MB: limit_mb.to_s) { example.run }
-    end
-
     before do
+      stub_const('ScratchAsset::MAX_SIZE_MB', limit_mb)
       authenticated_in_hydra_as(teacher)
     end
 
@@ -542,6 +539,57 @@ def make_request(body = upload)
     end
   end
 
+  describe 'enforcing the project asset size limit' do
+    let(:limit_mb) { 1 }
+    let(:half_limit) { 'a' * (limit_mb.megabytes / 2) }
+    let(:request_path) { '/api/scratch/assets/test_image_1.png' }
+    let(:request_headers) do
+      { 'Content-Type' => 'application/octet-stream', 'X-Project-ID' => project.identifier }.merge(auth_headers)
+    end
+
+    def make_request(body)
+      post request_path, headers: request_headers, params: body
+    end
+
+    before do
+      stub_const('ScratchAsset::MAX_PROJECT_SIZE_MB', limit_mb)
+      authenticated_in_hydra_as(teacher)
+    end
+
+    it 'responds with a 422 status code when the upload would take the project over the limit' do
+      create_uploaded_scratch_asset(filename: 'existing.png', project:, body: half_limit)
+
+      make_request("#{half_limit}a")
+
+      expect(response).to have_http_status(:unprocessable_content)
+    end
+
+    it 'does not create an asset or file when the project would go over the limit' do
+      create_uploaded_scratch_asset(filename: 'existing.png', project:, body: half_limit)
+
+      expect { make_request("#{half_limit}a") }.not_to change(ScratchAsset, :count)
+      expect { make_request("#{half_limit}a") }.not_to change(ActiveStorage::Blob, :count)
+    end
+
+    it 'accepts an upload that brings the project exactly up to the limit' do
+      create_uploaded_scratch_asset(filename: 'existing.png', project:, body: half_limit)
+
+      make_request(half_limit)
+
+      expect(response).to have_http_status(:created)
+    end
+
+    it 'skips the check when the asset already has a file attached' do
+      create_uploaded_scratch_asset(filename: 'existing.png', project:, body: half_limit)
+      existing_asset = create_uploaded_scratch_asset(filename:, project:, body: half_limit)
+
+      make_request("#{half_limit}a")
+
+      expect(response).to have_http_status(:created)
+      expect(existing_asset.reload.file.download).to eq(half_limit)
+    end
+  end
+
   describe 'visibility of a migrated project asset' do
     let(:student) { create(:student, school:) }
     let(:class_teacher) { create(:teacher, school:) }

spec/models/scratch_asset_spec.rb

@@ -57,4 +57,29 @@ def create_scratch_asset(project:, uploaded_user_id:)
       create(:scratch_asset, filename:, project:, uploaded_user_id:)
     end
   end
+
+  describe '.total_byte_size' do
+    let(:user_id) { SecureRandom.uuid }
+    let(:project) { create(:project, user_id:, project_type: Project::Types::CODE_EDITOR_SCRATCH, locale: nil) }
+    let(:other_project) { create(:project, user_id:, project_type: Project::Types::CODE_EDITOR_SCRATCH, locale: nil) }
+
+    it 'sums the file sizes of assets uploaded by the user to the given projects' do
+      create_attached_asset(filename: 'a.png', project:, uploaded_user_id: user_id, body: 'aaa')
+      create_attached_asset(filename: 'b.png', project: other_project, uploaded_user_id: user_id, body: 'bbbbb')
+
+      expect(described_class.project_byte_size(project_ids: [project.id, other_project.id], uploaded_user_id: user_id)).to eq(8)
+    end
+
+    it 'ignores assets without an attached file' do
+      create(:scratch_asset, filename: 'a.png', project:, uploaded_user_id: user_id)
+
+      expect(described_class.project_byte_size(project_ids: [project.id], uploaded_user_id: user_id)).to eq(0)
+    end
+
+    def create_attached_asset(filename:, project:, uploaded_user_id:, body:)
+      described_class.create!(filename:, project:, uploaded_user_id:).tap do |asset|
+        asset.file.attach(io: StringIO.new(body), filename:)
+      end
+    end
+  end
 end