From 21428f5f30d753bcac48d30c11250f28842a8e2e Mon Sep 17 00:00:00 2001
From: Joao Salles <jvsalles@amazon.de>
Date: Tue, 30 Apr 2024 11:19:39 +0200
Subject: [PATCH] fix(featureDev): Prevent crash on large repos or unsupported
 encoding during file collection

---
 ...-f5decdc1-1202-45f7-891c-b89643214089.json |  4 ++
 .../core/src/amazonqFeatureDev/constants.ts   |  3 +
 .../core/src/amazonqFeatureDev/util/files.ts  | 60 +++++++++++++------
 3 files changed, 49 insertions(+), 18 deletions(-)
 create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-f5decdc1-1202-45f7-891c-b89643214089.json

diff --git a/packages/amazonq/.changes/next-release/Bug Fix-f5decdc1-1202-45f7-891c-b89643214089.json b/packages/amazonq/.changes/next-release/Bug Fix-f5decdc1-1202-45f7-891c-b89643214089.json
new file mode 100644
index 00000000000..657281df6ee
--- /dev/null
+++ b/packages/amazonq/.changes/next-release/Bug Fix-f5decdc1-1202-45f7-891c-b89643214089.json	
@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fixed a crash when trying to use Q /dev on large projects or projects containing files with unsupported encoding."
+}
diff --git a/packages/core/src/amazonqFeatureDev/constants.ts b/packages/core/src/amazonqFeatureDev/constants.ts
index e20023373aa..bfa05eadbe3 100644
--- a/packages/core/src/amazonqFeatureDev/constants.ts
+++ b/packages/core/src/amazonqFeatureDev/constants.ts
@@ -14,6 +14,9 @@ export const featureDevChat = 'featureDevChat'
 
 export const featureName = 'Amazon Q feature development'
 
+// Max allowed size for file collection
+export const maxRepoSizeBytes = 200 * 1024 * 1024
+
 // License text that's used in codewhisperer reference log
 export const referenceLogText = (reference: CodeReference) =>
     `[${new Date().toLocaleString()}] Accepted recommendation from Amazon Q. Code provided with reference under <a href="${LicenseUtil.getLicenseHtml(
diff --git a/packages/core/src/amazonqFeatureDev/util/files.ts b/packages/core/src/amazonqFeatureDev/util/files.ts
index 74a023fb90e..fb650d274d0 100644
--- a/packages/core/src/amazonqFeatureDev/util/files.ts
+++ b/packages/core/src/amazonqFeatureDev/util/files.ts
@@ -12,7 +12,7 @@ import { Uri } from 'vscode'
 import { GitIgnoreFilter } from './gitignore'
 
 import AdmZip from 'adm-zip'
-import { PrepareRepoFailedError } from '../errors'
+import { ContentLengthError, PrepareRepoFailedError } from '../errors'
 import { getLogger } from '../../shared/logger/logger'
 import { maxFileSizeBytes } from '../limits'
 import { createHash } from 'crypto'
@@ -21,6 +21,7 @@ import { ToolkitError } from '../../shared/errors'
 import { AmazonqCreateUpload, Metric } from '../../shared/telemetry/telemetry'
 import { TelemetryHelper } from './telemetryHelper'
 import { sanitizeFilename } from '../../shared/utilities/textUtilities'
+import { maxRepoSizeBytes } from '../constants'
 
 export function getExcludePattern(additionalPatterns: string[] = []) {
     const globAlwaysExcludedDirs = getGlobDirExcludedPatterns().map(pattern => `**/${pattern}/*`)
@@ -94,6 +95,7 @@ export async function collectFiles(
         return prefix === '' ? path : `${prefix}/${path}`
     }
 
+    let totalSizeBytes = 0
     for (const rootPath of sourcePaths) {
         const allFiles = await vscode.workspace.findFiles(
             new vscode.RelativePattern(rootPath, '**'),
@@ -102,29 +104,48 @@ export async function collectFiles(
         const files = respectGitIgnore ? await filterOutGitignoredFiles(rootPath, allFiles) : allFiles
 
         for (const file of files) {
-            try {
-                const fileContent = await SystemUtilities.readFile(file, new TextDecoder('utf8', { fatal: true }))
-                const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
+            const relativePath = getWorkspaceRelativePath(file.fsPath, { workspaceFolders })
+            if (!relativePath) {
+                continue
+            }
 
-                if (relativePath) {
-                    storage.push({
-                        workspaceFolder: relativePath.workspaceFolder,
-                        relativeFilePath: relativePath.relativePath,
-                        fileUri: file,
-                        fileContent: fileContent,
-                        zipFilePath: prefixWithFolderPrefix(relativePath.workspaceFolder, relativePath.relativePath),
-                    })
-                }
-            } catch (error) {
-                getLogger().debug(
-                    `featureDev: Failed to read file ${file.fsPath} when collecting repository: ${error}. Skipping the file`
-                )
+            const fileStat = await vscode.workspace.fs.stat(file)
+            if (totalSizeBytes + fileStat.size > maxRepoSizeBytes) {
+                throw new ContentLengthError()
             }
+
+            const fileContent = await readFile(file)
+            if (fileContent === undefined) {
+                continue
+            }
+
+            // Now that we've read the file, increase our usage
+            totalSizeBytes += fileStat.size
+            storage.push({
+                workspaceFolder: relativePath.workspaceFolder,
+                relativeFilePath: relativePath.relativePath,
+                fileUri: file,
+                fileContent: fileContent,
+                zipFilePath: prefixWithFolderPrefix(relativePath.workspaceFolder, relativePath.relativePath),
+            })
         }
     }
     return storage
 }
 
+const readFile = async (file: vscode.Uri) => {
+    try {
+        const fileContent = await SystemUtilities.readFile(file, new TextDecoder('utf8', { fatal: false }))
+        return fileContent
+    } catch (error) {
+        getLogger().debug(
+            `featureDev: Failed to read file ${file.fsPath} when collecting repository. Skipping the file`
+        )
+    }
+
+    return undefined
+}
+
 const getSha256 = (file: Buffer) => createHash('sha256').update(file).digest('base64')
 
 /**
@@ -137,9 +158,9 @@ export async function prepareRepoData(
     span: Metric<AmazonqCreateUpload>
 ) {
     try {
+        const files = await collectFiles(repoRootPaths, workspaceFolders, true)
         const zip = new AdmZip()
 
-        const files = await collectFiles(repoRootPaths, workspaceFolders, true)
         let totalBytes = 0
         for (const file of files) {
             const fileSize = (await vscode.workspace.fs.stat(file.fileUri)).size
@@ -163,6 +184,9 @@ export async function prepareRepoData(
         }
     } catch (error) {
         getLogger().debug(`featureDev: Failed to prepare repo: ${error}`)
+        if (error instanceof ContentLengthError) {
+            throw error
+        }
         throw new PrepareRepoFailedError()
     }
 }