@ictechgy/context-guard 0.4.7 → 0.4.8

package/context-guard-kit/cost_guard.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import argparse
 import base64
 import binascii
+import errno
 try:
     import fcntl
 except ImportError:  # pragma: no cover - fcntl is unavailable on Windows.
@@ -49,6 +50,8 @@ DEFAULT_SAFETY_FACTOR = 1.25
 DEFAULT_LARGE_SECTION_BYTES = 64_000
 MAX_LEDGER_ROWS = 20_000
 LEDGER_TAIL_INITIAL_BYTES = 64 * 1024
+LEDGER_OPEN_RETRY_ATTEMPTS = 5
+LEDGER_OPEN_RETRY_SECONDS = 0.01
 TTL_SECONDS = {"5m": 5 * 60, "1h": 60 * 60}
 ANTHROPIC_DOCS_URL = "https://docs.anthropic.com/en/build-with-claude/prompt-caching"
 ANTHROPIC_PRICING_URL = "https://platform.claude.com/docs/en/about-claude/pricing"
@@ -58,6 +61,10 @@ ALLOWED_FIRST_COMPONENT_SYMLINKS = {
 }
 DIR_FD_OPEN_SUPPORTED = os.open in getattr(os, "supports_dir_fd", set())
 NO_FOLLOW_SUPPORTED = hasattr(os, "O_NOFOLLOW")
