[Workplace Search] Persist OAuth token package during OAuth connect flow (#93210)

* Store session data sent from Enterprise Search server

This modifies the EnterpriseSearchRequestHandler to remove any data in a
response under the _sessionData key and instead persist it on the server
side.

Ultimately, this data will be persisted in the login session, but for
now we'll just store it in a cookie. https://github.com/elastic/kibana/issues/92558

Also uses this functionality to persist Workplace Search's OAuth token
package.

* Only return a modified response body if _sessionData was found

The destructuring I'm doing to remove _sessionData from the response is
breaking routes that currently expect an empty response body. This
change just leaves those response bodies alone.

* Refactor from initial feedback & add tests

* Decrease levity

* Changes from PR feedback

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/enterprise_search/common/constants.ts

@@ -81,3 +81,5 @@ export const JSON_HEADER = {
 };
 
 export const READ_ONLY_MODE_HEADER = 'x-ent-search-read-only-mode';
+
+export const ENTERPRISE_SEARCH_KIBANA_COOKIE = '_enterprise_search';

x-pack/plugins/enterprise_search/server/lib/enterprise_search_request_handler.test.ts

@@ -7,7 +7,11 @@
 
 import { mockConfig, mockLogger } from '../__mocks__';
 
-import { JSON_HEADER, READ_ONLY_MODE_HEADER } from '../../common/constants';
+import {
+  ENTERPRISE_SEARCH_KIBANA_COOKIE,
+  JSON_HEADER,
+  READ_ONLY_MODE_HEADER,
+} from '../../common/constants';
 
 import { EnterpriseSearchRequestHandler } from './enterprise_search_request_handler';
 
@@ -171,6 +175,28 @@ describe('EnterpriseSearchRequestHandler', () => {
           headers: mockExpectedResponseHeaders,
         });
       });
+
+      it('filters out any _sessionData passed back from Enterprise Search', async () => {
+        const jsonWithSessionData = {
+          _sessionData: {
+            secrets: 'no peeking',
+          },
+          regular: 'data',
+        };
+
+        EnterpriseSearchAPI.mockReturn(jsonWithSessionData, { headers: JSON_HEADER });
+
+        const requestHandler = enterpriseSearchRequestHandler.createRequest({ path: '/api/prep' });
+        await makeAPICall(requestHandler);
+
+        expect(responseMock.custom).toHaveBeenCalledWith({
+          statusCode: 200,
+          body: {
+            regular: 'data',
+          },
+          headers: mockExpectedResponseHeaders,
+        });
+      });
     });
   });
 
@@ -378,6 +404,33 @@ describe('EnterpriseSearchRequestHandler', () => {
     });
   });
 
