From 6f00252211aec0fda38dd01bfad685fc0af1b52d Mon Sep 17 00:00:00 2001
From: Chenlei Hu <hcl@comfy.org>
Date: Thu, 14 Nov 2024 20:17:47 -0500
Subject: [PATCH] Disallow selecting existing ComfyUI dir as install path
 (#256)

* Disallow selecting existing ComfyUI dir as install path

* Fix tests

* nit
---
 src/config/comfyConfigManager.ts      | 11 +++--------
 src/handlers/pathHandlers.ts          | 15 +++++++++++----
 src/main.ts                           |  3 ++-
 tests/unit/comfyConfigManager.test.ts | 25 +++++--------------------
 4 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/src/config/comfyConfigManager.ts b/src/config/comfyConfigManager.ts
index 0bbb0ba2..6198b17c 100644
--- a/src/config/comfyConfigManager.ts
+++ b/src/config/comfyConfigManager.ts
@@ -54,18 +54,13 @@ export class ComfyConfigManager {
     'Comfy.Workflow.ShowMissingModelsWarning': true,
   };
 
-  public static setUpComfyUI(localComfyDirectory: string): string {
-    if (!this.isComfyUIDirectory(localComfyDirectory)) {
-      log.info(
-        `Selected directory ${localComfyDirectory} is not a ComfyUI directory. Appending ComfyUI to install path.`
-      );
-      localComfyDirectory = path.join(localComfyDirectory, 'ComfyUI');
+  public static setUpComfyUI(localComfyDirectory: string) {
+    if (fs.existsSync(localComfyDirectory)) {
+      throw new Error(`Selected directory ${localComfyDirectory} already exists`);
     }
-
     this.createComfyDirectories(localComfyDirectory);
     const userSettingsPath = path.join(localComfyDirectory, 'user', 'default');
     this.createComfyConfigFile(userSettingsPath, true);
-    return localComfyDirectory;
   }
 
   public static createComfyConfigFile(userSettingsPath: string, overwrite: boolean = false): void {
diff --git a/src/handlers/pathHandlers.ts b/src/handlers/pathHandlers.ts
index 57272812..66591efe 100644
--- a/src/handlers/pathHandlers.ts
+++ b/src/handlers/pathHandlers.ts
@@ -7,6 +7,7 @@ import type { SystemPaths } from '../preload';
 import fs from 'fs';
 import si from 'systeminformation';
 import { ComfyConfigManager } from '../config/comfyConfigManager';
+import path from 'path';
 
 export class PathHandlers {
   static readonly REQUIRED_SPACE = 10 * 1024 * 1024 * 1024; // 10GB in bytes
@@ -45,23 +46,29 @@ export class PathHandlers {
      */
     ipcMain.handle(
       IPC_CHANNELS.VALIDATE_INSTALL_PATH,
-      async (event, path: string): Promise<{ isValid: boolean; error?: string }> => {
+      async (event, inputPath: string): Promise<{ isValid: boolean; error?: string }> => {
         try {
           // Check if path exists
-          if (!fs.existsSync(path)) {
+          if (!fs.existsSync(inputPath)) {
             return { isValid: false, error: 'Path does not exist' };
           }
 
+          // Check if `path/ComfyUI` exists
+          // We are going to create a ComfyUI directory in the selected path
+          if (fs.existsSync(path.join(inputPath, 'ComfyUI'))) {
+            return { isValid: false, error: 'Path already contains ComfyUI/' };
+          }
+
           // Check if path is writable
           try {
-            fs.accessSync(path, fs.constants.W_OK);
+            fs.accessSync(inputPath, fs.constants.W_OK);
           } catch (err) {
             return { isValid: false, error: 'Path is not writable' };
           }
 
           // Check available disk space (require at least 10GB free)
           const disks = await si.fsSize();
-          const disk = disks.find((disk) => path.startsWith(disk.mount));
+          const disk = disks.find((disk) => inputPath.startsWith(disk.mount));
           if (disk && disk.available < PathHandlers.REQUIRED_SPACE) {
             return {
               isValid: false,
diff --git a/src/main.ts b/src/main.ts
index f5caf629..328390ca 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -473,7 +473,8 @@ async function handleInstall(installOptions: InstallOptions) {
   const migrationSource = installOptions.migrationSourcePath;
   const migrationItemIds = new Set<string>(installOptions.migrationItemIds ?? []);
 
-  const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(installOptions.installPath);
+  const actualComfyDirectory = path.join(installOptions.installPath, 'ComfyUI');
+  ComfyConfigManager.setUpComfyUI(actualComfyDirectory);
 
   const { comfyui: comfyuiConfig, ...extraConfigs } = await ComfyServerConfig.getMigrationConfig(
     migrationSource,
diff --git a/tests/unit/comfyConfigManager.test.ts b/tests/unit/comfyConfigManager.test.ts
index 2532abc4..838ca6e7 100644
--- a/tests/unit/comfyConfigManager.test.ts
+++ b/tests/unit/comfyConfigManager.test.ts
@@ -20,35 +20,20 @@ describe('ComfyConfigManager', () => {
   });
 
   describe('setUpComfyUI', () => {
-    it('should use existing directory when it contains ComfyUI structure', () => {
-      // Mock isComfyUIDirectory to return true for the input path
-      (fs.existsSync as jest.Mock).mockImplementation((path: string) => {
-        const requiredDirs = [
-          '/existing/ComfyUI/models',
-          '/existing/ComfyUI/input',
-          '/existing/ComfyUI/user',
-          '/existing/ComfyUI/output',
-          '/existing/ComfyUI/custom_nodes',
-        ];
-        return requiredDirs.includes(path);
-      });
-
-      const result = ComfyConfigManager.setUpComfyUI('/existing/ComfyUI');
-
-      expect(result).toBe('/existing/ComfyUI');
+    it('should reject existing directory when it contains ComfyUI structure', () => {
+      (fs.existsSync as jest.Mock).mockReturnValue(true);
+      expect(() => ComfyConfigManager.setUpComfyUI('/existing/ComfyUI')).toThrow();
     });
 
     it('should create ComfyUI subdirectory when it is missing', () => {
       (fs.existsSync as jest.Mock).mockImplementationOnce((path: string) => {
-        if (path === '/some/base/path/ComfyUI') {
+        if (['/some/base/path/ComfyUI'].includes(path)) {
           return false;
         }
         return true;
       });
 
-      const result = ComfyConfigManager.setUpComfyUI('/some/base/path');
-
-      expect(result).toBe('/some/base/path/ComfyUI');
+      ComfyConfigManager.setUpComfyUI('/some/base/path/ComfyUI');
     });
   });