[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,171 @@
  * 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)).toBeUndefined();
+  });
+
+  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,
+      });
+
+      await credStore.update({
+        credential: {
+          authorization: `ApiKey ${base64ApiKey}`,
+        },
+        reindexOp: {
+          ...reindexOpMock,
+          attributes: {
+            ...reindexOpMock.attributes,
+            status: 1,
+          },
+        },
+        security: securityStartMock,
+      });
+
+      expect(securityStartMock.authc.apiKeys.invalidateAsInternalUser).toHaveBeenCalled();
+    });
+
+    it('falls back to user credentials when error granting API key', async () => {
+      const credStore = credentialStoreFactory(logMock);
+
+      securityStartMock.authc.apiKeys.grantAsInternalUser.mockRejectedValue(
+        new Error('Error granting API key')
+      );
 
-    const credStore = credentialStoreFactory();
-    credStore.set(reindexOp, creds);
+      await credStore.set({
+        request: requestMock,
+        reindexOp: reindexOpMock,
+        security: securityStartMock,
+      });
 
-    reindexOp.attributes.lastCompletedStep = 0;
-    expect(credStore.get(reindexOp)).not.toBeDefined();
+      expect(credStore.get(reindexOpMock)).toEqual({
+        authorization: basicAuthHeader,
+      });
+    });
   });
 });

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

@@ -8,10 +8,73 @@
 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>;
 
+// 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');
+
+// Returns a base64-encoded API key string or undefined
+const getApiKey = async ({
+  request,
+  security,
+  reindexOpId,
+  apiKeysMap,
+}: {
+  request: KibanaRequest;
+  security: SecurityPluginStart;
+  reindexOpId: string;
+  apiKeysMap: Map<string, string>;
+}): Promise<string | undefined> => {
+  try {
+    const apiKeyResult = await security.authc.apiKeys.grantAsInternalUser(request, {
+      name: `ua_reindex_${reindexOpId}`,
+      role_descriptors: {},
+      metadata: {
+        description:
+          'Created by the Upgrade Assistant for a reindex operation; this can be safely deleted after Kibana is upgraded.',
+      },
+    });
+
+    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');
+    }
+  } catch (error) {
+    // There are a few edge cases were granting an API key could fail,
+    // in which case we fall back to using the requestor's credentials in memory
+    return undefined;
+  }
+};
+
+const invalidateApiKey = async ({
+  apiKeyId,
+  security,
+  log,
+}: {
+  apiKeyId: string;
+  security?: SecurityPluginStart;
+  log: Logger;
+}) => {
+  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}`);
+  }
+};
+
 /**
  * An in-memory cache for user credentials to be used for reindexing operations. When looking up
  * credentials, the reindex operation must be in the same state it was in when the credentials
@@ -20,25 +83,82 @@ 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>();
-
-  // 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 apiKeysMap = new Map<string, string>();
+  const log = logger.get('credential_store');
 
   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;
+
+      if (areApiKeysEnabled) {
+        const apiKey = await getApiKey({
+          request,
+          security: security!,
+          reindexOpId: reindexOp.id,
+          apiKeysMap,
+        });
+
+        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...
+      if (reindexOp.attributes.status === ReindexStatus.completed) {
+        // ...and an API key is being used, invalidate it
+        const apiKeyId = apiKeysMap.get(reindexOp.id);
+        if (apiKeyId) {
+          await invalidateApiKey({ apiKeyId, security, log });
+          apiKeysMap.delete(reindexOp.id);
+          return;
+        }
+      }
+
+      // Otherwise, re-associate the credentials
       credMap.set(getHash(reindexOp), credential);
     },