+DIR_FD_STAT_NOFOLLOW_SUPPORTED = (
+    os.stat in getattr(os, "supports_dir_fd", set())
+    and os.stat in getattr(os, "supports_follow_symlinks", set())
+)
 
 SECRET_RE = re.compile(
     r"(?is)("
@@ -148,25 +155,68 @@ def token_proxy_obj(data: Any) -> int:
     return token_proxy_text(json_bytes(data))
 
 
+def read_bounded_regular_path(path: str | Path, *, max_bytes: int, label: str) -> tuple[str, bool]:
+    if max_bytes < 1 or max_bytes > MAX_MAX_BYTES:
+        fail(f"max bytes must be between 1 and {MAX_MAX_BYTES}")
+    p = reject_symlink_components(Path(path), label=label)
+    leaf_name = _private_leaf_name(p, label=label)
+    parent_fd = -1
+    fd = -1
+    try:
+        parent_fd = open_directory_no_follow(p.parent, label=f"{label} parent")
+        if not DIR_FD_STAT_NOFOLLOW_SUPPORTED:
+            fail(f"{label} requires dir_fd stat support for symlink-safe regular-file validation")
+        try:
+            pre_st = os.stat(leaf_name, dir_fd=parent_fd, follow_symlinks=False)
+        except OSError as exc:
+            fail(f"could not inspect {label}: {os_error_detail(exc)}")
+        if not stat.S_ISREG(pre_st.st_mode):
+            fail(f"{label} must be a regular file")
+        flags = _base_open_flags() | _no_follow_flag(label=label)
+        if hasattr(os, "O_NONBLOCK"):
+            flags |= os.O_NONBLOCK
+        if hasattr(os, "O_NOCTTY"):
+            flags |= os.O_NOCTTY
+        fd = os.open(leaf_name, flags, dir_fd=parent_fd)
+        if not stat.S_ISREG(os.fstat(fd).st_mode):
+            fail(f"{label} must be a regular file")
+        chunks: list[bytes] = []
+        remaining = max_bytes + 1
+        while remaining > 0:
+            chunk = os.read(fd, min(64 * 1024, remaining))
+            if not chunk:
+                break
+            chunks.append(chunk)
+            remaining -= len(chunk)
+        raw = b"".join(chunks)
+    except CostGuardError:
+        raise
+    except OSError as exc:
+        fail(f"could not read {label}: {os_error_detail(exc)}")
+    finally:
+        if fd >= 0:
+            try:
+                os.close(fd)
+            except OSError:
+                pass
+        if parent_fd >= 0:
+            try:
+                os.close(parent_fd)
+            except OSError:
+                pass
+    truncated = len(raw) > max_bytes
+    if truncated:
+        raw = raw[:max_bytes]
+    return raw.decode("utf-8", errors="replace"), truncated
+
+
 def read_text_path(path: str, *, max_bytes: int = DEFAULT_MAX_BYTES) -> tuple[str, bool]:
     if max_bytes < 1 or max_bytes > MAX_MAX_BYTES:
         fail(f"max bytes must be between 1 and {MAX_MAX_BYTES}")
     if path == "-":
         raw = sys.stdin.buffer.read(max_bytes + 1)
     else:
-        p = Path(path)
-        try:
-            st = p.stat()
-        except OSError as exc:
-            fail(f"could not read input file: {exc}")
-        if not stat.S_ISREG(st.st_mode):
-            fail("input path must be a regular file")
-        if st.st_size > max_bytes + 1:
-            # Read only the bounded prefix so large requests cannot exhaust memory.
-            with p.open("rb") as fh:
-                raw = fh.read(max_bytes + 1)
-        else:
-            raw = p.read_bytes()
+        return read_bounded_regular_path(path, max_bytes=max_bytes, label="input file")
     truncated = len(raw) > max_bytes
     if truncated:
         raw = raw[:max_bytes]
@@ -494,20 +544,20 @@ def _base_open_flags() -> int:
     return flags
 
 
-def _no_follow_flag() -> int:
+def _no_follow_flag(*, label: str = "private local cost storage") -> int:
     if not NO_FOLLOW_SUPPORTED:
-        fail("private local cost storage requires O_NOFOLLOW support")
+        fail(f"{label} requires O_NOFOLLOW support")
     return os.O_NOFOLLOW
 
 
-def _directory_open_flags(*, follow_final: bool = False) -> int:
+def _directory_open_flags(*, follow_final: bool = False, label: str = "private local cost storage") -> int:
     flags = os.O_RDONLY
     if hasattr(os, "O_CLOEXEC"):
         flags |= os.O_CLOEXEC
     if hasattr(os, "O_DIRECTORY"):
         flags |= os.O_DIRECTORY
     if not follow_final:
-        flags |= _no_follow_flag()
+        flags |= _no_follow_flag(label=label)
     return flags
 
 
@@ -572,18 +622,18 @@ def reject_symlink_components(path: Path, *, label: str) -> Path:
     return path
 
 
-def open_private_directory(path: Path, *, label: str) -> int:
+def open_directory_no_follow(path: Path, *, label: str) -> int:
     """Open an existing directory without following symlink path components."""
 
     if not dir_fd_open_supported():
-        fail(f"{label} requires dir_fd support for symlink-safe private storage")
+        fail(f"{label} requires dir_fd support for symlink-safe directory traversal")
     path = reject_symlink_components(path, label=label)
-    flags = _directory_open_flags()
+    flags = _directory_open_flags(label=label)
     if path.is_absolute():
         anchor = path.anchor or os.sep
         parts = path.parts[1:]
         try:
-            current_fd = os.open(anchor, _directory_open_flags(follow_final=True))
+            current_fd = os.open(anchor, _directory_open_flags(follow_final=True, label=label))
         except OSError as exc:
             fail(f"could not inspect {label}: {os_error_detail(exc)}")
     else:
@@ -635,6 +685,12 @@ def open_private_directory(path: Path, *, label: str) -> int:
                 pass
 
 
+def open_private_directory(path: Path, *, label: str) -> int:
+    """Open an existing private-storage directory without following symlinks."""
+
+    return open_directory_no_follow(path, label=label)
+
+
 def fsync_directory_fd(fd: int) -> None:
     if os.name != "posix":
         return
@@ -676,7 +732,7 @@ def open_private_regular_fd_for_read(path: Path, *, label: str) -> int:
     fd = -1
     try:
         parent_fd = open_private_directory(path.parent, label=f"{label} parent")
-        fd = os.open(leaf_name, _base_open_flags() | _no_follow_flag(), dir_fd=parent_fd)
+        fd = os.open(leaf_name, _base_open_flags() | _no_follow_flag(label=label), dir_fd=parent_fd)
         st = os.fstat(fd)
         if not stat.S_ISREG(st.st_mode):
             fail(f"{label} must be a regular file")
@@ -1138,41 +1194,47 @@ def open_private_regular_file_for_append(path: Path, *, label: str) -> int:
     flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND | _no_follow_flag()
     if hasattr(os, "O_CLOEXEC"):
         flags |= os.O_CLOEXEC
-    parent_fd = -1
-    fd = -1
-    try:
-        parent_fd = open_private_directory(path.parent, label=f"{label} parent")
-        fd = os.open(leaf_name, flags, 0o600, dir_fd=parent_fd)
-        st = os.fstat(fd)
-        if not stat.S_ISREG(st.st_mode):
-            fail(f"{label} must be a regular file")
-        try:
-            os.fchmod(fd, 0o600)
-        except (AttributeError, OSError):
-            pass
-        st = os.fstat(fd)
-        if os.name == "posix" and stat.S_IMODE(st.st_mode) != 0o600:
-            fail(f"could not verify {label} privacy: expected mode 0600")
-        owned_fd = fd
+    for attempt in range(LEDGER_OPEN_RETRY_ATTEMPTS):
+        parent_fd = -1
         fd = -1
-        return owned_fd
-    except CostGuardError:
-        raise
-    except OSError as exc:
-        fail(f"could not open {label}: {os_error_detail(exc)}")
-    finally:
-        if fd >= 0:
-            # Ownership transfers to the caller only on the successful return
-            # above. On errors, close before surfacing a deterministic message.
-            try:
-                os.close(fd)
-            except OSError:
-                pass
-        if parent_fd >= 0:
+        try:
+            parent_fd = open_private_directory(path.parent, label=f"{label} parent")
+            fd = os.open(leaf_name, flags, 0o600, dir_fd=parent_fd)
+            st = os.fstat(fd)
+            if not stat.S_ISREG(st.st_mode):
+                fail(f"{label} must be a regular file")
             try:
-                os.close(parent_fd)
-            except OSError:
+                os.fchmod(fd, 0o600)
+            except (AttributeError, OSError):
                 pass
+            st = os.fstat(fd)
+            if os.name == "posix" and stat.S_IMODE(st.st_mode) != 0o600:
+                fail(f"could not verify {label} privacy: expected mode 0600")
+            owned_fd = fd
+            fd = -1
+            return owned_fd
+        except CostGuardError:
+            raise
+        except OSError as exc:
+            if exc.errno == errno.ENOENT and attempt + 1 < LEDGER_OPEN_RETRY_ATTEMPTS:
+                time.sleep(LEDGER_OPEN_RETRY_SECONDS)
+                continue
+            fail(f"could not open {label}: {os_error_detail(exc)}")
+        finally:
+            if fd >= 0:
+                # Ownership transfers to the caller only on the successful
+                # return above. On errors, close before surfacing a
+                # deterministic message.
+                try:
+                    os.close(fd)
+                except OSError:
+                    pass
+            if parent_fd >= 0:
+                try:
+                    os.close(parent_fd)
+                except OSError:
+                    pass
+    raise AssertionError("unreachable: append retry loop exits via return or fail")
 
 
 def load_ledger(store_dir: Path) -> list[dict[str, Any]]:
@@ -1280,7 +1342,7 @@ def default_pricing_profile() -> dict[str, Any]:
     }
 
 
