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

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.

zerver/lib/email_mirror.py

@@ -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)

zerver/lib/upload.py

@@ -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):
+    def upload_message_image(self, uploaded_file_name, uploaded_file_size,
-        # type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
+                             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):
+    def upload_message_image(self, uploaded_file_name, uploaded_file_size,
-        # type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
+                             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):
+    def upload_message_image(self, uploaded_file_name, uploaded_file_size,
-        # type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
+                             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):
+def upload_message_image(uploaded_file_name, uploaded_file_size,
-    # type: (Text, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
+                         content_type, file_data, user_profile, target_realm=None):
-    return upload_backend.upload_message_image(uploaded_file_name, content_type, file_data,
+    # type: (Text, int, Optional[Text], binary_type, UserProfile, Optional[Realm]) -> Text
-                                               user_profile, target_realm=target_realm)
+    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):
+def create_attachment(file_name, path_id, user_profile, file_size):
-    # type: (Text, Text, UserProfile) -> bool
+    # type: (Text, Text, UserProfile, int) -> bool
-    Attachment.objects.create(file_name=file_name, path_id=path_id, owner=user_profile, realm=user_profile.realm)
+    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)
+    uploaded_file_name, uploaded_file_size, content_type = get_file_info(request, user_file)
-    return upload_message_image(uploaded_file_name, content_type, user_file.read(), user_profile)
+    return upload_message_image(uploaded_file_name, uploaded_file_size,
+                                content_type, user_file.read(), user_profile)

zerver/migrations/0055_attachment_size.py

@@ -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),
+        ),
+    ]

zerver/models.py

@@ -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

zerver/tests/test_export.py

@@ -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,

zerver/tests/test_messages.py

@@ -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'),
+            ('zulip.txt', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt', sample_size),
-            ('temp_file.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py'),
+            ('temp_file.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py', sample_size),
-            ('abc.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py')
+            ('abc.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py', sample_size)
         ]
 
-        for file_name, path_id in dummy_files:
+        for file_name, path_id, size in dummy_files:
-            create_attachment(file_name, path_id, user_profile)
+            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())
 

zerver/tests/test_upload.py

@@ -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))