[Workspace]Optimize workspace permission validation for bulk operatio…

…ns (#7516) * Optimize workspace permission validation for bulk operations Signed-off-by: Lin Wang <wonglam@amazon.com> * Changeset file for PR #7516 created/updated * Remove isSavedObjectsCacheActive in permission control client Signed-off-by: Lin Wang <wonglam@amazon.com> * Update src/plugins/workspace/server/permission_control/client.ts Co-authored-by: SuZhou-Joe <suzhou@amazon.com> Signed-off-by: Lin Wang <wonglam@amazon.com> * Rename cacheSavedObjects to addToCacheAllowlist Signed-off-by: Lin Wang <wonglam@amazon.com> * Use request uuid as cache Signed-off-by: Lin Wang <wonglam@amazon.com> --------- Signed-off-by: Lin Wang <wonglam@amazon.com> Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com> Co-authored-by: SuZhou-Joe <suzhou@amazon.com>
opensearch-project · Aug 7, 2024 · ef47b5f · ef47b5f
1 parent a65a8aa
commit ef47b5f
Show file tree

Hide file tree

Showing 7 changed files with 255 additions and 10 deletions.
diff --git a/changelogs/fragments/7516.yml b/changelogs/fragments/7516.yml
@@ -0,0 +1,2 @@
+feat:
+- [Workspace]Optimize workspace permission validation for bulk operations ([#7516](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7516))
diff --git a/src/plugins/workspace/server/permission_control/client.test.ts b/src/plugins/workspace/server/permission_control/client.test.ts
@@ -102,7 +102,16 @@ describe('PermissionControl', () => {
     });
     const batchValidateResult = await permissionControlClient.batchValidate(
       httpServerMock.createOpenSearchDashboardsRequest(),
-      [],
+      [
+        {
+          id: 'foo',
+          type: 'dashboard',
+        },
+        {
+          id: 'bar',
+          type: 'dashboard',
+        },
+      ],
       ['read']
     );
     expect(batchValidateResult.success).toEqual(true);
@@ -142,7 +151,16 @@ describe('PermissionControl', () => {
     });
     const batchValidateResult = await permissionControlClient.batchValidate(
       httpServerMock.createOpenSearchDashboardsRequest(),
-      [],
+      [
+        {
+          id: 'foo',
+          type: 'dashboard',
+        },
+        {
+          id: 'bar',
+          type: 'dashboard',
+        },
+      ],
       ['read']
     );
     expect(batchValidateResult.success).toEqual(true);
@@ -197,4 +215,123 @@ describe('PermissionControl', () => {
       );
     });
   });
