From dfa74cbd094bf28000efb1432568dbfff304fe71 Mon Sep 17 00:00:00 2001
From: "liujunling.0" <liujunling.0@bytedance.com>
Date: Fri, 4 Sep 2026 11:14:27 +0800
Subject: [PATCH] fix(auth): trim userinfo stored in session cookies

---
 tests/auth/test_oauth2_auth.py       | 39 ++++++++++++++++++++++++++++
 veadk/auth/middleware/oauth2_auth.py | 36 ++++++++++++++++++++++---
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/tests/auth/test_oauth2_auth.py b/tests/auth/test_oauth2_auth.py
index 7beae213a..4d7f5a78c 100644
--- a/tests/auth/test_oauth2_auth.py
+++ b/tests/auth/test_oauth2_auth.py
@@ -205,6 +205,45 @@ def test_oauth2_callback_revalidates_stored_redirect(
     assert response.headers["location"] == expected
 
 
+@pytest.mark.asyncio
+async def test_fetch_user_info_keeps_only_cookie_safe_fields() -> None:
+    config = oauth2_config()
+    config.userinfo_url = "https://identity.example.com/userinfo"
+    config.user_id_field = "employee_id"
+    handler = OAuth2Handler(config)
+    response = Mock()
+    response.raise_for_status.return_value = None
+    response.json.return_value = {
+        "sub": "user-1",
+        "email": "user@example.com",
+        "name": "Example User",
+        "picture": "https://identity.example.com/avatar.png",
+        "employee_id": "employee-1",
+        "tenant_id": "tenant-1",
+        "roles": ["admin"],
+        "external.claims": {"opaque": "x" * 10_000},
+    }
+    handler._http_client.get = AsyncMock(return_value=response)
+
+    user_info = await handler._fetch_user_info("access-token")
+    cookie = handler.encode_session(
+        OAuth2Session(
+            access_token="access-token",
+            expires_at=time.time() + 3600,
+            user_info=user_info,
+        )
+    )
+
+    assert user_info == {
+        "sub": "user-1",
+        "email": "user@example.com",
+        "name": "Example User",
+        "picture": "https://identity.example.com/avatar.png",
+        "employee_id": "employee-1",
+    }
+    assert len(cookie) < 4096
+
+
 def test_refresh_token_keeps_browser_cookie_beyond_access_token_lifetime() -> None:
     handler = OAuth2Handler(oauth2_config())
     now = time.time()
diff --git a/veadk/auth/middleware/oauth2_auth.py b/veadk/auth/middleware/oauth2_auth.py
index 6db3b9e3e..ad4274c85 100644
--- a/veadk/auth/middleware/oauth2_auth.py
+++ b/veadk/auth/middleware/oauth2_auth.py
@@ -558,6 +558,22 @@ def to_authorization_header(self) -> str:
         return f"{self.token_type} {self.access_token}"
 
 
+_USERINFO_SESSION_FIELDS = frozenset(
+    {
+        "sub",
+        "email",
+        "email_verified",
+        "name",
+        "given_name",
+        "family_name",
+        "preferred_username",
+        "picture",
+        "locale",
+        "updated_at",
+    }
+)
+
+
 @runtime_checkable
 class StateStore(Protocol):
     """Protocol for OAuth2 state storage backends.
@@ -968,7 +984,7 @@ async def _refresh_access_token_once(
             return None
 
     async def _fetch_user_info(self, access_token: str) -> dict[str, Any]:
-        """Fetch user information from the userinfo endpoint."""
+        """Fetch the userinfo fields safe to store in a browser session cookie."""
         if not self.config.userinfo_url:
             raise ValueError("userinfo_url not configured")
 
@@ -985,8 +1001,22 @@ async def _fetch_user_info(self, access_token: str) -> dict[str, Any]:
             response.raise_for_status()
 
             user_info = response.json()
-            logger.debug("Fetched user info: %s", user_info)
-            return user_info
+            if not isinstance(user_info, dict):
+                raise ValueError("User info response is not a JSON object")
+
+            fields = _USERINFO_SESSION_FIELDS | {self.config.user_id_field}
+            filtered_user_info = {
+                field: value
+                for field, value in user_info.items()
+                if field in fields
+                and isinstance(value, (str, int, float, bool))
+                and value is not None
+            }
+            logger.debug(
+                "Fetched user info with cookie-safe fields: %s",
+                sorted(filtered_user_info),
+            )
+            return filtered_user_info
 
         except httpx.HTTPStatusError as e:
             logger.error("User info fetch failed: %s", e.response.text)