Limit concurrent access to download API + Replace with LRU cache (#72503)

* Limit concurrent access to  download API

* Replacing cache with LRU Cache

* Configure the LRU cache

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

Author: Alex Kahan

x-pack/plugins/security_solution/common/endpoint/constants.ts

@@ -9,3 +9,5 @@ export const alertsIndexPattern = 'logs-endpoint.alerts-*';
 export const metadataIndexPattern = 'metrics-endpoint.metadata-*';
 export const policyIndexPattern = 'metrics-endpoint.policy-*';
 export const telemetryIndexPattern = 'metrics-endpoint.telemetry-*';
+export const LIMITED_CONCURRENCY_ENDPOINT_ROUTE_TAG = 'endpoint:limited-concurrency';
+export const LIMITED_CONCURRENCY_ENDPOINT_COUNT = 100;

x-pack/plugins/security_solution/server/endpoint/lib/artifacts/cache.test.ts

@@ -4,7 +4,6 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-export * from './cache';
 export * from './common';
 export * from './lists';
 export * from './manifest';

x-pack/plugins/security_solution/server/endpoint/lib/artifacts/cache.ts

@@ -4,6 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 import { deflateSync, inflateSync } from 'zlib';
+import LRU from 'lru-cache';
 import {
   ILegacyClusterClient,
   IRouter,
@@ -22,7 +23,6 @@ import {
   httpServerMock,
   loggingSystemMock,
 } from 'src/core/server/mocks';
-import { ExceptionsCache } from '../../lib/artifacts/cache';
 import { ArtifactConstants } from '../../lib/artifacts';
 import { registerDownloadExceptionListRoute } from './download_exception_list';
 import { EndpointAppContextService } from '../../endpoint_app_context_services';
@@ -97,7 +97,7 @@ describe('test alerts route', () => {
   let routeConfig: RouteConfig<unknown, unknown, unknown, never>;
   let routeHandler: RequestHandler<unknown, unknown, unknown>;
   let endpointAppContextService: EndpointAppContextService;
-  let cache: ExceptionsCache;
+  let cache: LRU<string, Buffer>;
   let ingestSavedObjectClient: jest.Mocked<SavedObjectsClientContract>;
 
   beforeEach(() => {
@@ -108,7 +108,7 @@ describe('test alerts route', () => {
     mockClusterClient.asScoped.mockReturnValue(mockScopedClient);
     routerMock = httpServiceMock.createRouter();
     endpointAppContextService = new EndpointAppContextService();
-    cache = new ExceptionsCache(5);
+    cache = new LRU<string, Buffer>({ max: 10, maxAge: 1000 * 60 * 60 });
     const startContract = createMockEndpointAppContextServiceStartContract();
 
     // The authentication with the Fleet Plugin needs a separate scoped SO Client
@@ -164,7 +164,7 @@ describe('test alerts route', () => {
       path.startsWith('/api/endpoint/artifacts/download')
     )!;
 
-    expect(routeConfig.options).toEqual(undefined);
+    expect(routeConfig.options).toEqual({ tags: ['endpoint:limited-concurrency'] });
 
     await routeHandler(
       ({

x-pack/plugins/security_solution/server/endpoint/lib/artifacts/index.ts

@@ -11,11 +11,13 @@ import {
   IKibanaResponse,
   SavedObject,
 } from 'src/core/server';
+import LRU from 'lru-cache';
 // eslint-disable-next-line @kbn/eslint/no-restricted-paths
 import { authenticateAgentWithAccessToken } from '../../../../../ingest_manager/server/services/agents/authenticate';
 import { validate } from '../../../../common/validate';
+import { LIMITED_CONCURRENCY_ENDPOINT_ROUTE_TAG } from '../../../../common/endpoint/constants';
 import { buildRouteValidation } from '../../../utils/build_validation/route_validation';
-import { ArtifactConstants, ExceptionsCache } from '../../lib/artifacts';
+import { ArtifactConstants } from '../../lib/artifacts';
 import {
   DownloadArtifactRequestParamsSchema,
   downloadArtifactRequestParamsSchema,
@@ -32,7 +34,7 @@ const allowlistBaseRoute: string = '/api/endpoint/artifacts';
 export function registerDownloadExceptionListRoute(
   router: IRouter,
   endpointContext: EndpointAppContext,
-  cache: ExceptionsCache
+  cache: LRU<string, Buffer>
 ) {
   router.get(
     {
@@ -43,6 +45,7 @@ export function registerDownloadExceptionListRoute(
           DownloadArtifactRequestParamsSchema
         >(downloadArtifactRequestParamsSchema),
       },
+      options: { tags: [LIMITED_CONCURRENCY_ENDPOINT_ROUTE_TAG] },
     },
     async (context, req, res) => {
       let scopedSOClient: SavedObjectsClientContract;

x-pack/plugins/security_solution/server/endpoint/routes/artifacts/download_exception_list.test.ts

@@ -0,0 +1,72 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import {
+  CoreSetup,
+  KibanaRequest,
+  LifecycleResponseFactory,
+  OnPreAuthToolkit,
+} from 'kibana/server';
+import {
+  LIMITED_CONCURRENCY_ENDPOINT_ROUTE_TAG,
+  LIMITED_CONCURRENCY_ENDPOINT_COUNT,
+} from '../../../common/endpoint/constants';
+
+class MaxCounter {
+  constructor(private readonly max: number = 1) {}
+  private counter = 0;
+  valueOf() {
+    return this.counter;
+  }
+  increase() {
+    if (this.counter < this.max) {
+      this.counter += 1;
+    }
+  }
+  decrease() {
+    if (this.counter > 0) {
+      this.counter -= 1;
+    }
+  }
+  lessThanMax() {
+    return this.counter < this.max;
+  }
+}
+
+function shouldHandleRequest(request: KibanaRequest) {
+  const tags = request.route.options.tags;
+  return tags.includes(LIMITED_CONCURRENCY_ENDPOINT_ROUTE_TAG);
+}
+
+export function registerLimitedConcurrencyRoutes(core: CoreSetup) {
+  const counter = new MaxCounter(LIMITED_CONCURRENCY_ENDPOINT_COUNT);
+  core.http.registerOnPreAuth(function preAuthHandler(
+    request: KibanaRequest,
+    response: LifecycleResponseFactory,
+    toolkit: OnPreAuthToolkit
+  ) {
+    if (!shouldHandleRequest(request)) {
+      return toolkit.next();
+    }
+
+    if (!counter.lessThanMax()) {
+      return response.customError({
+        body: 'Too Many Requests',
+        statusCode: 429,
+      });
+    }
+
+    counter.increase();
+
+    // requests.events.aborted$ has a bug (but has test which explicitly verifies) where it's fired even when the request completes
+    // https://github.com/elastic/kibana/pull/70495#issuecomment-656288766
+    request.events.aborted$.toPromise().then(() => {
+      counter.decrease();
+    });
+
+    return toolkit.next();
+  });
+}

x-pack/plugins/security_solution/server/endpoint/routes/artifacts/download_exception_list.ts

@@ -9,7 +9,7 @@ import { Logger } from 'src/core/server';
 import { PackageConfigServiceInterface } from '../../../../../../ingest_manager/server';
 import { createPackageConfigServiceMock } from '../../../../../../ingest_manager/server/mocks';
 import { listMock } from '../../../../../../lists/server/mocks';
-import { ExceptionsCache } from '../../../lib/artifacts';
+import LRU from 'lru-cache';
 import { getArtifactClientMock } from '../artifact_client.mock';
 import { getManifestClientMock } from '../manifest_client.mock';
 import { ManifestManager } from './manifest_manager';
@@ -28,11 +28,11 @@ export enum ManifestManagerMockType {
 
 export const getManifestManagerMock = (opts?: {
   mockType?: ManifestManagerMockType;
-  cache?: ExceptionsCache;
+  cache?: LRU<string, Buffer>;
   packageConfigService?: jest.Mocked<PackageConfigServiceInterface>;
   savedObjectsClient?: ReturnType<typeof savedObjectsClientMock.create>;
 }): ManifestManager => {
-  let cache = new ExceptionsCache(5);
+  let cache = new LRU<string, Buffer>({ max: 10, maxAge: 1000 * 60 * 60 });
   if (opts?.cache !== undefined) {
     cache = opts.cache;
   }

x-pack/plugins/security_solution/server/endpoint/routes/limited_concurrency.ts

@@ -7,13 +7,10 @@
 import { inflateSync } from 'zlib';
 import { savedObjectsClientMock } from 'src/core/server/mocks';
 import { createPackageConfigServiceMock } from '../../../../../../ingest_manager/server/mocks';
-import {
-  ArtifactConstants,
-  ManifestConstants,
-  ExceptionsCache,
-  isCompleteArtifact,
-} from '../../../lib/artifacts';
+import { ArtifactConstants, ManifestConstants, isCompleteArtifact } from '../../../lib/artifacts';
+
 import { getManifestManagerMock } from './manifest_manager.mock';
+import LRU from 'lru-cache';
 
 describe('manifest_manager', () => {
   describe('ManifestManager sanity checks', () => {
@@ -41,7 +38,7 @@ describe('manifest_manager', () => {
     });
 
     test('ManifestManager populates cache properly', async () => {
-      const cache = new ExceptionsCache(5);
+      const cache = new LRU<string, Buffer>({ max: 10, maxAge: 1000 * 60 * 60 });
       const manifestManager = getManifestManagerMock({ cache });
       const oldManifest = await manifestManager.getLastComputedManifest(
         ManifestConstants.SCHEMA_VERSION

x-pack/plugins/security_solution/server/endpoint/services/artifacts/manifest_manager/manifest_manager.mock.ts

@@ -5,6 +5,7 @@
  */
 
 import { Logger, SavedObjectsClientContract } from 'src/core/server';
+import LRU from 'lru-cache';
 import { PackageConfigServiceInterface } from '../../../../../../ingest_manager/server';
 import { ExceptionListClient } from '../../../../../../lists/server';
 import { ManifestSchemaVersion } from '../../../../../common/endpoint/schema/common';
@@ -16,7 +17,6 @@ import {
   Manifest,
   buildArtifact,
   getFullEndpointExceptionList,
-  ExceptionsCache,
   ManifestDiff,
   getArtifactId,
 } from '../../../lib/artifacts';
@@ -33,7 +33,7 @@ export interface ManifestManagerContext {
   exceptionListClient: ExceptionListClient;
   packageConfigService: PackageConfigServiceInterface;
   logger: Logger;
-  cache: ExceptionsCache;
+  cache: LRU<string, Buffer>;
 }
 
 export interface ManifestSnapshotOpts {
@@ -51,7 +51,7 @@ export class ManifestManager {
   protected packageConfigService: PackageConfigServiceInterface;
   protected savedObjectsClient: SavedObjectsClientContract;
   protected logger: Logger;
-  protected cache: ExceptionsCache;
+  protected cache: LRU<string, Buffer>;
 
   constructor(context: ManifestManagerContext) {
     this.artifactClient = context.artifactClient;