Skip to content

Commit

Permalink
fix(middleware-sdk-sqs): add middleware to prioritize QueueUrl endpoi…
Browse files Browse the repository at this point in the history
…nt resolution (#5759)

* fix(middleware-sdk-sqs): add middleware to prioritize QueueUrl endpoint resolution

* fix(middleware-sdk-sqs): add option to use QueueUrl as endpoint

* test(middleware-sdk-sqs): initial unit test for queue-url middleware

* test(middleware-sdk-sqs): fix unit and add integ tests

* test(middleware-sdk-sqs): avoid using QueueUrl if custom endpoint given to client

* test(middleware-sdk-sqs): warn with exact error

---------

Co-authored-by: George Fu <kuhe@users.noreply.github.com>
  • Loading branch information
siddsriv and kuhe committed Feb 5, 2024
1 parent 4d6db90 commit 2da1b6c
Show file tree
Hide file tree
Showing 7 changed files with 276 additions and 5 deletions.
20 changes: 15 additions & 5 deletions clients/client-sqs/src/SQSClient.ts
Expand Up @@ -7,6 +7,12 @@ import {
} from "@aws-sdk/middleware-host-header";
import { getLoggerPlugin } from "@aws-sdk/middleware-logger";
import { getRecursionDetectionPlugin } from "@aws-sdk/middleware-recursion-detection";
import {
getQueueUrlPlugin,
QueueUrlInputConfig,
QueueUrlResolvedConfig,
resolveQueueUrlConfig,
} from "@aws-sdk/middleware-sdk-sqs";
import {
getUserAgentPlugin,
resolveUserAgentConfig,
Expand Down Expand Up @@ -311,6 +317,7 @@ export type SQSClientConfigType = Partial<__SmithyConfiguration<__HttpHandlerOpt
EndpointInputConfig<EndpointParameters> &
RetryInputConfig &
HostHeaderInputConfig &
QueueUrlInputConfig &
UserAgentInputConfig &
HttpAuthSchemeInputConfig &
ClientInputEndpointParameters;
Expand All @@ -331,6 +338,7 @@ export type SQSClientResolvedConfigType = __SmithyResolvedConfiguration<__HttpHa
EndpointResolvedConfig<EndpointParameters> &
RetryResolvedConfig &
HostHeaderResolvedConfig &
QueueUrlResolvedConfig &
UserAgentResolvedConfig &
HttpAuthSchemeResolvedConfig &
ClientResolvedEndpointParameters;
Expand Down Expand Up @@ -435,16 +443,18 @@ export class SQSClient extends __Client<
const _config_3 = resolveEndpointConfig(_config_2);
const _config_4 = resolveRetryConfig(_config_3);
const _config_5 = resolveHostHeaderConfig(_config_4);
const _config_6 = resolveUserAgentConfig(_config_5);
const _config_7 = resolveHttpAuthSchemeConfig(_config_6);
const _config_8 = resolveRuntimeExtensions(_config_7, configuration?.extensions || []);
super(_config_8);
this.config = _config_8;
const _config_6 = resolveQueueUrlConfig(_config_5);
const _config_7 = resolveUserAgentConfig(_config_6);
const _config_8 = resolveHttpAuthSchemeConfig(_config_7);
const _config_9 = resolveRuntimeExtensions(_config_8, configuration?.extensions || []);
super(_config_9);
this.config = _config_9;
this.middlewareStack.use(getRetryPlugin(this.config));
this.middlewareStack.use(getContentLengthPlugin(this.config));
this.middlewareStack.use(getHostHeaderPlugin(this.config));
this.middlewareStack.use(getLoggerPlugin(this.config));
this.middlewareStack.use(getRecursionDetectionPlugin(this.config));
this.middlewareStack.use(getQueueUrlPlugin(this.config));
this.middlewareStack.use(getUserAgentPlugin(this.config));
this.middlewareStack.use(
getHttpAuthSchemeEndpointRuleSetPlugin(this.config, {
Expand Down
Expand Up @@ -15,6 +15,7 @@

package software.amazon.smithy.aws.typescript.codegen;

import static software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin.Convention.HAS_CONFIG;
import static software.amazon.smithy.typescript.codegen.integration.RuntimeClientPlugin.Convention.HAS_MIDDLEWARE;

import java.util.Collections;
Expand Down Expand Up @@ -58,6 +59,14 @@ public List<RuntimeClientPlugin> getClientPlugins() {
.withConventions(AwsDependency.SQS_MIDDLEWARE.dependency, "ReceiveMessage",
HAS_MIDDLEWARE)
.operationPredicate((m, s, o) -> o.getId().getName(s).equals("ReceiveMessage") && isSqs(s))
.build(),
RuntimeClientPlugin.builder()
.withConventions(
AwsDependency.SQS_MIDDLEWARE.dependency,
"QueueUrl",
HAS_MIDDLEWARE, HAS_CONFIG
)
.servicePredicate((m, s) -> isSqs(s))
.build()
);
}
Expand Down
1 change: 1 addition & 0 deletions packages/middleware-sdk-sqs/package.json
Expand Up @@ -22,6 +22,7 @@
"license": "Apache-2.0",
"dependencies": {
"@aws-sdk/types": "*",
"@smithy/smithy-client": "^2.3.1",
"@smithy/types": "^2.9.1",
"@smithy/util-hex-encoding": "^2.1.1",
"@smithy/util-utf8": "^2.1.1",
Expand Down
1 change: 1 addition & 0 deletions packages/middleware-sdk-sqs/src/index.ts
@@ -1,3 +1,4 @@
export * from "./queue-url";
export * from "./receive-message";
export * from "./send-message";
export * from "./send-message-batch";
56 changes: 56 additions & 0 deletions packages/middleware-sdk-sqs/src/middleware-sdk-sqs.integ.spec.ts
Expand Up @@ -4,6 +4,7 @@ import crypto from "crypto";
import { Readable } from "stream";

import sqsModel from "../../../codegen/sdk-codegen/aws-models/sqs.json";
import { requireRequestsFrom } from "../../../private/aws-util-test/src";
const useAwsQuery = !!sqsModel.shapes["com.amazonaws.sqs#AmazonSQS"].traits["aws.protocols#awsQuery"];

let hashError = "";
Expand Down Expand Up @@ -289,4 +290,59 @@ describe("middleware-sdk-sqs", () => {
});
});
});

describe("queue-url", () => {
it("should override resolved endpoint by default", async () => {
const client = new SQS({
region: "us-west-2",
});

requireRequestsFrom(client).toMatch({
hostname: "abc.com",
protocol: "https:",
path: "/",
});

await client.sendMessage({
QueueUrl: "https://abc.com/123/MyQueue",
MessageBody: "hello",
});
});

it("does not override endpoint if shut off with useQueueUrlAsEndpoint=false", async () => {
const client = new SQS({
region: "us-west-2",
useQueueUrlAsEndpoint: false,
});

requireRequestsFrom(client).toMatch({
hostname: "sqs.us-west-2.amazonaws.com",
protocol: "https:",
path: "/",
});

await client.sendMessage({
QueueUrl: "https://abc.com/123/MyQueue",
MessageBody: "hello",
});
});

it("does not override endpoint if custom endpoint given to client", async () => {
const client = new SQS({
region: "us-west-2",
endpoint: "https://custom-endpoint.com/",
});

requireRequestsFrom(client).toMatch({
hostname: "custom-endpoint.com",
protocol: "https:",
path: "/",
});

await client.sendMessage({
QueueUrl: "https://abc.com/123/MyQueue",
MessageBody: "hello",
});
});
});
});
97 changes: 97 additions & 0 deletions packages/middleware-sdk-sqs/src/queue-url.spec.ts
@@ -0,0 +1,97 @@
import { HttpRequest } from "@aws-sdk/protocol-http";
import { FinalizeHandlerArguments, HandlerExecutionContext } from "@aws-sdk/types";

