diff --git a/src/discord-cluster-manager/api/main.py b/src/discord-cluster-manager/api/main.py index 3917da07..c26264cd 100644 --- a/src/discord-cluster-manager/api/main.py +++ b/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 diff --git a/src/discord-cluster-manager/env.py b/src/discord-cluster-manager/env.py index 2fcac11e..94595e17 100644 --- a/src/discord-cluster-manager/env.py +++ b/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") diff --git a/src/discord-cluster-manager/leaderboard_db.py b/src/discord-cluster-manager/leaderboard_db.py index 83b3f78b..56abde60 100644 --- a/src/discord-cluster-manager/leaderboard_db.py +++ b/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( diff --git a/src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py b/src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py index a5e9ef26..955f3b56 100644 --- a/src/discord-cluster-manager/migrations/20241226_01_ZQSOK-add_gpu_type_to_submission.py +++ b/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"), ] diff --git a/src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py b/src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py index ab1d4209..cde9d60e 100644 --- a/src/discord-cluster-manager/migrations/20250221_01_GA8ro-submission-collection.py +++ b/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 diff --git a/src/discord-cluster-manager/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py b/src/discord-cluster-manager/migrations/20250406_01_ZXjWK-user-info-add-cli-id.py new file mode 100644 index 00000000..b186db86 --- /dev/null +++ b/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 + ) +]