Repository: streamlit/streamlit
File: lib/streamlit/web/server/oidc_mixin.py (Python)

Code before:
        # Authlib 1.6.6+ always writes state to session even when cache is available.
        session: dict[str, Any] = {}

        claims_options = kwargs.pop("claims_options", None)
        state_data = self.framework.get_state_data(session, params.get("state"))
        self.framework.clear_state_data(session, params.get("state"))

        if not state_data:
            raise OAuthError(
                error="invalid_state",
                description="OAuth state not found or expired. Please try logging in again.",
            )

        params = self._format_state_params(state_data, params)  # type: ignore[attr-defined]
        token = self.fetch_access_token(**params, **kwargs)

        if "id_token" in token and "nonce" in state_data:
            userinfo = self.parse_id_token(
                token, nonce=state_data["nonce"], claims_options=claims_options
            )
            token = {**token, "userinfo": userinfo}
        return cast("dict[str, Any]", token)

    def _save_authorize_data(self, **kwargs: Any) -> None:
        """Authlib underlying uses the concept of "session" to store state data.
        In Tornado, we don't have a session, so we use an empty dict as a placeholder.
        Authlib 1.6.6+ always writes state to session even when cache is available.
        """
        state = kwargs.pop("state", None)
        if state:
            session: dict[str, Any] = {}
            self.framework.set_state_data(session, state, kwargs)
        else:
            raise RuntimeError("Missing state value")


class TornadoOAuth(BaseOAuth):
    oauth2_client_cls = TornadoOAuth2App
    framework_integration_cls = TornadoIntegration

    def __init__(
        self,
        config: dict[str, Any] | None = None,
        cache: AuthCache | None = None,
        fetch_token: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
        update_token: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    ):
        super().__init__(
            cache=cache, fetch_token=fetch_token, update_token=update_token
        )
        self.config = config

Review comment (performance):
```suggestion
        In Tornado, we don't have a session, so we use an empty dict as a placeholder. We also override state access to use the cache instead of session in `TornadoIntegration`.
```

Code after:
        # We also override state access to use the cache instead of session in `TornadoIntegration`.
        session: dict[str, Any] = {}

        claims_options = kwargs.pop("claims_options", None)
        state_data = self.framework.get_state_data(session, params.get("state"))
        self.framework.clear_state_data(session, params.get("state"))

        if not state_data:
            raise OAuthError(
                error="invalid_state",
                description="OAuth state not found or expired. Please try logging in again.",
            )

        params = self._format_state_params(state_data, params)  # type: ignore[attr-defined]
        token = self.fetch_access_token(**params, **kwargs)

        if "id_token" in token and "nonce" in state_data:
            userinfo = self.parse_id_token(
                token, nonce=state_data["nonce"], claims_options=claims_options
            )
            token = {**token, "userinfo": userinfo}
        return cast("dict[str, Any]", token)

    def _save_authorize_data(self, **kwargs: Any) -> None:
        """Authlib underlying uses the concept of "session" to store state data.
        In Tornado, we don't have a session, so we use an empty dict as a placeholder.
        We also override state access to use the cache instead of session in `TornadoIntegration`.
        Authlib 1.6.6+ always writes state to session even when cache is available.
        """
        state = kwargs.pop("state", None)
        if state:
            session: dict[str, Any] = {}
            self.framework.set_state_data(session, state, kwargs)
        else:
            raise RuntimeError("Missing state value")


class TornadoOAuth(BaseOAuth):
    oauth2_client_cls = TornadoOAuth2App
    framework_integration_cls = TornadoIntegration

    def __init__(
        self,
        config: dict[str, Any] | None = None,
        cache: AuthCache | None = None,
        fetch_token: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
        update_token: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    ):
        super().__init__(
            cache=cache, fetch_token=fetch_token, update_token=update_token
        )