Merge pull request #920 from LAION-AI/857-backend-token-decrypt

857 backend token decrypt

backend/oasst_backend/api/v1/api.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter
 from oasst_backend.api.v1 import (
     admin,
+    auth,
     frontend_messages,
     frontend_users,
     hugging_face,
@@ -23,3 +24,4 @@ api_router.include_router(stats.router, prefix="/stats", tags=["stats"])
 api_router.include_router(leaderboards.router, prefix="/leaderboards", tags=["leaderboards"])
 api_router.include_router(hugging_face.router, prefix="/hf", tags=["hugging_face"])
 api_router.include_router(admin.router, prefix="/admin", tags=["admin"])
+api_router.include_router(auth.router, prefix="/auth", tags=["auth"])

backend/oasst_backend/api/v1/auth.py

@@ -0,0 +1,46 @@
+from typing import Union
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from fastapi import APIRouter, Depends, Security
+from fastapi.security import APIKeyCookie
+from jose import jwe
+from oasst_backend.config import settings
+from pydantic import BaseModel, EmailStr
+
+router = APIRouter()
+
+oauth2_scheme = APIKeyCookie(name=settings.AUTH_COOKIE_NAME)
+
+
+class TokenData(BaseModel):
+    """
+    A minimal re-creation of the web's token type.  To be expanded later.
+    """
+
+    email: Union[EmailStr, None] = None
+
+
+async def get_current_user(token: str = Security(oauth2_scheme)):
+    """
+    Decrypts the user's JSON Web Token using HKDF encryption and returns the
+    TokenData.
+    """
+    # 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)
+    # Finally we have the real token JSON payload and can do whatever we want.
+    return TokenData.parse_raw(payload)
+
+
+@router.get("/check", response_model=str)
+async def auth_check(token_data: TokenData = Depends(get_current_user)):
+    """Returns the user's email if it can be decrypted."""
+    return token_data.email

backend/oasst_backend/config.py

@@ -124,6 +124,15 @@ class Settings(BaseSettings):
     API_V1_STR: str = "/api/v1"
     OFFICIAL_WEB_API_KEY: str = "1234"
 
+    # Encryption fields for handling the web generated JSON Web Tokens.
+    # These fields need to be shared with the web's auth settings in order to
+    # correctly decrypt the web tokens.
+    AUTH_INFO: bytes = b"NextAuth.js Generated Encryption Key"
+    AUTH_SALT: bytes = b""
+    AUTH_LENGTH: int = 32
+    AUTH_SECRET: bytes = b"O/M2uIbGj+lDD2oyNa8ax4jEOJqCPJzO53UbWShmq98="
+    AUTH_COOKIE_NAME: str = "next-auth.session-token"
+
     POSTGRES_HOST: str = "localhost"
     POSTGRES_PORT: str = "5432"
     POSTGRES_USER: str = "postgres"

backend/requirements.txt

@@ -1,4 +1,5 @@
 alembic==1.8.1
+cryptography==39.0.0
 fastapi==0.88.0
 fastapi-limiter==0.1.5
 fastapi-utils==0.2.1
@@ -6,7 +7,10 @@ loguru==0.6.0
 numpy==1.22.4
 psycopg2-binary==2.9.5
 pydantic==1.9.1
+pydantic[email]==1.9.1
 python-dotenv==0.21.0
+python-jose[cryptography]==3.3.0
+redis
 scipy==1.8.1
 SQLAlchemy==1.4.41
 sqlmodel==0.0.8