diff --git a/storages/tests.py b/storages/tests.py index 414fad1f..dbde4edb 100644 --- a/storages/tests.py +++ b/storages/tests.py @@ -624,6 +624,38 @@ class StorageUploadSecurityTests(TestCase): handle_uploaded_file(self.mock_ssh_conn, "/var/lib/libvirt/images", "../../../evil.iso", dummy_chunk, is_last_chunk=True, upload_id="sess-ssh-1") + def test_interrupted_first_chunk_can_be_retried_without_appending_partial_data(self): + file_name = "first-chunk-retry.iso" + chunk = SimpleUploadedFile(file_name, b"complete-data") + + def interrupted(): + yield b"partial" + raise OSError("Interrupted first write") + + chunk.chunks = interrupted + with self.assertRaises(OSError): + handle_uploaded_file( + self.mock_socket_conn, self.temp_dir, file_name, chunk, True, + upload_id="first-chunk-retry", + ) + handle_uploaded_file( + self.mock_socket_conn, self.temp_dir, file_name, + SimpleUploadedFile(file_name, b"complete-data"), True, + upload_id="first-chunk-retry", + ) + with open(os.path.join(self.temp_dir, file_name), "rb") as stream: + self.assertEqual(stream.read(), b"complete-data") + + def test_first_chunk_can_retry_after_failure_before_part_file_creation(self): + chunk = SimpleUploadedFile("start-retry.iso", b"complete-data") + with patch("storages.upload._UploadTarget.write_chunk", side_effect=OSError("Cannot open part")): + with self.assertRaises(OSError): + handle_uploaded_file(self.mock_socket_conn, self.temp_dir, chunk.name, chunk, True, upload_id="start-retry") + handle_uploaded_file(self.mock_socket_conn, self.temp_dir, chunk.name, chunk, True, upload_id="start-retry") + with open(os.path.join(self.temp_dir, chunk.name), "rb") as stream: + self.assertEqual(stream.read(), b"complete-data") + + class StorageUploadViewTests(TestCase): def setUp(self): self.admin = User.objects.create_superuser(username="storage_admin", password="password", email="storage_admin@example.com") diff --git a/storages/upload.py b/storages/upload.py index 50aad82d..6120fd15 100644 --- a/storages/upload.py +++ b/storages/upload.py @@ -658,11 +658,14 @@ def handle_uploaded_file( "file_name": clean_name, "total_chunks": total_chunks, "expected_chunk": 0, "current_size": 0, "created_at": time.time(), "user_id": user_id, - "compute_id": compute_id, "pool": pool, + "compute_id": compute_id, "pool": pool, "upload_id": upload_id, } + # Record ownership before creating the part file, so an interrupted + # first write can be retried without adopting an unrelated file. + target.write_meta(meta) else: _matches_session(meta, clean_name, total_chunks, user_id, compute_id, pool) - if target.info(target.part) is None: + if meta["expected_chunk"] > 0 and target.info(target.part) is None: raise ValueError(_("Upload session not found or expired")) expected = meta["expected_chunk"] @@ -679,7 +682,9 @@ def handle_uploaded_file( % {"exp": expected, "got": chunk_index} ) - meta["current_size"] = target.write_chunk(file_chunk, meta["current_size"], new=expected == 0) + meta["current_size"] = target.write_chunk( + file_chunk, meta["current_size"], new=target.info(target.part) is None + ) meta["expected_chunk"] = expected + 1 meta["updated_at"] = time.time() target.write_meta(meta)