fix(middleware-sdk-sqs): add middleware to prioritize QueueUrl endpoint 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 <[email protected]>

Author: siddsriv

clients/client-sqs/src/SQSClient.ts

@@ -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,
@@ -311,6 +317,7 @@ export type SQSClientConfigType = Partial<__SmithyConfiguration<__HttpHandlerOpt
   EndpointInputConfig<EndpointParameters> &
   RetryInputConfig &
   HostHeaderInputConfig &
+  QueueUrlInputConfig &
   UserAgentInputConfig &
   HttpAuthSchemeInputConfig &
   ClientInputEndpointParameters;
@@ -331,6 +338,7 @@ export type SQSClientResolvedConfigType = __SmithyResolvedConfiguration<__HttpHa
   EndpointResolvedConfig<EndpointParameters> &
   RetryResolvedConfig &
   HostHeaderResolvedConfig &
+  QueueUrlResolvedConfig &
   UserAgentResolvedConfig &
   HttpAuthSchemeResolvedConfig &
   ClientResolvedEndpointParameters;
@@ -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, {

codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddSqsDependency.java

@@ -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;
@@ -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()
         );
     }

packages/middleware-sdk-sqs/package.json

@@ -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",

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";

packages/middleware-sdk-sqs/src/middleware-sdk-sqs.integ.spec.ts

@@ -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 = "";
@@ -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",
+      });
+    });
+  });
 });

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();
+  });
+});

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);
+  },
+});