refactor upload_partial view (998b3c54) · Commits · certic / Jama

src/rpc/tests.py

+18 −0

Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@ from io import BytesIO
		from django.urls import reverse
		from openpyxl import load_workbook
		from unittest.mock import patch
		import hashlib
		import json

		object_classes = [
		@@ -471,6 +472,23 @@ class ServiceTestCase(TestCase):

		self.assertEqual(response.status_code, 403)

		def test_upload_partial_accepts_single_chunk_upload(self):
		data = b"test"
		file_hash = hashlib.sha256(data).hexdigest()

		with tempfile.TemporaryDirectory() as tmpdir:
		with patch.object(views.settings, "PARTIAL_UPLOADS_DIR", tmpdir):
		with patch.object(views, "_enqueue_file_tasks") as enqueue_file_tasks:
		response = views.upload_partial(
		self._upload_partial_request(file_hash, "1/1", data)
		)

		self.assertEqual(response.status_code, 200)
		file_instance = models.File.objects.get(pk=int(response.content))
		self.assertEqual(file_instance.hash, file_hash)
		self.assertEqual(file_instance.original_name, "test.png")
		enqueue_file_tasks.assert_called_once_with(file_instance.pk)

		def test_upload_lock_creation_is_exclusive(self):
		with tempfile.TemporaryDirectory() as tmpdir:
		lock_file_path = os.path.join(tmpdir, ".lock")

src/rpc/views.py

+283 −148

Original line number	Diff line number	Diff line
		@@ -6,6 +6,7 @@ from django.contrib.auth.models import AbstractUser, User
		from django.contrib.auth import authenticate
		from django.shortcuts import render
		from django import forms
		from dataclasses import dataclass
		from urllib.parse import unquote_plus
		import json
		import binascii
		@@ -69,6 +70,23 @@ class BatchTooLarge(ValueError):
		pass


		@dataclass(frozen=True)
		class UploadRequestContext:
		user: User
		project: models.Project
		acl: UserAccess


		@dataclass(frozen=True)
		class PartialUploadRequest:
		origin_dir: str \| None
		file_hash: str
		file_name: str
		file_extension: str
		chunk_number: int
		total_chunks: int


		def _rpc_request_body(request: HttpRequest) -> bytes:
		max_size = settings.JAMA_RPC_MAX_BODY_SIZE
		content_length_header = request.headers.get("Content-Length")
		@@ -301,6 +319,246 @@ def _try_create_upload_lock(lock_file_path: str) -> bool:
		return True


		def _upload_request_context(
		request: HttpRequest,
		) -> Union[UploadRequestContext, HttpResponse]:
		"""Authenticate the upload request and build the authorization context."""
		user = _get_user_from_request(request)
		if not user:
		return HttpResponse("Forbidden", status=403)
		project = _get_project_from_request(request)
		if not project:
		return HttpResponse("Bad Request", status=400)

		acl = UserAccess(user, project)
		try:
		acl.check_create("resource")
		except ServiceException:
		return HttpResponse("Forbidden", status=403)
		return UploadRequestContext(user=user, project=project, acl=acl)


		def _parse_partial_upload_request(request: HttpRequest) -> PartialUploadRequest:
		"""Decode and validate the headers that identify an upload chunk."""
		origin_dir = request.headers.get("X-origin-dir", None)
		if origin_dir:
		origin_dir = _decode_base64_urlquoted_header(origin_dir)

		file_hash = request.headers["X-file-hash"]
		if re.fullmatch("[A-Fa-f0-9]{64}", file_hash) is None:
		raise ValueError

		file_name = _decode_base64_urlquoted_header(request.headers["X-file-name"])
		_, file_extension = os.path.splitext(file_name)
		if not file_extension:
		raise UnknownFileType

		chunk_number, total_chunks = _parse_upload_chunk_header(
		request.headers["X-file-chunk"]
		)
		return PartialUploadRequest(
		origin_dir=origin_dir,
		file_hash=file_hash,
		file_name=file_name,
		file_extension=file_extension,
		chunk_number=chunk_number,
		total_chunks=total_chunks,
		)


		def _restore_deleted_file_if_needed(
		file_instance: models.File,
		file_name: str,
		user: User,
		) -> None:
		"""Restore an existing soft-deleted file so the upload can reuse it."""
		if not file_instance.deleted_at:
		return

		file_instance.deleted_at = None
		file_instance.title = file_name
		file_instance.original_name = file_name
		file_instance.save()
		logger.info(
		"User {}({}) respawned file {}".format(user.username, user.pk, file_instance.pk)
		)


		def _attach_file_to_origin_collection(
		upload: PartialUploadRequest,
		file_instance: models.File,
		context: UploadRequestContext,
		) -> None:
		"""Attach the uploaded file to the collection named by X-origin-dir."""
		if not upload.origin_dir:
		return

		collection = _add_file_to_origin_collection(
		upload.origin_dir,
		file_instance,
		context.project,
		context.acl,
		)
		logger.info(
		"User {}({}) added file({}) to collection({})".format(
		context.user.username, context.user.pk, file_instance.pk, collection.pk
		)
		)


		def _existing_upload_response(
		upload: PartialUploadRequest, context: UploadRequestContext
		) -> Union[HttpResponse, None]:
		"""Return a response for an already-known file, or None if upload must continue."""
		try:
		file_instance = models.File.objects.get(
		hash=upload.file_hash, project=context.project
		)
		except models.File.DoesNotExist:
		return None

		_restore_deleted_file_if_needed(file_instance, upload.file_name, context.user)
		if upload.origin_dir:
		try:
		with transaction.atomic():
		_attach_file_to_origin_collection(upload, file_instance, context)
		except ServiceException:
		return HttpResponse("Forbidden", status=403)
		return HttpResponse(file_instance.pk, status=200)


		def _partial_upload_part_path(partials_dir: str, upload: PartialUploadRequest) -> str:
		"""Build the on-disk path for this chunk's temporary part file."""
		return "{}/{}-{}.part".format(
		partials_dir,
		upload.total_chunks,
		str(upload.chunk_number).zfill(len(str(upload.total_chunks))),
		)


		def _partial_upload_parts_glob(partials_dir: str, total_chunks: int) -> str:
		"""Build the glob pattern used to find part files for one upload."""
		return "{}/{}-*.part".format(partials_dir, total_chunks)


		def _store_partial_upload_chunk(
		request: HttpRequest,
		partials_dir: str,
		upload: PartialUploadRequest,
		) -> Union[HttpResponse, None]:
		"""Persist the request body as a complete chunk, rejecting size mismatches."""
		try:
		content_length = _request_content_length(request)
		if not _write_request_body_to_file(
		request,
		_partial_upload_part_path(partials_dir, upload),
		content_length,
		):
		return HttpResponse("Content-Length does not match body size", status=400)
		except KeyError:
		return HttpResponse("Length Required", status=411)
		except UploadLimitExceeded:
		return HttpResponse("Payload Too Large", status=413)
		except ValueError:
		return HttpResponse("Bad Request", status=400)
		return None


		def _partial_upload_is_complete(
		partials_dir: str, complete_file_path: str, upload: PartialUploadRequest
		) -> bool:
		"""Check whether all chunks are present or a previous join already completed."""
		return len(
		glob(_partial_upload_parts_glob(partials_dir, upload.total_chunks))
		) == upload.total_chunks or os.path.exists(complete_file_path)


		def _completed_upload_checksum_ok(
		partials_dir: str,
		complete_file_path: str,
		upload: PartialUploadRequest,
		) -> bool:
		"""Join chunks if needed and report whether the complete file matches the hash."""
		if not os.path.exists(complete_file_path):
		return _join_parts(
		partials_dir, complete_file_path, upload.file_hash, upload.total_chunks
		)
		return upload.file_hash == _file_hash256(complete_file_path)


		def _import_completed_partial_upload(
		complete_file_path: str,
		partials_dir: str,
		upload: PartialUploadRequest,
		context: UploadRequestContext,
		) -> int:
		"""Create or restore the File row from the completed upload."""
		with open(complete_file_path, "rb") as f:
		with transaction.atomic():
		file_id = handle_uploaded_file(
		File(f),
		context.project,
		force_file_name=upload.file_name,
		)
		_silent_rmdir(partials_dir)
		logger.info(
		"User {}({}) handled file {}".format(
		context.user.username, context.user.pk, file_id
		)
		)
		if upload.origin_dir:
		file_instance = models.File.objects.get(id=file_id)
		_attach_file_to_origin_collection(upload, file_instance, context)
		return file_id


		def _enqueue_file_tasks(file_id: int) -> None:
		"""Schedule all post-upload processing tasks for a file."""
		iiif_task.enqueue(file_id)
		exif_task.enqueue(file_id)
		ocr_task.enqueue(file_id)
		hls_task.enqueue(file_id)


		def _completed_partial_upload_response(
		partials_dir: str,
		lock_file_path: str,
		upload: PartialUploadRequest,
		context: UploadRequestContext,
		) -> HttpResponse:
		"""Serialize completion, import the file, and map upload errors to responses."""
		if not _try_create_upload_lock(lock_file_path):
		return HttpResponse(status=202)

		complete_file_path = "{}/complete{}".format(partials_dir, upload.file_extension)
		checksum_ok = _completed_upload_checksum_ok(
		partials_dir, complete_file_path, upload
		)
		try:
		try:
		file_id = _import_completed_partial_upload(
		complete_file_path, partials_dir, upload, context
		)
		_enqueue_file_tasks(file_id)
		if not checksum_ok:
		return HttpResponse(file_id, status=210)
		return HttpResponse(file_id, status=200)
		except UnknownFileType:
		_silent_rmdir(partials_dir)
		return HttpResponse("Type de fichier inconnu", status=400)
		except ResourceError:
		_silent_rmdir(partials_dir)
		return HttpResponse("Impossible d'enregistrer la ressource", status=400)
		except FileAlreadyExists as already_exists:
		return HttpResponse(str(already_exists), status=409)
		except ConcurrencyError:
		return HttpResponse("Too Many Requests", status=429)
		except ServiceException:
		return HttpResponse("Forbidden", status=403)
		except FileNotFoundError:
		return HttpResponse("Too Many Requests", status=429)


		def _rpc_method_accepts_user(fn) -> bool:
		try:
		params = list(signature(fn).parameters.values())
		@@ -448,167 +706,44 @@ def rpc(request: HttpRequest) -> HttpResponse:
		@csrf_exempt
		@require_http_methods(["POST"])
		def upload_partial(request: HttpRequest) -> HttpResponse:
		user = _get_user_from_request(request)
		if not user:
		return HttpResponse("Forbidden", status=403)
		project = _get_project_from_request(request)
		if not project:
		return HttpResponse("Bad Request", status=400)
		acl = UserAccess(user, project)
		try:
		acl.check_create("resource")
		except ServiceException:
		return HttpResponse("Forbidden", status=403)
		# Step 1: authenticate the caller and check project-level create access.
		context = _upload_request_context(request)
		if isinstance(context, HttpResponse):
		return context

		##
		## First pass is trying to find an existing file
		##
		# Step 2: decode the upload headers shared by existing-file and chunk paths.
		try:
		# optional header, used to create collections
		origin_dir = request.headers.get("X-origin-dir", None)
		if origin_dir:
		origin_dir = _decode_base64_urlquoted_header(origin_dir)
		file_hash = request.headers["X-file-hash"]
		if re.fullmatch("[A-Fa-f0-9]{64}", file_hash) is None:
		raise ValueError
		file_name = _decode_base64_urlquoted_header(request.headers["X-file-name"])
		_, f_extension = os.path.splitext(file_name)
		if not f_extension:
		return HttpResponse("Type de fichier inconnu", status=400)
		chunk_number, total_chunks = _parse_upload_chunk_header(
		request.headers["X-file-chunk"]
		)
		file_instance = models.File.objects.get(hash=file_hash, project=project)
		# Rise from your grave !
		if file_instance.deleted_at:
		file_instance.deleted_at = None
		file_instance.title = file_name
		file_instance.original_name = file_name
		file_instance.save()
		logger.info(
		"User {}({}) respawned file {}".format(
		user.username, user.pk, file_instance.pk
		)
		)
		if origin_dir:
		try:
		with transaction.atomic():
		collection = _add_file_to_origin_collection(
		origin_dir, file_instance, project, acl
		)
		except ServiceException:
		return HttpResponse("Forbidden", status=403)
		logger.info(
		"User {}({}) added file({}) to collection({})".format(
		user.username, user.pk, file_instance.pk, collection.pk
		)
		)
		return HttpResponse(file_instance.pk, status=200)
		upload = _parse_partial_upload_request(request)
		except UploadLimitExceeded:
		return HttpResponse("Payload Too Large", status=413)
		except UnknownFileType:
		return HttpResponse("Type de fichier inconnu", status=400)
		except (KeyError, ValueError, UnicodeDecodeError, binascii.Error):
		return HttpResponse(status=400)
		except models.File.DoesNotExist:
		pass # File does not exist, just go on with the upload

		##
		## Existing file not found, go on with the upload
		##
		partials_dir = _partial_upload_dir(user, project, file_hash)
		# Step 3: short-circuit when the same file hash already exists in this project.
		existing_file_response = _existing_upload_response(upload, context)
		if existing_file_response is not None:
		return existing_file_response

		# Step 4: initialize the partial upload directory unless another request is joining.
		partials_dir = _partial_upload_dir(context.user, context.project, upload.file_hash)
		lock_file_path = "{}/.lock".format(partials_dir)
		if os.path.isfile(lock_file_path):
		return HttpResponse(status=202)
		os.makedirs(partials_dir, exist_ok=True)

		# We don't want to store an incomplete chunk.
		# (prevent client disconnect)
		try:
		content_length = _request_content_length(request)
		partial_destination = "{}/{}-{}.part".format(
		partials_dir,
		total_chunks,
		str(chunk_number).zfill(len(str(total_chunks))),
		)
		if not _write_request_body_to_file(
		request, partial_destination, content_length
		):
		return HttpResponse("Content-Length does not match body size", status=400)
		except KeyError: # Content-length header is mandatory
		return HttpResponse("Length Required", status=411)
		except UploadLimitExceeded:
		return HttpResponse("Payload Too Large", status=413)
		except ValueError:
		return HttpResponse("Bad Request", status=400)

		# All chunks are complete, time to assemble
		f_name = "{}/complete{}".format(partials_dir, f_extension)
		if len(glob("{}/{}-*.part".format(partials_dir, total_chunks))) == int(
		total_chunks
		) or os.path.exists(f_name):
		if not _try_create_upload_lock(lock_file_path):
		return HttpResponse(status=202)
		# Step 5: store this chunk only after confirming the body matches Content-Length.
		chunk_response = _store_partial_upload_chunk(request, partials_dir, upload)
		if chunk_response is not None:
		return chunk_response

		if not os.path.exists(f_name):
		checksum_ok = _join_parts(partials_dir, f_name, file_hash, total_chunks)
		elif file_hash == _file_hash256(f_name):
		checksum_ok = True
		else:
		checksum_ok = False
		try:
		with open(f_name, "rb") as f:
		try:
		file_id = None
		with transaction.atomic():
		file_id = handle_uploaded_file(
		File(f),
		project,
		force_file_name=file_name,
		# Step 6: if this was the final chunk, join/import the upload; otherwise wait.
		complete_file_path = "{}/complete{}".format(partials_dir, upload.file_extension)
		if _partial_upload_is_complete(partials_dir, complete_file_path, upload):
		return _completed_partial_upload_response(
		partials_dir, lock_file_path, upload, context
		)
		f.close()
		_silent_rmdir(partials_dir)
		logger.info(
		"User {}({}) handled file {}".format(
		user.username, user.pk, file_id
		)
		)
		if origin_dir:
		file_instance = models.File.objects.get(id=file_id)
		collection = _add_file_to_origin_collection(
		origin_dir, file_instance, project, acl
		)
		logger.info(
		"User {}({}) added file({}) to collection({})".format(
		user.username, user.pk, file_id, collection.pk
		)
		)
		iiif_task.enqueue(file_id)
		exif_task.enqueue(file_id)
		ocr_task.enqueue(file_id)
		hls_task.enqueue(file_id)
		if not checksum_ok:
		return HttpResponse(file_id, status=210) # content different
		else:
		return HttpResponse(file_id, status=200)
		except UnknownFileType:
		_silent_rmdir(partials_dir)
		return HttpResponse("Type de fichier inconnu", status=400)
		except ResourceError:
		_silent_rmdir(partials_dir)
		return HttpResponse(
		"Impossible d'enregistrer la ressource", status=400
		)
		except FileAlreadyExists as already_exists:
		return HttpResponse(str(already_exists), status=409)
		except ConcurrencyError:
		# Concurrency problem with clients sending multiple chunks in parallel.
		# Nothing to do here, work has already been done.
		return HttpResponse("Too Many Requests", status=429)
		except ServiceException:
		return HttpResponse("Forbidden", status=403)

		# another race condition
		except FileNotFoundError:
		return HttpResponse("Too Many Requests", status=429)
		return HttpResponse(status=202)