Adding Response headers to CodeScan and reusing the uploadCode function

Author: laileni-aws

packages/core/src/amazonqFeatureDev/session/sessionState.ts

@@ -12,7 +12,7 @@ import { getLogger } from '../../shared/logger'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { VirtualFileSystem } from '../../shared/virtualFilesystem'
 import { VirtualMemoryFile } from '../../shared/virtualMemoryFile'
-import { featureDevScheme } from '../constants'
+import { featureDevScheme, featureName } from '../constants'
 import { IllegalStateTransition, UserMessageNotFoundError } from '../errors'
 import {
     CurrentWsFolders,
@@ -73,8 +73,17 @@ export class PrepareRefinementState implements Omit<SessionState, 'uploadId'> {
                 zipFileChecksum,
                 zipFileBuffer.length
             )
+            const requestHeaders = {
+                'Content-Type': 'application/zip',
+                'Content-Length': String(zipFileBuffer.length),
+                'x-amz-checksum-sha256': zipFileChecksum,
+                ...(kmsKeyArn && {
+                    'x-amz-server-side-encryption-aws-kms-key-id': kmsKeyArn,
+                    'x-amz-server-side-encryption': 'aws:kms',
+                }),
+            }
 
-            await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn)
+            await uploadCode(uploadUrl, zipFileBuffer, requestHeaders, featureName)
             return uploadId
         })
         const nextState = new RefinementState({ ...this.config, uploadId }, this.approach, this.tabID, 0)
@@ -472,8 +481,17 @@ export class PrepareCodeGenState implements SessionState {
                 zipFileChecksum,
                 zipFileBuffer.length
             )
+            const requestHeaders = {
+                'Content-Type': 'application/zip',
+                'Content-Length': String(zipFileBuffer.length),
+                'x-amz-checksum-sha256': zipFileChecksum,
+                ...(kmsKeyArn && {
+                    'x-amz-server-side-encryption-aws-kms-key-id': kmsKeyArn,
+                    'x-amz-server-side-encryption': 'aws:kms',
+                }),
+            }
 
-            await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn)
+            await uploadCode(uploadUrl, zipFileBuffer, requestHeaders, featureName)
             return uploadId
         })
         this.uploadId = uploadId

packages/core/src/amazonqFeatureDev/util/upload.ts

@@ -5,7 +5,6 @@
 
 import request, { RequestError } from '../../common/request'
 import { getLogger } from '../../shared/logger/logger'
-import { featureName } from '../constants'
 
 import { UploadCodeError } from '../errors'
 
@@ -14,19 +13,11 @@ import { UploadCodeError } from '../errors'
  *
  * uses a presigned url and files checksum to transfer data to s3 through http.
  */
-export async function uploadCode(url: string, buffer: Buffer, checksumSha256: string, kmsKeyArn?: string) {
+export async function uploadCode(url: string, buffer: Buffer, requestHeaders: any, featureName: string) {
     try {
         await request.fetch('PUT', url, {
             body: buffer,
-            headers: {
-                'Content-Type': 'application/zip',
-                'Content-Length': String(buffer.length),
-                'x-amz-checksum-sha256': checksumSha256,
-                ...(kmsKeyArn && {
-                    'x-amz-server-side-encryption-aws-kms-key-id': kmsKeyArn,
-                    'x-amz-server-side-encryption': 'aws:kms',
-                }),
-            },
+            headers: requestHeaders,
         }).response
     } catch (e: any) {
         getLogger().error(`${featureName}: failed to upload code to s3: ${(e as Error).message}`)

packages/core/src/codewhisperer/client/service-2.json

@@ -583,6 +583,9 @@
                 },
                 "kmsKeyArn": {
                     "shape": "ResourceArn"
+                },
+                "requestHeaders": {
+                    "shape": "RequestHeaders"
                 }
             }
         },
@@ -1086,6 +1089,27 @@
             "max": 1224,
             "min": 0
         },
