Merge master into feature/LSP-gamma

Author: aws-toolkit-automation

packages/core/src/awsService/sagemaker/uriHandlers.ts

@@ -9,10 +9,28 @@ import { deeplinkConnect } from './commands'
 import { ExtContext } from '../../shared/extensions'
 import { telemetry } from '../../shared/telemetry/telemetry'
 
+const amzHeaders = [
+    'X-Amz-Security-Token',
+    'X-Amz-Algorithm',
+    'X-Amz-Date',
+    'X-Amz-SignedHeaders',
+    'X-Amz-Credential',
+    'X-Amz-Expires',
+    'X-Amz-Signature',
+] as const
+
 export function register(ctx: ExtContext) {
     async function connectHandler(params: ReturnType<typeof parseConnectParams>) {
         await telemetry.sagemaker_deeplinkConnect.run(async () => {
-            const wsUrl = `${params.ws_url}&cell-number=${params['cell-number']}`
+            let wsUrl = `${params.ws_url}&cell-number=${encodeURIComponent(params['cell-number'])}`
+
+            for (const header of amzHeaders) {
+                const value = params[header]
+                if (value) {
+                    wsUrl += `&${header}=${encodeURIComponent(value)}`
+                }
+            }
+
             await deeplinkConnect(
                 ctx,
                 params.connection_identifier,
@@ -54,6 +72,7 @@ export function parseHyperpodConnectParams(query: SearchParams) {
     const optionalParams = query.getFromKeys('workspaceName', 'namespace', 'eksClusterArn')
     return { ...requiredParams, ...optionalParams }
 }
+
 export function parseConnectParams(query: SearchParams) {
     const requiredParams = query.getFromKeysOrThrow(
         'connection_identifier',
@@ -66,5 +85,6 @@ export function parseConnectParams(query: SearchParams) {
     )
     const optionalParams = query.getFromKeys('app_type')
 
-    return { ...requiredParams, ...optionalParams }
+    const amzHeaderParams = query.getFromKeys(...amzHeaders)
+    return { ...requiredParams, ...optionalParams, ...amzHeaderParams }
 }

packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts

@@ -76,4 +76,77 @@ describe('SageMaker URI handler', function () {
         assert.ok(deeplinkConnectStub.calledOnce)
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined)
     })
+
+    it('properly encodes cell-number with spaces and special characters', async function () {
+        const params = {
+            connection_identifier: 'abc123',
+            domain: 'my-domain',
+            user_profile: 'me',
+            session: 'sess-xyz',
+            ws_url: 'wss://example.com',
+            'cell-number': 'test/data with spaces',
+            token: 'my-token',
+        }
+
+        const uri = createConnectUri(params)
+        await handler.handleUri(uri)
+
+        assert.ok(deeplinkConnectStub.calledOnce)
+        // Verify cell-number is properly encoded
+        const expectedUrl = 'wss://example.com&cell-number=test%2Fdata%20with%20spaces'
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], expectedUrl)
+    })
+
+    it('includes AMZ headers in WebSocket URL when provided', async function () {
+        const params = {
+            connection_identifier: 'abc123',
+            domain: 'my-domain',
+            user_profile: 'me',
+            session: 'sess-xyz',
+            ws_url: 'wss://example.com',
+            'cell-number': 'test123',
+            token: 'my-token',
+            'X-Amz-Security-Token': 'fake/token+with=special',
+            'X-Amz-Algorithm': 'AWS4-HMAC-SHA256',
+            'X-Amz-Date': '20240101T120000Z',
+            'X-Amz-SignedHeaders': 'host',
+            'X-Amz-Credential': 'AKIATEST/20240101/us-west-2/ssmmessages/aws4_request',
+            'X-Amz-Expires': '60',
+            'X-Amz-Signature': 'fakesignature123',
+        }
+
+        const uri = createConnectUri(params)
+        await handler.handleUri(uri)
+
+        assert.ok(deeplinkConnectStub.calledOnce)
+        const actualUrl = deeplinkConnectStub.firstCall.args[3]
+
+        // Verify all AMZ headers are included and properly encoded
+        assert.ok(actualUrl.includes('cell-number=test123'))
+        assert.ok(actualUrl.includes('X-Amz-Security-Token=fake%2Ftoken%2Bwith%3Dspecial'))
+        assert.ok(actualUrl.includes('X-Amz-Algorithm=AWS4-HMAC-SHA256'))
+        assert.ok(actualUrl.includes('X-Amz-Date=20240101T120000Z'))
+        assert.ok(actualUrl.includes('X-Amz-SignedHeaders=host'))
+        assert.ok(actualUrl.includes('X-Amz-Credential=AKIATEST%2F20240101%2Fus-west-2%2Fssmmessages%2Faws4_request'))
+        assert.ok(actualUrl.includes('X-Amz-Expires=60'))
+        assert.ok(actualUrl.includes('X-Amz-Signature=fakesignature123'))
+    })
+
+    it('works without AMZ headers', async function () {
+        const params = {
+            connection_identifier: 'abc123',
+            domain: 'my-domain',
+            user_profile: 'me',
+            session: 'sess-xyz',
+            ws_url: 'wss://example.com',
+            'cell-number': 'simple',
+            token: 'my-token',
+        }
+
+        const uri = createConnectUri(params)
+        await handler.handleUri(uri)
+
+        assert.ok(deeplinkConnectStub.calledOnce)
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], 'wss://example.com&cell-number=simple')
+    })
 })