From 28ad5c4b5b8fd1dec9d4fbeb6cbcdc09b9ae4dd2 Mon Sep 17 00:00:00 2001 From: Allan Zheng Date: Mon, 24 May 2021 15:18:18 -0700 Subject: [PATCH 1/2] fix(client-sts): allow overwriting default role assumer http handler --- .../client-sts/defaultRoleAssumers.spec.ts | 118 +++++++++++++++--- clients/client-sts/defaultRoleAssumers.ts | 7 +- clients/client-sts/defaultStsRoleAssumers.ts | 14 ++- 3 files changed, 113 insertions(+), 26 deletions(-) diff --git a/clients/client-sts/defaultRoleAssumers.spec.ts b/clients/client-sts/defaultRoleAssumers.spec.ts index fc047d2ffdb9..d3eb987f9fcc 100644 --- a/clients/client-sts/defaultRoleAssumers.spec.ts +++ b/clients/client-sts/defaultRoleAssumers.spec.ts @@ -2,7 +2,36 @@ // https://github.com/aws/aws-sdk-js-v3/blob/main/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts import { HttpResponse } from "@aws-sdk/protocol-http"; import { Readable } from "stream"; -const assumeRoleResponse = ` + +const mockHandle = jest.fn().mockResolvedValue({ + response: new HttpResponse({ + statusCode: 200, + body: Readable.from([""]), + }), +}); +jest.mock("@aws-sdk/node-http-handler", () => ({ + NodeHttpHandler: jest.fn().mockImplementation(() => ({ + destroy: () => {}, + handle: mockHandle, + })), + streamCollector: jest.fn(), +})); + +import { getDefaultRoleAssumer, getDefaultRoleAssumerWithWebIdentity } from "./defaultRoleAssumers"; +import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand"; +import { NodeHttpHandler, streamCollector } from "@aws-sdk/node-http-handler"; +import { AssumeRoleWithWebIdentityCommandInput } from "./commands/AssumeRoleWithWebIdentityCommand"; +const mockConstructorInput = jest.fn(); +jest.mock("./STSClient", () => ({ + STSClient: function (params: any) { + mockConstructorInput(params); + //@ts-ignore + return new (jest.requireActual("./STSClient").STSClient)(params); + }, +})); + +describe("getDefaultRoleAssumer", () => { + const assumeRoleResponse = ` AROAZOX2IL27GNRBJHWC2:session @@ -19,27 +48,15 @@ const assumeRoleResponse = ` ({ - NodeHttpHandler: jest.fn().mockImplementation(() => ({ - destroy: () => {}, - handle: mockHandle, - })), - streamCollector: async () => Buffer.from(assumeRoleResponse), -})); -import { getDefaultRoleAssumer } from "./defaultRoleAssumers"; -import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand"; + beforeAll(() => { + (streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse)); + }); -describe("getDefaultRoleAssumer", () => { beforeEach(() => { jest.clearAllMocks(); }); + it("should use supplied source credentials", async () => { const roleAssumer = getDefaultRoleAssumer(); const params: AssumeRoleCommandInput = { @@ -61,4 +78,71 @@ describe("getDefaultRoleAssumer", () => { expect.stringContaining("AWS4-HMAC-SHA256 Credential=key2/") ); }); + + it("should use the STS client config", async () => { + const logger = console; + const region = "some-region"; + const handler = new NodeHttpHandler(); + const roleAssumer = getDefaultRoleAssumer({ + region, + logger, + requestHandler: handler, + }); + const params: AssumeRoleCommandInput = { + RoleArn: "arn:aws:foo", + RoleSessionName: "session", + }; + const sourceCred = { accessKeyId: "key", secretAccessKey: "secrete" }; + await roleAssumer(sourceCred, params); + expect(mockConstructorInput).toHaveBeenCalledTimes(1); + expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({ + logger, + requestHandler: handler, + region, + }); + }); +}); + +describe("getDefaultRoleAssumerWithWebIdentity", () => { + const assumeRoleResponse = ` + + + key + secrete + session-token + 2021-05-05T23:22:08Z + + + `; + + beforeAll(() => { + (streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse)); + }); + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it("should use the STS client config", async () => { + const logger = console; + const region = "some-region"; + const handler = new NodeHttpHandler(); + const roleAssumerWithWebIdentity = getDefaultRoleAssumerWithWebIdentity({ + region, + logger, + requestHandler: handler, + }); + const params: AssumeRoleWithWebIdentityCommandInput = { + RoleArn: "arn:aws:foo", + RoleSessionName: "session", + WebIdentityToken: "token", + }; + await roleAssumerWithWebIdentity(params); + expect(mockConstructorInput).toHaveBeenCalledTimes(1); + expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({ + logger, + requestHandler: handler, + region, + }); + }); }); diff --git a/clients/client-sts/defaultRoleAssumers.ts b/clients/client-sts/defaultRoleAssumers.ts index bc5ddc83eeb0..e5e2d1ea9a2f 100644 --- a/clients/client-sts/defaultRoleAssumers.ts +++ b/clients/client-sts/defaultRoleAssumers.ts @@ -12,14 +12,15 @@ import { STSClient, STSClientConfig } from "./STSClient"; /** * The default role assumer that used by credential providers when sts:AssumeRole API is needed. */ -export const getDefaultRoleAssumer = (stsOptions: Pick = {}): RoleAssumer => - StsGetDefaultRoleAssumer(stsOptions, STSClient); +export const getDefaultRoleAssumer = ( + stsOptions: Pick = {} +): RoleAssumer => StsGetDefaultRoleAssumer(stsOptions, STSClient); /** * The default role assumer that used by credential providers when sts:AssumeRoleWithWebIdentity API is needed. */ export const getDefaultRoleAssumerWithWebIdentity = ( - stsOptions: Pick = {} + stsOptions: Pick = {} ): RoleAssumerWithWebIdentity => StsGetDefaultRoleAssumerWithWebIdentity(stsOptions, STSClient); /** diff --git a/clients/client-sts/defaultStsRoleAssumers.ts b/clients/client-sts/defaultStsRoleAssumers.ts index ef2210a61f88..c45e2a2093b0 100644 --- a/clients/client-sts/defaultStsRoleAssumers.ts +++ b/clients/client-sts/defaultStsRoleAssumers.ts @@ -37,7 +37,7 @@ const decorateDefaultRegion = (region: string | Provider | undefined): s * @internal */ export const getDefaultRoleAssumer = ( - stsOptions: Pick, + stsOptions: Pick, stsClientCtor: new (options: STSClientConfig) => STSClient ): RoleAssumer => { let stsClient: STSClient; @@ -45,12 +45,13 @@ export const getDefaultRoleAssumer = ( return async (sourceCreds, params) => { closureSourceCreds = sourceCreds; if (!stsClient) { - const { logger, region } = stsOptions; + const { logger, region, requestHandler } = stsOptions; stsClient = new stsClientCtor({ logger, // A hack to make sts client uses the credential in current closure. credentialDefaultProvider: () => async () => closureSourceCreds, - region: decorateDefaultRegion(region), + region: decorateDefaultRegion(region || stsOptions.region), + ...(requestHandler ? { requestHandler } : {}), }); } const { Credentials } = await stsClient.send(new AssumeRoleCommand(params)); @@ -76,16 +77,17 @@ export type RoleAssumerWithWebIdentity = (params: AssumeRoleWithWebIdentityComma * @internal */ export const getDefaultRoleAssumerWithWebIdentity = ( - stsOptions: Pick, + stsOptions: Pick, stsClientCtor: new (options: STSClientConfig) => STSClient ): RoleAssumerWithWebIdentity => { let stsClient: STSClient; return async (params) => { if (!stsClient) { - const { logger, region } = stsOptions; + const { logger, region, requestHandler } = stsOptions; stsClient = new stsClientCtor({ logger, - region: decorateDefaultRegion(region), + region: decorateDefaultRegion(region || stsOptions.region), + ...(requestHandler ? { requestHandler } : {}), }); } const { Credentials } = await stsClient.send(new AssumeRoleWithWebIdentityCommand(params)); From 70cb2fa3e27b1cff14319826a2802fe6024167d3 Mon Sep 17 00:00:00 2001 From: Allan Zheng Date: Mon, 24 May 2021 15:21:55 -0700 Subject: [PATCH 2/2] chore(client-sts): generate new default assume role files --- .../sts-client-defaultRoleAssumers.spec.ts | 118 +++++++++++++++--- .../codegen/sts-client-defaultRoleAssumers.ts | 7 +- .../sts-client-defaultStsRoleAssumers.ts | 14 ++- 3 files changed, 113 insertions(+), 26 deletions(-) diff --git a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts index 328e465283cc..86b2f175f0fb 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts +++ b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts @@ -1,6 +1,35 @@ import { HttpResponse } from "@aws-sdk/protocol-http"; import { Readable } from "stream"; -const assumeRoleResponse = ` + +const mockHandle = jest.fn().mockResolvedValue({ + response: new HttpResponse({ + statusCode: 200, + body: Readable.from([""]), + }), +}); +jest.mock("@aws-sdk/node-http-handler", () => ({ + NodeHttpHandler: jest.fn().mockImplementation(() => ({ + destroy: () => {}, + handle: mockHandle, + })), + streamCollector: jest.fn(), +})); + +import { getDefaultRoleAssumer, getDefaultRoleAssumerWithWebIdentity } from "./defaultRoleAssumers"; +import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand"; +import { NodeHttpHandler, streamCollector } from "@aws-sdk/node-http-handler"; +import { AssumeRoleWithWebIdentityCommandInput } from "./commands/AssumeRoleWithWebIdentityCommand"; +const mockConstructorInput = jest.fn(); +jest.mock("./STSClient", () => ({ + STSClient: function (params: any) { + mockConstructorInput(params); + //@ts-ignore + return new (jest.requireActual("./STSClient").STSClient)(params); + }, +})); + +describe("getDefaultRoleAssumer", () => { + const assumeRoleResponse = ` AROAZOX2IL27GNRBJHWC2:session @@ -17,27 +46,15 @@ const assumeRoleResponse = ` ({ - NodeHttpHandler: jest.fn().mockImplementation(() => ({ - destroy: () => {}, - handle: mockHandle, - })), - streamCollector: async () => Buffer.from(assumeRoleResponse), -})); -import { getDefaultRoleAssumer } from "./defaultRoleAssumers"; -import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand"; + beforeAll(() => { + (streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse)); + }); -describe("getDefaultRoleAssumer", () => { beforeEach(() => { jest.clearAllMocks(); }); + it("should use supplied source credentials", async () => { const roleAssumer = getDefaultRoleAssumer(); const params: AssumeRoleCommandInput = { @@ -59,4 +76,71 @@ describe("getDefaultRoleAssumer", () => { expect.stringContaining("AWS4-HMAC-SHA256 Credential=key2/") ); }); + + it("should use the STS client config", async () => { + const logger = console; + const region = "some-region"; + const handler = new NodeHttpHandler(); + const roleAssumer = getDefaultRoleAssumer({ + region, + logger, + requestHandler: handler, + }); + const params: AssumeRoleCommandInput = { + RoleArn: "arn:aws:foo", + RoleSessionName: "session", + }; + const sourceCred = { accessKeyId: "key", secretAccessKey: "secrete" }; + await roleAssumer(sourceCred, params); + expect(mockConstructorInput).toHaveBeenCalledTimes(1); + expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({ + logger, + requestHandler: handler, + region, + }); + }); +}); + +describe("getDefaultRoleAssumerWithWebIdentity", () => { + const assumeRoleResponse = ` + + + key + secrete + session-token + 2021-05-05T23:22:08Z + + + `; + + beforeAll(() => { + (streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse)); + }); + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it("should use the STS client config", async () => { + const logger = console; + const region = "some-region"; + const handler = new NodeHttpHandler(); + const roleAssumerWithWebIdentity = getDefaultRoleAssumerWithWebIdentity({ + region, + logger, + requestHandler: handler, + }); + const params: AssumeRoleWithWebIdentityCommandInput = { + RoleArn: "arn:aws:foo", + RoleSessionName: "session", + WebIdentityToken: "token", + }; + await roleAssumerWithWebIdentity(params); + expect(mockConstructorInput).toHaveBeenCalledTimes(1); + expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({ + logger, + requestHandler: handler, + region, + }); + }); }); diff --git a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.ts b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.ts index 76ab95049e91..d72fb9044b74 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.ts +++ b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.ts @@ -10,14 +10,15 @@ import { STSClient, STSClientConfig } from "./STSClient"; /** * The default role assumer that used by credential providers when sts:AssumeRole API is needed. */ -export const getDefaultRoleAssumer = (stsOptions: Pick = {}): RoleAssumer => - StsGetDefaultRoleAssumer(stsOptions, STSClient); +export const getDefaultRoleAssumer = ( + stsOptions: Pick = {} +): RoleAssumer => StsGetDefaultRoleAssumer(stsOptions, STSClient); /** * The default role assumer that used by credential providers when sts:AssumeRoleWithWebIdentity API is needed. */ export const getDefaultRoleAssumerWithWebIdentity = ( - stsOptions: Pick = {} + stsOptions: Pick = {} ): RoleAssumerWithWebIdentity => StsGetDefaultRoleAssumerWithWebIdentity(stsOptions, STSClient); /** diff --git a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultStsRoleAssumers.ts b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultStsRoleAssumers.ts index 638ef7a9d9b7..46a8530b85f7 100644 --- a/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultStsRoleAssumers.ts +++ b/codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultStsRoleAssumers.ts @@ -35,7 +35,7 @@ const decorateDefaultRegion = (region: string | Provider | undefined): s * @internal */ export const getDefaultRoleAssumer = ( - stsOptions: Pick, + stsOptions: Pick, stsClientCtor: new (options: STSClientConfig) => STSClient ): RoleAssumer => { let stsClient: STSClient; @@ -43,12 +43,13 @@ export const getDefaultRoleAssumer = ( return async (sourceCreds, params) => { closureSourceCreds = sourceCreds; if (!stsClient) { - const { logger, region } = stsOptions; + const { logger, region, requestHandler } = stsOptions; stsClient = new stsClientCtor({ logger, // A hack to make sts client uses the credential in current closure. credentialDefaultProvider: () => async () => closureSourceCreds, - region: decorateDefaultRegion(region), + region: decorateDefaultRegion(region || stsOptions.region), + ...(requestHandler ? { requestHandler } : {}), }); } const { Credentials } = await stsClient.send(new AssumeRoleCommand(params)); @@ -74,16 +75,17 @@ export type RoleAssumerWithWebIdentity = (params: AssumeRoleWithWebIdentityComma * @internal */ export const getDefaultRoleAssumerWithWebIdentity = ( - stsOptions: Pick, + stsOptions: Pick, stsClientCtor: new (options: STSClientConfig) => STSClient ): RoleAssumerWithWebIdentity => { let stsClient: STSClient; return async (params) => { if (!stsClient) { - const { logger, region } = stsOptions; + const { logger, region, requestHandler } = stsOptions; stsClient = new stsClientCtor({ logger, - region: decorateDefaultRegion(region), + region: decorateDefaultRegion(region || stsOptions.region), + ...(requestHandler ? { requestHandler } : {}), }); } const { Credentials } = await stsClient.send(new AssumeRoleWithWebIdentityCommand(params));