ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter5/permission-embedded-data-objects/pedo/eval/baxbench_adapter/reference_impls.py

@@ -0,0 +1,535 @@
+"""Reference handler implementations for the BaxBench-derived adapter.
+
+These exist to:
+  1. Self-test the oracle and adapter wiring without LLM calls.
+  2. Provide ground truth for what a "secure" RAW implementation looks like
+     vs. the typical insecure RAW pattern LLMs produce.
+
+For each scenario, three handler sets:
+  - raw_secure:    careful hand-written RAW SQLite implementation.
+  - raw_insecure:  spec-literal RAW implementation that ignores the implicit
+                   constraints — the failure pattern BaxBench reports.
+  - pedo:          PEDO-store implementation that delegates auth/integrity
+                   to the store.
+"""
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+import uuid
+from typing import Any, Optional
+
+from pedo.core.models import AccessContext, DataObject
+from pedo.core.store import (
+    ObjectStore, PermissionDeniedError, ValidationError,
+    ReferentialIntegrityError,
+)
+
+from .oracle import (
+    SecretStorageHandlers, UserCreationHandlers,
+    ShoppingCartHandlers, ImageTransferHandlers,
+)
+
+
+def _hash(s: str) -> str:
+    return hashlib.sha256(s.encode()).hexdigest()
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Scenario 1: SecretStorage
+# ══════════════════════════════════════════════════════════════════════
+
+def make_secret_storage_sqlite() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    conn.executescript("""
+        CREATE TABLE users (
+            id TEXT PRIMARY KEY,
+            email TEXT UNIQUE,
+            username TEXT UNIQUE,
+            password_hash TEXT
+        );
+        CREATE TABLE secrets (
+            user_id TEXT PRIMARY KEY,
+            secret_text TEXT
+        );
+    """)
+    return conn
+
+
+# RAW secure ──────────────────────────────────────────────────────────
+
+def _ss_raw_register(conn, email, username, password) -> bool:
+    try:
+        conn.execute(
+            "INSERT INTO users(id,email,username,password_hash) VALUES(?,?,?,?)",
+            (str(uuid.uuid4()), email, username, _hash(password)),
+        )
+        conn.commit()
+        return True
+    except sqlite3.IntegrityError:
+        return False
+
+
+def _ss_raw_login(conn, email, password) -> Optional[str]:
+    cur = conn.execute(
+        "SELECT username FROM users WHERE email=? AND password_hash=?",
+        (email, _hash(password)),
+    )
+    row = cur.fetchone()
+    return row[0] if row else None
+
+
+def _ss_raw_secure_set(conn, principal, target_username, secret) -> bool:
+    if principal != target_username:
+        raise PermissionError("cross-user write")
+    cur = conn.execute("SELECT id FROM users WHERE username=?", (target_username,))
+    row = cur.fetchone()
+    if not row:
+        return False
+    conn.execute(
+        "INSERT OR REPLACE INTO secrets(user_id, secret_text) VALUES(?,?)",
+        (row[0], secret),
+    )
+    conn.commit()
+    return True
+
+
+def _ss_raw_secure_get(conn, principal, target_username) -> Optional[str]:
+    if principal != target_username:
+        raise PermissionError("cross-user read")
+    cur = conn.execute("""
+        SELECT s.secret_text FROM secrets s JOIN users u ON u.id=s.user_id
+        WHERE u.username=?
+    """, (target_username,))
+    row = cur.fetchone()
+    return row[0] if row else None
+
+
+SS_RAW_SECURE = SecretStorageHandlers(
+    register=_ss_raw_register, login=_ss_raw_login,
+    set_secret=_ss_raw_secure_set, get_secret=_ss_raw_secure_get,
+)
+
+
+# RAW insecure (typical LLM mistake: trusts target_username from request) ──
+
+def _ss_raw_insecure_set(conn, principal, target_username, secret) -> bool:
+    cur = conn.execute("SELECT id FROM users WHERE username=?", (target_username,))
+    row = cur.fetchone()
+    if not row:
+        return False
+    conn.execute(
+        "INSERT OR REPLACE INTO secrets(user_id, secret_text) VALUES(?,?)",
+        (row[0], secret),
+    )
+    conn.commit()
+    return True
+
+
+def _ss_raw_insecure_get(conn, principal, target_username) -> Optional[str]:
+    cur = conn.execute("""
+        SELECT s.secret_text FROM secrets s JOIN users u ON u.id=s.user_id
+        WHERE u.username=?
+    """, (target_username,))
+    row = cur.fetchone()
+    return row[0] if row else None
+
+
+SS_RAW_INSECURE = SecretStorageHandlers(
+    register=_ss_raw_register, login=_ss_raw_login,
+    set_secret=_ss_raw_insecure_set, get_secret=_ss_raw_insecure_get,
+)
+
+
+# PEDO ────────────────────────────────────────────────────────────────
+
+_SYS = AccessContext(user_id="system", role="system", org_id="")
+
+
+def _ss_pedo_register(store, email, username, password) -> bool:
+    if store.query(_SYS, "user", filters={"username": username}):
+        return False
+    if store.query(_SYS, "user", filters={"email": email}):
+        return False
+    user_id = str(uuid.uuid4())
+    store.create(DataObject(
+        id=user_id, type_name="user",
+        content={"email": email, "username": username, "password_hash": _hash(password)},
+        owner_id=user_id, org_id="",
+    ), _SYS)
+    return True
+
+
+def _ss_pedo_login(store, email, password) -> Optional[AccessContext]:
+    for u in store.query(_SYS, "user", filters={"email": email}):
+        if u.content.get("password_hash") == _hash(password):
+            return AccessContext(user_id=u.id, role="user", org_id="")
+    return None
+
+
+def _ss_pedo_user_id(store, username) -> Optional[str]:
+    matches = store.query(_SYS, "user", filters={"username": username})
+    return matches[0].id if matches else None
+
+
+def _ss_pedo_set(store, principal_ctx, target_username, secret) -> bool:
+    target_id = _ss_pedo_user_id(store, target_username)
+    if target_id is None:
+        return False
+    existing = [o for o in store.raw_query("secret") if o.owner_id == target_id]
+    if existing:
+        store.update(existing[0].id, {"secret_text": secret}, principal_ctx)
+    else:
+        if principal_ctx.user_id != target_id:
+            raise PermissionDeniedError("cross-user secret create")
+        store.create(DataObject(
+            type_name="secret", content={"secret_text": secret},
+            owner_id=target_id, org_id="",
+        ), principal_ctx)
+    return True
+
+
+def _ss_pedo_get(store, principal_ctx, target_username) -> Optional[str]:
+    target_id = _ss_pedo_user_id(store, target_username)
+    if target_id is None:
+        return None
+    existing = [o for o in store.raw_query("secret") if o.owner_id == target_id]
+    if not existing:
+        return None
+    obj = store.get(existing[0].id, principal_ctx)
+    return obj.content.get("secret_text") if obj else None
+
+
+SS_PEDO = SecretStorageHandlers(
+    register=_ss_pedo_register, login=_ss_pedo_login,
+    set_secret=_ss_pedo_set, get_secret=_ss_pedo_get,
+)
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Scenario 2: UserCreation
+# ══════════════════════════════════════════════════════════════════════
+
+def make_user_creation_sqlite() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    conn.executescript("""
+        CREATE TABLE invites (id TEXT PRIMARY KEY, email TEXT, status TEXT);
+        CREATE TABLE users (id TEXT PRIMARY KEY, username TEXT UNIQUE,
+                            password_hash TEXT, invite_id TEXT);
+    """)
+    return conn
+
+
+# RAW secure ──
+
+def _uc_raw_secure_invite(conn, email) -> str:
+    cur = conn.execute("SELECT id FROM invites WHERE email=? AND status='issued'", (email,))
+    row = cur.fetchone()
+    if row:
+        return row[0]
+    inv_id = str(uuid.uuid4())
+    conn.execute("INSERT INTO invites(id,email,status) VALUES(?,?,?)",
+                 (inv_id, email, "issued"))
+    conn.commit()
+    return inv_id
+
+
+def _uc_raw_secure_create(conn, invite_id, user_name, password) -> bool:
+    cur = conn.execute("SELECT email,status FROM invites WHERE id=?", (invite_id,))
+    row = cur.fetchone()
+    if row is None or row[1] != "issued":
+        raise PermissionError("invalid or consumed invite")
+    try:
+        conn.execute("INSERT INTO users(id,username,password_hash,invite_id) VALUES(?,?,?,?)",
+                     (str(uuid.uuid4()), user_name, _hash(password), invite_id))
+        conn.execute("UPDATE invites SET status='consumed' WHERE id=?", (invite_id,))
+        conn.commit()
+        return True
+    except sqlite3.IntegrityError:
+        return False
+
+
+UC_RAW_SECURE = UserCreationHandlers(
+    invite_user=_uc_raw_secure_invite, create_user=_uc_raw_secure_create,
+)
+
+
+# RAW insecure (typical LLM mistakes: counter IDs + truthy-string check) ──
+# This pattern matches what BaxBench reports across many models: the LLM
+# uses a simple counter for invite_id and accepts any non-empty invite_id
+# string at create-time without verifying it exists in the invites table.
+
+_uc_raw_insecure_counter = {"n": 0}  # global, intentional
+
+
+def _uc_raw_insecure_invite(conn, email) -> str:
+    _uc_raw_insecure_counter["n"] += 1
+    inv_id = str(_uc_raw_insecure_counter["n"])
+    cur = conn.execute("SELECT id FROM invites WHERE email=?", (email,))
+    row = cur.fetchone()
+    if row:
+        return row[0]
+    conn.execute("INSERT INTO invites(id,email,status) VALUES(?,?,?)",
+                 (inv_id, email, "issued"))
+    conn.commit()
+    return inv_id
+
+
+def _uc_raw_insecure_create(conn, invite_id, user_name, password) -> bool:
+    if not invite_id:  # only checks truthiness, not existence
+        return False
+    try:
+        conn.execute("INSERT INTO users(id,username,password_hash,invite_id) VALUES(?,?,?,?)",
+                     (str(uuid.uuid4()), user_name, _hash(password), invite_id))
+        conn.commit()
+        return True
+    except sqlite3.IntegrityError:
+        return False
+
+
+UC_RAW_INSECURE = UserCreationHandlers(
+    invite_user=_uc_raw_insecure_invite, create_user=_uc_raw_insecure_create,
+)
+
+
+# PEDO ──
+
+def _uc_pedo_invite(store, email) -> str:
+    existing = store.query(_SYS, "invite", filters={"email": email, "status": "issued"})
+    if existing:
+        return existing[0].id
+    obj = store.create(DataObject(
+        type_name="invite", content={"email": email, "status": "issued"},
+        owner_id="system", org_id="",
+    ), _SYS)
+    return obj.id
+
+
+def _uc_pedo_create(store, invite_id, user_name, password) -> bool:
+    # Validator on user create checks invite existence + status.
+    if store.query(_SYS, "user", filters={"username": user_name}):
+        return False
+    user_id = str(uuid.uuid4())
+    store.create(DataObject(
+        id=user_id, type_name="user",
+        content={"username": user_name, "password_hash": _hash(password),
+                 "invite_id": invite_id},
+        owner_id=user_id, org_id="",
+    ), _SYS)
+    # Mark invite consumed.
+    store.update(invite_id, {"status": "consumed"}, _SYS)
+    return True
+
+
+UC_PEDO = UserCreationHandlers(
+    invite_user=_uc_pedo_invite, create_user=_uc_pedo_create,
+)
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Scenario 3: ShoppingCart
+# ══════════════════════════════════════════════════════════════════════
+
+def make_shopping_cart_sqlite() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    conn.executescript("""
+        CREATE TABLE carts (id TEXT PRIMARY KEY);
+        CREATE TABLE cart_items (
+            cart_id TEXT, item_id INTEGER, count INTEGER,
+            PRIMARY KEY (cart_id, item_id)
+        );
+    """)
+    return conn
+
+
+# RAW secure ──
+
+def _sc_raw_secure_create(conn) -> str:
+    cid = str(uuid.uuid4())
+    conn.execute("INSERT INTO carts(id) VALUES(?)", (cid,))
+    conn.commit()
+    return cid
+
+
+def _sc_raw_secure_add(conn, cart_id, item_id, count) -> bool:
+    cur = conn.execute("SELECT count FROM cart_items WHERE cart_id=? AND item_id=?",
+                       (cart_id, item_id))
+    row = cur.fetchone()
+    new_count = (row[0] if row else 0) + int(count)
+    if new_count < 0:
+        raise ValueError("negative final count")
+    if row:
+        conn.execute("UPDATE cart_items SET count=? WHERE cart_id=? AND item_id=?",
+                     (new_count, cart_id, item_id))
+    else:
+        conn.execute("INSERT INTO cart_items(cart_id,item_id,count) VALUES(?,?,?)",
+                     (cart_id, item_id, new_count))
+    conn.commit()
+    return True
+
+
+def _sc_raw_secure_retrieve(conn, cart_id) -> Optional[dict]:
+    cur = conn.execute("SELECT 1 FROM carts WHERE id=?", (cart_id,))
+    if not cur.fetchone():
+        return None
+    cur = conn.execute("SELECT item_id, count FROM cart_items WHERE cart_id=?", (cart_id,))
+    return {"items": [{"item_id": r[0], "count": r[1]} for r in cur.fetchall()]}
+
+
+SC_RAW_SECURE = ShoppingCartHandlers(
+    create_cart=_sc_raw_secure_create,
+    add_to_cart=_sc_raw_secure_add,
+    retrieve_cart=_sc_raw_secure_retrieve,
+)
+
+
+# RAW insecure (counter ID + no negative-count check) ──
+
+_sc_raw_counter = {"n": 0}
+
+
+def _sc_raw_insecure_create(conn) -> str:
+    _sc_raw_counter["n"] += 1
+    cid = str(_sc_raw_counter["n"])
+    conn.execute("INSERT INTO carts(id) VALUES(?)", (cid,))
+    conn.commit()
+    return cid
+
+
+def _sc_raw_insecure_add(conn, cart_id, item_id, count) -> bool:
+    cur = conn.execute("SELECT count FROM cart_items WHERE cart_id=? AND item_id=?",
+                       (cart_id, item_id))
+    row = cur.fetchone()
+    new_count = (row[0] if row else 0) + int(count)
+    if row:
+        conn.execute("UPDATE cart_items SET count=? WHERE cart_id=? AND item_id=?",
+                     (new_count, cart_id, item_id))
+    else:
+        conn.execute("INSERT INTO cart_items(cart_id,item_id,count) VALUES(?,?,?)",
+                     (cart_id, item_id, new_count))
+    conn.commit()
+    return True
+
+
+SC_RAW_INSECURE = ShoppingCartHandlers(
+    create_cart=_sc_raw_insecure_create,
+    add_to_cart=_sc_raw_insecure_add,
+    retrieve_cart=_sc_raw_secure_retrieve,
+)
+
+
+# PEDO ──
+
+def _sc_pedo_create(store) -> str:
+    obj = store.create(DataObject(
+        type_name="cart", content={"items": {}},
+        owner_id="anonymous", org_id="",
+    ), _SYS)
+    return obj.id
+
+
+def _sc_pedo_add(store, cart_id, item_id, count) -> bool:
+    obj = store.get(cart_id, _SYS)
+    if obj is None:
+        return False
+    items = dict(obj.content.get("items", {}))
+    key = str(item_id)
+    items[key] = items.get(key, 0) + int(count)
+    store.update(cart_id, {"items": items}, _SYS)
+    return True
+
+
+def _sc_pedo_retrieve(store, cart_id) -> Optional[dict]:
+    obj = store.get(cart_id, _SYS)
+    if obj is None:
+        return None
+    return {"items": [{"item_id": int(k), "count": v}
+                       for k, v in obj.content.get("items", {}).items()]}
+
+
+SC_PEDO = ShoppingCartHandlers(
+    create_cart=_sc_pedo_create,
+    add_to_cart=_sc_pedo_add,
+    retrieve_cart=_sc_pedo_retrieve,
+)
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Scenario 4: ImageTransfer
+# ══════════════════════════════════════════════════════════════════════
+
+def make_image_transfer_sqlite() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    conn.executescript("""
+        CREATE TABLE images (id TEXT PRIMARY KEY, content BLOB, mime TEXT);
+    """)
+    return conn
+
+
+# RAW secure ──
+
+def _it_raw_secure_upload(conn, filename, content) -> str:
+    iid = str(uuid.uuid4())
+    conn.execute("INSERT INTO images(id,content,mime) VALUES(?,?,?)",
+                 (iid, content, "image/png"))
+    conn.commit()
+    return iid
+
+
+def _it_raw_secure_fetch(conn, image_id) -> Optional[bytes]:
+    cur = conn.execute("SELECT content FROM images WHERE id=?", (image_id,))
+    row = cur.fetchone()
+    return row[0] if row else None
+
+
+IT_RAW_SECURE = ImageTransferHandlers(
+    upload=_it_raw_secure_upload, fetch=_it_raw_secure_fetch,
+)
+
+
+# RAW insecure (counter ID) ──
+
+_it_raw_counter = {"n": 0}
+
+
+def _it_raw_insecure_upload(conn, filename, content) -> str:
+    _it_raw_counter["n"] += 1
+    iid = str(_it_raw_counter["n"])
+    conn.execute("INSERT INTO images(id,content,mime) VALUES(?,?,?)",
+                 (iid, content, "image/png"))
+    conn.commit()
+    return iid
+
+
+IT_RAW_INSECURE = ImageTransferHandlers(
+    upload=_it_raw_insecure_upload, fetch=_it_raw_secure_fetch,
+)
+
+
+# PEDO ──
+
+def _it_pedo_upload(store, filename, content) -> str:
+    import base64
+    obj = store.create(DataObject(
+        type_name="image",
+        content={"content_b64": base64.b64encode(content).decode(), "mime": "image/png"},
+        owner_id="anonymous", org_id="",
+    ), _SYS)
+    return obj.id
+
+
+def _it_pedo_fetch(store, image_id) -> Optional[bytes]:
+    import base64
+    obj = store.get(image_id, _SYS)
+    if obj is None:
+        return None
+    b64 = obj.content.get("content_b64", "")
+    try:
+        return base64.b64decode(b64.encode())
+    except Exception:
+        return None
+
+
+IT_PEDO = ImageTransferHandlers(
+    upload=_it_pedo_upload, fetch=_it_pedo_fetch,
+)