+  describe('setSessionData', () => {
+    it('sets the value of wsOAuthTokenPackage in a cookie', async () => {
+      const tokenPackage = 'some_encrypted_secrets';
+
+      const mockNow = 'Thu, 04 Mar 2021 22:40:32 GMT';
+      const mockInAnHour = 'Thu, 04 Mar 2021 23:40:32 GMT';
+      jest.spyOn(global.Date, 'now').mockImplementationOnce(() => {
+        return new Date(mockNow).valueOf();
+      });
+
+      const sessionDataBody = {
+        _sessionData: { wsOAuthTokenPackage: tokenPackage },
+        regular: 'data',
+      };
+
+      EnterpriseSearchAPI.mockReturn(sessionDataBody, { headers: JSON_HEADER });
+
+      const requestHandler = enterpriseSearchRequestHandler.createRequest({ path: '/' });
+      await makeAPICall(requestHandler);
+
+      expect(enterpriseSearchRequestHandler.headers).toEqual({
+        ['set-cookie']: `${ENTERPRISE_SEARCH_KIBANA_COOKIE}=${tokenPackage}; Path=/; Expires=${mockInAnHour}; SameSite=Lax; HttpOnly`,
+        ...mockExpectedResponseHeaders,
+      });
+    });
+  });
+
   it('isEmptyObj', async () => {
     expect(enterpriseSearchRequestHandler.isEmptyObj({})).toEqual(true);
     expect(enterpriseSearchRequestHandler.isEmptyObj({ empty: false })).toEqual(false);

x-pack/plugins/enterprise_search/server/lib/enterprise_search_request_handler.ts

@@ -16,7 +16,12 @@ import {
   Logger,
 } from 'src/core/server';
 
-import { JSON_HEADER, READ_ONLY_MODE_HEADER } from '../../common/constants';
+import {
+  ENTERPRISE_SEARCH_KIBANA_COOKIE,
+  JSON_HEADER,
+  READ_ONLY_MODE_HEADER,
+} from '../../common/constants';
+
 import { ConfigType } from '../index';
 
 interface ConstructorDependencies {
@@ -113,11 +118,17 @@ export class EnterpriseSearchRequestHandler {
           return this.handleInvalidDataError(response, url, json);
         }
 
+        // Intercept data that is meant for the server side session
+        const { _sessionData, ...responseJson } = json;
+        if (_sessionData) {
+          this.setSessionData(_sessionData);
+        }
+
         // Pass successful responses back to the front-end
         return response.custom({
           statusCode: status,
           headers: this.headers,
-          body: json,
+          body: _sessionData ? responseJson : json,
         });
       } catch (e) {
         // Catch connection/auth errors
@@ -270,6 +281,27 @@ export class EnterpriseSearchRequestHandler {
     this.headers[READ_ONLY_MODE_HEADER] = readOnlyMode as 'true' | 'false';
   }
 
+  /**
+   * Extract Session Data
+   *
+   * In the future, this will set the keys passed back from Enterprise Search
+   * into the Kibana login session.
+   * For now we'll explicity look for the Workplace Search OAuth token package
+   * and stuff it into a cookie so it can be picked up later when we proxy the
+   * OAuth callback.
+   */
+  setSessionData(sessionData: { [key: string]: string }) {
+    if (sessionData.wsOAuthTokenPackage) {
+      const anHourFromNow = new Date(Date.now());
+      anHourFromNow.setHours(anHourFromNow.getHours() + 1);
+
+      const cookiePayload = `${ENTERPRISE_SEARCH_KIBANA_COOKIE}=${sessionData.wsOAuthTokenPackage};`;
+      const cookieRestrictions = `Path=/; Expires=${anHourFromNow.toUTCString()}; SameSite=Lax; HttpOnly`;
+
+      this.headers['set-cookie'] = `${cookiePayload} ${cookieRestrictions}`;
+    }
+  }
+
   /**
    * Misc helpers
    */

x-pack/plugins/enterprise_search/server/lib/get_oauth_token_package_params.test.ts

@@ -0,0 +1,55 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { ENTERPRISE_SEARCH_KIBANA_COOKIE } from '../../common/constants';
+
+import { getOAuthTokenPackageParams } from './get_oauth_token_package_params';
+
+describe('getOAuthTokenPackage', () => {
+  const tokenPackage = 'some_encrypted_secrets';
+  const tokenPackageCookie = `${ENTERPRISE_SEARCH_KIBANA_COOKIE}=${tokenPackage}`;
+  const tokenPackageParams = { token_package: tokenPackage };
+
+  describe('when there are no cookie headers', () => {
+    it('returns an empty parameter set', () => {
+      expect(getOAuthTokenPackageParams(undefined)).toEqual({});
+    });
+  });
+
+  describe('when there is a single cookie header', () => {
+    it('returns an empty parameter set when our cookie is not there', () => {
+      const cookieHeader = '_st_fruit=banana';
+
+      expect(getOAuthTokenPackageParams(cookieHeader)).toEqual({});
+    });
+
+    it('returns the token package when our cookie is the only one', () => {
+      const cookieHeader = `${tokenPackageCookie}`;
+
+      expect(getOAuthTokenPackageParams(cookieHeader)).toEqual(tokenPackageParams);
+    });
+
+    it('returns the token package when there are other cookies in the header', () => {
+      const cookieHeader = `_chocolate=chip; ${tokenPackageCookie}; _oatmeal=raisin`;
+
+      expect(getOAuthTokenPackageParams(cookieHeader)).toEqual(tokenPackageParams);
+    });
+  });
+
+  describe('when there are multiple cookie headers', () => {
+    it('returns an empty parameter set when none of them include our cookie', () => {
+      const cookieHeaders = ['_st_fruit=banana', '_sid=12345'];
+
+      expect(getOAuthTokenPackageParams(cookieHeaders)).toEqual({});
+    });
+
+    it('returns the token package when our cookie is present', () => {
+      const cookieHeaders = ['_st_fruit=banana', `_heat=spicy; ${tokenPackageCookie}`];
+
+      expect(getOAuthTokenPackageParams(cookieHeaders)).toEqual(tokenPackageParams);
+    });
+  });
+});

x-pack/plugins/enterprise_search/server/lib/get_oauth_token_package_params.ts

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ENTERPRISE_SEARCH_KIBANA_COOKIE } from '../../common/constants';
+
+export const getOAuthTokenPackageParams = (rawCookieHeader: string | string[] | undefined) => {
+  // In the future the token package will be stored in the login session. For now it's in a cookie.
+
+  if (!rawCookieHeader) {
+    return {};
+  }
+
+  /**
+   * A request can have multiple cookie headers and each header can hold multiple cookies.
+   * Within a header, cookies are separated by '; '. Here we are splitting out the individual
+   * cookies from the header(s) and looking for the specific one that holds our token package.
+   */
+
+  const cookieHeaders = Array.isArray(rawCookieHeader) ? rawCookieHeader : [rawCookieHeader];
+
+  let tokenPackage: string | undefined;
+
+  cookieHeaders
+    .flatMap((rawHeader) => rawHeader.split('; '))
+    .forEach((rawCookie) => {
+      const [cookieName, cookieValue] = rawCookie.split('=');
+      if (cookieName === ENTERPRISE_SEARCH_KIBANA_COOKIE) tokenPackage = cookieValue;
+    });
+
+  if (tokenPackage) {
+    return { token_package: tokenPackage };
+  } else {
+    return {};
+  }
+};

x-pack/plugins/enterprise_search/server/routes/workplace_search/sources.test.ts

@@ -7,6 +7,8 @@
 
 import { MockRouter, mockRequestHandler, mockDependencies } from '../../__mocks__';
 
+import { ENTERPRISE_SEARCH_KIBANA_COOKIE } from '../../../common/constants';
+
 import {
   registerAccountSourcesRoute,
   registerAccountSourcesStatusRoute,
@@ -1249,6 +1251,15 @@ describe('sources routes', () => {
   });
 
   describe('GET /api/workplace_search/sources/create', () => {
+    const tokenPackage = 'some_encrypted_secrets';
+
+    const mockRequest = {
+      headers: {
+        authorization: 'BASIC 123',
+        cookie: `${ENTERPRISE_SEARCH_KIBANA_COOKIE}=${tokenPackage}`,
+      },
+    };
+
     let mockRouter: MockRouter;
 
     beforeEach(() => {
@@ -1265,8 +1276,11 @@ describe('sources routes', () => {
     });
 
     it('creates a request handler', () => {
+      mockRouter.callRoute(mockRequest as any);
+
       expect(mockRequestHandler.createRequest).toHaveBeenCalledWith({
         path: '/ws/sources/create',
+        params: { token_package: tokenPackage },
       });
     });
   });

x-pack/plugins/enterprise_search/server/routes/workplace_search/sources.ts

@@ -7,6 +7,8 @@
 
 import { schema } from '@kbn/config-schema';
 
+import { getOAuthTokenPackageParams } from '../../lib/get_oauth_token_package_params';
+
 import { RouteDependencies } from '../../plugin';
 
 const schemaValuesSchema = schema.recordOf(
@@ -862,9 +864,12 @@ export function registerOauthConnectorParamsRoute({
         }),
       },
     },
-    enterpriseSearchRequestHandler.createRequest({
-      path: '/ws/sources/create',
-    })
+    async (context, request, response) => {
+      return enterpriseSearchRequestHandler.createRequest({
+        path: '/ws/sources/create',
+        params: getOAuthTokenPackageParams(request.headers.cookie),
+      })(context, request, response);
+    }
   );
 }