fix(core/httpAuthSchemes): allow extensions to set signer credentials

kuhe · kuhe · commit 4411f4031c5b · 2025-03-25T12:15:34.000-04:00
diff --git a/packages/core/src/submodules/httpAuthSchemes/aws_sdk/resolveAwsSdkSigV4Config.spec.ts b/packages/core/src/submodules/httpAuthSchemes/aws_sdk/resolveAwsSdkSigV4Config.spec.ts
@@ -46,10 +46,79 @@ describe(resolveAwsSdkSigV4Config.name, () => {
 
     await config.credentials(arg);
     expect(fn).toHaveBeenCalledWith(expect.objectContaining(arg));
-    // todo: callerClientConfig should be `config` after https://github.com/aws/aws-sdk-js-v3/pull/6959.
-    // expect(fn).toHaveBeenCalledWith({
-    //   ...arg,
-    //   callerClientConfig: input,
-    // });
+    expect(fn).toHaveBeenCalledWith({
+      ...arg,
+      callerClientConfig: config,
+    });
+  });
+
+  it("should use a credentials getter/setter to normalize the memoization and config binding transform", async () => {
+    const myCredentialsProvider: AwsCredentialIdentityProvider = async (arg) => {
+      return {
+        accessKeyId: "unit-test",
+        secretAccessKey: "unit-test",
+      };
+    };
+
+    const input = {
+      credentials: myCredentialsProvider,
+      region: "us-east-1",
+      sha256: vi.fn(),
+      serviceId: "",
+      useDualstackEndpoint: async () => false,
+      useFipsEndpoint: async () => false,
+    };
+
+    const config = resolveAwsSdkSigV4Config(input);
+
+    expect(config.credentials).not.toBe(myCredentialsProvider);
+    expect(config.credentials.memoized).toBe(true);
+    expect(config.credentials.configBound).toBe(true);
+    expect(config.credentials.attributed).toBe(true);
+
+    // consistent getter retrieval
+    expect(config.credentials).toBe(config.credentials);
+
+    // no transform applied if set to itself.
+    const snapshot = config.credentials;
+    config.credentials = (() => config.credentials)();
+    expect(config.credentials).toBe(snapshot);
+
+    // re-normalizes input
+    config.credentials = myCredentialsProvider;
+    expect(config.credentials).not.toBe(myCredentialsProvider);
+    expect(config.credentials.memoized).toBe(true);
+    expect(config.credentials.configBound).toBe(true);
+    expect(config.credentials.attributed).toBe(true);
+    expect(await config.credentials()).toEqual({
+      accessKeyId: "unit-test",
+      secretAccessKey: "unit-test",
+      $source: {
+        CREDENTIALS_CODE: "e",
+      },
+    });
+
+    {
+      // no transforms applied if they are already present according to function state variables.
+      const fn = Object.assign(
+        async () => {
+          return {
+            accessKeyId: "unit-test-2",
+            secretAccessKey: "unit-test-2",
+          };
+        },
+        {
+          memoized: true,
+          configBound: true,
+          attributed: true,
+        }
+      ) as any;
+      config.credentials = fn;
+      expect(config.credentials).toBe(fn);
+      expect(await config.credentials()).toEqual({
+        accessKeyId: "unit-test-2",
+        secretAccessKey: "unit-test-2",
+      });
+    }
   });
 });
diff --git a/packages/core/src/submodules/httpAuthSchemes/aws_sdk/resolveAwsSdkSigV4Config.ts b/packages/core/src/submodules/httpAuthSchemes/aws_sdk/resolveAwsSdkSigV4Config.ts
@@ -59,6 +59,25 @@ export interface AwsSdkSigV4AuthInputConfig {
   signerConstructor?: new (options: SignatureV4Init & SignatureV4CryptoInit) => RequestSigner;
 }
 
