Skip to content

Commit 7985995

Browse files
AllanZhengYPAllanZhengYP
authored
fix(client-sts): allow overwriting default role assumer http handler (#2426)
* fix(client-sts): allow overwriting default role assumer http handler * chore(client-sts): generate new default assume role files
1 parent 8ef104d commit 7985995

File tree

6 files changed

+226
-52
lines changed

6 files changed

+226
-52
lines changed

Diff for: ‎clients/client-sts/defaultRoleAssumers.spec.ts

+101-17
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,36 @@
22
// 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
33
import { HttpResponse } from "@aws-sdk/protocol-http";
44
import { Readable } from "stream";
5-
const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
5+
6+
const mockHandle = jest.fn().mockResolvedValue({
7+
response: new HttpResponse({
8+
statusCode: 200,
9+
body: Readable.from([""]),
10+
}),
11+
});
12+
jest.mock("@aws-sdk/node-http-handler", () => ({
13+
NodeHttpHandler: jest.fn().mockImplementation(() => ({
14+
destroy: () => {},
15+
handle: mockHandle,
16+
})),
17+
streamCollector: jest.fn(),
18+
}));
19+
20+
import { getDefaultRoleAssumer, getDefaultRoleAssumerWithWebIdentity } from "./defaultRoleAssumers";
21+
import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand";
22+
import { NodeHttpHandler, streamCollector } from "@aws-sdk/node-http-handler";
23+
import { AssumeRoleWithWebIdentityCommandInput } from "./commands/AssumeRoleWithWebIdentityCommand";
24+
const mockConstructorInput = jest.fn();
25+
jest.mock("./STSClient", () => ({
26+
STSClient: function (params: any) {
27+
mockConstructorInput(params);
28+
//@ts-ignore
29+
return new (jest.requireActual("./STSClient").STSClient)(params);
30+
},
31+
}));
32+
33+
describe("getDefaultRoleAssumer", () => {
34+
const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
635
<AssumeRoleResult>
736
<AssumedRoleUser>
837
<AssumedRoleId>AROAZOX2IL27GNRBJHWC2:session</AssumedRoleId>
@@ -19,27 +48,15 @@ const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com
1948
<RequestId>12345678id</RequestId>
2049
</ResponseMetadata>
2150
</AssumeRoleResponse>`;
22-
const mockHandle = jest.fn().mockResolvedValue({
23-
response: new HttpResponse({
24-
statusCode: 200,
25-
body: Readable.from([""]),
26-
}),
27-
});
28-
jest.mock("@aws-sdk/node-http-handler", () => ({
29-
NodeHttpHandler: jest.fn().mockImplementation(() => ({
30-
destroy: () => {},
31-
handle: mockHandle,
32-
})),
33-
streamCollector: async () => Buffer.from(assumeRoleResponse),
34-
}));
3551

36-
import { getDefaultRoleAssumer } from "./defaultRoleAssumers";
37-
import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand";
52+
beforeAll(() => {
53+
(streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse));
54+
});
3855

39-
describe("getDefaultRoleAssumer", () => {
4056
beforeEach(() => {
4157
jest.clearAllMocks();
4258
});
59+
4360
it("should use supplied source credentials", async () => {
4461
const roleAssumer = getDefaultRoleAssumer();
4562
const params: AssumeRoleCommandInput = {
@@ -61,4 +78,71 @@ describe("getDefaultRoleAssumer", () => {
6178
expect.stringContaining("AWS4-HMAC-SHA256 Credential=key2/")
6279
);
6380
});
81+
82+
it("should use the STS client config", async () => {
83+
const logger = console;
84+
const region = "some-region";
85+
const handler = new NodeHttpHandler();
86+
const roleAssumer = getDefaultRoleAssumer({
87+
region,
88+
logger,
89+
requestHandler: handler,
90+
});
91+
const params: AssumeRoleCommandInput = {
92+
RoleArn: "arn:aws:foo",
93+
RoleSessionName: "session",
94+
};
95+
const sourceCred = { accessKeyId: "key", secretAccessKey: "secrete" };
96+
await roleAssumer(sourceCred, params);
97+
expect(mockConstructorInput).toHaveBeenCalledTimes(1);
98+
expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({
99+
logger,
100+
requestHandler: handler,
101+
region,
102+
});
103+
});
104+
});
105+
106+
describe("getDefaultRoleAssumerWithWebIdentity", () => {
107+
const assumeRoleResponse = `<Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
108+
<AssumeRoleWithWebIdentityResult>
109+
<Credentials>
110+
<AccessKeyId>key</AccessKeyId>
111+
<SecretAccessKey>secrete</SecretAccessKey>
112+
<SessionToken>session-token</SessionToken>
113+
<Expiration>2021-05-05T23:22:08Z</Expiration>
114+
</Credentials>
115+
</AssumeRoleWithWebIdentityResult>
116+
</Response>`;
117+
118+
beforeAll(() => {
119+
(streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse));
120+
});
121+
122+
beforeEach(() => {
123+
jest.clearAllMocks();
124+
});
125+
126+
it("should use the STS client config", async () => {
127+
const logger = console;
128+
const region = "some-region";
129+
const handler = new NodeHttpHandler();
130+
const roleAssumerWithWebIdentity = getDefaultRoleAssumerWithWebIdentity({
131+
region,
132+
logger,
133+
requestHandler: handler,
134+
});
135+
const params: AssumeRoleWithWebIdentityCommandInput = {
136+
RoleArn: "arn:aws:foo",
137+
RoleSessionName: "session",
138+
WebIdentityToken: "token",
139+
};
140+
await roleAssumerWithWebIdentity(params);
141+
expect(mockConstructorInput).toHaveBeenCalledTimes(1);
142+
expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({
143+
logger,
144+
requestHandler: handler,
145+
region,
146+
});
147+
});
64148
});

Diff for: ‎clients/client-sts/defaultRoleAssumers.ts

+4-3
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,15 @@ import { STSClient, STSClientConfig } from "./STSClient";
1212
/**
1313
* The default role assumer that used by credential providers when sts:AssumeRole API is needed.
1414
*/
15-
export const getDefaultRoleAssumer = (stsOptions: Pick<STSClientConfig, "logger" | "region"> = {}): RoleAssumer =>
16-
StsGetDefaultRoleAssumer(stsOptions, STSClient);
15+
export const getDefaultRoleAssumer = (
16+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler"> = {}
17+
): RoleAssumer => StsGetDefaultRoleAssumer(stsOptions, STSClient);
1718

1819
/**
1920
* The default role assumer that used by credential providers when sts:AssumeRoleWithWebIdentity API is needed.
2021
*/
2122
export const getDefaultRoleAssumerWithWebIdentity = (
22-
stsOptions: Pick<STSClientConfig, "logger" | "region"> = {}
23+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler"> = {}
2324
): RoleAssumerWithWebIdentity => StsGetDefaultRoleAssumerWithWebIdentity(stsOptions, STSClient);
2425

2526
/**

Diff for: ‎clients/client-sts/defaultStsRoleAssumers.ts

+8-6
Original file line numberDiff line numberDiff line change
@@ -37,20 +37,21 @@ const decorateDefaultRegion = (region: string | Provider<string> | undefined): s
3737
* @internal
3838
*/
3939
export const getDefaultRoleAssumer = (
40-
stsOptions: Pick<STSClientConfig, "logger" | "region">,
40+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler">,
4141
stsClientCtor: new (options: STSClientConfig) => STSClient
4242
): RoleAssumer => {
4343
let stsClient: STSClient;
4444
let closureSourceCreds: Credentials;
4545
return async (sourceCreds, params) => {
4646
closureSourceCreds = sourceCreds;
4747
if (!stsClient) {
48-
const { logger, region } = stsOptions;
48+
const { logger, region, requestHandler } = stsOptions;
4949
stsClient = new stsClientCtor({
5050
logger,
5151
// A hack to make sts client uses the credential in current closure.
5252
credentialDefaultProvider: () => async () => closureSourceCreds,
53-
region: decorateDefaultRegion(region),
53+
region: decorateDefaultRegion(region || stsOptions.region),
54+
...(requestHandler ? { requestHandler } : {}),
5455
});
5556
}
5657
const { Credentials } = await stsClient.send(new AssumeRoleCommand(params));
@@ -76,16 +77,17 @@ export type RoleAssumerWithWebIdentity = (params: AssumeRoleWithWebIdentityComma
7677
* @internal
7778
*/
7879
export const getDefaultRoleAssumerWithWebIdentity = (
79-
stsOptions: Pick<STSClientConfig, "logger" | "region">,
80+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler">,
8081
stsClientCtor: new (options: STSClientConfig) => STSClient
8182
): RoleAssumerWithWebIdentity => {
8283
let stsClient: STSClient;
8384
return async (params) => {
8485
if (!stsClient) {
85-
const { logger, region } = stsOptions;
86+
const { logger, region, requestHandler } = stsOptions;
8687
stsClient = new stsClientCtor({
8788
logger,
88-
region: decorateDefaultRegion(region),
89+
region: decorateDefaultRegion(region || stsOptions.region),
90+
...(requestHandler ? { requestHandler } : {}),
8991
});
9092
}
9193
const { Credentials } = await stsClient.send(new AssumeRoleWithWebIdentityCommand(params));

Diff for: ‎codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.spec.ts

+101-17
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,35 @@
11
import { HttpResponse } from "@aws-sdk/protocol-http";
22
import { Readable } from "stream";
3-
const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
3+
4+
const mockHandle = jest.fn().mockResolvedValue({
5+
response: new HttpResponse({
6+
statusCode: 200,
7+
body: Readable.from([""]),
8+
}),
9+
});
10+
jest.mock("@aws-sdk/node-http-handler", () => ({
11+
NodeHttpHandler: jest.fn().mockImplementation(() => ({
12+
destroy: () => {},
13+
handle: mockHandle,
14+
})),
15+
streamCollector: jest.fn(),
16+
}));
17+
18+
import { getDefaultRoleAssumer, getDefaultRoleAssumerWithWebIdentity } from "./defaultRoleAssumers";
19+
import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand";
20+
import { NodeHttpHandler, streamCollector } from "@aws-sdk/node-http-handler";
21+
import { AssumeRoleWithWebIdentityCommandInput } from "./commands/AssumeRoleWithWebIdentityCommand";
22+
const mockConstructorInput = jest.fn();
23+
jest.mock("./STSClient", () => ({
24+
STSClient: function (params: any) {
25+
mockConstructorInput(params);
26+
//@ts-ignore
27+
return new (jest.requireActual("./STSClient").STSClient)(params);
28+
},
29+
}));
30+
31+
describe("getDefaultRoleAssumer", () => {
32+
const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
433
<AssumeRoleResult>
534
<AssumedRoleUser>
635
<AssumedRoleId>AROAZOX2IL27GNRBJHWC2:session</AssumedRoleId>
@@ -17,27 +46,15 @@ const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com
1746
<RequestId>12345678id</RequestId>
1847
</ResponseMetadata>
1948
</AssumeRoleResponse>`;
20-
const mockHandle = jest.fn().mockResolvedValue({
21-
response: new HttpResponse({
22-
statusCode: 200,
23-
body: Readable.from([""]),
24-
}),
25-
});
26-
jest.mock("@aws-sdk/node-http-handler", () => ({
27-
NodeHttpHandler: jest.fn().mockImplementation(() => ({
28-
destroy: () => {},
29-
handle: mockHandle,
30-
})),
31-
streamCollector: async () => Buffer.from(assumeRoleResponse),
32-
}));
3349

34-
import { getDefaultRoleAssumer } from "./defaultRoleAssumers";
35-
import type { AssumeRoleCommandInput } from "./commands/AssumeRoleCommand";
50+
beforeAll(() => {
51+
(streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse));
52+
});
3653

37-
describe("getDefaultRoleAssumer", () => {
3854
beforeEach(() => {
3955
jest.clearAllMocks();
4056
});
57+
4158
it("should use supplied source credentials", async () => {
4259
const roleAssumer = getDefaultRoleAssumer();
4360
const params: AssumeRoleCommandInput = {
@@ -59,4 +76,71 @@ describe("getDefaultRoleAssumer", () => {
5976
expect.stringContaining("AWS4-HMAC-SHA256 Credential=key2/")
6077
);
6178
});
79+
80+
it("should use the STS client config", async () => {
81+
const logger = console;
82+
const region = "some-region";
83+
const handler = new NodeHttpHandler();
84+
const roleAssumer = getDefaultRoleAssumer({
85+
region,
86+
logger,
87+
requestHandler: handler,
88+
});
89+
const params: AssumeRoleCommandInput = {
90+
RoleArn: "arn:aws:foo",
91+
RoleSessionName: "session",
92+
};
93+
const sourceCred = { accessKeyId: "key", secretAccessKey: "secrete" };
94+
await roleAssumer(sourceCred, params);
95+
expect(mockConstructorInput).toHaveBeenCalledTimes(1);
96+
expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({
97+
logger,
98+
requestHandler: handler,
99+
region,
100+
});
101+
});
102+
});
103+
104+
describe("getDefaultRoleAssumerWithWebIdentity", () => {
105+
const assumeRoleResponse = `<Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
106+
<AssumeRoleWithWebIdentityResult>
107+
<Credentials>
108+
<AccessKeyId>key</AccessKeyId>
109+
<SecretAccessKey>secrete</SecretAccessKey>
110+
<SessionToken>session-token</SessionToken>
111+
<Expiration>2021-05-05T23:22:08Z</Expiration>
112+
</Credentials>
113+
</AssumeRoleWithWebIdentityResult>
114+
</Response>`;
115+
116+
beforeAll(() => {
117+
(streamCollector as jest.Mock).mockImplementation(async () => Buffer.from(assumeRoleResponse));
118+
});
119+
120+
beforeEach(() => {
121+
jest.clearAllMocks();
122+
});
123+
124+
it("should use the STS client config", async () => {
125+
const logger = console;
126+
const region = "some-region";
127+
const handler = new NodeHttpHandler();
128+
const roleAssumerWithWebIdentity = getDefaultRoleAssumerWithWebIdentity({
129+
region,
130+
logger,
131+
requestHandler: handler,
132+
});
133+
const params: AssumeRoleWithWebIdentityCommandInput = {
134+
RoleArn: "arn:aws:foo",
135+
RoleSessionName: "session",
136+
WebIdentityToken: "token",
137+
};
138+
await roleAssumerWithWebIdentity(params);
139+
expect(mockConstructorInput).toHaveBeenCalledTimes(1);
140+
expect(mockConstructorInput.mock.calls[0][0]).toMatchObject({
141+
logger,
142+
requestHandler: handler,
143+
region,
144+
});
145+
});
62146
});

Diff for: ‎codegen/smithy-aws-typescript-codegen/src/main/resources/software/amazon/smithy/aws/typescript/codegen/sts-client-defaultRoleAssumers.ts

+4-3
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,15 @@ import { STSClient, STSClientConfig } from "./STSClient";
1010
/**
1111
* The default role assumer that used by credential providers when sts:AssumeRole API is needed.
1212
*/
13-
export const getDefaultRoleAssumer = (stsOptions: Pick<STSClientConfig, "logger" | "region"> = {}): RoleAssumer =>
14-
StsGetDefaultRoleAssumer(stsOptions, STSClient);
13+
export const getDefaultRoleAssumer = (
14+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler"> = {}
15+
): RoleAssumer => StsGetDefaultRoleAssumer(stsOptions, STSClient);
1516

1617
/**
1718
* The default role assumer that used by credential providers when sts:AssumeRoleWithWebIdentity API is needed.
1819
*/
1920
export const getDefaultRoleAssumerWithWebIdentity = (
20-
stsOptions: Pick<STSClientConfig, "logger" | "region"> = {}
21+
stsOptions: Pick<STSClientConfig, "logger" | "region" | "requestHandler"> = {}
2122
): RoleAssumerWithWebIdentity => StsGetDefaultRoleAssumerWithWebIdentity(stsOptions, STSClient);
2223

2324
/**

0 commit comments

Comments
 (0)