Upload managed_auth.py

managed_auth.py

@@ -4,19 +4,16 @@
 from __future__ import annotations
 
 import base64
+import hashlib
 import html
 import io
-import json
 import os
 import time
 from datetime import datetime, timezone
 from typing import Any, Callable
-from urllib.parse import unquote
-
 import qrcode
 import requests
 import streamlit as st
-import streamlit.components.v1 as components
 from qrcode.constants import ERROR_CORRECT_M
 
 try:
@@ -26,6 +23,11 @@ except ImportError:  # Safe fallback during a delayed dependency rebuild.
 
 from i18n import tr
 
+try:
+    from smartpaper_restore_bridge import restore_bridge
+except Exception:  # Fail open to the approved QR login page.
+    restore_bridge = None
+
 
 # The authentication endpoint should normally respond within a few seconds.
 # Use a direct connection first (ignoring inherited proxy variables), then one
@@ -37,8 +39,8 @@ AUTO_LOGIN_POLL_INTERVAL_SECONDS = 2.0
 TRANSIENT_HTTP_STATUS = {502, 503, 504}
 DEFAULT_PUBLIC_URL = "https://sqhcn-smartpaper.hf.space"
 ACCESS_QR_FILENAME = "SmartPaper_Access_QR.png"
-RESTORE_COOKIE_NAME = "__Host-smartpaper_restore"
-RESTORE_COOKIE_DEFAULT_DAYS = 30
+RESTORE_BRIDGE_COMPONENT_KEY = "sp_restore_bridge"
+RESTORE_DELETE_PENDING_KEY = "sp_restore_delete_pending"
 
 _DIRECT_SESSION = requests.Session()
 _DIRECT_SESSION.trust_env = False
@@ -203,96 +205,93 @@ def _clear_session(*, clear_restore: bool = False) -> None:
         st.session_state.pop("sp_restore_expires_at", None)
 
 
