Skip to content

Commit

Permalink
Merge master into feature/emr
Browse files Browse the repository at this point in the history
  • Loading branch information
aws-toolkit-automation authored Jan 31, 2025
2 parents e67e92e + 21c7f12 commit 7ea6253
Show file tree
Hide file tree
Showing 7 changed files with 47 additions and 22 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "Bug Fix",
"description": "Amazon Q: Fix code upload error when using /dev or /doc on Remote SSH"
}
14 changes: 6 additions & 8 deletions packages/amazonq/test/unit/amazonqFeatureDev/util/files.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,11 @@ import {
maxRepoSizeBytes,
} from 'aws-core-vscode/amazonqFeatureDev'
import { assertTelemetry, getWorkspaceFolder, TestFolder } from 'aws-core-vscode/test'
import { fs, AmazonqCreateUpload } from 'aws-core-vscode/shared'
import { fs, AmazonqCreateUpload, ZipStream } from 'aws-core-vscode/shared'
import { MetricName, Span } from 'aws-core-vscode/telemetry'
import sinon from 'sinon'
import { CodeWhispererSettings } from 'aws-core-vscode/codewhisperer'

import AdmZip from 'adm-zip'

const testDevfilePrepareRepo = async (devfileEnabled: boolean) => {
const files: Record<string, string> = {
'file.md': 'test content',
Expand All @@ -38,8 +36,8 @@ const testDevfilePrepareRepo = async (devfileEnabled: boolean) => {
}

const expectedFiles = !devfileEnabled
? ['./file.md', './.gitignore']
: ['./devfile.yaml', './file.md', './.gitignore', './abc.jar', 'data/logo.ico']
? ['file.md', '.gitignore']
: ['devfile.yaml', 'file.md', '.gitignore', 'abc.jar', 'data/logo.ico']

const workspace = getWorkspaceFolder(folder.path)
sinon
Expand Down Expand Up @@ -71,8 +69,8 @@ const testPrepareRepoData = async (
}

// Unzip the buffer and compare the entry names
const zip = new AdmZip(result.zipFileBuffer)
const actualZipEntries = zip.getEntries().map((entry) => entry.entryName)
const zipEntries = await ZipStream.unzip(result.zipFileBuffer)
const actualZipEntries = zipEntries.map((entry) => entry.filename)
actualZipEntries.sort((a, b) => a.localeCompare(b))
assert.deepStrictEqual(actualZipEntries, expectedFiles)
}
Expand All @@ -89,7 +87,7 @@ describe('file utils', () => {
await folder.write('file2.md', 'test content')
const workspace = getWorkspaceFolder(folder.path)

await testPrepareRepoData(workspace, ['./file1.md', './file2.md'])
await testPrepareRepoData(workspace, ['file1.md', 'file2.md'])
})

it('prepareRepoData ignores denied file extensions', async function () {
Expand Down
19 changes: 9 additions & 10 deletions packages/core/src/amazonqFeatureDev/util/files.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,9 @@ import * as vscode from 'vscode'
import * as path from 'path'
import { collectFiles } from '../../shared/utilities/workspaceUtils'

import AdmZip from 'adm-zip'
import { ContentLengthError, PrepareRepoFailedError } from '../errors'
import { getLogger } from '../../shared/logger/logger'
import { maxFileSizeBytes } from '../limits'
import { createHash } from 'crypto'
import { CurrentWsFolders } from '../types'
import { hasCode, ToolkitError } from '../../shared/errors'
import { AmazonqCreateUpload, Span, telemetry as amznTelemetry } from '../../shared/telemetry/telemetry'
Expand All @@ -20,8 +18,7 @@ import { maxRepoSizeBytes } from '../constants'
import { isCodeFile } from '../../shared/filetypes'
import { fs } from '../../shared'
import { CodeWhispererSettings } from '../../codewhisperer'

const getSha256 = (file: Buffer) => createHash('sha256').update(file).digest('base64')
import { ZipStream } from '../../shared/utilities/zipStream'

export async function checkForDevFile(root: string) {
const devFilePath = root + '/devfile.yaml'
Expand All @@ -37,7 +34,7 @@ export async function prepareRepoData(
workspaceFolders: CurrentWsFolders,
telemetry: TelemetryHelper,
span: Span<AmazonqCreateUpload>,
zip: AdmZip = new AdmZip()
zip: ZipStream = new ZipStream()
) {
try {
const autoBuildSetting = CodeWhispererSettings.instance.getAutoBuildSetting()
Expand Down Expand Up @@ -77,11 +74,12 @@ export async function prepareRepoData(
continue
}
totalBytes += fileSize

const zipFolderPath = path.dirname(file.zipFilePath)
// Paths in zip should be POSIX compliant regardless of OS
// Reference: https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT
const posixPath = file.zipFilePath.split(path.sep).join(path.posix.sep)

try {
zip.addLocalFile(file.fileUri.fsPath, zipFolderPath)
zip.writeFile(file.fileUri.fsPath, posixPath)
} catch (error) {
if (error instanceof Error && error.message.includes('File not found')) {
// No-op: Skip if file was deleted or does not exist
Expand Down Expand Up @@ -111,11 +109,12 @@ export async function prepareRepoData(

telemetry.setRepositorySize(totalBytes)
span.record({ amazonqRepositorySize: totalBytes })
const zipResult = await zip.finalize()

const zipFileBuffer = zip.toBuffer()
const zipFileBuffer = zipResult.streamBuffer.getContents() || Buffer.from('')
return {
zipFileBuffer,
zipFileChecksum: getSha256(zipFileBuffer),
zipFileChecksum: zipResult.hash,
}
} catch (error) {
getLogger().debug(`featureDev: Failed to prepare repo: ${error}`)
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/shared/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ export * from './localizedText'
export * as env from './vscode/env'
export * from './vscode/commands2'
export * from './utilities/pathUtils'
export * from './utilities/zipStream'
export * from './errors'
export * as messages from './utilities/messages'
export * as errors from './errors'
Expand Down
15 changes: 14 additions & 1 deletion packages/core/src/shared/utilities/zipStream.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import { WritableStreamBuffer } from 'stream-buffers'
import crypto from 'crypto'
import { readFileAsString } from '../filesystemUtilities'
// Use require instead of import since this package doesn't support commonjs
const { ZipWriter, TextReader } = require('@zip.js/zip.js')
const { ZipWriter, TextReader, ZipReader, Uint8ArrayReader } = require('@zip.js/zip.js')
import { getLogger } from '../logger/logger'

export interface ZipStreamResult {
Expand All @@ -15,6 +15,10 @@ export interface ZipStreamResult {
streamBuffer: WritableStreamBuffer
}

export type ZipReaderResult = {
filename: string
}

export type ZipStreamProps = {
hashAlgorithm: 'md5' | 'sha256'
maxNumberOfFileStreams: number
Expand Down Expand Up @@ -150,4 +154,13 @@ export class ZipStream {
streamBuffer: this._streamBuffer,
}
}

public static async unzip(zipBuffer: Buffer): Promise<ZipReaderResult[]> {
const reader = new ZipReader(new Uint8ArrayReader(new Uint8Array(zipBuffer)))
try {
return await reader.getEntries()
} finally {
await reader.close()
}
}
}
11 changes: 11 additions & 0 deletions packages/core/src/test/shared/utilities/zipStream.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,4 +60,15 @@ describe('zipStream', function () {
assert.strictEqual(result.hash, expectedMd5)
assert.strictEqual(result.sizeInBytes, (await fs.stat(zipPath)).size)
})

it('should unzip from a buffer', async function () {
const zipStream = new ZipStream()
await zipStream.writeString('foo bar', 'file.txt')
const result = await zipStream.finalize()

const zipBuffer = result.streamBuffer.getContents()
assert.ok(zipBuffer)
const zipEntries = await ZipStream.unzip(zipBuffer)
assert.strictEqual(zipEntries[0].filename, 'file.txt')
})
})
5 changes: 2 additions & 3 deletions packages/core/src/testInteg/perf/prepareRepoData.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ function performanceTestWrapper(numFiles: number, fileSize: number) {
getEqualOSTestOptions({
userCpuUsage: 200,
systemCpuUsage: 35,
heapTotal: 4,
heapTotal: 20,
}),
`handles ${numFiles} files of size ${fileSize} bytes`,
function () {
Expand Down Expand Up @@ -63,8 +63,7 @@ function verifyResult(setup: setupResult, result: resultType, telemetry: Telemet
assert.strictEqual(Buffer.isBuffer(result.zipFileBuffer), true)
assert.strictEqual(telemetry.repositorySize, expectedSize)
assert.strictEqual(result.zipFileChecksum.length, 44)

assert.ok(getFsCallsUpperBound(setup.fsSpy) <= setup.numFiles * 4, 'total system calls should be under 4 per file')
assert.ok(getFsCallsUpperBound(setup.fsSpy) <= setup.numFiles * 8, 'total system calls should be under 8 per file')
}

describe('prepareRepoData', function () {
Expand Down

0 comments on commit 7ea6253

Please sign in to comment.