tusd: Allow user to upload files during preregistration.

Allow user to upload files during preregistration. This will be used
to upload import data.

Co-authored-by: Alex Vandiver <alexmv@zulip.com>

zerver/tests/test_tusd.py

@@ -3,17 +3,20 @@ import os
 import orjson
 from django.conf import settings
 from django.test import override_settings
+from typing_extensions import ParamSpec
 
 from zerver.lib.cache import cache_delete, get_realm_used_upload_space_cache_key
 from zerver.lib.test_classes import ZulipTestCase
-from zerver.lib.test_helpers import create_s3_buckets, use_s3_backend
+from zerver.lib.test_helpers import create_s3_buckets, find_key_by_email, use_s3_backend
 from zerver.lib.upload import sanitize_name, upload_backend, upload_message_attachment
 from zerver.lib.upload.s3 import S3UploadBackend
 from zerver.lib.utils import assert_is_not_none
-from zerver.models import Attachment, Realm
+from zerver.models import Attachment, PreregistrationRealm, Realm
 from zerver.models.realms import get_realm
 from zerver.views.tusd import TusEvent, TusHook, TusHTTPRequest, TusUpload
 
+ParamT = ParamSpec("ParamT")
+
 
 class TusdHooksTest(ZulipTestCase):
     def test_non_localhost(self) -> None:
@@ -98,7 +101,7 @@ class TusdHooksTest(ZulipTestCase):
 
 
 class TusdPreCreateTest(ZulipTestCase):
-    def request(self) -> TusHook:
+    def request(self, key: str = "") -> TusHook:
         return TusHook(
             type="pre-create",
             event=TusEvent(
@@ -114,6 +117,7 @@ class TusdPreCreateTest(ZulipTestCase):
                         "filetype": "text/plain",
                         "name": "zulip.txt",
                         "type": "text/plain",
+                        "key": key,
                     },
                     offset=0,
                     partial_uploads=None,
@@ -295,6 +299,162 @@ class TusdPreCreateTest(ZulipTestCase):
         )
         self.assertEqual(result_json["RejectUpload"], True)
 
