Skip to content

Commit 95457c7

Browse files
laileni-awslaileni-aws
authored
Adding request headers to uploadArtifactToS3 in CodeScan (aws#4336)
1 parent 70ea000 commit 95457c7

File tree

9 files changed

+54
-24
lines changed

9 files changed

+54
-24
lines changed

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import org.junit.Before
2020
import org.junit.Rule
2121
import org.junit.Test
2222
import org.mockito.kotlin.mock
23+
import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
2324
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
2425
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevException
2526
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,12 @@ import org.mockito.kotlin.mock
2121
import org.mockito.kotlin.times
2222
import org.mockito.kotlin.verify
2323
import org.mockito.kotlin.whenever
24+
import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
25+
import software.aws.toolkits.jetbrains.services.amazonq.ZipCreationResult
2426
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
2527
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase
2628
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
2729
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswerPart
28-
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.model.ZipCreationResult
2930
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.deleteUploadArtifact
3031
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.exportTaskAssistArchiveResult
3132
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getTaskAssistCodeGeneration

‎plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/service-2.json‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -672,7 +672,10 @@
672672
"members":{
673673
"uploadId":{"shape":"UploadId"},
674674
"uploadUrl":{"shape":"PreSignedUrl"},
675-
"kmsKeyArn":{"shape":"ResourceArn"}
675+
"kmsKeyArn":{"shape":"ResourceArn"},
676+
"requestHeaders": {
677+
"shape": "RequestHeaders"
678+
}
676679
}
677680
},
678681
"CursorState":{
@@ -1248,6 +1251,27 @@
12481251
"min":0,
12491252
"pattern":"arn:([-.a-z0-9]{1,63}:){2}([-.a-z0-9]{0,63}:){2}([a-zA-Z0-9-_:/]){1,1023}"
12501253
},
1254+
"RequestHeaderKey": {
1255+
"type": "string",
1256+
"max": 64,
1257+
"min": 1
1258+
},
1259+
"RequestHeaderValue": {
1260+
"type": "string",
1261+
"max": 256,
1262+
"min": 1
1263+
},
1264+
"RequestHeaders": {
1265+
"type": "map",
1266+
"key": {
1267+
"shape": "RequestHeaderKey"
1268+
},
1269+
"value": {
1270+
"shape": "RequestHeaderValue"
1271+
},
1272+
"max": 16,
1273+
"min": 1
1274+
},
12511275
"ResourceNotFoundException":{
12521276
"type":"structure",
12531277
"required":["message"],

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanException.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,6 @@ internal fun fileFormatNotSupported(format: String): Nothing =
2424

2525
internal fun fileTooLarge(presentableSize: String): Nothing =
2626
throw CodeWhispererCodeScanException(message("codewhisperer.codescan.file_too_large", presentableSize))
27+
28+
internal fun uploadArtifactFailedError(): Nothing =
29+
throw CodeWhispererCodeScanException(message("codewhisperer.codescan.upload_to_s3_failed"))

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -213,13 +213,12 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
213213
LOG.debug { "Fetching presigned URL for uploading $artifactType." }
214214
val createUploadUrlResponse = createUploadUrl(fileMd5, artifactType)
215215
LOG.debug { "Successfully fetched presigned URL for uploading $artifactType." }
216-
val url = createUploadUrlResponse.uploadUrl()
217216
LOG.debug { "Uploading $artifactType using the presigned URL." }
218-
uploadArtifactToS3(url, createUploadUrlResponse.uploadId(), zipFile, fileMd5, createUploadUrlResponse.kmsKeyArn())
217+
uploadArtifactToS3(createUploadUrlResponse, createUploadUrlResponse.uploadId(), zipFile, fileMd5)
219218
createUploadUrlResponse
220219
} catch (e: Exception) {
221220
LOG.error { "Security scan failed. Something went wrong uploading artifacts: ${e.message}" }
222-
throw e
221+
uploadArtifactFailedError()
223222
}
224223

225224
fun createUploadUrl(md5Content: String, artifactType: String): CreateUploadUrlResponse = clientAdaptor.createUploadUrl(
@@ -230,16 +229,22 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
230229
)
231230

232231
@Throws(IOException::class)
233-
fun uploadArtifactToS3(url: String, uploadId: String, fileToUpload: File, md5: String, kmsArn: String?) {
232+
fun uploadArtifactToS3(createUploadUrlResponse: CreateUploadUrlResponse, uploadId: String, fileToUpload: File, md5: String) {
234233
val uploadIdJson = """{"uploadId":"$uploadId"}"""
235-
HttpRequests.put(url, "application/zip").userAgent(AwsClientManager.userAgent).tuner {
236-
it.setRequestProperty(CONTENT_MD5, md5)
237-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS)
238-
it.setRequestProperty(CONTENT_TYPE, APPLICATION_ZIP)
239-
if (kmsArn?.isNotEmpty() == true) {
240-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, kmsArn)
234+
HttpRequests.put(createUploadUrlResponse.uploadUrl(), "application/zip").userAgent(AwsClientManager.userAgent).tuner {
235+
if (createUploadUrlResponse.requestHeaders() != null && createUploadUrlResponse.requestHeaders().isNotEmpty()) {
236+
for ((key, value) in createUploadUrlResponse.requestHeaders()) {
237+
it.setRequestProperty(key, value)
238+
}
239+
} else {
240+
it.setRequestProperty(CONTENT_MD5, md5)
241+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS)
242+
it.setRequestProperty(CONTENT_TYPE, APPLICATION_ZIP)
243+
if (createUploadUrlResponse.kmsKeyArn()?.isNotEmpty() == true) {
244+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, createUploadUrlResponse.kmsKeyArn())
245+
}
246+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_CONTEXT, Base64.getEncoder().encodeToString(uploadIdJson.toByteArray()))
241247
}
242-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_CONTEXT, Base64.getEncoder().encodeToString(uploadIdJson.toByteArray()))
243248
}.connect {
244249
val connection = it.connection as HttpURLConnection
245250
connection.setFixedLengthStreamingMode(fileToUpload.length())

‎plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeFileScanTest.kt‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import org.mockito.kotlin.argumentCaptor
1616
import org.mockito.kotlin.doNothing
1717
import org.mockito.kotlin.eq
1818
import org.mockito.kotlin.inOrder
19-
import org.mockito.kotlin.isNull
2019
import org.mockito.kotlin.spy
2120
import org.mockito.kotlin.stub
2221
import org.mockito.kotlin.verify
@@ -96,7 +95,7 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
9695
// Mock CodeWhispererClient needs to be setup before initializing CodeWhispererCodeScanSession
9796
codeScanSessionContext = CodeScanSessionContext(project, sessionConfigSpy)
9897
codeScanSessionSpy = spy(CodeWhispererCodeScanSession(codeScanSessionContext))
99-
doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any(), isNull())
98+
doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any())
10099

