uploads: Add a method to copy attachment contents out.

This commit is contained in:
Alex Vandiver 2023-03-14 16:16:41 +00:00 committed by Tim Abbott
parent 885334a3ad
commit e408f069fe
6 changed files with 46 additions and 6 deletions

View File

@ -3,7 +3,7 @@ import logging
import urllib
from datetime import datetime
from mimetypes import guess_type
from typing import IO, Any, Callable, Iterator, List, Optional, Tuple
from typing import IO, Any, BinaryIO, Callable, Iterator, List, Optional, Tuple
from urllib.parse import urljoin
from django.conf import settings
@ -107,6 +107,10 @@ def upload_message_attachment_from_request(
)
def save_attachment_contents(path_id: str, filehandle: BinaryIO) -> None:
return upload_backend.save_attachment_contents(path_id, filehandle)
def delete_message_attachment(path_id: str) -> bool:
return upload_backend.delete_message_attachment(path_id)

View File

@ -3,7 +3,7 @@ import os
import re
import unicodedata
from datetime import datetime
from typing import IO, Any, Callable, Iterator, List, Optional, Tuple
from typing import IO, Any, BinaryIO, Callable, Iterator, List, Optional, Tuple
from django.utils.translation import gettext as _
from PIL import GifImagePlugin, Image, ImageOps, PngImagePlugin
@ -198,6 +198,9 @@ class ZulipUploadBackend:
) -> str:
raise NotImplementedError
def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:
raise NotImplementedError
def delete_message_attachment(self, path_id: str) -> bool:
raise NotImplementedError

View File

@ -4,7 +4,7 @@ import random
import secrets
import shutil
from datetime import datetime
from typing import IO, Any, Callable, Iterator, Literal, Optional, Tuple
from typing import IO, Any, BinaryIO, Callable, Iterator, Literal, Optional, Tuple
from django.conf import settings
@ -98,6 +98,9 @@ class LocalUploadBackend(ZulipUploadBackend):
create_attachment(uploaded_file_name, path, user_profile, target_realm, uploaded_file_size)
return "/user_uploads/" + path
def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:
filehandle.write(read_local_file("files", path_id))
def delete_message_attachment(self, path_id: str) -> bool:
return delete_local_file("files", path_id)

View File

@ -4,7 +4,7 @@ import secrets
import urllib
from datetime import datetime
from mimetypes import guess_type
from typing import IO, Any, Callable, Iterator, List, Optional, Tuple
from typing import IO, Any, BinaryIO, Callable, Iterator, List, Optional, Tuple
import boto3
import botocore
@ -231,6 +231,10 @@ class S3UploadBackend(ZulipUploadBackend):
)
return url
def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:
for chunk in self.uploads_bucket.Object(path_id).get()["Body"]:
filehandle.write(chunk)
def delete_message_attachment(self, path_id: str) -> bool:
return self.delete_file_from_s3(path_id, self.uploads_bucket)

View File

@ -1,7 +1,7 @@
import os
import re
import urllib
from io import StringIO
from io import BytesIO, StringIO
from urllib.parse import urlparse
from django.conf import settings
@ -20,6 +20,7 @@ from zerver.lib.upload import (
delete_export_tarball,
delete_message_attachment,
delete_message_attachments,
save_attachment_contents,
upload_emoji_image,
upload_export_tarball,
upload_message_attachment,
@ -56,6 +57,17 @@ 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_save_attachment_contents(self) -> None:
user_profile = self.example_user("hamlet")
uri = upload_message_attachment(
"dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile
)
path_id = re.sub("/user_uploads/", "", uri)
output = BytesIO()
save_attachment_contents(path_id, output)
self.assertEqual(output.getvalue(), b"zulip!")
def test_upload_message_attachment_local_cross_realm_path(self) -> None:
"""
Verifies that the path of a file uploaded by a cross-realm bot to another

View File

@ -2,7 +2,7 @@ import io
import os
import re
import urllib
from io import StringIO
from io import BytesIO, StringIO
from unittest.mock import patch
import botocore.exceptions
@ -25,6 +25,7 @@ from zerver.lib.upload import (
delete_export_tarball,
delete_message_attachment,
delete_message_attachments,
save_attachment_contents,
upload_export_tarball,
upload_message_attachment,
)
@ -67,6 +68,19 @@ class S3Test(ZulipTestCase):
body = f"First message ...[zulip.txt](http://{user_profile.realm.host}{uri})"
self.send_stream_message(self.example_user("hamlet"), "Denmark", body, "test")
@use_s3_backend
def test_save_attachment_contents(self) -> None:
create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)
user_profile = self.example_user("hamlet")
uri = upload_message_attachment(
"dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile
)
path_id = re.sub("/user_uploads/", "", uri)
output = BytesIO()
save_attachment_contents(path_id, output)
self.assertEqual(output.getvalue(), b"zulip!")
@use_s3_backend
def test_upload_message_attachment_s3_cross_realm_path(self) -> None:
"""