diff --git a/zerver/actions/uploads.py b/zerver/actions/uploads.py new file mode 100644 index 0000000000..ac5b2ba354 --- /dev/null +++ b/zerver/actions/uploads.py @@ -0,0 +1,95 @@ +import logging +from typing import Any, Dict, List + +from zerver.lib.markdown import MessageRenderingResult +from zerver.lib.upload import claim_attachment, delete_message_image +from zerver.models import ( + Attachment, + Message, + Stream, + UserProfile, + get_old_unclaimed_attachments, + validate_attachment_request, +) +from zerver.tornado.django_api import send_event + + +def notify_attachment_update( + user_profile: UserProfile, op: str, attachment_dict: Dict[str, Any] +) -> None: + event = { + "type": "attachment", + "op": op, + "attachment": attachment_dict, + "upload_space_used": user_profile.realm.currently_used_upload_space_bytes(), + } + send_event(user_profile.realm, event, [user_profile.id]) + + +def do_claim_attachments(message: Message, potential_path_ids: List[str]) -> bool: + claimed = False + for path_id in potential_path_ids: + user_profile = message.sender + is_message_realm_public = False + is_message_web_public = False + if message.is_stream_message(): + stream = Stream.objects.get(id=message.recipient.type_id) + is_message_realm_public = stream.is_public() + is_message_web_public = stream.is_web_public + + if not validate_attachment_request(user_profile, path_id): + # Technically, there are 2 cases here: + # * The user put something in their message that has the form + # of an upload, but doesn't correspond to a file that doesn't + # exist. validate_attachment_request will return None. + # * The user is trying to send a link to a file they don't have permission to + # access themselves. validate_attachment_request will return False. + # + # Either case is unusual and suggests a UI bug that got + # the user in this situation, so we log in these cases. + logging.warning( + "User %s tried to share upload %s in message %s, but lacks permission", + user_profile.id, + path_id, + message.id, + ) + continue + + claimed = True + attachment = claim_attachment( + user_profile, path_id, message, is_message_realm_public, is_message_web_public + ) + notify_attachment_update(user_profile, "update", attachment.to_dict()) + return claimed + + +def do_delete_old_unclaimed_attachments(weeks_ago: int) -> None: + old_unclaimed_attachments = get_old_unclaimed_attachments(weeks_ago) + + for attachment in old_unclaimed_attachments: + delete_message_image(attachment.path_id) + attachment.delete() + + +def check_attachment_reference_change( + message: Message, rendering_result: MessageRenderingResult +) -> bool: + # For a unsaved message edit (message.* has been updated, but not + # saved to the database), adjusts Attachment data to correspond to + # the new content. + prev_attachments = {a.path_id for a in message.attachment_set.all()} + new_attachments = set(rendering_result.potential_attachment_path_ids) + + if new_attachments == prev_attachments: + return bool(prev_attachments) + + to_remove = list(prev_attachments - new_attachments) + if len(to_remove) > 0: + attachments_to_update = Attachment.objects.filter(path_id__in=to_remove).select_for_update() + message.attachment_set.remove(*attachments_to_update) + + to_add = list(new_attachments - prev_attachments) + if len(to_add) > 0: + do_claim_attachments(message, to_add) + + return message.attachment_set.exists() diff --git a/zerver/lib/actions.py b/zerver/lib/actions.py index 7ae95fe7f0..46b3402c5d 100644 --- a/zerver/lib/actions.py +++ b/zerver/lib/actions.py @@ -45,6 +45,7 @@ from zerver.actions.default_streams import ( get_default_streams_for_realm, ) from zerver.actions.invites import notify_invites_changed, revoke_invites_generated_by_user +from zerver.actions.uploads import check_attachment_reference_change, do_claim_attachments from zerver.actions.user_activity import update_user_activity_interval from zerver.actions.user_groups import ( do_send_user_group_members_update_event, @@ -172,7 +173,7 @@ from zerver.lib.types import ( SubscriptionInfo, SubscriptionStreamDict, ) -from zerver.lib.upload import claim_attachment, delete_avatar_image, delete_message_image +from zerver.lib.upload import delete_avatar_image from zerver.lib.user_counts import realm_user_count, realm_user_count_by_role from zerver.lib.user_groups import ( create_system_user_groups_for_realm, @@ -233,7 +234,6 @@ from zerver.models import ( get_client, get_fake_email_domain, get_huddle_user_ids, - get_old_unclaimed_attachments, get_realm, get_realm_domains, get_stream, @@ -243,7 +243,6 @@ from zerver.models import ( get_user_profile_by_id, is_cross_realm_bot_email, query_for_ids, - validate_attachment_request, ) from zerver.tornado.django_api import send_event @@ -6993,87 +6992,6 @@ def do_remove_realm_domain( transaction.on_commit(lambda: send_event(realm, event, active_user_ids(realm.id))) -def notify_attachment_update( - user_profile: UserProfile, op: str, attachment_dict: Dict[str, Any] -) -> None: - event = { - "type": "attachment", - "op": op, - "attachment": attachment_dict, - "upload_space_used": user_profile.realm.currently_used_upload_space_bytes(), - } - send_event(user_profile.realm, event, [user_profile.id]) - - -def do_claim_attachments(message: Message, potential_path_ids: List[str]) -> bool: - claimed = False - for path_id in potential_path_ids: - user_profile = message.sender - is_message_realm_public = False - is_message_web_public = False - if message.is_stream_message(): - stream = Stream.objects.get(id=message.recipient.type_id) - is_message_realm_public = stream.is_public() - is_message_web_public = stream.is_web_public - - if not validate_attachment_request(user_profile, path_id): - # Technically, there are 2 cases here: - # * The user put something in their message that has the form - # of an upload, but doesn't correspond to a file that doesn't - # exist. validate_attachment_request will return None. - # * The user is trying to send a link to a file they don't have permission to - # access themselves. validate_attachment_request will return False. - # - # Either case is unusual and suggests a UI bug that got - # the user in this situation, so we log in these cases. - logging.warning( - "User %s tried to share upload %s in message %s, but lacks permission", - user_profile.id, - path_id, - message.id, - ) - continue - - claimed = True - attachment = claim_attachment( - user_profile, path_id, message, is_message_realm_public, is_message_web_public - ) - notify_attachment_update(user_profile, "update", attachment.to_dict()) - return claimed - - -def do_delete_old_unclaimed_attachments(weeks_ago: int) -> None: - old_unclaimed_attachments = get_old_unclaimed_attachments(weeks_ago) - - for attachment in old_unclaimed_attachments: - delete_message_image(attachment.path_id) - attachment.delete() - - -def check_attachment_reference_change( - message: Message, rendering_result: MessageRenderingResult -) -> bool: - # For a unsaved message edit (message.* has been updated, but not - # saved to the database), adjusts Attachment data to correspond to - # the new content. - prev_attachments = {a.path_id for a in message.attachment_set.all()} - new_attachments = set(rendering_result.potential_attachment_path_ids) - - if new_attachments == prev_attachments: - return bool(prev_attachments) - - to_remove = list(prev_attachments - new_attachments) - if len(to_remove) > 0: - attachments_to_update = Attachment.objects.filter(path_id__in=to_remove).select_for_update() - message.attachment_set.remove(*attachments_to_update) - - to_add = list(new_attachments - prev_attachments) - if len(to_add) > 0: - do_claim_attachments(message, to_add) - - return message.attachment_set.exists() - - def notify_realm_custom_profile_fields(realm: Realm) -> None: fields = custom_profile_fields_for_realm(realm.id) event = dict(type="custom_profile_fields", fields=[f.as_dict() for f in fields]) diff --git a/zerver/lib/upload.py b/zerver/lib/upload.py index 7707715994..bcbee4496d 100644 --- a/zerver/lib/upload.py +++ b/zerver/lib/upload.py @@ -1142,7 +1142,7 @@ def create_attachment( realm=realm, size=file_size, ) - from zerver.lib.actions import notify_attachment_update + from zerver.actions.uploads import notify_attachment_update notify_attachment_update(user_profile, "add", attachment.to_dict()) return True diff --git a/zerver/management/commands/delete_old_unclaimed_attachments.py b/zerver/management/commands/delete_old_unclaimed_attachments.py index c241fbff81..b8a5757062 100644 --- a/zerver/management/commands/delete_old_unclaimed_attachments.py +++ b/zerver/management/commands/delete_old_unclaimed_attachments.py @@ -3,7 +3,7 @@ from typing import Any from django.core.management.base import BaseCommand, CommandError -from zerver.lib.actions import do_delete_old_unclaimed_attachments +from zerver.actions.uploads import do_delete_old_unclaimed_attachments from zerver.models import get_old_unclaimed_attachments diff --git a/zerver/tests/test_message_fetch.py b/zerver/tests/test_message_fetch.py index 277f5e94aa..f311b2d9ad 100644 --- a/zerver/tests/test_message_fetch.py +++ b/zerver/tests/test_message_fetch.py @@ -13,12 +13,8 @@ from sqlalchemy.types import Integer from analytics.lib.counts import COUNT_STATS from analytics.models import RealmCount -from zerver.lib.actions import ( - do_claim_attachments, - do_deactivate_user, - do_set_realm_property, - do_update_message, -) +from zerver.actions.uploads import do_claim_attachments +from zerver.lib.actions import do_deactivate_user, do_set_realm_property, do_update_message from zerver.lib.avatar import avatar_url from zerver.lib.exceptions import JsonableError from zerver.lib.mention import MentionBackend, MentionData @@ -3860,7 +3856,9 @@ class MessageHasKeywordsTest(ZulipTestCase): msg_id = self.send_stream_message(hamlet, "Denmark", body, "test") msg = Message.objects.get(id=msg_id) - with mock.patch("zerver.lib.actions.do_claim_attachments", wraps=do_claim_attachments) as m: + with mock.patch( + "zerver.actions.uploads.do_claim_attachments", wraps=do_claim_attachments + ) as m: self.update_message( msg, f"[link](http://{hamlet.realm.host}/user_uploads/{dummy_path_ids[0]})" ) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index c6f6475d75..441e86d360 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -21,10 +21,10 @@ from PIL import Image import zerver.lib.upload from zerver.actions.realm_icon import do_change_icon_source from zerver.actions.realm_logo import do_change_logo_source +from zerver.actions.uploads import do_delete_old_unclaimed_attachments from zerver.lib.actions import ( do_change_realm_plan_type, do_create_realm, - do_delete_old_unclaimed_attachments, do_set_realm_property, internal_send_private_message, ) diff --git a/zerver/views/attachments.py b/zerver/views/attachments.py index 8b10ce7584..e23bacd5c8 100644 --- a/zerver/views/attachments.py +++ b/zerver/views/attachments.py @@ -1,6 +1,6 @@ from django.http import HttpRequest, HttpResponse -from zerver.lib.actions import notify_attachment_update +from zerver.actions.uploads import notify_attachment_update from zerver.lib.attachments import access_attachment_by_id, remove_attachment, user_attachments from zerver.lib.response import json_success from zerver.models import UserProfile