import { queueUrlMiddleware } from "./queue-url";

describe("queueUrlMiddleware", () => {
const mockNextHandler = jest.fn();

const mockContext: HandlerExecutionContext = {
logger: {
...console,
warn: jest.fn(),
},
endpointV2: void 0,
};

beforeEach(() => {
mockNextHandler.mockReset();
mockNextHandler.mockResolvedValue({ output: {}, response: {} });
mockContext.endpointV2 = {
url: new URL("https://sqs.us-east-1.amazonaws.com"),
};
});

afterEach(() => {
jest.resetAllMocks();
});

it("should use the QueueUrl hostname as the endpoint if useQueueUrlAsEndpoint is true", async () => {
const middleware = queueUrlMiddleware({ useQueueUrlAsEndpoint: true });
const input = { QueueUrl: "https://xyz.com/123/MyQueue" };
const request = new HttpRequest({
hostname: "sqs.us-east-1.amazonaws.com",
protocol: "https:",
path: "/",
headers: {},
method: "GET",
});
const args: FinalizeHandlerArguments<any> = { input, request };

await middleware(mockNextHandler, mockContext)(args);

// Verify that the resolvedEndpoint.url has been modified to match QueueUrl
expect(mockContext.endpointV2?.url.href).toEqual("https://xyz.com/");
expect(mockNextHandler).toHaveBeenCalled();
expect(mockContext.logger?.warn).toHaveBeenCalled();
});

it("should not modify the endpoint if useQueueUrlAsEndpoint is false", async () => {
const middleware = queueUrlMiddleware({ useQueueUrlAsEndpoint: false });
const input = { QueueUrl: "https://xyz.com/123/MyQueue" };
const request = new HttpRequest({
hostname: "sqs.us-east-1.amazonaws.com",
protocol: "https:",
path: "/",
});
const args: FinalizeHandlerArguments<any> = { input, request };

await middleware(mockNextHandler, mockContext)(args);

expect(mockNextHandler).toHaveBeenCalledWith(args);
expect(mockContext.endpointV2?.url.href).toEqual("https://sqs.us-east-1.amazonaws.com/");
});

it("should not modify the endpoint when QueueUrl is not provided", async () => {
const middleware = queueUrlMiddleware({ useQueueUrlAsEndpoint: true });
const input = {}; // No QueueUrl provided
const request = new HttpRequest({
hostname: "sqs.us-east-1.amazonaws.com",
protocol: "https:",
path: "/",
});
const args: FinalizeHandlerArguments<any> = { input, request };

await middleware(mockNextHandler, mockContext)(args);

expect(mockNextHandler).toHaveBeenCalledWith(args);
expect(mockContext.endpointV2?.url.href).toEqual("https://sqs.us-east-1.amazonaws.com/");
});

it("should not modify the endpoint when a custom endpoint is provided in config", async () => {
const middleware = queueUrlMiddleware({ useQueueUrlAsEndpoint: true, endpoint: "https://my-endpoint.com/" });
const input = { QueueUrl: "https://xyz.com/123/MyQueue" };
const request = new HttpRequest({
hostname: "my-endpoint.com",
protocol: "https:",
path: "/",
});
const args: FinalizeHandlerArguments<any> = { input, request };

await middleware(mockNextHandler, mockContext)(args);

expect(mockNextHandler).toHaveBeenCalledWith(args);
expect(mockContext.endpointV2?.url.href).toEqual("https://sqs.us-east-1.amazonaws.com/");
expect(mockContext.logger?.warn).not.toHaveBeenCalled();
});
});
97 changes: 97 additions & 0 deletions packages/middleware-sdk-sqs/src/queue-url.ts
@@ -0,0 +1,97 @@
import {
Endpoint,
EndpointV2,
FinalizeHandlerArguments,
FinalizeHandlerOutput,
HandlerExecutionContext,
MiddlewareStack,
Pluggable,
Provider,
RelativeMiddlewareOptions,
} from "@aws-sdk/types";
import { NoOpLogger } from "@smithy/smithy-client";