-def load_pricing_profile(raw: str | None) -> dict[str, Any]:
+def load_pricing_profile(raw: str | None, *, max_bytes: int = DEFAULT_MAX_BYTES) -> dict[str, Any]:
     profile = default_pricing_profile()
     if not raw:
         return profile
@@ -1288,7 +1350,12 @@ def load_pricing_profile(raw: str | None) -> dict[str, Any]:
         if raw.lstrip().startswith("{"):
             override = json.loads(raw, parse_constant=reject_json_constant)
         else:
-            override = json.loads(Path(raw).read_text(encoding="utf-8"), parse_constant=reject_json_constant)
+            text, truncated = read_bounded_regular_path(raw, max_bytes=max_bytes, label="pricing profile")
+            if truncated:
+                fail("pricing profile exceeded max bytes")
+            override = json.loads(text, parse_constant=reject_json_constant)
+    except CostGuardError:
+        raise
     except (OSError, json.JSONDecodeError, ValueError) as exc:
         fail(f"could not load pricing profile: {exc}")
     if not isinstance(override, dict):
@@ -1542,7 +1609,7 @@ def annotate_cache_state(
 def preflight_command(args: argparse.Namespace) -> int:
     request_raw, _truncated = load_json_input(args.request, max_bytes=args.max_bytes)
     request = require_json_object(request_raw, "request")
-    profile = load_pricing_profile(args.pricing_profile)
+    profile = load_pricing_profile(args.pricing_profile, max_bytes=args.max_bytes)
     if args.usd_to_krw is not None:
         profile["usd_to_krw"] = usd_to_krw(profile, args.usd_to_krw)
     if args.budget_usd is not None:
@@ -1809,7 +1876,7 @@ def observe_command(args: argparse.Namespace) -> int:
         usage = usage_raw
     if not isinstance(usage, dict):
         fail("usage must be a JSON object or an object containing a usage object")
-    profile = load_pricing_profile(args.pricing_profile)
+    profile = load_pricing_profile(args.pricing_profile, max_bytes=args.max_bytes)
     if args.usd_to_krw is not None:
         profile["usd_to_krw"] = usd_to_krw(profile, args.usd_to_krw)
     model = str(args.model or (usage_raw.get("model") if isinstance(usage_raw, dict) else "") or "unknown")
@@ -2217,7 +2284,7 @@ def emit(data: dict[str, Any], *, json_mode: bool) -> None:
 def add_common_cost_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--pricing-profile", help="JSON string or file with input/output rates, cache multipliers, and usd_to_krw")
     parser.add_argument("--usd-to-krw", type=float, help="override USD→KRW exchange rate used for estimates")
-    parser.add_argument("--max-bytes", type=int, default=DEFAULT_MAX_BYTES, help=f"maximum JSON input bytes (default: {DEFAULT_MAX_BYTES})")
+    parser.add_argument("--max-bytes", type=int, default=DEFAULT_MAX_BYTES, help=f"maximum JSON input and pricing profile file bytes (default: {DEFAULT_MAX_BYTES})")
     parser.add_argument("--json", action="store_true", help="emit machine-readable JSON")