+/**
+ * Used to indicate whether a credential provider function was memoized by this resolver.
+ * @public
+ */
+export type AwsSdkSigV4Memoized = {
+  /**
+   * The credential provider has been memoized by the AWS SDK SigV4 config resolver.
+   */
+  memoized?: boolean;
+  /**
+   * The credential provider has the caller client config object bound to its arguments.
+   */
+  configBound?: boolean;
+  /**
+   * Function is wrapped with attribution transform.
+   */
+  attributed?: boolean;
+};
+
 /**
  * @internal
  */
@@ -82,7 +101,8 @@ export interface AwsSdkSigV4AuthResolvedConfig {
    * Resolved value for input config {@link AwsSdkSigV4AuthInputConfig.credentials}
    * This provider MAY memoize the loaded credentials for certain period.
    */
-  credentials: MergeFunctions<AwsCredentialIdentityProvider, MemoizedProvider<AwsCredentialIdentity>>;
+  credentials: MergeFunctions<AwsCredentialIdentityProvider, MemoizedProvider<AwsCredentialIdentity>> &
+    AwsSdkSigV4Memoized;
   /**
    * Resolved value for input config {@link AwsSdkSigV4AuthInputConfig.signer}
    */
@@ -103,33 +123,42 @@ export interface AwsSdkSigV4AuthResolvedConfig {
 export const resolveAwsSdkSigV4Config = <T>(
   config: T & AwsSdkSigV4AuthInputConfig & AwsSdkSigV4PreviouslyResolved
 ): T & AwsSdkSigV4AuthResolvedConfig => {
-  let isUserSupplied = false;
-  // Normalize credentials
-  let credentialsProvider: AwsCredentialIdentityProvider | undefined;
-  if (config.credentials) {
-    isUserSupplied = true;
-    credentialsProvider = memoizeIdentityProvider(config.credentials, isIdentityExpired, doesIdentityRequireRefresh);
-  }
-  if (!credentialsProvider) {
-    // credentialDefaultProvider should always be populated, but in case
-    // it isn't, set a default identity provider that throws an error
-    if (config.credentialDefaultProvider) {
-      credentialsProvider = normalizeProvider(
-        config.credentialDefaultProvider(
-          Object.assign({}, config as any, {
-            parentClientConfig: config,
-          })
-        )
-      );
-    } else {
-      credentialsProvider = async () => {
-        throw new Error("`credentials` is missing");
-      };
-    }
-  }
+  let inputCredentials = config.credentials;
+  let isUserSupplied = !!config.credentials;
+  let resolvedCredentials: AwsSdkSigV4AuthResolvedConfig["credentials"] | undefined = undefined;
+
+  Object.defineProperty(config, "credentials", {
+    set(credentials: AwsSdkSigV4AuthInputConfig["credentials"]) {
+      if (credentials && credentials !== inputCredentials && credentials !== resolvedCredentials) {
+        isUserSupplied = true;
+      }
+      inputCredentials = credentials;
+      const memoizedProvider = normalizeCredentialProvider(config, {
+        credentials: inputCredentials,
+        credentialDefaultProvider: config.credentialDefaultProvider,
+      });
+      const boundProvider = bindCallerConfig(config, memoizedProvider);
+      if (isUserSupplied && !boundProvider.attributed) {
+        resolvedCredentials = async (options: Record<string, any> | undefined) =>
+          boundProvider(options).then((creds: AttributedAwsCredentialIdentity) =>
+            setCredentialFeature(creds, "CREDENTIALS_CODE", "e")
+          );
+        resolvedCredentials.memoized = boundProvider.memoized;
+        resolvedCredentials.configBound = boundProvider.configBound;
+        resolvedCredentials.attributed = true;
+      } else {
+        resolvedCredentials = boundProvider;
+      }
+    },
+    get(): AwsSdkSigV4AuthResolvedConfig["credentials"] {
+      return resolvedCredentials!;
+    },
+    enumerable: true,
+    configurable: true,
+  });
 
-  const boundCredentialsProvider = async (options: Record<string, any> | undefined) =>
-    credentialsProvider!({ ...options, callerClientConfig: config });
+  // invoke setter so that resolvedCredentials is set.
+  config.credentials = inputCredentials;
 
   // Populate sigv4 arguments
   const {
@@ -172,7 +201,7 @@ export const resolveAwsSdkSigV4Config = <T>(
 
           const params: SignatureV4Init & SignatureV4CryptoInit = {
             ...config,
-            credentials: boundCredentialsProvider,
+            credentials: config.credentials as AwsSdkSigV4AuthResolvedConfig["credentials"],
             region: config.signingRegion,
             service: config.signingName,
             sha256,
@@ -208,7 +237,7 @@ export const resolveAwsSdkSigV4Config = <T>(
 
       const params: SignatureV4Init & SignatureV4CryptoInit = {
         ...config,
-        credentials: boundCredentialsProvider,
+        credentials: config.credentials as AwsSdkSigV4AuthResolvedConfig["credentials"],
         region: config.signingRegion,
         service: config.signingName,
         sha256,
@@ -220,17 +249,16 @@ export const resolveAwsSdkSigV4Config = <T>(
     };
   }
 
-  return Object.assign(config, {
+  const resolvedConfig = Object.assign(config, {
     systemClockOffset,
     signingEscapePath,
-    credentials: isUserSupplied
-      ? async (options: Record<string, any> | undefined) =>
-          boundCredentialsProvider!(options).then((creds: AttributedAwsCredentialIdentity) =>
-            setCredentialFeature(creds, "CREDENTIALS_CODE", "e")
-          )
-      : boundCredentialsProvider!,
     signer,
   });
+
+  return resolvedConfig as typeof resolvedConfig & {
+    // this was set earlier with Object.defineProperty.
+    credentials: AwsSdkSigV4AuthResolvedConfig["credentials"];
+  };
 };
 
 /**
@@ -256,3 +284,63 @@ export interface AWSSDKSigV4AuthResolvedConfig extends AwsSdkSigV4AuthResolvedCo
  * @deprecated renamed to {@link resolveAwsSdkSigV4Config}
  */
 export const resolveAWSSDKSigV4Config = resolveAwsSdkSigV4Config;
+
+/**
+ * Normalizes the credentials to a memoized provider and sets memoized=true on the function
+ * object. This prevents multiple layering of the memoization process.
+ */
+function normalizeCredentialProvider(
+  config: Parameters<typeof resolveAwsSdkSigV4Config>[0],
+  {
+    credentials,
+    credentialDefaultProvider,
+  }: Pick<Parameters<typeof resolveAwsSdkSigV4Config>[0], "credentials" | "credentialDefaultProvider">
+): AwsSdkSigV4AuthResolvedConfig["credentials"] {
+  let credentialsProvider: AwsSdkSigV4AuthResolvedConfig["credentials"] | undefined;
+
+  if (credentials) {
+    if (!(credentials as typeof credentials & AwsSdkSigV4Memoized)?.memoized) {
+      credentialsProvider = memoizeIdentityProvider(credentials, isIdentityExpired, doesIdentityRequireRefresh)!;
+    } else {
+      credentialsProvider = credentials as AwsSdkSigV4AuthResolvedConfig["credentials"];
+    }
+  } else {
+    // credentialDefaultProvider should always be populated, but in case
+    // it isn't, set a default identity provider that throws an error
+    if (credentialDefaultProvider) {
+      credentialsProvider = normalizeProvider(
+        credentialDefaultProvider(
+          Object.assign({}, config as any, {
+            parentClientConfig: config,
+          })
+        )
+      );
+    } else {
+      credentialsProvider = async () => {
+        throw new Error(
+          "@aws-sdk/core::resolveAwsSdkSigV4Config - `credentials` not provided and no credentialDefaultProvider was configured."
+        );
+      };
+    }
+  }
+  credentialsProvider.memoized = true;
+  return credentialsProvider;
+}
+
+/**
+ * Binds the caller client config as an argument to the credentialsProvider function.
+ * Uses a state marker on the function to avoid doing this more than once.
+ */
+function bindCallerConfig(
+  config: Parameters<typeof resolveAwsSdkSigV4Config>[0],
+  credentialsProvider: AwsSdkSigV4AuthResolvedConfig["credentials"]
+): AwsSdkSigV4AuthResolvedConfig["credentials"] {
+  if (credentialsProvider.configBound) {
+    return credentialsProvider;
+  }
+  const fn: typeof credentialsProvider = async (options: Parameters<typeof credentialsProvider>[0]) =>
+    credentialsProvider({ ...options, callerClientConfig: config });
+  fn.memoized = credentialsProvider.memoized;
+  fn.configBound = true;
+  return fn;
+}
diff --git a/packages/credential-provider-node/src/credential-provider-node.integ.spec.ts b/packages/credential-provider-node/src/credential-provider-node.integ.spec.ts
@@ -1,4 +1,4 @@
-import { STS } from "@aws-sdk/client-sts";
+import { STS, STSExtensionConfiguration } from "@aws-sdk/client-sts";
 import * as credentialProviderHttp from "@aws-sdk/credential-provider-http";
 import { fromCognitoIdentity, fromCognitoIdentityPool, fromIni, fromWebToken } from "@aws-sdk/credential-providers";
 import { HttpResponse } from "@smithy/protocol-http";
@@ -1267,6 +1267,78 @@ describe("credential-provider-node integration test", () => {
     });
   });
 
+  describe("extension provided credentials", () => {
+    class OverrideCredentialsExtension {
+      private invocation = 0;
+      configure(extensionConfiguration: STSExtensionConfiguration): void {
+        extensionConfiguration.setCredentials(async () => ({
+          accessKeyId: "STS_AK" + ++this.invocation,
+          secretAccessKey: "STS_SAK" + this.invocation,
+        }));
+      }
+    }
+
+    it("allows an extension to modify client config credentials", async () => {
+      const client = new STS({
+        extensions: [new OverrideCredentialsExtension()],
+      });
+
+      const credentials = await client.config.credentials({});
+
+      expect(credentials).toEqual({
+        accessKeyId: "STS_AK1",
+        secretAccessKey: "STS_SAK1",
+        $source: {
+          CREDENTIALS_CODE: "e",
+        },
+      });
+    });
+
+    it("the extension provided credentials are still memoized", async () => {
+      const client = new STS({
+        extensions: [new OverrideCredentialsExtension()],
+      });
+
+      const credentials1 = await client.config.credentials({});
+      expect(credentials1).toEqual({
+        accessKeyId: "STS_AK1",
+        secretAccessKey: "STS_SAK1",
+        $source: {
+          CREDENTIALS_CODE: "e",
+        },
+      });
+
+      const credentials2 = await client.config.credentials({});
+      expect(credentials2).toEqual({
+        accessKeyId: "STS_AK1",
+        secretAccessKey: "STS_SAK1",
+        $source: {
+          CREDENTIALS_CODE: "e",
+        },
+      });
+
+      const credentials3 = await client.config.credentials({
+        forceRefresh: true,
+      });
+      expect(credentials3).toEqual({
+        accessKeyId: "STS_AK2",
+        secretAccessKey: "STS_SAK2",
+        $source: {
+          CREDENTIALS_CODE: "e",
+        },
+      });
+
+      const credentials4 = await client.config.credentials({});
+      expect(credentials4).toEqual({
+        accessKeyId: "STS_AK2",
+        secretAccessKey: "STS_SAK2",
+        $source: {
+          CREDENTIALS_CODE: "e",
+        },
+      });
+    });
+  });
+
   describe("No credentials available", () => {
     it("should throw CredentialsProviderError", async () => {
       process.env.AWS_EC2_METADATA_DISABLED = "true";
