Feat: Authenticate from CLI (#230)

* Feat: start work on the auth flow

* Feat: make authentication work

* Fix: decode before creating in db

* Fix: typo

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix: deduplicate state checks

* Fix: correct exception types

* Fix: new migration to keep linear history

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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,80 @@ 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")
+    user_name = user_json.get("username")
+
+    try:
+        cli_id = base64.b64decode(state).decode("utf-8")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid state parameter") from None
+
+    with bot_instance.leaderboard_db as db:
+        try:
+            db.create_user_from_cli(user_id, user_name, cli_id)
+        except Exception:
+            raise HTTPException(status_code=400, detail="Failed to create user") from None
+
+    return {"status": "success", "user_id": user_id, "cli_id": cli_id, "user_name": user_name}
+
+
 @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 these 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/leaderboard_db.py

@@ -744,6 +744,26 @@ def get_leaderboard_submission_count(
         self.cursor.execute(query, args)
         return self.cursor.fetchone()[0]
 
+    def create_user_from_cli(self, user_id: str, user_name: str, cli_id: str):
+        """
+        Method to create a user from the CLI. Shouldn't be used for Discord.
+        """
+        try:
+            self.cursor.execute(
+                """
+                INSERT INTO leaderboard.user_info (id, user_name, cli_id)
+                VALUES (%s, %s, %s)
+                ON CONFLICT (id) DO UPDATE
+                SET user_name = %s, cli_id = %s
+                """,
+                (user_id, user_name, cli_id, user_name, cli_id),
+            )
+            self.connection.commit()
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            logger.exception("Could not create/update user %s from CLI.", user_id, exc_info=e)
+            raise e
+
 
 if __name__ == "__main__":
     print(

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

src/discord-cluster-manager/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py

@@ -0,0 +1,13 @@
+"""
+user-info-add-cli-id
+"""
+
+from yoyo import step
+
+__depends__ = {"20250329_01_7VjJJ-add-a-secret-seed-column"}
+
+steps = [
+    step(
+        "ALTER TABLE leaderboard.user_info ADD COLUMN IF NOT EXISTS cli_id VARCHAR(255) DEFAULT NULL;"  # noqa: E501
+    )
+]