Skip to content

Commit de45a9b

Browse files
authored
feat(codewhisperer): golang support (#4138)
* feat(codewhisperer): golang support * add changelog * fix windows test * address comments
1 parent eedf096 commit de45a9b

File tree

11 files changed

+463
-0
lines changed

11 files changed

+463
-0
lines changed
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "CodeWhisperer security scans now support Go files."
4+
}

src/codewhisperer/models/constants.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,8 @@ export const codeScanJavaPayloadSizeLimitBytes = Math.pow(2, 20) // 1 MB
200200

201201
export const codeScanCsharpPayloadSizeLimitBytes = Math.pow(2, 20) // 1 MB
202202

203+
export const codeScanGoPayloadSizeLimitBytes = Math.pow(2, 20) // 1 MB
204+
203205
export const codeScanPythonPayloadSizeLimitBytes = 200 * Math.pow(2, 10) // 200 KB
204206

205207
export const codeScanCFPayloadSizeLimitBytes = 200 * Math.pow(2, 10) // 200 KB

src/codewhisperer/util/dependencyGraph/dependencyGraph.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ export const DependencyGraphConstants = {
5555
ymlExt: '.yml',
5656
tfExt: '.tf',
5757
hclExt: '.hcl',
58+
goExt: '.go',
5859
}
5960

