From a4a84e78bc5da7b272f15795b72d2357cae56394 Mon Sep 17 00:00:00 2001
From: Jochen Kressin <jochen.kressin-gh@eliatra.com>
Date: Mon, 11 Sep 2023 15:51:09 +0200
Subject: [PATCH] The cookie splitter should be able to use cookie values that
 have been set within the current request

Signed-off-by: Jochen Kressin <jochen.kressin-gh@eliatra.com>
---
 server/session/cookie_splitter.test.ts | 65 ++++++++++++++++++++++++++
 server/session/cookie_splitter.ts      | 24 +++++++++-
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/server/session/cookie_splitter.test.ts b/server/session/cookie_splitter.test.ts
index 833987427..73745d59a 100644
--- a/server/session/cookie_splitter.test.ts
+++ b/server/session/cookie_splitter.test.ts
@@ -14,6 +14,7 @@
  */
 import { Request as HapiRequest, ResponseObject as HapiResponseObject } from '@hapi/hapi';
 import { httpServerMock } from '../../../../src/core/server/http/http_server.mocks';
+import { merge } from 'lodash';
 import {
   clearSplitCookies,
   getExtraAuthStorageValue,
@@ -171,4 +172,68 @@ describe('Test extra auth storage', () => {
 
     expect(unsplitValue).toEqual('abcdefghi');
   });
+
+  test('should check for cookie values updated in the same request', async () => {
+    const cookiePrefix = 'testcookie';
+    const additionalCookies = 5;
+
+    const mockRequest = httpServerMock.createRawRequest();
+
+    const extendedMockRequest = merge(mockRequest, {
+      _states: {
+        [cookiePrefix + '1']: {
+          name: cookiePrefix + '1',
+          value: 'abc',
+        },
+        [cookiePrefix + '2']: {
+          name: cookiePrefix + '2',
+          value: 'def',
+        },
+        [cookiePrefix + '3']: {
+          name: cookiePrefix + '3',
+          value: 'ghi',
+        },
+      },
+    }) as HapiRequest;
+
+    const osRequest = OpenSearchDashboardsRequest.from(extendedMockRequest);
+    const unsplitValue = unsplitCookiesIntoValue(osRequest, cookiePrefix, additionalCookies);
+
+    expect(unsplitValue).toEqual('abcdefghi');
+  });
+
+  test('should not mix cookie values updated in the same request with previous cookie values', async () => {
+    const cookiePrefix = 'testcookie';
+    const additionalCookies = 5;
+
+    const mockRequest = httpServerMock.createRawRequest({
+      state: {
+        [cookiePrefix + '1']: 'abc',
+        [cookiePrefix + '2']: 'def',
+        [cookiePrefix + '3']: 'ghi',
+      },
+    });
+
+    const extendedMockRequest = merge(mockRequest, {
+      _states: {
+        [cookiePrefix + '1']: {
+          name: cookiePrefix + '1',
+          value: 'jkl',
+        },
+        [cookiePrefix + '2']: {
+          name: cookiePrefix + '2',
+          value: 'mno',
+        },
+        [cookiePrefix + '3']: {
+          name: cookiePrefix + '3',
+          value: 'pqr',
+        },
+      },
+    }) as HapiRequest;
+
+    const osRequest = OpenSearchDashboardsRequest.from(extendedMockRequest);
+    const unsplitValue = unsplitCookiesIntoValue(osRequest, cookiePrefix, additionalCookies);
+
+    expect(unsplitValue).toEqual('jklmnopqr');
+  });
 });
diff --git a/server/session/cookie_splitter.ts b/server/session/cookie_splitter.ts
index c6b563902..33b3ca12d 100644
--- a/server/session/cookie_splitter.ts
+++ b/server/session/cookie_splitter.ts
@@ -29,6 +29,15 @@ export interface ExtraAuthStorageOptions {
 
 type CookieAuthWithResponseObject = HapiRequest['cookieAuth'] & { h: HapiResponseObject };
 
+interface HapiStates {
+  [cookieName: string]: {
+    name: string;
+    value: string;
+  };
+}
+
+export type HapiRequestWithStates = HapiRequest & { _states: HapiStates };
+
 export function getExtraAuthStorageValue(
   request: OpenSearchDashboardsRequest,
   options: ExtraAuthStorageOptions
@@ -134,12 +143,23 @@ export function unsplitCookiesIntoValue(
   cookiePrefix: string,
   additionalCookies: number
 ): string {
-  const rawRequest: HapiRequest = ensureRawRequest(request);
+  const rawRequest: HapiRequestWithStates = ensureRawRequest(request) as HapiRequestWithStates;
   let fullCookieValue = '';
 
+  // We don't want to mix and match between _states and .state.
+  // If we find the first additional cookie in _states, we
+  // use _states for all subsequent additional cookies
+  const requestHasNewerCookieState = rawRequest._states && rawRequest._states[cookiePrefix + 1];
+
   for (let i = 1; i <= additionalCookies; i++) {
     const cookieName = cookiePrefix + i;
-    if (rawRequest.state[cookieName]) {
+    if (
+      requestHasNewerCookieState &&
+      rawRequest._states[cookieName] &&
+      rawRequest._states[cookieName].value
+    ) {
+      fullCookieValue = fullCookieValue + rawRequest._states[cookieName].value;
+    } else if (!requestHasNewerCookieState && rawRequest.state[cookieName]) {
       fullCookieValue = fullCookieValue + rawRequest.state[cookieName];
     }
   }