+    def test_realm_import_data_upload(self) -> None:
+        # Choose import from slack
+        email = "ete-slack-import@zulip.com"
+        self.submit_realm_creation_form(
+            email,
+            realm_subdomain="ete-slack-import",
+            realm_name="Slack import end to end",
+            # TODO: Uncomment after adding support to form.
+            # import_from="slack",
+        )
+        prereg_realm = PreregistrationRealm.objects.get(email=email)
+        # TODO: Uncomment after adding support to form.
+        # self.assertEqual(prereg_realm.data_import_metadata["import_from"], "slack")
+        prereg_realm.data_import_metadata["import_from"] = "slack"
+        prereg_realm.save()
+
+        confirmation_key = find_key_by_email(email)
+        assert confirmation_key is not None
+
+        with self.assertLogs(level="WARNING") as warn_log:
+            result = self.client_post(
+                "/api/internal/tusd",
+                self.request(key=confirmation_key).model_dump(),
+                content_type="application/json",
+            )
+        self.assertEqual(result.status_code, 200)
+        # Verify if we tried to remove any existing upload.
+        self.assertEqual(
+            warn_log.output,
+            ["WARNING:root:slack.zip does not exist. Its entry in the database will be removed."],
+        )
+        result_json = result.json()
+        self.assertEqual(result_json.get("HttpResponse", None), None)
+        self.assertEqual(result_json.get("RejectUpload", False), False)
+        self.assertEqual(list(result_json["ChangeFileInfo"].keys()), ["ID"])
+        filename = f"import/{prereg_realm.id}/slack.zip"
+        self.assertTrue(result_json["ChangeFileInfo"]["ID"].endswith(filename))
+
+        info = TusUpload(
+            id=filename,
+            size=len("zulip!"),
+            offset=0,
+            size_is_deferred=False,
+            meta_data={
+                "filename": filename,
+                "filetype": "text/plain",
+                "name": "zulip.zip",
+                "type": "text/plain",
+                "key": confirmation_key,
+            },
+            is_final=False,
+            is_partial=False,
+            partial_uploads=None,
+            storage=None,
+        )
+        request = TusHook(
+            type="pre-finish",
+            event=TusEvent(
+                upload=info,
+                http_request=TusHTTPRequest(
+                    method="PATCH",
+                    uri=f"/api/v1/tus/{info.id}",
+                    remote_addr="12.34.56.78",
+                    header={},
+                ),
+            ),
+        )
+
+        # Post the hook saying the file is in place
+        result = self.client_post(
+            "/api/internal/tusd",
+            request.model_dump(),
+            content_type="application/json",
+        )
+
+        self.assertEqual(result.status_code, 200)
+        prereg_realm.refresh_from_db()
+        self.assertTrue(
+            prereg_realm.data_import_metadata["uploaded_import_file_name"].endswith(filename)
+        )
+
+    def test_realm_import_data_upload_size_is_deferred(self) -> None:
+        # Choose import from slack
+        email = "ete-slack-import@zulip.com"
+        self.submit_realm_creation_form(
+            email,
+            realm_subdomain="ete-slack-import",
+            realm_name="Slack import end to end",
+            # TODO: Uncomment after adding support to form.
+            # import_from="slack",
+        )
+        prereg_realm = PreregistrationRealm.objects.get(email=email)
+        # TODO: Uncomment after adding support to form.
+        # self.assertEqual(prereg_realm.data_import_metadata["import_from"], "slack")
+        prereg_realm.data_import_metadata["import_from"] = "slack"
+        prereg_realm.save()
+
+        confirmation_key = find_key_by_email(email)
+        assert confirmation_key is not None
+
+        request = self.request(key=confirmation_key)
+        request.event.upload.size = None
+        request.event.upload.size_is_deferred = True
+        result = self.client_post(
+            "/api/internal/tusd",
+            request.model_dump(),
+            content_type="application/json",
+        )
+        self.assertEqual(result.status_code, 200)
+        result_json = result.json()
+        self.assertEqual(result_json["HttpResponse"]["StatusCode"], 411)
+        self.assertEqual(
+            orjson.loads(result_json["HttpResponse"]["Body"]),
+            {"message": "SizeIsDeferred is not supported"},
+        )
+        self.assertEqual(result_json["RejectUpload"], True)
+
+    def test_realm_import_data_upload_size_limit_exceeded(self) -> None:
+        # Choose import from slack
+        email = "ete-slack-import@zulip.com"
+        self.submit_realm_creation_form(
+            email,
+            realm_subdomain="ete-slack-import",
+            realm_name="Slack import end to end",
+            # TODO: Uncomment after adding support to form.
+            # import_from="slack",
+        )
+        prereg_realm = PreregistrationRealm.objects.get(email=email)
+        # TODO: Uncomment after adding support to form.
+        # self.assertEqual(prereg_realm.data_import_metadata["import_from"], "slack")
+        prereg_realm.data_import_metadata["import_from"] = "slack"
+        prereg_realm.save()
+
+        confirmation_key = find_key_by_email(email)
+        assert confirmation_key is not None
+
+        request = self.request(key=confirmation_key)
+
+        max_upload_size = settings.MAX_WEB_DATA_IMPORT_SIZE_MB * 1024 * 1024
+        request.event.upload.size = max_upload_size + 1
+        result = self.client_post(
+            "/api/internal/tusd",
+            request.model_dump(),
+            content_type="application/json",
+        )
+        self.assertEqual(result.status_code, 200)
+        result_json = result.json()
+        self.assertEqual(result_json["HttpResponse"]["StatusCode"], 413)
+        self.assertEqual(
+            orjson.loads(result_json["HttpResponse"]["Body"]),
+            {
+                "message": f"Uploaded file is larger than the allowed limit of {settings.MAX_WEB_DATA_IMPORT_SIZE_MB} MiB"
+            },
+        )
+        self.assertEqual(result_json["RejectUpload"], True)
+
 
 class TusdPreFinishTest(ZulipTestCase):
     def request(self, info: TusUpload) -> TusHook:

zerver/views/tusd.py

@@ -11,6 +11,7 @@ from django.views.decorators.csrf import csrf_exempt
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic.alias_generators import to_pascal
 
+from confirmation.models import Confirmation, ConfirmationKeyError, get_object_from_key
 from zerver.decorator import get_basic_credentials, validate_api_key
 from zerver.lib.exceptions import AccessDeniedError, JsonableError
 from zerver.lib.mime_types import guess_type
@@ -21,11 +22,12 @@ from zerver.lib.upload import (
     attachment_vips_source,
     check_upload_within_quota,
     create_attachment,
+    delete_message_attachment,
     sanitize_name,
     upload_backend,
 )
 from zerver.lib.upload.base import INLINE_MIME_TYPES
