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

website/.env

@@ -7,6 +7,9 @@ DATABASE_URL=postgres://postgres:postgres@localhost:5433/oasst_web
 FASTAPI_URL=http://localhost:8080
 FASTAPI_KEY=1234
 
+# Used to expose the backend url to the clientside javascript
+NEXT_PUBLIC_BACKEND_URL=$FASTAPI_URL
+
 # A dev Auth Secret.  Can be exposed if we never use this publicly.
 NEXTAUTH_SECRET=O/M2uIbGj+lDD2oyNa8ax4jEOJqCPJzO53UbWShmq98=
 

website/src/lib/api.ts

@@ -6,8 +6,11 @@ const headers = {
   "Content-Type": "application/json",
 };
 
+// Create Axios such that we always send credential cookies along with the
+// request.  This allows the Backend services to authenticate the user.
 const api = axios.create({
   headers,
+  withCredentials: true,
 });
 
 export const get = (url: string) => api.get(url).then((res) => res.data);

website/src/pages/api/auth/[...nextauth].ts

@@ -148,6 +148,21 @@ export const authOptions: AuthOptions = {
       }
     },
   },
+  /*
+   * We maybe need this, we maybe don't.  Checking in this uncommented until
+   * it's confirmed we can drop this.
+  cookies: {
+    sessionToken: {
+      name: `next-auth.session-token`,
+      options: {
+        httpOnly: true,
+        sameSite: "none",
+        path: "/",
+        secure: true,
+      },
+    },
+  },
+  */
   session: {
     strategy: "jwt",
   },

website/src/pages/dashboard.tsx

@@ -11,6 +11,9 @@ export { getDefaultStaticProps as getStaticProps } from "src/lib/default_static_
 import useSWR from "swr";
 
 const Dashboard = () => {
+  // Adding a demonstrative call to the backend that includes the web's JWT.
+  useSWR(`${process.env.NEXT_PUBLIC_BACKEND_URL}/api/v1/auth/check`, get);
+
   const {
     t,
     i18n: { language },