From cfc5848d140d02f3f01de4d52a3a0a17e3c34190 Mon Sep 17 00:00:00 2001
From: mike12345567 <me@michaeldrury.co.uk>
Date: Thu, 17 Oct 2024 16:10:32 +0100
Subject: [PATCH] Improving how traversal is performed for role inheritance.

---
 packages/backend-core/src/security/roles.ts | 109 +++++++++++---------
 1 file changed, 61 insertions(+), 48 deletions(-)

diff --git a/packages/backend-core/src/security/roles.ts b/packages/backend-core/src/security/roles.ts
index fdea09944a..44abc64d00 100644
--- a/packages/backend-core/src/security/roles.ts
+++ b/packages/backend-core/src/security/roles.ts
@@ -10,6 +10,7 @@ import { getAppDB } from "../context"
 import { Screen, Role as RoleDoc, RoleUIMetadata } from "@budibase/types"
 import cloneDeep from "lodash/fp/cloneDeep"
 import { RoleColor, helpers } from "@budibase/shared-core"
+import { uniqBy } from "lodash"
 
 export const BUILTIN_ROLE_IDS = {
   ADMIN: "ADMIN",
@@ -38,6 +39,14 @@ export const RoleIDVersion = {
   NAME: "name",
 }
 
+function rolesInList(roleIds: string[], ids: string | string[]) {
+  if (Array.isArray(ids)) {
+    return ids.filter(id => roleIds.includes(id)).length === ids.length
+  } else {
+    return roleIds.includes(ids)
+  }
+}
+
 export class Role implements RoleDoc {
   _id: string
   _rev?: string
@@ -76,6 +85,54 @@ export class Role implements RoleDoc {
   }
 }
 
+export class RoleHierarchyTraversal {
+  allRoles: RoleDoc[]
+  opts?: { defaultPublic?: boolean }
+
+  constructor(allRoles: RoleDoc[], opts?: { defaultPublic?: boolean }) {
+    this.allRoles = allRoles
+    this.opts = opts
+  }
+
+  walk(role: RoleDoc): RoleDoc[] {
+    const opts = this.opts,
+      allRoles = this.allRoles
+    // this will be a full walked list of roles - which may contain duplicates
+    const roleList: RoleDoc[] = []
+    if (!role || !role._id) {
+      return roleList
+    }
+    roleList.push(role)
+    if (Array.isArray(role.inherits)) {
+      for (let roleId of role.inherits) {
+        const foundRole = findRole(roleId, allRoles, opts)
+        if (foundRole) {
+          return this.walk(foundRole)
+        }
+      }
+    } else {
+      const foundRoleIds: string[] = []
+      let currentRole: RoleDoc | undefined = role
+      while (
+        currentRole &&
+        currentRole.inherits &&
+        !rolesInList(foundRoleIds, currentRole.inherits)
+      ) {
+        if (Array.isArray(currentRole.inherits)) {
+          return this.walk(currentRole)
+        } else {
+          foundRoleIds.push(currentRole.inherits)
+          currentRole = findRole(currentRole.inherits, allRoles, opts)
+          if (role) {
+            roleList.push(role)
+          }
+        }
+      }
+    }
+    return uniqBy(roleList, role => role._id)
+  }
+}
+
 const BUILTIN_ROLES = {
   ADMIN: new Role(
     BUILTIN_IDS.ADMIN,
@@ -309,59 +366,15 @@ async function getAllUserRoles(
   if (userRoleId === BUILTIN_IDS.ADMIN) {
     return allRoles
   }
-  const rolesFound = (ids: string | string[]) => {
-    if (Array.isArray(ids)) {
-      return ids.filter(id => roleIds.includes(id)).length === ids.length
-    } else {
-      return roleIds.includes(ids)
-    }
-  }
-
-  const roleIds = [userRoleId]
-  const roles: RoleDoc[] = []
-  const iterateInherited = (role: RoleDoc | undefined) => {
-    if (!role || !role._id) {
-      return
-    }
-    roleIds.push(role._id)
-    roles.push(role)
-    if (Array.isArray(role.inherits)) {
-      role.inherits.forEach(roleId => {
-        const foundRole = findRole(roleId, allRoles, opts)
-        if (foundRole) {
-          iterateInherited(foundRole)
-        }
-      })
-    } else {
-      while (role && role.inherits && !rolesFound(role.inherits)) {
-        if (Array.isArray(role.inherits)) {
-          iterateInherited(role)
-          break
-        } else {
-          roleIds.push(role.inherits)
-          role = findRole(role.inherits, allRoles, opts)
-          if (role) {
-            roles.push(role)
-          }
-        }
-      }
-    }
-  }
 
   // get all the inherited roles
   const foundRole = findRole(userRoleId, allRoles, opts)
+  let roles: RoleDoc[] = []
   if (foundRole) {
-    iterateInherited(foundRole)
+    const traversal = new RoleHierarchyTraversal(allRoles, opts)
+    roles = traversal.walk(foundRole)
   }
-  const foundRoleIds: string[] = []
-  return roles.filter(role => {
-    if (role._id && !foundRoleIds.includes(role._id)) {
-      foundRoleIds.push(role._id)
-      return true
-    } else {
-      return false
-    }
-  })
+  return roles
 }
 
 export async function getUserRoleIdHierarchy(