From bbf33a4f60f12a56aaaa8b33d22b21700685a2cb Mon Sep 17 00:00:00 2001
From: James Brunton <james@stirlingpdf.com>
Date: Tue, 16 Dec 2025 14:27:39 +0000
Subject: [PATCH] Initial commit of self-hosted SSO

---
 .../SetupWizard/DesktopOAuthButtons.tsx       | 59 ++++++++++++--
 .../SetupWizard/SaaSLoginScreen.tsx           |  5 +-
 .../SetupWizard/SelfHostedLoginScreen.tsx     | 79 ++++++++++++++++++-
 .../SetupWizard/ServerSelection.tsx           | 21 +++--
 frontend/src/desktop/services/authService.ts  | 18 +++++
 .../desktop/services/connectionModeService.ts |  6 +-
 .../src/proprietary/routes/AuthCallback.tsx   | 60 +++++++++++---
 7 files changed, 217 insertions(+), 31 deletions(-)

diff --git a/frontend/src/desktop/components/SetupWizard/DesktopOAuthButtons.tsx b/frontend/src/desktop/components/SetupWizard/DesktopOAuthButtons.tsx
index 49b55a3869..d26711df47 100644
--- a/frontend/src/desktop/components/SetupWizard/DesktopOAuthButtons.tsx
+++ b/frontend/src/desktop/components/SetupWizard/DesktopOAuthButtons.tsx
@@ -7,12 +7,24 @@ import '@app/routes/authShared/auth.css';
 
 export type OAuthProvider = 'google' | 'github' | 'keycloak' | 'azure' | 'apple' | 'oidc';
 
