From a7398d235bcf530d9e84f2d391edd33262301698 Mon Sep 17 00:00:00 2001
From: heyzec <61238538+heyzec@users.noreply.github.com>
Date: Tue, 1 Jul 2025 23:52:41 +0800
Subject: [PATCH] feat: SAML

---
 package.json                 |  5 +++++
 src/commands/index.tsx       |  5 ++++-
 src/commands/showPanel.tsx   |  6 ++++--
 src/extension.ts             | 24 ++++++++++++++++++------
 src/utils/config/schema.ts   |  5 +++++
 src/utils/messageHandler.tsx |  8 +++++---
 6 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/package.json b/package.json
index 3fff572..654e661 100644
--- a/package.json
+++ b/package.json
@@ -42,6 +42,11 @@
           "default": "https://sourceacademy.nus.edu.sg",
           "description": "URL to the Source Academy frontend"
         },
+        "source-academy.backendBaseUrl": {
+          "type": "string",
+          "default": "https://api.sourceacademy.nus.edu.sg",
+          "description": "URL to the Source Academy backend (only needed for SAML logins)"
+        },
         "source-academy.workspaceFolder": {
           "type": "string",
           "default": ".sourceacademy",
diff --git a/src/commands/index.tsx b/src/commands/index.tsx
index e087a43..6a8bbc5 100644
--- a/src/commands/index.tsx
+++ b/src/commands/index.tsx
@@ -3,6 +3,7 @@ import { runLanguagePicker } from "./language";
 import { evalEditor } from "./evalEditor";
 import { showPanel } from "./showPanel";
 import { navigate } from "./navigate";
+import { login } from "./login";
 
 const EXTENSION_ID = "source-academy";
 
@@ -12,9 +13,11 @@ const EXTENSION_ID = "source-academy";
  */
 const commands = (context: vscode.ExtensionContext) => ({
   pick: () => runLanguagePicker(context),
-  "show-panel": (route?: string) => showPanel(context, route),
+  "show-panel": (route?: string, altUrl?: string) =>
+    showPanel(context, route, altUrl),
   "eval-editor": () => evalEditor(context),
   navigate: (route: string) => navigate(context, route),
+  login: () => login(context),
 });
 
 /**
diff --git a/src/commands/showPanel.tsx b/src/commands/showPanel.tsx
index 02e9311..1742d3e 100644
--- a/src/commands/showPanel.tsx
+++ b/src/commands/showPanel.tsx
@@ -16,6 +16,7 @@ let messageHandler = MessageHandler.getInstance();
 export async function showPanel(
   context: vscode.ExtensionContext,
   route?: string,
+  altUrl?: string,
 ) {
   let language: string | undefined = context.workspaceState.get("language");
   if (!language) {
@@ -41,9 +42,10 @@ export async function showPanel(
     );
   }
 
-  const iframeUrl = new URL(route ?? "/playground", config.frontendBaseUrl)
-    .href;
+  const iframeUrl =
+    altUrl ?? new URL(route ?? "/playground", config.frontendBaseUrl).href;
 
+  // equivalent to panel.webview.html = ...
   setWebviewContent(
     messageHandler.panel,
     context,
diff --git a/src/extension.ts b/src/extension.ts
index 1b0f2a9..2d94297 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -55,14 +55,26 @@ export function activate(context: vscode.ExtensionContext) {
   vscode.window.registerUriHandler({
     handleUri(uri: vscode.Uri) {
       const searchParams = new URLSearchParams(uri.query);
-      // The following two params are available when logging in via OAuth providers
+
       const code = searchParams.get("code");
-      const clientRequestId = searchParams.get("client-request-id");
+      // The following params are available conditionally based on provider
+      const clientRequestId = searchParams.get("client-request-id"); // OAuth (NUS's IdP)
+      const provider = searchParams.get("provider"); // SAML
+
+      let route, url;
+      if (clientRequestId) {
+        // OAuth
+        route = `/login/vscode_callback?code=${code}&client-request-id=${clientRequestId}`;
+        url = null;
+      } else if (provider) {
+        // SAML
+        route = null;
+        // TODO: Let the frontend handle the contacting of backend, instead of us.
+        // Then, remove backendBaseUrl from schema and altUrl from showPanel
+        url = `${config.backendBaseUrl}/v2/auth/exchange?code=${code}&provider=${provider}`;
+      }
 
-      vscode.commands.executeCommand(
-        "source-academy.show-panel",
-        `/login/vscode_callback?code=${code}&client-request-id=${clientRequestId}`,
-      );
+      vscode.commands.executeCommand("source-academy.show-panel", route, url);
     },
   });
 }
diff --git a/src/utils/config/schema.ts b/src/utils/config/schema.ts
index a34157a..398aaeb 100644
--- a/src/utils/config/schema.ts
+++ b/src/utils/config/schema.ts
@@ -5,6 +5,11 @@ export default {
     type: "string",
     default: "https://sourceacademy.nus.edu.sg",
   },
+  // TODO: Remove this config in the future. See the login URI handler for details.
+  backendBaseUrl: {
+    type: "string",
+    default: "https://api.sourceacademy.nus.edu.sg",
+  },
   workspaceFolder: {
     type: "string",
     default: ".sourceacademy",
diff --git a/src/utils/messageHandler.tsx b/src/utils/messageHandler.tsx
index 3dd25a9..57b8834 100644
--- a/src/utils/messageHandler.tsx
+++ b/src/utils/messageHandler.tsx
@@ -173,9 +173,11 @@ export class MessageHandler {
             this.panel?.reveal(vscode.ViewColumn.Two);
           }
           break;
-      case MessageTypeNames.LoginWithBrowser:
-        const { route } = message;
-        vscode.env.openExternal(vscode.Uri.parse(route));
+        case MessageTypeNames.LoginWithBrowser:
+          let { route } = message;
+          // TODO: Remove this hack! This should be changed in the frontend
+          route = route.replace("saml_redirect", "saml_redirect_vscode");
+          vscode.env.openExternal(vscode.Uri.parse(route));
       }
       console.log(`${Date.now()} Finish handleMessage: ${message.type}`);
     }