[Response Ops][Task Manager] Pass fakeRequest for tasks with associated API key instead of actual key (#217327)

Resolves #216808

## Summary

Instead of directly passing the `apiKey` and `userScope` (when
available) into each task type runner, we create a fake Kibana Request
with an Authorization ApiKey header in the task manager and pass that
request into the task type runner.

---------

Co-authored-by: kibanamachine <[email protected]>
Co-authored-by: Elastic Machine <[email protected]>

Author: 3 people

x-pack/examples/triggers_actions_ui_example/server/plugin.ts

@@ -5,12 +5,10 @@
  * 2.0.
  */
 
-import { Plugin, CoreSetup, FakeRawRequest } from '@kbn/core/server';
+import { Plugin, CoreSetup } from '@kbn/core/server';
 
 import { PluginSetupContract as ActionsSetup } from '@kbn/actions-plugin/server';
 import { AlertingServerSetup } from '@kbn/alerting-plugin/server';
-import { addSpaceIdToPath } from '@kbn/spaces-plugin/common';
-import { kibanaRequestFactory } from '@kbn/core-http-server-utils';
 import {
   TaskManagerSetupContract,
   TaskManagerStartContract,
@@ -140,22 +138,8 @@ export class TriggersActionsUiExamplePlugin
     taskManager.registerTaskDefinitions({
       taskWithApiKey: {
         title: 'taskWithApiKey',
-        createTaskRunner: ({ taskInstance }) => ({
+        createTaskRunner: ({ taskInstance, fakeRequest }) => ({
           run: async () => {
-            const services = await core.getStartServices();
-            const fakeRawRequest: FakeRawRequest = {
-              headers: {
-                authorization: `ApiKey ${taskInstance?.apiKey}`,
-              },
-              path: '/',
-            };
-
-            const path = addSpaceIdToPath('/', taskInstance.userScope?.spaceId || 'default');
-
-            // Fake request from the API key
-            const fakeRequest = kibanaRequestFactory(fakeRawRequest);
-            services[0].http.basePath.set(fakeRequest, path);
-
             return {
               state: {},
             };

x-pack/examples/triggers_actions_ui_example/tsconfig.json

@@ -34,8 +34,6 @@
     "@kbn/field-formats-plugin",
     "@kbn/licensing-plugin",
     "@kbn/response-ops-rule-form",
-    "@kbn/spaces-plugin",
-    "@kbn/core-http-server-utils",
     "@kbn/task-manager-plugin",
     "@kbn/fields-metadata-plugin",
   ]

x-pack/platform/plugins/shared/task_manager/server/plugin.ts

@@ -287,7 +287,7 @@ export class TaskManagerPlugin
   }
 
   public start(
-    { savedObjects, elasticsearch, executionContext, security }: CoreStart,
+    { http, savedObjects, elasticsearch, executionContext, security }: CoreStart,
     { cloud, spaces }: TaskManagerPluginsStart
   ): TaskManagerStartContract {
     const savedObjectsRepository = savedObjects.createInternalRepository([
@@ -367,6 +367,7 @@ export class TaskManagerPlugin
       });
 
       this.taskPollingLifecycle = new TaskPollingLifecycle({
+        basePathService: http.basePath,
         config: this.config!,
         definitions: this.definitions,
         logger: this.logger,

x-pack/platform/plugins/shared/task_manager/server/polling_lifecycle.test.ts

@@ -23,7 +23,7 @@ import type {
 import type { Err, Ok } from './lib/result_type';
 import { asOk, isErr, isOk } from './lib/result_type';
 import { FillPoolResult } from './lib/fill_pool';
-import { executionContextServiceMock } from '@kbn/core/server/mocks';
+import { executionContextServiceMock, httpServiceMock } from '@kbn/core/server/mocks';
 import { TaskCost } from './task';
 import { CLAIM_STRATEGY_MGET, DEFAULT_KIBANAS_PER_PARTITION } from './config';
 import { TaskPartitioner } from './lib/task_partitioner';
@@ -105,6 +105,7 @@ describe('TaskPollingLifecycle', () => {
       },
       auto_calculate_default_ech_capacity: false,
     },
+    basePathService: httpServiceMock.createBasePath(),
     taskStore: mockTaskStore,
     logger: taskManagerLogger,
     definitions: new TaskTypeDictionary(taskManagerLogger),

x-pack/platform/plugins/shared/task_manager/server/polling_lifecycle.ts

@@ -12,7 +12,7 @@ import { pipe } from 'fp-ts/lib/pipeable';
 import { map as mapOptional, none } from 'fp-ts/lib/Option';
 import { tap } from 'rxjs';
 import type { UsageCounter } from '@kbn/usage-collection-plugin/server';
-import type { Logger, ExecutionContextStart } from '@kbn/core/server';
+import type { Logger, ExecutionContextStart, IBasePath } from '@kbn/core/server';
 
 import type { Result } from './lib/result_type';
 import { asErr, mapErr, asOk, map, mapOk, isOk } from './lib/result_type';
@@ -69,6 +69,7 @@ export interface ITaskEventEmitter<T> {
 }
 
 export interface TaskPollingLifecycleOpts {
+  basePathService: IBasePath;
   logger: Logger;
   definitions: TaskTypeDictionary;
   taskStore: TaskStore;
@@ -99,6 +100,7 @@ export class TaskPollingLifecycle implements ITaskEventEmitter<TaskLifecycleEven
   private store: TaskStore;
   private taskClaiming: TaskClaiming;
   private bufferedStore: BufferedTaskStore;
+  private readonly basePathService: IBasePath;
   private readonly executionContext: ExecutionContextStart;
 
   private logger: Logger;
@@ -126,6 +128,7 @@ export class TaskPollingLifecycle implements ITaskEventEmitter<TaskLifecycleEven
    * mechanism.
    */
   constructor({
+    basePathService,
     logger,
     middleware,
     config,
@@ -138,6 +141,7 @@ export class TaskPollingLifecycle implements ITaskEventEmitter<TaskLifecycleEven
     taskPartitioner,
     startingCapacity,
   }: TaskPollingLifecycleOpts) {
+    this.basePathService = basePathService;
     this.logger = logger;
     this.middleware = middleware;
     this.definitions = definitions;
@@ -251,6 +255,7 @@ export class TaskPollingLifecycle implements ITaskEventEmitter<TaskLifecycleEven
 
   private createTaskRunnerForTask = (instance: ConcreteTaskInstance) => {
     return new TaskManagerRunner({
+      basePathService: this.basePathService,
       logger: this.logger,
       instance,
       store: this.bufferedStore,

x-pack/platform/plugins/shared/task_manager/server/task.ts

@@ -55,6 +55,12 @@ export interface RunContext {
    * The document describing the task instance, its params, state, id, etc.
    */
   taskInstance: ConcreteTaskInstance;
+
+  /**
+   * If an API key is associated with the task, a fake KibanaRequest object
+   * is generated using the API key and passed as part of the run context.
+   */
+  fakeRequest?: KibanaRequest;
 }
 
 /**

x-pack/platform/plugins/shared/task_manager/server/task_running/task_runner.test.ts

@@ -32,7 +32,7 @@ import { TaskTypeDictionary } from '../task_type_dictionary';
 import { mockLogger } from '../test_utils';
 import { throwRetryableError, throwUnrecoverableError } from './errors';
 import apm from 'elastic-apm-node';
-import { executionContextServiceMock } from '@kbn/core/server/mocks';
+import { executionContextServiceMock, httpServiceMock } from '@kbn/core/server/mocks';
 import { usageCountersServiceMock } from '@kbn/usage-collection-plugin/server/usage_counters/usage_counters_service.mock';
 import { bufferedTaskStoreMock } from '../buffered_task_store.mock';
 import {
@@ -984,6 +984,31 @@ describe('TaskManagerRunner', () => {
         expect.any(Function)
       );
     });
+    test('provides fakeRequest when task has apiKey', async () => {
+      const createTaskRunnerFn = jest.fn();
+      const instance = mockInstance();
+      const { runner } = await readyToRunStageSetup({
+        instance: {
+          ...instance,
+          apiKey: 'aw4badfg333',
+          userScope: {
+            apiKeyId: 'abcdefg',
+            spaceId: 'default',
+            apiKeyCreatedByUser: false,
+          },
+        },
+        definitions: {
+          bar: {
+            title: 'Bar!',
+            createTaskRunner: createTaskRunnerFn,
+          },
+        },
+      });
+      await runner.run();
+      const createTaskRunnerParams = createTaskRunnerFn.mock.calls[0][0];
+      expect(createTaskRunnerParams.fakeRequest).toBeDefined();
+      expect(createTaskRunnerParams.taskInstance).toEqual(instance);
+    });
     test('queues a reattempt if the task fails', async () => {
       const initialAttempts = _.random(0, 2);
       const id = Date.now().toString();
@@ -2565,6 +2590,7 @@ describe('TaskManagerRunner', () => {
     }
 
     const runner = new TaskManagerRunner({
+      basePathService: httpServiceMock.createBasePath(),
       defaultMaxAttempts: 5,
       beforeRun: (context) => Promise.resolve(context),
       beforeMarkRunning: (context) => Promise.resolve(context),

x-pack/platform/plugins/shared/task_manager/server/task_running/task_runner.ts

@@ -15,9 +15,18 @@ import apm from 'elastic-apm-node';
 import { v4 as uuidv4 } from 'uuid';
 import { withSpan } from '@kbn/apm-utils';
 import { flow, identity, omit } from 'lodash';
-import type { ExecutionContextStart, Logger } from '@kbn/core/server';
+import type {
+  ExecutionContextStart,
+  FakeRawRequest,
+  Headers,
+  IBasePath,
+  KibanaRequest,
+  Logger,
+} from '@kbn/core/server';
 import { SavedObjectsErrorHelpers } from '@kbn/core/server';
 import type { UsageCounter } from '@kbn/usage-collection-plugin/server';
+import { addSpaceIdToPath } from '@kbn/spaces-plugin/server';
+import { kibanaRequestFactory } from '@kbn/core-http-server-utils';
 import type { Middleware } from '../lib/middleware';
 import type { Result } from '../lib/result_type';
 import {
@@ -103,6 +112,7 @@ export interface Updatable {
 }
 
 type Opts = {
+  basePathService: IBasePath;
   logger: Logger;
   definitions: TaskTypeDictionary;
   instance: ConcreteTaskInstance;
@@ -162,6 +172,7 @@ export class TaskManagerRunner implements TaskRunner {
   private onTaskEvent: (event: TaskRun | TaskMarkRunning | TaskManagerStat) => void;
   private defaultMaxAttempts: number;
   private uuid: string;
+  private readonly basePathService: IBasePath;
   private readonly executionContext: ExecutionContextStart;
   private usageCounter?: UsageCounter;
   private config: TaskManagerConfig;
@@ -180,6 +191,7 @@ export class TaskManagerRunner implements TaskRunner {
    * @memberof TaskManagerRunner
    */
   constructor({
+    basePathService,
     instance,
     definitions,
     logger,
@@ -195,6 +207,7 @@ export class TaskManagerRunner implements TaskRunner {
     strategy,
     getPollInterval,
   }: Opts) {
+    this.basePathService = basePathService;
     this.instance = asPending(sanitizeInstance(instance));
     this.definitions = definitions;
     this.logger = logger;
@@ -374,7 +387,12 @@ export class TaskManagerRunner implements TaskRunner {
     );
 
     try {
-      this.task = definition.createTaskRunner(modifiedContext);
+      const sanitizedTaskInstance = omit(modifiedContext.taskInstance, ['apiKey', 'userScope']);
+      const fakeRequest = this.getFakeKibanaRequest(
+        modifiedContext.taskInstance.apiKey,
+        modifiedContext.taskInstance.userScope?.spaceId
+      );
+      this.task = definition.createTaskRunner({ taskInstance: sanitizedTaskInstance, fakeRequest });
 
       const ctx = {
         type: 'task manager',
@@ -825,6 +843,25 @@ export class TaskManagerRunner implements TaskRunner {
   private getMaxAttempts() {
     return this.definition?.maxAttempts ?? this.defaultMaxAttempts;
   }
+
+  private getFakeKibanaRequest(apiKey?: string, spaceId?: string): KibanaRequest | undefined {
+    if (apiKey) {
+      const requestHeaders: Headers = {};
+
+      requestHeaders.authorization = `ApiKey ${apiKey}`;
+      const path = addSpaceIdToPath('/', spaceId || 'default');
+
+      const fakeRawRequest: FakeRawRequest = {
+        headers: requestHeaders,
+        path: '/',
+      };
+
+      const fakeRequest = kibanaRequestFactory(fakeRawRequest);
+      this.basePathService.set(fakeRequest, path);
+
+      return fakeRequest;
+    }
+  }
 }
 
 function sanitizeInstance(instance: ConcreteTaskInstance): ConcreteTaskInstance {

x-pack/platform/plugins/shared/task_manager/tsconfig.json

@@ -34,6 +34,7 @@
     "@kbn/spaces-plugin",
     "@kbn/encrypted-saved-objects-shared",
     "@kbn/core-saved-objects-api-server-mocks",
+    "@kbn/core-http-server-utils",
   ],
   "exclude": ["target/**/*"]
 }