Added nodes limit parameter to the API

kfirpeled · kfirpeled · commit b6ddab9a14dc · 2024-11-07T14:56:11.000Z
diff --git a/x-pack/packages/kbn-cloud-security-posture/common/schema/graph/v1.ts b/x-pack/packages/kbn-cloud-security-posture/common/schema/graph/v1.ts
@@ -8,9 +8,10 @@
 import { schema } from '@kbn/config-schema';
 
 export const graphRequestSchema = schema.object({
+  nodesLimit: schema.maybe(schema.number()),
+  showUnknownTarget: schema.maybe(schema.boolean()),
   query: schema.object({
     eventIds: schema.arrayOf(schema.string()),
-    showUnknownTarget: schema.maybe(schema.boolean()),
     // TODO: use zod for range validation instead of config schema
     start: schema.oneOf([schema.number(), schema.string()]),
     end: schema.oneOf([schema.number(), schema.string()]),
diff --git a/x-pack/plugins/cloud_security_posture/server/routes/graph/route.ts b/x-pack/plugins/cloud_security_posture/server/routes/graph/route.ts
@@ -38,13 +38,8 @@ export const defineGraphRoute = (router: CspRouter) =>
         },
       },
       async (context, request, response) => {
-        const {
-          eventIds,
-          start,
-          end,
-          esQuery,
-          showUnknownTarget = false,
-        } = request.body.query as GraphRequest['query'];
+        const { nodesLimit, showUnknownTarget = false } = request.body;
+        const { eventIds, start, end, esQuery } = request.body.query as GraphRequest['query'];
         const cspContext = await context.csp;
         const spaceId = (await cspContext.spaces?.spacesService?.getActiveSpace(request))?.id;
 
@@ -55,13 +50,14 @@ export const defineGraphRoute = (router: CspRouter) =>
               esClient: cspContext.esClient,
             },
             query: {
-              showUnknownTarget,
               eventIds,
               spaceId,
               start,
               end,
               esQuery,
             },
+            showUnknownTarget,
+            nodesLimit,
           });
 
           return response.ok({ body: { nodes, edges } });
diff --git a/x-pack/plugins/cloud_security_posture/server/routes/graph/v1.ts b/x-pack/plugins/cloud_security_posture/server/routes/graph/v1.ts
@@ -43,18 +43,21 @@ interface LabelEdges {
 interface GetGraphParams {
   services: GraphContextServices;
   query: {
-    showUnknownTarget: boolean;
     eventIds: string[];
     spaceId?: string;
     start: string | number;
     end: string | number;
     esQuery?: EsQuery;
   };
+  showUnknownTarget: boolean;
+  nodesLimit?: number;
 }
 
 export const getGraph = async ({
   services: { esClient, logger },
-  query: { eventIds, showUnknownTarget, spaceId = 'default', start, end, esQuery },
+  query: { eventIds, spaceId = 'default', start, end, esQuery },
+  showUnknownTarget,
+  nodesLimit,
 }: GetGraphParams): Promise<{
   nodes: NodeDataModel[];
   edges: EdgeDataModel[];
@@ -72,22 +75,31 @@ export const getGraph = async ({
   });
 
   // Convert results into set of nodes and edges
-  const graphContext = parseRecords(logger, results.records);
+  const graphContext = parseRecords(logger, results.records, nodesLimit);
 
   return { nodes: graphContext.nodes, edges: graphContext.edges };
 };
 
 interface ParseContext {
-  nodesMap: Record<string, NodeDataModel>;
-  edgesMap: Record<string, EdgeDataModel>;
-  edgeLabelsNodes: Record<string, string[]>;
-  labelEdges: Record<string, LabelEdges>;
+  readonly nodesLimit?: number;
+  readonly nodesMap: Record<string, NodeDataModel>;
+  readonly edgesMap: Record<string, EdgeDataModel>;
+  readonly edgeLabelsNodes: Record<string, string[]>;
+  readonly labelEdges: Record<string, LabelEdges>;
+  readonly logger: Logger;
 }
 
-const parseRecords = (logger: Logger, records: GraphEdge[]): GraphContext => {
-  const ctx: ParseContext = { nodesMap: {}, edgeLabelsNodes: {}, edgesMap: {}, labelEdges: {} };
+const parseRecords = (logger: Logger, records: GraphEdge[], nodesLimit?: number): GraphContext => {
+  const ctx: ParseContext = {
+    nodesLimit,
+    logger,
+    nodesMap: {},
+    edgeLabelsNodes: {},
+    edgesMap: {},
+    labelEdges: {},
+  };
 
-  logger.trace(`Parsing records [length: ${records.length}]`);
+  logger.trace(`Parsing records [length: ${records.length}] [nodesLimit: ${nodesLimit ?? 'none'}]`);
 
   createNodes(records, ctx);
   createEdgesAndGroups(ctx);
@@ -205,6 +217,15 @@ const createNodes = (records: GraphEdge[], context: Omit<ParseContext, 'edgesMap
   const { nodesMap, edgeLabelsNodes, labelEdges } = context;
 
   for (const record of records) {
+    if (context.nodesLimit !== undefined && Object.keys(nodesMap).length >= context.nodesLimit) {
+      context.logger.debug(
+        `Reached nodes limit [limit: ${context.nodesLimit}] [current: ${
+          Object.keys(nodesMap).length
+        }]`
+      );
+      break;
+    }
+
     const { ips, hosts, users, actorIds, action, targetIds, isAlert, eventOutcome } = record;
     const actorIdsArray = castArray(actorIds);
     const targetIdsArray = castArray(targetIds);
diff --git a/x-pack/test/cloud_security_posture_api/routes/graph.ts b/x-pack/test/cloud_security_posture_api/routes/graph.ts
@@ -416,11 +416,11 @@ export default function (providerContext: FtrProviderContext) {
 
       it('Should return unknown targets', async () => {
         const response = await postGraph(supertest, {
+          showUnknownTarget: true,
           query: {
             eventIds: [],
             start: '2024-09-01T00:00:00Z',
             end: '2024-09-02T00:00:00Z',
-            showUnknownTarget: true,
             esQuery: {
               bool: {
                 filter: [
@@ -438,6 +438,31 @@ export default function (providerContext: FtrProviderContext) {
         expect(response.body).to.have.property('nodes').length(3);
         expect(response.body).to.have.property('edges').length(2);
       });
+
+      it('Should limit number of nodes', async () => {
+        const response = await postGraph(supertest, {
+          nodesLimit: 1,
+          query: {
+            eventIds: [],
+            start: '2024-09-01T00:00:00Z',
+            end: '2024-09-02T00:00:00Z',
+            esQuery: {
+              bool: {
+                filter: [
+                  {
+                    exists: {
+                      field: 'actor.entity.id',
+                    },
+                  },
+                ],
+              },
+            },
+          },
+        }).expect(result(200));
+
+        expect(response.body).to.have.property('nodes').length(3); // Minimal number of nodes in a single relationship
+        expect(response.body).to.have.property('edges').length(2);
+      });
     });
   });
 }