101100
mockClient.stub {
102101
onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
@@ -119,11 +118,10 @@ class CodeWhispererCodeFileScanTest : CodeWhispererCodeScanTestBase(PythonCodeIn
119118
val inOrder = inOrder(codeScanSessionSpy)
120119
inOrder.verify(codeScanSessionSpy).createUploadUrl(eq(fileMd5), eq("artifactType"))
121120
inOrder.verify(codeScanSessionSpy).uploadArtifactToS3(
122-
eq(fakeCreateUploadUrlResponse.uploadUrl()),
121+
eq(fakeCreateUploadUrlResponse),
123122
eq(fakeCreateUploadUrlResponse.uploadId()),
124123
eq(file),
125124
eq(fileMd5),
126-
eq(null)
127125
)
128126
}
129127

‎plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTest.kt‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import org.mockito.kotlin.argumentCaptor
1616
import org.mockito.kotlin.doNothing
1717
import org.mockito.kotlin.eq
1818
import org.mockito.kotlin.inOrder
19-
import org.mockito.kotlin.isNull
2019
import org.mockito.kotlin.spy
2120
import org.mockito.kotlin.stub
2221
import org.mockito.kotlin.verify
@@ -76,7 +75,7 @@ class CodeWhispererCodeScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsigh
7675
// Mock CodeWhispererClient needs to be setup before initializing CodeWhispererCodeScanSession
7776
codeScanSessionContext = CodeScanSessionContext(project, sessionConfigSpy)
7877
codeScanSessionSpy = spy(CodeWhispererCodeScanSession(codeScanSessionContext))
79-
doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any(), isNull())
78+
doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any())
8079

8180
mockClient.stub {
8281
onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
@@ -99,11 +98,10 @@ class CodeWhispererCodeScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsigh
9998
val inOrder = inOrder(codeScanSessionSpy)
10099
inOrder.verify(codeScanSessionSpy).createUploadUrl(eq(fileMd5), eq("artifactType"))
101100
inOrder.verify(codeScanSessionSpy).uploadArtifactToS3(
102-
eq(fakeCreateUploadUrlResponse.uploadUrl()),
101+
eq(fakeCreateUploadUrlResponse),
103102
eq(fakeCreateUploadUrlResponse.uploadId()),
104103
eq(file),
105-
eq(fileMd5),
106-
eq(null)
104+
eq(fileMd5)
107105
)
108106
}
109107

‎plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import org.junit.jupiter.api.assertThrows
2424
import org.mockito.kotlin.any
2525
import org.mockito.kotlin.doNothing
2626
import org.mockito.kotlin.doReturn
27-
import org.mockito.kotlin.isNull
2827
import org.mockito.kotlin.mock
2928
import org.mockito.kotlin.spy
3029
import org.mockito.kotlin.stub
@@ -276,7 +275,7 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule
276275
) {
277276
val codeScanContext = CodeScanSessionContext(project, sessionConfigSpy)
278277
val sessionMock = spy(CodeWhispererCodeScanSession(codeScanContext))
279-
doNothing().`when`(sessionMock).uploadArtifactToS3(any(), any(), any(), any(), isNull())
278+
doNothing().`when`(sessionMock).uploadArtifactToS3(any(), any(), any(), any())
280279
doNothing().`when`(sessionMock).sleepThread()
281280

282281
ToolWindowManager.getInstance(project).registerToolWindow(

‎plugins/toolkit/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -704,6 +704,7 @@ codewhisperer.codescan.stop_scan_confirm_message=Are you sure you want to stop o
704704
codewhisperer.codescan.stopping_scan=Stopping Security Scan...
705705
codewhisperer.codescan.suggested_fix_description=Why are we recommending this?
706706
codewhisperer.codescan.suggested_fix_label=Suggested code fix preview
707+
codewhisperer.codescan.upload_to_s3_failed=<html>Amazon Q is unable to upload workspace artifacts to Amazon S3 for security scans. For more information, see the Amazon Q documentation or contact your network or organization administrator.<br/> https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/security_iam_manage-access-with-policies.html</html>
707708
codewhisperer.codescan.view_scanned_files=View {0} scanned files
708709
codewhisperer.credential.login.dialog.exception.cancel_login=Login cancelled
709710
codewhisperer.credential.login.dialog.iam.description=Not supported by CodeWhisperer.

0 commit comments

Comments
 (0)