+export interface DesktopOAuthProvider {
+  id: string;
+  label?: string;
+  file?: string;
+  url?: string;
+}
+
 interface DesktopOAuthButtonsProps {
   onOAuthSuccess: (userInfo: UserInfo) => Promise<void>;
   onError: (error: string) => void;
   isDisabled: boolean;
   serverUrl: string;
-  providers: OAuthProvider[];
+  providers: DesktopOAuthProvider[];
+  /**
+   * Optional override to handle provider click (used for self-hosted SSO flow).
+   * When provided, the caller is responsible for invoking onOAuthSuccess/onError.
+   */
+  onProviderClick?: (provider: DesktopOAuthProvider) => Promise<void>;
 }
 
 export const DesktopOAuthButtons: React.FC<DesktopOAuthButtonsProps> = ({
@@ -21,6 +33,7 @@ export const DesktopOAuthButtons: React.FC<DesktopOAuthButtonsProps> = ({
   isDisabled,
   serverUrl,
   providers,
+  onProviderClick,
 }) => {
   const { t } = useTranslation();
   const [oauthLoading, setOauthLoading] = useState(false);
@@ -64,7 +77,25 @@ export const DesktopOAuthButtons: React.FC<DesktopOAuthButtonsProps> = ({
     }
   };
 
-  const providerConfig: Record<OAuthProvider, { label: string; file: string }> = {
+  const handleCustomClick = async (provider: DesktopOAuthProvider) => {
+    if (!onProviderClick) {
+      return;
+    }
+
+    try {
+      setOauthLoading(true);
+      await onProviderClick(provider);
+    } catch (error) {
+      const errorMessage = error instanceof Error
+        ? error.message
+        : t('setup.login.error.oauthFailed', 'OAuth login failed. Please try again.');
+      onError(errorMessage);
+    } finally {
+      setOauthLoading(false);
+    }
+  };
+
+  const providerConfig: Record<string, { label: string; file: string }> = {
     google: { label: 'Google', file: 'google.svg' },
     github: { label: 'GitHub', file: 'github.svg' },
     keycloak: { label: 'Keycloak', file: 'keycloak.svg' },
@@ -80,19 +111,31 @@ export const DesktopOAuthButtons: React.FC<DesktopOAuthButtonsProps> = ({
   return (
     <div className="oauth-container-vertical">
       {providers
-        .filter((providerId) => providerId in providerConfig)
-        .map((providerId) => {
-          const provider = providerConfig[providerId];
+        .map((providerMeta) => {
+          const provider =
+            (providerMeta.id in providerConfig
+              ? providerConfig[providerMeta.id as OAuthProvider]
+              : {
+                  label: providerMeta.label || providerMeta.id,
+                  file: providerMeta.file || 'oidc.svg',
+                });
+
           return (
             <button
-              key={providerId}
-              onClick={() => handleOAuthLogin(providerId)}
+              key={providerMeta.id}
+              onClick={() => {
+                if (onProviderClick) {
+                  void handleCustomClick(providerMeta);
+                } else {
+                  void handleOAuthLogin(providerMeta.id as OAuthProvider);
+                }
+              }}
               disabled={isDisabled || oauthLoading}
               className="oauth-button-vertical"
               title={provider.label}
             >
               <img
-                src={`${BASE_PATH}/Login/${provider.file}`}
+                src={`${BASE_PATH}/Login/${provider.file || 'oidc.svg'}`}
                 alt={provider.label}
                 className="oauth-icon-tiny"
               />
diff --git a/frontend/src/desktop/components/SetupWizard/SaaSLoginScreen.tsx b/frontend/src/desktop/components/SetupWizard/SaaSLoginScreen.tsx
index ab82bde1ed..6fb7a85498 100644
--- a/frontend/src/desktop/components/SetupWizard/SaaSLoginScreen.tsx
+++ b/frontend/src/desktop/components/SetupWizard/SaaSLoginScreen.tsx
@@ -64,7 +64,10 @@ export const SaaSLoginScreen: React.FC<SaaSLoginScreenProps> = ({
         onError={handleOAuthError}
         isDisabled={loading}
         serverUrl={serverUrl}
-        providers={['google', 'github']}
+        providers={[
+          { id: 'google' },
+          { id: 'github' },
+        ]}
       />
 
       <DividerWithText
diff --git a/frontend/src/desktop/components/SetupWizard/SelfHostedLoginScreen.tsx b/frontend/src/desktop/components/SetupWizard/SelfHostedLoginScreen.tsx
index 9a68b91561..70a44d9090 100644
--- a/frontend/src/desktop/components/SetupWizard/SelfHostedLoginScreen.tsx
+++ b/frontend/src/desktop/components/SetupWizard/SelfHostedLoginScreen.tsx
@@ -5,13 +5,13 @@ import LoginHeader from '@app/routes/login/LoginHeader';
 import ErrorMessage from '@app/routes/login/ErrorMessage';
 import EmailPasswordForm from '@app/routes/login/EmailPasswordForm';
 import DividerWithText from '@app/components/shared/DividerWithText';
-import { DesktopOAuthButtons, OAuthProvider } from '@app/components/SetupWizard/DesktopOAuthButtons';
-import { UserInfo } from '@app/services/authService';
+import { DesktopOAuthButtons, DesktopOAuthProvider } from '@app/components/SetupWizard/DesktopOAuthButtons';
+import { authService, UserInfo } from '@app/services/authService';
 import '@app/routes/authShared/auth.css';
 
 interface SelfHostedLoginScreenProps {
   serverUrl: string;
-  enabledOAuthProviders?: string[];
+  enabledOAuthProviders?: DesktopOAuthProvider[];
   onLogin: (username: string, password: string) => Promise<void>;
   onOAuthSuccess: (userInfo: UserInfo) => Promise<void>;
   loading: boolean;
@@ -51,6 +51,76 @@ export const SelfHostedLoginScreen: React.FC<SelfHostedLoginScreenProps> = ({
     setValidationError(errorMessage);
   };
 
+  const waitForSsoCompletion = (popup: Window): Promise<string> => {
+    return new Promise((resolve, reject) => {
+      const messageHandler = (event: MessageEvent) => {
+        if (event.origin !== window.location.origin || typeof event.data !== 'object' || event.data === null) {
+          return;
+        }
+
+        const { type, token, error } = event.data as { type?: string; token?: string; error?: string };
+        if (type === 'stirling-sso-success' && token) {
+          cleanup();
+          resolve(token);
+        } else if (type === 'stirling-sso-error') {
+          cleanup();
+          reject(new Error(error || 'SSO login failed'));
+        }
+      };
+
+      const interval = window.setInterval(() => {
+        if (popup.closed) {
+          cleanup();
+          reject(new Error('Login window was closed before authentication completed'));
+        }
+      }, 500);
+
+      const cleanup = () => {
+        window.clearInterval(interval);
+        window.removeEventListener('message', messageHandler);
+      };
+
+      window.addEventListener('message', messageHandler);
+    });
+  };
+
+  const handleSelfHostedOAuthLogin = async (provider: DesktopOAuthProvider) => {
+    setValidationError(null);
+
+    if (!provider.url) {
+      handleOAuthError(t('setup.login.error.configFetch', 'Failed to fetch server configuration. Please check the URL and try again.'));
+      return;
+    }
+
+    // Mark SSO flow so the callback can short-circuit verification
+    localStorage.setItem('desktop_sso_in_progress', JSON.stringify({ mode: 'selfhosted' }));
+
+    const popup = window.open(provider.url, '_blank', 'width=480,height=720');
+    if (!popup) {
+      localStorage.removeItem('desktop_sso_in_progress');
+      handleOAuthError(t('setup.login.error.oauthFailed', 'OAuth login failed. Please try again.'));
+      return;
+    }
+
+    try {
+      const token = await waitForSsoCompletion(popup);
+      localStorage.removeItem('desktop_sso_in_progress');
+
+      await authService.applyExternalToken(token);
+      await onOAuthSuccess({ username: provider.label || provider.id });
+    } catch (err) {
+      localStorage.removeItem('desktop_sso_in_progress');
+      const message = err instanceof Error ? err.message : t('setup.login.error.oauthFailed', 'OAuth login failed. Please try again.');
+      handleOAuthError(message);
+    } finally {
+      try {
+        popup.close();
+      } catch (_) {
+        // ignore
+      }
+    }
+  };
+
   const displayError = error || validationError;
 
   return (
@@ -74,7 +144,8 @@ export const SelfHostedLoginScreen: React.FC<SelfHostedLoginScreenProps> = ({
             onError={handleOAuthError}
             isDisabled={loading}
             serverUrl={serverUrl}
-            providers={enabledOAuthProviders as OAuthProvider[]}
+            providers={enabledOAuthProviders}
+            onProviderClick={handleSelfHostedOAuthLogin}
           />
 
           <DividerWithText
diff --git a/frontend/src/desktop/components/SetupWizard/ServerSelection.tsx b/frontend/src/desktop/components/SetupWizard/ServerSelection.tsx
index f8a0d4f9e1..829cc7ccbb 100644
--- a/frontend/src/desktop/components/SetupWizard/ServerSelection.tsx
+++ b/frontend/src/desktop/components/SetupWizard/ServerSelection.tsx
@@ -4,6 +4,7 @@ import { useTranslation } from 'react-i18next';
 import { ServerConfig } from '@app/services/connectionModeService';
 import { connectionModeService } from '@app/services/connectionModeService';
 import LocalIcon from '@app/components/shared/LocalIcon';
+import { DesktopOAuthProvider } from '@app/components/SetupWizard/DesktopOAuthButtons';
 
 interface ServerSelectionProps {
   onSelect: (config: ServerConfig) => void;
@@ -43,7 +44,7 @@ export const ServerSelection: React.FC<ServerSelectionProps> = ({ onSelect, load
       }
 
       // Fetch OAuth providers and check if login is enabled
-      let enabledProviders: string[] = [];
+      let enabledProviders: DesktopOAuthProvider[] = [];
       try {
         const response = await fetch(`${url}/api/v1/proprietary/ui-data/login`);
 
@@ -74,11 +75,19 @@ export const ServerSelection: React.FC<ServerSelectionProps> = ({ onSelect, load
           return;
         }
 
-        // Extract provider IDs from authorization URLs
-        // Example: "/oauth2/authorization/google" → "google"
-        enabledProviders = Object.keys(data.providerList || {})
-          .map(key => key.split('/').pop())
-          .filter((id): id is string => id !== undefined);
+        // Extract provider metadata from authorization URLs
+        // Example: "/oauth2/authorization/google" → { id: "google", url: "https://server/oauth2/authorization/google" }
+        enabledProviders = Object.entries(data.providerList || {})
+          .map(([path, label]) => {
+            const normalizedPath = path.startsWith('/') ? path : `/${path}`;
+            const id = normalizedPath.split('/').pop() || normalizedPath;
+            const trimmedUrl = url.replace(/\/$/, '');
+            return {
+              id,
+              label: typeof label === 'string' ? label : undefined,
+              url: `${trimmedUrl}${normalizedPath}`,
+            } satisfies DesktopOAuthProvider;
+          });
 
         console.log('[ServerSelection] Detected OAuth providers:', enabledProviders);
       } catch (err) {
diff --git a/frontend/src/desktop/services/authService.ts b/frontend/src/desktop/services/authService.ts
index ed8be911ff..1ddf75bb0e 100644
--- a/frontend/src/desktop/services/authService.ts
+++ b/frontend/src/desktop/services/authService.ts
@@ -69,6 +69,11 @@ export class AuthService {
     // Notify other parts of the system
     window.dispatchEvent(new CustomEvent('jwt-available'));
     console.log('[Desktop AuthService] Dispatched jwt-available event');
+
+    // Keep auth status in sync if we already have user info
+    if (this.userInfo) {
+      this.setAuthStatus('authenticated', this.userInfo);
+    }
   }
 
   /**
@@ -100,6 +105,19 @@ export class AuthService {
     return localStorageToken;
   }
 
+  /**
+   * Apply a JWT obtained from an external flow (e.g., browser SSO).
+   */
+  async applyExternalToken(token: string, userInfo?: UserInfo | null): Promise<void> {
+    await this.saveTokenEverywhere(token);
+
+    if (userInfo) {
+      this.userInfo = userInfo;
+    }
+
+    this.setAuthStatus('authenticated', userInfo ?? this.userInfo);
+  }
+
   /**
    * Clear token from all storage locations
    */
diff --git a/frontend/src/desktop/services/connectionModeService.ts b/frontend/src/desktop/services/connectionModeService.ts
index cf454a50b3..e91a526473 100644
--- a/frontend/src/desktop/services/connectionModeService.ts
+++ b/frontend/src/desktop/services/connectionModeService.ts
@@ -5,7 +5,11 @@ export type ConnectionMode = 'saas' | 'selfhosted';
 
 export interface ServerConfig {
   url: string;
-  enabledOAuthProviders?: string[];
+  enabledOAuthProviders?: Array<{
+    id: string;
+    label?: string;
+    url?: string;
+  }>;
 }
 
 export interface ConnectionConfig {
diff --git a/frontend/src/proprietary/routes/AuthCallback.tsx b/frontend/src/proprietary/routes/AuthCallback.tsx
index 488a541464..daa137b1f9 100644
--- a/frontend/src/proprietary/routes/AuthCallback.tsx
+++ b/frontend/src/proprietary/routes/AuthCallback.tsx
@@ -21,9 +21,20 @@ export default function AuthCallback() {
         const hash = window.location.hash.substring(1); // Remove '#'
         const params = new URLSearchParams(hash);
         const token = params.get('access_token');
+        const desktopSsoState = localStorage.getItem('desktop_sso_in_progress');
+        const isDesktopSso = Boolean(desktopSsoState);
 
         if (!token) {
           console.error('[AuthCallback] No access_token in URL fragment');
+          if (isDesktopSso && window.opener) {
+            localStorage.removeItem('desktop_sso_in_progress');
+            window.opener.postMessage(
+              { type: 'stirling-sso-error', error: 'OAuth login failed - no token received.' },
+              window.location.origin
+            );
+            window.close();
+            return;
+          }
           navigate('/login', {
             replace: true,
             state: { error: 'OAuth login failed - no token received.' }
@@ -38,26 +49,53 @@ export default function AuthCallback() {
         // Dispatch custom event for other components to react to JWT availability
         window.dispatchEvent(new CustomEvent('jwt-available'));
 
-        // Validate the token and load user info
-        // This calls /api/v1/auth/me with the JWT to get user details
-        const { data, error } = await springAuth.getSession();
+        // Desktop SSO flow relies on the opener to finalize setup, so skip server validation here
+        if (!isDesktopSso) {
+          // Validate the token and load user info
+          // This calls /api/v1/auth/me with the JWT to get user details
+          const { data, error } = await springAuth.getSession();
 
-        if (error || !data.session) {
-          console.error('[AuthCallback] Failed to validate token:', error);
-          localStorage.removeItem('stirling_jwt');
-          navigate('/login', {
-            replace: true,
-            state: { error: 'OAuth login failed - invalid token.' }
-          });
-          return;
+          if (error || !data.session) {
+            console.error('[AuthCallback] Failed to validate token:', error);
+            localStorage.removeItem('stirling_jwt');
+            navigate('/login', {
+              replace: true,
+              state: { error: 'OAuth login failed - invalid token.' }
+            });
+            return;
+          }
+        }
+
+        // Cleanup flag for desktop flow
+        if (isDesktopSso) {
+          localStorage.removeItem('desktop_sso_in_progress');
         }
 
         console.log('[AuthCallback] Token validated, redirecting to home');
 
+        if (isDesktopSso && window.opener) {
+          window.opener.postMessage(
+            { type: 'stirling-sso-success', token },
+            window.location.origin
+          );
+          window.close();
+          return;
+        }
+
         // Clear the hash from URL and redirect to home page
         navigate('/', { replace: true });
       } catch (error) {
         console.error('[AuthCallback] Error:', error);
+        const desktopSsoState = localStorage.getItem('desktop_sso_in_progress');
+        if (desktopSsoState && window.opener) {
+          localStorage.removeItem('desktop_sso_in_progress');
+          window.opener.postMessage(
+            { type: 'stirling-sso-error', error: 'OAuth login failed. Please try again.' },
+            window.location.origin
+          );
+          window.close();
+          return;
+        }
         navigate('/login', {
           replace: true,
           state: { error: 'OAuth login failed. Please try again.' }