Update auth.py

auth.py

@@ -126,18 +126,46 @@ async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends
     logger.info(f"User logged in: {user['email']}")
     return {"access_token": create_access_token(user['email']), "refresh_token": create_refresh_token(user['email']), "token_type": "bearer", "name": user['name'], "avatar": avatar_url}
 
-# Public endpoint: fetch user data by email, no auth
 @router.get("/user/data", response_model=User)
-async def get_user_data_public(email: str = Query(..., description="User email")):
-    logger.info(f"[PUBLIC] Fetching data for email: {email}")
-    user = get_user(email)
-    if not user:
-        logger.warning(f"[PUBLIC] No user found for email: {email}")
-        raise HTTPException(status_code=404, detail="User not found")
-    avatar_url = f"/auth/avatar/{user['avatar']}" if user.get('avatar') else None
-    return {"name": user['name'], "email": user['email'], "avatar": avatar_url, "chat_histories": user.get('chat_histories', [])}
+async def get_user_data(
+    request: Request,
+    email: Optional[str] = Query(None, description="User email")
+):
+    # Determine if Authorization header is provided
+    auth: str = request.headers.get("Authorization", "")
+    user = None
+    if auth.startswith("Bearer "):
+        token = auth.split(" ", 1)[1]
+        # Authenticated path
+        try:
+            payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
+            user_email: str = payload.get("sub")
+            if not user_email:
+                raise HTTPException(status_code=401, detail="Invalid token payload")
+            user = get_user(user_email)
+            if not user:
+                raise HTTPException(status_code=404, detail="User not found")
+        except JWTError as e:
+            logger.error(f"[USER DATA] JWT decode failed: {e}")
+            raise HTTPException(status_code=401, detail="Invalid token")
+    elif email:
+        # Public path
+        user = get_user(email)
+        if not user:
+            raise HTTPException(status_code=404, detail="User not found")
+    else:
+        raise HTTPException(status_code=422, detail="Provide either Authorization header or 'email' query parameter")
 
-@router.put("/user/update")
+    # Build response
+    avatar_url = f"/auth/avatar/{user['avatar']}" if user.get('avatar') else None
+    return {
+        "name": user['name'],
+        "email": user['email'],
+        "avatar": avatar_url,
+        "chat_histories": user.get('chat_histories', [])
+    }
+
+@router.put("/user/update")("/user/update")
 async def update_user(
     request: Request,
     name: Optional[str] = Form(None),