Introduce users and Discord auth to inference server (#1772)

olliestanley · yk · web-flow · commit 0b9f4ac98212 · 2023-02-22T06:07:23.000Z
Resolves #1471. The idea here is to introduce a flexible system which allows arbitrarily many authentication methods to be added in future, starting with Discord. We use a Discord OAuth2 flow to login a user (or register a new user), then issue a JSON Web Token (JWT) which can be decoded to obtain a user ID directly. Therefore any new auth method can follow the same pattern and auth can be handled entirely on the backend. This additionally includes a separate Redis instance for the inference stack, resolves the pre-commit failure in `main`, and adds missing dependencies for the inference server. - Ideally we would use `fastapi-discord` which would handle the Discord side for us, however there is no version of it which is compatible with FastAPI versions above `0.84.0`, and we depend on FastAPI `0.88.0` Some code relating to creating/decoding JWTs is duplicated from existing backend code and may be possible to move to `oasst_shared`. We also really need to clean up `main.py` in inference. After this initial inclusion of users to the server is merged, we will require two follow-on issues: - Modify DB schemas and code to associate client actions with DB users - Lock relevant API calls behind authentication by validating provided JWTs The new settings will also need configuring in the environments for the server to operate properly. --------- Co-authored-by: Yannic Kilcher <yk@users.noreply.github.com>
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -153,6 +153,21 @@ services:
       retries: 10
     profiles: ["inference"]
 
+  inference-redis:
+    image: redis
+    restart: always
+    profiles: ["inference"]
+    ports:
+      - 6389:6379
+    healthcheck:
+      test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
+      interval: 2s
+      timeout: 2s
+      retries: 10
+    command: redis-server /usr/local/etc/redis/redis.conf
+    volumes:
+      - ./redis.conf:/usr/local/etc/redis/redis.conf
+
   inference-server:
     build:
       dockerfile: docker/inference/Dockerfile.server
@@ -161,7 +176,7 @@ services:
     image: oasst-inference-server:dev
     environment:
       PORT: 8000
-      REDIS_HOST: redis
+      REDIS_HOST: inference-redis
       POSTGRES_HOST: inference-db
       POSTGRES_DB: oasst_inference
       DEBUG_API_KEYS: '["0000"]'
@@ -172,7 +187,7 @@ services:
     ports:
       - "8000:8000"
     depends_on:
-      redis:
+      inference-redis:
         condition: service_healthy
       inference-db:
         condition: service_healthy
diff --git a/inference/server/alembic/versions/2023_02_21_2205-b365a18db6fd_added_users_table.py b/inference/server/alembic/versions/2023_02_21_2205-b365a18db6fd_added_users_table.py
@@ -0,0 +1,41 @@
+"""added users table
+
+Revision ID: b365a18db6fd
+Revises: 4bead1c4cf52
+Create Date: 2023-02-21 22:05:52.620014
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b365a18db6fd"
+down_revision = "4bead1c4cf52"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "user",
+        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("provider", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("provider_account_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("display_name", sqlmodel.sql.sqltypes.AutoString(length=256), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(op.f("ix_user_provider"), "user", ["provider"], unique=False)
+    op.create_index(op.f("ix_user_provider_account_id"), "user", ["provider_account_id"], unique=False)
+    op.create_index("provider", "user", ["provider_account_id"], unique=True)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("provider", table_name="user")
+    op.drop_index(op.f("ix_user_provider_account_id"), table_name="user")
+    op.drop_index(op.f("ix_user_provider"), table_name="user")
+    op.drop_table("user")
+    # ### end Alembic commands ###
diff --git a/inference/server/main.py b/inference/server/main.py
@@ -1,20 +1,27 @@
 import time
+from datetime import datetime, timedelta
 from pathlib import Path
 
+import aiohttp
 import alembic.command
 import alembic.config
 import fastapi
 import sqlmodel
-from fastapi import Depends
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from fastapi import Depends, HTTPException, Security
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.security import APIKeyCookie
+from jose import jwe, jwt
 from loguru import logger
 from oasst_inference_server import client_handler, deps, interface, models, worker_handler
 from oasst_inference_server.chat_repository import ChatRepository
 from oasst_inference_server.settings import settings
-from oasst_shared.schemas import inference
+from oasst_shared.schemas import inference, protocol
 from prometheus_fastapi_instrumentator import Instrumentator
 
 app = fastapi.FastAPI()
+oauth2_scheme = APIKeyCookie(name=settings.auth_cookie_name)
 
 
 # add prometheus metrics at /metrics
@@ -48,7 +55,7 @@ def get_root_token(token: str = Depends(get_bearer_token)) -> str:
     root_token = settings.root_token
     if token == root_token:
         return token
-    raise fastapi.HTTPException(
+    raise HTTPException(
         status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
         detail="Invalid token",
     )
@@ -106,6 +113,74 @@ def maybe_add_debug_api_keys():
         raise
 
 
+@app.get("/auth/login/discord")
+async def login_discord():
+    redirect_uri = f"{settings.api_root}/auth/callback/discord"
+    auth_url = f"https://discord.com/api/oauth2/authorize?client_id={settings.auth_discord_client_id}&redirect_uri={redirect_uri}&response_type=code&scope=identify"
+    raise HTTPException(status_code=302, headers={"location": auth_url})
+
+
+@app.get("/auth/callback/discord", response_model=protocol.Token)
+async def callback_discord(
+    code: str,
+    db: sqlmodel.Session = Depends(deps.create_session),
+):
+    redirect_uri = f"{settings.api_root}/auth/callback/discord"
+
+    async with aiohttp.ClientSession(raise_for_status=True) as session:
+        # Exchange the auth code for a Discord access token
+        async with session.post(
+            "https://discord.com/api/oauth2/token",
+            data={
+                "client_id": settings.auth_discord_client_id,
+                "client_secret": settings.auth_discord_client_secret,
+                "grant_type": "authorization_code",
+                "code": code,
+                "redirect_uri": redirect_uri,
+                "scope": "identify",
+            },
+        ) as token_response:
+            token_response_json = await token_response.json()
+
+        try:
+            access_token = token_response_json["access_token"]
+        except KeyError:
+            raise HTTPException(status_code=400, detail="Invalid access token response from Discord")
+
+        # Retrieve user's Discord information using access token
+        async with session.get(
+            "https://discord.com/api/users/@me", headers={"Authorization": f"Bearer {access_token}"}
+        ) as user_response:
+            user_response_json = await user_response.json()
+
+    try:
+        discord_id = user_response_json["id"]
+        discord_username = user_response_json["username"]
+    except KeyError:
+        raise HTTPException(status_code=400, detail="Invalid user info response from Discord")
+
+    # Try to find a user in our DB linked to the Discord user
+    user: models.DbUser = query_user_by_provider_id(db, discord_id=discord_id)
+
+    # Create if no user exists
+    if not user:
+        user = models.DbUser(provider="discord", provider_account_id=discord_id, display_name=discord_username)
+
+        db.add(user)
+        db.commit()
+        db.refresh(user)
+
+    # Discord account is authenticated and linked to a user; create JWT
+    access_token = create_access_token(
+        {"user_id": user.id},
+        settings.auth_secret,
+        settings.auth_algorithm,
+        settings.auth_access_token_expire_minutes,
+    )
+
+    return protocol.Token(access_token=access_token, token_type="bearer")
+
+
 @app.get("/chat")
 async def list_chats(cr: ChatRepository = Depends(deps.create_chat_repository)) -> interface.ListChatsResponse:
     """Lists all chats."""
@@ -142,13 +217,11 @@ async def get_chat(id: str, cr: ChatRepository = Depends(deps.create_chat_reposi
 @app.put("/worker")
 def create_worker(
     request: interface.CreateWorkerRequest,
-    root_token: str = fastapi.Depends(get_root_token),
-    session: sqlmodel.Session = fastapi.Depends(deps.create_session),
+    root_token: str = Depends(get_root_token),
+    session: sqlmodel.Session = Depends(deps.create_session),
 ):
     """Allows a client to register a worker."""
-    worker = models.DbWorker(
-        name=request.name,
-    )
+    worker = models.DbWorker(name=request.name)
     session.add(worker)
     session.commit()
     session.refresh(worker)
@@ -157,8 +230,8 @@ def create_worker(
 
 @app.get("/worker")
 def list_workers(
-    root_token: str = fastapi.Depends(get_root_token),
-    session: sqlmodel.Session = fastapi.Depends(deps.create_session),
+    root_token: str = Depends(get_root_token),
+    session: sqlmodel.Session = Depends(deps.create_session),
 ):
     """Lists all workers."""
     workers = session.exec(sqlmodel.select(models.DbWorker)).all()
@@ -168,11 +241,52 @@ def list_workers(
 @app.delete("/worker/{worker_id}")
 def delete_worker(
     worker_id: str,
-    root_token: str = fastapi.Depends(get_root_token),
-    session: sqlmodel.Session = fastapi.Depends(deps.create_session),
+    root_token: str = Depends(get_root_token),
+    session: sqlmodel.Session = Depends(deps.create_session),
 ):
     """Deletes a worker."""
     worker = session.get(models.DbWorker, worker_id)
     session.delete(worker)
     session.commit()
     return fastapi.Response(status_code=200)
+
+
+def query_user_by_provider_id(db: sqlmodel.Session, discord_id: str | None = None) -> models.DbUser | None:
+    """Returns the user associated with a given provider ID if any."""
+    user_qry = db.query(models.DbUser)
+
+    if discord_id:
+        user_qry = user_qry.filter(models.DbUser.provider == "discord").filter(
+            models.DbUser.provider_account_id == discord_id
+        )
+    # elif other IDs...
+    else:
+        return None
+
+    user: models.DbUser = user_qry.first()
+    return user
+
+
+def create_access_token(data: dict, secret: str, algorithm: str, expire_minutes: int) -> str:
+    """Create encoded JSON Web Token (JWT) using the given data."""
+    expires_delta = timedelta(minutes=expire_minutes)
+    to_encode = data.copy()
+    expire = datetime.utcnow() + expires_delta
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, secret, algorithm=algorithm)
+    return encoded_jwt
+
+
+def decode_user_access_token(token: str = Security(oauth2_scheme)) -> dict:
+    """Decode the current user JWT token and return the payload."""
+    # We first generate a key from the auth secret
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=settings.auth_length,
+        salt=settings.auth_salt,
+        info=settings.auth_info,
+    )
+    key = hkdf.derive(settings.auth_secret)
+    # Next we decrypt the JWE token
+    payload = jwe.decrypt(token, key)
+    return payload
diff --git a/inference/server/oasst_inference_server/models.py b/inference/server/oasst_inference_server/models.py
@@ -6,7 +6,7 @@
 import sqlalchemy.dialects.postgresql as pg
 from oasst_inference_server import interface
 from oasst_shared.schemas import inference
-from sqlmodel import Field, Relationship, SQLModel
+from sqlmodel import Field, Index, Relationship, SQLModel
 
 
 class DbMessage(SQLModel, table=True):
@@ -81,5 +81,16 @@ class DbWorker(SQLModel, table=True):
 
     in_compliance_check: bool = Field(default=False, sa_column=sa.Column(sa.Boolean, server_default=sa.text("false")))
     next_compliance_check: datetime.datetime | None = Field(None)
-
     events: list[DbWorkerEvent] = Relationship(back_populates="worker")
+
+
+class DbUser(SQLModel, table=True):
+    __tablename__ = "user"
+    __table_args__ = (Index("provider", "provider_account_id", unique=True),)
+
+    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
+
+    provider: str = Field(..., index=True)
+    provider_account_id: str = Field(..., index=True)
+
+    display_name: str = Field(nullable=False, max_length=256)
diff --git a/inference/server/oasst_inference_server/settings.py b/inference/server/oasst_inference_server/settings.py
@@ -43,5 +43,20 @@ def assemble_db_connection(cls, v: str | None, values: dict[str, Any]) -> Any:
     do_compliance_checks: bool = True
     compliance_check_interval: int = 60
 
+    api_root: str = "https://inference.prod.open-assistant.io"
+
+    use_auth: bool = True
+
+    auth_info: bytes = b"NextAuth.js Generated Encryption Key"
+    auth_salt: bytes = b""
+    auth_length: int = 32
+    auth_secret: str = ""
+    auth_algorithm: str = "HS256"
+    auth_access_token_expire_minutes: int = 60
+    auth_cookie_name: str = "temp"
+
+    auth_discord_client_id: str = ""
+    auth_discord_client_secret: str = ""
+
 
 settings = Settings()
diff --git a/inference/server/requirements.txt b/inference/server/requirements.txt
@@ -1,11 +1,15 @@
+aiohttp
 alembic
+cryptography==39.0.0
 fastapi[all]==0.88.0
 loguru
 nvidia-ml-py
 prometheus-fastapi-instrumentator
 psutil
 psycopg2-binary
 pydantic
+pynvml
+python-jose[cryptography]==3.3.0
 redis
 sqlmodel
 sse-starlette
diff --git a/oasst-shared/oasst_shared/schemas/inference.py b/oasst-shared/oasst_shared/schemas/inference.py
@@ -64,9 +64,7 @@ def __init__(self, **data):
 
 class WorkerConfig(pydantic.BaseModel):
     model_name: str = DEFAULT_MODEL_NAME
-    hardware_info: WorkerHardwareInfo = pydantic.Field(
-        default_factory=WorkerHardwareInfo
-    )
+    hardware_info: WorkerHardwareInfo = pydantic.Field(default_factory=WorkerHardwareInfo)
 
     @property
     def compat_hash(self) -> str:
@@ -81,9 +79,7 @@ class WorkParameters(pydantic.BaseModel):
     top_p: float = 0.9
     temperature: float = 1.0
     repetition_penalty: float | None = None
-    seed: int = pydantic.Field(
-        default_factory=lambda: random.randint(0, 0xFFFF_FFFF_FFFF_FFFF - 1)
-    )
+    seed: int = pydantic.Field(default_factory=lambda: random.randint(0, 0xFFFF_FFFF_FFFF_FFFF - 1))
 
 
 class MessageState(str, enum.Enum):
@@ -111,9 +107,7 @@ class Thread(pydantic.BaseModel):
 
 class WorkRequest(pydantic.BaseModel):
     thread: Thread = pydantic.Field(..., repr=False)
-    created_at: datetime.datetime = pydantic.Field(
-        default_factory=datetime.datetime.utcnow
-    )
+    created_at: datetime.datetime = pydantic.Field(default_factory=datetime.datetime.utcnow)
     parameters: WorkParameters = pydantic.Field(default_factory=WorkParameters)
 
 
