Refactor getClient and getLegacyClient to use authentication method r…

…egistry

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>

src/plugins/data_source/common/data_sources/types.ts

@@ -14,6 +14,7 @@ export interface DataSourceAttributes extends SavedObjectAttributes {
     credentials: UsernamePasswordTypedContent | SigV4Content | undefined | AuthTypeContent;
   };
   lastUpdatedTime?: string;
+  name: AuthType | string;
 }
 
 export interface AuthTypeContent {
@@ -30,6 +31,7 @@ export interface SigV4Content extends SavedObjectAttributes {
   secretKey: string;
   region: string;
   service?: SigV4ServiceName;
+  sessionToken?: string;
 }
 
 export interface UsernamePasswordTypedContent extends SavedObjectAttributes {

src/plugins/data_source/server/client/configure_client.ts

@@ -7,7 +7,7 @@ import { Client, ClientOptions } from '@opensearch-project/opensearch';
 import { Client as LegacyClient } from 'elasticsearch';
 import { Credentials } from 'aws-sdk';
 import { AwsSigv4Signer } from '@opensearch-project/opensearch/aws';
-import { Logger } from '../../../../../src/core/server';
+import { Logger, OpenSearchDashboardsRequest } from '../../../../../src/core/server';
 import {
   AuthType,
   DataSourceAttributes,
@@ -27,6 +27,8 @@ import {
   getDataSource,
   generateCacheKey,
 } from './configure_client_utils';
+import { IAuthenticationMethodRegistery } from '../auth_registry';
+import { authRegistryCredentialProvider } from '../util/credential_provider';
 
 export const configureClient = async (
   {
@@ -35,6 +37,8 @@ export const configureClient = async (
     cryptography,
     testClientDataSourceAttr,
     customApiSchemaRegistryPromise,
+    request,
+    authRegistry,
   }: DataSourceClientParams,
   openSearchClientPoolSetup: OpenSearchClientPoolSetup,
   config: DataSourcePluginConfigType,
@@ -80,6 +84,8 @@ export const configureClient = async (
       cryptography,
       rootClient,
       dataSourceId,
+      request,
+      authRegistry,
       requireDecryption
     );
   } catch (error: any) {
@@ -101,6 +107,8 @@ export const configureClient = async (
  * @param config data source config
  * @param addClientToPool function to add client to client pool
  * @param dataSourceId id of data source saved Object
+ * @param request OpenSearch Dashboards incoming request to read client parameters from header.
+ * @param authRegistry registry to retrieve the credentials provider for the authentication method in order to return the client
  * @param requireDecryption false when creating test client before data source exists
  * @returns Promise of query client
  */
@@ -112,15 +120,31 @@ const getQueryClient = async (
   cryptography?: CryptographyServiceSetup,
   rootClient?: Client,
   dataSourceId?: string,
+  request?: OpenSearchDashboardsRequest,
+  authRegistry?: IAuthenticationMethodRegistery,
   requireDecryption: boolean = true
 ): Promise<Client> => {
-  const {
+  let credential;
+  let {
     auth: { type },
-    endpoint,
+    name,
   } = dataSourceAttr;
+  const { endpoint } = dataSourceAttr;
+  name = name ?? type;
   const clientOptions = parseClientOptions(config, endpoint, registeredSchema);
   const cacheKey = generateCacheKey(dataSourceAttr, dataSourceId);
 
+  const authenticationMethod = authRegistry?.getAuthenticationMethod(name);
+  if (authenticationMethod !== undefined) {
+    const credentialProvider = await authRegistryCredentialProvider(authenticationMethod, {
+      dataSourceAttr,
+      request,
+      cryptography,
+    });
+    credential = credentialProvider.credential;
+    type = credentialProvider.type;
+  }
+
   switch (type) {
     case AuthType.NoAuth:
       if (!rootClient) rootClient = new Client(clientOptions);
@@ -129,21 +153,25 @@ const getQueryClient = async (
       return rootClient.child();
 
     case AuthType.UsernamePasswordType:
-      const credential = requireDecryption
-        ? await getCredential(dataSourceAttr, cryptography!)
-        : (dataSourceAttr.auth.credentials as UsernamePasswordTypedContent);
+      credential =
+        (credential as UsernamePasswordTypedContent) ??
+        (requireDecryption
+          ? await getCredential(dataSourceAttr, cryptography!)
+          : (dataSourceAttr.auth.credentials as UsernamePasswordTypedContent));
 
       if (!rootClient) rootClient = new Client(clientOptions);
       addClientToPool(cacheKey, type, rootClient);
 
       return getBasicAuthClient(rootClient, credential);
 
     case AuthType.SigV4:
-      const awsCredential = requireDecryption
-        ? await getAWSCredential(dataSourceAttr, cryptography!)
-        : (dataSourceAttr.auth.credentials as SigV4Content);
+      credential =
+        (credential as SigV4Content) ??
+        (requireDecryption
+          ? await getAWSCredential(dataSourceAttr, cryptography!)
+          : (dataSourceAttr.auth.credentials as SigV4Content));
 
-      const awsClient = rootClient ? rootClient : getAWSClient(awsCredential, clientOptions);
+      const awsClient = rootClient ? rootClient : getAWSClient(credential, clientOptions);
       addClientToPool(cacheKey, type, awsClient);
 
       return awsClient;

src/plugins/data_source/server/legacy/configure_legacy_client.ts

@@ -14,6 +14,7 @@ import {
   LegacyCallAPIOptions,
   LegacyOpenSearchErrorHelpers,
   Logger,
+  OpenSearchDashboardsRequest,
 } from '../../../../../src/core/server';
 import {
   AuthType,
@@ -34,13 +35,17 @@ import {
   getDataSource,
   generateCacheKey,
 } from '../client/configure_client_utils';
+import { IAuthenticationMethodRegistery } from '../auth_registry';
+import { authRegistryCredentialProvider } from '../util/credential_provider';
 
 export const configureLegacyClient = async (
   {
     dataSourceId,
     savedObjects,
     cryptography,
     customApiSchemaRegistryPromise,
+    request,
+    authRegistry,
   }: DataSourceClientParams,
   callApiParams: LegacyClientCallAPIParams,
   openSearchClientPoolSetup: OpenSearchClientPoolSetup,
@@ -65,7 +70,9 @@ export const configureLegacyClient = async (
       config,
       registeredSchema,
       rootClient,
-      dataSourceId
+      dataSourceId,
+      request,
+      authRegistry
     );
   } catch (error: any) {
     logger.debug(
@@ -96,15 +103,31 @@ const getQueryClient = async (
   config: DataSourcePluginConfigType,
   registeredSchema: any[],
   rootClient?: LegacyClient,
-  dataSourceId?: string
+  dataSourceId?: string,
+  request?: OpenSearchDashboardsRequest,
+  authRegistry?: IAuthenticationMethodRegistery
 ) => {
-  const {
+  let credential;
+  let {
     auth: { type },
-    endpoint: nodeUrl,
+    name,
   } = dataSourceAttr;
+  const { endpoint: nodeUrl } = dataSourceAttr;
+  name = name ?? type;
   const clientOptions = parseClientOptions(config, nodeUrl, registeredSchema);
   const cacheKey = generateCacheKey(dataSourceAttr, dataSourceId);
 
+  const authenticationMethod = authRegistry?.getAuthenticationMethod(name);
+  if (authenticationMethod !== undefined) {
+    const credentialProvider = await authRegistryCredentialProvider(authenticationMethod, {
+      dataSourceAttr,
+      request,
+      cryptography,
+    });
+    credential = credentialProvider.credential;
+    type = credentialProvider.type;
+  }
+
   switch (type) {
     case AuthType.NoAuth:
       if (!rootClient) rootClient = new LegacyClient(clientOptions);
@@ -117,17 +140,20 @@ const getQueryClient = async (
       );
 
     case AuthType.UsernamePasswordType:
-      const credential = await getCredential(dataSourceAttr, cryptography);
+      credential =
+        (credential as UsernamePasswordTypedContent) ??
+        (await getCredential(dataSourceAttr, cryptography));
 
       if (!rootClient) rootClient = new LegacyClient(clientOptions);
       addClientToPool(cacheKey, type, rootClient);
 
       return getBasicAuthClient(rootClient, { endpoint, clientParams, options }, credential);
 
     case AuthType.SigV4:
-      const awsCredential = await getAWSCredential(dataSourceAttr, cryptography);
+      credential =
+        (credential as SigV4Content) ?? (await getAWSCredential(dataSourceAttr, cryptography));
 
-      const awsClient = rootClient ? rootClient : getAWSClient(awsCredential, clientOptions);
+      const awsClient = rootClient ? rootClient : getAWSClient(credential, clientOptions);
       addClientToPool(cacheKey, type, awsClient);
 
       return await (callAPI.bind(null, awsClient) as LegacyAPICaller)(

src/plugins/data_source/server/plugin.ts

@@ -168,7 +168,8 @@ export class DataSourcePlugin implements Plugin<DataSourcePluginSetup, DataSourc
     authRegistryPromise: Promise<IAuthenticationMethodRegistery>,
     customApiSchemaRegistryPromise: Promise<CustomApiSchemaRegistry>
   ): IContextProvider<RequestHandler<unknown, unknown, unknown>, 'dataSource'> => {
-    return (context, req) => {
+    return async (context, req) => {
+      const authRegistry = await authRegistryPromise;
       return {
         opensearch: {
           getClient: (dataSourceId: string) => {
@@ -181,6 +182,8 @@ export class DataSourcePlugin implements Plugin<DataSourcePluginSetup, DataSourc
               savedObjects: context.core.savedObjects.client,
               cryptography,
               customApiSchemaRegistryPromise,
+              request: req,
+              authRegistry,
             });
           },
           legacy: {
@@ -190,6 +193,8 @@ export class DataSourcePlugin implements Plugin<DataSourcePluginSetup, DataSourc
                 savedObjects: context.core.savedObjects.client,
                 cryptography,
                 customApiSchemaRegistryPromise,
+                request: req,
+                authRegistry,
               });
             },
           },

src/plugins/data_source/server/routes/test_connection.ts

@@ -11,12 +11,13 @@ import { DataSourceServiceSetup } from '../data_source_service';
 import { CryptographyServiceSetup } from '../cryptography_service';
 import { IAuthenticationMethodRegistery } from '../auth_registry';
 
-export const registerTestConnectionRoute = (
+export const registerTestConnectionRoute = async (
   router: IRouter,
   dataSourceServiceSetup: DataSourceServiceSetup,
   cryptography: CryptographyServiceSetup,
   authRegistryPromise: Promise<IAuthenticationMethodRegistery>
 ) => {
+  const authRegistry = await authRegistryPromise;
   router.post(
     {
       path: '/internal/data-source-management/validate',
@@ -65,6 +66,8 @@ export const registerTestConnectionRoute = (
             cryptography,
             dataSourceId,
             testClientDataSourceAttr: dataSourceAttr as DataSourceAttributes,
+            request,
+            authRegistry,
           }
         );
 

src/plugins/data_source/server/types.ts

@@ -37,6 +37,10 @@ export interface DataSourceClientParams {
   testClientDataSourceAttr?: DataSourceAttributes;
   // custom API schema registry promise, required for getting registered custom API schema
   customApiSchemaRegistryPromise: Promise<CustomApiSchemaRegistry>;
+  // When client parameters are required to be retrieved from the request header, the caller should provide the request.
+  request?: OpenSearchDashboardsRequest;
+  // To retrieve the credentials provider for the authentication method from the registry in order to return the client.
+  authRegistry?: IAuthenticationMethodRegistery;
 }
 
 export interface DataSourceCredentialsProviderOptions {

src/plugins/data_source/server/util/credential_provider.ts

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { DataSourceCredentialsProviderOptions, AuthenticationMethod } from '../types';
+
+export const authRegistryCredentialProvider = async (
+  authenticationMethod: AuthenticationMethod,
+  options: DataSourceCredentialsProviderOptions
+) => ({
+  credential: await authenticationMethod.credentialProvider(options),
+  type: authenticationMethod.authType,
+});