fix(keycloak): customize quantity of parallel Keycloak requests

Signed-off-by: Oleksandr Andriienko <[email protected]>
backstage · Dec 2, 2024 · 37986c6 · 37986c6
1 parent 056a5ca
commit 37986c6
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 59 deletions.
diff --git a/workspaces/keycloak/plugins/catalog-backend-module-keycloak/package.json b/workspaces/keycloak/plugins/catalog-backend-module-keycloak/package.json
@@ -48,6 +48,7 @@
     "@keycloak/keycloak-admin-client": "24.0.5",
     "inclusion": "^1.0.1",
     "lodash": "^4.17.21",
+    "p-limit": "^6.1.0",
     "pg-format": "^1.0.4",
     "uuid": "^9.0.1"
   },

diff --git a/workspaces/keycloak/plugins/catalog-backend-module-keycloak/src/lib/config.ts b/workspaces/keycloak/plugins/catalog-backend-module-keycloak/src/lib/config.ts
@@ -91,6 +91,11 @@ export type KeycloakProviderConfig = {
    * @see https://www.keycloak.org/docs-api/11.0/rest-api/index.html#_groups_resource
    */
   groupQuerySize?: number;
+
+  /**
+   * Maximum request concurrency to prevent DoS attacks on the Keycloak server.
+   */
+  maxConcurrency?: number;
 };
 
 const readProviderConfig = (

diff --git a/workspaces/keycloak/plugins/catalog-backend-module-keycloak/src/lib/read.ts b/workspaces/keycloak/plugins/catalog-backend-module-keycloak/src/lib/read.ts
@@ -22,6 +22,7 @@ import type GroupRepresentation from '@keycloak/keycloak-admin-client/lib/defs/g
 import type UserRepresentation from '@keycloak/keycloak-admin-client/lib/defs/userRepresentation';
 import type { Groups } from '@keycloak/keycloak-admin-client/lib/resources/groups';
 import type { Users } from '@keycloak/keycloak-admin-client/lib/resources/users';
+import pLimit from 'p-limit';
 
 import { KeycloakProviderConfig } from './config';
 import {
@@ -112,6 +113,7 @@ export async function getEntities<T extends Users | Groups>(
   config: KeycloakProviderConfig,
   logger: LoggerService,
   entityQuerySize: number = KEYCLOAK_ENTITY_QUERY_SIZE,
+  concurrency: number = Number.POSITIVE_INFINITY,
 ): Promise<Awaited<ReturnType<T['find']>>> {
   const rawEntityCount = await entities.count({ realm: config.realm });
   const entityCount =
@@ -120,18 +122,19 @@ export async function getEntities<T extends Users | Groups>(
   const pageCount = Math.ceil(entityCount / entityQuerySize);
 
   // The next line acts like range in python
-  const entityPromises = Array.from(
-    { length: pageCount },
-    (_, i) =>
+  const entityPromises = Array.from({ length: pageCount }, (_, i) =>
+    pLimit(concurrency)(() =>
       entities
         .find({
           realm: config.realm,
           max: entityQuerySize,
           first: i * entityQuerySize,
         })
-        .catch(err =>
-          logger.warn('Failed to retieve Keycloak entities.', err),
-        ) as ReturnType<T['find']>,
+        .catch(err => {
+          logger.warn('Failed to retrieve Keycloak entities.', err);
+          return [];
+        }),
+    ),
   );
 
   const entityResults = (await Promise.all(entityPromises)).flat() as Awaited<
@@ -227,18 +230,22 @@ export const readKeycloakRealm = async (
   users: UserEntity[];
   groups: GroupEntity[];
 }> => {
+  const concurrency = config.maxConcurrency ?? Number.POSITIVE_INFINITY;
+
   const kUsers = await getEntities(
     client.users,
     config,
     logger,
     options?.userQuerySize,
+    concurrency,
   );
 
   const topLevelKGroups = (await getEntities(
     client.groups,
     config,
     logger,
     options?.groupQuerySize,
+    concurrency,
   )) as GroupRepresentationWithParent[];
 
   let serverVersion: number;
@@ -269,65 +276,78 @@ export const readKeycloakRealm = async (
       [] as GroupRepresentationWithParent[],
     );
   }
+  const limit = pLimit(concurrency);
   const kGroups = await Promise.all(
-    rawKGroups.map(async g => {
-      g.members = await getAllGroupMembers(
-        client.groups as Groups,
-        g.id!,
-        config,
-        options,
-      );
-
-      if (isVersion23orHigher) {
-        if (g.subGroupCount! > 0) {
-          g.subGroups = await client.groups.listSubGroups({
-            parentId: g.id!,
-            first: 0,
-            max: g.subGroupCount,
-            briefRepresentation: false,
-            realm: config.realm,
-          });
+    rawKGroups.map(g =>
+      limit(async () => {
+        g.members = await getAllGroupMembers(
+          client.groups as Groups,
+          g.id!,
+          config,
+          options,
+        );
+
+        if (isVersion23orHigher) {
+          if (g.subGroupCount! > 0) {
+            g.subGroups = await client.groups.listSubGroups({
+              parentId: g.id!,
+              first: 0,
+              max: g.subGroupCount,
+              briefRepresentation: false,
+              realm: config.realm,
+            });
+          }
+          if (g.parentId) {
+            const groupParent = await client.groups.findOne({
+              id: g.parentId,
+              realm: config.realm,
+            });
+            g.parent = groupParent?.name;
+          }
         }
-        if (g.parentId) {
-          const groupParent = await client.groups.findOne({
-            id: g.parentId,
-            realm: config.realm,
-          });
-          g.parent = groupParent?.name;
-        }
-      }
 
-      return g;
-    }),
+        return g;
+      }),
+    ),
   );
 
-  const parsedGroups = await kGroups.reduce(async (promise, g) => {
-    const partial = await promise;
-    const entity = await parseGroup(g, config.realm, options?.groupTransformer);
-    if (entity) {
-      const group = {
-        ...g,
-        entity,
-      } as GroupRepresentationWithParentAndEntity;
-      partial.push(group);
-    }
-    return partial;
-  }, Promise.resolve([] as GroupRepresentationWithParentAndEntity[]));
+  const parsedGroups = await kGroups.reduce(
+    async (promise, g) => {
+      const partial = await promise;
+      const entity = await parseGroup(
+        g,
+        config.realm,
+        options?.groupTransformer,
+      );
+      if (entity) {
+        const group = {
+          ...g,
+          entity,
+        } as GroupRepresentationWithParentAndEntity;
+        partial.push(group);
+      }
+      return partial;
+    },
+    Promise.resolve([] as GroupRepresentationWithParentAndEntity[]),
+  );
 
-  const parsedUsers = await kUsers.reduce(async (promise, u) => {
-    const partial = await promise;
-    const entity = await parseUser(
-      u,
-      config.realm,
-      parsedGroups,
-      options?.userTransformer,
-    );
-    if (entity) {
-      const user = { ...u, entity } as UserRepresentationWithEntity;
-      partial.push(user);
-    }
-    return partial;
-  }, Promise.resolve([] as UserRepresentationWithEntity[]));
+  const parsedUsers = await kUsers.reduce(
+    async (promise, u) => {
+      const partial = await promise;
+      const entity = await parseUser(
+        u,
+        config.realm,
+        parsedGroups,
+        options?.userTransformer,
+      );
+      if (entity) {
+        const user = { ...u, entity } as UserRepresentationWithEntity;
+        partial.push(user);
+      }
+      return partial;
+    },
+    Promise.resolve([] as UserRepresentationWithEntity[]),
+  );
 
   const groups = parsedGroups.map(g => {
     const entity = g.entity;

diff --git a/workspaces/keycloak/yarn.lock b/workspaces/keycloak/yarn.lock
@@ -2481,6 +2481,7 @@ __metadata:
     deepmerge: 4.3.1
     inclusion: ^1.0.1
     lodash: ^4.17.21
+    p-limit: ^6.1.0
     pg-format: ^1.0.4
     prettier: 3.3.3
     uuid: ^9.0.1
@@ -4185,6 +4186,7 @@ __metadata:
     "@spotify/prettier-config": ^12.0.0
     knip: ^5.27.4
     node-gyp: ^9.0.0
+    p-limit: ^6.1.0
     prettier: ^2.3.2
     typescript: ~5.3.0
   languageName: unknown
@@ -18041,6 +18043,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"p-limit@npm:^6.1.0":
+  version: 6.1.0
+  resolution: "p-limit@npm:6.1.0"
+  dependencies:
+    yocto-queue: ^1.1.1
+  checksum: 0c98d8fc1006b70fc7423232a47e8d026dc69279b06fe7ff8b4c0cc8023de2b6bb8991b609d93c3dec691a7a362ab0f0157df521d931a01fec192a5e404b9ee5
+  languageName: node
+  linkType: hard
+
 "p-locate@npm:^3.0.0":
   version: 3.0.0
   resolution: "p-locate@npm:3.0.0"
@@ -23286,6 +23297,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"yocto-queue@npm:^1.1.1":
+  version: 1.1.1
+  resolution: "yocto-queue@npm:1.1.1"
+  checksum: f2e05b767ed3141e6372a80af9caa4715d60969227f38b1a4370d60bffe153c9c5b33a862905609afc9b375ec57cd40999810d20e5e10229a204e8bde7ef255c
+  languageName: node
+  linkType: hard
+
 "zip-stream@npm:^6.0.1":
   version: 6.0.1
   resolution: "zip-stream@npm:6.0.1"