Fix login issue (#180)

* fix https cookies & secure

* fix websocked db hoarding

Author: mariusandra

backend/app/api/auth.py

@@ -26,6 +26,13 @@
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/login", auto_error=False)
 
+
+def _should_use_secure_cookie(request: Request) -> bool:
+    if request.url.scheme == "https":
+        return True
+    forwarded_proto = request.headers.get("x-forwarded-proto", "").split(",", 1)[0].strip().lower()
+    return forwarded_proto == "https"
+
 def create_access_token(data: dict, expires_delta: Optional[datetime.timedelta] = None):
     to_encode = data.copy()
     if expires_delta:
@@ -168,12 +175,12 @@ async def login(
         max_age=max_age,
         httponly=True,
         samesite="lax",
-        secure=not (config.DEBUG or config.TEST),
+        secure=_should_use_secure_cookie(request),
     )
     return {"access_token": access_token, "token_type": "bearer"}
 
 @api_no_auth.post("/signup")
-async def signup(data: UserSignup, response: Response, db: Session = Depends(get_db)):
+async def signup(request: Request, data: UserSignup, response: Response, db: Session = Depends(get_db)):
     if config.HASSIO_RUN_MODE is not None:
         raise HTTPException(status_code=401, detail="Signup not allowed with HASSIO_RUN_MODE")
 
@@ -218,7 +225,7 @@ async def signup(data: UserSignup, response: Response, db: Session = Depends(get
         max_age=max_age,
         httponly=True,
         samesite="lax",
-        secure=not (config.DEBUG or config.TEST),
+        secure=_should_use_secure_cookie(request),
     )
     return {"success": True, "access_token": access_token, "token_type": "bearer"}
 

backend/app/api/tests/test_auth.py

@@ -173,3 +173,25 @@ async def test_cookie_auth_on_protected_route(no_auth_client, db: Session):
     no_auth_client.headers.pop("Authorization", None)
     response = await no_auth_client.get("/api/system/metrics")
     assert response.status_code == HTTP_200_OK
+
+
+@pytest.mark.asyncio
+async def test_login_cookie_secure_flag_depends_on_request_scheme(no_auth_client, db: Session):
+    user = User(email="securecookie@example.com")
+    user.set_password("testpassword")
+    db.add(user)
+    db.commit()
+
+    login_data = {"username": "securecookie@example.com", "password": "testpassword"}
+
+    http_response = await no_auth_client.post("/api/login", data=login_data)
+    assert http_response.status_code == HTTP_200_OK
+    assert "Secure" not in http_response.headers.get("set-cookie", "")
+
+    https_response = await no_auth_client.post(
+        "/api/login",
+        data=login_data,
+        headers={"x-forwarded-proto": "https"},
+    )
+    assert https_response.status_code == HTTP_200_OK
+    assert "Secure" in https_response.headers.get("set-cookie", "")

backend/app/websockets.py

@@ -2,9 +2,9 @@
 import json
 from typing import List
 from redis.asyncio import from_url as create_redis, Redis
-from fastapi import WebSocket, WebSocketDisconnect, Depends
-from sqlalchemy.orm import Session
-from app.database import get_db
+from fastapi import WebSocket, WebSocketDisconnect
+
+from app.database import SessionLocal
 
 from app.config import config
 
@@ -69,12 +69,17 @@ async def publish_message(redis: Redis, event: str, data: dict):
 
 def register_ws_routes(app):
     @app.websocket("/ws")
-    async def websocket_endpoint(websocket: WebSocket, db: Session = Depends(get_db)):
+    async def websocket_endpoint(websocket: WebSocket):
         # Full access in the HASSIO ingress mode
         if config.HASSIO_RUN_MODE != "ingress":
             from app.api.auth import get_current_user_from_websocket
 
-            user, error_reason = get_current_user_from_websocket(websocket, db)
+            db = SessionLocal()
+            try:
+                user, error_reason = get_current_user_from_websocket(websocket, db)
+            finally:
+                db.close()
+
             if user is None:
                 await websocket.close(code=1008, reason=error_reason or "Could not validate credentials")
                 return

backend/app/ws/agent_ws.py

@@ -14,11 +14,10 @@
 
 from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status
 from starlette.websockets import WebSocketState
-from sqlalchemy.orm import Session
 from arq import ArqRedis as Redis
 
 # ── project locals ──────────────────────────────────────────────────────────
-from app.database import get_db
+from app.database import SessionLocal
 from app.redis import get_redis, create_redis_connection   # create_… gives a raw conn
 from app.websockets import publish_message
 from app.models.frame import Frame
@@ -35,6 +34,14 @@
 active_sockets_by_frame: dict[int, list[WebSocket]] = {}
 active_sockets: set[WebSocket] = set()
 
+
+async def write_log(redis: Redis, frame_id: int, type: str, line: str, ip: str | None = None):
+    db = SessionLocal()
+    try:
+        await log(db, redis, frame_id, type, line, ip=ip)
+    finally:
+        db.close()
+
 # ────────────────────────────────────────────────────────────────────────────
 # tiny helpers
 # ────────────────────────────────────────────────────────────────────────────
@@ -264,7 +271,6 @@ async def pump_commands(
 @router.websocket("/ws/agent")
 async def ws_agent_endpoint(
     ws: WebSocket,
-    db: Session = Depends(get_db),
     redis: Redis = Depends(get_redis),
 ):
     # ----- rudimentary DoS guard (per-worker) ------------------------------
@@ -286,7 +292,12 @@ async def ws_agent_endpoint(
         return
 
     server_api_key = str(hello_msg.get("serverApiKey", "")) or ""
-    frame = db.query(Frame).filter_by(server_api_key=server_api_key).first()
+    db = SessionLocal()
+    try:
+        frame = db.query(Frame).filter_by(server_api_key=server_api_key).first()
+    finally:
+        db.close()
+
     if frame is None:
         await ws.close(code=status.WS_1008_POLICY_VIOLATION, reason="unknown frame")
         return
@@ -340,7 +351,7 @@ async def ws_agent_endpoint(
         {"active_connections": await number_of_connections_for_frame(redis, frame.id),
          "id": frame.id}
     )
-    await log(db, redis, frame.id, "agent", f'☎️ Frame "{frame.name}" connected ☎️', ip=client_ip)
+    await write_log(redis, frame.id, "agent", f'☎️ Frame "{frame.name}" connected ☎️', ip=client_ip)
 
     # =======================================================================
     #                           RECEIVE LOOP
@@ -433,7 +444,7 @@ async def ws_agent_endpoint(
 
                 for line in data.splitlines():
                     if line:
-                        await log(db, redis, frame.id, stream, line, ip=client_ip)
+                        await write_log(redis, frame.id, stream, line, ip=client_ip)
                         await redis.rpush(STREAM_KEY.format(id=pl["id"]),
                                           json.dumps({"stream": stream, "data": line}).encode())
                 await redis.expire(STREAM_KEY.format(id=pl["id"]), 300)
@@ -463,4 +474,4 @@ async def ws_agent_endpoint(
             {"active_connections": await number_of_connections_for_frame(redis, frame.id),
              "id": frame.id}
         )
-        await log(db, redis, frame.id, "agent", f'👋 Frame "{frame.name}" disconnected 👋', ip=client_ip)
+        await write_log(redis, frame.id, "agent", f'👋 Frame "{frame.name}" disconnected 👋', ip=client_ip)

backend/app/ws/terminal_ws.py

@@ -6,7 +6,7 @@
 from sqlalchemy.orm import Session
 from arq import ArqRedis as Redis
 
-from app.database import get_db
+from app.database import SessionLocal
 from app.redis import get_redis
 from app.models.frame import Frame
 from app.api.auth import get_current_user_from_websocket
@@ -18,22 +18,36 @@
 async def ssh_terminal(
     websocket: WebSocket,
     frame_id: int,
-    db: Session = Depends(get_db),
     redis: Redis = Depends(get_redis),
 ):
-    user, error_reason = get_current_user_from_websocket(websocket, db)
+    db: Session = SessionLocal()
+    try:
+        user, error_reason = get_current_user_from_websocket(websocket, db)
+    finally:
+        db.close()
+
     if user is None:
         await websocket.close(code=1008, reason=error_reason or "Could not validate credentials")
         return
 
     await websocket.accept()
 
-    frame = db.query(Frame).filter(Frame.id == frame_id).first()
+    db = SessionLocal()
+    try:
+        frame = db.query(Frame).filter(Frame.id == frame_id).first()
+    finally:
+        db.close()
+
     if frame is None:
         await websocket.close(code=1008, reason="Frame not found")
         return
 
-    ssh = await get_ssh_connection(db, redis, frame)
+    db = SessionLocal()
+    try:
+        ssh = await get_ssh_connection(db, redis, frame)
+    finally:
+        db.close()
+
     proc = await ssh.create_process(term_type="xterm", encoding="utf-8")
 
     async def pipe(reader):
@@ -61,4 +75,9 @@ async def pipe(reader):
         with contextlib.suppress(Exception):
             proc.stdin.write_eof()
             await proc.wait_closed()
-        await remove_ssh_connection(db, redis, ssh, frame)
+
+        db = SessionLocal()
+        try:
+            await remove_ssh_connection(db, redis, ssh, frame)
+        finally:
+            db.close()