From 309cbd4dc4d8dcedec317d477ed4a782f1e789a6 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Sat, 21 Jun 2025 20:37:31 -0700
Subject: [PATCH] Prevent access without login.

---
 src/router.ts      | 44 ++++++++++++++++++++++++++++++++++++++++---
 src/scripts/api.ts | 47 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 87 insertions(+), 4 deletions(-)

diff --git a/src/router.ts b/src/router.ts
index 69f1a249a..e04f10ebd 100644
--- a/src/router.ts
+++ b/src/router.ts
@@ -6,11 +6,12 @@ import {
   createWebHistory
 } from 'vue-router'
 
+import { useDialogService } from '@/services/dialogService'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
+import { useUserStore } from '@/stores/userStore'
+import { isElectron } from '@/utils/envUtil'
 import LayoutDefault from '@/views/layouts/LayoutDefault.vue'
 
-import { useUserStore } from './stores/userStore'
-import { isElectron } from './utils/envUtil'
-
 const isFileProtocol = window.location.protocol === 'file:'
 const basePath = isElectron() ? '/' : window.location.pathname
 
@@ -130,4 +131,41 @@ const router = createRouter({
   }
 })
 
+// Global authentication guard
+router.beforeEach(async (_to, _from, next) => {
+  const authStore = useFirebaseAuthStore()
+
+  // Wait for Firebase auth to initialize
+  if (!authStore.isInitialized) {
+    await new Promise<void>((resolve) => {
+      const unwatch = authStore.$subscribe((_, state) => {
+        if (state.isInitialized) {
+          unwatch()
+          resolve()
+        }
+      })
+    })
+  }
+
+  // Check if user is authenticated (Firebase or API key)
+  const authHeader = await authStore.getAuthHeader()
+
+  if (!authHeader) {
+    // User is not authenticated, show sign-in dialog
+    const dialogService = useDialogService()
+    const loginSuccess = await dialogService.showSignInDialog()
+
+    if (loginSuccess) {
+      // After successful login, proceed to the intended route
+      next()
+    } else {
+      // User cancelled login, stay on current page or redirect to home
+      next(false)
+    }
+  } else {
+    // User is authenticated, proceed
+    next()
+  }
+})
+
 export default router
diff --git a/src/scripts/api.ts b/src/scripts/api.ts
index e83eb9f8b..7502ea144 100644
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -36,6 +36,7 @@ import type {
 } from '@/schemas/comfyWorkflowSchema'
 import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
 import type { NodeExecutionId } from '@/types/nodeIdentification'
+import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
 import { WorkflowTemplates } from '@/types/workflowTemplateTypes'
 
 interface QueuePromptRequestBody {
@@ -311,7 +312,27 @@ export class ComfyApi extends EventTarget {
     return this.api_base + route
   }
 
-  fetchApi(route: string, options?: RequestInit) {
+  /**
+   * Waits for Firebase auth to be initialized before proceeding
+   */
+  async #waitForAuthInitialization(): Promise<void> {
+    const authStore = useFirebaseAuthStore()
+
+    if (authStore.isInitialized) {
+      return
+    }
+
+    return new Promise<void>((resolve) => {
+      const unwatch = authStore.$subscribe((_, state) => {
+        if (state.isInitialized) {
+          unwatch()
+          resolve()
+        }
+      })
+    })
+  }
+
+  async fetchApi(route: string, options?: RequestInit) {
     if (!options) {
       options = {}
     }
@@ -322,6 +343,30 @@ export class ComfyApi extends EventTarget {
       options.cache = 'no-cache'
     }
 
+    // Wait for Firebase auth to be initialized before making any API request
+    await this.#waitForAuthInitialization()
+
+    // Add Firebase JWT token if user is logged in
+    try {
+      const authHeader = await useFirebaseAuthStore().getAuthHeader()
+      if (authHeader) {
+        if (Array.isArray(options.headers)) {
+          for (const [key, value] of Object.entries(authHeader)) {
+            options.headers.push([key, value])
+          }
+        } else if (options.headers instanceof Headers) {
+          for (const [key, value] of Object.entries(authHeader)) {
+            options.headers.set(key, value)
+          }
+        } else {
+          Object.assign(options.headers, authHeader)
+        }
+      }
+    } catch (error) {
+      // Silently ignore auth errors to avoid breaking API calls
+      console.warn('Failed to get auth header:', error)
+    }
+
     if (Array.isArray(options.headers)) {
       options.headers.push(['Comfy-User', this.user])
     } else if (options.headers instanceof Headers) {