upload: Enforce per-realm quota.

zerver/lib/exceptions.py

@@ -27,7 +27,7 @@ class ErrorCode(AbstractEnum):
     REQUEST_VARIABLE_MISSING = ()
     REQUEST_VARIABLE_INVALID = ()
     BAD_IMAGE = ()
-    QUOTA_EXCEEDED = ()
+    REALM_UPLOAD_QUOTA = ()
     BAD_NARROW = ()
     UNAUTHORIZED_PRINCIPAL = ()
     BAD_EVENT_QUEUE_ID = ()

zerver/lib/upload.py

@@ -50,6 +50,9 @@ DEFAULT_EMOJI_SIZE = 64
 # "file name" is the original filename provided by the user run
 # through a sanitization function.
 
+class RealmUploadQuotaError(JsonableError):
+    code = ErrorCode.REALM_UPLOAD_QUOTA
+
 attachment_url_re = re.compile('[/\-]user[\-_]uploads[/\.-].*?(?=[ )]|\Z)')
 
 def attachment_url_to_path_id(attachment_url: Text) -> Text:
@@ -181,6 +184,19 @@ def upload_image_to_s3(
 
     key.set_contents_from_string(contents, headers=headers)  # type: ignore # https://github.com/python/typeshed/issues/1552
 
+def currently_used_upload_space(realm: Realm) -> int:
+    used_space = Attachment.objects.filter(realm=realm).aggregate(Sum('size'))['size__sum']
+    if used_space is None:
+        return 0
+    return used_space
+
+def check_upload_within_quota(realm: Realm, uploaded_file_size: int) -> None:
+    if realm.upload_quota_bytes() is None:
+        return
+    used_space = currently_used_upload_space(realm)
+    if (used_space + uploaded_file_size) > realm.upload_quota_bytes():
+        raise RealmUploadQuotaError(_("Upload would exceed your organization's upload quota."))
+
 def get_file_info(request: HttpRequest, user_file: File) -> Tuple[Text, int, Optional[Text]]:
 
     uploaded_file_name = user_file.name

zerver/tests/test_upload.py

@@ -389,6 +389,44 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
             result = self.client_post("/json/user_uploads", {'f1': fp})
             assert sanitize_name(expected) in result.json()['uri']
 
+    def test_realm_quota(self) -> None:
+        """
+        Realm quota for uploading should not be exceeded.
+        """
+        self.login(self.example_email("hamlet"))
+
+        d1 = StringIO("zulip!")
+        d1.name = "dummy_1.txt"
+        result = self.client_post("/json/user_uploads", {'file': d1})
+        d1_path_id = re.sub('/user_uploads/', '', result.json()['uri'])
+        d1_attachment = Attachment.objects.get(path_id = d1_path_id)
+        self.assert_json_success(result)
+
+        realm = get_realm("zulip")
+        realm.upload_quota_gb = 1
+        realm.save(update_fields=['upload_quota_gb'])
+
+        # The size of StringIO("zulip!") is 6 bytes. Setting the size of
+        # d1_attachment to realm.upload_quota_bytes() - 11 should allow
+        # us to upload only one more attachment.
+        d1_attachment.size = realm.upload_quota_bytes() - 11
+        d1_attachment.save(update_fields=['size'])
+
+        d2 = StringIO("zulip!")
+        d2.name = "dummy_2.txt"
+        result = self.client_post("/json/user_uploads", {'file': d2})
+        self.assert_json_success(result)
+
+        d3 = StringIO("zulip!")
+        d3.name = "dummy_3.txt"
+        result = self.client_post("/json/user_uploads", {'file': d3})
+        self.assert_json_error(result, "Upload would exceed your organization's upload quota.")
+
+        realm.upload_quota_gb = None
+        realm.save(update_fields=['upload_quota_gb'])
+        result = self.client_post("/json/user_uploads", {'file': d3})
+        self.assert_json_success(result)
+
     def test_cross_realm_file_access(self) -> None:
 
         def create_user(email: Text, realm_id: Text) -> UserProfile:

zerver/views/upload.py

@@ -8,7 +8,7 @@ from django.utils.translation import ugettext as _
 from zerver.lib.request import has_request_variables, REQ
 from zerver.lib.response import json_success, json_error
 from zerver.lib.upload import upload_message_image_from_request, get_local_file_path, \
-    get_signed_upload_url, get_realm_for_filename
+    get_signed_upload_url, get_realm_for_filename, check_upload_within_quota
 from zerver.lib.validator import check_bool
 from zerver.models import UserProfile, validate_attachment_request
 from django.conf import settings
@@ -55,6 +55,7 @@ def upload_file_backend(request: HttpRequest, user_profile: UserProfile) -> Http
     if settings.MAX_FILE_UPLOAD_SIZE * 1024 * 1024 < file_size:
         return json_error(_("Uploaded file is larger than the allowed limit of %s MB") % (
             settings.MAX_FILE_UPLOAD_SIZE))
+    check_upload_within_quota(user_profile.realm, file_size)
 
     if not isinstance(user_file.name, str):
         # It seems that in Python 2 unicode strings containing bytes are