diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 728373d..3bad1bf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -69,7 +69,7 @@ jobs: # tests need. numpy is required because constants.py imports it; requests # because Authlib's Flask client imports it. - name: Install test deps - run: python -m pip install --upgrade pip pytest werkzeug SQLAlchemy Flask-SQLAlchemy alembic argon2-cffi Authlib requests flask-cors numpy python-dotenv + run: python -m pip install --upgrade pip pytest werkzeug SQLAlchemy Flask-SQLAlchemy alembic argon2-cffi Authlib requests flask-cors numpy nibabel python-dotenv # Catches syntax errors anywhere in the backend (api_blueprint.py etc.) # without installing the heavy runtime deps - nothing else in CI parses @@ -80,6 +80,11 @@ jobs: - name: Path-safety unit tests run: python -m pytest tests/unit/test_path_safety.py -v + # User-dataset admission gatekeeper: quotas, dedup, CT/segmentation validity, + # promotion into the PanTS-mirroring layout. + - name: User-dataset gatekeeper unit tests + run: python -m pytest tests/unit/test_user_dataset.py -v + # Chunked-upload staging area: which chunks the server still holds (what a # resuming client checks its cursor against) and the stale-upload sweep. - name: Chunk-store unit tests diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index db854e2..939d07e 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -1032,6 +1032,11 @@ def _start_auto_segmentation(session_id, model_name, ct_file=None, server_input_ user = current_user() if user is None: return jsonify({"error": "Sign in to run inference"}), 401 + # Captured in the request context so the background worker (which has none) can + # attribute the scan for per-IP quotas in the user-dataset gatekeeper. Use + # remote_addr (set by the trusted reverse proxy) rather than a client-supplied + # X-Forwarded-For, which an attacker could rotate per request to defeat the cap. + _collector_ip = request.remote_addr or "" blocked = plan_store.check_inference(user["id"], model_name) if blocked is not None: # 402 Payment Required: the request is well-formed and the user is @@ -1118,6 +1123,20 @@ def _on_gpu_slot(): _set_inference_job(session_id, status="completed", error=None, zip_path=zip_path, output_mask_dir=output_mask_dir) print(f"✅ Finished segmentation and zipping for session {session_id}") + + # Non-blocking: offer this scan+mask to the user-dataset gatekeeper, + # which decides (async) whether it's worth keeping. The result is + # already delivered above; this never affects the user, and is a no-op + # unless USER_DATASET_PATH is configured. + try: + from services.user_dataset import collect_user_scan_async + collect_user_scan_async( + ct_path=input_path, output_mask_dir=output_mask_dir, + model=model_name, user_id=user.get("id"), ip=_collector_ip, + session_id=session_id, + ) + except Exception as _ude: + print(f"[user_dataset] hook error (non-fatal): {_ude}") except Exception as e: # A killed subprocess surfaces here as CalledProcessError/RuntimeError; # if the user cancelled, keep "cancelled" rather than reporting failure. diff --git a/flask-server/constants.py b/flask-server/constants.py index 4128499..fc56200 100644 --- a/flask-server/constants.py +++ b/flask-server/constants.py @@ -33,6 +33,11 @@ class Constants: CANCERVERSE_PATH = os.environ.get('CANCERVERSE_PATH') CANCERVERSE_LOWRES_PATH = os.environ.get('CANCERVERSE_LOWRES_PATH', '/home/visitor/cancerverse_lowres') DATASET_PREFIXES = {'PanTS': 'PanTS', 'CancerVerse': 'CV'} + # Where accepted user scans (CT + mask + sublabels) are collected, in a + # PanTS-mirroring layout (image_only/, mask_only/). Unset => the collection + # gatekeeper (services/user_dataset.py) is a no-op. Point it at a writable + # staging dir now; relocate beside PanTS/CancerVerse once write access lands. + USER_DATASET_PATH = os.environ.get('USER_DATASET_PATH') PERMISSIONS_DIR = os.environ.get('PERMISSIONS_DIR', "/home/visitor/data") MESH_PATH = PERMISSIONS_DIR + "/render_only" CASE_QUALITY_MANIFEST = os.environ.get('BODYMAPS_CASE_QUALITY_MANIFEST') diff --git a/flask-server/services/user_dataset.py b/flask-server/services/user_dataset.py new file mode 100644 index 0000000..d05c7da --- /dev/null +++ b/flask-server/services/user_dataset.py @@ -0,0 +1,457 @@ +"""User-data collection with an admission gatekeeper. + +When a user runs inference we may keep their CT + our segmentation to grow an +in-house dataset — but ONLY if it passes an admission gate. This is deliberately +decoupled from the user's experience: it runs in a background thread AFTER the +result is already delivered, and any failure here is swallowed. A user who +uploads garbage still gets their result; the garbage just doesn't enter the +dataset. + +The gate has four jobs (see `evaluate`): + 1. Is it a real CT? -- valid 3D volume, plausible dims, CT-like HU range. + 2. Is it a duplicate? -- exact + near-duplicate fingerprint vs the registry. + 3. Is it usable? -- our own segmentation found plausible organs. + 4. Is it abuse? -- per-user / per-IP rolling quotas + a size cap, so a + competitor dumping thousands of files can't flood us. + +Accepted scans are promoted into a PanTS-mirroring layout under +``Constants.USER_DATASET_PATH`` (so it can sit beside PanTS/CancerVerse once a +writable location is granted): + + UserData/ + image_only/USER_00000001/ct.nii.gz + mask_only/USER_00000001/combined_labels.nii.gz + mask_only/USER_00000001/segmentations/.nii.gz + mask_only/USER_00000001/metadata.json + registry.json # case-id counter, fingerprints, per-user counts + rejections.jsonl # audit trail of what was turned away and why + +The whole feature is inert unless ``USER_DATASET_PATH`` is configured, so merging +this changes nothing until it is switched on in the environment. +""" +from __future__ import annotations + +import contextlib +import hashlib +import json +import os +import shutil +import threading +import time +import traceback +from datetime import datetime, timezone +from typing import Optional + +# Heavy/optional imports (nibabel, numpy) are done lazily inside the worker so +# importing this module in the request path stays cheap and never fails a route. + +# Viewer label id -> organ name, for splitting combined_labels into per-organ +# sublabels. Mirrors _VIEWER_LABELS in auto_segmentor.py / the frontend scheme. +_LABEL_NAMES = { + 1: "adrenal_gland_left", 2: "adrenal_gland_right", 3: "aorta", 4: "bladder", + 5: "celiac_artery", 6: "colon", 7: "common_bile_duct", 8: "duodenum", + 9: "femur_left", 10: "femur_right", 11: "gall_bladder", 12: "kidney_left", + 13: "kidney_right", 14: "liver", 15: "lung_left", 16: "lung_right", + 17: "pancreas", 18: "pancreas_body", 19: "pancreas_head", 20: "pancreas_tail", + 21: "pancreatic_duct", 22: "pancreatic_lesion", 23: "postcava", 24: "prostate", + 25: "spleen", 26: "stomach", 27: "superior_mesenteric_artery", 28: "veins", + 29: "intestine", 30: "renal_vein_left", 31: "renal_vein_right", 32: "cbd_stent", + 33: "liver_lesion", 34: "kidney_lesion", 35: "colon_lesion", +} + +# --- tunables (all overridable via env) --- +MAX_CT_BYTES = int(os.environ.get("USER_DATASET_MAX_CT_BYTES", str(600 * 1024 * 1024))) # 600 MB +# The .nii.gz cap is on the COMPRESSED bytes; a header can still declare a volume +# that decodes to tens of GB (a near-constant "gzip bomb" fits well under 600 MB). +# Bound the DECODED voxel count too, before any np.asarray, so a crafted scan +# can't OOM-kill the worker. 400M voxels ~= 0.8 GB int16 / 1.6 GB float32. +MAX_VOXELS = int(os.environ.get("USER_DATASET_MAX_VOXELS", str(400_000_000))) +DAILY_PER_USER = int(os.environ.get("USER_DATASET_DAILY_PER_USER", "50")) +DAILY_PER_IP = int(os.environ.get("USER_DATASET_DAILY_PER_IP", "50")) +DAILY_GLOBAL = int(os.environ.get("USER_DATASET_DAILY_GLOBAL", "2000")) +MIN_ORGAN_VOXELS = int(os.environ.get("USER_DATASET_MIN_ORGAN_VOXELS", "20000")) +MIN_DISTINCT_ORGANS = int(os.environ.get("USER_DATASET_MIN_DISTINCT_ORGANS", "3")) +NEAR_DUP_GRID = 24 # downsample edge for the perceptual (near-duplicate) fingerprint + +_WINDOW_SECONDS = 24 * 3600 +_registry_lock = threading.Lock() + + +@contextlib.contextmanager +def _locked(root: str): + """Serialize the registry read-modify-write across BOTH threads (in-process) + and processes. Prod is gunicorn --workers 1, where the threading lock alone + suffices; the fcntl file lock makes it correct even if that ever changes. + Falls back to threads-only where fcntl is unavailable (e.g. Windows/CI).""" + with _registry_lock: + lf = None + try: + import fcntl + lf = open(os.path.join(root, ".registry.lock"), "w") + fcntl.flock(lf, fcntl.LOCK_EX) + except Exception: + if lf is not None: + try: + lf.close() + except Exception: + pass + lf = None + try: + yield + finally: + if lf is not None: + try: + import fcntl + fcntl.flock(lf, fcntl.LOCK_UN) + lf.close() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Root / registry helpers +# --------------------------------------------------------------------------- +def _root() -> Optional[str]: + """Configured dataset root, or None when the feature is switched off.""" + p = os.environ.get("USER_DATASET_PATH", "").strip() + return p or None + + +def _registry_path(root: str) -> str: + return os.path.join(root, "registry.json") + + +def _load_registry(root: str) -> dict: + path = _registry_path(root) + if os.path.exists(path): + try: + with open(path, "r", encoding="utf-8") as f: + reg = json.load(f) + except Exception: + reg = {} + else: + reg = {} + reg.setdefault("next_id", 1) + reg.setdefault("sha256", {}) # exact-dup: fingerprint -> case_id + reg.setdefault("phash", {}) # near-dup: fingerprint -> case_id + reg.setdefault("events", []) # [{ts, user_id, ip}] for rolling quotas + return reg + + +def _save_registry(root: str, reg: dict) -> None: + path = _registry_path(root) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(reg, f) + os.replace(tmp, path) # atomic + + +def _prune_events(reg: dict, now: float) -> None: + cutoff = now - _WINDOW_SECONDS + reg["events"] = [e for e in reg.get("events", []) if e.get("ts", 0) >= cutoff] + + +# --------------------------------------------------------------------------- +# Gatekeeper checks (pure-ish; unit-tested) +# --------------------------------------------------------------------------- +def validate_ct(ct_path: str): + """(ok, reason). A real abdominal/body CT: valid 3D volume, plausible dims, + CT-like HU range (air near -1000, tissue/bone above), not constant.""" + try: + import numpy as np + import nibabel as nib + except Exception as e: # pragma: no cover - deps present in prod env + return False, f"deps_missing:{e}" + if not os.path.exists(ct_path): + return False, "missing_file" + try: + size = os.path.getsize(ct_path) + except OSError: + return False, "unstatable" + if size > MAX_CT_BYTES: + return False, "too_large" + if size < 5 * 1024: + return False, "too_small" + try: + img = nib.load(ct_path) + shape = img.shape + except Exception as e: + return False, f"unreadable:{type(e).__name__}" + dims = [d for d in shape if d > 1] + if len(dims) != 3: + return False, f"not_3d:{list(shape)}" + if any(d < 16 or d > 2048 for d in dims): + return False, f"implausible_dims:{dims}" + # Reject on the DECODED size before materializing anything (anti-OOM, C1). + voxels = 1 + for d in dims: + voxels *= int(d) + if voxels > MAX_VOXELS: + return False, f"too_many_voxels:{voxels}" + try: + # Native dtype (usually int16) rather than forcing float32 -- half the + # memory, and CT intensities are integral anyway. + arr = np.asarray(img.dataobj) + except Exception as e: + return False, f"undecodable:{type(e).__name__}" + # Require ALL values finite: nanmin/nanmax ignore NaN but not +/-inf, which + # would otherwise sail through the range check and poison the fingerprint. + if arr.dtype.kind == "f" and not np.isfinite(arr).all(): + return False, "non_finite_values" + lo, hi = float(np.min(arr)), float(np.max(arr)) + if lo == hi: + return False, "constant_volume" + # CT signature: air present (well below 0) and tissue/bone present (above 0). + if lo > -200 or hi < 100: + return False, f"non_ct_intensity:[{lo:.0f},{hi:.0f}]" + return True, "ok" + + +def fingerprints(ct_path: str): + """(sha256, phash). sha256 = exact-duplicate key over the raw file; phash = + a coarse content hash (downsampled + quantized) that catches re-uploads that + differ only in compression/metadata.""" + import numpy as np + import nibabel as nib + h = hashlib.sha256() + with open(ct_path, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + sha = h.hexdigest() + try: + img = nib.load(ct_path) + vox = 1 + for d in img.shape: + vox *= int(d) + if vox > MAX_VOXELS: # defensive: validate_ct already rejects these + return sha, sha + arr = np.squeeze(np.asarray(img.dataobj)) # native dtype + g = NEAR_DUP_GRID + # block-average to a fixed GxGxG grid, quantize, hash + idx = [np.linspace(0, s, g + 1).astype(int) for s in arr.shape[:3]] + small = np.zeros((g, g, g), dtype="float32") + for i in range(g): + for j in range(g): + for k in range(g): + blk = arr[idx[0][i]:idx[0][i + 1], idx[1][j]:idx[1][j + 1], idx[2][k]:idx[2][k + 1]] + small[i, j, k] = float(blk.mean()) if blk.size else 0.0 + q = np.round(small / 50.0).astype("int16") # 50 HU buckets + ph = hashlib.sha256(q.tobytes()).hexdigest() + except Exception: + ph = sha # fall back to exact key if the perceptual pass fails + return sha, ph + + +def segmentation_quality_ok(combined_labels_path: str): + """(ok, reason, stats). Our own mask is the judge: a real abdominal CT yields + plausible organs. Reject empty/degenerate segmentations (garbage in -> nothing + the model can find).""" + try: + import numpy as np + import nibabel as nib + except Exception as e: # pragma: no cover + return False, f"deps_missing:{e}", {} + if not os.path.exists(combined_labels_path): + return False, "no_mask", {} + try: + data = np.asarray(nib.load(combined_labels_path).dataobj) + except Exception as e: + return False, f"mask_unreadable:{type(e).__name__}", {} + labels = [int(v) for v in np.unique(data) if int(v) != 0] + organ_voxels = int((data != 0).sum()) + stats = {"organ_voxels": organ_voxels, "distinct_organs": len(labels)} + if organ_voxels < MIN_ORGAN_VOXELS: + return False, f"too_few_organ_voxels:{organ_voxels}", stats + if len(labels) < MIN_DISTINCT_ORGANS: + return False, f"too_few_organs:{len(labels)}", stats + return True, "ok", stats + + +def check_quota(reg: dict, user_id: Optional[str], ip: Optional[str], now: float): + """(ok, reason). Rolling 24h caps: per-user, per-IP, and global. Anti-flood: + a burst from one account/IP is bounded regardless of dedup.""" + events = reg.get("events", []) + cutoff = now - _WINDOW_SECONDS + recent = [e for e in events if e.get("ts", 0) >= cutoff] + if len(recent) >= DAILY_GLOBAL: + return False, "global_quota" + if user_id and sum(1 for e in recent if e.get("user_id") == user_id) >= DAILY_PER_USER: + return False, "user_quota" + if ip and sum(1 for e in recent if e.get("ip") == ip) >= DAILY_PER_IP: + return False, "ip_quota" + return True, "ok" + + +# --------------------------------------------------------------------------- +# Promotion +# --------------------------------------------------------------------------- +def _write_sublabels(combined_labels_path: str, seg_dir: str, existing_seg_dir: Optional[str]) -> list: + """Per-organ binary masks. Prefer masks a model already produced; otherwise + split combined_labels by label id.""" + import numpy as np + import nibabel as nib + os.makedirs(seg_dir, exist_ok=True) + written = [] + if existing_seg_dir and os.path.isdir(existing_seg_dir): + for fn in sorted(os.listdir(existing_seg_dir)): + if fn.endswith(".nii.gz"): + shutil.copy2(os.path.join(existing_seg_dir, fn), os.path.join(seg_dir, fn)) + written.append(fn[:-len(".nii.gz")]) + if written: + return written + img = nib.load(combined_labels_path) + data = np.asarray(img.dataobj) + for lab in [int(v) for v in np.unique(data) if int(v) != 0]: + name = _LABEL_NAMES.get(lab, f"label_{lab}") + binary = (data == lab).astype("uint8") + # Fresh header (affine only) with an explicit uint8 dtype, so the source + # label map's datatype / scl_slope / scl_inter can't misencode the mask. + out = nib.Nifti1Image(binary, img.affine) + out.set_data_dtype(np.uint8) + nib.save(out, os.path.join(seg_dir, f"{name}.nii.gz")) + written.append(name) + return written + + +def _promote(root: str, case_id: str, ct_path: str, combined_labels_path: str, + existing_seg_dir: Optional[str], metadata: dict) -> None: + """Write all of a case's files into `.partial` staging dirs, then publish each + with an atomic rename. A partial failure (disk full, bad NIfTI, killed thread) + leaves only staging dirs behind -- never a half-written case that the next + admission could merge into.""" + img_final = os.path.join(root, "image_only", case_id) + mask_final = os.path.join(root, "mask_only", case_id) + img_stage, mask_stage = img_final + ".partial", mask_final + ".partial" + for d in (img_stage, mask_stage): + shutil.rmtree(d, ignore_errors=True) # clear leftovers from a prior crash + try: + os.makedirs(img_stage, exist_ok=True) + os.makedirs(mask_stage, exist_ok=True) + shutil.copy2(ct_path, os.path.join(img_stage, "ct.nii.gz")) + shutil.copy2(combined_labels_path, os.path.join(mask_stage, "combined_labels.nii.gz")) + organs = _write_sublabels(combined_labels_path, os.path.join(mask_stage, "segmentations"), + existing_seg_dir) + meta = {**metadata, "case_id": case_id, "organs": organs} + with open(os.path.join(mask_stage, "metadata.json"), "w", encoding="utf-8") as f: + json.dump(meta, f, indent=2) + os.replace(img_stage, img_final) # atomic publish (final can't pre-exist: + os.replace(mask_stage, mask_final) # case_id is a freshly reserved counter) + except Exception: + shutil.rmtree(img_stage, ignore_errors=True) + shutil.rmtree(mask_stage, ignore_errors=True) + raise + + +def _record_rejection(root: str, reason: str, ctx: dict) -> None: + try: + with open(os.path.join(root, "rejections.jsonl"), "a", encoding="utf-8") as f: + f.write(json.dumps({"ts": time.time(), "reason": reason, **ctx}) + "\n") + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- +def evaluate(ct_path: str, combined_labels_path: str, reg: dict, + user_id: Optional[str], ip: Optional[str], now: float): + """Run the four gates. Returns (accepted: bool, reason: str, extra: dict).""" + ok, reason = check_quota(reg, user_id, ip, now) + if not ok: + return False, reason, {} + ok, reason = validate_ct(ct_path) + if not ok: + return False, reason, {} + ok, reason, stats = segmentation_quality_ok(combined_labels_path) + if not ok: + return False, reason, {"stats": stats} + sha, ph = fingerprints(ct_path) + if sha in reg["sha256"]: + return False, "duplicate_exact", {"sha256": sha, "of": reg["sha256"][sha]} + if ph in reg["phash"]: + return False, "duplicate_near", {"phash": ph, "of": reg["phash"][ph]} + return True, "ok", {"sha256": sha, "phash": ph, "stats": stats} + + +def _admit_and_store(ct_path: str, output_mask_dir: str, model: str, + user_id: Optional[str], ip: Optional[str], session_id: Optional[str]) -> None: + root = _root() + if not root: + return # feature off + if not os.path.isfile(ct_path): + return # e.g. ShapeKit is handed a segmentation directory, not a CT file + combined = os.path.join(output_mask_dir, "combined_labels.nii.gz") + existing_seg = os.path.join(output_mask_dir, "segmentations") + ctx = {"user_id": user_id, "ip": ip, "model": model, "session_id": session_id} + try: + os.makedirs(root, exist_ok=True) + now = time.time() + + # 1) Quota + attempt accounting (cheap, under the lock). The attempt is + # recorded regardless of the outcome below, so a flood of REJECTS is + # rate-limited too -- not just accepted scans. + with _locked(root): + reg = _load_registry(root) + _prune_events(reg, now) + ok, qreason = check_quota(reg, user_id, ip, now) + reg["events"].append({"ts": now, "user_id": user_id, "ip": ip}) + _save_registry(root, reg) + if not ok: + _record_rejection(root, qreason, ctx) + return + + # 2) Expensive gates OUTSIDE the lock (they touch no registry state), so a + # single large scan doesn't serialize every other collection thread. + ok, reason = validate_ct(ct_path) + if not ok: + _record_rejection(root, reason, ctx) + return + ok, reason, stats = segmentation_quality_ok(combined) + if not ok: + _record_rejection(root, reason, {**ctx, "stats": stats}) + return + sha, ph = fingerprints(ct_path) + + # 3) Dedup + reserve id + atomic promote + commit, under the lock (short). + with _locked(root): + reg = _load_registry(root) + if sha in reg["sha256"]: + _record_rejection(root, "duplicate_exact", {**ctx, "of": reg["sha256"][sha]}) + return + if ph in reg["phash"]: + _record_rejection(root, "duplicate_near", {**ctx, "of": reg["phash"][ph]}) + return + case_id = "USER_%08d" % reg["next_id"] + metadata = { + "user_id": user_id, "source_ip": ip, "model": model, + "session_id": session_id, + "collected_at": datetime.now(timezone.utc).isoformat(), + "sha256": sha, "phash": ph, "segmentation_stats": stats, + } + _promote(root, case_id, ct_path, combined, + existing_seg if os.path.isdir(existing_seg) else None, metadata) + # Commit only after the files are published (atomic renames done). + reg["next_id"] += 1 + reg["sha256"][sha] = case_id + reg["phash"][ph] = case_id + _save_registry(root, reg) + print(f"[user_dataset] admitted {case_id} (model={model}, user={user_id})", flush=True) + except Exception: + # Never let dataset collection affect the request or crash the worker. + print(f"[user_dataset] collection error (non-fatal):\n{traceback.format_exc()}", flush=True) + + +def collect_user_scan_async(ct_path: str, output_mask_dir: str, model: str, + user_id: Optional[str] = None, ip: Optional[str] = None, + session_id: Optional[str] = None) -> None: + """Fire-and-forget entry called after a result is delivered. Returns + immediately; the gatekeeper + promotion run on a daemon thread. No-op unless + USER_DATASET_PATH is set.""" + if not _root(): + return + t = threading.Thread( + target=_admit_and_store, + args=(ct_path, output_mask_dir, model, user_id, ip, session_id), + name=f"user-dataset-{session_id or 'x'}", daemon=True, + ) + t.start() diff --git a/flask-server/tests/unit/test_user_dataset.py b/flask-server/tests/unit/test_user_dataset.py new file mode 100644 index 0000000..b102197 --- /dev/null +++ b/flask-server/tests/unit/test_user_dataset.py @@ -0,0 +1,221 @@ +"""Unit tests for the user-dataset admission gatekeeper (services/user_dataset.py). + +Pure-logic gates (quota, dedup, registry, case-id, rejection audit) run always. +The CT/mask-loading gates need nibabel and are guarded with importorskip. +""" +import importlib +import json +import os + +import pytest + + +@pytest.fixture() +def ud(tmp_path, monkeypatch): + """Fresh module bound to a temp dataset root.""" + monkeypatch.setenv("USER_DATASET_PATH", str(tmp_path / "UserData")) + # small quotas so the tests can trip them + monkeypatch.setenv("USER_DATASET_DAILY_PER_USER", "2") + monkeypatch.setenv("USER_DATASET_DAILY_PER_IP", "3") + monkeypatch.setenv("USER_DATASET_DAILY_GLOBAL", "5") + monkeypatch.setenv("USER_DATASET_MIN_ORGAN_VOXELS", "100") + monkeypatch.setenv("USER_DATASET_MIN_DISTINCT_ORGANS", "2") + import services.user_dataset as m + importlib.reload(m) + return m + + +# --------------------------- pure logic --------------------------- +def test_quota_per_user_ip_global(ud): + now = 1000.0 + reg = {"events": []} + # per-user cap = 2 + reg["events"] = [{"ts": now, "user_id": "u1", "ip": "a"} for _ in range(2)] + ok, reason = ud.check_quota(reg, "u1", "a", now) + assert not ok and reason == "user_quota" + # a different user is fine (but ip cap = 3 not yet hit) + ok, _ = ud.check_quota(reg, "u2", "b", now) + assert ok + # per-ip cap = 3 + reg["events"] = [{"ts": now, "user_id": f"u{i}", "ip": "x"} for i in range(3)] + ok, reason = ud.check_quota(reg, "u9", "x", now) + assert not ok and reason == "ip_quota" + # global cap = 5 + reg["events"] = [{"ts": now, "user_id": f"u{i}", "ip": f"ip{i}"} for i in range(5)] + ok, reason = ud.check_quota(reg, "new", "new", now) + assert not ok and reason == "global_quota" + + +def test_quota_prunes_old_events(ud): + now = 1_000_000.0 + reg = {"events": [{"ts": now - 48 * 3600, "user_id": "u1", "ip": "a"} for _ in range(9)]} + ud._prune_events(reg, now) + assert reg["events"] == [] + ok, _ = ud.check_quota(reg, "u1", "a", now) + assert ok # stale events don't count + + +def test_registry_roundtrip_atomic(ud, tmp_path): + root = ud._root() + os.makedirs(root, exist_ok=True) + reg = ud._load_registry(root) + assert reg["next_id"] == 1 and reg["sha256"] == {} + reg["next_id"] = 7 + ud._save_registry(root, reg) + assert ud._load_registry(root)["next_id"] == 7 + + +def test_evaluate_dedup_and_admit_flow(ud, monkeypatch): + """Drive the full admit flow with the NIfTI-dependent gates stubbed, so the + quota/dedup/promotion/registry bookkeeping is exercised without nibabel.""" + monkeypatch.setattr(ud, "validate_ct", lambda p: (True, "ok")) + monkeypatch.setattr(ud, "segmentation_quality_ok", + lambda p: (True, "ok", {"organ_voxels": 999, "distinct_organs": 5})) + monkeypatch.setattr(ud, "fingerprints", lambda p: ("SHA", "PH")) + monkeypatch.setattr(ud, "_promote", lambda *a, **k: None) # skip file copies + + root = ud._root(); os.makedirs(root, exist_ok=True) + now = 2000.0 + reg = ud._load_registry(root) + accepted, reason, extra = ud.evaluate("ct", "mask", reg, "u1", "1.2.3.4", now) + assert accepted and extra["sha256"] == "SHA" + # simulate the commit + reg["sha256"]["SHA"] = "USER_00000001"; reg["phash"]["PH"] = "USER_00000001" + # a re-upload of the same content is now an exact duplicate + accepted, reason, _ = ud.evaluate("ct", "mask", reg, "u1", "1.2.3.4", now) + assert not accepted and reason == "duplicate_exact" + + +# --------------------------- NIfTI-dependent --------------------------- +def _write_nifti(path, arr, affine=None): + import numpy as np + import nibabel as nib + if affine is None: + affine = np.eye(4) + nib.save(nib.Nifti1Image(arr, affine), path) + + +def test_validate_ct_accepts_real_ct_and_rejects_garbage(ud, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + # a plausible CT: air background (-1000) with a tissue/bone blob + vol = np.full((64, 64, 64), -1000.0, dtype="float32") + vol[20:40, 20:40, 20:40] = 60.0 + vol[30:34, 30:34, 30:34] = 400.0 + ct = str(tmp_path / "ct.nii.gz"); _write_nifti(ct, vol) + ok, reason = ud.validate_ct(ct) + assert ok, reason + + # constant volume -> rejected (a constant volume compresses tiny, so it trips + # too_small before the constant/intensity checks -- all are valid rejections). + flat = str(tmp_path / "flat.nii.gz"); _write_nifti(flat, np.zeros((64, 64, 64), "float32")) + ok, reason = ud.validate_ct(flat) + assert not ok and reason in ("constant_volume", "too_small", "non_ct_intensity:[0,0]") + + # non-CT intensities (all positive, no air) -> rejected + pos = np.full((64, 64, 64), 50.0, dtype="float32"); pos[0, 0, 0] = 90.0 + posf = str(tmp_path / "pos.nii.gz"); _write_nifti(posf, pos) + ok, reason = ud.validate_ct(posf) + assert not ok + + # 2D image -> rejected as not_3d (random data so it clears the size floor and + # reaches the dimensionality check rather than tripping too_small first) + twod_arr = np.random.RandomState(1).uniform(-1000, 500, (256, 256, 1)).astype("float32") + twod = str(tmp_path / "2d.nii.gz"); _write_nifti(twod, twod_arr) + ok, reason = ud.validate_ct(twod) + assert not ok and reason.startswith("not_3d") + + +def test_segmentation_quality_gate(ud, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + # a mask with several organs, plenty of voxels + mask = np.zeros((64, 64, 64), "uint8") + mask[10:30, 10:30, 10:30] = 14 # liver + mask[35:45, 35:45, 35:45] = 17 # pancreas + good = str(tmp_path / "m.nii.gz"); _write_nifti(good, mask) + ok, reason, stats = ud.segmentation_quality_ok(good) + assert ok and stats["distinct_organs"] == 2 + + empty = str(tmp_path / "e.nii.gz"); _write_nifti(empty, np.zeros((64, 64, 64), "uint8")) + ok, reason, _ = ud.segmentation_quality_ok(empty) + assert not ok + + +def test_fingerprints_deterministic(ud, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + vol = np.random.RandomState(0).uniform(-1000, 500, (48, 48, 48)).astype("float32") + ct = str(tmp_path / "ct.nii.gz"); _write_nifti(ct, vol) + a = ud.fingerprints(ct) + b = ud.fingerprints(ct) + assert a == b and len(a[0]) == 64 + + +def test_admit_and_store_end_to_end(ud, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + vol = np.full((64, 64, 64), -1000.0, dtype="float32") + vol[20:40, 20:40, 20:40] = 60.0 + vol[30:34, 30:34, 30:34] = 400.0 + ct = str(tmp_path / "ct.nii.gz"); _write_nifti(ct, vol) + + out = tmp_path / "out"; out.mkdir() + mask = np.zeros((64, 64, 64), "uint8") + mask[10:30, 10:30, 10:30] = 14 + mask[35:45, 35:45, 35:45] = 17 + _write_nifti(str(out / "combined_labels.nii.gz"), mask) + + ud._admit_and_store(ct, str(out), "ePAI", "u1", "1.2.3.4", "sess1") + + root = ud._root() + assert os.path.exists(os.path.join(root, "image_only", "USER_00000001", "ct.nii.gz")) + md = os.path.join(root, "mask_only", "USER_00000001", "metadata.json") + assert os.path.exists(md) + meta = json.load(open(md)) + assert meta["model"] == "ePAI" and "liver" in meta["organs"] + # per-organ sublabels written + assert os.path.exists(os.path.join(root, "mask_only", "USER_00000001", "segmentations", "liver.nii.gz")) + + # a second identical scan is rejected as a duplicate (logged, not stored) + ud._admit_and_store(ct, str(out), "ePAI", "u1", "1.2.3.4", "sess2") + assert not os.path.exists(os.path.join(root, "image_only", "USER_00000002")) + rej = os.path.join(root, "rejections.jsonl") + assert os.path.exists(rej) + reasons = [json.loads(l)["reason"] for l in open(rej)] + assert "duplicate_exact" in reasons + + +def test_noop_when_unset(tmp_path, monkeypatch): + """With USER_DATASET_PATH unset the feature is fully inert: no root, and the + async entry spawns nothing.""" + monkeypatch.delenv("USER_DATASET_PATH", raising=False) + import services.user_dataset as m + importlib.reload(m) + assert m._root() is None + before = __import__("threading").active_count() + m.collect_user_scan_async("ct", "out", "ePAI", "u1", "ip", "s") + assert __import__("threading").active_count() == before # no thread started + + +def test_voxel_guard_blocks_oom(ud, tmp_path, monkeypatch): + """A volume whose decoded voxel count exceeds MAX_VOXELS is rejected BEFORE + np.asarray, so a decompression bomb can't allocate/OOM the worker.""" + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + monkeypatch.setattr(ud, "MAX_VOXELS", 1000) # tiny cap + vol = np.full((64, 64, 64), -1000.0, dtype="float32") + vol[20:40, 20:40, 20:40] = 60.0 + ct = str(tmp_path / "big.nii.gz"); _write_nifti(ct, vol) + ok, reason = ud.validate_ct(ct) + assert not ok and reason.startswith("too_many_voxels") + + +def test_rejects_non_finite(ud, tmp_path): + np = pytest.importorskip("numpy") + pytest.importorskip("nibabel") + vol = np.random.RandomState(2).uniform(-1000, 500, (64, 64, 64)).astype("float32") + vol[0, 0, 0] = np.inf + ct = str(tmp_path / "inf.nii.gz"); _write_nifti(ct, vol) + ok, reason = ud.validate_ct(ct) + assert not ok and reason == "non_finite_values"