Update rft_flightrecorder.py

rft_flightrecorder.py

@@ -19,7 +19,15 @@ EVENT_SPEC = "rft-flight-event-v0"
 ROOT_SPEC = "rft-flight-session-root-v0"
 DEFAULT_LOG_PATH = "flightlog.jsonl"
 
-LOCK_TIMEOUT_S = 10  # keep small; we only lock for quick file IO
+# ---------- Locking ----------
+LOCK_TIMEOUT_S = 10  # short on purpose; UI should not hang
+
+# ---------- Public abuse limits (Import Bundle) ----------
+MAX_BUNDLE_BYTES = 15 * 1024 * 1024          # 15 MB zip upload cap
+MAX_ZIP_MEMBER_BYTES = 25 * 1024 * 1024      # 25 MB decompressed member cap
+MAX_EVENTS_PER_BUNDLE = 50_000
+MAX_EVENT_LINE_BYTES = 200_000               # 200 KB per JSONL line (prevents huge single lines)
+MAX_ZIP_RATIO = 200.0                        # zip-bomb heuristic (file_size / compress_size)
 
 
 # ============================================================
@@ -54,7 +62,6 @@ def short(h: str, n: int = 12) -> str:
 # ============================================================
 
 def _lock_for(path: str) -> FileLock:
-    # one lock file per jsonl
     return FileLock(path + ".lock", timeout=LOCK_TIMEOUT_S)
 
 
@@ -175,14 +182,17 @@ def events_for_session(all_events: List[Dict[str, Any]], session_id: str) -> Lis
         e for e in all_events
         if isinstance(e, dict)
         and e.get("session_id") == sid
-        and "event_hash_sha256" in e
         and e.get("spec") == EVENT_SPEC
+        and "event_hash_sha256" in e
     ]
 
 
 def last_event_for_session(all_events: List[Dict[str, Any]], session_id: str) -> Optional[Dict[str, Any]]:
     evs = events_for_session(all_events, session_id)
-    return evs[-1] if evs else None
+    if not evs:
+        return None
+    evs.sort(key=lambda x: int(x.get("seq", 0)))
+    return evs[-1]
 
 
 def next_seq(all_events: List[Dict[str, Any]], session_id: str) -> int:
