fix(toolkit-lib): MFA token cannot be provided through IoHost (#508)

Relates to #396 

Considering this a bug fix, since it is currently not possible to
integrate an IoHost with this. While this does affect the CLI, the DX is
virtually unchanged. I run the auth test suite to ensure everything is
still working as expected.

Before:

<img width="589" alt="before"
src="https://github.com/user-attachments/assets/cb5d7d39-ed81-4bef-b859-3eb26d964f59"
/>

After:

<img width="619" alt="after"
src="https://github.com/user-attachments/assets/bf410b53-44c5-4657-ab6b-890ca63fb508"
/>


---
By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache-2.0 license

Author: mrgrain

.github/workflows/integ.yml

@@ -1706,6 +1706,7 @@ new CdkCliIntegTestsWorkflow(repo, {
     cdkAliasPackage.name,
     cliInteg.name,
     toolkitLib.name,
+    cliPluginContract.name,
   ],
 
   allowUpstreamVersions: [

.projenrc.ts

@@ -102,3 +102,4 @@ group: Documents
 | `CDK_SDK_I0000` | An SDK debug message. | `debug` | n/a |
 | `CDK_SDK_W0000` | An SDK warning message. | `warn` | n/a |
 | `CDK_SDK_I0100` | An SDK trace. SDK traces are emitted as traces to the IoHost, but contain the original SDK logging level. | `trace` | {@link SdkTrace} |
+| `CDK_SDK_I1100` | Get an MFA token for an MFA device. | `info` | {@link MfaTokenRequest} |

packages/@aws-cdk/toolkit-lib/docs/message-registry.md

@@ -4,7 +4,6 @@ import { createCredentialChain, fromEnv, fromIni, fromNodeProviderChain } from '
 import { MetadataService } from '@aws-sdk/ec2-metadata-service';
 import type { NodeHttpHandlerOptions } from '@smithy/node-http-handler';
 import { loadSharedConfigFiles } from '@smithy/shared-ini-file-loader';
-import * as promptly from 'promptly';
 import { makeCachingProvider } from './provider-caching';
 import { ProxyAgentProvider } from './proxy-agent';
 import type { ISdkLogger } from './sdk-logger';
@@ -219,18 +218,18 @@ export class AwsCliCompatible {
   }
 
   /**
-   * Ask user for MFA token for given serial
+   * Ask user for MFA token for given MFA device
    *
    * Result is send to callback function for SDK to authorize the request
    */
-  private async tokenCodeFn(serialArn: string): Promise<string> {
+  private async tokenCodeFn(deviceArn: string): Promise<string> {
     const debugFn = (msg: string, ...args: any[]) => this.ioHelper.notify(IO.DEFAULT_SDK_DEBUG.msg(format(msg, ...args)));
-    await debugFn('Require MFA token for serial ARN', serialArn);
+    await debugFn('Require MFA token from MFA device with ARN', deviceArn);
     try {
-      const token: string = await promptly.prompt(`MFA token for ${serialArn}: `, {
-        trim: true,
-        default: '',
-      });
+      const token: string = await this.ioHelper.requestResponse(IO.CDK_SDK_I1100.req(`MFA token for ${deviceArn}`, {
+        deviceArn,
+      }, ''));
+
       await debugFn('Successfully got MFA token from user');
       return token;
     } catch (err: any) {

packages/@aws-cdk/toolkit-lib/lib/api/aws-auth/awscli-compatible.ts

@@ -113,7 +113,11 @@ export interface IoRequestMaker<T, U> extends MessageInfo {
   /**
    * Create a message for this code, with or without payload.
    */
-  req: [T] extends [AbsentData] ? (message: string) => ActionLessMessage<AbsentData> : (message: string, data: T) => ActionLessRequest<T, U>;
+  req: [T] extends [AbsentData]
+    ? (message: string) => ActionLessMessage<AbsentData>
+    : [U] extends [boolean]
+      ? (message: string, data: T) => ActionLessRequest<T, U>
+      : (message: string, data: T, defaultResponse: U) => ActionLessRequest<T, U>;
 }
 
 /**
@@ -143,3 +147,24 @@ export const confirm = <T extends object = ImpossibleType>(details: Required<Omi
   ...details,
   defaultResponse: true,
 });
+
+/**
+ * An open ended question with a string answer, typically provided on-demand by a user.
+ */
+export function question<T>(details: CodeInfo): IoRequestMaker<T, string> {
+  const level = 'info';
+  const maker = (text: string, data: T, defaultResponse: string) => ({
+    time: new Date(),
+    level,
+    code: details.code,
+    message: text,
+    data,
+    defaultResponse,
+  } as ActionLessRequest<T, string>);
+
+  return {
+    ...details,
+    level,
+    req: maker as any,
+  };
+}

packages/@aws-cdk/toolkit-lib/lib/api/io/private/message-maker.ts

@@ -11,7 +11,7 @@ import type { StackDetailsPayload } from '../../../payloads/list';
 import type { CloudWatchLogEvent, CloudWatchLogMonitorControlEvent } from '../../../payloads/logs-monitor';
 import type { RefactorResult } from '../../../payloads/refactor';
 import type { StackRollbackProgress } from '../../../payloads/rollback';
-import type { SdkTrace } from '../../../payloads/sdk-trace';
+import type { MfaTokenRequest, SdkTrace } from '../../../payloads/sdk';
 import type { StackActivity, StackMonitoringControlEvent } from '../../../payloads/stack-activity';
 import type { StackSelectionDetails } from '../../../payloads/synth';
 import type { AssemblyData, ConfirmationRequest, ContextProviderMessageSource, Duration, ErrorPayload, StackAndAssemblyData } from '../../../payloads/types';
@@ -517,6 +517,11 @@ export const IO = {
     description: 'An SDK trace. SDK traces are emitted as traces to the IoHost, but contain the original SDK logging level.',
     interface: 'SdkTrace',
   }),
+  CDK_SDK_I1100: make.question<MfaTokenRequest>({
+    code: 'CDK_SDK_I1100',
+    description: 'Get an MFA token for an MFA device.',
+    interface: 'MfaTokenRequest',
+  }),
 };
 
 //////////////////////////////////////////////////////////////////////////////////////////

packages/@aws-cdk/toolkit-lib/lib/api/io/private/messages.ts

@@ -2,7 +2,7 @@ export * from './bootstrap-environment-progress';
 export * from './deploy';
 export * from './destroy';
 export * from './list';
-export * from './sdk-trace';
+export * from './sdk';
 export * from './context';
 export * from './rollback';
 export * from './stack-activity';

packages/@aws-cdk/toolkit-lib/lib/payloads/index.ts

@@ -1,3 +1,5 @@
+import type { DataRequest } from './types';
+
 /**
  * An SDK logging trace.
  *
@@ -19,3 +21,13 @@ export interface SdkTrace {
    */
   readonly content: any;
 }
+
+/**
+ * Get an MFA token for an MFA device.
+ */
+export interface MfaTokenRequest extends DataRequest {
+  /**
+   * The ARN of the MFA device a token is required for.
+   */
+  readonly deviceArn: string;
+}

packages/@aws-cdk/toolkit-lib/lib/payloads/sdk-trace.ts renamed to packages/@aws-cdk/toolkit-lib/lib/payloads/sdk.ts

@@ -99,6 +99,18 @@ export interface ConfirmationRequest {
   readonly concurrency?: number;
 }
 
+/**
+ * A generic request for data
+ */
+export interface DataRequest {
+  /**
+   * An optional description of the expected response
+   * Provides additional details on what the response can be.
+   * This can be treated as a direct instruction to end-users when prompting for input.
+   */
+  responseDescription?: string;
+}
+
 export interface ContextProviderMessageSource {
   /**
    * The name of the context provider sending the message

packages/@aws-cdk/toolkit-lib/lib/payloads/types.ts

@@ -2,6 +2,9 @@ import type { IIoHost, IoMessage, IoMessageLevel, IoRequest } from '../../lib/ap
 import type { IoHelper } from '../../lib/api/io/private';
 import { asIoHelper, isMessageRelevantForLevel } from '../../lib/api/io/private';
 
+type MessageMock = jest.Mock<void, [IoMessage<any>]>;
+type RequestMock<U extends any> = jest.Mock<U, [IoRequest<any, U>]>;
+
 /**
  * An implementation of `IIoHost` that records messages,
  * lets you assert on what was logged and can be spied on.
@@ -22,8 +25,8 @@ import { asIoHelper, isMessageRelevantForLevel } from '../../lib/api/io/private'
  */
 export class TestIoHost implements IIoHost {
   public messages: Array<IoMessage<unknown>> = [];
-  public readonly notifySpy: jest.Mock<any, any, any>;
-  public readonly requestSpy: jest.Mock<any, any, any>;
+  public readonly notifySpy: MessageMock;
+  public readonly requestSpy: RequestMock<any>;
 
   constructor(public level: IoMessageLevel = 'info') {
     this.notifySpy = jest.fn();
@@ -49,10 +52,11 @@ export class TestIoHost implements IIoHost {
   }
 
   public async requestResponse<T, U>(msg: IoRequest<T, U>): Promise<U> {
+    let spyResponse;
     if (isMessageRelevantForLevel(msg, this.level)) {
-      this.requestSpy(msg);
+      spyResponse = await this.requestSpy(msg);
     }
-    return msg.defaultResponse;
+    return spyResponse ?? msg.defaultResponse;
   }
 
   public expectMessage(m: { containing: string; level?: IoMessageLevel }) {

packages/@aws-cdk/toolkit-lib/test/_helpers/test-io-host.ts

@@ -15,7 +15,6 @@
 import * as os from 'os';
 import * as cxapi from '@aws-cdk/cx-api';
 import * as fromEnv from '@aws-sdk/credential-provider-env';
-import * as promptly from 'promptly';
 import * as uuid from 'uuid';
 import type { RegisterRoleOptions, RegisterUserOptions } from './fake-sts';
 import { FakeSts } from './fake-sts';
@@ -335,7 +334,7 @@ describe('with intercepted network calls', () => {
 
     test.each(providersForMfa)('mfa_serial in profile will ask user for token', async (metaProvider: () => Promise<SdkProvider>) => {
       // GIVEN
-      const mockPrompt = jest.spyOn(promptly, 'prompt').mockResolvedValue('1234');
+      ioHost.requestSpy.mockImplementation((msg) => msg.code === 'CDK_SDK_I1100' ? '1234' : undefined);
 
       prepareCreds({
         fakeSts,
@@ -362,7 +361,10 @@ describe('with intercepted network calls', () => {
 
       // Make sure the MFA mock was called during this test, only once
       // (Credentials need to remain cached)
-      expect(mockPrompt).toHaveBeenCalledTimes(1);
+      expect(ioHost.requestSpy).toHaveBeenCalledTimes(1);
+      expect(ioHost.requestSpy).toHaveBeenCalledWith(expect.objectContaining({
+        code: 'CDK_SDK_I1100',
+      }));
     });
   });
 

packages/@aws-cdk/toolkit-lib/test/api/aws-auth/sdk-provider.test.ts

@@ -357,10 +357,12 @@ export class CliIoHost implements IIoHost {
       const data: {
         motivation?: string;
         concurrency?: number;
+        responseDescription?: string;
       } = msg.data ?? {};
 
       const motivation = data.motivation ?? 'User input is needed';
       const concurrency = data.concurrency ?? 0;
+      const responseDescription = data.responseDescription;
 
       // only talk to user if STDIN is a terminal (otherwise, fail)
       if (!this.isTTY) {
@@ -392,8 +394,10 @@ export class CliIoHost implements IIoHost {
 
       // Asking for a specific value
       const prompt = extractPromptInfo(msg);
-      const answer = await promptly.prompt(`${chalk.cyan(msg.message)} (${prompt.default})`, {
+      const desc = responseDescription ?? prompt.default;
+      const answer = await promptly.prompt(`${chalk.cyan(msg.message)}${desc ? ` (${desc})` : ''}`, {
         default: prompt.default,
+        trim: true,
       });
       return prompt.convertAnswer(answer);
     });