Cookie compression and splitting for JWT (#1651)

jochen-kressin · github-actions[bot] · commit f265cdea26ef · 2024-01-04T21:17:50.000Z
Signed-off-by: Jochen Kressin <jochen.kressin-gh@eliatra.com> Co-authored-by: Craig Perkins <cwperx@amazon.com> (cherry picked from commit 7cad47c)
diff --git a/server/auth/types/jwt/jwt_auth.ts b/server/auth/types/jwt/jwt_auth.ts
@@ -25,10 +25,21 @@ import {
   AuthToolkit,
   IOpenSearchDashboardsResponse,
 } from 'opensearch-dashboards/server';
+import { ServerStateCookieOptions } from '@hapi/hapi';
 import { SecurityPluginConfigType } from '../../..';
 import { SecuritySessionCookie } from '../../../session/security_cookie';
 import { AuthenticationType } from '../authentication_type';
 import { JwtAuthRoutes } from './routes';
+import {
+  ExtraAuthStorageOptions,
+  getExtraAuthStorageValue,
+  setExtraAuthStorage,
+} from '../../../session/cookie_splitter';
+
+export const JWT_DEFAULT_EXTRA_STORAGE_OPTIONS: ExtraAuthStorageOptions = {
+  cookiePrefix: 'security_authentication_jwt',
+  additionalCookies: 5,
+};
 
 export class JwtAuthentication extends AuthenticationType {
   public readonly type: string = 'jwt';
@@ -48,10 +59,47 @@ export class JwtAuthentication extends AuthenticationType {
   }
 
   public async init() {
-    const routes = new JwtAuthRoutes(this.router, this.sessionStorageFactory);
+    this.createExtraStorage();
+    const routes = new JwtAuthRoutes(this.router, this.sessionStorageFactory, this.config);
     routes.setupRoutes();
   }
 
+  createExtraStorage() {
+    // @ts-ignore
+    const hapiServer: Server = this.sessionStorageFactory.asScoped({}).server;
+
+    const { cookiePrefix, additionalCookies } = this.getExtraAuthStorageOptions();
+    const extraCookieSettings: ServerStateCookieOptions = {
+      isSecure: this.config.cookie.secure,
+      isSameSite: this.config.cookie.isSameSite,
+      password: this.config.cookie.password,
+      domain: this.config.cookie.domain,
+      path: this.coreSetup.http.basePath.serverBasePath || '/',
+      clearInvalid: false,
+      isHttpOnly: true,
+      ignoreErrors: true,
+      encoding: 'iron', // Same as hapi auth cookie
+    };
+
+    for (let i = 1; i <= additionalCookies; i++) {
+      hapiServer.states.add(cookiePrefix + i, extraCookieSettings);
+    }
+  }
+
+  private getExtraAuthStorageOptions(): ExtraAuthStorageOptions {
+    const extraAuthStorageOptions: ExtraAuthStorageOptions = {
+      cookiePrefix:
+        this.config.jwt?.extra_storage.cookie_prefix ||
+        JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+      additionalCookies:
+        this.config.jwt?.extra_storage.additional_cookies ||
+        JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+      logger: this.logger,
+    };
+
+    return extraAuthStorageOptions;
+  }
+
   private getTokenFromUrlParam(request: OpenSearchDashboardsRequest): string | undefined {
     const urlParamName = this.config.jwt?.url_param;
     if (urlParamName) {
@@ -77,6 +125,7 @@ export class JwtAuthentication extends AuthenticationType {
     if (request.headers[this.authHeaderName]) {
       return true;
     }
+
     const urlParamName = this.config.jwt?.url_param;
     if (urlParamName && request.url.searchParams.get(urlParamName)) {
       return true;
@@ -100,22 +149,29 @@ export class JwtAuthentication extends AuthenticationType {
     request: OpenSearchDashboardsRequest<unknown, unknown, unknown, any>,
     authInfo: any
   ): SecuritySessionCookie {
+    setExtraAuthStorage(
+      request,
+      this.getBearerToken(request) || '',
+      this.getExtraAuthStorageOptions()
+    );
     return {
       username: authInfo.user_name,
       credentials: {
-        authHeaderValue: this.getBearerToken(request),
+        authHeaderValueExtra: true,
       },
       authType: this.type,
       expiryTime: Date.now() + this.config.session.ttl,
     };
   }
 
-  async isValidCookie(cookie: SecuritySessionCookie): Promise<boolean> {
+  async isValidCookie(
+    cookie: SecuritySessionCookie,
+    request: OpenSearchDashboardsRequest
+  ): Promise<boolean> {
+    const hasAuthHeaderValue =
+      cookie.credentials?.authHeaderValue || this.getExtraAuthStorageValue(request, cookie);
     return (
-      cookie.authType === this.type &&
-      cookie.username &&
-      cookie.expiryTime &&
-      cookie.credentials?.authHeaderValue
+      cookie.authType === this.type && cookie.username && cookie.expiryTime && hasAuthHeaderValue
     );
   }
 
@@ -127,8 +183,35 @@ export class JwtAuthentication extends AuthenticationType {
     return response.unauthorized();
   }
 
-  buildAuthHeaderFromCookie(cookie: SecuritySessionCookie): any {
+  getExtraAuthStorageValue(request: OpenSearchDashboardsRequest, cookie: SecuritySessionCookie) {
+    let extraValue = '';
+    if (!cookie.credentials?.authHeaderValueExtra) {
+      return extraValue;
+    }
+
+    try {
+      extraValue = getExtraAuthStorageValue(request, this.getExtraAuthStorageOptions());
+    } catch (error) {
+      this.logger.info(error);
+    }
+
+    return extraValue;
+  }
+
+  buildAuthHeaderFromCookie(
+    cookie: SecuritySessionCookie,
+    request: OpenSearchDashboardsRequest
+  ): any {
     const header: any = {};
+    if (cookie.credentials.authHeaderValueExtra) {
+      try {
+        const extraAuthStorageValue = this.getExtraAuthStorageValue(request, cookie);
+        header.authorization = extraAuthStorageValue;
+        return header;
+      } catch (error) {
+        this.logger.error(error);
+      }
+    }
     const authHeaderValue = cookie.credentials?.authHeaderValue;
     if (authHeaderValue) {
       header[this.authHeaderName] = authHeaderValue;
diff --git a/server/auth/types/jwt/jwt_helper.test.ts b/server/auth/types/jwt/jwt_helper.test.ts
@@ -14,18 +14,58 @@
  */
 
 import { getAuthenticationHandler } from '../../auth_handler_factory';
+import { JWT_DEFAULT_EXTRA_STORAGE_OPTIONS } from './jwt_auth';
+import {
+  CoreSetup,
+  ILegacyClusterClient,
+  IRouter,
+  Logger,
+  OpenSearchDashboardsRequest,
+  SessionStorageFactory,
+} from '../../../../../../src/core/server';
+import { SecuritySessionCookie } from '../../../session/security_cookie';
+import { SecurityPluginConfigType } from '../../../index';
+import { httpServerMock } from '../../../../../../src/core/server/http/http_server.mocks';
+import { deflateValue } from '../../../utils/compression';
 
 describe('test jwt auth library', () => {
-  const router: IRouter = { post: (body) => {} };
-  let core: CoreSetup;
+  const router: Partial<IRouter> = { post: (body) => {} };
+  const core = {
+    http: {
+      basePath: {
+        serverBasePath: '/',
+      },
+    },
+  } as CoreSetup;
   let esClient: ILegacyClusterClient;
-  let sessionStorageFactory: SessionStorageFactory<SecuritySessionCookie>;
+  const sessionStorageFactory: SessionStorageFactory<SecuritySessionCookie> = {
+    asScoped: jest.fn().mockImplementation(() => {
+      return {
+        server: {
+          states: {
+            add: jest.fn(),
+          },
+        },
+      };
+    }),
+  };
   let logger: Logger;
 
+  const cookieConfig: Partial<SecurityPluginConfigType> = {
+    cookie: {
+      secure: false,
+      name: 'test_cookie_name',
+      password: 'secret',
+      ttl: 60 * 60 * 1000,
+      domain: null,
+      isSameSite: false,
+    },
+  };
+
   function getTestJWTAuthenticationHandlerWithConfig(config: SecurityPluginConfigType) {
     return getAuthenticationHandler(
       'jwt',
-      router,
+      router as IRouter,
       config,
       core,
       esClient,
@@ -36,9 +76,14 @@ describe('test jwt auth library', () => {
 
   test('test getTokenFromUrlParam', async () => {
     const config = {
+      ...cookieConfig,
       jwt: {
         header: 'Authorization',
         url_param: 'authorization',
+        extra_storage: {
+          cookie_prefix: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+          additional_cookies: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+        },
       },
     };
     const auth = await getTestJWTAuthenticationHandlerWithConfig(config);
@@ -55,9 +100,14 @@ describe('test jwt auth library', () => {
 
   test('test getTokenFromUrlParam incorrect url_param', async () => {
     const config = {
+      ...cookieConfig,
       jwt: {
         header: 'Authorization',
         url_param: 'urlParamName',
+        extra_storage: {
+          cookie_prefix: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+          additional_cookies: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+        },
       },
     };
     const auth = await getTestJWTAuthenticationHandlerWithConfig(config);
@@ -71,4 +121,83 @@ describe('test jwt auth library', () => {
     const token = auth.getTokenFromUrlParam(request);
     expect(token).toEqual(expectedToken);
   });
+
+  test('make sure that cookies with authHeaderValue instead of split cookies are still valid', async () => {
+    const config = {
+      ...cookieConfig,
+      jwt: {
+        header: 'Authorization',
+        url_param: 'authorization',
+        extra_storage: {
+          cookie_prefix: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+          additional_cookies: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+        },
+      },
+    } as SecurityPluginConfigType;
+
+    const jwtAuthentication = await getTestJWTAuthenticationHandlerWithConfig(config);
+
+    const mockRequest = httpServerMock.createRawRequest();
+    const osRequest = OpenSearchDashboardsRequest.from(mockRequest);
+
+    const cookie: SecuritySessionCookie = {
+      credentials: {
+        authHeaderValue: 'Bearer eyToken',
+      },
+    };
+
+    const expectedHeaders = {
+      authorization: 'Bearer eyToken',
+    };
+
+    const headers = jwtAuthentication.buildAuthHeaderFromCookie(cookie, osRequest);
+
+    expect(headers).toEqual(expectedHeaders);
+  });
+
+  test('get authHeaderValue from split cookies', async () => {
+    const config = {
+      ...cookieConfig,
+      jwt: {
+        header: 'Authorization',
+        url_param: 'authorization',
+        extra_storage: {
+          cookie_prefix: 'testcookie',
+          additional_cookies: 2,
+        },
+      },
+    } as SecurityPluginConfigType;
+
+    const jwtAuthentication = await getTestJWTAuthenticationHandlerWithConfig(config);
+
+    const testString = 'Bearer eyCombinedToken';
+    const testStringBuffer: Buffer = deflateValue(testString);
+    const cookieValue = testStringBuffer.toString('base64');
+    const cookiePrefix = config.jwt!.extra_storage.cookie_prefix;
+    const splitValueAt = Math.ceil(
+      cookieValue.length / config.jwt!.extra_storage.additional_cookies
+    );
+    const mockRequest = httpServerMock.createRawRequest({
+      state: {
+        [cookiePrefix + '1']: cookieValue.substring(0, splitValueAt),
+        [cookiePrefix + '2']: cookieValue.substring(splitValueAt),
+      },
+    });
+
+    const osRequest = OpenSearchDashboardsRequest.from(mockRequest);
+
+    const cookie: SecuritySessionCookie = {
+      credentials: {
+        authHeaderValueExtra: true,
+      },
+    };
+
+    const expectedHeaders = {
+      authorization: testString,
+    };
+
+    const headers = jwtAuthentication.buildAuthHeaderFromCookie(cookie, osRequest);
+
+    expect(headers).toEqual(expectedHeaders);
+  });
 });
diff --git a/server/auth/types/jwt/routes.ts b/server/auth/types/jwt/routes.ts
@@ -13,16 +13,34 @@
  *   permissions and limitations under the License.
  */
 
-import { IRouter, SessionStorageFactory } from 'opensearch-dashboards/server';
+import { IRouter, Logger, SessionStorageFactory } from 'opensearch-dashboards/server';
 import { SecuritySessionCookie } from '../../../session/security_cookie';
 import { API_AUTH_LOGOUT, API_PREFIX } from '../../../../common';
+import { clearSplitCookies, ExtraAuthStorageOptions } from '../../../session/cookie_splitter';
+import { JWT_DEFAULT_EXTRA_STORAGE_OPTIONS } from './jwt_auth';
+import { SecurityPluginConfigType } from '../../../index';
 
 export class JwtAuthRoutes {
   constructor(
     private readonly router: IRouter,
-    private readonly sessionStorageFactory: SessionStorageFactory<SecuritySessionCookie>
+    private readonly sessionStorageFactory: SessionStorageFactory<SecuritySessionCookie>,
+    private readonly config: SecurityPluginConfigType
   ) {}
 
+  private getExtraAuthStorageOptions(logger?: Logger): ExtraAuthStorageOptions {
+    const extraAuthStorageOptions: ExtraAuthStorageOptions = {
+      cookiePrefix:
+        this.config.jwt?.extra_storage.cookie_prefix ||
+        JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+      additionalCookies:
+        this.config.jwt?.extra_storage.additional_cookies ||
+        JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+      logger,
+    };
+
+    return extraAuthStorageOptions;
+  }
+
   public setupRoutes() {
     this.router.post(
       {
@@ -33,6 +51,7 @@ export class JwtAuthRoutes {
         },
       },
       async (context, request, response) => {
+        await clearSplitCookies(request, this.getExtraAuthStorageOptions());
         this.sessionStorageFactory.asScoped(request).clear();
         return response.ok();
       }
diff --git a/server/index.ts b/server/index.ts
@@ -17,6 +17,7 @@ import { schema, TypeOf } from '@osd/config-schema';
 import { PluginInitializerContext, PluginConfigDescriptor } from '../../../src/core/server';
 import { SecurityPlugin } from './plugin';
 import { AuthType } from '../common';
+import { JWT_DEFAULT_EXTRA_STORAGE_OPTIONS } from './auth/types/jwt/jwt_auth';
 
 const validateAuthType = (value: string[]) => {
   const supportedAuthTypes = [
@@ -233,6 +234,16 @@ export const configSchema = schema.object({
       login_endpoint: schema.maybe(schema.string()),
       url_param: schema.string({ defaultValue: 'authorization' }),
       header: schema.string({ defaultValue: 'Authorization' }),
+      extra_storage: schema.object({
+        cookie_prefix: schema.string({
+          defaultValue: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.cookiePrefix,
+          minLength: 2,
+        }),
+        additional_cookies: schema.number({
+          min: 1,
+          defaultValue: JWT_DEFAULT_EXTRA_STORAGE_OPTIONS.additionalCookies,
+        }),
+      }),
     })
   ),
   ui: schema.object({
diff --git a/test/jest_integration/jwt_auth.test.ts b/test/jest_integration/jwt_auth.test.ts
