From 6f5c153a3f7e9cef96dbfeaa1d9d181e1d27c888 Mon Sep 17 00:00:00 2001
From: siddsriv <sidddsri@amazon.com>
Date: Wed, 6 Sep 2023 17:50:54 +0000
Subject: [PATCH] feat(middleware-sdk-s3): add middleware for following region
 redirects

---
 packages/middleware-sdk-s3/src/index.ts       |   1 +
 .../src/region-redirect-middleware.ts         | 117 ++++++++++++++++++
 .../middleware-sdk-s3/src/s3Configuration.ts  |   9 +-
 3 files changed, 126 insertions(+), 1 deletion(-)
 create mode 100644 packages/middleware-sdk-s3/src/region-redirect-middleware.ts

diff --git a/packages/middleware-sdk-s3/src/index.ts b/packages/middleware-sdk-s3/src/index.ts
index 92c9f7291a67..02aeff027f0f 100644
--- a/packages/middleware-sdk-s3/src/index.ts
+++ b/packages/middleware-sdk-s3/src/index.ts
@@ -1,4 +1,5 @@
 export * from "./check-content-length-header";
+export * from "./region-redirect-middleware";
 export * from "./s3Configuration";
 export * from "./throw-200-exceptions";
 export * from "./validate-bucket-name";
diff --git a/packages/middleware-sdk-s3/src/region-redirect-middleware.ts b/packages/middleware-sdk-s3/src/region-redirect-middleware.ts
new file mode 100644
index 000000000000..ab56bc6d9899
--- /dev/null
+++ b/packages/middleware-sdk-s3/src/region-redirect-middleware.ts
@@ -0,0 +1,117 @@
+import {
+  HandlerExecutionContext,
+  InitializeHandler,
+  InitializeHandlerArguments,
+  InitializeHandlerOptions,
+  InitializeHandlerOutput,
+  InitializeMiddleware,
+  MetadataBearer,
+  Pluggable,
+  Provider,
+  RelativeMiddlewareOptions,
+  SerializeHandler,
+  SerializeHandlerArguments,
+  SerializeHandlerOutput,
+  SerializeMiddleware,
+} from "@smithy/types";
+
+/**
+ * @internal
+ */
+interface PreviouslyResolved {
+  region: Provider<string>;
+  followRegionRedirects: boolean;
+}
+
+/**
+ * @internal
+ */
+export function regionRedirectMiddleware(clientConfig: PreviouslyResolved): InitializeMiddleware<any, any> {
+  return <Output extends MetadataBearer>(
+      next: InitializeHandler<any, Output>,
+      context: HandlerExecutionContext
+    ): InitializeHandler<any, Output> =>
+    async (args: InitializeHandlerArguments<any>): Promise<InitializeHandlerOutput<Output>> => {
+      try {
+        return next(args);
+      } catch (err) {
+        if (
+          clientConfig.followRegionRedirects &&
+          err.Code === "PermanentRedirect" &&
+          err.$metadata.httpStatusCode === 301
+        ) {
+          try {
+            const actualRegion = err.$response.headers["x-amz-bucket-region"];
+            context.logger?.debug(`Redirecting from ${await clientConfig.region()} to ${actualRegion}`);
+            context.__s3RegionRedirect = actualRegion;
+          } catch (e) {
+            throw new Error("Region redirect failed: " + e);
+          }
+          return next(args);
+        } else {
+          throw err;
+        }
+      }
+    };
+}
+
+/**
+ * @internal
+ */
+export const regionRedirectEndpointMiddleware = (config: PreviouslyResolved): SerializeMiddleware<any, any> => {
+  return <Output extends MetadataBearer>(
+      next: SerializeHandler<any, Output>,
+      context: HandlerExecutionContext
+    ): SerializeHandler<any, Output> =>
+    async (args: SerializeHandlerArguments<any>): Promise<SerializeHandlerOutput<Output>> => {
+      const originalRegion = await config.region();
+      if (context.__s3RegionRedirect) {
+        const regionProviderRef = config.region;
+        config.region = async () => {
+          config.region = regionProviderRef;
+          return context.__s3RegionRedirect;
+        };
+      }
+      const result = await next({
+        ...args,
+      });
+      if (context.__s3RegionRedirect) {
+        const region = await config.region();
+        if (originalRegion !== region) {
+          throw new Error("Region was not restored following S3 region redirect.");
+        }
+      }
+      return result;
+    };
+};
+
+/**
+ * @internal
+ */
+export const regionRedirectMiddlewareOptions: InitializeHandlerOptions = {
+  step: "initialize",
+  tags: ["REGION_REDIRECT", "S3"],
+  name: "regionRedirectMiddleware",
+  override: true,
+};
+
+/**
+ * @internal
+ */
+export const regionRedirectEndpointMiddlewareOptions: RelativeMiddlewareOptions = {
+  tags: ["REGION_REDIRECT", "S3"],
+  name: "regionRedirectEndpointMiddleware",
+  override: true,
+  relation: "before",
+  toMiddleware: "endpointV2Middleware",
+};
+
+/**
+ * @internal
+ */
+export const getRegionRedirectMiddlewarePlugin = (clientConfig: PreviouslyResolved): Pluggable<any, any> => ({
+  applyToStack: (clientStack) => {
+    clientStack.add(regionRedirectMiddleware(clientConfig), regionRedirectMiddlewareOptions);
+    clientStack.addRelativeTo(regionRedirectEndpointMiddleware(clientConfig), regionRedirectEndpointMiddlewareOptions);
+  },
+});
diff --git a/packages/middleware-sdk-s3/src/s3Configuration.ts b/packages/middleware-sdk-s3/src/s3Configuration.ts
index ba91f3f16293..6aee45ceffd9 100644
--- a/packages/middleware-sdk-s3/src/s3Configuration.ts
+++ b/packages/middleware-sdk-s3/src/s3Configuration.ts
@@ -1,6 +1,6 @@
 /**
  * @public
- * 
+ *
  * All endpoint parameters with built-in bindings of AWS::S3::*
  */
 export interface S3InputConfig {
@@ -17,12 +17,18 @@ export interface S3InputConfig {
    * Whether multi-region access points (MRAP) should be disabled.
    */
   disableMultiregionAccessPoints?: boolean;
+  /**
+   * If you receive a permanent redirect with status 301,
+   * the client will retry your request with the corrected region.
+   */
+  followRegionRedirects?: boolean;
 }
 
 export interface S3ResolvedConfig {
   forcePathStyle: boolean;
   useAccelerateEndpoint: boolean;
   disableMultiregionAccessPoints: boolean;
+  followRegionRedirects: boolean;
 }
 
 export const resolveS3Config = <T>(input: T & S3InputConfig): T & S3ResolvedConfig => ({
@@ -30,4 +36,5 @@ export const resolveS3Config = <T>(input: T & S3InputConfig): T & S3ResolvedConf
   forcePathStyle: input.forcePathStyle ?? false,
   useAccelerateEndpoint: input.useAccelerateEndpoint ?? false,
   disableMultiregionAccessPoints: input.disableMultiregionAccessPoints ?? false,
+  followRegionRedirects: input.followRegionRedirects ?? false,
 });