6061
export abstract class DependencyGraph {

src/codewhisperer/util/dependencyGraph/dependencyGraphFactory.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { cloudformationDependencyGraph } from './cloudformationDependencyGraph'
1212
import { DependencyGraphConstants } from './dependencyGraph'
1313
import * as vscode from 'vscode'
1414
import { terraformDependencyGraph } from './terraformDependencyGraph'
15+
import { GoDependencyGraph } from './goDependencyGraph'
1516

1617
const languageMap = {
1718
java: JavaDependencyGraph,
@@ -21,6 +22,7 @@ const languageMap = {
2122
csharp: CsharpDependencyGraph,
2223
cloudformation: cloudformationDependencyGraph,
2324
terraform: terraformDependencyGraph,
25+
go: GoDependencyGraph,
2426
} as const
2527

2628
type LanguageMap = typeof languageMap
@@ -53,6 +55,8 @@ export class DependencyGraphFactory {
5355
return new languageMap['csharp']('csharp' satisfies CodeWhispererConstants.PlatformLanguageId)
5456
case 'yaml' satisfies CodeWhispererConstants.PlatformLanguageId:
5557
return new languageMap['cloudformation']('yaml' satisfies CodeWhispererConstants.PlatformLanguageId)
58+
case 'go' satisfies CodeWhispererConstants.PlatformLanguageId:
59+
return new languageMap['go']('go' satisfies CodeWhispererConstants.PlatformLanguageId)
5660
default:
5761
return this.getDependencyGraphFromFileExtensions(editor.document.fileName)
5862
}
Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import * as vscode from 'vscode'
7+
import { Uri } from 'vscode'
8+
import { DependencyGraph, DependencyGraphConstants, Truncation } from './dependencyGraph'
9+
import * as CodeWhispererConstants from '../../models/constants'
10+
import { existsSync, readdirSync, statSync } from 'fs'
11+
import { sleep } from '../../../shared/utilities/timeoutUtils'
12+
import { getLogger } from '../../../shared/logger'
13+
import path from 'path'
14+
import { readFileAsString } from '../../../shared/filesystemUtilities'
15+
import { ToolkitError } from '../../../shared/errors'
16+
17+
const importRegex = /^\s*import\s+([^(]+?$|\([^)]+\))/gm
18+
const moduleRegex = /"[^"\r\n]+"/gm
19+
const packageRegex = /^package\s+(.+)/gm
20+
21+
export class GoDependencyGraph extends DependencyGraph {
22+
override async getSourceDependencies(uri: Uri, content: string): Promise<string[]> {
23+
const imports = this.readImports(content)
24+
const dependencies = this.getDependencies(uri, imports)
25+
return dependencies
26+
}
27+
28+
// Returns file paths of other .go files in the same directory declared with the same package statement.
29+
override async getSamePackageFiles(uri: Uri): Promise<string[]> {
30+
const fileList: string[] = []
31+
const packagePath = path.dirname(uri.fsPath)
32+
const fileName = path.basename(uri.fsPath)
33+
const content = await readFileAsString(uri.fsPath)
34+
const packageName = this.readPackageName(content)
35+
36+
const files = readdirSync(packagePath, { withFileTypes: true })
37+
for (const file of files) {
38+
if (file.isDirectory() || !file.name.endsWith(DependencyGraphConstants.goExt) || file.name === fileName) {
39+
continue
40+
}
41+
const filePath = path.join(packagePath, file.name)
42+
const content = await readFileAsString(filePath)
43+
if (this.readPackageName(content) !== packageName) {
44+
continue
45+
}
46+
fileList.push(filePath)
47+
}
48+
49+
return fileList
50+
}
51+
52+
override async isTestFile(content: string): Promise<boolean> {
53+
const imports = this.readImports(content)
54+
const filteredImports = imports.filter(importStr => {
55+
return importStr.includes('"testing"')
56+
})
57+
return filteredImports.length > 0
58+
}
59+
60+
override async generateTruncation(uri: Uri): Promise<Truncation> {
61+
try {
62+
const workspaceFolder = vscode.workspace.getWorkspaceFolder(uri)
63+
if (workspaceFolder === undefined) {
64+
this._pickedSourceFiles.add(uri.fsPath)
65+
} else {
66+
await this.searchDependency(uri)
67+
await this.traverseDir(this.getProjectPath(uri))
68+
}
69+
await sleep(1000)
70+
getLogger().verbose(`CodeWhisperer: Picked source files: [${[...this._pickedSourceFiles].join(', ')}]`)
71+
const truncDirPath = this.getTruncDirPath(uri)
72+
this.copyFilesToTmpDir(this._pickedSourceFiles, truncDirPath)
73+
const zipFilePath = this.zipDir(truncDirPath, CodeWhispererConstants.codeScanZipExt)
74+
const zipFileSize = statSync(zipFilePath).size
75+
return {
76+
rootDir: truncDirPath,
77+
zipFilePath,
78+
scannedFiles: new Set(this._pickedSourceFiles),
79+
srcPayloadSizeInBytes: this._totalSize,
80+
zipFileSizeInBytes: zipFileSize,
81+
buildPayloadSizeInBytes: 0,
82+
lines: this._totalLines,
83+
}
84+
} catch (error) {
85+
getLogger().error('Go dependency graph error caused by:', error)
86+
throw ToolkitError.chain(error, 'Go context processing failed.')
87+
}
88+
}
89+
90+
override async searchDependency(uri: Uri): Promise<Set<string>> {
91+
const filePath = uri.fsPath
92+
const q: string[] = []
93+
q.push(filePath)
94+
const siblings = await this.getSamePackageFiles(uri)
95+
siblings.forEach(sibling => {
96+
q.push(sibling)
97+
})
98+
while (q.length > 0) {
99+
let count: number = q.length
100+
while (count > 0) {
101+
if (this.reachSizeLimit(this._totalSize)) {
102+
return this._pickedSourceFiles
103+
}
104+
count -= 1
105+
const currentFilePath = q.shift()
106+
if (currentFilePath === undefined) {
107+
throw new Error('"undefined" is invalid for queued file.')
108+
}
109+
if (this._pickedSourceFiles.has(currentFilePath)) {
110+
continue
111+
}
112+
this._pickedSourceFiles.add(currentFilePath)
113+
this._totalSize += statSync(currentFilePath).size
114+
const uri = vscode.Uri.file(currentFilePath)
115+
const content: string = await readFileAsString(uri.fsPath)
116+
const dependencies = await this.getSourceDependencies(uri, content)
117+
dependencies.forEach(dependency => {
118+
q.push(dependency)
119+
})
120+
}
121+
}
122+
123+
return this._pickedSourceFiles
124+
}
125+
126+
override async traverseDir(dirPath: string): Promise<void> {
127+
if (this.reachSizeLimit(this._totalSize)) {
128+
return
129+
}
130+
readdirSync(dirPath, { withFileTypes: true }).forEach(async file => {
131+
const absPath = path.join(dirPath, file.name)
132+
if (file.name.charAt(0) === '.' || !existsSync(absPath)) {
133+
return
134+
}
135+
if (file.isDirectory()) {
136+
await this.traverseDir(absPath)
137+
} else if (file.isFile()) {
138+
if (
139+
file.name.endsWith(DependencyGraphConstants.goExt) &&
140+
!this.reachSizeLimit(this._totalSize) &&
141+
!this.willReachSizeLimit(this._totalSize, statSync(absPath).size)
142+
) {
143+
await this.searchDependency(vscode.Uri.file(absPath))
144+
}
145+
}
146+
})
147+
}
148+
149+
override parseImport(importStr: string, dirPaths: string[]): string[] {
150+
if (this._parsedStatements.has(importStr)) {
151+
return []
152+
}
153+
154+
this._parsedStatements.add(importStr)
155+
const modulePaths = this.extractModulePaths(importStr)
156+
const dependencies = this.generateSourceFilePaths(modulePaths, dirPaths)
157+
return dependencies
158+
}
159+
160+
override updateSysPaths(uri: Uri): void {
161+
this.getDirPaths(uri).forEach(dirPath => {
162+
this._sysPaths.add(dirPath)
163+
})
164+
}
165+
166+
override getDependencies(uri: Uri, imports: string[]): string[] {
167+
const dependencies: string[] = []
168+
imports.forEach(importStr => {
169+
this.updateSysPaths(uri)
170+
const findings = this.parseImport(importStr, Array.from(this._sysPaths.values()))
171+
const validSourceFiles = findings.filter(finding => !this._pickedSourceFiles.has(finding))
172+
validSourceFiles.forEach(file => {
173+
if (existsSync(file) && !this.willReachSizeLimit(this._totalSize, statSync(file).size)) {
174+
dependencies.push(file)
175+
}
176+
})
177+
})
178+
return dependencies
179+
}
180+
181+
override getPayloadSizeLimitInBytes(): number {
182+
return CodeWhispererConstants.codeScanGoPayloadSizeLimitBytes
183+
}
184+
185+
private generateSourceFilePaths(modulePaths: string[], dirPaths: string[]): string[] {
186+
const filePaths: string[] = []
187+
modulePaths.forEach(modulePath => {
188+
dirPaths.forEach(dirPath => {
189+
const packageDir = this.generateSourceFilePath(modulePath, dirPath)
190+
if (packageDir !== '') {
191+
readdirSync(packageDir, { withFileTypes: true }).forEach(file => {
192+
if (file.name.endsWith(DependencyGraphConstants.goExt)) {
193+
filePaths.push(path.join(packageDir, file.name))
194+
}
195+
})
196+
}
197+
})
198+
})
199+
return filePaths
200+
}
201+
202+
private generateSourceFilePath(modulePath: string, dirPath: string): string {
203+
if (modulePath.length === 0) {
204+
return ''
205+
}
206+
const packageDir = path.join(dirPath, modulePath)
207+
const slashPos = modulePath.indexOf('/')
208+
const newModulePath = slashPos !== -1 ? modulePath.substring(slashPos + 1) : ''
209+
210+
return existsSync(packageDir) ? packageDir : this.generateSourceFilePath(newModulePath, dirPath)
211+
}
212+
213+
private extractModulePaths(importStr: string): string[] {
214+
const matches = importStr.match(moduleRegex)
215+
if (matches) {
216+
return matches.map(match => match.substring(1, match.length - 1))
217+
}
218+
return []
219+
}
220+
221+
private readImports(content: string) {
222+
this._totalLines += content.split(DependencyGraphConstants.newlineRegex).length
223+
const regExp = new RegExp(importRegex)
224+
return content.match(regExp) ?? []
225+
}
226+
227+
private readPackageName(content: string) {
228+
const regExp = new RegExp(packageRegex)
229+
const matches = regExp.exec(content)
230+
if (matches && matches.length > 1) {
231+
return matches[1]
232+
}
233+
return ''
234+
}
235+
}
236+
237+
export class GoDependencyGraphError extends Error {}

src/test/codewhisperer/testUtil.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { getLogger } from '../../shared/logger'
1717
import { CodeWhispererCodeCoverageTracker } from '../../codewhisperer/tracker/codewhispererCodeCoverageTracker'
1818
import globals from '../../shared/extensionGlobals'
1919
import { session } from '../../codewhisperer/util/codeWhispererSession'
20+
import fs from 'fs'
2021
import { DefaultAWSClientBuilder, ServiceOptions } from '../../shared/awsClientBuilder'
2122
import { FakeAwsContext } from '../utilities/fakeAwsContext'
2223
import { spy } from '../utilities/mockito'
@@ -198,3 +199,10 @@ export function createCodeActionContext(): vscode.CodeActionContext {
198199
triggerKind: vscode.CodeActionTriggerKind.Automatic,
199200
}
200201
}
202+
203+
export function createMockDirentFile(fileName: string): fs.Dirent {
204+
const dirent = new fs.Dirent()
205+
dirent.isFile = () => true
206+
dirent.name = fileName
207+
return dirent
208+
}

0 commit comments

Comments
 (0)