Restrict the file access to only current workspaces directory

Author: jkuruba

packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts

@@ -325,7 +325,7 @@ export class Connector extends BaseConnector {
 
         if (
             !this.onChatAnswerUpdated ||
-            !['accept-code-diff', 'reject-code-diff', 'confirm-tool-use'].includes(action.id)
+            !['accept-code-diff', 'reject-code-diff', 'confirm-tool-use', 'reject-tool-use'].includes(action.id)
         ) {
             return
         }
@@ -375,6 +375,18 @@ export class Connector extends BaseConnector {
                     },
                 ]
                 break
+            case 'reject-tool-use':
+                answer.buttons = [
+                    {
+                        keepCardAfterClick: true,
+                        text: 'Rejected',
+                        id: 'rejected-tool-use',
+                        status: 'error',
+                        position: 'outside',
+                        disabled: true,
+                    },
+                ]
+                break
             default:
                 break
         }

packages/core/src/codewhisperer/client/codewhisperer.ts

@@ -31,8 +31,8 @@ export interface CodeWhispererConfig {
 }
 
 export const defaultServiceConfig: CodeWhispererConfig = {
-    region: 'us-east-1',
-    endpoint: 'https://codewhisperer.us-east-1.amazonaws.com/',
+    region: 'us-west-2',
+    endpoint: 'https://rts.alpha-us-west-2.codewhisperer.ai.aws.dev/',
 }
 
 export function getCodewhispererConfig(): CodeWhispererConfig {

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -762,6 +762,7 @@ export class ChatController {
                 await this.processToolUseMessage(message)
                 break
             case 'reject-code-diff':
+            case 'reject-tool-use':
                 await this.closeDiffView()
                 break
             default:

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -45,7 +45,7 @@ import { ChatHistoryManager } from '../../../storages/chatHistory'
 import { ToolType, ToolUtils } from '../../../tools/toolUtils'
 import { ChatStream } from '../../../tools/chatStream'
 import path from 'path'
-import { CommandValidation } from '../../../tools/executeBash'
+import { CommandValidation } from '../../../tools/toolShared'
 import { Change } from 'diff'
 import { FsWriteParams } from '../../../tools/fsWrite'
 
@@ -499,6 +499,17 @@ export class Messenger {
             if (validation.warning) {
                 message = validation.warning + message
             }
+        } else if (validation.requiresAcceptance && toolUse?.name === ToolType.ListDirectory) {
+            buttons.push({
+                id: 'reject-tool-use',
+                text: 'Reject',
+                status: 'info',
+            })
+            buttons.push({
+                id: 'confirm-tool-use',
+                text: 'Confirm',
+                status: 'info',
+            })
         } else if (toolUse?.name === ToolType.FsWrite) {
             const input = toolUse.input as unknown as FsWriteParams
             const fileName = path.basename(input.path)

packages/core/src/codewhispererChat/tools/chatStream.ts

@@ -7,7 +7,7 @@ import { Writable } from 'stream'
 import { getLogger } from '../../shared/logger/logger'
 import { Messenger } from '../controllers/chat/messenger/messenger'
 import { ToolUse } from '@amzn/codewhisperer-streaming'
-import { CommandValidation } from './executeBash'
+import { CommandValidation } from './toolShared'
 import { Change } from 'diff'
 
 /**

packages/core/src/codewhispererChat/tools/executeBash.ts

@@ -7,7 +7,7 @@ import { Writable } from 'stream'
 import { getLogger } from '../../shared/logger/logger'
 import { fs } from '../../shared/fs/fs'
 import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
-import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { CommandValidation, InvokeOutput, OutputKind, sanitizePath } from './toolShared'
 import { split } from 'shlex'
 
 export enum CommandCategory {
@@ -126,11 +126,6 @@ export interface ExecuteBashParams {
     cwd?: string
 }
 
-export interface CommandValidation {
-    requiresAcceptance: boolean
-    warning?: string
-}
-
 export class ExecuteBash {
     private readonly command: string
     private readonly workingDirectory?: string

packages/core/src/codewhispererChat/tools/listDirectory.ts

@@ -6,7 +6,8 @@ import * as vscode from 'vscode'
 import { getLogger } from '../../shared/logger/logger'
 import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
 import fs from '../../shared/fs/fs'
-import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { CommandValidation, InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { isInDirectory } from '../../shared/filesystemUtilities'
 import { Writable } from 'stream'
 import path from 'path'
 
@@ -61,6 +62,24 @@ export class ListDirectory {
         updates.end()
     }
 
+    public requiresAcceptance(): CommandValidation {
+        // Check if the path is within any of the workspace folders
+        const workspaceFolders = vscode.workspace.workspaceFolders
+
+        if (!workspaceFolders || workspaceFolders.length === 0) {
+            return { requiresAcceptance: true }
+        }
+
+        // Check if the path is within any workspace folder
+        const isInWorkspace = workspaceFolders.some((folder) => isInDirectory(folder.uri.fsPath, this.fsPath))
+
+        if (!isInWorkspace) {
+            return { requiresAcceptance: true }
+        }
+
+        return { requiresAcceptance: false }
+    }
+
     public async invoke(updates?: Writable): Promise<InvokeOutput> {
         try {
             const fileUri = vscode.Uri.file(this.fsPath)

packages/core/src/codewhispererChat/tools/toolShared.ts

@@ -21,6 +21,11 @@ export interface InvokeOutput {
     }
 }
 
+export interface CommandValidation {
+    requiresAcceptance: boolean
+    warning?: string
+}
+
 export function sanitizePath(inputPath: string): string {
     let sanitized = inputPath.trim()
 

packages/core/src/codewhispererChat/tools/toolUtils.ts

@@ -5,10 +5,11 @@
 import { Writable } from 'stream'
 import { FsRead, FsReadParams } from './fsRead'
 import { FsWrite, FsWriteParams } from './fsWrite'
-import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash'
+import { ExecuteBash, ExecuteBashParams } from './executeBash'
 import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
 import { InvokeOutput } from './toolShared'
 import { ListDirectory, ListDirectoryParams } from './listDirectory'
+import { CommandValidation } from './toolShared'
 
 export enum ToolType {
     FsRead = 'fsRead',
@@ -46,7 +47,7 @@ export class ToolUtils {
             case ToolType.ExecuteBash:
                 return tool.tool.requiresAcceptance()
             case ToolType.ListDirectory:
-                return { requiresAcceptance: false }
+                return tool.tool.requiresAcceptance()
         }
     }
 

packages/core/src/test/codewhispererChat/tools/listDirectory.test.ts

@@ -6,12 +6,22 @@ import assert from 'assert'
 import { ListDirectory } from '../../../codewhispererChat/tools/listDirectory'
 import { TestFolder } from '../../testUtil'
 import path from 'path'
+import * as vscode from 'vscode'
+import * as sinon from 'sinon'
 
 describe('ListDirectory Tool', () => {
     let testFolder: TestFolder
+    let workspaceFoldersStub: sinon.SinonStub
 
     before(async () => {
         testFolder = await TestFolder.create()
+        // Initialize the stub
+        workspaceFoldersStub = sinon.stub(vscode.workspace, 'workspaceFolders')
+    })
+
+    after(() => {
+        // Restore the stub after tests
+        workspaceFoldersStub.restore()
     })
 
     it('throws if path is empty', async () => {
@@ -86,4 +96,48 @@ describe('ListDirectory Tool', () => {
         assert.strictEqual(result.output.kind, 'text')
         assert.ok(result.output.content.length > 0)
     })
+
+    it('set requiresAcceptance=true if path is outside workspace', () => {
+        // Mock empty workspace folders
+        workspaceFoldersStub.value([])
+
+        const listDirectory = new ListDirectory({ path: '/some/outside/path' })
+        const validation = listDirectory.requiresAcceptance()
+
+        assert.strictEqual(validation.requiresAcceptance, true, 'Should require acceptance for path outside workspace')
+    })
+
+    it('set requiresAcceptance=true if path is not within any workspace folder', () => {
+        // Mock workspace folders that don't contain the target path
+        workspaceFoldersStub.value([
+            { uri: { fsPath: '/workspace/folder1' } },
+            { uri: { fsPath: '/workspace/folder2' } },
+        ])
+
+        const listDirectory = new ListDirectory({ path: '/some/outside/path' })
+        const validation = listDirectory.requiresAcceptance()
+
+        assert.strictEqual(
+            validation.requiresAcceptance,
+            true,
+            'Should require acceptance for path outside workspace folders'
+        )
+    })
+
+    it('set requiresAcceptance=false if path is within a workspace folder', () => {
+        // Mock workspace folders with one containing the target path
+        const workspacePath = '/workspace/folder1'
+        const targetPath = `${workspacePath}/subfolder`
+
+        workspaceFoldersStub.value([{ uri: { fsPath: workspacePath } }, { uri: { fsPath: '/workspace/folder2' } }])
+
+        const listDirectory = new ListDirectory({ path: targetPath })
+        const validation = listDirectory.requiresAcceptance()
+
+        assert.strictEqual(
+            validation.requiresAcceptance,
+            false,
+            'Should not require acceptance for path within workspace folder'
+        )
+    })
 })