mirror of
https://github.com/retspen/webvirtcloud
synced 2026-11-08 11:16:47 +00:00
158d4db6c4
The upload session was recorded only after the first chunk had been written, and a retried first chunk was refused when its part file was missing. The session (with its upload id) is now written before the part file is created, the part file is created by whichever chunk finds none, and a missing part file is an error only after the first chunk. An interrupted first chunk can be sent again without adopting an unrelated file.
843 lines
35 KiB
Python
843 lines
35 KiB
Python
import os
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from django.core.files.uploadedfile import SimpleUploadedFile
|
|
from django.test import TestCase
|
|
from django.urls import reverse
|
|
from django.contrib.auth import get_user_model
|
|
|
|
from computes.models import Compute
|
|
from storages.upload import (
|
|
handle_uploaded_file,
|
|
_acquire_upload_lock,
|
|
_atomic_finalize_ssh,
|
|
_cleanup_stale_local_uploads,
|
|
_get_upload_lock_directory,
|
|
_is_session_completed,
|
|
_sftp_replace_file,
|
|
)
|
|
from vrtManager.connection import CONN_SOCKET, CONN_SSH, CONN_TCP, CONN_TLS
|
|
|
|
User = get_user_model()
|
|
|
|
|
|
class LibvirtUploadTests(TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.lock_patch = patch("storages.upload._get_upload_lock_directory", return_value=self.temp_dir)
|
|
self.lock_patch.start()
|
|
self.conn = MagicMock()
|
|
self.conn.conn = CONN_TCP
|
|
self.volumes = {}
|
|
self.conn.get_volumes.side_effect = lambda: list(self.volumes)
|
|
self.conn.get_volume.side_effect = lambda name: self.volumes[name]
|
|
|
|
def create_volume(xml, flags):
|
|
from xml.etree import ElementTree
|
|
root = ElementTree.fromstring(xml)
|
|
name = root.findtext("name")
|
|
self.assertNotIn(name, self.volumes)
|
|
self.assertEqual(root.find("capacity").get("unit"), "bytes")
|
|
self.assertEqual(root.find("target/format").get("type"), "raw")
|
|
self.volumes[name] = MagicMock()
|
|
|
|
self.conn.pool.createXML.side_effect = create_volume
|
|
self.conn.wvm.newStream.return_value.send.side_effect = lambda data: len(data)
|
|
|
|
def tearDown(self):
|
|
self.lock_patch.stop()
|
|
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
|
|
|
def upload(self, data, index=0, total=1, upload_id="tcp-session", file_size=None):
|
|
return handle_uploaded_file(
|
|
self.conn, "/var/lib/libvirt/images", "ubuntu.iso",
|
|
SimpleUploadedFile("ubuntu.iso", data), index == total - 1,
|
|
chunk_index=index, total_chunks=total, upload_id=upload_id,
|
|
user_id=1, compute_id=2, pool="default", file_size=file_size,
|
|
)
|
|
|
|
def test_tcp_upload_retries_and_stream_offsets(self):
|
|
self.upload(b"abc", index=0, total=2, file_size=5)
|
|
self.upload(b"abc", index=0, total=2, file_size=5)
|
|
self.upload(b"de", index=1, total=2, file_size=5)
|
|
self.upload(b"de", index=1, total=2, file_size=5)
|
|
self.conn.pool.createXML.assert_called_once()
|
|
self.assertEqual(self.volumes["ubuntu.iso"].upload.call_count, 2)
|
|
self.assertEqual(self.volumes["ubuntu.iso"].upload.call_args.args[1:], (3, 2, 0))
|
|
self.assertEqual(self.conn.wvm.newStream.return_value.finish.call_count, 2)
|
|
|
|
def test_tls_uses_libvirt_stream(self):
|
|
self.conn.conn = CONN_TLS
|
|
self.upload(b"iso", file_size=3)
|
|
self.volumes["ubuntu.iso"].upload.assert_called_once()
|
|
|
|
def test_existing_volume_is_preserved(self):
|
|
self.volumes["ubuntu.iso"] = MagicMock()
|
|
with self.assertRaises(FileExistsError):
|
|
self.upload(b"iso", file_size=3)
|
|
self.conn.pool.createXML.assert_not_called()
|
|
|
|
def test_size_and_sequence_are_validated(self):
|
|
with self.assertRaises(ValueError):
|
|
self.upload(b"iso", file_size=None)
|
|
with self.assertRaises(ValueError):
|
|
self.upload(b"iso", file_size=4)
|
|
self.conn.pool.createXML.assert_not_called()
|
|
self.upload(b"ab", index=0, total=2, file_size=4)
|
|
with self.assertRaises(ValueError):
|
|
self.upload(b"cd", index=1, total=2, file_size=5)
|
|
with self.assertRaises(ValueError):
|
|
self.upload(b"cd", index=1, total=3, file_size=4)
|
|
|
|
def test_stream_failure_aborts_and_can_retry(self):
|
|
stream = self.conn.wvm.newStream.return_value
|
|
stream.send.side_effect = [1, OSError("connection lost")]
|
|
with self.assertRaises(OSError):
|
|
self.upload(b"iso", file_size=3)
|
|
stream.abort.assert_called_once()
|
|
self.assertEqual(stream.finish.call_count, 0)
|
|
stream.send.side_effect = lambda data: len(data)
|
|
self.upload(b"iso", file_size=3)
|
|
self.conn.pool.createXML.assert_called_once()
|
|
|
|
|
|
|
|
|
|
class StorageUploadSecurityTests(TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.mock_socket_conn = MagicMock()
|
|
self.mock_socket_conn.conn = CONN_SOCKET
|
|
|
|
self.mock_ssh_conn = MagicMock()
|
|
self.mock_ssh_conn.conn = CONN_SSH
|
|
self.mock_ssh_conn.host = "127.0.0.1"
|
|
self.mock_ssh_conn.login = "test"
|
|
self.mock_ssh_conn.passwd = "test"
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
|
|
|
def test_handle_uploaded_file_missing_or_invalid_upload_id(self):
|
|
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
|
|
# Missing upload_id
|
|
with self.assertRaises(ValueError) as ctx:
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, "test.iso", dummy_chunk, is_last_chunk=True)
|
|
self.assertIn("Missing required upload_id", str(ctx.exception))
|
|
|
|
# Invalid upload_id format
|
|
for bad_id in ["", " ", "sess/../id", "id with spaces", "a" * 65]:
|
|
with self.assertRaises(ValueError):
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
"test.iso",
|
|
dummy_chunk,
|
|
is_last_chunk=True,
|
|
upload_id=bad_id,
|
|
)
|
|
|
|
def test_handle_uploaded_file_invalid_names(self):
|
|
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
|
|
invalid_names = ["", " ", ".", "..", "test\x00evil.iso", "../evil.iso", "../../etc/shadow"]
|
|
|
|
for bad_name in invalid_names:
|
|
clean_base = os.path.basename(bad_name).strip()
|
|
if not clean_base or clean_base in (".", "..") or "\x00" in clean_base:
|
|
with self.assertRaises(ValueError):
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
bad_name,
|
|
dummy_chunk,
|
|
is_last_chunk=True,
|
|
upload_id="sess-valid-1",
|
|
)
|
|
|
|
def test_only_iso_extension_is_accepted(self):
|
|
chunk = SimpleUploadedFile("image.img", b"not checked")
|
|
for conn in (self.mock_socket_conn, self.mock_ssh_conn):
|
|
with self.assertRaisesRegex(ValueError, r"Only \.iso files"):
|
|
handle_uploaded_file(
|
|
conn, self.temp_dir, "image.img", chunk,
|
|
is_last_chunk=True, upload_id="iso-extension-test",
|
|
)
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn, self.temp_dir, "image.ISO",
|
|
SimpleUploadedFile("image.ISO", b"contents are not inspected"),
|
|
is_last_chunk=True, upload_id="uppercase-iso-test",
|
|
)
|
|
self.assertTrue(os.path.exists(os.path.join(self.temp_dir, "image.ISO")))
|
|
|
|
def test_handle_uploaded_file_successful_socket_upload(self):
|
|
dummy_chunk = SimpleUploadedFile("install.iso", b"ISODATA_BLOCK_123")
|
|
file_name = "install.iso"
|
|
upload_id = "sess-install-1"
|
|
|
|
# Chunk 1 (is_last_chunk=False, chunk_index=0, total_chunks=2)
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
file_name,
|
|
dummy_chunk,
|
|
is_last_chunk=False,
|
|
chunk_index=0,
|
|
total_chunks=2,
|
|
upload_id=upload_id,
|
|
)
|
|
part_path = os.path.join(self.temp_dir, f"{file_name}.wvc_upload_{upload_id}.part")
|
|
final_path = os.path.join(self.temp_dir, file_name)
|
|
self.assertTrue(os.path.exists(part_path))
|
|
self.assertFalse(os.path.exists(final_path))
|
|
|
|
# Chunk 2 (is_last_chunk=True, chunk_index=1, total_chunks=2)
|
|
dummy_chunk2 = SimpleUploadedFile("install.iso", b"_PART2")
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
file_name,
|
|
dummy_chunk2,
|
|
is_last_chunk=True,
|
|
chunk_index=1,
|
|
total_chunks=2,
|
|
upload_id=upload_id,
|
|
)
|
|
self.assertFalse(os.path.exists(part_path))
|
|
self.assertTrue(os.path.exists(final_path))
|
|
|
|
with open(final_path, "rb") as f:
|
|
self.assertEqual(f.read(), b"ISODATA_BLOCK_123_PART2")
|
|
|
|
def test_handle_uploaded_file_idempotent_retry(self):
|
|
file_name = "retry_test.iso"
|
|
upload_id = "sess-retry-123"
|
|
c0 = SimpleUploadedFile(file_name, b"CHUNK0_")
|
|
c1 = SimpleUploadedFile(file_name, b"CHUNK1_")
|
|
c2 = SimpleUploadedFile(file_name, b"CHUNK2")
|
|
|
|
# Chunk 0 initial
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c0, is_last_chunk=False, chunk_index=0, total_chunks=3, upload_id=upload_id)
|
|
# Chunk 0 retry (idempotent no-op!)
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c0, is_last_chunk=False, chunk_index=0, total_chunks=3, upload_id=upload_id)
|
|
|
|
# Chunk 1 initial
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c1, is_last_chunk=False, chunk_index=1, total_chunks=3, upload_id=upload_id)
|
|
# Chunk 1 retry (idempotent no-op!)
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c1, is_last_chunk=False, chunk_index=1, total_chunks=3, upload_id=upload_id)
|
|
|
|
# Chunk 2 initial (final chunk)
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c2, is_last_chunk=True, chunk_index=2, total_chunks=3, upload_id=upload_id)
|
|
|
|
# Chunk 2 retry (idempotent no-op, file already finalized!)
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c2, is_last_chunk=True, chunk_index=2, total_chunks=3, upload_id=upload_id)
|
|
|
|
final_file = os.path.join(self.temp_dir, file_name)
|
|
self.assertTrue(os.path.exists(final_file))
|
|
with open(final_file, "rb") as f:
|
|
self.assertEqual(f.read(), b"CHUNK0_CHUNK1_CHUNK2")
|
|
|
|
def test_handle_uploaded_file_interrupted_finalization(self):
|
|
file_name = "interrupted.iso"
|
|
upload_id = "sess-interrupted-1"
|
|
c0 = SimpleUploadedFile(file_name, b"PART1")
|
|
c1 = SimpleUploadedFile(file_name, b"PART2")
|
|
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c0, is_last_chunk=False, chunk_index=0, total_chunks=2, upload_id=upload_id)
|
|
|
|
# Simulate state where chunk 1 was recorded in .meta and .part, but finalization was interrupted
|
|
meta_file = os.path.join(self.temp_dir, f"{file_name}.wvc_upload_{upload_id}.meta")
|
|
part_file = os.path.join(self.temp_dir, f"{file_name}.wvc_upload_{upload_id}.part")
|
|
with open(part_file, "ab") as f:
|
|
f.write(b"PART2")
|
|
import json
|
|
with open(meta_file, "w") as f:
|
|
json.dump({
|
|
"file_name": file_name,
|
|
"total_chunks": 2,
|
|
"expected_chunk": 2,
|
|
"current_size": 10,
|
|
"created_at": time.time(),
|
|
"updated_at": time.time()
|
|
}, f)
|
|
|
|
# Client retries final chunk (chunk 1)
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c1, is_last_chunk=True, chunk_index=1, total_chunks=2, upload_id=upload_id)
|
|
|
|
final_file = os.path.join(self.temp_dir, file_name)
|
|
self.assertTrue(os.path.exists(final_file))
|
|
with open(final_file, "rb") as f:
|
|
self.assertEqual(f.read(), b"PART1PART2")
|
|
|
|
|
|
def test_handle_uploaded_file_inconsistent_total_chunks(self):
|
|
file_name = "inconsistent.iso"
|
|
upload_id = "sess-inc-1"
|
|
c0 = SimpleUploadedFile(file_name, b"part0")
|
|
c1 = SimpleUploadedFile(file_name, b"part1")
|
|
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c0, is_last_chunk=False, chunk_index=0, total_chunks=3, upload_id=upload_id)
|
|
|
|
with self.assertRaises(ValueError) as ctx:
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, c1, is_last_chunk=False, chunk_index=1, total_chunks=4, upload_id=upload_id)
|
|
self.assertIn("Inconsistent total_chunks", str(ctx.exception))
|
|
|
|
def test_handle_uploaded_file_missing_previous_chunk(self):
|
|
dummy_chunk = SimpleUploadedFile("broken.iso", b"data")
|
|
with self.assertRaises(ValueError):
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, "broken.iso", dummy_chunk, is_last_chunk=False, chunk_index=1, total_chunks=2, upload_id="sess-broken-1")
|
|
|
|
def test_handle_uploaded_file_destination_exists_raises(self):
|
|
file_name = "existing.iso"
|
|
final_path = os.path.join(self.temp_dir, file_name)
|
|
with open(final_path, "wb") as f:
|
|
f.write(b"already exists")
|
|
|
|
dummy_chunk = SimpleUploadedFile(file_name, b"new data")
|
|
with self.assertRaises(Exception):
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, dummy_chunk, is_last_chunk=False, chunk_index=0, total_chunks=2, upload_id="sess-exist-1")
|
|
|
|
def test_handle_uploaded_file_symlink_defense(self):
|
|
file_name = "symlink_test.iso"
|
|
upload_id = "sess-symlink-1"
|
|
target_file = os.path.join(self.temp_dir, "real.iso")
|
|
with open(target_file, "wb") as f:
|
|
f.write(b"real content")
|
|
|
|
symlink_part = os.path.join(self.temp_dir, f"{file_name}.wvc_upload_{upload_id}.part")
|
|
os.symlink(target_file, symlink_part)
|
|
|
|
dummy_chunk = SimpleUploadedFile(file_name, b"malicious chunk")
|
|
with self.assertRaises(PermissionError):
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, file_name, dummy_chunk, is_last_chunk=False, chunk_index=0, total_chunks=2, upload_id=upload_id)
|
|
|
|
def test_handle_uploaded_file_jail_containment(self):
|
|
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
|
|
with patch("os.path.basename", side_effect=lambda x: x):
|
|
with self.assertRaises(PermissionError):
|
|
handle_uploaded_file(self.mock_socket_conn, self.temp_dir, "../../../evil.iso", dummy_chunk, is_last_chunk=True, upload_id="sess-jail-1")
|
|
|
|
def test_cleanup_stale_local_uploads_strict_regex(self):
|
|
unrelated_files = [
|
|
"unrelated.part",
|
|
"important.lock",
|
|
"data.meta",
|
|
"short.123.part",
|
|
"backup.abcdefgh.meta",
|
|
"disk.12345678.part",
|
|
]
|
|
for fname in unrelated_files:
|
|
p = os.path.join(self.temp_dir, fname)
|
|
with open(p, "wb") as f:
|
|
f.write(b"keep me")
|
|
old_time = time.time() - 172800
|
|
os.utime(p, (old_time, old_time))
|
|
|
|
stale_part = os.path.join(self.temp_dir, "ubuntu.iso.wvc_upload_sess-12345678.part")
|
|
stale_meta = os.path.join(self.temp_dir, "ubuntu.iso.wvc_upload_sess-12345678.meta")
|
|
fresh_part = os.path.join(self.temp_dir, "ubuntu.iso.wvc_upload_sess-87654321.part")
|
|
|
|
with open(stale_part, "wb") as f:
|
|
f.write(b"stale part")
|
|
import json
|
|
with open(stale_meta, "w") as f:
|
|
json.dump({"upload_id": "sess-12345678", "expected_chunk": 1}, f)
|
|
with open(fresh_part, "wb") as f:
|
|
f.write(b"fresh part")
|
|
|
|
old_time = time.time() - 172800
|
|
os.utime(stale_part, (old_time, old_time))
|
|
os.utime(stale_meta, (old_time, old_time))
|
|
|
|
_cleanup_stale_local_uploads(self.temp_dir, max_age_seconds=86400)
|
|
|
|
for fname in unrelated_files:
|
|
self.assertTrue(os.path.exists(os.path.join(self.temp_dir, fname)), f"{fname} was unexpectedly deleted")
|
|
self.assertTrue(os.path.exists(fresh_part))
|
|
|
|
self.assertFalse(os.path.exists(stale_part))
|
|
self.assertFalse(os.path.exists(stale_meta))
|
|
|
|
def test_sftp_replace_file_requires_atomic_replace(self):
|
|
mock_sftp = MagicMock()
|
|
mock_sftp.posix_rename.return_value = None
|
|
_sftp_replace_file(mock_sftp, "/tmp/src", "/tmp/dst")
|
|
mock_sftp.posix_rename.assert_called_once_with("/tmp/src", "/tmp/dst")
|
|
mock_sftp.rename.assert_not_called()
|
|
|
|
mock_sftp.reset_mock()
|
|
mock_sftp.posix_rename.side_effect = IOError("posix_rename unsupported")
|
|
with self.assertRaises(OSError):
|
|
_sftp_replace_file(mock_sftp, "/tmp/src", "/tmp/dst")
|
|
mock_sftp.posix_rename.assert_called_once_with("/tmp/src", "/tmp/dst")
|
|
mock_sftp.remove.assert_not_called()
|
|
mock_sftp.rename.assert_not_called()
|
|
|
|
# posix_rename with other error must NOT fallback to remove
|
|
mock_sftp.reset_mock()
|
|
mock_sftp.posix_rename.side_effect = IOError("permission denied")
|
|
with self.assertRaises(IOError):
|
|
_sftp_replace_file(mock_sftp, "/tmp/src", "/tmp/dst")
|
|
mock_sftp.remove.assert_not_called()
|
|
|
|
mock_sftp = MagicMock(spec=["remove", "rename"])
|
|
with self.assertRaises(OSError):
|
|
_sftp_replace_file(mock_sftp, "/tmp/src", "/tmp/dst")
|
|
mock_sftp.remove.assert_not_called()
|
|
mock_sftp.rename.assert_not_called()
|
|
|
|
mock_ssh = MagicMock()
|
|
stdout = MagicMock()
|
|
stdout.channel.recv_exit_status.return_value = 0
|
|
mock_ssh.exec_command.return_value = (MagicMock(), stdout, MagicMock())
|
|
_sftp_replace_file(mock_sftp, "/tmp/src", "/tmp/dst", ssh=mock_ssh)
|
|
mock_sftp.remove.assert_not_called()
|
|
|
|
@patch("storages.upload.time.sleep")
|
|
@patch("storages.upload.fcntl.flock", side_effect=BlockingIOError)
|
|
@patch("storages.upload.LOCK_TIMEOUT_SECONDS", 0)
|
|
def test_upload_lock_times_out(self, _flock, _sleep):
|
|
with self.assertRaises(TimeoutError):
|
|
_acquire_upload_lock(123)
|
|
|
|
def test_atomic_finalize_ssh_success_and_overwrite_prevention(self):
|
|
mock_sftp = MagicMock()
|
|
|
|
_atomic_finalize_ssh(mock_sftp, "/tmp/temp.part", "/tmp/final.iso")
|
|
mock_sftp.rename.assert_called_once_with("/tmp/temp.part", "/tmp/final.iso")
|
|
|
|
# A destination created concurrently remains untouched.
|
|
mock_sftp.rename.side_effect = IOError("File exists")
|
|
mock_sftp.reset_mock()
|
|
with self.assertRaises(FileExistsError):
|
|
_atomic_finalize_ssh(mock_sftp, "/tmp/temp.part", "/tmp/final.iso")
|
|
mock_sftp.remove.assert_not_called()
|
|
|
|
def test_is_session_completed_crash_recovery(self):
|
|
upload_id = "sess-crash-rec-1"
|
|
file_name = "recovered.iso"
|
|
final_file = os.path.join(self.temp_dir, file_name)
|
|
with open(final_file, "wb") as f:
|
|
f.write(b"COMPLETE_CONTENT_12345")
|
|
file_size = os.path.getsize(final_file)
|
|
|
|
meta_file = os.path.join(self.temp_dir, f"{file_name}.wvc_upload_{upload_id}.meta")
|
|
import json
|
|
with open(meta_file, "w") as mf:
|
|
json.dump({
|
|
"file_name": file_name,
|
|
"expected_chunk": 3,
|
|
"total_chunks": 3,
|
|
"current_size": file_size,
|
|
"created_at": time.time(),
|
|
"updated_at": time.time(),
|
|
}, mf)
|
|
|
|
# No .done marker exists yet; _is_session_completed should recover the session from .meta
|
|
lock_dir = _get_upload_lock_directory()
|
|
marker_path = os.path.join(lock_dir, f"upload_{upload_id}.done")
|
|
if os.path.exists(marker_path):
|
|
os.unlink(marker_path)
|
|
|
|
recovered = _is_session_completed(upload_id, file_name, path=self.temp_dir)
|
|
self.assertTrue(recovered)
|
|
# .meta should now be cleaned up and .done marker created
|
|
self.assertFalse(os.path.exists(meta_file))
|
|
self.assertTrue(os.path.exists(marker_path))
|
|
if os.path.exists(marker_path):
|
|
os.unlink(marker_path)
|
|
|
|
def test_lock_file_not_unlinked_on_completion(self):
|
|
upload_id = "sess-lock-persist-1"
|
|
file_name = "lock_persist.iso"
|
|
dummy_chunk = SimpleUploadedFile(file_name, b"data")
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
file_name,
|
|
dummy_chunk,
|
|
is_last_chunk=True,
|
|
chunk_index=0,
|
|
total_chunks=1,
|
|
upload_id=upload_id,
|
|
)
|
|
lock_dir = _get_upload_lock_directory()
|
|
lock_path = os.path.join(lock_dir, f"upload_session_{upload_id}.lock")
|
|
self.assertTrue(os.path.exists(lock_path))
|
|
|
|
def test_destination_exists_requires_completed_session(self):
|
|
file_name = "already_here.iso"
|
|
dest_path = os.path.join(self.temp_dir, file_name)
|
|
with open(dest_path, "wb") as f:
|
|
f.write(b"existing content")
|
|
|
|
dummy_chunk = SimpleUploadedFile(file_name, b"new chunk")
|
|
with self.assertRaises(FileExistsError):
|
|
handle_uploaded_file(
|
|
self.mock_socket_conn,
|
|
self.temp_dir,
|
|
file_name,
|
|
dummy_chunk,
|
|
is_last_chunk=True,
|
|
chunk_index=0,
|
|
total_chunks=1,
|
|
upload_id="sess-rogue-123",
|
|
)
|
|
|
|
@patch("paramiko.SSHClient")
|
|
def test_sftp_upload_multi_chunk_and_retry(self, mock_ssh_cls):
|
|
mock_ssh = MagicMock()
|
|
mock_ssh_cls.return_value = mock_ssh
|
|
mock_sftp = MagicMock()
|
|
mock_ssh.open_sftp.return_value = mock_sftp
|
|
|
|
remote_files = {}
|
|
|
|
def mock_lstat(path):
|
|
if path in remote_files:
|
|
stat_obj = MagicMock()
|
|
stat_obj.st_mode = 0o100644
|
|
return stat_obj
|
|
raise FileNotFoundError(f"{path} not found")
|
|
|
|
mock_sftp.lstat.side_effect = mock_lstat
|
|
|
|
class FakeRemoteFile:
|
|
def __init__(self, path, mode):
|
|
self.path = path
|
|
self.mode = mode
|
|
if "w" in mode or path not in remote_files:
|
|
remote_files[path] = b""
|
|
self.buf = bytearray(remote_files[path])
|
|
self.pos = 0
|
|
|
|
def write(self, data):
|
|
if isinstance(data, str):
|
|
data = data.encode("utf-8")
|
|
self.buf[self.pos:self.pos + len(data)] = data
|
|
self.pos += len(data)
|
|
remote_files[self.path] = bytes(self.buf)
|
|
|
|
def read(self):
|
|
return bytes(self.buf)
|
|
|
|
def seek(self, pos):
|
|
self.pos = pos
|
|
|
|
def truncate(self, size):
|
|
del self.buf[size:]
|
|
remote_files[self.path] = bytes(self.buf)
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
mock_sftp.open.side_effect = lambda p, m: FakeRemoteFile(p, m)
|
|
|
|
def mock_rename(src, dst):
|
|
if src not in remote_files:
|
|
raise FileNotFoundError(f"{src} not found")
|
|
remote_files[dst] = remote_files.pop(src)
|
|
|
|
mock_sftp.rename.side_effect = mock_rename
|
|
mock_sftp.posix_rename.side_effect = mock_rename
|
|
mock_sftp.remove.side_effect = lambda p: remote_files.pop(p, None)
|
|
|
|
def mock_exec_command(command, **kwargs):
|
|
if "os.replace(" in command:
|
|
import shlex
|
|
src, dst = shlex.split(command)[-2:]
|
|
mock_rename(src, dst)
|
|
if "os.link(" in command:
|
|
import shlex
|
|
src, dst = shlex.split(command)[-2:]
|
|
if dst in remote_files:
|
|
raise FileExistsError(dst)
|
|
remote_files[dst] = remote_files[src]
|
|
stdout = MagicMock()
|
|
stdout.channel.recv_exit_status.return_value = 0
|
|
return MagicMock(), stdout, MagicMock()
|
|
|
|
mock_ssh.exec_command.side_effect = mock_exec_command
|
|
|
|
upload_id = "sess-sftp-multi-1"
|
|
file_name = "debian.iso"
|
|
c0 = SimpleUploadedFile(file_name, b"DEBIAN_CHUNK0_")
|
|
c1 = SimpleUploadedFile(file_name, b"DEBIAN_CHUNK1")
|
|
|
|
handle_uploaded_file(
|
|
self.mock_ssh_conn,
|
|
"/var/lib/libvirt/images",
|
|
file_name,
|
|
c0,
|
|
is_last_chunk=False,
|
|
chunk_index=0,
|
|
total_chunks=2,
|
|
upload_id=upload_id,
|
|
)
|
|
self.assertEqual(mock_ssh.connect.call_args.kwargs["timeout"], 10)
|
|
mock_sftp.get_channel.return_value.settimeout.assert_called_with(30)
|
|
|
|
# Simulate bytes written before a worker crashed without advancing metadata.
|
|
part_remote = f"/var/lib/libvirt/images/{file_name}.wvc_upload_{upload_id}.part"
|
|
remote_files[part_remote] += b"UNCOMMITTED"
|
|
|
|
handle_uploaded_file(
|
|
self.mock_ssh_conn,
|
|
"/var/lib/libvirt/images",
|
|
file_name,
|
|
c1,
|
|
is_last_chunk=True,
|
|
chunk_index=1,
|
|
total_chunks=2,
|
|
upload_id=upload_id,
|
|
)
|
|
|
|
final_remote = "/var/lib/libvirt/images/debian.iso"
|
|
self.assertIn(final_remote, remote_files)
|
|
self.assertEqual(remote_files[final_remote], b"DEBIAN_CHUNK0_DEBIAN_CHUNK1")
|
|
|
|
# Retry Chunk 1
|
|
handle_uploaded_file(
|
|
self.mock_ssh_conn,
|
|
"/var/lib/libvirt/images",
|
|
file_name,
|
|
c1,
|
|
is_last_chunk=True,
|
|
chunk_index=1,
|
|
total_chunks=2,
|
|
upload_id=upload_id,
|
|
)
|
|
self.assertEqual(remote_files[final_remote], b"DEBIAN_CHUNK0_DEBIAN_CHUNK1")
|
|
|
|
def test_handle_uploaded_file_ssh_traversal_check(self):
|
|
dummy_chunk = SimpleUploadedFile("test.iso", b"data")
|
|
with patch("os.path.basename", side_effect=lambda x: x):
|
|
with self.assertRaises(PermissionError):
|
|
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")
|
|
self.client.login(username="storage_admin", password="password")
|
|
self.compute = Compute.objects.create(name="LocalCompute", hostname="localhost", type=1, login="u", password="p")
|
|
|
|
@patch("storages.views.wvmStorage")
|
|
def test_iso_upload_rejects_other_extensions(self, mock_wvm):
|
|
mock_conn = mock_wvm.return_value
|
|
mock_conn.conn = CONN_SOCKET
|
|
mock_conn.get_target_path.return_value = "/tmp"
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
response = self.client.post(url, {
|
|
"iso_upload": "1", "file": SimpleUploadedFile("notes.txt", b"notes"),
|
|
"file_name": "notes.txt", "chunk_index": "0", "total_chunks": "1",
|
|
"upload_id": "invalid-extension-test",
|
|
})
|
|
self.assertEqual(response.status_code, 400)
|
|
self.assertIn("Only .iso files", response.json()["error"])
|
|
|
|
@patch("storages.views.handle_uploaded_file")
|
|
@patch("storages.views.wvmStorage")
|
|
def test_tcp_upload_passes_declared_size(self, mock_wvm, mock_upload):
|
|
mock_conn = mock_wvm.return_value
|
|
mock_conn.conn = CONN_TCP
|
|
mock_conn.get_target_path.return_value = "/var/lib/libvirt/images"
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
payload = {
|
|
"iso_upload": "1", "file_name": "test.iso", "chunk_index": "0",
|
|
"total_chunks": "1", "upload_id": "tcp-view-test",
|
|
}
|
|
response = self.client.post(url, {
|
|
**payload, "file": SimpleUploadedFile("test.iso", b"iso"),
|
|
})
|
|
self.assertEqual(response.status_code, 400)
|
|
mock_upload.assert_not_called()
|
|
|
|
response = self.client.post(url, {
|
|
**payload, "file": SimpleUploadedFile("test.iso", b"iso"),
|
|
"file_size": "3",
|
|
})
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(mock_upload.call_args.kwargs["file_size"], 3)
|
|
|
|
@patch("storages.views.wvmStorage")
|
|
def test_iso_upload_missing_chunk(self, mock_wvm):
|
|
mock_conn = MagicMock()
|
|
mock_conn.get_status.return_value = 1
|
|
mock_conn.get_target_path.return_value = "/tmp"
|
|
mock_conn.get_type.return_value = "dir"
|
|
mock_conn.get_autostart.return_value = 1
|
|
mock_conn.get_size.return_value = [1000, 500]
|
|
mock_conn.update_volumes.return_value = []
|
|
mock_wvm.return_value = mock_conn
|
|
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
res = self.client.post(url, {"iso_upload": "1", "file_name": "test.iso", "upload_id": "sess-view-1"})
|
|
self.assertEqual(res.status_code, 400)
|
|
self.assertIn("No file chunk was submitted", res.json().get("error", ""))
|
|
|
|
@patch("storages.views.wvmStorage")
|
|
def test_iso_upload_invalid_filename(self, mock_wvm):
|
|
mock_conn = MagicMock()
|
|
mock_conn.get_status.return_value = 1
|
|
mock_conn.get_target_path.return_value = "/tmp"
|
|
mock_conn.get_type.return_value = "dir"
|
|
mock_conn.get_autostart.return_value = 1
|
|
mock_conn.get_size.return_value = [1000, 500]
|
|
mock_conn.update_volumes.return_value = []
|
|
mock_wvm.return_value = mock_conn
|
|
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
chunk = SimpleUploadedFile("test.iso", b"chunkdata")
|
|
res = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk,
|
|
"file_name": "..\x00..",
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
"upload_id": "sess-view-2",
|
|
})
|
|
self.assertEqual(res.status_code, 400)
|
|
self.assertIn("Invalid file name", res.json().get("error", ""))
|
|
|
|
@patch("storages.views.wvmStorage")
|
|
def test_iso_upload_missing_upload_id(self, mock_wvm):
|
|
mock_conn = MagicMock()
|
|
mock_conn.get_status.return_value = 1
|
|
mock_conn.get_target_path.return_value = "/tmp"
|
|
mock_conn.get_type.return_value = "dir"
|
|
mock_conn.get_autostart.return_value = 1
|
|
mock_conn.get_size.return_value = [1000, 500]
|
|
mock_conn.update_volumes.return_value = []
|
|
mock_wvm.return_value = mock_conn
|
|
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
chunk = SimpleUploadedFile("test.iso", b"chunkdata")
|
|
res = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk,
|
|
"file_name": "test.iso",
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
})
|
|
self.assertEqual(res.status_code, 400)
|
|
self.assertIn("Missing upload_id", res.json().get("error", ""))
|
|
|
|
@patch("storages.views.wvmStorage")
|
|
def test_iso_upload_final_chunk_retry_idempotency(self, mock_wvm):
|
|
mock_conn = MagicMock()
|
|
mock_conn.conn = CONN_SOCKET
|
|
mock_conn.get_status.return_value = 1
|
|
temp_dir = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
|
|
mock_conn.get_target_path.return_value = temp_dir
|
|
mock_conn.get_type.return_value = "dir"
|
|
mock_conn.get_autostart.return_value = 1
|
|
mock_conn.get_size.return_value = [1000, 500]
|
|
mock_conn.get_volumes.return_value = []
|
|
mock_conn.update_volumes.return_value = []
|
|
mock_wvm.return_value = mock_conn
|
|
|
|
url = reverse("storage", args=[self.compute.id, "default"])
|
|
upload_id = "sess-view-retry-123"
|
|
file_name = "test_retry.iso"
|
|
chunk = SimpleUploadedFile(file_name, b"test payload")
|
|
|
|
# 1. First upload (single chunk, total_chunks=1)
|
|
res1 = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk,
|
|
"file_name": file_name,
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
"upload_id": upload_id,
|
|
})
|
|
self.assertEqual(res1.status_code, 200)
|
|
self.assertTrue(res1.json().get("success"))
|
|
|
|
# Volume now appears in libvirt pool
|
|
mock_conn.get_volumes.return_value = [file_name]
|
|
|
|
# 2. Retrying final chunk with same upload_id returns 200 idempotent success
|
|
chunk_retry = SimpleUploadedFile(file_name, b"test payload")
|
|
res2 = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk_retry,
|
|
"file_name": file_name,
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
"upload_id": upload_id,
|
|
})
|
|
self.assertEqual(res2.status_code, 200)
|
|
self.assertTrue(res2.json().get("success"))
|
|
mock_conn.get_storages.assert_not_called()
|
|
mock_conn.update_volumes.assert_not_called()
|
|
mock_conn.refresh.assert_not_called()
|
|
self.assertTrue(res2.json().get("reload"))
|
|
|
|
# 3. Upload with different upload_id returns 400 error
|
|
chunk_attacker = SimpleUploadedFile(file_name, b"evil payload")
|
|
res3 = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk_attacker,
|
|
"file_name": file_name,
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
"upload_id": "sess-attacker-999",
|
|
})
|
|
self.assertEqual(res3.status_code, 400)
|
|
self.assertIn("ISO image already exists", res3.json().get("error", ""))
|
|
|
|
# 4. Upload with same upload_id but by a different superuser cannot bypass/claim
|
|
User.objects.create_superuser(username="other_admin", password="password", email="other@example.com")
|
|
self.client.login(username="other_admin", password="password")
|
|
res4 = self.client.post(url, {
|
|
"iso_upload": "1",
|
|
"file": chunk_retry,
|
|
"file_name": file_name,
|
|
"chunk_index": "0",
|
|
"total_chunks": "1",
|
|
"upload_id": upload_id,
|
|
})
|
|
self.assertEqual(res4.status_code, 400)
|
|
self.assertIn("ISO image already exists", res4.json().get("error", ""))
|