[BUG/Question] Fixing TypeError during WebSocket Authentication Migration from FastAPI 0.96 to 0.97

[monometa]

Describe the bug
I am trying to migrate from FastAPI 0.96 to 0.97 and have azure_scheme as a dependency to protect my WebSocket connection. I think some changes were introduced in FastAPI related to WebSocket dependency and now I am having conflicts between connect from React and API. I'm more of a person who supports change, so if the initial setup was incorrect, then I will be glad to receive advice on a more correct setup or any workaround

To Reproduce

Steps to reproduce the behavior:

1. React app tries to listen to WS using this snippet

import { endpoints } from '../../api';
import { IWebSocketNotification } from './types';
import { *************** } from './***************';
import { useAccessToken } from '../../components/auth';

export const useNotificationService = () => {
  const notifyOnRun = ***************();
  const accessToken = useAccessToken();

  useEffect(() => {
    if (!accessToken) {
      return;
    }

    const webSocket = new WebSocket(endpoints.services.notification(accessToken), []);

    webSocket.onmessage = (event) => {
      const stringObject = JSON.parse(event.data);
      const { data }: IWebSocketNotification = JSON.parse(stringObject);

      notifyOnRun(data);
    };

    return () => {
      webSocket.close();
    };
  }, [notifyOnRun, accessToken]);
};

2. I defined the websocket endpoint in FastAPI

from fastapi import APIRouter, Security
from fastapi_azure_auth import SingleTenantAzureAuthorizationCodeBearer
from .settings import settings

azure_scheme = SingleTenantAzureAuthorizationCodeBearer(
    app_client_id=settings.app_client_id,
    tenant_id=settings.tenant_id,
    allow_guest_users=settings.allow_guest_users,
    scopes={
        f"api://{settings.app_client_id}/user_impersonation": "user_impersonation",
    },
)


@router.websocket("/")
async def websocket_endpoint(
    websocket: WebSocket,
    token: Annotated[str, Query()] = "",
):
    await websocket.accept()

    user = await get_current_user_from_token(token)
    client_queue: Queue = Queue()

    user_id = user.claims["oid"]
    await global_notification_queue.subscribe(queue=client_queue)
    try:
        # While the websocket is open, send notifications to the client
        while True:
            message = await client_queue.get()
            parsed = json.loads(message)
            if parsed["data"].get("user_id") == user_id:
                await websocket.send_json(message)
    except WebSocketDisconnect:
        pass


main_router = APIRouter()
main_router.include_router(router.router, dependencies=[Security(azure_scheme)])

3. In Chrome dev tools under "Console" I see corresponding error
   

