From 8fb7d188b3dcc5f43116022bb923db1cea870cf6 Mon Sep 17 00:00:00 2001
From: Adam Moussa <adam@seahavenind.com>
Date: Tue, 23 Jun 2026 12:29:25 -0400
Subject: [PATCH] fix(ws5): allowlist save_memory names + symlink-safe write

Security follow-up from the per-PR review (non-blocking, defense-in-depth):
- Replace the save_memory name blocklist with an allowlist regex
  (^[A-Za-z0-9][A-Za-z0-9._-]*$, max 128) so dot-only/hidden/backslash/NUL/
  over-long names are rejected outright, not written as malformed-but-contained
  files.
- Write via os.open(..., O_NOFOLLOW): the open fails (ELOOP) if the final
  path component is a pre-planted symlink, closing the TOCTOU where a symlink
  in _box-drafts/ could redirect the write outside the dir. O_CREAT|O_TRUNC
  keeps overwrite-on-resave for regular files.

Tests: adds allowlist-rejection + symlink-refusal cases (21 pass).
---
 agent-team/tests/test_ws5_memory_handbook.py | 30 ++++++++++++++++++
 retriever.py                                 | 32 +++++++++++++++++---
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/agent-team/tests/test_ws5_memory_handbook.py b/agent-team/tests/test_ws5_memory_handbook.py
index 52d5cc1..9c10c9d 100644
--- a/agent-team/tests/test_ws5_memory_handbook.py
+++ b/agent-team/tests/test_ws5_memory_handbook.py
@@ -57,6 +57,36 @@ def test_save_memory_rejects_path_traversal(tmp_path: Path) -> None:
         save_memory("", "bad", memory_dir=tmp_path)
 
 
+def test_save_memory_rejects_unsafe_allowlist_names(tmp_path: Path) -> None:
+    from retriever import save_memory
+
+    # Allowlist rejects dot-only / hidden / backslash / NUL / over-long names
+    # that the old blocklist let through as malformed-but-contained files.
+    for bad in (".", "..", ".hidden", "a\\b", "a\x00b", "x" * 200, "-leading"):
+        with pytest.raises(ValueError):
+            save_memory(bad, "bad", memory_dir=tmp_path)
+
+
+def test_save_memory_does_not_follow_symlink_out_of_drafts(tmp_path: Path) -> None:
+    import os
+
+    from retriever import save_memory
+
+    drafts = tmp_path / "_box-drafts"
+    drafts.mkdir()
+    outside = tmp_path / "outside.txt"
+    outside.write_text("original")
+    # Pre-plant a symlink in the drafts dir pointing outside it.
+    (drafts / "evil.md").symlink_to(outside)
+
+    # O_NOFOLLOW must refuse to write through the symlink (ELOOP).
+    with pytest.raises(OSError):
+        save_memory("evil", "overwrite attempt", memory_dir=tmp_path)
+    # The outside target is untouched.
+    assert outside.read_text() == "original"
+    assert os.path.islink(drafts / "evil.md")
+
+
 def test_save_memory_does_not_appear_in_live_dir(tmp_path: Path) -> None:
     from retriever import load_memories, save_memory
 
diff --git a/retriever.py b/retriever.py
index 4df0fdd..7001f9d 100644
--- a/retriever.py
+++ b/retriever.py
@@ -14,6 +14,7 @@ from __future__ import annotations
 import json
 import math
 import os
+import re
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -160,6 +161,13 @@ def retrieve(
 # so the retriever never auto-indexes a not-yet-reviewed draft.
 _DRAFTS_SUBDIR = "_box-drafts"
 
+# Allowlist for save_memory() names (defense-in-depth over the old blocklist).
+# Must start alphanumeric, then alphanumerics / dot / underscore / hyphen. This
+# rejects path separators, NUL, leading-dot hidden files, and dot-only names
+# (".", "..") outright — a name cannot escape the drafts dir or be malformed.
+_SAFE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+_MAX_NAME_LEN = 128
+
 
 def save_memory(name: str, content: str, *, memory_dir: Path | None = None) -> Path:
     """Write a memory draft to the review queue (``_box-drafts/`` subdir).
@@ -170,17 +178,31 @@ def save_memory(name: str, content: str, *, memory_dir: Path | None = None) -> P
     never indexes ``_box-drafts/`` entries, so writing here never
     auto-activates a draft.
 
-    ``memory_dir`` defaults to :data:`MEMORY_DIR`. ``name`` must be a plain
-    filename stem (no path separators or null bytes) to prevent traversal.
-    Raises ``ValueError`` on unsafe names. Creates the subdir if absent.
+    ``memory_dir`` defaults to :data:`MEMORY_DIR`. ``name`` must match the safe
+    allowlist (alphanumeric start; then ``A-Z a-z 0-9 . _ -``; max 128 chars) so
+    it cannot contain path separators, NUL, ``..`` traversal, or be a dot-only /
+    hidden name. Raises ``ValueError`` on unsafe names. Creates the subdir if
+    absent.
     """
-    if not name or any(c in name for c in ("/", "\\", "\x00", "..")):
+    if (
+        not name
+        or len(name) > _MAX_NAME_LEN
+        or ".." in name
+        or not _SAFE_NAME_RE.match(name)
+    ):
         raise ValueError(f"unsafe memory name: {name!r}")
     effective_dir = memory_dir if memory_dir is not None else MEMORY_DIR
     drafts_dir = effective_dir / _DRAFTS_SUBDIR
     drafts_dir.mkdir(parents=True, exist_ok=True)
     dest = drafts_dir / f"{name}.md"
-    dest.write_text(content)
+    # Symlink-safe write: O_NOFOLLOW makes the open fail (ELOOP) if the final
+    # path component is a pre-planted symlink, closing the TOCTOU where a symlink
+    # in _box-drafts/ could redirect the write outside the dir. O_CREAT|O_TRUNC
+    # preserves the overwrite-on-resave behavior for a regular file.
+    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW
+    fd = os.open(dest, flags, 0o600)
+    with os.fdopen(fd, "w", encoding="utf-8") as fh:
+        fh.write(content)
     return dest