-def _read_restore_cookie() -> str:
-    """Read the restore cookie supplied with the initial Streamlit request."""
-    try:
-        context = getattr(st, "context", None)
-        cookies = getattr(context, "cookies", {}) if context is not None else {}
-        raw_value = str(cookies.get(RESTORE_COOKIE_NAME, "") or "").strip()
-        return unquote(raw_value)
-    except Exception:
-        # Compatibility fallback: never interrupt the approved login page.
-        return ""
+def _restore_command_id(action: str, token: str = "") -> str:
+    """Return a non-secret identifier used to reject stale component values."""
+    normalized_action = str(action or "read").strip().lower()
+    if normalized_action != "write":
+        return normalized_action
+    digest = hashlib.sha256(str(token or "").encode("utf-8")).hexdigest()[:16]
+    return f"write:{digest}"
 
 
-def _render_restore_cookie_command(
-    token: str | None,
+def _call_restore_bridge(
     *,
-    expires_at: str | None = None,
-    reload_parent: bool = False,
-) -> None:
-    """Write or delete the recovery cookie without changing visible UI."""
-    safe_name = json.dumps(RESTORE_COOKIE_NAME)
-    safe_token = json.dumps(str(token or ""))
-    max_age = RESTORE_COOKIE_DEFAULT_DAYS * 24 * 60 * 60
-    parsed_expiry = _parse_iso_timestamp(expires_at)
-    if parsed_expiry is not None:
-        max_age = max(
-            0,
-            int((parsed_expiry - datetime.now(timezone.utc)).total_seconds()),
+    action: str,
+    token: str = "",
+    expires_at: str = "",
+) -> dict[str, Any] | None:
+    """Call the invisible browser bridge and ignore stale/asynchronous values."""
+    if restore_bridge is None:
+        return None
+
+    command_id = _restore_command_id(action, token)
+    try:
+        result = restore_bridge(
+            action=action,
+            command_id=command_id,
+            token=token,
+            expires_at=expires_at,
+            key=RESTORE_BRIDGE_COMPONENT_KEY,
         )
-    if not token:
-        max_age = 0
-
-    script = f"""
-    <script>
-    (() => {{
-      const name = {safe_name};
-      const value = {safe_token};
-      const maxAge = {int(max_age)};
-      const reloadParent = {str(bool(reload_parent)).lower()};
-      const cookieText = `${{name}}=${{encodeURIComponent(value)}}; Max-Age=${{maxAge}}; Path=/; Secure; SameSite=Strict`;
-
-      let targetDocument = document;
-      let cookieActionSucceeded = false;
-
-      try {{
-        if (window.parent && window.parent.document) {{
-          targetDocument = window.parent.document;
-        }}
-        targetDocument.cookie = cookieText;
-        const cookiePresent = targetDocument.cookie
-          .split("; ")
-          .some((entry) => entry.startsWith(`${{name}}=`));
-        cookieActionSucceeded = maxAge === 0
-          ? !cookiePresent
-          : cookiePresent;
-      }} catch (error) {{
-        console.error("SmartPaper restore-cookie action failed", error);
-      }}
-
-      if (!cookieActionSucceeded) {{
-        console.error("SmartPaper restore-cookie verification failed");
-      }}
-
-      if (reloadParent && cookieActionSucceeded) {{
-        window.setTimeout(() => {{
-          try {{
-            window.parent.location.reload();
-          }} catch (error) {{
-            window.location.reload();
-          }}
-        }}, 180);
-      }}
-    }})();
-    </script>
-    """
-    components.html(script, height=0, width=0)
+    except Exception:
+        return None
+
+    if not isinstance(result, dict):
+        return None
+    if str(result.get("command_id") or "") != command_id:
+        return None
+    if str(result.get("action") or "") != action:
+        return None
+    return result
 
 
-def _sync_restore_cookie_to_browser(*, reload_parent: bool = False) -> None:
-    token = str(st.session_state.get("sp_restore_token") or "").strip()
-    if token:
-        _render_restore_cookie_command(
-            token,
-            expires_at=st.session_state.get("sp_restore_expires_at"),
-            reload_parent=reload_parent,
+def _queue_restore_delete() -> None:
+    st.session_state[RESTORE_DELETE_PENDING_KEY] = True
+
+
+def _sync_restore_bridge() -> tuple[str, bool]:
+    """Synchronize recovery-token state with browser localStorage.
+
+    Returns ``(token, suppress_restore)``. ``suppress_restore`` is true while
+    an explicit logout/delete command is pending, preventing a stale browser
+    token from restoring the just-closed session.
+    """
+    if bool(st.session_state.get(RESTORE_DELETE_PENDING_KEY)):
+        result = _call_restore_bridge(action="delete")
+        if restore_bridge is None or (
+            result
+            and str(result.get("status") or "")
+            in {"deleted", "missing", "error"}
+        ):
+            st.session_state.pop(RESTORE_DELETE_PENDING_KEY, None)
+        return "", True
+
+    session_token = str(
+        st.session_state.get("sp_restore_token") or ""
+    ).strip()
+    if session_token:
+        _call_restore_bridge(
+            action="write",
+            token=session_token,
+            expires_at=str(
+                st.session_state.get("sp_restore_expires_at") or ""
+            ),
         )
+        return session_token, False
 
+    result = _call_restore_bridge(action="read")
+    if not result or str(result.get("status") or "") != "found":
+        return "", False
 
-def _delete_restore_cookie(*, reload_parent: bool = False) -> None:
-    _render_restore_cookie_command(None, reload_parent=reload_parent)
+    browser_token = str(result.get("token") or "").strip()
+    if not browser_token:
+        return "", False
 
+    st.session_state["sp_restore_token"] = browser_token
+    st.session_state["sp_restore_expires_at"] = str(
+        result.get("expires_at") or ""
+    )
+    return browser_token, False
 
 def _verify_existing_session() -> bool:
     session_token = st.session_state.get("sp_session_token")
@@ -379,35 +378,41 @@ def _store_authenticated_session(result: dict[str, Any]) -> None:
         st.session_state["sp_restore_expires_at"] = result.get(
             "restore_expires_at"
         )
-        _sync_restore_cookie_to_browser()
+        # A newly authenticated session supersedes any older pending delete.
+        st.session_state.pop(RESTORE_DELETE_PENDING_KEY, None)
     _clear_login_request()
 
 
-def _try_restore_authenticated_session() -> bool:
+def _try_restore_authenticated_session(restore_token: str = "") -> bool:
     """Restore a short-lived session after refresh without another QR scan."""
-    restore_token = str(
-        st.session_state.get("sp_restore_token") or _read_restore_cookie()
+    token = str(
+        restore_token or st.session_state.get("sp_restore_token") or ""
     ).strip()
-    if not restore_token:
+    if not token:
         return False
     try:
-        result = _post("/session/restore", {"restore_token": restore_token})
+        result = _post("/session/restore", {"restore_token": token})
     except (AuthServiceTimeoutError, AuthServiceConnectionError):
+        # Preserve the browser token during a temporary network interruption.
         return False
     except Exception:
         st.session_state.pop("sp_restore_token", None)
         st.session_state.pop("sp_restore_expires_at", None)
-        _delete_restore_cookie()
+        _queue_restore_delete()
         return False
-    if not result.get("ok") or not result.get("session_token"):
+    if (
+        not result.get("ok")
+        or not result.get("session_token")
+        or not result.get("restore_token")
+        or not result.get("user")
+    ):
         st.session_state.pop("sp_restore_token", None)
         st.session_state.pop("sp_restore_expires_at", None)
-        _delete_restore_cookie()
+        _queue_restore_delete()
         return False
     _store_authenticated_session(result)
     return True
 
-
 def _public_url() -> str:
     """Return the permanent, shareable SmartPaper address."""
     value = _runtime_value("SMARTPAPER_PUBLIC_URL") or DEFAULT_PUBLIC_URL
@@ -633,10 +638,9 @@ def _poll_login_once(
 
     if status == "authenticated":
         # Store the authenticated Streamlit session first and rerun within the
-        # same browser session. On the authenticated rerun, the hidden cookie
-        # command is rendered without forcing another page reload, so the
-        # writing interface opens immediately and the restore token has time
-        # to persist for future manual refreshes.
+        # same browser session. On the authenticated rerun, the invisible
+        # local-storage bridge persists the recovery token without forcing a
+        # full page reload, so the writing interface opens immediately.
         _store_authenticated_session(result)
         status_box.success(tr(language, "login_auto_success"))
         st.rerun()
@@ -746,7 +750,7 @@ def logout_alipay_session(
 ) -> None:
     session_token = str(st.session_state.get("sp_session_token") or "").strip()
     restore_token = str(
-        st.session_state.get("sp_restore_token") or _read_restore_cookie()
+        st.session_state.get("sp_restore_token") or ""
     ).strip()
     user_id = str((st.session_state.get("sp_user") or {}).get("id") or "")
     task = st.session_state.get("sp_generation_task") or {}
@@ -776,13 +780,14 @@ def logout_alipay_session(
         except Exception:
             pass
 
-    # Delete the browser token first, then let browser JavaScript reload the
-    # page. This avoids interrupting the cookie deletion with st.rerun().
-    _delete_restore_cookie(reload_parent=True)
+    # Revoke server state first, then queue a browser-storage deletion. The
+    # next run renders the normal login page while the invisible component
+    # removes localStorage; no full-page JavaScript reload is required.
+    _queue_restore_delete()
     _clear_session(clear_restore=True)
     _clear_login_request()
     _clear_generation_session_state()
-    st.stop()
+    st.rerun()
 
 
 def render_logged_in_account(
@@ -923,11 +928,14 @@ def require_alipay_login(
     language: str = "zh",
 ) -> dict[str, Any]:
     """Require Alipay authentication and return the verified user."""
-    _sync_restore_cookie_to_browser()
+    bridge_token, suppress_restore = _sync_restore_bridge()
     if _verify_existing_session():
         return dict(st.session_state.get("sp_user", {}))
-    if _try_restore_authenticated_session():
-        return dict(st.session_state.get("sp_user", {}))
+    if not suppress_restore and _try_restore_authenticated_session(bridge_token):
+        # The Worker rotates restore tokens. Rerun within the same Streamlit
+        # session so the bridge receives and persists the newly issued token
+        # before any later browser refresh.
+        st.rerun()
 
     st.markdown(
         f"<h3 style='text-align:center;'>"