ws ticket

.dockerignore

@@ -0,0 +1,24 @@
+.git/
+.gitea/
+.gitignore
+
+.venv/
+__pycache__/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.uv/
+
+*.py[cod]
+*.pyo
+*.log
+
+.env
+.env.*
+.secrets.toml
+
+db.sqlite3
+
+Dockerfile
+docker-compose*.yml
+README.md

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))

loon/web/utils/jwt.py

@@ -1,4 +1,5 @@
 import time
+from typing import Literal
 
 import jwt
 
@@ -8,13 +9,23 @@ JWT_SECRET = settings.jwt.secret
 JWT_ALGORITHM = "HS256"
 
 
-def generate_token(user_id: int) -> str:
+def _encode_jwt(user_id: int, ttl: int, typ: Literal["ws", "web"]) -> str:
     payload = {
         "user_id": user_id,
-        "exp": int(time.time()) + 60 * 60 * 24,
+        "exp": int(time.time()) + ttl,
+        "typ": typ,
     }
     return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
 
+
+def generate_token(user_id: int) -> str:
+    return _encode_jwt(user_id, settings.jwt.access_ttl, "web")
+
+
+def generate_ws_token(user_id: int) -> str:
+    return _encode_jwt(user_id, settings.jwt.ws_ttl, "ws")
+
+
 def decode_token(token: str) -> dict:
     payload = jwt.decode(token, JWT_SECRET)
     return payload

settings.toml

@@ -24,3 +24,5 @@ log_data = false
 
 [jwt]
 secret = 'secret'
+access_ttl = 86400 # 24h
+ws_ttl = 45