@@ -309,17 +319,17 @@ def finalise_session(
         return None, "Missing session_id."
 
     try:
-        # Lock while reading events + appending session_end (append_event is locked too,
-        # but we want a consistent anchor over a stable snapshot of the file)
+        # Lock the entire finalise operation so nothing can append between:
+        # anchor computation -> session_end append.
         with _lock_for(log_path):
             all_events, _corrupt = read_jsonl(log_path)
             evs = events_for_session(all_events, sid)
             if not evs:
                 return None, "No events found for this session."
 
-            # Ensure correct order
             evs.sort(key=lambda x: int(x.get("seq", 0)))
 
+            # Anchor must NOT include the session_end event itself (avoids circularity).
             first_hash = evs[0]["event_hash_sha256"]
             last_hash = evs[-1]["event_hash_sha256"]
             count = len(evs)
@@ -349,15 +359,13 @@ def finalise_session(
                 except Exception:
                     return None, "Failed to sign anchor (invalid private key?)."
 
-            # Append session_end inside the same lock to prevent new events slipping in
+            # Append session_end event (manually, inside the same lock)
             payload_text = json.dumps({"anchor": anchor}, ensure_ascii=False)
-
-            # Build the session_end event manually so it stays inside our lock
-            # (we avoid calling append_event here to prevent nested lock complexity)
             meta = {
                 "model_id": (model_id or "unknown").strip(),
                 "run_mode": (run_mode or "unknown").strip(),
             }
+
             seq = int(evs[-1].get("seq", 0)) + 1
             prev_hash = evs[-1]["event_hash_sha256"]
             parent = prev_hash
@@ -376,8 +384,6 @@ def finalise_session(
             event["event_hash_sha256"] = event_hash
 
             if sign_anchor:
-                if not sk_hex or not sk_hex.strip():
-                    return None, "Anchor signing enabled but private key is missing."
                 try:
                     sk = load_sk(sk_hex)
                     sig = sk.sign(bytes.fromhex(event_hash))
@@ -387,7 +393,7 @@ def finalise_session(
 
             append_jsonl(log_path, event)
 
-        return anchor, f"OK. Session finalised. Root hash: {root_hash} (last event hash: {event_hash})"
+        return anchor, f"OK. Session finalised. Root hash: {root_hash} (session_end hash: {event_hash})"
 
     except Timeout:
         return None, "Busy: log is locked (try again)."
@@ -399,6 +405,26 @@ def finalise_session(
 # Verification
 # ============================================================
 
+def _anchor_expected_from_events(evs_sorted: List[Dict[str, Any]], session_end_index: int) -> str:
+    """
+    Compute the expected anchor root_hash for the chain *before* the session_end event at session_end_index.
+    This avoids circularity (session_end payload contains anchor).
+    """
+    chain = evs_sorted[:session_end_index]  # exclude the session_end itself
+    first_hash = chain[0]["event_hash_sha256"]
+    last_hash = chain[-1]["event_hash_sha256"]
+    count = len(chain)
+
+    root_core = {
+        "spec": ROOT_SPEC,
+        "session_id": chain[0]["session_id"],
+        "first_event_hash_sha256": first_hash,
+        "last_event_hash_sha256": last_hash,
+        "event_count": count,
+    }
+    return sha256_hex(canon(root_core))
+
+
 def verify_session_from_events(
     evs: List[Dict[str, Any]],
     session_id: str,
@@ -424,6 +450,10 @@ def verify_session_from_events(
         except Exception:
             return "Invalid public key.", False, ""
 
+    # Ensure order by seq
+    evs = [e for e in evs if isinstance(e, dict)]
+    evs.sort(key=lambda x: int(x.get("seq", 0)))
+
     expected_prev = "0" * 64
     expected_seq = 1
 
@@ -439,6 +469,11 @@ def verify_session_from_events(
             report.append(f"[FAIL] Bad spec at seq {e.get('seq')}.")
             continue
 
+        if e.get("session_id") != sid:
+            ok = False
+            report.append(f"[FAIL] session_id mismatch at seq {e.get('seq')}.")
+            continue
+
         if int(e.get("seq", -1)) != expected_seq:
             ok = False
             report.append(f"[FAIL] Seq mismatch: got {e.get('seq')} expected {expected_seq}.")
@@ -480,28 +515,20 @@ def verify_session_from_events(
 
         expected_prev = e["event_hash_sha256"]
 
-    end_events = [e for e in evs if e.get("event_type") == "session_end"]
-    if end_events:
-        se = end_events[-1]
+    # Anchor check: validate the LAST session_end event (if present)
+    end_idxs = [i for i, e in enumerate(evs) if e.get("event_type") == "session_end"]
+    if end_idxs:
+        se_idx = end_idxs[-1]
+        se = evs[se_idx]
         anchor = (se.get("payload") or {}).get("anchor")
-        if isinstance(anchor, dict):
-            first_hash = evs[0]["event_hash_sha256"]
-            last_hash = evs[-1]["event_hash_sha256"]
-            count = len(evs)
 
-            root_core = {
-                "spec": ROOT_SPEC,
-                "session_id": sid,
-                "first_event_hash_sha256": first_hash,
-                "last_event_hash_sha256": last_hash,
-                "event_count": count,
-            }
-            root_hash = sha256_hex(canon(root_core))
-            if anchor.get("root_hash_sha256") != root_hash:
+        if isinstance(anchor, dict) and se_idx > 0:
+            expected_root = _anchor_expected_from_events(evs, se_idx)
+            if anchor.get("root_hash_sha256") != expected_root:
                 ok = False
-                report.append("[FAIL] Session anchor root hash does not match current event chain.")
+                report.append("[FAIL] Session anchor root hash does not match the pre-session_end event chain.")
             else:
-                report.append("[OK] Session anchor matches current event chain.")
+                report.append("[OK] Session anchor matches the pre-session_end event chain.")
         else:
             report.append("[WARN] session_end found but anchor payload is missing/invalid.")
 
@@ -518,8 +545,8 @@ def verify_session(log_path: str, session_id: str, pk_hex: str, require_signatur
         all_events, _corrupt = _read_jsonl_locked(log_path)
     except Timeout:
         return "FAIL", False, "Busy: log is locked (try again)."
+
     evs = events_for_session(all_events, session_id)
-    evs.sort(key=lambda x: int(x.get("seq", 0)))
     return verify_session_from_events(evs, session_id, pk_hex=pk_hex, require_signatures=require_signatures)
 
 
@@ -650,7 +677,6 @@ def export_session_bundle(log_path: str, session_id: str) -> Tuple[Optional[str]
         return None, "No events found."
 
     evs.sort(key=lambda x: int(x.get("seq", 0)))
-
     status, ok, report = verify_session_from_events(evs, sid, pk_hex="", require_signatures=False)
 
     zip_name = f"rft_flight_bundle_{sid}.zip"
@@ -681,19 +707,48 @@ def export_session_bundle(log_path: str, session_id: str) -> Tuple[Optional[str]
 # Import bundle (third-party verification)
 # ============================================================
 
+def _zip_member_safe(z: zipfile.ZipFile, member: str) -> Tuple[bool, str]:
+    try:
+        info = z.getinfo(member)
+    except Exception:
+        return False, "Missing events member in zip."
+
+    # Decompressed size cap
+    if info.file_size > MAX_ZIP_MEMBER_BYTES:
+        return False, "Events file too large (decompressed)."
+
+    # Zip-bomb heuristic: insane compression ratio
+    if info.compress_size and info.compress_size > 0:
+        ratio = float(info.file_size) / float(info.compress_size)
+        if ratio > MAX_ZIP_RATIO:
+            return False, "Suspicious zip compression ratio (possible zip bomb)."
+
+    return True, "OK"
+
+
 def _read_jsonl_from_zip(z: zipfile.ZipFile, member: str) -> List[Dict[str, Any]]:
     out: List[Dict[str, Any]] = []
     raw_text = z.read(member).decode("utf-8", errors="replace")
+
     for raw in raw_text.splitlines():
         raw = raw.strip()
         if not raw:
             continue
+
+        # single-line cap
+        if len(raw.encode("utf-8", errors="ignore")) > MAX_EVENT_LINE_BYTES:
+            continue
+
         try:
             obj = json.loads(raw)
             if isinstance(obj, dict):
                 out.append(obj)
         except Exception:
             continue
+
+        if len(out) > MAX_EVENTS_PER_BUNDLE:
+            break
+
     return out
 
 
@@ -707,9 +762,17 @@ def import_bundle_verify(
     if not bundle_path or not os.path.exists(bundle_path):
         return "Missing bundle file.", False, "", None
 
+    # zip upload cap
+    try:
+        if os.path.getsize(bundle_path) > MAX_BUNDLE_BYTES:
+            return "FAIL", False, "Bundle too large.", None
+    except Exception:
+        return "FAIL", False, "Unable to read bundle size.", None
+
     try:
         with zipfile.ZipFile(bundle_path, "r") as z:
             members = z.namelist()
+
             events_member = None
             for m in members:
                 if m.endswith("_events.jsonl"):
@@ -723,11 +786,22 @@ def import_bundle_verify(
             if not events_member:
                 return "No .jsonl events file found in bundle.", False, "", None
 
+            safe, why = _zip_member_safe(z, events_member)
+            if not safe:
+                return "FAIL", False, why, None
+
             evs = _read_jsonl_from_zip(z, events_member)
 
     except Exception:
         return "Failed to read bundle (invalid zip?).", False, "", None
 
+    if not evs:
+        return "Bundle contains no usable events.", False, "", None
+
+    if len(evs) > MAX_EVENTS_PER_BUNDLE:
+        return "FAIL", False, "Too many events in bundle.", None
+
+    # Determine session_id
     sid = ""
     for e in evs:
         if e.get("spec") == EVENT_SPEC and e.get("session_id"):
@@ -744,11 +818,18 @@ def import_bundle_verify(
     )
 
     if store_into_log and ok:
-        # Store under a single lock to avoid interleaving with live writers.
+        # Store under a lock + dedupe by event_hash_sha256 to avoid log pollution
         try:
             with _lock_for(log_path):
+                existing, _ = read_jsonl(log_path)
+                seen = {e.get("event_hash_sha256") for e in existing if isinstance(e, dict)}
+
                 for e in evs:
-                    append_jsonl(log_path, e)
+                    h = e.get("event_hash_sha256")
+                    if h and h not in seen:
+                        append_jsonl(log_path, e)
+                        seen.add(h)
+
         except Timeout:
             return "FAIL", False, "Busy: log is locked (try again).", None