diff --git a/SlideServer.py b/SlideServer.py index 540396d..443a69a 100644 --- a/SlideServer.py +++ b/SlideServer.py @@ -41,7 +41,7 @@ import pymongo from bson.objectid import ObjectId - +from collections import defaultdict try: from io import BytesIO @@ -61,6 +61,7 @@ app.config['SECRET_KEY'] = os.urandom(24) app.config['ROI_FOLDER'] = "/images/roiDownload" + download_folder = os.getenv('DOWNLOAD_FOLDER', app.config['UPLOAD_FOLDER']) app.config['DOWNLOAD_FOLDER'] = download_folder @@ -72,6 +73,18 @@ os.mkdir(app.config['TEMP_FOLDER']) +# Per-token locks to prevent parallel/out-of-order chunk corruption +token_locks = defaultdict(threading.Lock) +# Move token-masking helper to module level (used for logs; keeps request handlers small) +def _mask_token(token, keep_prefix=2, keep_suffix=2): + try: + if not token or len(token) <= keep_prefix + keep_suffix: + return "****" + return token[:keep_prefix] + "*" * (len(token) - keep_prefix - keep_suffix) + token[-keep_suffix:] + except Exception: + return "****" + + # should be used instead of secure_filename to create new files whose extensions are important. # use secure_filename to access previous files. # secure_filename ensures security but may result in invalid filenames. @@ -172,25 +185,48 @@ def start_upload(): @app.route('/upload/continue/', methods=['POST']) def continue_file(token): token = secure_filename(token) - print(token, file=sys.stderr) + masked = _mask_token(token) + app.logger.info(f"[upload] continue called for token={masked}") + tmppath = os.path.join(app.config['TEMP_FOLDER'], token) - if os.path.isfile(tmppath): - body = flask.request.get_json() - if not body: - return flask.Response(json.dumps({"error": "Missing JSON body"}), status=400, mimetype='text/json') - offset = body['offset'] or 0 - if not 'data' in body: - return flask.Response(json.dumps({"error": "File data not found in body"}), status=400, mimetype='text/json') - else: - data = base64.b64decode(body['data']) - f = open(tmppath, "ab") - f.seek(int(offset)) - f.write(data) - f.close() - return flask.Response(json.dumps({"status": "OK"}), status=200, mimetype='text/json') - else: + if not os.path.isfile(tmppath): return flask.Response(json.dumps({"error": "Token Not Recognised"}), status=400, mimetype='text/json') + body = flask.request.get_json() + if not body: + return flask.Response(json.dumps({"error": "Missing JSON body"}), status=400, mimetype='text/json') + + # Validate offset presence and type + try: + offset = int(body.get('offset', 0)) + except: + return flask.Response(json.dumps({"error": "Invalid offset"}), status=400, mimetype='text/json') + + if 'data' not in body: + return flask.Response(json.dumps({"error": "File data not found in body"}), status=400, mimetype='text/json') + + # decode payload + try: + data = base64.b64decode(body['data']) + except Exception: + return flask.Response(json.dumps({"error": "Invalid base64 data"}), status=400, mimetype='text/json') + + # Acquire per-token lock and validate the offset before writing. + lock = token_locks[token] + with lock: + current_size = os.path.getsize(tmppath) + + if offset != current_size: + return flask.Response(json.dumps({ + "error": "Offset mismatch", + "expected_offset": current_size + }), status=409, mimetype='text/json') + + with open(tmppath, "ab") as f: + f.write(data) + + return flask.Response(json.dumps({"status": "OK", "written": len(data)}), status=200, mimetype='text/json') + # end the upload, by removing the in progress indication; locks further modification @app.route('/upload/finish/', methods=['POST', "GET"]) @@ -202,6 +238,39 @@ def finish_upload(token): tmppath = os.path.join(app.config['TEMP_FOLDER'], token) if not os.path.isfile(tmppath): return flask.Response(json.dumps({"error": "Token Not Recognised"}), status=400, mimetype='text/json') + + # Optional client-provided integrity hints + expected_sha256 = body.get('sha256') + expected_size = body.get('size') + + # Quick size check if client provided expected size + if expected_size is not None: + try: + expected_size = int(expected_size) + actual_size = os.path.getsize(tmppath) + if actual_size != expected_size: + app.logger.warning(f"[upload/finish] size mismatch token={token} expected={expected_size} actual={actual_size}") + return flask.Response(json.dumps({"error": "Size mismatch", "expected_size": expected_size, "actual_size": actual_size}), status=409, mimetype='text/json') + except Exception: + # ignore parse errors, fall through to normal flow + pass + + # Optional SHA256 verification + if expected_sha256: + try: + h = hashlib.sha256() + with open(tmppath, "rb") as f: + for chunk in iter(lambda: f.read(8192), b""): + h.update(chunk) + actual_sha = h.hexdigest().upper() + provided = expected_sha256.upper() + if actual_sha != provided: + app.logger.warning(f"[upload/finish] sha mismatch token={token} provided={provided} actual={actual_sha}") + return flask.Response(json.dumps({"error": "SHA256 mismatch", "expected_sha256": provided, "actual_sha256": actual_sha}), status=400, mimetype='text/json') + except Exception as e: + app.logger.error(f"[upload/finish] sha computation failed token={token}: {e}") + return flask.Response(json.dumps({"error": "Failed to verify file integrity", "detail": str(e)}), status=500, mimetype='text/json') + filename = body['filename'] if filename and verify_extension(filename): filename = secure_filename_strict(filename) @@ -215,7 +284,11 @@ def finish_upload(token): relpath = filename filepath = os.path.join(app.config['UPLOAD_FOLDER'], relpath) if not os.path.isfile(filepath): - shutil.move(tmppath, filepath) + try: + shutil.move(tmppath, filepath) + except Exception as e: + app.logger.error(f"[upload/finish] move failed token={token}: {e}") + return flask.Response(json.dumps({"error": "Failed to move finished upload", "detail": str(e)}), status=500, mimetype='text/json') return flask.Response(json.dumps({"ended": token, "filepath": filepath, "filename": filename, "relpath": relpath}), status=200, mimetype='text/json') else: return flask.Response(json.dumps({"error": "File with name '" + filename + "' already exists", "filepath": filepath, "filename": filename}), status=400, mimetype='text/json') diff --git a/test/Out-of-order-check.py b/test/Out-of-order-check.py new file mode 100644 index 0000000..701692d --- /dev/null +++ b/test/Out-of-order-check.py @@ -0,0 +1,125 @@ +import requests +import base64 +import hashlib +import threading +import time +import os + +SERVER = "http://localhost:5000" # Change if needed +TESTFILE = "Path_of_file" + +# --------- Create a large test file (20MB) ---------- +if not os.path.exists(TESTFILE): + with open(TESTFILE, "wb") as f: + f.write(os.urandom(20 * 1024 * 1024)) +print(f"Using test file: {TESTFILE}") + +# Compute SHA256 +with open(TESTFILE, "rb") as f: + ORIGINAL_SHA = hashlib.sha256(f.read()).hexdigest().upper() + +print("Original SHA256:", ORIGINAL_SHA) + + +# ---------------- Helper functions ----------------------- + +def start_upload(): + r = requests.post(f"{SERVER}/upload/start", json={"filename": TESTFILE}) + token = r.json()["upload_token"] + return token + +def send_chunk(token, offset, data): + body = { + "offset": offset, + "data": base64.b64encode(data).decode() + } + r = requests.post(f"{SERVER}/upload/continue/{token}", json=body) + return r.status_code, r.json() + + +# ------------------------- TEST 1 ------------------------- +print("\n=== TEST 1: Out-of-order upload should FAIL ===") +token = start_upload() + +with open(TESTFILE, "rb") as f: + data = f.read() + +chunk_size = 1024 * 1024 +chunks = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)] + +# Send chunk 0 (correct) +status, resp = send_chunk(token, 0, chunks[0]) +print("Chunk 0:", status, resp) + +# Send chunk 2 BEFORE chunk 1 → MUST FAIL (409) +status, resp = send_chunk(token, chunk_size * 2, chunks[2]) +print("Out-of-order chunk (expected 409):", status, resp) + +if status == 409: + print("✔ Server correctly rejected out-of-order chunk") +else: + print("❌ SERVER IS STILL VULNERABLE!") + exit() + +# ------------------------- TEST 2 ------------------------- +print("\n=== TEST 2: Parallel uploads MUST trigger 409 ===") + +token = start_upload() + +threads = [] +results = [] + +def worker(i, blob): + # ALL threads intentionally send WRONG offset (0) + status, resp = send_chunk(token, 0, blob) + results.append((i, status)) + +for i in range(10): + t = threading.Thread(target=worker, args=(i, chunks[0])) + threads.append(t) + t.start() + +for t in threads: + t.join() + +print("Parallel responses:", results) + +if any(s == 409 for _, s in results): + print("✔ Server correctly rejects parallel uploads") +else: + print("❌ Server allowed multiple writes → still broken") + + +# ------------------------- TEST 3 ------------------------- +print("\n=== TEST 3: Full Upload With Resume Must Match SHA ===") + +token = start_upload() + +offset = 0 +i = 0 + +while offset < len(data): + chunk = chunks[i] + status, resp = send_chunk(token, offset, chunk) + + if status == 409: + # Server tells us "expected_offset" + offset = resp["expected_offset"] + continue + + offset += len(chunk) + i += 1 + +# FINISH +r = requests.post(f"{SERVER}/upload/finish/{token}", json={ + "filename": "uploaded_test.bin", + "sha256": ORIGINAL_SHA, + "size": len(data) +}) + +print("Finish response:", r.status_code, r.json()) + +if r.status_code == 200: + print("✔ Full upload completed with correct SHA") +else: + print("❌ Final SHA mismatch → file corrupted")