Feat: start work on the auth flow

Author: S1ro1

src/discord-cluster-manager/api/main.py

@@ -1,10 +1,14 @@
 import asyncio
+import base64
+import os
 import time
 from dataclasses import asdict
 
+import requests
 from cogs.submit_cog import SubmitCog
 from consts import _GPU_LOOKUP, SubmissionMode, get_gpu_by_name
 from discord import app_commands
+from env import CLI_DISCORD_CLIENT_ID, CLI_DISCORD_CLIENT_SECRET, CLI_TOKEN_URL
 from fastapi import FastAPI, HTTPException, UploadFile
 from utils import LeaderboardItem, build_task_config
 
@@ -53,6 +57,75 @@ async def update(self, message: str):
         pass
 
 
+@app.get("/auth/cli")
+async def cli_auth(code: str, state: str = None):
+    """
+    Handle Discord OAuth redirect. This endpoint receives the authorization code
+    and state parameter from Discord's OAuth flow.
+
+    Args:
+        code (str): Authorization code from Discord OAuth
+        state (str): Base64 encoded client ID from CLI
+    """
+
+    if not code or not state:
+        raise HTTPException(status_code=400, detail="Missing authorization code or state")
+
+    client_id = CLI_DISCORD_CLIENT_ID
+    client_secret = CLI_DISCORD_CLIENT_SECRET
+    redirect_uri = os.environ.get("HEROKU_APP_DEFAULT_DOMAIN_NAME") or os.getenv("POPCORN_API_URL")
+    token_url = CLI_TOKEN_URL
+
+    if not client_id or not client_secret:
+        raise HTTPException(status_code=500, detail="Discord client ID or secret not configured.")
+
+    if not token_url:
+        raise HTTPException(status_code=500, detail="Discord token URL not configured.")
+
+    if not redirect_uri:
+        raise HTTPException(
+            status_code=500,
+            detail="Redirect URI not configured. "
+            "If running locally, set env variable `POPCORN_API_URL` to your local API URL.",
+        )
+
+    token_data = {
+        "client_id": client_id,
+        "client_secret": client_secret,
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": redirect_uri + "/auth/cli",
+    }
+
+    token_response = requests.post(token_url, data=token_data)
+    if token_response.status_code != 200:
+        raise HTTPException(
+            status_code=401, detail=f"Failed to authenticate with Discord: {token_response.text}"
+        )
+
+    token_json = token_response.json()
+    access_token = token_json.get("access_token")
+
+    user_url = "https://discord.com/api/users/@me"
+    headers = {"Authorization": f"Bearer {access_token}"}
+
+    user_response = requests.get(user_url, headers=headers)
+    if user_response.status_code != 200:
+        raise HTTPException(status_code=401, detail="Failed to retrieve user information")
+
+    user_json = user_response.json()
+    user_id = user_json.get("id")
+
+    cli_id = ""
+    if state:
+        try:
+            cli_id = base64.b64decode(state).decode("utf-8")
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid state parameter: {str(e)}") from e
+
+    return {"status": "success", "user_id": user_id, "cli_id": cli_id}
+
+
 @app.post("/{leaderboard_name}/{gpu_type}/{submission_mode}")
 async def run_submission(
     leaderboard_name: str, gpu_type: str, submission_mode: str, file: UploadFile

src/discord-cluster-manager/env.py

@@ -21,6 +21,12 @@ def init_environment():
 DISCORD_CLUSTER_STAGING_ID = os.getenv("DISCORD_CLUSTER_STAGING_ID")
 DISCORD_DEBUG_CLUSTER_STAGING_ID = os.getenv("DISCORD_DEBUG_CLUSTER_STAGING_ID")
 
+# Only required to run the CLI against this instance
+# setting theses is required only to run the CLI against local instance
+CLI_DISCORD_CLIENT_ID = os.getenv("CLI_DISCORD_CLIENT_ID", "")
+CLI_DISCORD_CLIENT_SECRET = os.getenv("CLI_DISCORD_CLIENT_SECRET", "")
+CLI_TOKEN_URL = os.getenv("CLI_TOKEN_URL", "")
+
 # GitHub-specific constants
 GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
 GITHUB_REPO = os.getenv("GITHUB_REPO")

src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py

@@ -5,17 +5,14 @@
 
 from yoyo import step
 
-__depends__ = {'20241224_01_Pg4FX-delete-cascade'}
+__depends__ = {"20241224_01_Pg4FX-delete-cascade"}
 
 steps = [
     step("DROP TABLE leaderboard.runinfo"),
-
     step("""
          ALTER TABLE leaderboard.submission
          ADD COLUMN gpu_type TEXT NOT NULL DEFAULT 'nvidia'
          """),
-
     step("ALTER TABLE leaderboard.submission ADD COLUMN stdout TEXT"),
-
     step("ALTER TABLE leaderboard.submission ADD COLUMN profiler_output TEXT"),
 ]

src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py

@@ -12,7 +12,6 @@
     step("DROP TABLE IF EXISTS leaderboard.submission;"),
     step("DROP TABLE IF EXISTS leaderboard.code_files;"),
     step("DROP TABLE IF EXISTS leaderboard.runs;"),
-
     # create three new tables: One for deduplicating submitted code files,
     # one for the submission itself, and one for individual runs
     # The submission itself contains the code and the targeted leaderboard
@@ -25,7 +24,6 @@
              hash TEXT GENERATED ALWAYS AS (encode(sha256(code::bytea), 'hex')) STORED
          )
          """),
-
     step("""
          CREATE TABLE IF NOT EXISTS leaderboard.submission (
              id SERIAL PRIMARY KEY,
@@ -37,7 +35,6 @@
              done BOOLEAN DEFAULT FALSE
          )
          """),
-
     # the runs themselves contain information about a particular execution of that code.
     # This includes start and end time
     # Note that `score` can be NULL for non-ranked submissions