/**
* @public
*/
export interface QueueUrlInputConfig {
/**
* In cases where a QueueUrl is given as input, that
* will be preferred as the request endpoint.
*
* Set this value to false to ignore the QueueUrl and use the
* client's resolved endpoint, which may be a custom endpoint.
*/
useQueueUrlAsEndpoint?: boolean;
}

export interface QueueUrlResolvedConfig {
useQueueUrlAsEndpoint: boolean;
}

export interface PreviouslyResolved {
endpoint?: string | Endpoint | Provider<Endpoint> | EndpointV2 | Provider<EndpointV2>;
}

export const resolveQueueUrlConfig = <T>(
config: T & PreviouslyResolved & QueueUrlInputConfig
): T & QueueUrlResolvedConfig => {
return {
...config,
useQueueUrlAsEndpoint: config.useQueueUrlAsEndpoint ?? true,
};
};

/**
* @internal
*/
export function queueUrlMiddleware({ useQueueUrlAsEndpoint, endpoint }: QueueUrlResolvedConfig & PreviouslyResolved) {
return <Output extends object>(
next: (args: FinalizeHandlerArguments<any>) => Promise<FinalizeHandlerOutput<Output>>,
context: HandlerExecutionContext
): ((args: FinalizeHandlerArguments<any>) => Promise<FinalizeHandlerOutput<Output>>) => {
return async (args: FinalizeHandlerArguments<any>): Promise<FinalizeHandlerOutput<Output>> => {
const { input } = args;
const resolvedEndpoint = context.endpointV2;

if (!endpoint && input.QueueUrl && resolvedEndpoint && useQueueUrlAsEndpoint) {
const logger = context.logger instanceof NoOpLogger || !context.logger?.warn ? console : context.logger;
try {
const queueUrl: URL = new URL(input.QueueUrl);
const queueUrlOrigin: URL = new URL(queueUrl.origin);
if (resolvedEndpoint.url.origin !== queueUrlOrigin.origin) {
logger.warn(
`QueueUrl=${
input.QueueUrl
} differs from SQSClient resolved endpoint=${resolvedEndpoint.url.toString()}, using QueueUrl host as endpoint.
Set [endpoint=string] or [useQueueUrlAsEndpoint=false] on the SQSClient.`
);
resolvedEndpoint.url = queueUrlOrigin;
}
} catch (e: unknown) {
logger.warn(e);
}
}
return next(args);
};
};
}

/**
* @internal
*/
export const queueUrlMiddlewareOptions: RelativeMiddlewareOptions = {
name: "queueUrlMiddleware",
relation: "after",
toMiddleware: "endpointV2Middleware",
override: true,
};

/**
* @internal
*/
export const getQueueUrlPlugin = (config: QueueUrlResolvedConfig): Pluggable<any, any> => ({
applyToStack: (clientStack: MiddlewareStack<any, any>) => {
clientStack.addRelativeTo(queueUrlMiddleware(config), queueUrlMiddlewareOptions);
},
});

0 comments on commit 2da1b6c

Please sign in to comment.