ui: add API for command filtering for embedders

ui/src/base/registry.ts

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {assertExists, assertFalse} from './logging';
+
 export interface HasKind {
   kind: string;
 }
@@ -26,6 +28,8 @@ export class RegistryError extends Error {
 export class Registry<T> {
   private key: (t: T) => string;
   protected registry: Map<string, T>;
+  private _keyFilter?: (key: string) => boolean;
+  private readonly keyFilter = (key: string) => this._keyFilter?.(key) ?? true;
 
   static kindRegistry<T extends HasKind>(): Registry<T> {
     return new Registry<T>((t) => t.kind);
@@ -36,8 +40,33 @@ export class Registry<T> {
     this.key = key;
   }
 
+  /**
+   * Set a filter to allow services to be registered only under certain matching keys.
+   * This is intended for applications embedding the Perfetto UI to exclude services
+   * that are inappropriate or otherwise unwanted in their contexts. Initially, a
+   * registry has no filter.
+   *
+   * **Note** that a filter may only be set once. An attempt to replace or clear the
+   * filter will throw an error.
+   */
+  setFilter(filter: (key: string) => boolean): void {
+    assertFalse(this._keyFilter !== undefined, 'A key filter is already set.');
+    this._keyFilter = assertExists(filter);
+
+    // Run the filter to knock out anything already registered that does not pass it
+    [...this.registry.keys()]
+      .filter((key) => !filter(key))
+      .forEach((key) => this.registry.delete(key));
+  }
+
   register(registrant: T): Disposable {
     const kind = this.key(registrant);
+    if (!this.keyFilter(kind)) {
+      // Simply refuse to register the entry
+      return {
+        [Symbol.dispose]: () => undefined,
+      };
+    }
     if (this.registry.has(kind)) {
       throw new RegistryError(
         `Registrant ${kind} already exists in the registry`,
@@ -85,30 +114,48 @@ export class Registry<T> {
     // A proxy is not sufficient because we need non-overridden
     // methods to delegate to overridden methods.
     const result = new (class ChildRegistry extends Registry<T> {
-      constructor (private readonly parent: Registry<T>) {
+      constructor(private readonly parent: Registry<T>) {
         super(parent.key);
       }
 
+      override setFilter(filter: (key: string) => boolean): void {
+        // Dyamically delegate to whatever the parent filter is at the
+        // time of filtering
+        const parentFilter = this.parent.keyFilter.bind(this.parent);
+        const combinedFilter = (key: string) =>
+          filter(key) && parentFilter(key);
+
+        super.setFilter(combinedFilter);
+      }
+
       override has(kind: string): boolean {
-        return this.registry.has(kind) || this.parent.has(kind);
+        return (
+          this.keyFilter(kind) &&
+          (this.registry.has(kind) || this.parent.has(kind))
+        );
       }
 
       override get(kind: string): T {
+        if (!this.keyFilter(kind)) {
+          return super.get(kind); // This will throw a consistent Error type
+        }
         return this.tryGet(kind) ?? this.parent.get(kind);
       }
 
       override tryGet(kind: string): T | undefined {
-        return this.registry.get(kind) ?? this.parent.tryGet(kind);
+        return !this.keyFilter(kind)
+          ? undefined
+          : this.registry.get(kind) ?? this.parent.tryGet(kind);
       }
 
       override *values() {
         // Yield own values first
         yield* this.registry.values();
 
-        // Then yield parent values not shadowed by my keys
+        // Then yield parent values not shadowed by my keys and that pass my filter
         for (const value of this.parent.values()) {
           const kind = this.key(value);
-          if (!this.registry.has(kind)) {
+          if (!this.registry.has(kind) && this.keyFilter(kind)) {
             yield value;
           }
         }

ui/src/base/registry_unittest.ts

@@ -59,6 +59,43 @@ test('registry allows iteration', () => {
   expect(values.includes(b)).toBe(true);
 });
 
+describe('registry filter', () => {
+  test('prevents registration of non-matching kinds', () => {
+    const registry = Registry.kindRegistry<Registrant>();
+    registry.setFilter((key) => key.startsWith('a'));
+
+    const alpha: Registrant = {kind: 'alpha', n: 1};
+    const beta: Registrant = {kind: 'beta', n: 2};
+    registry.register(alpha);
+    registry.register(beta);
+
+    expect(registry.has('alpha')).toBe(true);
+    expect(registry.has('beta')).toBe(false);
+    expect(() => registry.get('beta')).toThrow();
+    expect(registry.get('alpha')).toBe(alpha);
+  });
+
+  test('removes extant non-matching registrations', () => {
+    const registry = Registry.kindRegistry<Registrant>();
+    const a: Registrant = {kind: 'alpha', n: 1};
+    const b: Registrant = {kind: 'beta', n: 2};
+    registry.register(a);
+    registry.register(b);
+
+    registry.setFilter((key) => key.startsWith('a'));
+
+    expect(registry.has('alpha')).toBe(true);
+    expect(registry.has('beta')).toBe(false);
+    expect(() => registry.get('beta')).toThrow();
+  });
+});
+
+test('cannot replace a filter', () => {
+  const registry = Registry.kindRegistry<Registrant>();
+  registry.setFilter(() => true);
+  expect(() => registry.setFilter(() => false)).toThrow();
+});
+
 describe('Hierarchical (child) registries', () => {
   test('inheritance of registrations', () => {
     const parent = Registry.kindRegistry<Registrant>();
@@ -176,4 +213,65 @@ describe('Hierarchical (child) registries', () => {
 
     expect(child).not.toHaveProperty('id');
   });
+
+  describe('registry filter', () => {
+    test('child does not inherit parent filter', () => {
+      const parent = Registry.kindRegistry<Registrant>();
+      parent.setFilter((key) => key.startsWith('x'));
+
+      const child = parent.createChild();
+
+      const child1: Registrant = {kind: 'xyz', n: 1};
+      const child2: Registrant = {kind: 'abc', n: 2};
+      child.register(child1);
+      child.register(child2);
+
+      const parentOk: Registrant = {kind: 'xenon', n: 3};
+      const parentNok: Registrant = {kind: 'beta', n: 4};
+      parent.register(parentOk);
+      parent.register(parentNok);
+
+      expect(child.has('xyz')).toBe(true);
+      expect(child.get('xyz')).toBe(child1);
+      expect(child.has('abc')).toBe(true);
+      expect(child.get('abc')).toBe(child2);
+
+      // Other registrations still accessible (or not) as usual via inheritance
+      expect(child.get('xenon')).toBe(parentOk);
+      expect(child.has('beta')).toBe(false);
+      expect(() => child.get('beta')).toThrow();
+    });
+
+    test('setting filter on child does not affect parent', () => {
+      const parent = Registry.kindRegistry<Registrant>();
+      const a: Registrant = {kind: 'a', n: 1};
+      const b: Registrant = {kind: 'b', n: 2};
+      parent.register(a);
+      parent.register(b);
+
+      const child = parent.createChild();
+      const c: Registrant = {kind: 'c', n: 3};
+      const d: Registrant = {kind: 'd', n: 4};
+      child.register(c);
+      child.register(d);
+
+      child.setFilter((key) => key === 'b' || key === 'd');
+
+      // Parent not pruned
+      expect(parent.has('a')).toBe(true);
+      expect(parent.has('b')).toBe(true);
+
+      // Child pruned
+      expect(child.has('c')).toBe(false);
+      expect(child.has('d')).toBe(true);
+
+      // Other registrations still accessible (or not) as usual
+      expect(child.get('b')).toBe(b);
+
+      // But inheritance needs to respect the child's filtering
+      expect(child.has('a')).toBe(false);
+      expect(() => child.get('a')).toThrow();
+      expect(child.valuesAsArray()).toStrictEqual([d, b]);
+    });
+  });
 });

ui/src/core/command_manager.ts

@@ -82,6 +82,23 @@ export class CommandManagerImpl implements CommandManager {
     return new CommandManagerImpl(this.registry);
   }
 
+  /**
+   * Set a filter to screen command registrations. Command IDs that do not
+   * pass the filter are not registered. This is distinct from the
+   * `allowlistCheck` function, which screens out start-up commands that
+   * should be skipped.
+   *
+   * This is intended for applications embedding the Perfetto UI to exclude services
+   * that are inappropriate or otherwise unwanted in their contexts. Initially, a
+   * registry has no filter.
+   *
+   * **Note** that a filter may only be set once. An attempt to replace or clear the
+   * filter will throw an error.
+   */
+  setFilter(filter: (commandId: string) => boolean) {
+    this.registry.setFilter(filter);
+  }
+
   getCommand(commandId: string): Command {
     return this.registry.get(commandId);
   }

ui/src/core/command_manager_unittest.ts

@@ -38,3 +38,31 @@ describe('CommandManagerImpl child manager', () => {
     expect(parent.hasCommand('child')).toBe(false);
   });
 });
+
+describe('CommandManagerImpl filtering', () => {
+  test('filters existing and future registrations', () => {
+    const parent = new CommandManagerImpl();
+
+    const allowCmd = TestCommand('allow');
+    const denyCmd = TestCommand('deny');
+
+    parent.registerCommand(allowCmd);
+    parent.registerCommand(denyCmd);
+
+    const mgr = parent.createChild();
+    expect(mgr.hasCommand('allow')).toBe(true);
+    expect(mgr.hasCommand('deny')).toBe(true);
+
+    mgr.setFilter((id) => id.startsWith('allow'));
+
+    // Extant non-matching commands are purged
+    expect(mgr.hasCommand('allow')).toBe(true);
+    expect(mgr.hasCommand('deny')).toBe(false);
+
+    // New registrations are screened by the filter.
+    mgr.registerCommand(TestCommand('deny'));
+    expect(mgr.hasCommand('deny')).toBe(false);
+    mgr.registerCommand(TestCommand('allowed'));
+    expect(mgr.hasCommand('allowed')).toBe(true);
+  });
+});