[Upgrade Assistant] Add support for API keys when reindexing

x-pack/plugins/upgrade_assistant/kibana.json

@@ -9,6 +9,6 @@
   },
   "configPath": ["xpack", "upgrade_assistant"],
   "requiredPlugins": ["management", "data", "licensing", "features", "infra", "share"],
-  "optionalPlugins": ["usageCollection", "cloud"],
+  "optionalPlugins": ["usageCollection", "cloud", "security"],
   "requiredBundles": ["esUiShared", "kibanaReact", "kibanaUtils"]
 }

x-pack/plugins/upgrade_assistant/server/lib/reindexing/credential_store.test.ts

@@ -5,33 +5,153 @@
  * 2.0.
  */
 
+import { KibanaRequest } from 'src/core/server';
+import { loggingSystemMock, httpServerMock } from 'src/core/server/mocks';
+import { securityMock } from '../../../../security/server/mocks';
 import { ReindexSavedObject } from '../../../common/types';
-import { Credential, credentialStoreFactory } from './credential_store';
+import { credentialStoreFactory } from './credential_store';
+
+const basicAuthHeader = 'Basic abc';
+
+const logMock = loggingSystemMock.create().get();
+const requestMock = KibanaRequest.from(
+  httpServerMock.createRawRequest({
+    headers: {
+      authorization: basicAuthHeader,
+    },
+  })
+);
+const securityStartMock = securityMock.createStart();
+
+const reindexOpMock = {
+  id: 'asdf',
+  attributes: { indexName: 'test', lastCompletedStep: 1, locked: null },
+} as ReindexSavedObject;
 
 describe('credentialStore', () => {
-  it('retrieves the same credentials for the same state', () => {
-    const creds = { key: '1' } as Credential;
-    const reindexOp = {
-      id: 'asdf',
-      attributes: { indexName: 'test', lastCompletedStep: 1, locked: null },
-    } as ReindexSavedObject;
-
-    const credStore = credentialStoreFactory();
-    credStore.set(reindexOp, creds);
-    expect(credStore.get(reindexOp)).toEqual(creds);
+  it('retrieves the same credentials for the same state', async () => {
+    const credStore = credentialStoreFactory(logMock);
+
+    await credStore.set({
+      request: requestMock,
+      reindexOp: reindexOpMock,
+      security: securityStartMock,
+    });
+
+    expect(credStore.get(reindexOpMock)).toEqual({
+      authorization: basicAuthHeader,
+    });
+  });
+
+  it('does not retrieve credentials if the state changed', async () => {
+    const credStore = credentialStoreFactory(logMock);
+
+    await credStore.set({
+      request: requestMock,
+      reindexOp: reindexOpMock,
+      security: securityStartMock,
+    });
+
+    reindexOpMock.attributes.lastCompletedStep = 0;
+
+    expect(credStore.get(reindexOpMock)).not.toBeDefined();
+  });
+
+  it('retrieves credentials after update', async () => {
+    const credStore = credentialStoreFactory(logMock);
+
+    await credStore.set({
+      request: requestMock,
+      reindexOp: reindexOpMock,
+      security: securityStartMock,
+    });
+
+    const updatedReindexOp = {
+      ...reindexOpMock,
+      attributes: {
+        ...reindexOpMock.attributes,
+        status: 0,
+      },
+    };
+
+    await credStore.update({
+      credential: {
+        authorization: basicAuthHeader,
+      },
+      reindexOp: updatedReindexOp,
+      security: securityStartMock,
+    });
+
+    expect(credStore.get(updatedReindexOp)).toEqual({
+      authorization: basicAuthHeader,
+    });
   });
 
-  it('does retrieve credentials if the state is changed', () => {
-    const creds = { key: '1' } as Credential;
-    const reindexOp = {
-      id: 'asdf',
-      attributes: { indexName: 'test', lastCompletedStep: 1, locked: null },
-    } as ReindexSavedObject;
+  describe('API keys enabled', () => {
+    const apiKeyResultMock = {
+      id: 'api_key_id',
+      name: 'api_key_name',
+      api_key: '123',
+    };
+
+    const invalidateApiKeyResultMock = {
+      invalidated_api_keys: [apiKeyResultMock.api_key],
+      previously_invalidated_api_keys: [],
+      error_count: 0,
+    };
+
+    const base64ApiKey = Buffer.from(`${apiKeyResultMock.id}:${apiKeyResultMock.api_key}`).toString(
+      'base64'
+    );
+
+    beforeEach(() => {
+      securityStartMock.authc.apiKeys.areAPIKeysEnabled.mockReturnValue(Promise.resolve(true));
+      securityStartMock.authc.apiKeys.grantAsInternalUser.mockReturnValue(
+        Promise.resolve(apiKeyResultMock)
+      );
+      securityStartMock.authc.apiKeys.invalidateAsInternalUser.mockReturnValue(
+        Promise.resolve(invalidateApiKeyResultMock)
+      );
+    });
+
+    it('sets API key in authorization header', async () => {
+      const credStore = credentialStoreFactory(logMock);
+
+      await credStore.set({
+        request: requestMock,
+        reindexOp: reindexOpMock,
+        security: securityStartMock,
+      });
+
+      expect(credStore.get(reindexOpMock)).toEqual({
+        authorization: `ApiKey ${base64ApiKey}`,
+      });
+    });
+
+    it('invalidates API keys when a reindex operation is complete', async () => {
+      const credStore = credentialStoreFactory(logMock);
+
+      await credStore.set({
+        request: requestMock,
+        reindexOp: reindexOpMock,
+        security: securityStartMock,
+      });
 
-    const credStore = credentialStoreFactory();
-    credStore.set(reindexOp, creds);
+      await credStore.update({
+        credential: {
+          authorization: `ApiKey ${base64ApiKey}`,
+        },
+        reindexOp: {
+          ...reindexOpMock,
+          attributes: {
+            ...reindexOpMock.attributes,
+            status: 1,
+          },
+        },
+        security: securityStartMock,
+      });
 
-    reindexOp.attributes.lastCompletedStep = 0;
-    expect(credStore.get(reindexOp)).not.toBeDefined();
+      expect(securityStartMock.authc.apiKeys.invalidateAsInternalUser).toHaveBeenCalled();
+    });
   });
 });

x-pack/plugins/upgrade_assistant/server/lib/reindexing/credential_store.ts

@@ -8,7 +8,10 @@
 import { createHash } from 'crypto';
 import stringify from 'json-stable-stringify';
 
-import { ReindexSavedObject } from '../../../common/types';
+import { KibanaRequest, Logger } from 'src/core/server';
+
+import { SecurityPluginStart } from '../../../../security/server';
+import { ReindexSavedObject, ReindexStatus } from '../../../common/types';
 
 export type Credential = Record<string, any>;
 
@@ -20,25 +23,119 @@ export type Credential = Record<string, any>;
  */
 export interface CredentialStore {
   get(reindexOp: ReindexSavedObject): Credential | undefined;
-  set(reindexOp: ReindexSavedObject, credential: Credential): void;
+  set(params: {
+    reindexOp: ReindexSavedObject;
+    request: KibanaRequest;
+    security?: SecurityPluginStart;
+  }): Promise<void>;
+  update(params: {
+    reindexOp: ReindexSavedObject;
+    security?: SecurityPluginStart;
+    credential: Credential;
+  }): Promise<void>;
   clear(): void;
 }
 
-export const credentialStoreFactory = (): CredentialStore => {
+export const credentialStoreFactory = (logger: Logger): CredentialStore => {
   const credMap = new Map<string, Credential>();
+  const apiKeysMap = new Map<string, string>();
+  const log = logger.get('credential_store');
 
   // Generates a stable hash for the reindex operation's current state.
   const getHash = (reindexOp: ReindexSavedObject) =>
     createHash('sha256')
       .update(stringify({ id: reindexOp.id, ...reindexOp.attributes }))
       .digest('base64');
 
+  const getApiKey = async ({
+    request,
+    security,
+    reindexOpId,
+  }: {
+    request: KibanaRequest;
+    security?: SecurityPluginStart;
+    reindexOpId: string;
+  }): Promise<string | undefined> => {
+    const apiKeyResult = await security?.authc.apiKeys.grantAsInternalUser(request, {
+      name: `ua_reindex_${reindexOpId}`,
+      role_descriptors: {},
+    });
+
+    if (apiKeyResult) {
+      const { api_key: apiKey, id } = apiKeyResult;
+      // Store each API key per reindex operation so that we can later invalidate it when the reindex operation is complete
+      apiKeysMap.set(reindexOpId, id);
+      // Returns the base64 encoding of `id:api_key`
+      // This can be used when sending a request with an "Authorization: ApiKey xxx" header
+      return Buffer.from(`${id}:${apiKey}`).toString('base64');
+    }
+  };
+
+  const invalidateApiKey = async ({
+    apiKeyId,
+    security,
+  }: {
+    apiKeyId: string;
+    security?: SecurityPluginStart;
+  }) => {
+    try {
+      await security?.authc.apiKeys.invalidateAsInternalUser({ ids: [apiKeyId] });
+    } catch (error) {
+      // Swallow error if there's a problem invalidating API key
+      log.debug(`Error invalidating API key for id ${apiKeyId}: ${error.message}`);
+    }
+  };
+
   return {
     get(reindexOp: ReindexSavedObject) {
       return credMap.get(getHash(reindexOp));
     },
 
-    set(reindexOp: ReindexSavedObject, credential: Credential) {
+    async set({
+      reindexOp,
+      request,
+      security,
+    }: {
+      reindexOp: ReindexSavedObject;
+      request: KibanaRequest;
+      security?: SecurityPluginStart;
+    }) {
+      const areApiKeysEnabled = (await security?.authc.apiKeys.areAPIKeysEnabled()) ?? false;
+
+      const apiKey =
+        areApiKeysEnabled && (await getApiKey({ request, security, reindexOpId: reindexOp.id }));
+
+      if (apiKey) {
+        credMap.set(getHash(reindexOp), {
+          ...request.headers,
+          authorization: `ApiKey ${apiKey}`,
+        });
+        return;
+      }
+
+      // Set the requestor's credentials in memory if apiKeys are not enabled
+      credMap.set(getHash(reindexOp), request.headers);
+    },
+
+    async update({
+      reindexOp,
+      security,
+      credential,
+    }: {
+      reindexOp: ReindexSavedObject;
+      security?: SecurityPluginStart;
+      credential: Credential;
+    }) {
+      // If the reindex operation is completed, and an API key is being used, invalidate it
+      if (reindexOp.attributes.status === ReindexStatus.completed) {
+        const apiKeyId = apiKeysMap.get(reindexOp.id);
+        if (apiKeyId) {
+          await invalidateApiKey({ apiKeyId, security });
+          return;
+        }
+      }
+
+      // Otherwise, re-associate the credentials
       credMap.set(getHash(reindexOp), credential);
     },
 

x-pack/plugins/upgrade_assistant/server/lib/reindexing/worker.ts

@@ -7,6 +7,7 @@
 
 import { IClusterClient, Logger, SavedObjectsClientContract, FakeRequest } from 'src/core/server';
 import moment from 'moment';
+import { SecurityPluginStart } from '../../../../security/server';
 import { ReindexSavedObject, ReindexStatus } from '../../../common/types';
 import { Credential, CredentialStore } from './credential_store';
 import { reindexActionsFactory } from './reindex_actions';
@@ -46,15 +47,19 @@ export class ReindexWorker {
   private inProgressOps: ReindexSavedObject[] = [];
   private readonly reindexService: ReindexService;
   private readonly log: Logger;
+  private readonly security: SecurityPluginStart;
 
   constructor(
     private client: SavedObjectsClientContract,
     private credentialStore: CredentialStore,
     private clusterClient: IClusterClient,
     log: Logger,
-    private licensing: LicensingPluginSetup
+    private licensing: LicensingPluginSetup,
+    security: SecurityPluginStart
   ) {
     this.log = log.get('reindex_worker');
+    this.security = security;
+
     if (ReindexWorker.workerSingleton) {
       throw new Error(`More than one ReindexWorker cannot be created.`);
     }
@@ -171,7 +176,11 @@ export class ReindexWorker {
             firstOpInQueue.attributes.indexName
           );
           // Re-associate the credentials
-          this.credentialStore.set(firstOpInQueue, credential);
+          this.credentialStore.update({
+            reindexOp: firstOpInQueue,
+            security: this.security,
+            credential,
+          });
         }
       }
 
@@ -230,7 +239,7 @@ export class ReindexWorker {
     reindexOp = await swallowExceptions(service.processNextStep, this.log)(reindexOp);
 
     // Update credential store with most recent state.
-    this.credentialStore.set(reindexOp, credential);
+    this.credentialStore.update({ reindexOp, security: this.security, credential });
   };
 }