feat(smithy-typescript-codegen): allow deferred resolution for api key config (#588)

Co-authored-by: Eduardo Rodrigues <[email protected]>

Author: eduardomourar

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/TypeScriptDependency.java

@@ -85,6 +85,9 @@ public enum TypeScriptDependency implements SymbolDependencyContainer {
     AWS_SDK_FETCH_HTTP_HANDLER("dependencies", "@aws-sdk/fetch-http-handler", false),
     AWS_SDK_NODE_HTTP_HANDLER("dependencies", "@aws-sdk/node-http-handler", false),
 
+    // Conditionally added when setting the auth middleware.
+    AWS_SDK_UTIL_MIDDLEWARE("dependencies", "@aws-sdk/util-middleware", false),
+
     // Conditionally added if a event stream shape is found anywhere in the model
     AWS_SDK_EVENTSTREAM_SERDE_CONFIG_RESOLVER("dependencies", "@aws-sdk/eventstream-serde-config-resolver",
             false),

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddHttpApiKeyAuthPlugin.java

@@ -32,6 +32,7 @@
 import software.amazon.smithy.model.traits.OptionalAuthTrait;
 import software.amazon.smithy.typescript.codegen.CodegenUtils;
 import software.amazon.smithy.typescript.codegen.TypeScriptCodegenContext;
+import software.amazon.smithy.typescript.codegen.TypeScriptDependency;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings;
 import software.amazon.smithy.typescript.codegen.TypeScriptWriter;
 import software.amazon.smithy.utils.IoUtils;
@@ -134,6 +135,7 @@ private void writeAdditionalFiles(
         writerFactory.accept(
                 Paths.get(CodegenUtils.SOURCE_FOLDER, "middleware", INTEGRATION_NAME, "index.ts").toString(),
                 writer -> {
+                        writer.addDependency(TypeScriptDependency.AWS_SDK_UTIL_MIDDLEWARE);
                         String source = IoUtils.readUtf8Resource(getClass(), "http-api-key-auth.ts");
                         writer.write("$L$L", noTouchNoticePrefix, "http-api-key-auth.ts");
                         writer.write("$L", source);

smithy-typescript-codegen/src/main/resources/software/amazon/smithy/typescript/codegen/integration/http-api-key-auth.spec.ts

@@ -1,14 +1,16 @@
+import { HttpRequest } from "@aws-sdk/protocol-http";
+import { MiddlewareStack } from "@aws-sdk/types";
+
 import {
   getHttpApiKeyAuthPlugin,
   httpApiKeyAuthMiddleware,
   resolveHttpApiKeyAuthConfig,
 } from "./index";
-import { HttpRequest } from "@aws-sdk/protocol-http";
 
 describe("resolveHttpApiKeyAuthConfig", () => {
   it("should return the input unchanged", () => {
     const config = {
-      apiKey: "exampleApiKey",
+      apiKey: () => Promise.resolve("example-api-key"),
     };
     expect(resolveHttpApiKeyAuthConfig(config)).toEqual(config);
   });
@@ -18,32 +20,32 @@ describe("getHttpApiKeyAuthPlugin", () => {
   it("should apply the middleware to the stack", () => {
     const plugin = getHttpApiKeyAuthPlugin(
       {
-        apiKey: "exampleApiKey",
+        apiKey: () => Promise.resolve("example-api-key"),
       },
       {
         in: "query",
         name: "key",
       }
     );
 
-    const mockAdd = jest.fn();
+    const mockApplied = jest.fn();
     const mockOther = jest.fn();
 
     // TODO there's got to be a better way to do this mocking
     plugin.applyToStack({
-      add: mockAdd,
+      addRelativeTo: mockApplied,
       // We don't expect any of these others to be called.
-      addRelativeTo: mockOther,
+      add: mockOther,
       concat: mockOther,
       resolve: mockOther,
       applyToStack: mockOther,
       use: mockOther,
       clone: mockOther,
       remove: mockOther,
       removeByTag: mockOther,
-    });
+    } as unknown as MiddlewareStack<any, any>);
 
-    expect(mockAdd.mock.calls.length).toEqual(1);
+    expect(mockApplied.mock.calls.length).toEqual(1);
     expect(mockOther.mock.calls.length).toEqual(0);
   });
 });
@@ -59,7 +61,7 @@ describe("httpApiKeyAuthMiddleware", () => {
     it("should set the query parameter if the location is `query`", async () => {
       const middleware = httpApiKeyAuthMiddleware(
         {
-          apiKey: "exampleApiKey",
+          apiKey: () => Promise.resolve("example-api-key"),
         },
         {
           in: "query",
@@ -77,7 +79,7 @@ describe("httpApiKeyAuthMiddleware", () => {
       expect(mockNextHandler.mock.calls.length).toEqual(1);
       expect(
         mockNextHandler.mock.calls[0][0].request.query.key
-      ).toBe("exampleApiKey");
+      ).toBe("example-api-key");
     });
 
     it("should skip if the api key has not been set", async () => {
@@ -122,7 +124,7 @@ describe("httpApiKeyAuthMiddleware", () => {
     it("should set the API key in the lower-cased named header", async () => {
       const middleware = httpApiKeyAuthMiddleware(
         {
-          apiKey: "exampleApiKey",
+          apiKey: () => Promise.resolve("example-api-key"),
         },
         {
           in: "header",
@@ -140,13 +142,13 @@ describe("httpApiKeyAuthMiddleware", () => {
       expect(mockNextHandler.mock.calls.length).toEqual(1);
       expect(
         mockNextHandler.mock.calls[0][0].request.headers.authorization
-      ).toBe("exampleApiKey");
+      ).toBe("example-api-key");
     });
 
     it("should set the API key in the named header with the provided scheme", async () => {
       const middleware = httpApiKeyAuthMiddleware(
         {
-          apiKey: "exampleApiKey",
+          apiKey: () => Promise.resolve("example-api-key"),
         },
         {
           in: "header",
@@ -164,7 +166,7 @@ describe("httpApiKeyAuthMiddleware", () => {
       expect(mockNextHandler.mock.calls.length).toEqual(1);
       expect(
         mockNextHandler.mock.calls[0][0].request.headers.authorization
-      ).toBe("exampleScheme exampleApiKey");
+      ).toBe("exampleScheme example-api-key");
     });
   });
 });

smithy-typescript-codegen/src/main/resources/software/amazon/smithy/typescript/codegen/integration/http-api-key-auth.ts

@@ -1,12 +1,8 @@
 // derived from https://github.com/aws/aws-sdk-js-v3/blob/e35f78c97fa6710ff9c444351893f0f06755e771/packages/middleware-endpoint-discovery/src/endpointDiscoveryMiddleware.ts
 
 import { HttpRequest } from "@aws-sdk/protocol-http";
-import {
-  AbsoluteLocation,
-  BuildHandlerOptions,
-  BuildMiddleware,
-  Pluggable,
-} from "@aws-sdk/types";
+import { BuildMiddleware, Pluggable, Provider, RelativeMiddlewareOptions } from "@aws-sdk/types";
+import { normalizeProvider } from "@aws-sdk/util-middleware";
 
 interface HttpApiKeyAuthMiddlewareConfig {
   /**
@@ -36,28 +32,31 @@ export interface HttpApiKeyAuthInputConfig {
    *
    * This is optional because some operations may not require an API key.
    */
-  apiKey?: string;
+  apiKey?: string | Provider<string>;
 }
 
-interface PreviouslyResolved {}
+export interface ApiKeyPreviouslyResolved {}
 
 export interface HttpApiKeyAuthResolvedConfig {
   /**
    * The API key to use when making requests.
    *
    * This is optional because some operations may not require an API key.
    */
-  apiKey?: string;
+  apiKey?: Provider<string>;
 }
 
 // We have to provide a resolve function when we have config, even if it doesn't
 // actually do anything to the input value. "If any of inputConfig, resolvedConfig,
 // or resolveFunction are set, then all of inputConfig, resolvedConfig, and
 // resolveFunction must be set."
 export function resolveHttpApiKeyAuthConfig<T>(
-  input: T & PreviouslyResolved & HttpApiKeyAuthInputConfig
+  input: T & ApiKeyPreviouslyResolved & HttpApiKeyAuthInputConfig,
 ): T & HttpApiKeyAuthResolvedConfig {
-  return input;
+  return {
+    ...input,
+    apiKey: input.apiKey ? normalizeProvider(input.apiKey) : undefined,
+  };
 }
 
 /**
@@ -69,23 +68,25 @@ export function resolveHttpApiKeyAuthConfig<T>(
  * prefixed with a scheme. If the trait says to put the API key into a named
  * query parameter, that query parameter will be used.
  *
- * @param resolvedConfig the client configuration. Must include the API key value.
- * @param options the plugin options (location of the parameter, name, and optional scheme)
+ * @param pluginConfig the client configuration. Includes the function that will return the API key value.
+ * @param middlewareConfig the plugin options (location of the parameter, name, and optional scheme)
  * @returns a function that processes the HTTP request and passes it on to the next handler
  */
 export const httpApiKeyAuthMiddleware =
   <Input extends object, Output extends object>(
     pluginConfig: HttpApiKeyAuthResolvedConfig,
-    middlewareConfig: HttpApiKeyAuthMiddlewareConfig
+    middlewareConfig: HttpApiKeyAuthMiddlewareConfig,
   ): BuildMiddleware<Input, Output> =>
   (next) =>
   async (args) => {
     if (!HttpRequest.isInstance(args.request)) return next(args);
 
+    const apiKey = pluginConfig.apiKey && (await pluginConfig.apiKey());
+
     // This middleware will not be injected if the operation has the @optionalAuth trait.
     // We don't know if we're the only auth middleware, so let the service deal with the
     // absence of the API key (or let other middleware do its job).
-    if (!pluginConfig.apiKey) {
+    if (!apiKey) {
       return next(args);
     }
 
@@ -98,36 +99,35 @@ export const httpApiKeyAuthMiddleware =
           ...(middlewareConfig.in === "header" && {
             // Set the header, even if it's already been set.
             [middlewareConfig.name.toLowerCase()]: middlewareConfig.scheme
-              ? `${middlewareConfig.scheme} ${pluginConfig.apiKey}`
-              : pluginConfig.apiKey,
+              ? `${middlewareConfig.scheme} ${apiKey}`
+              : apiKey,
           }),
         },
         query: {
           ...args.request.query,
           // Set the query parameter, even if it's already been set.
-          ...(middlewareConfig.in === "query" && { [middlewareConfig.name]: pluginConfig.apiKey }),
+          ...(middlewareConfig.in === "query" && { [middlewareConfig.name]: apiKey }),
         },
       },
     });
   };
 
-export const httpApiKeyAuthMiddlewareOptions: BuildHandlerOptions &
-  AbsoluteLocation = {
+export const httpApiKeyAuthMiddlewareOptions: RelativeMiddlewareOptions = {
   name: "httpApiKeyAuthMiddleware",
-  step: "build",
-  priority: "low",
-  tags: ["AUTHORIZATION"],
+  tags: ["APIKEY", "AUTH"],
+  relation: "after",
+  toMiddleware: "retryMiddleware",
   override: true,
 };
 
 export const getHttpApiKeyAuthPlugin = (
   pluginConfig: HttpApiKeyAuthResolvedConfig,
-  middlewareConfig: HttpApiKeyAuthMiddlewareConfig
+  middlewareConfig: HttpApiKeyAuthMiddlewareConfig,
 ): Pluggable<any, any> => ({
   applyToStack: (clientStack) => {
-    clientStack.add(
+    clientStack.addRelativeTo(
       httpApiKeyAuthMiddleware(pluginConfig, middlewareConfig),
-      httpApiKeyAuthMiddlewareOptions
+      httpApiKeyAuthMiddlewareOptions,
     );
   },
 });