feat(cloudformation): enforce s3 upload for large templates and artifacts

Author: Zee2413

packages/core/src/awsService/cloudformation/commands/cfnCommands.ts

@@ -402,6 +402,15 @@ type UserInputtedTemplateParameters = {
     s3Key?: string
 }
 
+async function promptForS3Key(templateUri: string): Promise<string | undefined> {
+    const fileName = templateUri.split('/').pop()
+    const timestamp = Date.now()
+    const defaultKey = fileName
+        ? `${fileName.split('.')[0]}-${timestamp}.${fileName.split('.').pop()}`
+        : `template-${timestamp}.yaml`
+    return getS3Key(defaultKey)
+}
+
 async function changeSetSteps(
     client: LanguageClient,
     documentManager: DocumentManager,
@@ -432,31 +441,38 @@ async function changeSetSteps(
     // Ask user if they want to upload to S3
     let s3Bucket: string | undefined
     let s3Key: string | undefined
-    const uploadChoice = await shouldUploadToS3()
-    if (uploadChoice === undefined) {
-        return // User chose to configure settings, exit command
-    }
-    if (uploadChoice) {
-        s3Bucket = await getS3Bucket()
+    const templateExceedsSizeLimit = documentManager.requiresS3Upload(templateUri)
+
+    if (templateExceedsSizeLimit || hasArtifacts) {
+        const reason = templateExceedsSizeLimit
+            ? 'S3 bucket is required because template exceeds the 51,200 byte CloudFormation limit'
+            : 'S3 bucket is required because template contains artifacts that need to be uploaded to S3'
+        s3Bucket = await getS3Bucket(reason)
         if (!s3Bucket) {
             return
         }
 
-        const fileName = templateUri.split('/').pop()
-        const timestamp = Date.now()
-        const fileNameWithTimestamp = fileName
-            ? `${fileName.split('.')[0]}-${timestamp}.${fileName.split('.').pop()}`
-            : `template-${timestamp}.yaml`
-        s3Key = await getS3Key(fileNameWithTimestamp)
-        if (!s3Key) {
-            return
+        if (templateExceedsSizeLimit) {
+            s3Key = await promptForS3Key(templateUri)
+            if (!s3Key) {
+                return
+            }
         }
-    } else if (hasArtifacts) {
-        s3Bucket = await getS3Bucket(
-            'S3 bucket is required because template contains artifacts that need to be uploaded to S3'
-        )
-        if (!s3Bucket) {
-            return
+    } else {
+        const uploadChoice = await shouldUploadToS3()
+        if (uploadChoice === undefined) {
+            return // User chose to configure settings, exit command
+        }
+        if (uploadChoice) {
+            s3Bucket = await getS3Bucket()
+            if (!s3Bucket) {
+                return
+            }
+
+            s3Key = await promptForS3Key(templateUri)
+            if (!s3Key) {
+                return
+            }
         }
     }
 

packages/core/src/awsService/cloudformation/documents/documentManager.ts

@@ -5,6 +5,7 @@
 
 import { NotificationType } from 'vscode-languageserver-protocol'
 import { LanguageClient } from 'vscode-languageclient/node'
+import { getLogger } from '../../../shared/logger/logger'
 
 export type DocumentMetadata = {
     uri: string
@@ -15,13 +16,15 @@ export type DocumentMetadata = {
     languageId: string
     version: number
     lineCount: number
+    sizeBytes: number
 }
 
 const DocumentsMetadataNotification = new NotificationType<DocumentMetadata[]>('aws/documents/metadata')
 
 type DocumentsChangeListener = (documents: DocumentMetadata[]) => void
 
 export class DocumentManager {
+    private static readonly cfnTemplateBodyMaxBytes = 51_200
     private documents: DocumentMetadata[] = []
     private readonly listeners: DocumentsChangeListener[] = []
 
@@ -41,4 +44,15 @@ export class DocumentManager {
     get() {
         return [...this.documents]
     }
+
+    requiresS3Upload(uri: string): boolean {
+        const doc = this.documents.find((d) => d.uri === uri)
+        if (!doc) {
+            getLogger('awsCfnLsp').warn(
+                `Document metadata not found for URI: ${uri}. Assuming no s3 upload required may lead to deployment failure.`
+            )
+            return false
+        }
+        return doc.sizeBytes > DocumentManager.cfnTemplateBodyMaxBytes
+    }
 }

packages/core/src/test/awsService/cloudformation/documents/documentManager.test.ts

@@ -0,0 +1,68 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as sinon from 'sinon'
+import { DocumentManager, DocumentMetadata } from '../../../../awsService/cloudformation/documents/documentManager'
+import { LanguageClient } from 'vscode-languageclient/node'
+
+function createMetadata(overrides: Partial<DocumentMetadata> = {}): DocumentMetadata {
+    return {
+        uri: 'file:///template.yaml',
+        fileName: 'template.yaml',
+        ext: 'yaml',
+        type: 'yaml',
+        cfnType: 'template',
+        languageId: 'yaml',
+        version: 1,
+        lineCount: 10,
+        sizeBytes: 100,
+        ...overrides,
+    }
+}
+
+describe('DocumentManager', function () {
+    let sandbox: sinon.SinonSandbox
+    let documentManager: DocumentManager
+    let notificationCallback: (docs: DocumentMetadata[]) => void
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+        const fakeClient = {
+            onNotification: (_type: unknown, callback: (docs: DocumentMetadata[]) => void) => {
+                notificationCallback = callback
+            },
+        } as unknown as LanguageClient
+        documentManager = new DocumentManager(fakeClient)
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+    })
+
+    describe('requiresS3Upload', function () {
+        it('returns false when document is not found', function () {
+            notificationCallback([createMetadata()])
+
+            assert.strictEqual(documentManager.requiresS3Upload('file:///nonexistent.yaml'), false)
+        })
+
+        it('returns false when template size is at the limit', function () {
+            notificationCallback([createMetadata({ sizeBytes: 51_200 })])
+
+            assert.strictEqual(documentManager.requiresS3Upload('file:///template.yaml'), false)
+        })
+
+        it('returns true when template size exceeds the limit', function () {
+            notificationCallback([createMetadata({ sizeBytes: 51_201 })])
+
+            assert.strictEqual(documentManager.requiresS3Upload('file:///template.yaml'), true)
+        })
+
+        it('returns false when no documents have been received', function () {
+            assert.strictEqual(documentManager.requiresS3Upload('file:///template.yaml'), false)
+        })
+    })
+})