chore(middleware-user-agent): use authScheme identity instead of credentials provider

Author: kuhe

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");
@@ -42,20 +51,14 @@ export async function checkFeatures(
     }
   }
 
-  if (typeof config.credentials === "function") {
-    try {
-      const credentials: AttributedAwsCredentialIdentity = await config.credentials?.();
-      if (credentials.accountId) {
-        setFeature(context, "RESOLVED_ACCOUNT_ID", "T");
-      }
-      for (const [key, value] of Object.entries(credentials.$source ?? {})) {
-        setFeature(context, key as keyof AwsSdkCredentialsFeatures, value);
-      }
-    } catch (e: unknown) {
-      // Sometimes config.credentials is a function but only throws
-      // as a way of informing users that something is missing.
-      // That error and any other credential retrieval errors are
-      // not relevant for feature-checking and should be ignored.
+  const identity = context.__smithy_context?.selectedHttpAuthScheme?.identity;
+  if ((identity as AttributedAwsCredentialIdentity).$source) {
+    const credentials = identity as AttributedAwsCredentialIdentity;
+    if (credentials.accountId) {
+      setFeature(context, "RESOLVED_ACCOUNT_ID", "T");
+    }
+    for (const [key, value] of Object.entries(credentials.$source ?? {})) {
+      setFeature(context, key as keyof AwsSdkCredentialsFeatures, value);
     }
   }
 }

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: {

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,