Initial commit

loon/web/__init__.py

@@ -0,0 +1,47 @@
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.middleware.cors import CORSMiddleware
+
+from config import settings
+from loon.mqtt.manager import MQTTManager
+from loon.web.auth.middleware import JWTAuthBackend, on_auth_error
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.mqtt_manager = MQTTManager.setup_and_start()
+    yield
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    AuthenticationMiddleware,
+    backend=JWTAuthBackend(),
+    on_error=on_auth_error,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.web.allowed_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def get_mqtt_manager() -> MQTTManager:
+    return app.state.mqtt_manager
+
+from loon.web.auth.router import router as auth_router
+from loon.web.users.router import router as users_router
+from loon.web.world.router import router as world_router
+from loon.web.server.router import router as server_router
+from loon.web.admin.router import router as admin_router
+app.include_router(auth_router)
+app.include_router(users_router)
+app.include_router(world_router)
+app.include_router(server_router)
+app.include_router(admin_router)

loon/web/admin/cli.py

@@ -0,0 +1,25 @@
+from sqlmodel import Session, select
+
+from loon.web.db import engine
+from loon.web.users.models import User
+
+
+def choose_admin(username: str) -> None:
+    with Session(engine) as session:
+        user = session.exec(
+            select(User).where(User.internal_username == username)
+        ).first()
+
+        if user is None:
+            print(f"No user named {username!r}.")
+            return
+
+        if user.admin:
+            print(f"{username} is already admin.")
+            return
+
+        user.admin = True
+        session.add(user)
+        session.commit()
+
+        print(f"Promoted {username} to admin.")

loon/web/admin/router.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+from starlette.requests import Request
+
+from loon.web.auth.middleware import admin
+
+router = APIRouter(prefix="/admin", tags=["admin"])
+

loon/web/auth/middleware.py

@@ -0,0 +1,72 @@
+import jwt
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    AuthenticationError,
+    BaseUser,
+    requires as _requires,
+)
+from starlette.requests import HTTPConnection
+from starlette.responses import JSONResponse, Response
+from sqlmodel import Session, select
+
+from loon.web.db import engine
+from loon.web.users.models import User
+from loon.web.utils.jwt import JWT_SECRET, JWT_ALGORITHM
+
+authenticated = _requires("authenticated", status_code=401)
+admin = _requires(["authenticated", "admin"], status_code=403)
+
+
+class AuthenticatedUser(BaseUser):
+    def __init__(self, user: User) -> None:
+        self.user = user
+
+    @property
+    def is_authenticated(self) -> bool:
+        return True
+
+class JWTAuthBackend(AuthenticationBackend):
+    def _extract_token(self, conn: HTTPConnection) -> str | None:
+        auth_header = conn.headers.get("Authorization")
+        if auth_header is not None:
+            scheme, _, credentials = auth_header.partition(" ")
+            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
+
+    async def authenticate(
+        self, conn: HTTPConnection
+    ) -> tuple[AuthCredentials, AuthenticatedUser] | None:
+        credentials = self._extract_token(conn)
+        if credentials is None:
+            return None
+
+        try:
+            payload = jwt.decode(credentials, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        except jwt.InvalidTokenError:
+            raise AuthenticationError("Invalid or expired token")
+
+        user_id = payload.get("user_id")
+        if not isinstance(user_id, int):
+            raise AuthenticationError("Invalid or expired token")
+
+        with Session(engine) as session:
+            user = session.exec(select(User).where(User.id == user_id)).first()
+
+        if user is None:
+            raise AuthenticationError("User not found")
+
+        scopes = ["authenticated"]
+
+        if user.admin:
+            scopes.append("admin")
+
+        return AuthCredentials(scopes), AuthenticatedUser(user)
+
+
+def on_auth_error(conn: HTTPConnection, exc: Exception) -> Response:
+    return JSONResponse({"detail": str(exc)}, status_code=401)

loon/web/auth/router.py

@@ -0,0 +1,26 @@
+from fastapi import APIRouter, HTTPException, status
+from sqlmodel import Session, select
+
+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.password import verify_password
+
+router = APIRouter(prefix="/auth", tags=["auth"])
+
+
+@router.post("/jwt/", response_model=JWTResponse)
+async def get_jwt(body: JWTRequest) -> JWTResponse:
+    with Session(engine) as session:
+        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):
+            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))

loon/web/auth/schema.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class JWTRequest(BaseModel):
+    username: str
+    password: str
+
+
+class JWTResponse(BaseModel):
+    access_token: str
+    token_type: str = "bearer"

loon/web/db/__init__.py

@@ -0,0 +1,24 @@
+from sqlmodel import SQLModel, create_engine
+
+from config import settings
+
+
+def database_url(url: str | None = None) -> str:
+    url = url or settings.database.url
+    for src, dst in (
+        ("mysql://", "mysql+pymysql://"),
+        ("postgres://", "postgresql+psycopg://"),
+        ("postgresql://", "postgresql+psycopg://"),
+    ):
+        if url.startswith(src):
+            return dst + url.removeprefix(src)
+    return url
+
+
+url = database_url()
+kwargs: dict = {}
+if url.startswith("sqlite"):
+    kwargs["connect_args"] = {"check_same_thread": False}
+
+engine = create_engine(url, echo=True, **kwargs)
+SQLModel.metadata.create_all(engine)

loon/web/server/router.py

@@ -0,0 +1,13 @@
+from fastapi import APIRouter, Request
+
+from loon.web import get_mqtt_manager
+from loon.web.auth.middleware import authenticated
+
+router = APIRouter(prefix="/server", tags=["server"])
+
+
+@router.post("/info/request")
+@authenticated
+async def request_server_info(request: Request):
+    get_mqtt_manager().publish("loon/server/info/request")
+    return 200

