Improve snapshot reliability and usability

snapshot_storage.py

@@ -11,6 +11,7 @@ operations.  Only get_full_record() reads a file from disk after warm-up.
 """
 
 import json
+import math
 import os
 import shutil
 import tempfile
@@ -37,6 +38,13 @@ _DATA_DIR = os.path.join(_USER_SM_DIR, "snapshots")
 _cache = {}
 _cache_warmed = set()  # workflow keys already loaded from disk
 
+_ALLOWED_SOURCES = {"auto", "manual", "initial", "node", "restore_guard"}
+_ALLOWED_META_FIELDS = {"label", "notes", "locked", "parentId"}
+_MAX_WORKFLOW_KEY_LENGTH = 4096
+_MAX_ID_LENGTH = 255
+_MAX_LABEL_LENGTH = 500
+_MAX_NOTES_LENGTH = 50000
+
 
 def _extract_meta(record):
     """Return a lightweight copy of *record* without graphData or thumbnail.
@@ -76,6 +84,8 @@ def _ensure_cached(workflow_key):
 def _workflow_dir(workflow_key):
     if not workflow_key or not isinstance(workflow_key, str):
         raise ValueError(f"Invalid workflow key: {workflow_key!r}")
+    if len(workflow_key) > _MAX_WORKFLOW_KEY_LENGTH:
+        raise ValueError("Workflow key is too long")
     encoded = urllib.parse.quote(workflow_key, safe="")
     path = os.path.normpath(os.path.join(_DATA_DIR, encoded))
     # Defense in depth: urllib.parse.quote() leaves "." and ".." unescaped, so a
@@ -88,10 +98,79 @@ def _workflow_dir(workflow_key):
 
 
 def _validate_id(snapshot_id):
-    if not snapshot_id or "/" in snapshot_id or "\\" in snapshot_id or ".." in snapshot_id:
+    if (
+        not isinstance(snapshot_id, str)
+        or not snapshot_id
+        or len(snapshot_id) > _MAX_ID_LENGTH
+        or "/" in snapshot_id
+        or "\\" in snapshot_id
+        or ".." in snapshot_id
+    ):
         raise ValueError(f"Invalid snapshot id: {snapshot_id!r}")
 
 
+def _validate_meta_fields(fields):
+    if not isinstance(fields, dict):
+        raise ValueError("Metadata fields must be an object")
+    unknown = set(fields) - _ALLOWED_META_FIELDS
+    if unknown:
+        raise ValueError(f"Unsupported metadata fields: {', '.join(sorted(unknown))}")
+    if "label" in fields:
+        label = fields["label"]
+        if not isinstance(label, str) or not label.strip() or len(label) > _MAX_LABEL_LENGTH:
+            raise ValueError("Invalid snapshot label")
+    if "notes" in fields:
+        notes = fields["notes"]
+        if notes is not None and (not isinstance(notes, str) or len(notes) > _MAX_NOTES_LENGTH):
+            raise ValueError("Invalid snapshot notes")
+    if "locked" in fields and not isinstance(fields["locked"], bool):
+        raise ValueError("Invalid locked value")
+    if "parentId" in fields and fields["parentId"] is not None:
+        _validate_id(fields["parentId"])
+
+
+def validate_record(record):
+    """Validate the persisted snapshot envelope and graph container."""
+    if not isinstance(record, dict):
+        raise ValueError("Snapshot record must be an object")
+
+    snapshot_id = record.get("id")
+    workflow_key = record.get("workflowKey")
+    _validate_id(snapshot_id)
+    _workflow_dir(workflow_key)
+
+    timestamp = record.get("timestamp")
+    if (
+        not isinstance(timestamp, (int, float))
+        or isinstance(timestamp, bool)
+        or not math.isfinite(timestamp)
+        or timestamp < 0
+    ):
+        raise ValueError("Invalid snapshot timestamp")
+
+    label = record.get("label")
+    if not isinstance(label, str) or not label.strip() or len(label) > _MAX_LABEL_LENGTH:
+        raise ValueError("Invalid snapshot label")
+
+    graph_data = record.get("graphData")
+    if not isinstance(graph_data, dict) or not isinstance(graph_data.get("nodes"), list):
+        raise ValueError("Invalid snapshot graphData")
+
+    source = record.get("source")
+    if source is not None and source not in _ALLOWED_SOURCES:
+        raise ValueError(f"Invalid snapshot source: {source!r}")
+    if "locked" in record and not isinstance(record["locked"], bool):
+        raise ValueError("Invalid locked value")
+    if record.get("parentId") is not None:
+        _validate_id(record["parentId"])
+    if "notes" in record and (
+        not isinstance(record["notes"], str)
+        or len(record["notes"]) > _MAX_NOTES_LENGTH
+    ):
+        raise ValueError("Invalid snapshot notes")
+    return record
+
+
 def _atomic_write_json(path, obj):
     """Write *obj* as JSON to *path* atomically (temp file + os.replace).
 
@@ -116,6 +195,7 @@ def _atomic_write_json(path, obj):
 
 def put(record):
     """Write one snapshot record to disk and update the cache."""
+    validate_record(record)
     snapshot_id = record["id"]
     workflow_key = record["workflowKey"]
     _validate_id(snapshot_id)
@@ -158,6 +238,7 @@ def update_meta(workflow_key, snapshot_id, fields):
     Returns True on success, False if the file does not exist.
     """
     _validate_id(snapshot_id)
+    _validate_meta_fields(fields)
     path = os.path.join(_workflow_dir(workflow_key), f"{snapshot_id}.json")
     if not os.path.isfile(path):
         return False
@@ -404,8 +485,32 @@ def _invalidate_profile_cache():
 
 def profile_put(profile):
     """Create or update a profile. profile must have 'id'."""
+    if not isinstance(profile, dict):
+        raise ValueError("Profile must be an object")
     pid = profile["id"]
     _validate_id(pid)
+    name = profile.get("name")
+    if not isinstance(name, str) or not name.strip() or len(name) > _MAX_LABEL_LENGTH:
+        raise ValueError("Invalid profile name")
+    timestamp = profile.get("timestamp")
+    if (
+        not isinstance(timestamp, (int, float))
+        or isinstance(timestamp, bool)
+        or not math.isfinite(timestamp)
+        or timestamp < 0
+    ):
+        raise ValueError("Invalid profile timestamp")
+    workflows = profile.get("workflows")
+    if not isinstance(workflows, list) or len(workflows) > 500:
+        raise ValueError("Invalid profile workflows")
+    for workflow in workflows:
+        if not isinstance(workflow, dict):
+            raise ValueError("Invalid profile workflow")
+        _workflow_dir(workflow.get("workflowKey"))
+        if workflow.get("snapshotId") is not None:
+            _validate_id(workflow["snapshotId"])
+    if profile.get("activeWorkflowKey") is not None:
+        _workflow_dir(profile["activeWorkflowKey"])
     _ensure_profiles_dir()
     path = os.path.join(_PROFILES_DIR, f"{pid}.json")
     _atomic_write_json(path, profile)