+        "RequestHeaderKey": {
+            "type": "string",
+            "max": 64,
+            "min": 1
+        },
+        "RequestHeaderValue": {
+            "type": "string",
+            "max": 256,
+            "min": 1
+        },
+        "RequestHeaders": {
+            "type": "map",
+            "key": {
+                "shape": "RequestHeaderKey"
+            },
+            "value": {
+                "shape": "RequestHeaderValue"
+            },
+            "max": 16,
+            "min": 1
+        },
         "ResourceNotFoundException": {
             "type": "structure",
             "required": ["message"],

packages/core/src/codewhisperer/client/user-service-2.json

@@ -598,7 +598,10 @@
             "members": {
                 "uploadId": { "shape": "UploadId" },
                 "uploadUrl": { "shape": "PreSignedUrl" },
-                "kmsKeyArn": { "shape": "ResourceArn" }
+                "kmsKeyArn": { "shape": "ResourceArn" },
+                "requestHeaders": {
+                    "shape": "RequestHeaders"
+                }
             }
         },
         "CursorState": {
@@ -1137,6 +1140,27 @@
             "min": 0,
             "pattern": "arn:([-.a-z0-9]{1,63}:){2}([-.a-z0-9]{0,63}:){2}([a-zA-Z0-9-_:/]){1,1023}"
         },
+        "RequestHeaderKey": {
+            "type": "string",
+            "max": 64,
+            "min": 1
+        },
+        "RequestHeaderValue": {
+            "type": "string",
+            "max": 256,
+            "min": 1
+        },
+        "RequestHeaders": {
+            "type": "map",
+            "key": {
+                "shape": "RequestHeaderKey"
+            },
+            "value": {
+                "shape": "RequestHeaderValue"
+            },
+            "max": 16,
+            "min": 1
+        },
         "ResourceNotFoundException": {
             "type": "structure",
             "required": ["message"],

packages/core/src/codewhisperer/service/securityScanHandler.ts

@@ -14,10 +14,10 @@ import { RawCodeScanIssue } from '../models/model'
 import * as crypto from 'crypto'
 import path = require('path')
 import { pageableToCollection } from '../../shared/utilities/collectionUtils'
-import { ArtifactMap, CreateUploadUrlRequest, CreateUploadUrlResponse } from '../client/codewhispereruserclient'
+import { ArtifactMap, CreateUploadUrlRequest } from '../client/codewhispereruserclient'
 import { Truncation } from '../util/dependencyGraph/dependencyGraph'
 import { TelemetryHelper } from '../util/telemetryHelper'
-import request from '../../common/request'
+import { uploadCode } from '../../amazonqFeatureDev/util/upload'
 
 export async function listScanResults(
     client: DefaultCodeWhispererClient,
@@ -154,7 +154,12 @@ export async function getPresignedUrlAndUpload(client: DefaultCodeWhispererClien
     getLogger().verbose(`Request id: ${srcResp.$response.requestId}`)
     getLogger().verbose(`Complete Getting presigned Url for uploading src context.`)
     getLogger().verbose(`Uploading src context...`)
-    await uploadArtifactToS3(truncation.zipFilePath, srcResp)
+    await uploadCode(
+        srcResp.uploadUrl,
+        readFileSync(truncation.zipFilePath),
+        srcResp.requestHeaders,
+        'QCA SecurityScan'
+    )
     getLogger().verbose(`Complete uploading src context.`)
     const artifactMap: ArtifactMap = {
         SourceCode: srcResp.uploadId,
@@ -173,23 +178,3 @@ export function throwIfCancelled() {
         throw new CodeScanStoppedError()
     }
 }
-
-export async function uploadArtifactToS3(fileName: string, resp: CreateUploadUrlResponse) {
-    const encryptionContext = `{"uploadId":"${resp.uploadId}"}`
-    const headersObj: Record<string, string> = {
-        'Content-MD5': getMd5(fileName),
-        'x-amz-server-side-encryption': 'aws:kms',
-        'Content-Type': 'application/zip',
-        'x-amz-server-side-encryption-context': Buffer.from(encryptionContext, 'utf8').toString('base64'),
-    }
-
-    if (resp.kmsKeyArn !== '' && resp.kmsKeyArn !== undefined) {
-        headersObj['x-amz-server-side-encryption-aws-kms-key-id'] = resp.kmsKeyArn
-    }
-
-    const response = await request.fetch('PUT', resp.uploadUrl, {
-        body: readFileSync(fileName),
-        headers: headersObj,
-    }).response
-    getLogger().debug(`StatusCode: ${response.status}, Text: ${response.statusText}`)
-}