From c8ed6765dbf751e1c9678794ebb6d7556dc217a0 Mon Sep 17 00:00:00 2001
From: Andrew Morris <voltrevo@gmail.com>
Date: Tue, 17 May 2022 11:49:55 +1000
Subject: [PATCH] Enable eth_setPreferredAggregator

---
 .../Controllers/Network/createEthMiddleware.ts | 12 ++++++++++++
 .../Controllers/Network/createJsonRpcClient.ts |  3 +--
 .../source/Controllers/QuillController.ts      | 18 ++++++++++++++++--
 .../source/Controllers/knownTransactions.ts    |  1 +
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/extension/source/Controllers/Network/createEthMiddleware.ts b/extension/source/Controllers/Network/createEthMiddleware.ts
index adb867fc..bbc20f94 100644
--- a/extension/source/Controllers/Network/createEthMiddleware.ts
+++ b/extension/source/Controllers/Network/createEthMiddleware.ts
@@ -24,6 +24,7 @@ export interface IProviderHandlers {
   getProviderState: (
     req: JRPCRequest<unknown>,
   ) => Promise<{ accounts: string[]; chainId: string; isUnlocked: boolean }>;
+  setPreferredAggregator: (req: JRPCRequest<unknown>) => Promise<string>;
   submitBatch: (req: JRPCRequest<SendTransactionParams>) => Promise<string>;
 }
 
@@ -32,6 +33,7 @@ export function createWalletMiddleware({
   getAccounts,
   requestAccounts,
   getProviderState,
+  setPreferredAggregator,
   submitBatch,
 }: IProviderHandlers): JRPCMiddleware<string, unknown> {
   if (!getAccounts) {
@@ -75,6 +77,13 @@ export function createWalletMiddleware({
     res.result = await getProviderState(req);
   }
 
+  async function setPreferredAggregatorWrapper(
+    req: JRPCRequest<SendTransactionParams>,
+    res: JRPCResponse<unknown>,
+  ): Promise<void> {
+    res.result = await setPreferredAggregator(req);
+  }
+
   async function submitTransaction(
     req: JRPCRequest<SendTransactionParams>,
     res: JRPCResponse<unknown>,
@@ -91,6 +100,9 @@ export function createWalletMiddleware({
     [PROVIDER_JRPC_METHODS.GET_PROVIDER_STATE]: createAsyncMiddleware(
       getProviderStateFromController,
     ),
+    eth_setPreferredAggregator: createAsyncMiddleware<any, any>(
+      setPreferredAggregatorWrapper,
+    ),
     eth_sendTransaction: createAsyncMiddleware<any, any>(submitTransaction),
   });
 }
diff --git a/extension/source/Controllers/Network/createJsonRpcClient.ts b/extension/source/Controllers/Network/createJsonRpcClient.ts
index b8bf9a5e..df3f4308 100644
--- a/extension/source/Controllers/Network/createJsonRpcClient.ts
+++ b/extension/source/Controllers/Network/createJsonRpcClient.ts
@@ -7,7 +7,6 @@ import {
   mergeMiddleware,
 } from '@toruslabs/openlogin-jrpc';
 import { Aggregator } from 'bls-wallet-clients';
-import { AGGREGATOR_URL } from '../../env';
 
 import PollingBlockTracker from '../Block/PollingBlockTracker';
 import { ProviderConfig } from '../constants';
@@ -98,7 +97,7 @@ function createAggregatorMiddleware(): JRPCMiddleware<unknown, unknown> {
       if (hash in knownTransactions) {
         const knownTx = knownTransactions[hash];
 
-        const aggregator = new Aggregator(AGGREGATOR_URL);
+        const aggregator = new Aggregator(knownTx.aggregatorUrl);
         const bundleReceipt = await aggregator.lookupReceipt(hash);
 
         if (bundleReceipt === undefined) {
diff --git a/extension/source/Controllers/QuillController.ts b/extension/source/Controllers/QuillController.ts
index 6addd661..1e5ef32c 100644
--- a/extension/source/Controllers/QuillController.ts
+++ b/extension/source/Controllers/QuillController.ts
@@ -112,6 +112,10 @@ export default class QuillController extends BaseController<
 
   private preferencesController!: PreferencesController;
 
+  // This is just kept in memory because it supports setting the preferred
+  // aggregator for the particular tab only.
+  private tabPreferredAggregators: Record<number, string> = {};
+
   getRequestAccountTabIds: () => Record<string, number>;
   getOpenQuillTabsIds: () => Record<number, boolean>;
 
@@ -522,9 +526,16 @@ export default class QuillController extends BaseController<
         };
       },
 
+      setPreferredAggregator: async (req: any) => {
+        // eslint-disable-next-line prefer-destructuring
+        this.tabPreferredAggregators[req.tabId] = req.params[0];
+
+        return 'ok';
+      },
+
       submitBatch: async (req: any) => {
         const txParams = getAllReqParam<SendTransactionParams[]>(req);
-        const from = txParams[0].from;
+        const { from } = txParams[0];
 
         const actions = txParams.map((tx) => {
           return {
@@ -541,7 +552,9 @@ export default class QuillController extends BaseController<
         };
 
         const bundle = await this.keyringController.signTransactions(from, tx);
-        const agg = new Aggregator(AGGREGATOR_URL);
+        const aggregatorUrl =
+          this.tabPreferredAggregators[req.tabId] ?? AGGREGATOR_URL;
+        const agg = new Aggregator(aggregatorUrl);
         const result = await agg.add(bundle);
 
         if ('failures' in result) {
@@ -552,6 +565,7 @@ export default class QuillController extends BaseController<
           ...txParams[0],
           nonce: nonce.toString(),
           value: txParams[0].value || '0',
+          aggregatorUrl,
         };
 
         return result.hash;
diff --git a/extension/source/Controllers/knownTransactions.ts b/extension/source/Controllers/knownTransactions.ts
index 097c8fae..60cba918 100644
--- a/extension/source/Controllers/knownTransactions.ts
+++ b/extension/source/Controllers/knownTransactions.ts
@@ -6,6 +6,7 @@ const knownTransactions: Record<
   SendTransactionParams & {
     nonce: string;
     value: BigNumberish;
+    aggregatorUrl: string;
   }
 > = {};