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. elastic#92558

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

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.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 {
@@ -27,6 +32,7 @@ interface RequestParams {
   path: string;
   params?: object;
   hasValidData?: Function;
+  computeExtraParams?: Function;
 }
 interface ErrorResponse {
   message: string;
@@ -56,7 +62,12 @@ export class EnterpriseSearchRequestHandler {
     this.enterpriseSearchUrl = config.host as string;
   }
 
-  createRequest({ path, params = {}, hasValidData = () => true }: RequestParams) {
+  createRequest({
+    path,
+    params = {},
+    hasValidData = () => true,
+    computeExtraParams = () => {},
+  }: RequestParams) {
     return async (
       _context: RequestHandlerContext,
       request: KibanaRequest<unknown, unknown, unknown>,
@@ -65,7 +76,11 @@ export class EnterpriseSearchRequestHandler {
       try {
         // Set up API URL
         const encodedPath = this.encodePathParams(path, request.params as Record<string, string>);
-        const queryParams = { ...(request.query as object), ...params };
+        const queryParams = {
+          ...(request.query as object),
+          ...params,
+          ...computeExtraParams(request),
+        };
         const queryString = !this.isEmptyObj(queryParams)
           ? `?${querystring.stringify(queryParams)}`
           : '';
@@ -113,11 +128,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: responseJson,
         });
       } catch (e) {
         // Catch connection/auth errors
@@ -270,6 +291,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();
+      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/routes/workplace_search/sources.ts

@@ -5,8 +5,12 @@
  * 2.0.
  */
 
+import { KibanaRequest } from 'kibana/server';
+
 import { schema } from '@kbn/config-schema';
 
+import { ENTERPRISE_SEARCH_KIBANA_COOKIE } from '../../../common/constants';
+
 import { RouteDependencies } from '../../plugin';
 
 const schemaValuesSchema = schema.recordOf(
@@ -864,6 +868,30 @@ export function registerOauthConnectorParamsRoute({
     },
     enterpriseSearchRequestHandler.createRequest({
       path: '/ws/sources/create',
+      computeExtraParams: (request: KibanaRequest) => {
+        // In the future the token package will be stored in the login session. For now it's in a cookie.
+        const cookieHeader = request.headers.cookie;
+
+        if (!cookieHeader) {
+          return {};
+        }
+
+        // Take any cookie headers and split the individual cookies out, e.g. "_my_cookie=chocolateChip"
+        const cookiePayloads = [cookieHeader].flat().flatMap((rawHeader) => rawHeader.split('; '));
+
+        // Split those raw cookies into [key, value] pairs, e.g. ["_my_cookie", "chocolateChip"]
+        const cookiePairs = cookiePayloads.map((rawCookie) => rawCookie.split('='));
+
+        const tokenPackageCookie = cookiePairs.find((cookiePair) => {
+          return cookiePair[0] === ENTERPRISE_SEARCH_KIBANA_COOKIE;
+        });
+
+        if (tokenPackageCookie) {
+          return { token_package: tokenPackageCookie[1] };
+        } else {
+          return {};
+        }
+      },
     })
   );
 }