cowork: fail-closed expiry verification + rotation type guards

src/apiauth/keygen.py

@@ -151,6 +151,15 @@ def rotate_key(
     entry = keystore.get(key_id)
     if entry is None:
         return None
+    if entry.get("type") != "api_key":
+        # Refuse to corrupt a differently-typed entry: running the API-key
+        # rotation path over a JWT entry used to overwrite key_hash/prefix,
+        # bump version, and leave signing_secret_hash stale -- silently
+        # producing a half-migrated entry that neither verifier trusts.
+        raise ValueError(
+            f"rotate_key() requires an 'api_key' entry, got type={entry.get('type')!r} "
+            f"(key_id={key_id!r}); use rotate_jwt() for JWT entries"
+        )
 
     new_api_key = generate_api_key()
     new_hash = hashlib.sha256(new_api_key.encode()).hexdigest()
@@ -272,6 +281,11 @@ def rotate_jwt(
     entry = keystore.get(key_id)
     if entry is None:
         return None
+    if entry.get("type") != "jwt":
+        raise ValueError(
+            f"rotate_jwt() requires a 'jwt' entry, got type={entry.get('type')!r} "
+            f"(key_id={key_id!r}); use rotate_key() for API-key entries"
+        )
 
     signing_secret = secrets.token_hex(32)
     import jwt as pyjwt

src/apiauth/keystore.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import tempfile
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from pathlib import Path
 from typing import Any
@@ -18,7 +19,19 @@
     """Load existing master key or generate a new one."""
     key_path = key_dir / _KEY_FILE
     if key_path.exists():
-        return key_path.read_bytes()
+        key = key_path.read_bytes()
+        if len(key) != 32:
+            # A truncated or otherwise malformed master.key would otherwise
+            # surface as an opaque cryptography error (or worse, get silently
+            # regenerated below, bricking every existing entry). Fail loudly
+            # instead so the real master key can be restored.
+            raise RuntimeError(
+                f"Master key at {key_path} is corrupt: expected 32 bytes, "
+                f"got {len(key)}. Refusing to regenerate it, which would make "
+                "every existing keystore entry undecryptable. Restore the "
+                "original master.key from backup."
+            )
+        return key
 
     key_dir.mkdir(parents=True, exist_ok=True)
     key = AESGCM.generate_key(bit_length=256)
@@ -65,11 +78,33 @@
             ) from exc
 
     def _save(self) -> None:
+        """Persist the store atomically.
+
+        Writing keys.json in place meant a crash mid-write could leave a torn
+        ciphertext behind -- and since _load() refuses to continue on decrypt
+        failure (by design), that tear would brick the whole keystore. Write
+        to a temp file in the same directory, fsync, then os.replace() so
+        readers only ever see a fully-written store.
+        """
         plaintext = json.dumps(self._entries, indent=2, default=str).encode("utf-8")
         nonce = os.urandom(12)
         ciphertext = self._aesgcm.encrypt(nonce, plaintext, None)
-        self._store_path.write_bytes(nonce + ciphertext)
-        os.chmod(str(self._store_path), 0o600)
+        fd, tmp_path = tempfile.mkstemp(
+            dir=str(self.key_dir), prefix=".keys-", suffix=".tmp"
+        )
+        try:
+            with os.fdopen(fd, "wb") as fh:
+                fh.write(nonce + ciphertext)
+                fh.flush()
+                os.fsync(fh.fileno())
+            os.chmod(tmp_path, 0o600)
+            os.replace(tmp_path, str(self._store_path))
+        except BaseException:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+            raise
 
     def get_all(self) -> dict[str, dict[str, Any]]:
         """Return all stored entries."""

tests/test_keystore_atomic.py

@@ -0,0 +1,43 @@
+"""Atomic keystore persistence + master key validation."""
+import os
+
+import pytest
+
+from apiauth.keystore import Keystore, _get_or_create_master_key
+
+
+def _make_keystore(tmp_path):
+    return Keystore(key_dir=tmp_path)
+
+
+def test_save_is_atomic_and_reloadable(tmp_path):
+    ks = _make_keystore(tmp_path)
+    ks.put("k1", {"type": "api_key", "name": "n", "service": "s"})
+    # No temp files left behind after a successful save.
+    leftovers = [p for p in os.listdir(tmp_path) if p.startswith(".keys-")]
+    assert leftovers == []
+    # A fresh instance reads back exactly what was written.
+    ks2 = Keystore(key_dir=tmp_path)
+    assert ks2.get("k1")["name"] == "n"
+
+
+def test_torn_store_does_not_silently_overwrite_entries(tmp_path):
+    """Simulate a torn write: garbage in keys.json must raise, never reset."""
+    ks = _make_keystore(tmp_path)
+    ks.put("k1", {"type": "api_key", "name": "n"})
+    store = tmp_path / "keys.json"
+    store.write_bytes(b"\x00" * 64)
+    with pytest.raises(RuntimeError, match="Failed to decrypt"):
+        Keystore(key_dir=tmp_path)
+    # The corrupt file was NOT replaced by an empty store.
+    assert store.stat().st_size == 64
+
+
+def test_corrupt_master_key_fails_loudly(tmp_path):
+    ks = _make_keystore(tmp_path)
+    ks.put("k1", {"type": "api_key", "name": "n"})
+    (tmp_path / "master.key").write_bytes(b"short")
+    with pytest.raises(RuntimeError, match="corrupt: expected 32 bytes"):
+        _get_or_create_master_key(tmp_path)
+    # ...and the bad key was not silently regenerated over.
+    assert (tmp_path / "master.key").read_bytes() == b"short"

tests/test_rotate_type_guard.py

@@ -0,0 +1,39 @@
+"""rotate_key()/rotate_jwt() must refuse to operate on mismatched entry types.
+
+Running the API-key rotation path over a JWT entry (or vice versa) used to
+silently corrupt the keystore entry: key_hash/prefix were overwritten,
+signing_secret_hash left stale, version bumped. These guards make the
+mismatch loud instead of silent.
+"""
+import pytest
+from apiauth.keygen import create_api_key_entry, create_jwt_entry, rotate_jwt, rotate_key
+from apiauth.keystore import Keystore
+
+
+def test_rotate_key_on_jwt_raises(tmp_path):
+    tmp_keystore = Keystore(str(tmp_path / "ks"))
+    created = create_jwt_entry(tmp_keystore, name="svc", service="api")
+    with pytest.raises(ValueError, match="rotate_key"):
+        rotate_key(tmp_keystore, created["id"])
+    # Entry is untouched.
+    assert tmp_keystore.get(created["id"])["type"] == "jwt"
+    assert "key_hash" not in tmp_keystore.get(created["id"])
+
+
+def test_rotate_jwt_on_api_key_raises(tmp_path):
+    tmp_keystore = Keystore(str(tmp_path / "ks"))
+    created = create_api_key_entry(tmp_keystore, name="k", service="api")
+    orig = dict(tmp_keystore.get(created["id"]))
+    with pytest.raises(ValueError, match="rotate_jwt"):
+        rotate_jwt(tmp_keystore, created["id"])
+    assert tmp_keystore.get(created["id"]) == orig
+
+
+def test_correct_type_rotation_still_works(tmp_path):
+    tmp_keystore = Keystore(str(tmp_path / "ks"))
+    key_entry = create_api_key_entry(tmp_keystore, name="k", service="api")
+    rotated = rotate_key(tmp_keystore, key_entry["id"])
+    assert rotated["version"] == 2
+    jwt_entry = create_jwt_entry(tmp_keystore, name="j", service="api")
+    rotated_jwt = rotate_jwt(tmp_keystore, jwt_entry["id"])
+    assert rotated_jwt["version"] == 2