Stack trace

ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/uvicorn/protocols/websockets/websockets_impl.py", line 254, in run_asgi
    result = await self.app(self.scope, self.asgi_receive, self.asgi_send)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/uvicorn/middleware/proxy_headers.py", line 78, in __call__
    return await self.app(scope, receive, send)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/fastapi/applications.py", line 282, in __call__
    await super().__call__(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/applications.py", line 122, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/middleware/errors.py", line 149, in __call__
    await self.app(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/middleware/cors.py", line 75, in __call__
    await self.app(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/opentelemetry/instrumentation/asgi/__init__.py", line 596, in __call__
    await self.app(scope, otel_receive, otel_send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/middleware/base.py", line 26, in __call__
    await self.app(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 79, in __call__
    raise exc
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 68, in __call__
    await self.app(scope, receive, sender)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 20, in __call__
    raise e
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 17, in __call__
    await self.app(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/routing.py", line 718, in __call__
    await route.handle(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/routing.py", line 341, in handle
    await self.app(scope, receive, send)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/starlette/routing.py", line 82, in app
    await func(session)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/fastapi/routing.py", line 283, in app
    solved_result = await solve_dependencies(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/fastapi/dependencies/utils.py", line 622, in solve_dependencies
    solved = await call(**sub_values)
                   ^^^^^^^^^^^^^^^^^^
TypeError: AzureAuthorizationCodeBearerBase.__call__() missing 1 required positional argument: 'request'
INFO:     connection open
ERROR:    closing handshake failed
Traceback (most recent call last):
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 959, in transfer_data
    message = await self.read_message()
              ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1029, in read_message
    frame = await self.read_data_frame(max_size=self.max_size)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1104, in read_data_frame
    frame = await self.read_frame(max_size)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1161, in read_frame
    frame = await Frame.read(
            ^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/framing.py", line 68, in read
    data = await reader(2)
           ^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.3/lib/python3.11/asyncio/streams.py", line 727, in readexactly
    raise exceptions.IncompleteReadError(incomplete, n)
asyncio.exceptions.IncompleteReadError: 0 bytes read on a total of 2 expected bytes

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/server.py", line 248, in handler
    await self.close()
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 766, in close
    await self.write_close_frame(Close(code, reason))
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1232, in write_close_frame
    await self.write_frame(True, OP_CLOSE, data, _state=State.CLOSING)
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1205, in write_frame
    await self.drain()
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 1194, in drain
    await self.ensure_open()
  File "/home/user/.pyenv/versions/3.11.3/envs/rgo2/lib/python3.11/site-packages/websockets/legacy/protocol.py", line 935, in ensure_open
    raise self.connection_closed_exc()
websockets.exceptions.ConnectionClosedError: sent 1000 (OK); no close frame received

Your configuration

Config from pyproject.toml:

[tool.poetry.dependencies]
python = "3.11.3"
fastapi = "0.97.0"
httpx = "0.23.3"
psycopg2-binary = "2.9.5"
asyncpg = "0.27.0"
sqlalchemy = "2.0.6"
pydantic = {extras = ["dotenv"], version = "1.10.6"}
asyncpg-listen = "^0.0.6"
fastapi-azure-auth = "^4.0.0"
greenlet = "^2.0.2"
opentelemetry-instrumentation-fastapi = "^0.40b0"
opentelemetry-instrumentation-logging = "^0.40b0"
opentelemetry-exporter-otlp = "^1.19.0"
opentelemetry-api = "^1.19.0"
opentelemetry-sdk = "^1.19.0"
prometheus-fastapi-instrumentator = "^6.1.0"
uvicorn = {extras = ["standard"], version = "0.21.1"}

[tool.poetry.group.dev.dependencies]
jupyter = "^1.0.0"
pytest = "^7.4.2"
pytest-cov = "^4.1.0"
pytest-mock = "^3.11.1"
pytest-asyncio = "^0.21.1"
pytest-alembic = "^0.10.7"
polyfactory = "^2.9.0"
schemathesis = "^3.18.0"
starlette-testclient = "0.2.0"

[JonasKs]

Hi, I haven't looked at websockets much, but I'll try to check this out tomorrow. Does it work on newer versions such as 0.103.2?

[monometa]

Big thanks. I tried to migrate to 0.102.0, and my first attempt resulted in the same error

[JonasKs]

I’m sorry, I didn’t get to this today, but as far as I can see, it’s been added a way to natively use dependencies and websockets: https://github.com/tiangolo/fastapi/releases/tag/0.97.0

I suspect you could validate tokens with a dependencies=[Depends(azure_scheme)]?

[monometa]

Thank you for the update. Unfortunately, I tried out the solution you suggested, but I encountered the same error as before

[monometa]

I was looking at the PR and corresponding description, and the issue might be deeper than I expected - fastapi/fastapi#4534

It seems that Auth for Websocket might have never worked and the 0.97 release helped to realize it 😕. Perhaps adding Security(azure_auth) for Websocket router at all was not the right choice initially

[JonasKs]

Ouf, I see. That's not great, good someone caught it.
We'll have to add custom support then, since the request object is essential for the current dependency.

There's some design questions that needs to be made though:

• How is auth normally handled? Token on initialize?
• What happens when token expires? Should the session terminate?

[itsmejamshid]

any news on that issue? I am facing the same issue right now.

[JonasKs]

No, not really. I haven't looked much into it, and probably won't find the time for a while.
Pull request very welcome 😊

[JonasKs]

#200 will fix this.