From 10c6af7e49563cda01bd562f4f7b88bd757506b5 Mon Sep 17 00:00:00 2001 From: shishu314 Date: Thu, 28 Aug 2025 17:45:47 -0400 Subject: [PATCH] Fix(trust) - Disable commands from untrusted directories when useFolderTrust is enabled (#7341) Co-authored-by: Richie Foreman Co-authored-by: Shi Shu --- .../src/services/FileCommandLoader.test.ts | 58 +++++++++++++++++++ .../cli/src/services/FileCommandLoader.ts | 8 +++ packages/cli/src/ui/App.test.tsx | 2 + 3 files changed, 68 insertions(+) diff --git a/packages/cli/src/services/FileCommandLoader.test.ts b/packages/cli/src/services/FileCommandLoader.test.ts index 37093a3654..5b4c91cd93 100644 --- a/packages/cli/src/services/FileCommandLoader.test.ts +++ b/packages/cli/src/services/FileCommandLoader.test.ts @@ -224,6 +224,8 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => '/path/to/project'), getExtensions: vi.fn(() => []), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -267,6 +269,8 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => process.cwd()), getExtensions: vi.fn(() => []), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -556,6 +560,8 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -607,6 +613,8 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -714,6 +722,8 @@ describe('FileCommandLoader', () => { path: extensionDir2, }, ]), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -750,6 +760,8 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -782,6 +794,8 @@ describe('FileCommandLoader', () => { getExtensions: vi.fn(() => [ { name: 'a', version: '1.0.0', isActive: true, path: extensionDir }, ]), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -1169,4 +1183,48 @@ describe('FileCommandLoader', () => { } }); }); + + describe('with folder trust enabled', () => { + it('loads multiple commands', async () => { + const mockConfig = { + getProjectRoot: vi.fn(() => '/path/to/project'), + getExtensions: vi.fn(() => []), + getFolderTrustFeature: vi.fn(() => true), + getFolderTrust: vi.fn(() => true), + } as unknown as Config; + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'test1.toml': 'prompt = "Prompt 1"', + 'test2.toml': 'prompt = "Prompt 2"', + }, + }); + + const loader = new FileCommandLoader(mockConfig); + const commands = await loader.loadCommands(signal); + + expect(commands).toHaveLength(2); + }); + + it('does not load when folder is not trusted', async () => { + const mockConfig = { + getProjectRoot: vi.fn(() => '/path/to/project'), + getExtensions: vi.fn(() => []), + getFolderTrustFeature: vi.fn(() => true), + getFolderTrust: vi.fn(() => false), + } as unknown as Config; + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'test1.toml': 'prompt = "Prompt 1"', + 'test2.toml': 'prompt = "Prompt 2"', + }, + }); + + const loader = new FileCommandLoader(mockConfig); + const commands = await loader.loadCommands(signal); + + expect(commands).toHaveLength(0); + }); + }); }); diff --git a/packages/cli/src/services/FileCommandLoader.ts b/packages/cli/src/services/FileCommandLoader.ts index 8b06d0e75b..a78c372aed 100644 --- a/packages/cli/src/services/FileCommandLoader.ts +++ b/packages/cli/src/services/FileCommandLoader.ts @@ -63,8 +63,12 @@ const TomlCommandDefSchema = z.object({ */ export class FileCommandLoader implements ICommandLoader { private readonly projectRoot: string; + private readonly folderTrustEnabled: boolean; + private readonly folderTrust: boolean; constructor(private readonly config: Config | null) { + this.folderTrustEnabled = !!config?.getFolderTrustFeature(); + this.folderTrust = !!config?.getFolderTrust(); this.projectRoot = config?.getProjectRoot() || process.cwd(); } @@ -97,6 +101,10 @@ export class FileCommandLoader implements ICommandLoader { cwd: dirInfo.path, }); + if (this.folderTrustEnabled && !this.folderTrust) { + return []; + } + const commandPromises = files.map((file) => this.parseAndAdaptFile( path.join(dirInfo.path, file), diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index adf3aaab3b..65b5994878 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -177,6 +177,8 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { })), isTrustedFolder: vi.fn(() => true), getScreenReader: vi.fn(() => false), + getFolderTrustFeature: vi.fn(() => false), + getFolderTrust: vi.fn(() => false), }; });