feat: token exchange as part of VSC login (#1240)

* create auth flow for token exchange

* clean up create_user_and_tokens

* update example env file

* schedule deletion job

* rename CodeExchange to TokenExchange

* fix formatting

* fix credo issues

* make "code" PK in token_exchange table

* Rename migration file to maintain total order

Author: heyzec

config/config.exs

@@ -25,7 +25,9 @@ config :cadet, Cadet.Jobs.Scheduler,
     # Compute rolling leaderboard every 2 hours
     {"0 */2 * * *", {Cadet.Assessments, :update_rolling_contest_leaderboards, []}},
     # Collate contest entries that close in the previous day at 00:01
-    {"1 0 * * *", {Cadet.Assessments, :update_final_contest_entries, []}}
+    {"1 0 * * *", {Cadet.Assessments, :update_final_contest_entries, []}},
+    # Clean up expired exchange tokens at 00:01
+    {"1 0 * * *", {Cadet.TokenExchange, :delete_expired, []}}
   ]
 
 # Configures the endpoint

config/dev.secrets.exs.example

@@ -36,6 +36,8 @@ config :cadet,
     #    %{
     #      assertion_extractor: Cadet.Auth.Providers.NusstuAssertionExtractor,
     #      client_redirect_url: "http://cadet.frontend:8000/login/callback"
+    #      vscode_redirect_url_prefix: "vscode://source-academy.source-academy/sso",
+    #      client_post_exchange_redirect_url: "http://cadet.frontend:8000/login/vscode_callback",
     #    }},
 
     "test" =>

lib/cadet/code_exchange.ex

@@ -0,0 +1,60 @@
+defmodule Cadet.TokenExchange do
+  @moduledoc """
+  The TokenExchange entity stores short-lived codes to be exchanged for long-lived auth tokens.
+  """
+  use Cadet, :model
+
+  import Ecto.Query
+
+  alias Cadet.Repo
+  alias Cadet.Accounts.User
+
+  @primary_key {:code, :string, []}
+  schema "token_exchange" do
+    field(:generated_at, :utc_datetime_usec)
+    field(:expires_at, :utc_datetime_usec)
+
+    belongs_to(:user, User)
+
+    timestamps()
+  end
+
+  @required_fields ~w(code generated_at expires_at user_id)a
+
+  def get_by_code(code) do
+    case Repo.get_by(__MODULE__, code: code) do
+      nil ->
+        {:error, "Not found"}
+
+      struct ->
+        if Timex.before?(struct.expires_at, Timex.now()) do
+          {:error, "Expired"}
+        else
+          struct = Repo.preload(struct, :user)
+          Repo.delete(struct)
+          {:ok, struct}
+        end
+    end
+  end
+
+  def delete_expired do
+    now = Timex.now()
+
+    Repo.delete_all(from(c in __MODULE__, where: c.expires_at < ^now))
+  end
+
+  def changeset(struct, attrs) do
+    struct
+    |> cast(attrs, @required_fields)
+    |> validate_required(@required_fields)
+  end
+
+  def insert(attrs) do
+    changeset =
+      %__MODULE__{}
+      |> changeset(attrs)
+
+    changeset
+    |> Repo.insert()
+  end
+end

lib/cadet_web/controllers/auth_controller.ex

@@ -5,9 +5,9 @@ defmodule CadetWeb.AuthController do
   use CadetWeb, :controller
   use PhoenixSwagger
 
-  alias Cadet.Accounts
-  alias Cadet.Accounts.User
+  alias Cadet.{Accounts, Accounts.User}
   alias Cadet.Auth.{Guardian, Provider}