+
+  describe('saved objects cache', () => {
+    it('should not call bulk get again if saved objects cached', async () => {
+      const permissionControlClient = new SavedObjectsPermissionControl(loggerMock.create());
+      const getScopedClient = jest.fn();
+      const clientMock = savedObjectsClientMock.create();
+      const requestMock = httpServerMock.createOpenSearchDashboardsRequest();
+      getScopedClient.mockImplementation((request) => {
+        return clientMock;
+      });
+      permissionControlClient.setup(getScopedClient, mockAuth);
+      permissionControlClient.addToCacheAllowlist(requestMock, [
+        {
+          type: 'workspace',
+          id: 'foo',
+        },
+      ]);
+      clientMock.bulkGet.mockResolvedValue({
+        saved_objects: [
+          {
+            type: 'workspace',
+            id: 'foo',
+            attributes: {},
+            references: [],
+          },
+        ],
+      });
+
+      await permissionControlClient.validate(requestMock, { id: 'foo', type: 'workspace' }, [
+        'read',
+      ]);
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(1);
+
+      await permissionControlClient.validate(requestMock, { id: 'foo', type: 'workspace' }, [
+        'read',
+      ]);
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(1);
+    });
+    it('should call bulk get again for different requests', async () => {
+      const permissionControlClient = new SavedObjectsPermissionControl(loggerMock.create());
+      const getScopedClient = jest.fn();
+      const clientMock = savedObjectsClientMock.create();
+      const requestMock = httpServerMock.createOpenSearchDashboardsRequest();
+      getScopedClient.mockImplementation((request) => {
+        return clientMock;
+      });
+      permissionControlClient.setup(getScopedClient, mockAuth);
+      permissionControlClient.addToCacheAllowlist(requestMock, [
+        {
+          type: 'workspace',
+          id: 'foo',
+        },
+      ]);
+      clientMock.bulkGet.mockResolvedValue({
+        saved_objects: [
+          {
+            type: 'workspace',
+            id: 'foo',
+            attributes: {},
+            references: [],
+          },
+        ],
+      });
+
+      await permissionControlClient.validate(requestMock, { id: 'foo', type: 'workspace' }, [
+        'read',
+      ]);
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(1);
+
+      await permissionControlClient.validate(
+        httpServerMock.createOpenSearchDashboardsRequest({
+          opensearchDashboardsRequestState: {
+            requestId: '123',
+            requestUuid: 'another-uuid',
+          },
+        }),
+        { id: 'foo', type: 'workspace' },
+        ['read']
+      );
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(2);
+    });
+    it('should call bulk get again after cache been cleared', async () => {
+      const permissionControlClient = new SavedObjectsPermissionControl(loggerMock.create());
+      const getScopedClient = jest.fn();
+      const clientMock = savedObjectsClientMock.create();
+      const requestMock = httpServerMock.createOpenSearchDashboardsRequest();
+      getScopedClient.mockImplementation((request) => {
+        return clientMock;
+      });
+      permissionControlClient.setup(getScopedClient, mockAuth);
+      permissionControlClient.addToCacheAllowlist(requestMock, [
+        {
+          type: 'workspace',
+          id: 'foo',
+        },
+      ]);
+      clientMock.bulkGet.mockResolvedValue({
+        saved_objects: [
+          {
+            type: 'workspace',
+            id: 'foo',
+            attributes: {},
+            references: [],
+          },
+        ],
+      });
+
+      await permissionControlClient.validate(requestMock, { id: 'foo', type: 'workspace' }, [
+        'read',
+      ]);
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(1);
+      permissionControlClient.clearSavedObjectsCache(requestMock);
+
+      await permissionControlClient.validate(requestMock, { id: 'foo', type: 'workspace' }, [
+        'read',
+      ]);
+      expect(clientMock.bulkGet).toHaveBeenCalledTimes(2);
+    });
+  });
 });
diff --git a/src/plugins/workspace/server/permission_control/client.ts b/src/plugins/workspace/server/permission_control/client.ts
@@ -13,7 +13,6 @@ import {
   Principals,
   SavedObject,
   WORKSPACE_TYPE,
-  Permissions,
   HttpAuth,
 } from '../../../../core/server';
 import { WORKSPACE_SAVED_OBJECTS_CLIENT_WRAPPER_ID } from '../../common/constants';
