Revert "Adding Response headers to CodeScan and reusing the uploadCode function"

This reverts commit be4f6a5.

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, featureName } from '../constants'
+import { featureDevScheme } from '../constants'
 import { IllegalStateTransition, UserMessageNotFoundError } from '../errors'
 import {
     CurrentWsFolders,
@@ -73,17 +73,8 @@ 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, requestHeaders, featureName)
+            await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn)
             return uploadId
         })
         const nextState = new RefinementState({ ...this.config, uploadId }, this.approach, this.tabID, 0)
@@ -481,17 +472,8 @@ 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, requestHeaders, featureName)
+            await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn)
             return uploadId
         })
         this.uploadId = uploadId

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

@@ -5,6 +5,7 @@
 
 import request, { RequestError } from '../../common/request'
 import { getLogger } from '../../shared/logger/logger'
+import { featureName } from '../constants'
 
 import { UploadCodeError } from '../errors'
 
@@ -13,11 +14,19 @@ 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, requestHeaders: any, featureName: string) {
+export async function uploadCode(url: string, buffer: Buffer, checksumSha256: string, kmsKeyArn?: string) {
     try {
         await request.fetch('PUT', url, {
             body: buffer,
-            headers: requestHeaders,
+            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',
+                }),
+            },
         }).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,9 +583,6 @@
                 },
                 "kmsKeyArn": {
                     "shape": "ResourceArn"
-                },
-                "requestHeaders": {
-                    "shape": "RequestHeaders"
                 }
             }
         },
@@ -1089,27 +1086,6 @@
             "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,10 +598,7 @@
             "members": {
                 "uploadId": { "shape": "UploadId" },
                 "uploadUrl": { "shape": "PreSignedUrl" },
-                "kmsKeyArn": { "shape": "ResourceArn" },
-                "requestHeaders": {
-                    "shape": "RequestHeaders"
-                }
+                "kmsKeyArn": { "shape": "ResourceArn" }
             }
         },
         "CursorState": {
@@ -1140,27 +1137,6 @@
             "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 } from '../client/codewhispereruserclient'
+import { ArtifactMap, CreateUploadUrlRequest, CreateUploadUrlResponse } from '../client/codewhispereruserclient'
 import { Truncation } from '../util/dependencyGraph/dependencyGraph'
 import { TelemetryHelper } from '../util/telemetryHelper'
-import { uploadCode } from '../../amazonqFeatureDev/util/upload'
+import request from '../../common/request'
 
 export async function listScanResults(
     client: DefaultCodeWhispererClient,
@@ -154,12 +154,7 @@ 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 uploadCode(
-        srcResp.uploadUrl,
-        readFileSync(truncation.zipFilePath),
-        srcResp.requestHeaders,
-        'QCA SecurityScan'
-    )
+    await uploadArtifactToS3(truncation.zipFilePath, srcResp)
     getLogger().verbose(`Complete uploading src context.`)
     const artifactMap: ArtifactMap = {
         SourceCode: srcResp.uploadId,
@@ -178,3 +173,23 @@ 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}`)
+}