paulmataruso/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2025-11-18 12:54:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

attachment: Add 'size' field tracking size of uploaded files.

Browse Source 
This tracking will make it possible in the future to limit the total
size of uploads on a per-user or per-organization basis.

Fixes #3774.
This commit is contained in:
PhilSk
2017-02-26 13:03:45 +03:00
committed by Tim Abbott
parent 8b003aa48d
commit 53f3d84af2
7 changed files with 68 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
zerver/lib/email_mirror.py
View File

@@ -259,7 +259,7 @@ def extract_and_upload_attachments(message, realm):
if filename:
attachment = part.get_payload(decode=True)
if isinstance(attachment, binary_type):
s3_url = upload_message_image(filename, content_type,
s3_url = upload_message_image(filename, len(attachment), content_type,
attachment,
user_profile,
target_realm=realm)

46
zerver/lib/upload.py
View File

@@ -98,8 +98,9 @@ def resize_avatar(image_data, size=DEFAULT_AVATAR_SIZE):
### Common
class ZulipUploadBackend(object):
def upload_message_image(self, uploaded_file_name, content_type, file_data, user_profile, target_realm=None):
# type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
def upload_message_image(self, uploaded_file_name, uploaded_file_size,
content_type, file_data, user_profile, target_realm=None):
# type: (Text, int, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
raise NotImplementedError()
def upload_avatar_image(self, user_file, user_profile, email):
@@ -165,7 +166,7 @@ def upload_image_to_s3(
key.set_contents_from_string(force_str(contents), headers=headers)
def get_file_info(request, user_file):
# type: (HttpRequest, File) -> Tuple[Text, Optional[Text]]
# type: (HttpRequest, File) -> Tuple[Text, int, Optional[Text]]
uploaded_file_name = user_file.name
assert isinstance(uploaded_file_name, str)
@@ -179,7 +180,9 @@ def get_file_info(request, user_file):
uploaded_file_name = uploaded_file_name + guess_extension(content_type)
uploaded_file_name = urllib.parse.unquote(uploaded_file_name)
return uploaded_file_name, content_type
uploaded_file_size = user_file.size
return uploaded_file_name, uploaded_file_size, content_type
def get_signed_upload_url(path):
@@ -197,8 +200,9 @@ def get_realm_for_filename(path):
return get_user_profile_by_id(key.metadata["user_profile_id"]).realm_id
class S3UploadBackend(ZulipUploadBackend):
def upload_message_image(self, uploaded_file_name, content_type, file_data, user_profile, target_realm=None):
# type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
def upload_message_image(self, uploaded_file_name, uploaded_file_size,
content_type, file_data, user_profile, target_realm=None):
# type: (Text, int, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
bucket_name = settings.S3_AUTH_UPLOADS_BUCKET
if target_realm is None:
target_realm = user_profile.realm
@@ -217,7 +221,7 @@ class S3UploadBackend(ZulipUploadBackend):
file_data
)
create_attachment(uploaded_file_name, s3_file_name, user_profile)
create_attachment(uploaded_file_name, s3_file_name, user_profile, uploaded_file_size)
return url
def delete_message_image(self, path_id):
@@ -355,8 +359,9 @@ def get_local_file_path(path_id):
return None
class LocalUploadBackend(ZulipUploadBackend):
def upload_message_image(self, uploaded_file_name, content_type, file_data, user_profile, target_realm=None):
# type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
def upload_message_image(self, uploaded_file_name, uploaded_file_size,
content_type, file_data, user_profile, target_realm=None):
# type: (Text, int, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
# Split into 256 subdirectories to prevent directories from getting too big
path = "/".join([
str(user_profile.realm_id),
@@ -366,7 +371,7 @@ class LocalUploadBackend(ZulipUploadBackend):
])
write_local_file('files', path, file_data)
create_attachment(uploaded_file_name, path, user_profile)
create_attachment(uploaded_file_name, path, user_profile, uploaded_file_size)
return '/user_uploads/' + path
def delete_message_image(self, path_id):
@@ -449,10 +454,11 @@ def upload_icon_image(user_file, user_profile):
# type: (File, UserProfile) -> None
upload_backend.upload_realm_icon_image(user_file, user_profile)
def upload_message_image(uploaded_file_name, content_type, file_data, user_profile, target_realm=None):
# type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
return upload_backend.upload_message_image(uploaded_file_name, content_type, file_data,
user_profile, target_realm=target_realm)
def upload_message_image(uploaded_file_name, uploaded_file_size,
content_type, file_data, user_profile, target_realm=None):
# type: (Text, int, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
return upload_backend.upload_message_image(uploaded_file_name, uploaded_file_size,
content_type, file_data, user_profile, target_realm=target_realm)
def claim_attachment(user_profile, path_id, message, is_message_realm_public):
# type: (UserProfile, Text, Message, bool) -> bool
@@ -470,12 +476,14 @@ def claim_attachment(user_profile, path_id, message, is_message_realm_public):
raise JsonableError(_("The upload was not successful. Please reupload the file again in a new message."))
return False
def create_attachment(file_name, path_id, user_profile):
# type: (Text, Text, UserProfile) -> bool
Attachment.objects.create(file_name=file_name, path_id=path_id, owner=user_profile, realm=user_profile.realm)
def create_attachment(file_name, path_id, user_profile, file_size):
# type: (Text, Text, UserProfile, int) -> bool
Attachment.objects.create(file_name=file_name, path_id=path_id, owner=user_profile,
realm=user_profile.realm, size=file_size)
return True
def upload_message_image_from_request(request, user_file, user_profile):
# type: (HttpRequest, File, UserProfile) -> Text
uploaded_file_name, content_type = get_file_info(request, user_file)
return upload_message_image(uploaded_file_name, content_type, user_file.read(), user_profile)
uploaded_file_name, uploaded_file_size, content_type = get_file_info(request, user_file)
return upload_message_image(uploaded_file_name, uploaded_file_size,
content_type, user_file.read(), user_profile)

20
zerver/migrations/0055_attachment_size.py Normal file
View File

@@ -0,0 +1,20 @@
# -*- coding: utf-8 -*-
# Generated by Django 1.10.5 on 2017-03-01 06:28
from __future__ import unicode_literals
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('zerver', '0054_realm_icon'),
]
operations = [
migrations.AddField(
model_name='attachment',
name='size',
field=models.IntegerField(null=True),
),
]