@@ -30,6 +29,8 @@ export class SavedObjectsPermissionControl {
   private readonly logger: Logger;
   private _getScopedClient?: SavedObjectsServiceStart['getScopedClient'];
   private auth?: HttpAuth;
+  private _savedObjectCache: Map<string, { [key: string]: SavedObject }> = new Map();
+  private _shouldCachedSavedObjects: Map<string, string[]> = new Map();
   /**
    * Returns a saved objects client that is able to:
    * 1. Read objects whose type is `workspace` because workspace is a hidden type and the permission control client will need to get the metadata of a specific workspace to do the permission check.
@@ -48,11 +49,46 @@ export class SavedObjectsPermissionControl {
     this.logger = logger;
   }
 
+  private generateSavedObjectKey = ({ type, id }: { type: string; id: string }) => {
+    return `${type}:${id}`;
+  };
+
   private async bulkGetSavedObjects(
     request: OpenSearchDashboardsRequest,
     savedObjects: SavedObjectsBulkGetObject[]
   ) {
-    return (await this.getScopedClient?.(request)?.bulkGet(savedObjects))?.saved_objects || [];
+    const requestKey = request.uuid;
+    const savedObjectsToGet = savedObjects.filter(
+      (savedObject) =>
+        !this._savedObjectCache.get(requestKey)?.[this.generateSavedObjectKey(savedObject)]
+    );
+    const retrievedSavedObjects =
+      savedObjectsToGet.length > 0
+        ? (await this.getScopedClient?.(request)?.bulkGet(savedObjectsToGet))?.saved_objects || []
+        : [];
+
+    const retrievedSavedObjectsMap: { [key: string]: SavedObject } = {};
+    retrievedSavedObjects.forEach((savedObject) => {
+      const savedObjectKey = this.generateSavedObjectKey(savedObject);
+      if (this._shouldCachedSavedObjects.get(requestKey)?.includes(savedObjectKey)) {
+        const cachedSavedObjectsMap = this._savedObjectCache.get(requestKey) || {};
+        cachedSavedObjectsMap[savedObjectKey] = savedObject;
+        this._savedObjectCache.set(requestKey, cachedSavedObjectsMap);
+      }
+      retrievedSavedObjectsMap[savedObjectKey] = savedObject;
+    });
+
+    const results: SavedObject[] = [];
+    savedObjects.forEach((savedObject) => {
+      const savedObjectKey = this.generateSavedObjectKey(savedObject);
+      const foundedSavedObject =
+        this._savedObjectCache.get(requestKey)?.[savedObjectKey] ||
+        retrievedSavedObjectsMap[savedObjectKey];
+      if (foundedSavedObject) {
+        results.push(foundedSavedObject);
+      }
+    });
+    return results;
   }
   public async setup(getScopedClient: SavedObjectsServiceStart['getScopedClient'], auth: HttpAuth) {
     this._getScopedClient = getScopedClient;
@@ -175,4 +211,30 @@ export class SavedObjectsPermissionControl {
       result: hasPermissionToAllObjects,
     };
   }
+
+  public addToCacheAllowlist(
+    request: OpenSearchDashboardsRequest,
+    savedObjects: Array<Pick<SavedObjectsBulkGetObject, 'type' | 'id'>>
+  ) {
+    const requestKey = request.uuid;
+    this._shouldCachedSavedObjects.set(
+      requestKey,
+      Array.from(
+        new Set([
+          ...(this._shouldCachedSavedObjects.get(requestKey) ?? []),
+          ...savedObjects.map(this.generateSavedObjectKey),
+        ])
+      )
+    );
+  }
+
+  public clearSavedObjectsCache(request: OpenSearchDashboardsRequest) {
+    const requestKey = request.uuid;
+    if (this._shouldCachedSavedObjects.has(requestKey)) {
+      this._shouldCachedSavedObjects.delete(requestKey);
+    }
+    if (this._savedObjectCache.has(requestKey)) {
+      this._savedObjectCache.delete(requestKey);
+    }
+  }
 }
diff --git a/src/plugins/workspace/server/plugin.test.ts b/src/plugins/workspace/server/plugin.test.ts
@@ -8,6 +8,7 @@ import { coreMock, httpServerMock } from '../../../core/server/mocks';
 import { WorkspacePlugin } from './plugin';
 import { getWorkspaceState, updateWorkspaceState } from '../../../core/server/utils';
 import * as utilsExports from './utils';
+import { SavedObjectsPermissionControl } from './permission_control/client';
 
 describe('Workspace server plugin', () => {
   it('#setup', async () => {
@@ -157,6 +158,22 @@ describe('Workspace server plugin', () => {
       );
       expect(toolKitMock.next).toBeCalledTimes(1);
     });
+
+    it('should clear saved objects cache', async () => {
+      jest.spyOn(utilsExports, 'getPrincipalsFromRequest').mockImplementation(() => ({}));
+      const clearSavedObjectsCacheMock = jest
+        .spyOn(SavedObjectsPermissionControl.prototype, 'clearSavedObjectsCache')
+        .mockImplementationOnce(() => {});
+
+      await workspacePlugin.setup(setupMock);
+      const toolKitMock = httpServerMock.createToolkit();
+
+      expect(setupMock.http.registerOnPreResponse).toHaveBeenCalled();
+      const preResponseFn = setupMock.http.registerOnPreResponse.mock.calls[0][0];
+
+      preResponseFn(requestWithWorkspaceInUrl, { statusCode: 200 }, toolKitMock);
+      expect(clearSavedObjectsCacheMock).toHaveBeenCalled();
+    });
   });
 
   it('#start', async () => {

diff --git a/src/plugins/workspace/server/plugin.ts b/src/plugins/workspace/server/plugin.ts
@@ -101,6 +101,11 @@ export class WorkspacePlugin implements Plugin<WorkspacePluginSetup, WorkspacePl
       WORKSPACE_SAVED_OBJECTS_CLIENT_WRAPPER_ID,
       this.workspaceSavedObjectsClientWrapper.wrapperFactory
     );
+
+    core.http.registerOnPreResponse((request, _response, toolkit) => {
+      this.permissionControl?.clearSavedObjectsCache(request);
+      return toolkit.next();
+    });
   }
 
   constructor(initializerContext: PluginInitializerContext) {

diff --git a/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.test.ts b/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.test.ts
@@ -129,6 +129,7 @@ const generateWorkspaceSavedObjectsClientWrapper = (role = NO_DASHBOARD_ADMIN) =
     getPrincipalsFromRequest: jest.fn().mockImplementation(() => {
       return { users: ['user-1'] };
     }),
+    addToCacheAllowlist: jest.fn(),
   };
 
   const wrapper = new WorkspaceSavedObjectsClientWrapper(permissionControlMock);

diff --git a/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts b/src/plugins/workspace/server/saved_objects/workspace_saved_objects_client_wrapper.ts
@@ -72,6 +72,18 @@ const generateOSDAdminPermissionError = () =>
     )
   );
 
+const getWorkspacesFromSavedObjects = (savedObjects: SavedObject[]) => {
+  return savedObjects
+    .reduce<string[]>(
+      (previous, { workspaces }) => Array.from(new Set([...previous, ...(workspaces ?? [])])),
+      []
+    )
+    .map((id) => ({
+      type: WORKSPACE_TYPE,
+      id,
+    }));
+};
+
 const getDefaultValuesForEmpty = <T>(values: T[] | undefined, defaultValues: T[]) => {
   return !values || values.length === 0 ? defaultValues : values;
 };
@@ -126,15 +138,12 @@ export class WorkspaceSavedObjectsClientWrapper {
       return false;
     }
     for (const workspaceId of workspaces) {
-      const validateResult = await this.permissionControl.validate(
+      const validateResult = await this.validateMultiWorkspacesPermissions(
+        [workspaceId],
         request,
-        {
-          type: WORKSPACE_TYPE,
-          id: workspaceId,
-        },
         permissionModes
       );
-      if (validateResult?.result) {
+      if (validateResult) {
         return true;
       }
     }
@@ -265,6 +274,10 @@ export class WorkspaceSavedObjectsClientWrapper {
       options?: SavedObjectsBulkUpdateOptions
     ): Promise<SavedObjectsBulkUpdateResponse<T>> => {
       const objectsToUpdate = await wrapperOptions.client.bulkGet<T>(objects, options);
+      this.permissionControl.addToCacheAllowlist(
+        wrapperOptions.request,
+        getWorkspacesFromSavedObjects(objectsToUpdate.saved_objects)
+      );
 
       for (const object of objectsToUpdate.saved_objects) {
         const permitted = await validateUpdateWithWorkspacePermission(object);
@@ -322,6 +335,10 @@ export class WorkspaceSavedObjectsClientWrapper {
                 throw error;
               }
             }
+            this.permissionControl.addToCacheAllowlist(
+              wrapperOptions.request,
+              getWorkspacesFromSavedObjects([rawObject])
+            );
             if (
               !(await this.validateWorkspacesAndSavedObjectsPermissions(
                 rawObject,
@@ -426,6 +443,10 @@ export class WorkspaceSavedObjectsClientWrapper {
       options: SavedObjectsBaseOptions = {}
     ): Promise<SavedObjectsBulkResponse<T>> => {
       const objectToBulkGet = await wrapperOptions.client.bulkGet<T>(objects, options);
+      this.permissionControl.addToCacheAllowlist(
+        wrapperOptions.request,
+        getWorkspacesFromSavedObjects(objectToBulkGet.saved_objects)
+      );
 
       for (const object of objectToBulkGet.saved_objects) {
         if (object.type === DATA_SOURCE_SAVED_OBJECT_TYPE) {