WIP Add support for DSQL iam authentication

common/lib/authentication/dsql_iam_authentication_plugin_factory.ts

@@ -0,0 +1,37 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { ConnectionPluginFactory } from "../plugin_factory";
+import { PluginService } from "../plugin_service";
+import { ConnectionPlugin } from "../connection_plugin";
+import { AwsWrapperError } from "../utils/errors";
+import { Messages } from "../utils/messages";
+import { DSQLTokenUtils } from "../utils/dsql_token_utils";
+
+export class DsqlIamAuthenticationPluginFactory extends ConnectionPluginFactory {
+  private static iamAuthenticationPlugin: any;
+
+  async getInstance(pluginService: PluginService, properties: object): Promise<ConnectionPlugin> {
+    try {
+      if (!DsqlIamAuthenticationPluginFactory.iamAuthenticationPlugin) {
+        DsqlIamAuthenticationPluginFactory.iamAuthenticationPlugin = await import("./iam_authentication_plugin");
+      }
+      return new DsqlIamAuthenticationPluginFactory.iamAuthenticationPlugin.IamAuthenticationPlugin(pluginService, new DSQLTokenUtils());
+    } catch (error: any) {
+      throw new AwsWrapperError(Messages.get("ConnectionPluginChainBuilder.errorImportingPlugin", error.message, "IamAuthenticationPlugin"));
+    }
+  }
+}

common/lib/authentication/iam_authentication_plugin.ts

@@ -25,17 +25,20 @@ import { WrapperProperties } from "../wrapper_property";
 import { IamAuthUtils, TokenInfo } from "../utils/iam_auth_utils";
 import { ClientWrapper } from "../client_wrapper";
 import { RegionUtils } from "../utils/region_utils";
+import { TokenUtils } from "../utils/token_utils";
 
 export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
   private static readonly SUBSCRIBED_METHODS = new Set<string>(["connect", "forceConnect"]);
   protected static readonly tokenCache = new Map<string, TokenInfo>();
   private readonly telemetryFactory;
   private readonly fetchTokenCounter;
+  private readonly tokenUtils: TokenUtils;
   private pluginService: PluginService;
 
