From 58ca4e64971433f4d9b8d13fd2cad6215b07b67f Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Tue, 8 Oct 2024 18:23:27 +0000
Subject: [PATCH] chore(middleware-user-agent): use authScheme identity instead
 of credentials provider

---
 .../src/check-features.ts                      |  9 +++++++++
 .../src/middleware-user-agent.integ.spec.ts    | 17 +++++++++--------
 .../src/user-agent-middleware.ts               | 18 +++++++++---------
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/packages/middleware-user-agent/src/check-features.ts b/packages/middleware-user-agent/src/check-features.ts
index 18b41d24ac4e..d76b6fceb06b 100644
--- a/packages/middleware-user-agent/src/check-features.ts
+++ b/packages/middleware-user-agent/src/check-features.ts
@@ -16,6 +16,11 @@ type PreviouslyResolved = Partial<{
   accountIdEndpointMode?: Provider<AccountIdEndpointMode>;
 }>;
 
+/**
+ * @internal
+ */
+const ACCOUNT_ID_ENDPOINT_REGEX = /\d{12}\.ddb/;
+
 /**
  * @internal
  * Check for features that don't have a middleware activation site but
@@ -29,6 +34,10 @@ export async function checkFeatures(
   // eslint-disable-next-line
   const request = args.request as IHttpRequest;
   if (typeof config.accountIdEndpointMode === "function") {
+    const endpointV2 = context.endpointV2;
+    if (String(endpointV2?.url?.hostname).match(ACCOUNT_ID_ENDPOINT_REGEX)) {
+      setFeature(context, "ACCOUNT_ID_ENDPOINT", "O");
+    }
     switch (await config.accountIdEndpointMode?.()) {
       case "disabled":
         setFeature(context, "ACCOUNT_ID_MODE_DISABLED", "Q");
diff --git a/packages/middleware-user-agent/src/middleware-user-agent.integ.spec.ts b/packages/middleware-user-agent/src/middleware-user-agent.integ.spec.ts
index ca98caeec642..04ba7015c472 100644
--- a/packages/middleware-user-agent/src/middleware-user-agent.integ.spec.ts
+++ b/packages/middleware-user-agent/src/middleware-user-agent.integ.spec.ts
@@ -27,24 +27,25 @@ describe("middleware-user-agent", () => {
   });
 
   describe("features", () => {
-    it("should detect DDB mapper, and account id mode", async () => {
+    it("should detect DDB mapper, account id, and account id mode", async () => {
       const client = new DynamoDB({
-        credentials: {
-          accessKeyId: "",
-          secretAccessKey: "",
-          accountId: "123",
-        },
-        accountIdEndpointMode: async () => "preferred" as const,
+        accountIdEndpointMode: async () => "required" as const,
       });
 
       const doc = DynamoDBDocument.from(client);
 
       requireRequestsFrom(doc).toMatch({
         headers: {
-          "user-agent": /(.*?) m\/d,P$/,
+          "user-agent": /(.*?) m\/d,O,R$/,
         },
       });
 
+      client.config.credentials = async () => ({
+        accessKeyId: "",
+        secretAccessKey: "",
+        accountId: "123456789012",
+      });
+
       await doc.get({
         TableName: "table",
         Key: {
diff --git a/packages/middleware-user-agent/src/user-agent-middleware.ts b/packages/middleware-user-agent/src/user-agent-middleware.ts
index 4438ede668d0..42c915d8526c 100644
--- a/packages/middleware-user-agent/src/user-agent-middleware.ts
+++ b/packages/middleware-user-agent/src/user-agent-middleware.ts
@@ -3,10 +3,10 @@ import { getUserAgentPrefix } from "@aws-sdk/util-endpoints";
 import { HttpRequest } from "@smithy/protocol-http";
 import {
   AbsoluteLocation,
-  BuildHandler,
-  BuildHandlerArguments,
-  BuildHandlerOptions,
-  BuildHandlerOutput,
+  FinalizeHandler,
+  FinalizeHandlerArguments,
+  FinalizeHandlerOutput,
+  FinalizeRequestHandlerOptions,
   HandlerExecutionContext,
   MetadataBearer,
   Pluggable,
@@ -41,10 +41,10 @@ import { encodeFeatures } from "./encode-features";
 export const userAgentMiddleware =
   (options: UserAgentResolvedConfig) =>
   <Output extends MetadataBearer>(
-    next: BuildHandler<any, any>,
+    next: FinalizeHandler<any, any>,
     context: HandlerExecutionContext | AwsHandlerExecutionContext
-  ): BuildHandler<any, any> =>
-  async (args: BuildHandlerArguments<any>): Promise<BuildHandlerOutput<Output>> => {
+  ): FinalizeHandler<any, any> =>
+  async (args: FinalizeHandlerArguments<any>): Promise<FinalizeHandlerOutput<Output>> => {
     const { request } = args;
     if (!HttpRequest.isInstance(request)) {
       return next(args);
@@ -130,9 +130,9 @@ const escapeUserAgent = (userAgentPair: UserAgentPair): string => {
     }, "") as string;
 };
 
-export const getUserAgentMiddlewareOptions: BuildHandlerOptions & AbsoluteLocation = {
+export const getUserAgentMiddlewareOptions: FinalizeRequestHandlerOptions & AbsoluteLocation = {
   name: "getUserAgentMiddleware",
-  step: "build",
+  step: "finalizeRequest",
   priority: "low",
   tags: ["SET_USER_AGENT", "USER_AGENT"],
   override: true,