Prevent resource state import without open cfn template or empty yaml (#306)

* prevent resource state import without open cfn template

* allow state insertion into empty documents

Author: Zee2413

src/handlers/ResourceHandler.ts

@@ -1,8 +1,9 @@
 import { randomUUID } from 'crypto';
-import { ServerRequestHandler } from 'vscode-languageserver';
+import { ResponseError, ServerRequestHandler } from 'vscode-languageserver';
 import { RequestHandler } from 'vscode-languageserver/node';
 import { TopLevelSection } from '../context/ContextType';
 import { getEntityMap } from '../context/SectionContextBuilder';
+import { CloudFormationFileType } from '../document/Document';
 import { parseResourceTypeName } from '../resourceState/ResourceStateParser';
 import {
     ResourceTypesResult,
@@ -79,9 +80,22 @@ export function listResourcesHandler(
 
 export function importResourceStateHandler(
     components: ServerComponents,
-): ServerRequestHandler<ResourceStateParams, ResourceStateResult, never, void> {
+): RequestHandler<ResourceStateParams, ResourceStateResult, void> {
     return async (params: ResourceStateParams): Promise<ResourceStateResult> => {
         components.usageTracker.track(EventType.DidImportResources);
+        const doc = components.documentManager.get(params.textDocument.uri);
+        if (!doc) {
+            const msg = `${params.purpose} failed: ${params.textDocument.uri} not found`;
+            log.error(msg);
+            throw new ResponseError(500, msg); // all open TextDocuments should be registered by protocol
+        }
+
+        if (!doc.isTemplate() && doc.cfnFileType !== CloudFormationFileType.Empty) {
+            throw new ResponseError(
+                400,
+                `${params.purpose} failed: ${params.textDocument.uri} is not a valid CloudFormation template`,
+            );
+        }
         return await components.resourceStateImporter.importResourceState(params);
     };
 }

tst/unit/handlers/ResourceHandler.test.ts

@@ -1,10 +1,18 @@
+import { stubInterface } from 'ts-sinon';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
-import { CancellationToken } from 'vscode-languageserver-protocol';
+import { CancellationToken, RequestHandler } from 'vscode-languageserver-protocol';
 import { getEntityMap } from '../../../src/context/SectionContextBuilder';
+import { CloudFormationFileType, Document } from '../../../src/document/Document';
 import {
     getManagedResourceStackTemplateHandler,
+    importResourceStateHandler,
     removeResourceTypeHandler,
 } from '../../../src/handlers/ResourceHandler';
+import {
+    ResourceStateParams,
+    ResourceStatePurpose,
+    ResourceStateResult,
+} from '../../../src/resourceState/ResourceStateTypes';
 import { GetStackTemplateParams } from '../../../src/stacks/StackRequestType';
 import { createMockComponents } from '../../utils/MockServerComponents';
 
@@ -156,3 +164,43 @@ describe('ResourceHandler - removeResourceTypeHandler', () => {
         expect(() => handler(undefined as any)).toThrow();
     });
 });
+
+describe('ResourceHandler - importResourceStateHandler', () => {
+    let mockComponents: ReturnType<typeof createMockComponents>;
+    let handler: RequestHandler<ResourceStateParams, ResourceStateResult, void>;
+    const params = {
+        textDocument: { uri: 'docUri' },
+        resourceSelections: [
+            {
+                resourceType: 'AWS::S3::Bucket',
+                resourceIdentifiers: ['bucket1234'],
+            },
+        ],
+        purpose: ResourceStatePurpose.IMPORT,
+    };
+
+    beforeEach(() => {
+        vi.clearAllMocks();
+        mockComponents = createMockComponents();
+        handler = importResourceStateHandler(mockComponents);
+    });
+
+    it('should throw error if document not found', async () => {
+        mockComponents.documentManager.get.returns(undefined);
+
+        await expect(async () => {
+            await handler(params, CancellationToken.None);
+        }).rejects.toThrow('Import failed: docUri not found');
+    });
+
+    it('should throw error if document is not a valid CloudFormation template', async () => {
+        const mockDoc = stubInterface<Document>();
+        mockDoc.isTemplate.returns(false);
+        Object.defineProperty(mockDoc, 'cfnFileType', { value: CloudFormationFileType.Other });
+        mockComponents.documentManager.get.returns(mockDoc);
+
+        await expect(async () => {
+            await handler(params, CancellationToken.None);
+        }).rejects.toThrow('Import failed: docUri is not a valid CloudFormation template');
+    });
+});