-  constructor(pluginService: PluginService) {
+  constructor(pluginService: PluginService, tokenUtils: TokenUtils) {
     super();
     this.pluginService = pluginService;
+    this.tokenUtils = tokenUtils;
     this.telemetryFactory = this.pluginService.getTelemetryFactory();
     this.fetchTokenCounter = this.telemetryFactory.createCounter("iam.fetchTokenCount");
   }
@@ -90,7 +93,7 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
       WrapperProperties.PASSWORD.set(props, tokenInfo.token);
     } else {
       const tokenExpiry: number = Date.now() + tokenExpirationSec * 1000;
-      const token = await IamAuthUtils.generateAuthenticationToken(
+      const token = await this.tokenUtils.generateAuthenticationToken(
         host,
         port,
         region,
@@ -117,7 +120,7 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
       // Try to generate a new token and try to connect again
 
       const tokenExpiry: number = Date.now() + tokenExpirationSec * 1000;
-      const token = await IamAuthUtils.generateAuthenticationToken(
+      const token = await this.tokenUtils.generateAuthenticationToken(
         host,
         port,
         region,

common/lib/authentication/iam_authentication_plugin_factory.ts

@@ -19,6 +19,7 @@ import { PluginService } from "../plugin_service";
 import { ConnectionPlugin } from "../connection_plugin";
 import { AwsWrapperError } from "../utils/errors";
 import { Messages } from "../utils/messages";
+import { RdsTokenUtils } from "../utils/rds_token_utils";
 
 export class IamAuthenticationPluginFactory extends ConnectionPluginFactory {
   private static iamAuthenticationPlugin: any;
@@ -28,7 +29,7 @@ export class IamAuthenticationPluginFactory extends ConnectionPluginFactory {
       if (!IamAuthenticationPluginFactory.iamAuthenticationPlugin) {
         IamAuthenticationPluginFactory.iamAuthenticationPlugin = await import("./iam_authentication_plugin");
       }
-      return new IamAuthenticationPluginFactory.iamAuthenticationPlugin.IamAuthenticationPlugin(pluginService);
+      return new IamAuthenticationPluginFactory.iamAuthenticationPlugin.IamAuthenticationPlugin(pluginService, new RdsTokenUtils());
     } catch (error: any) {
       throw new AwsWrapperError(Messages.get("ConnectionPluginChainBuilder.errorImportingPlugin", error.message, "IamAuthenticationPlugin"));
     }

common/lib/connection_plugin_chain_builder.ts

@@ -22,6 +22,7 @@ import { Messages } from "./utils/messages";
 import { logger } from "../logutils";
 import { DefaultPlugin } from "./plugins/default_plugin";
 import { IamAuthenticationPluginFactory } from "./authentication/iam_authentication_plugin_factory";
+import { DsqlIamAuthenticationPluginFactory } from "./authentication/dsql_iam_authentication_plugin_factory";
 import { ExecuteTimePluginFactory } from "./plugins/execute_time_plugin_factory";
 import { ConnectTimePluginFactory } from "./plugins/connect_time_plugin_factory";
 import { AwsSecretsManagerPluginFactory } from "./authentication/aws_secrets_manager_plugin_factory";
@@ -70,6 +71,7 @@ export class ConnectionPluginChainBuilder {
     ["fastestResponseStrategy", { factory: FastestResponseStrategyPluginFactory, weight: 900 }],
     ["limitless", { factory: LimitlessConnectionPluginFactory, weight: 950 }],
     ["iam", { factory: IamAuthenticationPluginFactory, weight: 1000 }],
+    ["iamDsql", { factory: DsqlIamAuthenticationPluginFactory, weight: 1010 }],
     ["secretsManager", { factory: AwsSecretsManagerPluginFactory, weight: 1100 }],
     ["federatedAuth", { factory: FederatedAuthPluginFactory, weight: 1200 }],
     ["okta", { factory: OktaAuthPluginFactory, weight: 1300 }],
@@ -90,6 +92,7 @@ export class ConnectionPluginChainBuilder {
     [HostMonitoring2PluginFactory, 810],
     [LimitlessConnectionPluginFactory, 950],
     [IamAuthenticationPluginFactory, 1000],
+    [DsqlIamAuthenticationPluginFactory, 1010],
     [AwsSecretsManagerPluginFactory, 1100],
     [FederatedAuthPluginFactory, 1200],
     [OktaAuthPluginFactory, 1300],

common/lib/plugins/federated_auth/federated_auth_plugin.ts

@@ -28,11 +28,13 @@ import { SamlUtils } from "../../utils/saml_utils";
 import { ClientWrapper } from "../../client_wrapper";
 import { TelemetryCounter } from "../../utils/telemetry/telemetry_counter";
 import { RegionUtils } from "../../utils/region_utils";
+import { TokenUtils } from "../../utils/token_utils";
 
 export class FederatedAuthPlugin extends AbstractConnectionPlugin {
   protected static readonly tokenCache = new Map<string, TokenInfo>();
   protected rdsUtils: RdsUtils = new RdsUtils();
   protected pluginService: PluginService;
+  private readonly tokenUtils: TokenUtils;
   private static readonly subscribedMethods = new Set<string>(["connect", "forceConnect"]);
   private readonly credentialsProviderFactory: CredentialsProviderFactory;
   private readonly fetchTokenCounter: TelemetryCounter;
@@ -41,10 +43,11 @@ export class FederatedAuthPlugin extends AbstractConnectionPlugin {
     return FederatedAuthPlugin.subscribedMethods;
   }
 
-  constructor(pluginService: PluginService, credentialsProviderFactory: CredentialsProviderFactory) {
+  constructor(pluginService: PluginService, credentialsProviderFactory: CredentialsProviderFactory, tokenUtils: TokenUtils) {
     super();
     this.credentialsProviderFactory = credentialsProviderFactory;
     this.pluginService = pluginService;
+    this.tokenUtils = tokenUtils;
     this.fetchTokenCounter = this.pluginService.getTelemetryFactory().createCounter("federatedAuth.fetchToken.count");
   }
 
@@ -109,7 +112,7 @@ export class FederatedAuthPlugin extends AbstractConnectionPlugin {
     }
     const tokenExpiry: number = Date.now() + tokenExpirationSec * 1000;
     const port = IamAuthUtils.getIamPort(props, hostInfo, this.pluginService.getDialect().getDefaultPort());
-    const token = await IamAuthUtils.generateAuthenticationToken(
+    const token = await this.tokenUtils.generateAuthenticationToken(
       iamHost,
       port,
       region,

common/lib/plugins/federated_auth/federated_auth_plugin_factory.ts

@@ -19,6 +19,7 @@ import { PluginService } from "../../plugin_service";
 import { ConnectionPlugin } from "../../connection_plugin";
 import { AwsWrapperError } from "../../utils/errors";
 import { Messages } from "../../utils/messages";
+import { RdsTokenUtils } from "../../utils/rds_token_utils";
 
 export class FederatedAuthPluginFactory extends ConnectionPluginFactory {
   private static federatedAuthPlugin: any;
@@ -35,7 +36,7 @@ export class FederatedAuthPluginFactory extends ConnectionPluginFactory {
       }
 
       const adfsCredentialsProviderFactory = new FederatedAuthPluginFactory.adfsCredentialsProvider.AdfsCredentialsProviderFactory(pluginService);
-      return new FederatedAuthPluginFactory.federatedAuthPlugin.FederatedAuthPlugin(pluginService, adfsCredentialsProviderFactory);
+      return new FederatedAuthPluginFactory.federatedAuthPlugin.FederatedAuthPlugin(pluginService, adfsCredentialsProviderFactory, new RdsTokenUtils());
     } catch (error: any) {
       throw new AwsWrapperError(Messages.get("ConnectionPluginChainBuilder.errorImportingPlugin", error.message, "FederatedAuthPlugin"));
     }

common/lib/plugins/federated_auth/okta_auth_plugin.ts

@@ -28,18 +28,21 @@ import { AwsWrapperError } from "../../utils/errors";
 import { ClientWrapper } from "../../client_wrapper";
 import { TelemetryCounter } from "../../utils/telemetry/telemetry_counter";
 import { RegionUtils } from "../../utils/region_utils";
+import { TokenUtils } from "../../utils/token_utils";
 
 export class OktaAuthPlugin extends AbstractConnectionPlugin {
   protected static readonly tokenCache = new Map<string, TokenInfo>();
   private static readonly subscribedMethods = new Set<string>(["connect", "forceConnect"]);
   protected pluginService: PluginService;
   protected rdsUtils = new RdsUtils();
+  private readonly tokenUtils: TokenUtils;
   private readonly credentialsProviderFactory: CredentialsProviderFactory;
   private readonly fetchTokenCounter: TelemetryCounter;
 
-  constructor(pluginService: PluginService, credentialsProviderFactory: CredentialsProviderFactory) {
+  constructor(pluginService: PluginService, credentialsProviderFactory: CredentialsProviderFactory, tokenUtils: TokenUtils) {
     super();
     this.pluginService = pluginService;
+    this.tokenUtils = tokenUtils;
     this.credentialsProviderFactory = credentialsProviderFactory;
     this.fetchTokenCounter = this.pluginService.getTelemetryFactory().createCounter("oktaAuth.fetchToken.count");
   }
@@ -111,7 +114,7 @@ export class OktaAuthPlugin extends AbstractConnectionPlugin {
     const tokenExpiry = Date.now() + tokenExpirationSec * 1000;
     const port = IamAuthUtils.getIamPort(props, hostInfo, this.pluginService.getDialect().getDefaultPort());
     this.fetchTokenCounter.inc();
-    const token = await IamAuthUtils.generateAuthenticationToken(
+    const token = await this.tokenUtils.generateAuthenticationToken(
       iamHost,
       port,
       region,

common/lib/plugins/federated_auth/okta_auth_plugin_factory.ts

@@ -19,6 +19,7 @@ import { PluginService } from "../../plugin_service";
 import { ConnectionPlugin } from "../../connection_plugin";
 import { AwsWrapperError } from "../../utils/errors";
 import { Messages } from "../../utils/messages";
+import { RdsTokenUtils } from "../../utils/rds_token_utils";
 
 export class OktaAuthPluginFactory extends ConnectionPluginFactory {
   private static oktaAuthPlugin: any;
@@ -34,7 +35,7 @@ export class OktaAuthPluginFactory extends ConnectionPluginFactory {
       }
 
       const oktaCredentialsProviderFactory = new OktaAuthPluginFactory.oktaCredentialsProviderFactory.OktaCredentialsProviderFactory(pluginService);
-      return new OktaAuthPluginFactory.oktaAuthPlugin.OktaAuthPlugin(pluginService, oktaCredentialsProviderFactory);
+      return new OktaAuthPluginFactory.oktaAuthPlugin.OktaAuthPlugin(pluginService, oktaCredentialsProviderFactory, new RdsTokenUtils());
     } catch (error: any) {
       throw new AwsWrapperError(Messages.get("ConnectionPluginChainBuilder.errorImportingPlugin", error.message, "OktaAuthPlugin"));
     }

common/lib/utils/dsql_token_utils.ts

@@ -0,0 +1,50 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { AwsCredentialIdentity, AwsCredentialIdentityProvider } from "@smithy/types/dist-types/identity/awsCredentialIdentity";
+import { PluginService } from "../plugin_service";
+import { TokenUtils } from "../utils/token_utils";
+import { DsqlSigner } from "@aws-sdk/dsql-signer";
+import { TelemetryTraceLevel } from "./telemetry/telemetry_trace_level";
+
+export class DSQLTokenUtils extends TokenUtils {
+    private static readonly TELEMETRY_FETCH_TOKEN = "fetch DSQL IAM token";
+
+    public async generateAuthenticationToken(
+      hostname: string,
+      port: number,
+      region: string,
+      user: string,
+      credentials: AwsCredentialIdentity | AwsCredentialIdentityProvider,
+      pluginService: PluginService
+    ): Promise<string> {
+      const telemetryFactory = pluginService.getTelemetryFactory();
+      const telemetryContext = telemetryFactory.openTelemetryContext(DSQLTokenUtils.TELEMETRY_FETCH_TOKEN, TelemetryTraceLevel.NESTED);
+      return await telemetryContext.start(async () => {
+        const signer = new DsqlSigner({
+        hostname: hostname,
+        region,
+        });
+
+        if (user === "admin") {
+          return signer.getDbConnectAdminAuthToken();
+        }
+        else {
+          return signer.getDbConnectAuthToken()
+        }
+      });
+    }
+}

common/lib/utils/iam_auth_utils.ts

@@ -64,29 +64,6 @@ export class IamAuthUtils {
   public static getCacheKey(port: number, user?: string, hostname?: string, region?: string): string {
     return `${region}:${hostname}:${port}:${user}`;
   }
-
-  public static async generateAuthenticationToken(
-    hostname: string,
-    port: number,
-    region: string,
-    user: string,
-    credentials: AwsCredentialIdentity | AwsCredentialIdentityProvider,
-    pluginService: PluginService
-  ): Promise<string> {
-    const telemetryFactory = pluginService.getTelemetryFactory();
-    const telemetryContext = telemetryFactory.openTelemetryContext(IamAuthUtils.TELEMETRY_FETCH_TOKEN, TelemetryTraceLevel.NESTED);
-    return await telemetryContext.start(async () => {
-      const signer = new Signer({
-        hostname: hostname,
-        port: port,
-        region: region,
-        credentials: credentials,
-        username: user
-      });
-
-      return signer.getAuthToken();
-    });
-  }
 }
 
 export class TokenInfo {

common/lib/utils/rds_token_utils.ts

@@ -0,0 +1,48 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { AwsCredentialIdentity, AwsCredentialIdentityProvider } from "@smithy/types/dist-types/identity/awsCredentialIdentity";
+import { PluginService } from "../plugin_service";
+import { TokenUtils } from "../utils/token_utils";
+import { Signer } from "@aws-sdk/rds-signer";
+import { TelemetryTraceLevel } from "./telemetry/telemetry_trace_level";
+
+export class RdsTokenUtils extends TokenUtils {
+    private static readonly TELEMETRY_FETCH_TOKEN = "fetch IAM token";
+
+    public async generateAuthenticationToken(
+      hostname: string,
+      port: number,
+      region: string,
+      user: string,
+      credentials: AwsCredentialIdentity | AwsCredentialIdentityProvider,
+      pluginService: PluginService
+    ): Promise<string> {
+      const telemetryFactory = pluginService.getTelemetryFactory();
+      const telemetryContext = telemetryFactory.openTelemetryContext(RdsTokenUtils.TELEMETRY_FETCH_TOKEN, TelemetryTraceLevel.NESTED);
+      return await telemetryContext.start(async () => {
+        const signer = new Signer({
+        hostname: hostname,
+        port: port,
+        region: region,
+        credentials: credentials,
+        username: user
+        });
+
+        return signer.getAuthToken();
+      });
+    }
+}

common/lib/utils/token_utils.ts

@@ -0,0 +1,30 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { AwsCredentialIdentity, AwsCredentialIdentityProvider } from "@smithy/types/dist-types/identity/awsCredentialIdentity";
+import { PluginService } from "../plugin_service";
+
+export abstract class TokenUtils {
+
+  public abstract  generateAuthenticationToken(
+      hostname: string,
+      port: number,
+      region: string,
+      user: string,
+      credentials: AwsCredentialIdentity | AwsCredentialIdentityProvider,
+      pluginService: PluginService
+    ): Promise<string>;
+}