-from zerver.models import Realm, UserProfile
+from zerver.models import PreregistrationRealm, Realm, UserProfile
 
 
 # See https://tus.github.io/tusd/advanced-topics/hooks/ for the spec
@@ -237,6 +239,40 @@ def authenticate_user(request: HttpRequest) -> UserProfile | AnonymousUser:
     return request.user
 
 
+def handle_preregistration_pre_create_hook(
+    request: HttpRequest, preregistration_realm: PreregistrationRealm, data: TusUpload
+) -> HttpResponse:
+    max_upload_size = settings.MAX_WEB_DATA_IMPORT_SIZE_MB * 1024 * 1024  # 1G
+    if data.size_is_deferred or data.size is None:
+        return reject_upload("SizeIsDeferred is not supported", 411)
+    if data.size > max_upload_size:
+        return reject_upload(
+            _("Uploaded file is larger than the allowed limit of {max_file_size} MiB").format(
+                max_file_size=settings.MAX_WEB_DATA_IMPORT_SIZE_MB
+            ),
+            413,
+        )
+
+    filename = f"import/{preregistration_realm.id}/slack.zip"
+
+    # Delete any existing upload, so tusd doesn't declare that there's nothing
+    # to do. This also has the nice benefit of deleting the previous upload.
+    delete_message_attachment(filename)
+
+    return tusd_json_response({"ChangeFileInfo": {"ID": filename}})
+
+
+def handle_preregistration_pre_finish_hook(
+    request: HttpRequest, preregistration_realm: PreregistrationRealm, data: TusUpload
+) -> HttpResponse:
+    # Save the filename to display the uploaded file to user. We need to store it in
+    # the database so that is available even after a refresh.
+    filename = data.meta_data["filename"]
+    preregistration_realm.data_import_metadata["uploaded_import_file_name"] = filename
+    preregistration_realm.save(update_fields=["data_import_metadata"])
+    return tusd_json_response({})
+
+
 @csrf_exempt
 @typed_endpoint
 def handle_tusd_hook(
@@ -248,14 +284,34 @@ def handle_tusd_hook(
     if not is_local_addr(request.META["REMOTE_ADDR"]):
         raise AccessDeniedError
 
+    hook_name = payload.type
     maybe_user = authenticate_user(request)
-    if isinstance(maybe_user, AnonymousUser):
+    if maybe_user.is_authenticated:
+        # Authenticated requests are file upload requests
+        if hook_name == "pre-create":
+            return handle_upload_pre_create_hook(request, maybe_user, payload.event.upload)
+        elif hook_name == "pre-finish":
+            return handle_upload_pre_finish_hook(request, maybe_user, payload.event.upload)
+        else:
+            return HttpResponseNotFound()
+
+    # Check if unauthenticated requests are for realm creation
+    key = payload.event.upload.meta_data.get("key")
+    if key is None:
+        return reject_upload("Unauthenticated upload", 401)
+    try:
+        prereg_object = get_object_from_key(
+            key, [Confirmation.REALM_CREATION], mark_object_used=False
+        )
+    except ConfirmationKeyError:
         return reject_upload("Unauthenticated upload", 401)
 
-    hook_name = payload.type
+    assert isinstance(prereg_object, PreregistrationRealm)
+    assert prereg_object.created_realm is None
+
     if hook_name == "pre-create":
-        return handle_upload_pre_create_hook(request, maybe_user, payload.event.upload)
+        return handle_preregistration_pre_create_hook(request, prereg_object, payload.event.upload)
     elif hook_name == "pre-finish":
-        return handle_upload_pre_finish_hook(request, maybe_user, payload.event.upload)
-    else:
+        return handle_preregistration_pre_finish_hook(request, prereg_object, payload.event.upload)
+    else:  # nocoverage
         return HttpResponseNotFound()

zproject/default_settings.py

@@ -674,3 +674,7 @@ CUSTOM_AUTHENTICATION_WRAPPER_FUNCTION: Callable[..., Any] | None = None
 # notification to a stream and also delete the previous counter
 # notification.
 RESOLVE_TOPIC_UNDO_GRACE_PERIOD_SECONDS = 60
+
+# For realm imports during registration, maximum size of file
+# that can be uploaded.
+MAX_WEB_DATA_IMPORT_SIZE_MB = 1024