diff --git a/app/files/rest.py b/app/files/rest.py index 792974cc55..d0ca144070 100644 --- a/app/files/rest.py +++ b/app/files/rest.py @@ -51,6 +51,8 @@ "SKIPPED": FILE_STATUS_VIRUS_SCAN_FAILED, } +MAX_TOTAL_FILE_SIZE = 6 * 1024 * 1024 # 6MB + @files_blueprint.route("", methods=["POST"]) def create_file(template_id): @@ -72,13 +74,30 @@ def create_file(template_id): check_service_has_permission(UPLOAD_DOCUMENT, service.permissions) validate_template_exists(template_id, service) + existing_files = dao_get_files_by_template_id(template_id) + existing_size = sum(f.file_size or 0 for f in existing_files) + filename = data["name"] mime_type = data["mime_type"] try: file_data = base64.b64decode(data["file_data"]) - uploaded_file = document_download_client.upload_template_attachment(service.id, file_data, filename, mime_type) except (binascii.Error, ValueError): raise InvalidRequest("file_data is not valid base64", status_code=400) + + actual_file_size = len(file_data) + total_size = existing_size + actual_file_size + if total_size > MAX_TOTAL_FILE_SIZE: + return jsonify( + { + "error": "over_file_limit", + "current_usage": existing_size, + "requested": actual_file_size, + "limit": MAX_TOTAL_FILE_SIZE, + } + ), 400 + + try: + uploaded_file = document_download_client.upload_template_attachment(service.id, file_data, filename, mime_type) except DocumentDownloadError as e: raise InvalidRequest(e.message, status_code=e.status_code) diff --git a/tests/app/files/test_rest.py b/tests/app/files/test_rest.py index 43e2246aa9..a865451cc2 100644 --- a/tests/app/files/test_rest.py +++ b/tests/app/files/test_rest.py @@ -194,6 +194,164 @@ def test_create_file_returns_403_when_manage_templates_removed( _expected_status=403, ) + def test_create_file_returns_400_when_total_size_exceeds_6mb( + self, notify_db, notify_db_session, admin_request, sample_service_full_permissions + ): + from app.dao.files_dao import dao_create_file + from app.files.rest import MAX_TOTAL_FILE_SIZE + from app.models import FILE_STATUS_PENDING_VIRUS_SCAN, Files + + sample_template = create_sample_template(notify_db, notify_db_session, service=sample_service_full_permissions) + current_user_id = str(sample_template.service.users[0].id) + + # Create an existing file that uses most of the 6MB budget + existing_file = Files( + template_id=sample_template.id, + service_id=sample_service_full_permissions.id, + document_id=uuid.uuid4(), + type="attach", + name="existing.pdf", + mime_type="application/pdf", + file_size=MAX_TOTAL_FILE_SIZE - 1000, + status=FILE_STATUS_PENDING_VIRUS_SCAN, + created_by_id=current_user_id, + ) + dao_create_file(existing_file) + + # file_data decodes to 1001 bytes, which exceeds the remaining 1000 byte budget + oversized_content = b"x" * 1001 + response = admin_request.post( + "files.create_file", + template_id=str(sample_template.id), + _data={ + "template_id": str(sample_template.id), + "type": "attach", + "name": "too_large.pdf", + "mime_type": "application/pdf", + "file_size": 1001, + "file_data": base64.b64encode(oversized_content).decode("utf-8"), + "created_by": current_user_id, + }, + _expected_status=400, + ) + assert response["error"] == "over_file_limit" + assert response["current_usage"] == MAX_TOTAL_FILE_SIZE - 1000 + assert response["requested"] == 1001 + assert response["limit"] == MAX_TOTAL_FILE_SIZE + + def test_create_file_succeeds_when_total_size_exactly_at_6mb( + self, mocker, notify_db, notify_db_session, admin_request, sample_service_full_permissions + ): + from app.dao.files_dao import dao_create_file + from app.files.rest import MAX_TOTAL_FILE_SIZE + from app.models import FILE_STATUS_PENDING_VIRUS_SCAN, Files + + _mock_upload_template_attachment(mocker) + sample_template = create_sample_template(notify_db, notify_db_session, service=sample_service_full_permissions) + current_user_id = str(sample_template.service.users[0].id) + + # Create an existing file + existing_file = Files( + template_id=sample_template.id, + service_id=sample_service_full_permissions.id, + document_id=uuid.uuid4(), + type="attach", + name="existing.pdf", + mime_type="application/pdf", + file_size=MAX_TOTAL_FILE_SIZE - 1000, + status=FILE_STATUS_PENDING_VIRUS_SCAN, + created_by_id=current_user_id, + ) + dao_create_file(existing_file) + + # Adding exactly 1000 bytes of actual content should succeed (total == 6MB exactly) + exact_content = b"x" * 1000 + admin_request.post( + "files.create_file", + template_id=str(sample_template.id), + _data={ + "template_id": str(sample_template.id), + "type": "attach", + "name": "fits.pdf", + "mime_type": "application/pdf", + "file_size": 1000, + "file_data": base64.b64encode(exact_content).decode("utf-8"), + "created_by": current_user_id, + }, + _expected_status=201, + ) + + def test_create_file_returns_400_when_single_file_exceeds_6mb( + self, notify_db, notify_db_session, admin_request, sample_service_full_permissions + ): + from app.files.rest import MAX_TOTAL_FILE_SIZE + + sample_template = create_sample_template(notify_db, notify_db_session, service=sample_service_full_permissions) + current_user_id = str(sample_template.service.users[0].id) + + # file_data decodes to MAX + 1 bytes, rejected even with no existing files + oversized_content = b"x" * (MAX_TOTAL_FILE_SIZE + 1) + response = admin_request.post( + "files.create_file", + template_id=str(sample_template.id), + _data={ + "template_id": str(sample_template.id), + "type": "attach", + "name": "huge.pdf", + "mime_type": "application/pdf", + "file_size": MAX_TOTAL_FILE_SIZE + 1, + "file_data": base64.b64encode(oversized_content).decode("utf-8"), + "created_by": current_user_id, + }, + _expected_status=400, + ) + assert response["error"] == "over_file_limit" + assert response["current_usage"] == 0 + assert response["requested"] == MAX_TOTAL_FILE_SIZE + 1 + assert response["limit"] == MAX_TOTAL_FILE_SIZE + + def test_create_file_ignores_archived_files_in_size_calculation( + self, mocker, notify_db, notify_db_session, admin_request, sample_service_full_permissions + ): + from app.dao.files_dao import dao_archive_file, dao_create_file + from app.files.rest import MAX_TOTAL_FILE_SIZE + from app.models import FILE_STATUS_PENDING_VIRUS_SCAN, Files + + _mock_upload_template_attachment(mocker) + sample_template = create_sample_template(notify_db, notify_db_session, service=sample_service_full_permissions) + current_user_id = str(sample_template.service.users[0].id) + + # Create a large file and archive it + archived_file = Files( + template_id=sample_template.id, + service_id=sample_service_full_permissions.id, + document_id=uuid.uuid4(), + type="attach", + name="archived.pdf", + mime_type="application/pdf", + file_size=MAX_TOTAL_FILE_SIZE, + status=FILE_STATUS_PENDING_VIRUS_SCAN, + created_by_id=current_user_id, + ) + dao_create_file(archived_file) + dao_archive_file(archived_file) + + # Should succeed because archived files don't count toward the limit + admin_request.post( + "files.create_file", + template_id=str(sample_template.id), + _data={ + "template_id": str(sample_template.id), + "type": "attach", + "name": "new.pdf", + "mime_type": "application/pdf", + "file_size": 1000, + "file_data": base64.b64encode(b"test content").decode("utf-8"), + "created_by": current_user_id, + }, + _expected_status=201, + ) + class TestGetFile: def test_get_file_status(self, admin_request, sample_file):