Skip to content

Commit 9945398

Browse files
patrickvbeurdenpatrickvbeurden
authored
feat(amazonq): add regionalization support to security scan server (#859)
* feat(amazonq): add regionalization support to security scan server * refactor(amazonq): improve region discovery logging
1 parent 00fe7e2 commit 9945398

File tree

9 files changed

+89
-58
lines changed

9 files changed

+89
-58
lines changed

‎server/aws-lsp-codewhisperer/src/constants.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,6 @@ export const DEFAULT_AWS_Q_REGION = 'us-east-1'
44
export const AWS_Q_ENDPOINTS = {
55
[DEFAULT_AWS_Q_REGION]: DEFAULT_AWS_Q_ENDPOINT_URL,
66
}
7+
8+
export const AWS_Q_REGION_ENV_VAR = 'AWS_Q_REGION'
9+
export const AWS_Q_ENDPOINT_URL_ENV_VAR = 'AWS_Q_ENDPOINT_URL'

‎server/aws-lsp-codewhisperer/src/language-server/amazonQServiceManager/AmazonQTokenServiceManager.test.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,13 @@ import {
1515
LSPErrorCodes,
1616
ResponseError,
1717
} from '@aws/language-server-runtimes/protocol'
18-
import { AWS_Q_ENDPOINTS, DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION } from '../../constants'
18+
import {
19+
AWS_Q_ENDPOINT_URL_ENV_VAR,
20+
AWS_Q_ENDPOINTS,
21+
AWS_Q_REGION_ENV_VAR,
22+
DEFAULT_AWS_Q_ENDPOINT_URL,
23+
DEFAULT_AWS_Q_REGION,
24+
} from '../../constants'
1925
import * as qDeveloperProfilesFetcherModule from './qDeveloperProfiles'
2026
import { setCredentialsForAmazonQTokenServiceManagerFactory } from '../testUtils'
2127
import { CodeWhispererStreaming } from '@amzn/codewhisperer-streaming'
@@ -215,8 +221,8 @@ describe('AmazonQTokenServiceManager', () => {
215221
})
216222

217223
it('should initialize service with region set by runtime if not set by client', async () => {
218-
features.runtime.getConfiguration.withArgs('AWS_Q_REGION').returns('eu-central-1')
219-
features.runtime.getConfiguration.withArgs('AWS_Q_ENDPOINT_URL').returns(TEST_ENDPOINT_EU_CENTRAL_1)
224+
features.runtime.getConfiguration.withArgs(AWS_Q_REGION_ENV_VAR).returns('eu-central-1')
225+
features.runtime.getConfiguration.withArgs(AWS_Q_ENDPOINT_URL_ENV_VAR).returns(TEST_ENDPOINT_EU_CENTRAL_1)
220226

221227
amazonQTokenServiceManager.getCodewhispererService()
222228
assert(codewhispererStubFactory.calledOnceWithExactly('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1))

‎server/aws-lsp-codewhisperer/src/language-server/amazonQServiceManager/AmazonQTokenServiceManager.ts‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,13 @@ import {
1414
CredentialsType,
1515
InitializeParams,
1616
} from '@aws/language-server-runtimes/server-interface'
17-
import { DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION, AWS_Q_ENDPOINTS } from '../../constants'
17+
import {
18+
DEFAULT_AWS_Q_ENDPOINT_URL,
19+
DEFAULT_AWS_Q_REGION,
20+
AWS_Q_ENDPOINTS,
21+
AWS_Q_REGION_ENV_VAR,
22+
AWS_Q_ENDPOINT_URL_ENV_VAR,
23+
} from '../../constants'
1824
import { CodeWhispererServiceToken } from '../codeWhispererService'
1925
import {
2026
AmazonQError,
@@ -447,14 +453,30 @@ export class AmazonQTokenServiceManager implements BaseAmazonQServiceManager {
447453

448454
private createCodewhispererServiceInstances(connectionType: 'builderId' | 'identityCenter', region?: string) {
449455
this.logServiceState('Initializing CodewhispererService')
450-
let awsQRegion = this.features.runtime.getConfiguration('AWS_Q_REGION') ?? DEFAULT_AWS_Q_REGION
451-
let awsQEndpoint: string | undefined =
452-
this.features.runtime.getConfiguration('AWS_Q_ENDPOINT_URL') ?? DEFAULT_AWS_Q_ENDPOINT_URL
456+
let awsQRegion: string
457+
let awsQEndpoint: string | undefined
453458

454459
if (region) {
460+
this.log(
461+
`Selecting region (found: ${region}) provided by ${connectionType === 'builderId' ? 'client' : 'profile'}`
462+
)
455463
awsQRegion = region
456464
// @ts-ignore
457-
awsQEndpoint = AWS_Q_ENDPOINTS[region]
465+
awsQEndpoint = AWS_Q_ENDPOINTS[awsQRegion]
466+
} else {
467+
const runtimeRegion = this.features.runtime.getConfiguration(AWS_Q_REGION_ENV_VAR)
468+
469+
if (runtimeRegion) {
470+
this.log(`Selecting region (found: ${runtimeRegion}) provided by runtime`)
471+
awsQRegion = runtimeRegion
472+
// prettier-ignore
473+
awsQEndpoint = // @ts-ignore
474+
this.features.runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? AWS_Q_ENDPOINTS[awsQRegion]
475+
} else {
476+
this.log('Region not provided by caller or runtime, falling back to default region and endpoint')
477+
awsQRegion = DEFAULT_AWS_Q_REGION
478+
awsQEndpoint = DEFAULT_AWS_Q_ENDPOINT_URL
479+
}
458480
}
459481

460482
if (!awsQEndpoint) {

‎server/aws-lsp-codewhisperer/src/language-server/codeWhispererSecurityScanServer.ts‎

Lines changed: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -20,30 +20,18 @@ import { getErrorMessage, parseJson } from './utils'
2020
import { getUserAgent } from './utilities/telemetryUtils'
2121
import { DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION } from '../constants'
2222
import { SDKInitializator } from '@aws/language-server-runtimes/server-interface'
23+
import { AmazonQTokenServiceManager } from './amazonQServiceManager/AmazonQTokenServiceManager'
2324

2425
const RunSecurityScanCommand = 'aws/codewhisperer/runSecurityScan'
2526
const CancelSecurityScanCommand = 'aws/codewhisperer/cancelSecurityScan'
2627

2728
export const SecurityScanServerToken =
28-
(
29-
service: (
30-
credentialsProvider: CredentialsProvider,
31-
workspace: Workspace,
32-
awsQRegion: string,
33-
awsQEndpointUrl: string,
34-
sdkInitializator: SDKInitializator
35-
) => CodeWhispererServiceToken
36-
): Server =>
29+
(): Server =>
3730
({ credentialsProvider, workspace, logging, lsp, telemetry, runtime, sdkInitializator }) => {
38-
const codewhispererclient = service(
39-
credentialsProvider,
40-
workspace,
41-
runtime.getConfiguration('AWS_Q_REGION') ?? DEFAULT_AWS_Q_REGION,
42-
runtime.getConfiguration('AWS_Q_ENDPOINT_URL') ?? DEFAULT_AWS_Q_ENDPOINT_URL,
43-
sdkInitializator
44-
)
31+
let amazonQServiceManager: AmazonQTokenServiceManager
32+
let scanHandler: SecurityScanHandler
33+
4534
const diagnosticsProvider = new SecurityScanDiagnosticsProvider(lsp, logging)
46-
const scanHandler = new SecurityScanHandler(codewhispererclient, workspace, logging)
4735

4836
const runSecurityScan = async (params: SecurityScanRequestParams, token: CancellationToken) => {
4937
logging.log(`Starting security scan`)
@@ -228,10 +216,6 @@ export const SecurityScanServerToken =
228216
return
229217
}
230218
const onInitializeHandler = (params: InitializeParams) => {
231-
codewhispererclient.updateClientConfig({
232-
customUserAgent: getUserAgent(params, runtime.serverInfo),
233-
})
234-
235219
return {
236220
capabilities: {
237221
executeCommandProvider: {
@@ -241,8 +225,21 @@ export const SecurityScanServerToken =
241225
}
242226
}
243227

228+
const onInitializedHandler = () => {
229+
amazonQServiceManager = AmazonQTokenServiceManager.getInstance({
230+
lsp,
231+
logging,
232+
runtime,
233+
credentialsProvider,
234+
sdkInitializator,
235+
workspace,
236+
})
237+
scanHandler = new SecurityScanHandler(amazonQServiceManager, workspace, logging)
238+
}
239+
244240
lsp.onExecuteCommand(onExecuteCommandHandler)
245241
lsp.addInitializer(onInitializeHandler)
242+
lsp.onInitialized(onInitializedHandler)
246243
lsp.onDidChangeTextDocument(async p => {
247244
const textDocument = await workspace.getTextDocument(p.textDocument.uri)
248245
const languageId = getSupportedLanguageId(textDocument, supportedSecurityScanLanguages)
@@ -265,6 +262,6 @@ export const SecurityScanServerToken =
265262

266263
return () => {
267264
// dispose function
268-
scanHandler.tokenSource.dispose()
265+
scanHandler?.tokenSource.dispose()
269266
}
270267
}

‎server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,12 @@ import { fetchSupplementalContext } from './utilities/supplementalContextUtil/su
3838
import { textUtils } from '@aws/lsp-core'
3939
import { TelemetryService } from './telemetryService'
4040
import { AcceptedSuggestionEntry, CodeDiffTracker } from './telemetry/codeDiffTracker'
41-
import { DEFAULT_AWS_Q_ENDPOINT_URL, DEFAULT_AWS_Q_REGION } from '../constants'
41+
import {
42+
AWS_Q_ENDPOINT_URL_ENV_VAR,
43+
AWS_Q_REGION_ENV_VAR,
44+
DEFAULT_AWS_Q_ENDPOINT_URL,
45+
DEFAULT_AWS_Q_REGION,
46+
} from '../constants'
4247
import { AmazonQTokenServiceManager } from './amazonQServiceManager/AmazonQTokenServiceManager'
4348
import { AmazonQError, AmazonQServiceInitializationError } from './amazonQServiceManager/errors'
4449
import { BaseAmazonQServiceManager } from './amazonQServiceManager/BaseAmazonQServiceManager'
@@ -248,8 +253,8 @@ export const CodewhispererServerFactory =
248253

249254
const sessionManager = SessionManager.getInstance()
250255

251-
const awsQRegion = runtime.getConfiguration('AWS_Q_REGION') ?? DEFAULT_AWS_Q_REGION
252-
const awsQEndpointUrl = runtime.getConfiguration('AWS_Q_ENDPOINT_URL') ?? DEFAULT_AWS_Q_ENDPOINT_URL
256+
const awsQRegion = runtime.getConfiguration(AWS_Q_REGION_ENV_VAR) ?? DEFAULT_AWS_Q_REGION
257+
const awsQEndpointUrl = runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? DEFAULT_AWS_Q_ENDPOINT_URL
253258
const fallbackCodeWhispererService = service(
254259
credentialsProvider,
255260
workspace,

‎server/aws-lsp-codewhisperer/src/language-server/netTransformServer.ts‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,12 @@ const PollTransformForPlanCommand = 'aws/qNetTransform/pollTransformForPlan'
4545
const GetTransformPlanCommand = 'aws/qNetTransform/getTransformPlan'
4646
const CancelTransformCommand = 'aws/qNetTransform/cancelTransform'
4747
const DownloadArtifactsCommand = 'aws/qNetTransform/downloadArtifacts'
48-
import { DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL } from '../constants'
48+
import {
49+
DEFAULT_AWS_Q_REGION,
50+
DEFAULT_AWS_Q_ENDPOINT_URL,
51+
AWS_Q_REGION_ENV_VAR,
52+
AWS_Q_ENDPOINT_URL_ENV_VAR,
53+
} from '../constants'
4954
import { SDKInitializator } from '@aws/language-server-runtimes/server-interface'
5055

5156
/**
@@ -67,8 +72,8 @@ export const QNetTransformServerToken =
6772
const codewhispererclient = service(
6873
credentialsProvider,
6974
workspace,
70-
runtime.getConfiguration('AWS_Q_REGION') ?? DEFAULT_AWS_Q_REGION,
71-
runtime.getConfiguration('AWS_Q_ENDPOINT_URL') ?? DEFAULT_AWS_Q_ENDPOINT_URL,
75+
runtime.getConfiguration(AWS_Q_REGION_ENV_VAR) ?? DEFAULT_AWS_Q_REGION,
76+
runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? DEFAULT_AWS_Q_ENDPOINT_URL,
7277
sdkInitializator
7378
)
7479
const transformHandler = new TransformHandler(codewhispererclient, workspace, logging, runtime)
@@ -126,8 +131,8 @@ export const QNetTransformServerToken =
126131
const cwStreamingClientInstance = new StreamingClient()
127132
const cwStreamingClient = await cwStreamingClientInstance.getStreamingClient(
128133
credentialsProvider,
129-
runtime.getConfiguration('AWS_Q_REGION') ?? DEFAULT_AWS_Q_REGION,
130-
runtime.getConfiguration('AWS_Q_ENDPOINT_URL') ?? DEFAULT_AWS_Q_ENDPOINT_URL,
134+
runtime.getConfiguration(AWS_Q_REGION_ENV_VAR) ?? DEFAULT_AWS_Q_REGION,
135+
runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? DEFAULT_AWS_Q_ENDPOINT_URL,
131136
sdkInitializator,
132137
customCWClientConfig
133138
)

‎server/aws-lsp-codewhisperer/src/language-server/proxy-server.ts‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -29,17 +29,7 @@ export const CodeWhispererServerIAMProxy = CodewhispererServerFactory(
2929
}
3030
)
3131

32-
export const CodeWhispererSecurityScanServerTokenProxy = SecurityScanServerToken(
33-
(credentialsProvider, workspace, awsQRegion, awsQEndpointUrl, sdkInitializator) => {
34-
return new CodeWhispererServiceToken(
35-
credentialsProvider,
36-
workspace,
37-
awsQRegion,
38-
awsQEndpointUrl,
39-
sdkInitializator
40-
)
41-
}
42-
)
32+
export const CodeWhispererSecurityScanServerTokenProxy = SecurityScanServerToken()
4333

4434
export const QNetTransformServerTokenProxy = QNetTransformServerToken(
4535
(credentialsProvider, workspace, awsQRegion, awsQEndpointUrl, sdkInitializator) => {

‎server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.test.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { StartCodeAnalysisRequest } from '../../client/token/codewhispererbearer
88
import { CodeWhispererServiceToken } from '../codeWhispererService'
99
import { SecurityScanHandler } from './securityScanHandler'
1010
import { RawCodeScanIssue } from './types'
11+
import { AmazonQTokenServiceManager } from '../amazonQServiceManager/AmazonQTokenServiceManager'
1112

1213
const mockCodeScanFindings = JSON.stringify([
1314
{
@@ -54,9 +55,11 @@ describe('securityScanHandler', () => {
5455
const mockedLogging = stubInterface<Logging>()
5556
beforeEach(async () => {
5657
// Set up the server with a mock service
58+
const serviceManager = stubInterface<AmazonQTokenServiceManager>()
5759
client = stubInterface<CodeWhispererServiceToken>()
60+
serviceManager.getCodewhispererService.returns(client)
5861
workspace = stubInterface<Workspace>()
59-
securityScanhandler = new SecurityScanHandler(client, workspace, mockedLogging)
62+
securityScanhandler = new SecurityScanHandler(serviceManager, workspace, mockedLogging)
6063
})
6164

6265
describe('Test createCodeResourcePresignedUrlHandler', () => {

‎server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.ts‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,18 @@ import {
1616
ListCodeAnalysisFindingsRequest,
1717
StartCodeAnalysisRequest,
1818
} from '../../client/token/codewhispererbearertokenclient'
19-
import { CodeWhispererServiceToken } from '../codeWhispererService'
2019
import { sleep } from '../dependencyGraph/commonUtil'
2120
import { AggregatedCodeScanIssue, RawCodeScanIssue } from './types'
21+
import { AmazonQTokenServiceManager } from '../amazonQServiceManager/AmazonQTokenServiceManager'
2222

2323
export class SecurityScanHandler {
24-
private client: CodeWhispererServiceToken
24+
private serviceManager: AmazonQTokenServiceManager
2525
private workspace: Workspace
2626
private logging: Logging
2727
public tokenSource: CancellationTokenSource
2828

29-
constructor(client: CodeWhispererServiceToken, workspace: Workspace, logging: Logging) {
30-
this.client = client
29+
constructor(serviceManager: AmazonQTokenServiceManager, workspace: Workspace, logging: Logging) {
30+
this.serviceManager = serviceManager
3131
this.workspace = workspace
3232
this.logging = logging
3333
this.tokenSource = new CancellationTokenSource()
@@ -56,7 +56,7 @@ export class SecurityScanHandler {
5656
}
5757
try {
5858
this.logging.log('Prepare for uploading src context...')
59-
const response = await this.client.createUploadUrl(request)
59+
const response = await this.serviceManager.getCodewhispererService().createUploadUrl(request)
6060
this.logging.log(`Request id: ${response.$response.requestId}`)
6161
this.logging.log(`Complete Getting presigned Url for uploading src context.`)
6262
this.logging.log(`Uploading src context...`)
@@ -107,7 +107,7 @@ export class SecurityScanHandler {
107107
}
108108
this.logging.log(`Creating scan job...`)
109109
try {
110-
const resp = await this.client.startCodeAnalysis(req)
110+
const resp = await this.serviceManager.getCodewhispererService().startCodeAnalysis(req)
111111
this.logging.log(`Request id: ${resp.$response.requestId}`)
112112
return resp
113113
} catch (error) {
@@ -126,7 +126,7 @@ export class SecurityScanHandler {
126126
const req: GetCodeAnalysisRequest = {
127127
jobId: jobId,
128128
}
129-
const resp = await this.client.getCodeAnalysis(req)
129+
const resp = await this.serviceManager.getCodewhispererService().getCodeAnalysis(req)
130130
this.logging.log(`Request id: ${resp.$response.requestId}`)
131131

132132
if (resp.status !== 'Pending') {
@@ -151,7 +151,7 @@ export class SecurityScanHandler {
151151
jobId,
152152
codeAnalysisFindingsSchema: 'codeanalysis/findings/1.0',
153153
}
154-
const response = await this.client.listCodeAnalysisFindings(request)
154+
const response = await this.serviceManager.getCodewhispererService().listCodeAnalysisFindings(request)
155155
this.logging.log(`Request id: ${response.$response.requestId}`)
156156

157157
const aggregatedCodeScanIssueList = await this.mapToAggregatedList(response.codeAnalysisFindings, projectPath)

0 commit comments

Comments
 (0)