diff --git a/zerver/lib/upload.py b/zerver/lib/upload.py index 97c82ef9ee..6f57aa9723 100644 --- a/zerver/lib/upload.py +++ b/zerver/lib/upload.py @@ -33,7 +33,14 @@ from PIL.Image import DecompressionBombError from zerver.lib.avatar_hash import user_avatar_path from zerver.lib.exceptions import ErrorCode, JsonableError from zerver.lib.utils import assert_is_not_none -from zerver.models import Attachment, Message, Realm, RealmEmoji, UserProfile +from zerver.models import ( + Attachment, + Message, + Realm, + RealmEmoji, + UserProfile, + is_cross_realm_bot_email, +) DEFAULT_AVATAR_SIZE = 100 MEDIUM_AVATAR_SIZE = 500 @@ -496,7 +503,9 @@ class S3UploadBackend(ZulipUploadBackend): file_data, ) - create_attachment(uploaded_file_name, s3_file_name, user_profile, uploaded_file_size) + create_attachment( + uploaded_file_name, s3_file_name, user_profile, target_realm, uploaded_file_size + ) return url def delete_message_image(self, path_id: str) -> bool: @@ -835,10 +844,13 @@ class LocalUploadBackend(ZulipUploadBackend): user_profile: UserProfile, target_realm: Optional[Realm] = None, ) -> str: - path = self.generate_message_upload_path(str(user_profile.realm_id), uploaded_file_name) + if target_realm is None: + target_realm = user_profile.realm + + path = self.generate_message_upload_path(str(target_realm.id), uploaded_file_name) write_local_file("files", path, file_data) - create_attachment(uploaded_file_name, path, user_profile, uploaded_file_size) + create_attachment(uploaded_file_name, path, user_profile, target_realm, uploaded_file_size) return "/user_uploads/" + path def delete_message_image(self, path_id: str) -> bool: @@ -1094,13 +1106,16 @@ def claim_attachment( def create_attachment( - file_name: str, path_id: str, user_profile: UserProfile, file_size: int + file_name: str, path_id: str, user_profile: UserProfile, realm: Realm, file_size: int ) -> bool: + assert (user_profile.realm_id == realm.id) or is_cross_realm_bot_email( + user_profile.delivery_email + ) attachment = Attachment.objects.create( file_name=file_name, path_id=path_id, owner=user_profile, - realm=user_profile.realm, + realm=realm, size=file_size, ) from zerver.lib.actions import notify_attachment_update diff --git a/zerver/tests/test_email_mirror.py b/zerver/tests/test_email_mirror.py index 002c478a46..993cc889f7 100644 --- a/zerver/tests/test_email_mirror.py +++ b/zerver/tests/test_email_mirror.py @@ -40,6 +40,7 @@ from zerver.lib.send_email import FromAddress from zerver.lib.test_classes import ZulipTestCase from zerver.lib.test_helpers import mock_queue_publish, most_recent_message, most_recent_usermessage from zerver.models import ( + Attachment, MissedMessageEmailAddress, Recipient, Stream, @@ -535,6 +536,46 @@ class TestEmailMirrorMessagesWithAttachments(ZulipTestCase): message = most_recent_message(user_profile) self.assertEqual(message.content, "Test body\n[image.png](https://test_url)") + def test_message_with_valid_attachment_model_attributes_set_correctly(self) -> None: + """ + Verifies that the Attachment attributes are set correctly. + """ + user_profile = self.example_user("hamlet") + self.login_user(user_profile) + self.subscribe(user_profile, "Denmark") + stream = get_stream("Denmark", user_profile.realm) + stream_to_address = encode_email_address(stream) + + incoming_valid_message = EmailMessage() + incoming_valid_message.set_content("Test body") + with open( + os.path.join(settings.DEPLOY_ROOT, "static/images/default-avatar.png"), "rb" + ) as f: + image_bytes = f.read() + + incoming_valid_message.add_attachment( + image_bytes, + maintype="image", + subtype="png", + filename="image.png", + ) + + incoming_valid_message["Subject"] = "TestStreamEmailMessages subject" + incoming_valid_message["From"] = self.example_email("hamlet") + incoming_valid_message["To"] = stream_to_address + incoming_valid_message["Reply-to"] = self.example_email("othello") + + process_message(incoming_valid_message) + + message = most_recent_message(user_profile) + attachment = Attachment.objects.last() + self.assertEqual(list(attachment.messages.values_list("id", flat=True)), [message.id]) + self.assertEqual( + message.sender, get_system_bot(settings.EMAIL_GATEWAY_BOT, stream.realm_id) + ) + self.assertEqual(attachment.realm, stream.realm) + self.assertEqual(attachment.is_realm_public, True) + def test_message_with_attachment_utf8_filename(self) -> None: user_profile = self.example_user("hamlet") self.login_user(user_profile) diff --git a/zerver/tests/test_message_fetch.py b/zerver/tests/test_message_fetch.py index a5eff2ba6a..2181d8efd7 100644 --- a/zerver/tests/test_message_fetch.py +++ b/zerver/tests/test_message_fetch.py @@ -3703,7 +3703,7 @@ class MessageHasKeywordsTest(ZulipTestCase): ] for file_name, path_id, size in dummy_files: - create_attachment(file_name, path_id, user_profile, size) + create_attachment(file_name, path_id, user_profile, user_profile.realm, size) # return path ids return [x[1] for x in dummy_files] diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py index 22c7aa46e0..0989ed89ee 100644 --- a/zerver/tests/test_retention.py +++ b/zerver/tests/test_retention.py @@ -174,7 +174,7 @@ class ArchiveMessagesTestingBase(RetentionTestingBase): ] for file_name, path_id, size in dummy_files: - create_attachment(file_name, path_id, user_profile, size) + create_attachment(file_name, path_id, user_profile, user_profile.realm, size) self.subscribe(user_profile, "Denmark") body = ( @@ -573,7 +573,7 @@ class MoveMessageToArchiveBase(RetentionTestingBase): ] user_profile = self.example_user("hamlet") for file_name, path_id, size in dummy_files: - create_attachment(file_name, path_id, user_profile, size) + create_attachment(file_name, path_id, user_profile, user_profile.realm, size) def _assert_archive_empty(self) -> None: self.assertFalse(ArchivedUserMessage.objects.exists()) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index 883a2d6043..d62b301004 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -1691,6 +1691,23 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase): uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id) self.assert_length(b"zulip!", uploaded_file.size) + def test_file_upload_local_cross_realm_path(self) -> None: + """ + Verifies that the path of a file uploaded by a cross-realm bot to another + realm is correct. + """ + + internal_realm = get_realm(settings.SYSTEM_BOT_REALM) + zulip_realm = get_realm("zulip") + user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT, internal_realm.id) + self.assertEqual(user_profile.realm, internal_realm) + + uri = upload_message_file( + "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile, zulip_realm + ) + # Ensure the correct realm id of the target realm is used instead of the bot's realm. + self.assertTrue(uri.startswith(f"/user_uploads/{zulip_realm.id}/")) + def test_delete_message_image_local(self) -> None: self.login("hamlet") fp = StringIO("zulip!") @@ -1837,6 +1854,25 @@ class S3Test(ZulipTestCase): body = "First message ...[zulip.txt](http://localhost:9991" + uri + ")" self.send_stream_message(self.example_user("hamlet"), "Denmark", body, "test") + @use_s3_backend + def test_file_upload_s3_cross_realm_path(self) -> None: + """ + Verifies that the path of a file uploaded by a cross-realm bot to another + realm is correct. + """ + create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET) + + internal_realm = get_realm(settings.SYSTEM_BOT_REALM) + zulip_realm = get_realm("zulip") + user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT, internal_realm.id) + self.assertEqual(user_profile.realm, internal_realm) + + uri = upload_message_file( + "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile, zulip_realm + ) + # Ensure the correct realm id of the target realm is used instead of the bot's realm. + self.assertTrue(uri.startswith(f"/user_uploads/{zulip_realm.id}/")) + @use_s3_backend def test_file_upload_s3_with_undefined_content_type(self) -> None: bucket = create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)[0]