loon/web/users/models.py

@@ -0,0 +1,9 @@
+from sqlmodel import SQLModel, Field
+
+
+class User(SQLModel, table=True):
+    id: int = Field(default=None, primary_key=True)
+    uuid: str = Field(unique=True)
+    internal_username: str
+    password: str
+    admin: bool = Field(default=False)

loon/web/users/router.py

@@ -0,0 +1,60 @@
+import asyncio
+import threading
+
+from fastapi import APIRouter, Request
+from starlette.websockets import WebSocket, WebSocketDisconnect
+
+from loon.web import get_mqtt_manager
+from loon.web.auth.middleware import authenticated
+from loon.web.users.state import user_threads
+
+router = APIRouter(prefix="/users", tags=["users"])
+
+
+@router.get("/me")
+@authenticated
+async def me(request: Request):
+    return request.user.user.model_dump(exclude={"password", "id"})
+
+@router.post("/inventory/request")
+@authenticated
+async def request_inventory(request: Request):
+    user = request.user.user
+    get_mqtt_manager().publish(f"loon/player/{user.uuid}/inventory/full/request")
+
+    return 200
+
+@router.post("/online/request")
+@authenticated
+async def request_online(request: Request):
+    user = request.user.user
+    get_mqtt_manager().publish(f"loon/player/{user.uuid}/online/request")
+
+    return 200
+
+@router.post("/position/request")
+@authenticated
+async def request_position(request: Request):
+    user = request.user.user
+    get_mqtt_manager().publish(f"loon/player/{user.uuid}/position/request")
+
+    return 200
+
+
+@router.websocket("/me")
+@authenticated
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    user = websocket.user.user
+
+    queue = asyncio.Queue()
+    user_threads[user.uuid] = queue
+
+    try:
+        while True:
+            message = await queue.get()
+            await websocket.send_text(message)
+    except WebSocketDisconnect:
+        pass
+    finally:
+        user_threads.pop(user.uuid, None)

loon/web/users/state.py

@@ -0,0 +1 @@
+user_threads = {}

loon/web/utils/jwt.py

@@ -0,0 +1,20 @@
+import time
+
+import jwt
+
+from config import settings
+
+JWT_SECRET = settings.jwt.secret
+JWT_ALGORITHM = "HS256"
+
+
+def generate_token(user_id: int) -> str:
+    payload = {
+        "user_id": user_id,
+        "exp": int(time.time()) + 60 * 60 * 24,
+    }
+    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
+
+def decode_token(token: str) -> dict:
+    payload = jwt.decode(token, JWT_SECRET)
+    return payload

loon/web/utils/password.py

@@ -0,0 +1,19 @@
+from argon2 import PasswordHasher
+from argon2.exceptions import InvalidHashError, VerifyMismatchError
+
+
+def hash_password(password: str) -> str:
+    ph = PasswordHasher()
+
+    hashed = ph.hash(password)
+
+    return str(hashed)
+
+
+def verify_password(password: str, hashed_password: str) -> bool:
+    ph = PasswordHasher()
+
+    try:
+        return ph.verify(hashed_password, password)
+    except (VerifyMismatchError, InvalidHashError):
+        return False

loon/web/world/router.py

@@ -0,0 +1,80 @@
+import json
+
+import redis
+from fastapi import APIRouter, Request
+
+from config import settings
+from loon.redis.chunk import get_chunk_cache, reset_chunk_cache
+from loon.web import get_mqtt_manager
+from loon.web.auth.middleware import authenticated
+from loon.web.users.state import user_threads
+
+router = APIRouter(prefix="/world", tags=["world"])
+
+
+@router.get("/request")
+@authenticated
+async def request_world(
+    request: Request,
+    x_start: int,
+    x_end: int,
+    z_start: int,
+    z_end: int,
+):
+    if not settings.redis.enabled:
+        get_mqtt_manager().publish(f"loon/world/chunks/{x_start}:{x_end}/{z_start}:{z_end}/request")
+        return 200
+
+    misses = []
+    thread = user_threads.get(request.user.user.uuid)
+
+    try:
+        for x in range(x_start, x_end + 1):
+            for z in range(z_start, z_end + 1):
+                cached_chunk = get_chunk_cache(x, z)
+
+                if cached_chunk is not None:
+                    if thread:
+                        await thread.put(
+                            json.dumps({"topic": f"world/chunk/{x}/{z}", "payload": cached_chunk})
+                        )
+                else:
+                    misses.append((x, z))
+    except redis.RedisError as e:
+        get_mqtt_manager().publish(f"loon/world/chunks/{x_start}:{x_end}/{z_start}:{z_end}/request")
+        return 200
+
+    for xa, xb, za, zb in miss_row_rects(misses):
+        get_mqtt_manager().publish(f"loon/world/chunks/{xa}:{xb}/{za}:{zb}/request")
+
+    return 200
+
+
+def miss_row_rects(misses: list):
+    rows = {}
+
+    for x, z in misses:
+        if z not in rows:
+            rows[z] = []
+        rows[z].append(x)
+
+    rects = []
+
+    for z, xs in rows.items():
+        xs.sort()
+
+        start = xs[0]
+        end = xs[0]
+
+        for x in xs[1:]:
+            if x == end + 1:
+                end = x
+            else:
+                rects.append((start, end, z, z))
+                start = x
+                end = x
+
+        rects.append((start, end, z, z))
+
+    return rects
+