1
zerver/models.py
View File

@@ -1160,6 +1160,7 @@ class Attachment(ModelReprMixin, models.Model):
is_realm_public = models.BooleanField(default=False) # type: bool
messages = models.ManyToManyField(Message) # type: Manager
create_time = models.DateTimeField(default=timezone.now, db_index=True) # type: datetime.datetime
size = models.IntegerField(null=True) # type: int
def __unicode__(self):
# type: () -> Text

2
zerver/tests/test_export.py
View File

@@ -224,7 +224,7 @@ class ExportTest(TestCase):
# type: () -> None
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_image(u'dummy.txt', u'text/plain', b'zulip!', user_profile)
url = upload_message_image(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
path_id = url.replace('/user_uploads/', '')
claim_attachment(
user_profile=user_profile,

13
zerver/tests/test_messages.py
View File

@@ -1759,14 +1759,15 @@ class AttachmentTest(ZulipTestCase):
# Create dummy DB entry
sender_email = "hamlet@zulip.com"
user_profile = get_user_profile_by_email(sender_email)
sample_size = 10
dummy_files = [
('zulip.txt', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt'),
('temp_file.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py'),
('abc.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py')
('zulip.txt', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt', sample_size),
('temp_file.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py', sample_size),
('abc.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py', sample_size)
]
for file_name, path_id in dummy_files:
create_attachment(file_name, path_id, user_profile)
for file_name, path_id, size in dummy_files:
create_attachment(file_name, path_id, user_profile, size)
# Send message referring the attachment
self.subscribe_to_stream(sender_email, "Denmark")
@@ -1777,7 +1778,7 @@ class AttachmentTest(ZulipTestCase):
self.send_message(sender_email, "Denmark", Recipient.STREAM, body, "test")
for file_name, path_id in dummy_files:
for file_name, path_id, size in dummy_files:
attachment = Attachment.objects.get(path_id=path_id)
self.assertTrue(attachment.is_claimed())

15
zerver/tests/test_upload.py
View File

@@ -638,7 +638,7 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
# type: () -> None
sender_email = "hamlet@zulip.com"
user_profile = get_user_profile_by_email(sender_email)
uri = upload_message_image(u'dummy.txt', u'text/plain', b'zulip!', user_profile)
uri = upload_message_image(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
base = '/user_uploads/'
self.assertEqual(base, uri[:len(base)])
@@ -646,6 +646,9 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'files', path_id)
self.assertTrue(os.path.isfile(file_path))
uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id)
self.assertEqual(len(b'zulip!'), uploaded_file.size)
def test_delete_message_image_local(self):
# type: () -> None
self.login("hamlet@zulip.com")
@@ -687,12 +690,16 @@ class S3Test(ZulipTestCase):
sender_email = "hamlet@zulip.com"
user_profile = get_user_profile_by_email(sender_email)
uri = upload_message_image(u'dummy.txt', u'text/plain', b'zulip!', user_profile)
uri = upload_message_image(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
base = '/user_uploads/'
self.assertEqual(base, uri[:len(base)])
path_id = re.sub('/user_uploads/', '', uri)
self.assertEqual(b"zulip!", bucket.get_key(path_id).get_contents_as_string())
content = bucket.get_key(path_id).get_contents_as_string()
self.assertEqual(b"zulip!", content)
uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id)
self.assertEqual(len(b"zulip!"), uploaded_file.size)
self.subscribe_to_stream("hamlet@zulip.com", "Denmark")
body = "First message ...[zulip.txt](http://localhost:9991" + uri + ")"
@@ -707,7 +714,7 @@ class S3Test(ZulipTestCase):
sender_email = "hamlet@zulip.com"
user_profile = get_user_profile_by_email(sender_email)
uri = upload_message_image(u'dummy.txt', u'text/plain', b'zulip!', user_profile)
uri = upload_message_image(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
path_id = re.sub('/user_uploads/', '', uri)
self.assertTrue(delete_message_image(path_id))