+  alias Cadet.TokenExchange
 
   @doc """
   Receives a /login request with valid attributes.
@@ -85,9 +85,87 @@ defmodule CadetWeb.AuthController do
     send_resp(conn, :bad_request, "Missing parameter")
   end
 
-  @spec create_user_and_tokens(Provider.authorise_params()) ::
-          {:ok, %{access_token: String.t(), refresh_token: String.t()}} | Plug.Conn.t()
-  defp create_user_and_tokens(
+  @doc """
+  Exchanges a short-lived code for access and refresh tokens.
+  """
+  def exchange(
+        conn,
+        %{
+          "code" => code,
+          "provider" => provider
+        }
+      ) do
+    case TokenExchange.get_by_code(code) do
+      {:error, _message} ->
+        conn
+        |> put_status(:forbidden)
+        |> text("Invalid code")
+
+      {:ok, struct} ->
+        tokens = generate_tokens(struct.user)
+
+        {_provider, %{client_post_exchange_redirect_url: client_post_exchange_redirect_url}} =
+          Application.get_env(:cadet, :identity_providers, %{})[provider]
+
+        conn
+        |> put_resp_header(
+          "location",
+          URI.encode(
+            client_post_exchange_redirect_url <>
+              "?access_token=" <> tokens.access_token <> "&refresh_token=" <> tokens.refresh_token
+          )
+        )
+        |> send_resp(302, "")
+        |> halt()
+    end
+  end
+
+  @doc """
+  Alternate callback URL which redirect to VSCode via deeplinking.
+  """
+  def saml_redirect_vscode(
+        conn,
+        %{
+          "provider" => provider
+        }
+      ) do
+    code_ttl = 60
+
+    case create_user(%{
+           conn: conn,
+           provider_instance: provider,
+           code: nil,
+           client_id: nil,
+           redirect_uri: nil
+         }) do
+      {:ok, user} ->
+        code = generate_code()
+
+        TokenExchange.insert(%{
+          code: code,
+          generated_at: Timex.now(),
+          expires_at: Timex.add(Timex.now(), Timex.Duration.from_seconds(code_ttl)),
+          user_id: user.id
+        })
+
+        {_provider, %{vscode_redirect_url_prefix: vscode_redirect_url_prefix}} =
+          Application.get_env(:cadet, :identity_providers, %{})[provider]
+
+        conn
+        |> put_resp_header(
+          "location",
+          vscode_redirect_url_prefix <> "?provider=" <> provider <> "&code=" <> code
+        )
+        |> send_resp(302, "")
+        |> halt()
+
+      conn ->
+        conn
+    end
+  end
+
+  @spec create_user(Provider.authorise_params()) :: {:ok, User.t()} | Plug.Conn.t()
+  defp create_user(
          params = %{
            conn: conn,
            provider_instance: provider
@@ -96,7 +174,7 @@ defmodule CadetWeb.AuthController do
     with {:authorise, {:ok, %{token: token, username: username}}} <-
            {:authorise, Provider.authorise(params)},
          {:signin, {:ok, user}} <- {:signin, Accounts.sign_in(username, token, provider)} do
-      {:ok, generate_tokens(user)}
+      {:ok, user}
     else
       {:authorise, {:error, :upstream, reason}} ->
         conn
@@ -121,6 +199,18 @@ defmodule CadetWeb.AuthController do
     end
   end
 
+  @spec create_user_and_tokens(Provider.authorise_params()) ::
+          {:ok, %{access_token: String.t(), refresh_token: String.t()}} | Plug.Conn.t()
+  defp create_user_and_tokens(params) do
+    case create_user(params) do
+      {:ok, user} ->
+        {:ok, generate_tokens(user)}
+
+      conn ->
+        conn
+    end
+  end
+
   @doc """
   Receives a /refresh request with valid attribute.
 
@@ -170,6 +260,14 @@ defmodule CadetWeb.AuthController do
     %{access_token: access_token, refresh_token: refresh_token}
   end
 
+  @spec generate_code :: String.t()
+  defp generate_code do
+    16
+    |> :crypto.strong_rand_bytes()
+    |> Base.url_encode64(padding: false)
+    |> String.slice(0, 22)
+  end
+
   swagger_path :create do
     post("/auth/login")
 

lib/cadet_web/router.ex

@@ -51,6 +51,8 @@ defmodule CadetWeb.Router do
     post("/auth/login", AuthController, :create)
     post("/auth/logout", AuthController, :logout)
     get("/auth/saml_redirect", AuthController, :saml_redirect)
+    get("/auth/saml_redirect_vscode", AuthController, :saml_redirect_vscode)
+    get("/auth/exchange", AuthController, :exchange)
   end
 
   scope "/v2", CadetWeb do

priv/repo/migrations/20250417093922_create_token_exchange_table.exs

@@ -0,0 +1,13 @@
+defmodule Cadet.Repo.Migrations.CreateTokenExchangeTable do
+  use Ecto.Migration
+
+  def change do
+    create table(:token_exchange) do
+      add(:code, :string, null: false)
+      add(:generated_at, :utc_datetime_usec, null: false)
+      add(:expires_at, :utc_datetime_usec, null: false)
+      add(:user_id, references(:users), null: false)
+      timestamps()
+    end
+  end
+end