From c4ea1ce29290c898d278fb590a3cb34cb8d80cdc Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 14 Apr 2023 13:22:29 +0100 Subject: [PATCH 1/3] Support excluding types from rate limiting in approval controller --- .../src/ApprovalController.test.ts | 24 +++++++++++++++++++ .../src/ApprovalController.ts | 11 ++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/packages/approval-controller/src/ApprovalController.test.ts b/packages/approval-controller/src/ApprovalController.test.ts index 2d300a1da0..ac116b6342 100644 --- a/packages/approval-controller/src/ApprovalController.test.ts +++ b/packages/approval-controller/src/ApprovalController.test.ts @@ -224,6 +224,30 @@ describe('approval controller', () => { }), ).toThrow(getOriginTypeCollisionError('bar.baz', 'myType')); }); + + it('does not throw on origin and type collision if type excluded', () => { + approvalController = new ApprovalController({ + messenger: getRestrictedMessenger(), + showApprovalRequest, + typesExcludedFromRateLimiting: ['myType'], + }); + + expect(() => + approvalController.add({ + id: 'foo', + origin: 'bar.baz', + type: 'myType', + }), + ).not.toThrow(); + + expect(() => + approvalController.add({ + id: 'foo1', + origin: 'bar.baz', + type: 'myType', + }), + ).not.toThrow(); + }); }); // otherwise tested by 'add' above diff --git a/packages/approval-controller/src/ApprovalController.ts b/packages/approval-controller/src/ApprovalController.ts index a434d6a1f4..21d9fbc213 100644 --- a/packages/approval-controller/src/ApprovalController.ts +++ b/packages/approval-controller/src/ApprovalController.ts @@ -156,6 +156,7 @@ type ApprovalControllerOptions = { messenger: ApprovalControllerMessenger; showApprovalRequest: ShowApprovalRequest; state?: Partial; + typesExcludedFromRateLimiting?: string[]; }; /** @@ -178,6 +179,8 @@ export class ApprovalController extends BaseControllerV2< private _showApprovalRequest: () => void; + private _typesExcludedFromRateLimiting: string[]; + /** * Construct an Approval controller. * @@ -186,11 +189,13 @@ export class ApprovalController extends BaseControllerV2< * the request can be displayed to the user. * @param options.messenger - The restricted controller messenger for the Approval controller. * @param options.state - The initial controller state. + * @param options.typesExcludedFromRateLimiting - Array of aproval types which allow multiple pending approval requests from the same origin. */ constructor({ messenger, showApprovalRequest, state = {}, + typesExcludedFromRateLimiting = [], }: ApprovalControllerOptions) { super({ name: controllerName, @@ -202,6 +207,7 @@ export class ApprovalController extends BaseControllerV2< this._approvals = new Map(); this._origins = new Map(); this._showApprovalRequest = showApprovalRequest; + this._typesExcludedFromRateLimiting = typesExcludedFromRateLimiting; this.registerMessageHandlers(); } @@ -487,7 +493,10 @@ export class ApprovalController extends BaseControllerV2< ): Promise { this._validateAddParams(id, origin, type, requestData, requestState); - if (this._origins.get(origin)?.has(type)) { + if ( + !this._typesExcludedFromRateLimiting.includes(type) && + this._origins.get(origin)?.has(type) + ) { throw ethErrors.rpc.resourceUnavailable( getAlreadyPendingMessage(origin, type), ); From e53fe5441ebe37675487191cb15bb531bf5fc06b Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 14 Apr 2023 15:09:15 +0100 Subject: [PATCH 2/3] Update origin and type data structure to support disabling rate limiting --- .../src/ApprovalController.test.ts | 52 ++++++++++++++----- .../src/ApprovalController.ts | 40 +++++++------- 2 files changed, 58 insertions(+), 34 deletions(-) diff --git a/packages/approval-controller/src/ApprovalController.test.ts b/packages/approval-controller/src/ApprovalController.test.ts index ac116b6342..ae1ac7c6c6 100644 --- a/packages/approval-controller/src/ApprovalController.test.ts +++ b/packages/approval-controller/src/ApprovalController.test.ts @@ -430,6 +430,21 @@ describe('approval controller', () => { approvalController.getApprovalCount({ type: 'type3' }), ).toStrictEqual(0); }); + + it('gets the count when specifying origin and type with type excluded from rate limiting', () => { + approvalController = new ApprovalController({ + messenger: getRestrictedMessenger(), + showApprovalRequest: sinon.spy(), + typesExcludedFromRateLimiting: [TYPE], + }); + + addWithCatch({ id: '1', origin: 'origin1', type: TYPE }); + addWithCatch({ id: '2', origin: 'origin1', type: TYPE }); + + expect( + approvalController.getApprovalCount({ origin: 'origin1', type: TYPE }), + ).toStrictEqual(2); + }); }); describe('getTotalApprovalCount', () => { @@ -458,6 +473,30 @@ describe('approval controller', () => { approvalController.clear(new EthereumRpcError(1, 'clear')); expect(approvalController.getTotalApprovalCount()).toStrictEqual(0); }); + + it('gets the total approval count with type excluded from rate limiting', () => { + const approvalController = new ApprovalController({ + messenger: getRestrictedMessenger(), + showApprovalRequest: sinon.spy(), + typesExcludedFromRateLimiting: ['type0'], + }); + expect(approvalController.getTotalApprovalCount()).toStrictEqual(0); + + const addWithCatch = (args: any) => + approvalController.add(args).catch(() => undefined); + + addWithCatch({ id: '1', origin: 'origin1', type: 'type0' }); + expect(approvalController.getTotalApprovalCount()).toStrictEqual(1); + + addWithCatch({ id: '2', origin: 'origin1', type: 'type0' }); + expect(approvalController.getTotalApprovalCount()).toStrictEqual(2); + + approvalController.reject('2', new Error('foo')); + expect(approvalController.getTotalApprovalCount()).toStrictEqual(1); + + approvalController.clear(new EthereumRpcError(1, 'clear')); + expect(approvalController.getTotalApprovalCount()).toStrictEqual(0); + }); }); describe('has', () => { @@ -865,19 +904,6 @@ describe('approval controller', () => { }); }); - // TODO: Stop using private methods in tests - describe('_isEmptyOrigin', () => { - it('handles non-existing origin', () => { - const approvalController = new ApprovalController({ - messenger: getRestrictedMessenger(), - showApprovalRequest: sinon.spy(), - }); - expect(() => - (approvalController as any)._isEmptyOrigin('kaplar'), - ).not.toThrow(); - }); - }); - describe('actions', () => { it('addApprovalRequest: shouldShowRequest = true', async () => { const messenger = new ControllerMessenger< diff --git a/packages/approval-controller/src/ApprovalController.ts b/packages/approval-controller/src/ApprovalController.ts index 21d9fbc213..563dc15e30 100644 --- a/packages/approval-controller/src/ApprovalController.ts +++ b/packages/approval-controller/src/ApprovalController.ts @@ -175,7 +175,7 @@ export class ApprovalController extends BaseControllerV2< > { private _approvals: Map; - private _origins: Map>; + private _origins: Map>; private _showApprovalRequest: () => void; @@ -339,11 +339,13 @@ export class ApprovalController extends BaseControllerV2< const { origin, type: _type } = opts; if (origin && _type) { - return Number(Boolean(this._origins.get(origin)?.has(_type))); + return this._origins.get(origin)?.get(_type) || 0; } if (origin) { - return this._origins.get(origin)?.size || 0; + return Array.from( + (this._origins.get(origin) || new Map()).values(), + ).reduce((total, value) => total + value, 0); } // Only "type" was specified @@ -401,7 +403,7 @@ export class ApprovalController extends BaseControllerV2< // Check origin and type pair if type also specified if (_type) { - return Boolean(this._origins.get(origin)?.has(_type)); + return Boolean(this._origins.get(origin)?.get(_type)); } return this._origins.has(origin); } @@ -495,7 +497,7 @@ export class ApprovalController extends BaseControllerV2< if ( !this._typesExcludedFromRateLimiting.includes(type) && - this._origins.get(origin)?.has(type) + this.has({ origin, type }) ) { throw ethErrors.rpc.resourceUnavailable( getAlreadyPendingMessage(origin, type), @@ -560,11 +562,13 @@ export class ApprovalController extends BaseControllerV2< * @param type - The type associated with the approval request. */ private _addPendingApprovalOrigin(origin: string, type: string): void { - const originSet = this._origins.get(origin) || new Set(); - originSet.add(type); + const originMap = this._origins.get(origin) || new Map(); + const currentValue = originMap.get(type) || 0; + + originMap.set(type, currentValue + 1); if (!this._origins.has(origin)) { - this._origins.set(origin, originSet); + this._origins.set(origin, originMap); } } @@ -619,9 +623,14 @@ export class ApprovalController extends BaseControllerV2< // eslint-disable-next-line @typescript-eslint/no-non-null-assertion const { origin, type } = this.state.pendingApprovals[id]!; - (this._origins.get(origin) as Set).delete(type); - if (this._isEmptyOrigin(origin)) { + const originMap = this._origins.get(origin) as Map; + const originTotalCount = this.getApprovalCount({ origin }); + const originTypeCount = originMap.get(type) as number; + + if (originTotalCount === 1) { this._origins.delete(origin); + } else { + originMap.set(type, originTypeCount - 1); } this.update((draftState) => { @@ -649,16 +658,5 @@ export class ApprovalController extends BaseControllerV2< this._delete(id); return callbacks; } - - /** - * Checks whether there are any approvals associated with the given - * origin. - * - * @param origin - The origin to check. - * @returns True if the origin has no approvals, false otherwise. - */ - private _isEmptyOrigin(origin: string): boolean { - return !this._origins.get(origin)?.size; - } } export default ApprovalController; From ff77411e7efc5adf491991a78d389aae9f529060 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 14 Apr 2023 16:20:21 +0100 Subject: [PATCH 3/3] Simplify origin increment logic --- .../approval-controller/src/ApprovalController.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/packages/approval-controller/src/ApprovalController.ts b/packages/approval-controller/src/ApprovalController.ts index 563dc15e30..61dfaf6c66 100644 --- a/packages/approval-controller/src/ApprovalController.ts +++ b/packages/approval-controller/src/ApprovalController.ts @@ -562,14 +562,15 @@ export class ApprovalController extends BaseControllerV2< * @param type - The type associated with the approval request. */ private _addPendingApprovalOrigin(origin: string, type: string): void { - const originMap = this._origins.get(origin) || new Map(); - const currentValue = originMap.get(type) || 0; - - originMap.set(type, currentValue + 1); + let originMap = this._origins.get(origin); - if (!this._origins.has(origin)) { + if (!originMap) { + originMap = new Map(); this._origins.set(origin, originMap); } + + const currentValue = originMap.get(type) || 0; + originMap.set(type, currentValue + 1); } /**