[Security solution] Elastic Assistant APIs - always use performChecks

x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts

@@ -38,7 +38,7 @@ import {
   EsAnonymizationFieldsSchema,
   UpdateAnonymizationFieldSchema,
 } from '../../ai_assistant_data_clients/anonymization_fields/types';
-import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers';
+import { performChecks } from '../helpers';
 
 export interface BulkOperationError {
   message: string;
@@ -162,22 +162,18 @@ export const bulkActionAnonymizationFieldsRoute = (
         request.events.completed$.subscribe(() => abortController.abort());
         try {
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
-          const license = ctx.licensing.license;
-          if (!hasAIAssistantLicense(license)) {
-            return response.forbidden({
-              body: {
-                message: UPGRADE_LICENSE_MESSAGE,
-              },
-            });
-          }
+          // Perform license and authenticated user checks
+          const checkResponse = performChecks({
+            context: ctx,
+            request,
+            response,
+          });
 
-          const authenticatedUser = ctx.elasticAssistant.getCurrentUser();
-          if (authenticatedUser == null) {
-            return assistantResponse.error({
-              body: `Authenticated user not found`,
-              statusCode: 401,
-            });
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
+          const authenticatedUser = checkResponse.currentUser;
+
           const dataClient =
             await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient();
 
@@ -199,20 +195,20 @@ export const bulkActionAnonymizationFieldsRoute = (
           }
 
           const writer = await dataClient?.getWriter();
-          const changedAt = new Date().toISOString();
+          const createdAt = new Date().toISOString();
           const {
             errors,
             docs_created: docsCreated,
             docs_updated: docsUpdated,
             docs_deleted: docsDeleted,
             // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
           } = await writer!.bulk({
-            documentsToCreate: body.create?.map((f) =>
-              transformToCreateScheme(authenticatedUser, changedAt, f)
+            documentsToCreate: body.create?.map((doc) =>
+              transformToCreateScheme(authenticatedUser, createdAt, doc)
             ),
             documentsToDelete: body.delete?.ids,
-            documentsToUpdate: body.update?.map((f) =>
-              transformToUpdateScheme(authenticatedUser, changedAt, f)
+            documentsToUpdate: body.update?.map((doc) =>
+              transformToUpdateScheme(authenticatedUser, createdAt, doc)
             ),
             getUpdateScript: (document: UpdateAnonymizationFieldSchema) =>
               getUpdateScript({ anonymizationField: document, isPatch: true }),

x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts

@@ -12,6 +12,7 @@ import { requestContextMock } from '../../__mocks__/request_context';
 import { getFindAnonymizationFieldsResultWithSingleHit } from '../../__mocks__/response';
 import { findAnonymizationFieldsRoute } from './find_route';
 import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
+import type { AuthenticatedUser } from '@kbn/core-security-common';
 
 describe('Find user anonymization fields route', () => {
   let server: ReturnType<typeof serverMock.create>;
@@ -25,13 +26,13 @@ describe('Find user anonymization fields route', () => {
     clients.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient.findDocuments.mockResolvedValue(
       Promise.resolve(getFindAnonymizationFieldsResultWithSingleHit())
     );
-    clients.elasticAssistant.getCurrentUser.mockResolvedValue({
+    context.elasticAssistant.getCurrentUser.mockReturnValue({
       username: 'my_username',
       authentication_realm: {
         type: 'my_realm_type',
         name: 'my_realm_name',
       },
-    });
+    } as AuthenticatedUser);
     logger = loggingSystemMock.createLogger();
 
     findAnonymizationFieldsRoute(server.router, logger);

x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts

@@ -22,7 +22,7 @@ import { ElasticAssistantPluginRouter } from '../../types';
 import { buildResponse } from '../utils';
 import { EsAnonymizationFieldsSchema } from '../../ai_assistant_data_clients/anonymization_fields/types';
 import { transformESSearchToAnonymizationFields } from '../../ai_assistant_data_clients/anonymization_fields/helpers';
-import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers';
+import { performChecks } from '../helpers';
 
 export const findAnonymizationFieldsRoute = (
   router: ElasticAssistantPluginRouter,
@@ -55,14 +55,16 @@ export const findAnonymizationFieldsRoute = (
         try {
           const { query } = request;
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
-          const license = ctx.licensing.license;
-          if (!hasAIAssistantLicense(license)) {
-            return response.forbidden({
-              body: {
-                message: UPGRADE_LICENSE_MESSAGE,
-              },
-            });
+          // Perform license and authenticated user checks
+          const checkResponse = performChecks({
+            context: ctx,
+            request,
+            response,
+          });
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
+
           const dataClient =
             await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient();
 

x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts

@@ -32,7 +32,17 @@ const actionsClient = actionsClientMock.create();
 jest.mock('../../lib/build_response', () => ({
   buildResponse: jest.fn().mockImplementation((x) => x),
 }));
-jest.mock('../helpers');
+
+jest.mock('../helpers', () => {
+  const original = jest.requireActual('../helpers');
+
+  return {
+    ...original,
+    appendAssistantMessageToConversation: jest.fn(),
+    createConversationWithUserInput: jest.fn(),
+    langChainExecute: jest.fn(),
+  };
+});
 const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock;
 
 const mockLangChainExecute = langChainExecute as jest.Mock;

x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts

@@ -71,14 +71,12 @@ export const chatCompleteRoute = (
 
           // Perform license and authenticated user checks
           const checkResponse = performChecks({
-            authenticatedUser: true,
             context: ctx,
-            license: true,
             request,
             response,
           });
-          if (checkResponse) {
-            return checkResponse;
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
 
           const conversationsDataClient =

x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts

@@ -48,15 +48,14 @@ export const getEvaluateRoute = (router: IRouter<ElasticAssistantRequestHandlerC
 
         // Perform license, authenticated user and evaluation FF checks
         const checkResponse = performChecks({
-          authenticatedUser: true,
           capability: 'assistantModelEvaluation',
           context: ctx,
-          license: true,
           request,
           response,
         });
-        if (checkResponse) {
-          return checkResponse;
+
+        if (!checkResponse.isSuccess) {
+          return checkResponse.response;
         }
 
         // Fetch datasets from LangSmith // TODO: plumb apiKey so this will work in cloud w/o env vars

x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts

@@ -95,15 +95,13 @@ export const postEvaluateRoute = (
 
         // Perform license, authenticated user and evaluation FF checks
         const checkResponse = performChecks({
-          authenticatedUser: true,
           capability: 'assistantModelEvaluation',
           context: ctx,
-          license: true,
           request,
           response,
         });
-        if (checkResponse) {
-          return checkResponse;
+        if (!checkResponse.isSuccess) {
+          return checkResponse.response;
         }
 
         try {

x-pack/plugins/elastic_assistant/server/routes/helpers.ts

@@ -7,6 +7,7 @@
 
 import {
   AnalyticsServiceSetup,
+  type AuthenticatedUser,
   IKibanaResponse,
   KibanaRequest,
   KibanaResponseFactory,
@@ -561,55 +562,66 @@ export const updateConversationWithUserInput = async ({
 };
 
 interface PerformChecksParams {
-  authenticatedUser?: boolean;
   capability?: AssistantFeatureKey;
   context: AwaitedProperties<
     Pick<ElasticAssistantRequestHandlerContext, 'elasticAssistant' | 'licensing' | 'core'>
   >;
-  license?: boolean;
   request: KibanaRequest;
   response: KibanaResponseFactory;
 }
 
 /**
- * Helper to perform checks for authenticated user, capability, and license. Perform all or one
- * of the checks by providing relevant optional params. Check order is license, authenticated user,
- * then capability.
+ * Helper to perform checks for authenticated user, license, and optionally capability.
+ * Check order is license, authenticated user, then capability.
+ *
+ * Returns either a successful check with an AuthenticatedUser or
+ * an unsuccessful check with an error IKibanaResponse.
  *
- * @param authenticatedUser - Whether to check for an authenticated user
  * @param capability - Specific capability to check if enabled, e.g. `assistantModelEvaluation`
  * @param context - Route context
- * @param license - Whether to check for a valid license
  * @param request - Route KibanaRequest
  * @param response - Route KibanaResponseFactory
+ * @returns PerformChecks
  */
+
+type PerformChecks =
+  | {
+      isSuccess: true;
+      currentUser: AuthenticatedUser;
+    }
+  | {
+      isSuccess: false;
+      response: IKibanaResponse;
+    };
 export const performChecks = ({
-  authenticatedUser,
   capability,
   context,
-  license,
   request,
   response,
-}: PerformChecksParams): IKibanaResponse | undefined => {
+}: PerformChecksParams): PerformChecks => {
   const assistantResponse = buildResponse(response);
 
-  if (license) {
-    if (!hasAIAssistantLicense(context.licensing.license)) {
-      return response.forbidden({
+  if (!hasAIAssistantLicense(context.licensing.license)) {
+    return {
+      isSuccess: false,
+      response: response.forbidden({
         body: {
           message: UPGRADE_LICENSE_MESSAGE,
         },
-      });
-    }
+      }),
+    };
   }
 
-  if (authenticatedUser) {
-    if (context.elasticAssistant.getCurrentUser() == null) {
-      return assistantResponse.error({
+  const currentUser = context.elasticAssistant.getCurrentUser();
+
+  if (currentUser == null) {
+    return {
+      isSuccess: false,
+      response: assistantResponse.error({
         body: `Authenticated user not found`,
         statusCode: 401,
-      });
-    }
+      }),
+    };
   }
 
   if (capability) {
@@ -619,11 +631,17 @@ export const performChecks = ({
     });
     const registeredFeatures = context.elasticAssistant.getRegisteredFeatures(pluginName);
     if (!registeredFeatures[capability]) {
-      return response.notFound();
+      return {
+        isSuccess: false,
+        response: response.notFound(),
+      };
     }
   }
 
-  return undefined;
+  return {
+    isSuccess: true,
+    currentUser,
+  };
 };
 
 /**

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts

@@ -143,15 +143,13 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
 
           // Perform license, authenticated user and FF checks
           const checkResponse = performChecks({
-            authenticatedUser: true,
             capability: 'assistantKnowledgeBaseByDefault',
             context: ctx,
-            license: true,
             request,
             response,
           });
-          if (checkResponse) {
-            return checkResponse;
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
 
           logger.debug(

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts

@@ -47,15 +47,13 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
 
           // Perform license, authenticated user and FF checks
           const checkResponse = performChecks({
-            authenticatedUser: true,
             capability: 'assistantKnowledgeBaseByDefault',
             context: ctx,
-            license: true,
             request,
             response,
           });
-          if (checkResponse) {
-            return checkResponse;
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
 
           // Check mappings and upgrade if necessary -- this route only supports v2 KB, so always `true`

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts

@@ -58,21 +58,19 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout
 
           // Perform license, authenticated user and FF checks
           const checkResponse = performChecks({
-            authenticatedUser: true,
             capability: 'assistantKnowledgeBaseByDefault',
             context: ctx,
-            license: true,
             request,
             response,
           });
-          if (checkResponse) {
-            return checkResponse;
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
 
           const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient({
             v2KnowledgeBaseEnabled: true,
           });
-          const currentUser = ctx.elasticAssistant.getCurrentUser();
+          const currentUser = checkResponse.currentUser;
           const userFilter = getKBUserFilter(currentUser);
           const systemFilter = ` AND (kb_resource:"user" OR type:"index")`;
           const additionalFilter = query.filter ? ` AND ${query.filter}` : '';