ws ticket

loon/web/auth/middleware.py

@@ -34,9 +34,9 @@ class JWTAuthBackend(AuthenticationBackend):
             if scheme.lower() == "bearer" and credentials:
                 return credentials
 
-        # Browser WebSocket APIs cannot set Authorization headers.
-        token = conn.query_params.get("token")
-        return token or None
+        if conn.scope["type"] == "websocket":
+            return conn.query_params.get("token") or None
+        return None
 
     async def authenticate(
         self, conn: HTTPConnection
@@ -54,6 +54,13 @@ class JWTAuthBackend(AuthenticationBackend):
         if not isinstance(user_id, int):
             raise AuthenticationError("Invalid or expired token")
 
+        typ = payload.get("typ")
+        if conn.scope["type"] == "websocket":
+            if typ != "ws":
+                raise AuthenticationError("Invalid or expired token")
+        elif typ != "web":
+            raise AuthenticationError("Invalid or expired token")
+
         with Session(engine) as session:
             user = session.exec(select(User).where(User.id == user_id)).first()
 

loon/web/auth/router.py

@@ -2,10 +2,11 @@ from fastapi import APIRouter, HTTPException, Request, status
 from sqlmodel import Session, select
 
 from loon.web import limiter
+from loon.web.auth.middleware import authenticated
 from loon.web.auth.schema import JWTResponse, JWTRequest
 from loon.web.db import engine
 from loon.web.users.models import User
-from loon.web.utils.jwt import generate_token
+from loon.web.utils.jwt import generate_token, generate_ws_token
 from loon.web.utils.password import verify_password
 
 router = APIRouter(prefix="/auth", tags=["auth"])
@@ -18,11 +19,22 @@ async def get_jwt(request: Request, body: JWTRequest) -> JWTResponse:
         statement = select(User).where(User.internal_username == body.username)
         user = session.exec(statement).first()
 
-        if user is None or not verify_password(body.password, user.password):
+        if user is None or user.id is None or not verify_password(body.password, user.password):
             raise HTTPException(
                 status_code=status.HTTP_401_UNAUTHORIZED,
                 detail="Invalid username or password",
             )
 
-        assert user.id is not None
         return JWTResponse(access_token=generate_token(user.id))
+
+
+@router.post("/ws-token", response_model=JWTResponse)
+@authenticated
+async def get_ws_token(request: Request) -> JWTResponse:
+    user_id = request.user.user.id
+    if user_id is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired token",
+        )
+    return JWTResponse(access_token=generate_ws_token(user_id))