Skip to content

Commit 937db18

Browse files
Mo-HussainMo-Hussain
committed
refactor(rebalancer): replace positional private key args with protocol-keyed map in execute()
Address PR review feedback: - IExternalBridge.execute() now takes Partial<Record<ProtocolType, string>> instead of positional privateKey/solanaPrivateKey args - LiFiBridge: extract provider setup into configureLiFiProviders() helper, throw on unsupported protocol - LiFiBridge: use constructor-built Map<number, ChainMetadata> for O(1) chainId lookups - InventoryRebalancer: build privateKeys map from all inventory signers at call site - Remove unused getInventorySignerKey method - Update all test call sites and MockExternalBridge signature
1 parent 9370668 commit 937db18

5 files changed

Lines changed: 153 additions & 86 deletions

File tree

‎typescript/rebalancer/src/bridges/LiFiBridge.test.ts‎

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { expect } from 'chai';
22
import { pino } from 'pino';
33

44
import type { LiFiStep } from '@lifi/sdk';
5+
import { ProtocolType } from '@hyperlane-xyz/utils';
56

67
import type {
78
BridgeQuote,
@@ -174,7 +175,9 @@ describe('LiFiBridge.execute() route validation', function () {
174175
const quote = createTestQuote();
175176

176177
try {
177-
await bridge.execute(quote, TEST_PRIVATE_KEY);
178+
await bridge.execute(quote, {
179+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
180+
});
178181
// If it resolves, validation definitely passed
179182
} catch (error: unknown) {
180183
const msg = (error as Error).message;
@@ -190,7 +193,9 @@ describe('LiFiBridge.execute() route validation', function () {
190193
const quote = createTestQuote({ fromChainId: 999 }, { fromChain: 42161 });
191194

192195
try {
193-
await bridge.execute(quote, TEST_PRIVATE_KEY);
196+
await bridge.execute(quote, {
197+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
198+
});
194199
expect.fail('Expected execute to throw');
195200
} catch (error: unknown) {
196201
const msg = (error as Error).message;
@@ -204,7 +209,9 @@ describe('LiFiBridge.execute() route validation', function () {
204209
const quote = createTestQuote({ toChainId: 888 }, { toChain: 1399811149 });
205210

206211
try {
207-
await bridge.execute(quote, TEST_PRIVATE_KEY);
212+
await bridge.execute(quote, {
213+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
214+
});
208215
expect.fail('Expected execute to throw');
209216
} catch (error: unknown) {
210217
const msg = (error as Error).message;
@@ -221,7 +228,9 @@ describe('LiFiBridge.execute() route validation', function () {
221228
);
222229

223230
try {
224-
await bridge.execute(quote, TEST_PRIVATE_KEY);
231+
await bridge.execute(quote, {
232+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
233+
});
225234
expect.fail('Expected execute to throw');
226235
} catch (error: unknown) {
227236
const msg = (error as Error).message;
@@ -238,7 +247,9 @@ describe('LiFiBridge.execute() route validation', function () {
238247
);
239248

240249
try {
241-
await bridge.execute(quote, TEST_PRIVATE_KEY);
250+
await bridge.execute(quote, {
251+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
252+
});
242253
expect.fail('Expected execute to throw');
243254
} catch (error: unknown) {
244255
const msg = (error as Error).message;
@@ -254,7 +265,9 @@ describe('LiFiBridge.execute() route validation', function () {
254265
);
255266

256267
try {
257-
await bridge.execute(quote, TEST_PRIVATE_KEY);
268+
await bridge.execute(quote, {
269+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
270+
});
258271
expect.fail('Expected execute to throw');
259272
} catch (error: unknown) {
260273
const msg = (error as Error).message;
@@ -269,7 +282,9 @@ describe('LiFiBridge.execute() route validation', function () {
269282
);
270283

271284
try {
272-
await bridge.execute(quote, TEST_PRIVATE_KEY);
285+
await bridge.execute(quote, {
286+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
287+
});
273288
expect.fail('Expected execute to throw');
274289
} catch (error: unknown) {
275290
const msg = (error as Error).message;
@@ -287,7 +302,9 @@ describe('LiFiBridge.execute() route validation', function () {
287302
);
288303

289304
try {
290-
await bridge.execute(quote, TEST_PRIVATE_KEY);
305+
await bridge.execute(quote, {
306+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
307+
});
291308
} catch (error: unknown) {
292309
const msg = (error as Error).message;
293310
expect(
@@ -304,7 +321,9 @@ describe('LiFiBridge.execute() route validation', function () {
304321
);
305322

306323
try {
307-
await bridge.execute(quote, TEST_PRIVATE_KEY);
324+
await bridge.execute(quote, {
325+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
326+
});
308327
expect.fail('Expected execute to throw');
309328
} catch (error: unknown) {
310329
const msg = (error as Error).message;
@@ -323,7 +342,9 @@ describe('LiFiBridge.execute() route validation', function () {
323342
);
324343

325344
try {
326-
await bridge.execute(quote, TEST_PRIVATE_KEY);
345+
await bridge.execute(quote, {
346+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
347+
});
327348
expect.fail('Expected execute to throw');
328349
} catch (error: unknown) {
329350
const msg = (error as Error).message;
@@ -355,7 +376,9 @@ describe('LiFiBridge.execute() route validation', function () {
355376
);
356377

357378
try {
358-
await bridge.execute(quote, TEST_PRIVATE_KEY);
379+
await bridge.execute(quote, {
380+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
381+
});
359382
} catch (error: unknown) {
360383
const msg = (error as Error).message;
361384
expect(
@@ -371,7 +394,9 @@ describe('LiFiBridge.execute() route validation', function () {
371394
// With `!== undefined`, a 0n request is correctly compared against the route amount.
372395
const quote = createTestQuote({ fromAmount: '1' }, { fromAmount: 0n });
373396
try {
374-
await bridge.execute(quote, TEST_PRIVATE_KEY);
397+
await bridge.execute(quote, {
398+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
399+
});
375400
expect.fail('Expected execute to throw');
376401
} catch (error: unknown) {
377402
const msg = (error as Error).message;
@@ -388,7 +413,9 @@ describe('LiFiBridge.execute() route validation', function () {
388413
{ fromAmount: undefined, toAmount: 5000000000n },
389414
);
390415
try {
391-
await bridge.execute(quote, TEST_PRIVATE_KEY);
416+
await bridge.execute(quote, {
417+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
418+
});
392419
} catch (error: unknown) {
393420
const msg = (error as Error).message;
394421
expect(
@@ -405,7 +432,9 @@ describe('LiFiBridge.execute() route validation', function () {
405432
{ fromAmount: undefined, toAmount: 123n },
406433
);
407434
try {
408-
await bridge.execute(quote, TEST_PRIVATE_KEY);
435+
await bridge.execute(quote, {
436+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
437+
});
409438
expect.fail('Expected execute to throw');
410439
} catch (error: unknown) {
411440
const msg = (error as Error).message;
@@ -424,7 +453,9 @@ describe('LiFiBridge.execute() route validation', function () {
424453
{ fromAmount: undefined, toAmount: 1000000n },
425454
);
426455
try {
427-
await bridge.execute(quote, TEST_PRIVATE_KEY);
456+
await bridge.execute(quote, {
457+
[ProtocolType.Ethereum]: TEST_PRIVATE_KEY,
458+
});
428459
} catch (error: unknown) {
429460
const msg = (error as Error).message;
430461
expect(

‎typescript/rebalancer/src/bridges/LiFiBridge.ts‎

Lines changed: 85 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
getStatus,
1313
config as lifiConfig,
1414
} from '@lifi/sdk';
15+
import type { ChainMetadata } from '@hyperlane-xyz/sdk';
1516
import { ProtocolType, ensure0x } from '@hyperlane-xyz/utils';
1617
import type { Logger } from 'pino';
1718
import { type Chain, createWalletClient, http } from 'viem';
@@ -88,10 +89,20 @@ export class LiFiBridge implements IExternalBridge {
8889
readonly logger: Logger;
8990
private initialized = false;
9091
private readonly config: ExternalBridgeConfig;
92+
private readonly chainMetadataByChainId: Map<number, ChainMetadata>;
9193

9294
constructor(config: ExternalBridgeConfig, logger: Logger) {
9395
this.config = config;
9496
this.logger = logger;
97+
// Build chainId -> metadata map for O(1) lookups
98+
this.chainMetadataByChainId = new Map();
99+
if (config.chainMetadata) {
100+
for (const metadata of Object.values(config.chainMetadata)) {
101+
if (metadata.chainId !== undefined) {
102+
this.chainMetadataByChainId.set(Number(metadata.chainId), metadata);
103+
}
104+
}
105+
}
95106
}
96107

97108
getNativeTokenAddress(): string {
@@ -118,23 +129,76 @@ export class LiFiBridge implements IExternalBridge {
118129
* Iterates metadata to find matching chainId and returns first HTTP RPC URL.
119130
*/
120131
private getRpcUrlForChainId(chainId: number): string | undefined {
121-
if (!this.config.chainMetadata) return undefined;
122-
for (const metadata of Object.values(this.config.chainMetadata)) {
123-
if (metadata.chainId === chainId && metadata.rpcUrls?.length) {
124-
return metadata.rpcUrls[0].http;
125-
}
126-
}
127-
return undefined;
132+
return this.chainMetadataByChainId.get(chainId)?.rpcUrls?.[0]?.http;
128133
}
129134

130135
private getProtocolTypeForChainId(chainId: number): ProtocolType | undefined {
131-
if (!this.config.chainMetadata) return undefined;
132-
for (const metadata of Object.values(this.config.chainMetadata)) {
133-
if (metadata.chainId === chainId) {
134-
return metadata.protocol;
136+
return this.chainMetadataByChainId.get(chainId)?.protocol;
137+
}
138+
139+
/**
140+
* Configure LiFi SDK providers from the given private keys.
141+
* Sets up wallet/signer for each protocol type present in the keys map.
142+
*/
143+
private configureLiFiProviders(
144+
privateKeys: Partial<Record<ProtocolType, string>>,
145+
fromChain: number,
146+
fromRpcUrl: string | undefined,
147+
): void {
148+
const providers: Parameters<typeof lifiConfig.setProviders>[0] = [];
149+
for (const [protocol, key] of Object.entries(privateKeys)) {
150+
switch (protocol) {
151+
case ProtocolType.Ethereum: {
152+
const account = privateKeyToAccount(ensure0x(key) as `0x${string}`);
153+
const chain = getViemChain(fromChain, fromRpcUrl);
154+
const walletClient = createWalletClient({
155+
account,
156+
chain,
157+
transport: http(fromRpcUrl),
158+
});
159+
providers.push(
160+
EVM({
161+
getWalletClient: async () => walletClient,
162+
switchChain: async (requiredChainId: number) => {
163+
const switchRpcUrl = this.getRpcUrlForChainId(requiredChainId);
164+
const requiredChain = getViemChain(
165+
requiredChainId,
166+
switchRpcUrl,
167+
);
168+
return createWalletClient({
169+
account,
170+
chain: requiredChain,
171+
transport: http(switchRpcUrl),
172+
});
173+
},
174+
}),
175+
);
176+
break;
177+
}
178+
case ProtocolType.Sealevel: {
179+
providers.push(
180+
Solana({
181+
getWalletAdapter: async () => new KeypairWalletAdapter(key),
182+
}),
183+
);
184+
break;
185+
}
186+
default:
187+
throw new Error(
188+
`Unsupported protocol type '${protocol}' for LiFi provider`,
189+
);
135190
}
136191
}
137-
return undefined;
192+
193+
lifiConfig.setProviders(providers);
194+
195+
this.logger.debug(
196+
{
197+
fromChain,
198+
protocols: Object.keys(privateKeys),
199+
},
200+
'Configured LiFi providers for route execution',
201+
);
138202
}
139203

140204
/**
@@ -337,12 +401,11 @@ export class LiFiBridge implements IExternalBridge {
337401
* Handles approvals, transaction signing, and execution automatically.
338402
*
339403
* @param quote - Quote obtained from quote()
340-
* @param privateKey - Private key hex string (0x-prefixed) for signing the transaction
404+
* @param privateKeys - Private keys keyed by protocol type for signing transactions
341405
*/
342406
async execute(
343407
quote: BridgeQuote<LiFiStep>,
344-
privateKey: string,
345-
solanaPrivateKey?: string,
408+
privateKeys: Partial<Record<ProtocolType, string>>,
346409
): Promise<BridgeTransferResult> {
347410
this.initialize();
348411

@@ -354,6 +417,10 @@ export class LiFiBridge implements IExternalBridge {
354417
const fromChain = route.fromChainId;
355418
const toChain = route.toChainId;
356419
const fromProtocol = this.getProtocolTypeForChainId(fromChain);
420+
assert(
421+
privateKeys[fromProtocol ?? ProtocolType.Ethereum],
422+
`Missing private key for source chain protocol ${fromProtocol ?? ProtocolType.Ethereum}`,
423+
);
357424

358425
this.logger.info(
359426
{
@@ -366,47 +433,14 @@ export class LiFiBridge implements IExternalBridge {
366433
'Executing LiFi bridge transfer',
367434
);
368435

369-
const account = privateKeyToAccount(ensure0x(privateKey) as `0x${string}`);
370-
const rpcUrl = this.getRpcUrlForChainId(fromChain);
371-
const chain = getViemChain(fromChain, rpcUrl);
372-
const walletClient = createWalletClient({
373-
account,
374-
chain,
375-
transport: http(rpcUrl),
376-
});
377-
const evmProvider = EVM({
378-
getWalletClient: async () => walletClient,
379-
switchChain: async (requiredChainId: number) => {
380-
const switchRpcUrl = this.getRpcUrlForChainId(requiredChainId);
381-
const requiredChain = getViemChain(requiredChainId, switchRpcUrl);
382-
return createWalletClient({
383-
account,
384-
chain: requiredChain,
385-
transport: http(switchRpcUrl),
386-
});
387-
},
388-
});
389-
390-
const providers: Parameters<typeof lifiConfig.setProviders>[0] = [
391-
evmProvider,
392-
];
393-
if (solanaPrivateKey) {
394-
providers.push(
395-
Solana({
396-
getWalletAdapter: async () =>
397-
new KeypairWalletAdapter(solanaPrivateKey),
398-
}),
399-
);
400-
}
436+
const fromRpcUrl = this.getRpcUrlForChainId(fromChain);
401437

402-
lifiConfig.setProviders(providers);
438+
this.configureLiFiProviders(privateKeys, fromChain, fromRpcUrl);
403439

404440
this.logger.debug(
405441
{
406442
fromChain,
407-
protocol: fromProtocol ?? ProtocolType.Ethereum,
408-
hasEvmProvider: true,
409-
hasSolanaProvider: !!solanaPrivateKey,
443+
protocols: Object.keys(privateKeys),
410444
},
411445
'Configured LiFi providers for route execution',
412446
);

0 commit comments

Comments
 (0)