148 lines
6.1 KiB
Python
148 lines
6.1 KiB
Python
import os
|
|
import shutil
|
|
import pytest
|
|
import hashlib
|
|
import tempfile
|
|
from pathlib import Path
|
|
from httpx import AsyncClient, ASGITransport
|
|
|
|
from app.main import app
|
|
from app.core.database import init_db, AsyncSessionLocal
|
|
from app.core.config import settings
|
|
from app.storage.local import storage_provider
|
|
from app.models.models import Client, ClientCredential, BackupSession, BackupFile
|
|
from app.core.security import hash_token
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chunk_upload_interruption_and_resumption():
|
|
"""
|
|
CRITICAL PROOF-OF-CONCEPT TEST (#34):
|
|
1. Create a large test binary file with random bytes and known SHA-256.
|
|
2. Start upload session.
|
|
3. Upload first 50% of chunks.
|
|
4. Simulate client disconnect / network drop.
|
|
5. Reconnect client, query session status, ensure server reports already received chunks.
|
|
6. Upload ONLY the remaining 50% of chunks.
|
|
7. Complete upload session.
|
|
8. Verify server-assembled file SHA-256 matches exact original hash.
|
|
"""
|
|
await init_db()
|
|
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as ac:
|
|
# 1. Setup client and credentials in database
|
|
async with AsyncSessionLocal() as db:
|
|
client = Client(
|
|
client_code="CLIENT-TEST-01",
|
|
name="Windows 11 Test Machine",
|
|
hostname="WIN11-PROD",
|
|
os_info="Windows 11 Pro",
|
|
agent_version="1.0.0",
|
|
status="ONLINE"
|
|
)
|
|
db.add(client)
|
|
await db.flush()
|
|
|
|
device_id = "dev-test-uuid-001"
|
|
device_token = "secret-token-test-12345"
|
|
token_hash = hash_token(device_token)
|
|
|
|
cred = ClientCredential(
|
|
client_id=client.id,
|
|
device_id=device_id,
|
|
token_hash=token_hash,
|
|
name="Test Agent Credential"
|
|
)
|
|
db.add(cred)
|
|
await db.commit()
|
|
|
|
headers = {
|
|
"X-Device-Id": device_id,
|
|
"X-Device-Token": device_token
|
|
}
|
|
|
|
# 2. Generate a 12 MB test file (with 4 MB chunks -> 3 chunks total)
|
|
chunk_size = 4 * 1024 * 1024
|
|
file_size = 12 * 1024 * 1024
|
|
test_bytes = os.urandom(file_size)
|
|
expected_sha256 = hashlib.sha256(test_bytes).hexdigest()
|
|
|
|
# 3. Initialize upload session
|
|
init_resp = await ac.post("/api/upload/session", json={
|
|
"filename": "database_production.bak",
|
|
"file_size": file_size,
|
|
"sha256": expected_sha256,
|
|
"chunk_size": chunk_size
|
|
}, headers=headers)
|
|
assert init_resp.status_code == 200, f"Init failed: {init_resp.text}"
|
|
session_data = init_resp.json()
|
|
session_code = session_data["session_code"]
|
|
assert session_data["total_chunks"] == 3
|
|
assert session_data["received_chunks"] == []
|
|
|
|
# 4. Upload Chunk 0 (0MB to 4MB)
|
|
chunk_0_data = test_bytes[0 : chunk_size]
|
|
chunk_0_hash = hashlib.sha256(chunk_0_data).hexdigest()
|
|
c0_resp = await ac.post(
|
|
f"/api/upload/{session_code}/chunk?chunk_index=0&chunk_sha256={chunk_0_hash}",
|
|
content=chunk_0_data,
|
|
headers={**headers, "Content-Type": "application/octet-stream"}
|
|
)
|
|
assert c0_resp.status_code == 200
|
|
assert c0_resp.json()["is_received"] == True
|
|
|
|
# 5. SIMULATE NETWORK CUT / INTERRUPTION!
|
|
# Client disconnects. We now re-query status as if reconnecting later.
|
|
status_resp = await ac.get(f"/api/upload/{session_code}/status", headers=headers)
|
|
assert status_resp.status_code == 200
|
|
status_data = status_resp.json()
|
|
assert status_data["received_chunks"] == [0]
|
|
assert status_data["missing_chunks"] == [1, 2]
|
|
|
|
# Or re-call /session endpoint: it should return already received chunks!
|
|
resume_resp = await ac.post("/api/upload/session", json={
|
|
"filename": "database_production.bak",
|
|
"file_size": file_size,
|
|
"sha256": expected_sha256,
|
|
"chunk_size": chunk_size
|
|
}, headers=headers)
|
|
assert resume_resp.status_code == 200
|
|
assert resume_resp.json()["received_chunks"] == [0]
|
|
|
|
# 6. Upload Chunk 1 (4MB to 8MB)
|
|
chunk_1_data = test_bytes[chunk_size : 2 * chunk_size]
|
|
chunk_1_hash = hashlib.sha256(chunk_1_data).hexdigest()
|
|
c1_resp = await ac.post(
|
|
f"/api/upload/{session_code}/chunk?chunk_index=1&chunk_sha256={chunk_1_hash}",
|
|
content=chunk_1_data,
|
|
headers={**headers, "Content-Type": "application/octet-stream"}
|
|
)
|
|
assert c1_resp.status_code == 200
|
|
|
|
# 7. Upload Chunk 2 (8MB to 12MB)
|
|
chunk_2_data = test_bytes[2 * chunk_size : 3 * chunk_size]
|
|
chunk_2_hash = hashlib.sha256(chunk_2_data).hexdigest()
|
|
c2_resp = await ac.post(
|
|
f"/api/upload/{session_code}/chunk?chunk_index=2&chunk_sha256={chunk_2_hash}",
|
|
content=chunk_2_data,
|
|
headers={**headers, "Content-Type": "application/octet-stream"}
|
|
)
|
|
assert c2_resp.status_code == 200
|
|
|
|
# 8. Complete session and trigger assembly & SHA-256 integrity verification
|
|
complete_resp = await ac.post(f"/api/upload/{session_code}/complete", headers=headers)
|
|
assert complete_resp.status_code == 200, f"Complete failed: {complete_resp.text}"
|
|
complete_data = complete_resp.json()
|
|
assert complete_data["status"] == "SUCCESS"
|
|
assert complete_data["sha256"] == expected_sha256
|
|
assert complete_data["file_size"] == file_size
|
|
|
|
# 9. Verify physical file on disk and its SHA-256 directly
|
|
assembled_file_path = await storage_provider.get_file_path(complete_data["relative_path"])
|
|
assert os.path.exists(assembled_file_path)
|
|
with open(assembled_file_path, "rb") as f:
|
|
disk_bytes = f.read()
|
|
assert len(disk_bytes) == file_size
|
|
assert hashlib.sha256(disk_bytes).hexdigest() == expected_sha256
|
|
print("\n>>> PoC Test Passed: Interrupted chunk upload successfully resumed and verified bit-for-bit!")
|