fix: add validations to getVp and getScoresDirect

src/utils.ts

@@ -4,7 +4,6 @@ import snapshot from '@snapshot-labs/snapshot.js';
 import { getDelegations } from './utils/delegation';
 import { createHash } from 'crypto';
 import { Protocol, Score, Snapshot, VotingPower } from './types';
-import { VALID_PROTOCOLS } from './constants';
 
 export function sha256(str) {
   return createHash('sha256').update(str).digest('hex');
@@ -24,11 +23,6 @@ async function callStrategy(
   ) {
     return {};
   }
-
-  if (!_strategies.hasOwnProperty(strategy.name)) {
-    throw new Error(`Invalid strategy: ${strategy.name}`);
-  }
-
   const score: Score = await _strategies[strategy.name].strategy(
     space,
     network,
@@ -37,12 +31,19 @@ async function callStrategy(
     strategy.params,
     snapshot
   );
-  const addressesLc = addresses.map((address) => address.toLowerCase());
-  return Object.fromEntries(
-    Object.entries(score).filter(
-      ([address, vp]) => vp > 0 && addressesLc.includes(address.toLowerCase())
-    )
+
+  const normalizedAddresses = new Set(
+    addresses.map((address) => address.toLowerCase())
   );
+  const filteredScore: Score = {};
+
+  for (const [address, vp] of Object.entries(score)) {
+    if (vp > 0 && normalizedAddresses.has(address.toLowerCase())) {
+      filteredScore[address] = vp;
+    }
+  }
+
+  return filteredScore;
 }
 
 export async function getScoresDirect(
@@ -54,16 +55,20 @@ export async function getScoresDirect(
   snapshot: Snapshot
 ): Promise<Score[]> {
   try {
-    const networks = strategies.map((s) => s.network || network);
-    const snapshots = await getSnapshots(network, snapshot, provider, networks);
-    // @ts-ignore
     if (addresses.length === 0) return strategies.map(() => ({}));
+
+    const addressesByProtocol = categorizeAddressesByProtocol(addresses);
+    validateStrategies(strategies);
+
+    const networks = [...new Set(strategies.map((s) => s.network || network))];
+    const snapshots = await getSnapshots(network, snapshot, provider, networks);
+
     return await Promise.all(
       strategies.map((strategy) =>
         callStrategy(
           space,
           strategy.network || network,
-          addresses,
+          filterAddressesForStrategy(addressesByProtocol, strategy.name),
           strategy,
           snapshots[strategy.network || network]
         )
@@ -81,7 +86,11 @@ export async function getVp(
   snapshot: Snapshot,
   space: string
 ): Promise<VotingPower> {
-  const scores: Score[] = await getScoresDirect(
+  if (!strategies.length) {
+    throw new Error('no strategies provided');
+  }
+
+  const scores = await getScoresDirect(
     space,
     strategies,
     network,
@@ -90,7 +99,13 @@ export async function getVp(
     snapshot
   );
 
-  const vpByStrategy = scores.map((score) => score[address] || 0);
+  const normalizedAddress = address.toLowerCase();
+  const vpByStrategy = scores.map((score) => {
+    const matchingKey = Object.keys(score).find(
+      (key) => key.toLowerCase() === normalizedAddress
+    );
+    return matchingKey ? score[matchingKey] : 0;
+  });
 
   return {
     vp: vpByStrategy.reduce((a, b) => a + b, 0),
@@ -124,66 +139,59 @@ export function customFetch(
   ]);
 }
 
-/**
- * Validates that protocols are non-empty and contain only valid protocol names.
- *
- * @param protocols - Array of protocol names to validate
- */
-function validateProtocols(protocols: Protocol[]): void {
-  if (!protocols.length) {
-    throw new Error('At least one protocol must be specified');
-  }
-
-  const invalidProtocols = protocols.filter(
-    (p) => !VALID_PROTOCOLS.includes(p)
-  );
-  if (invalidProtocols.length > 0) {
-    throw new Error(`Invalid protocol(s): ${invalidProtocols.join(', ')}`);
-  }
+function detectProtocol(address: string): Protocol | null {
+  if (/^0x[a-fA-F0-9]{40}$/.test(address)) return 'evm';
+  if (/^0x[a-fA-F0-9]{64}$/.test(address)) return 'starknet';
+  return null;
 }
 
-/**
- * Formats addresses according to the specified blockchain protocols.
- *
- * This function takes a list of addresses and formats them according to the provided
- * protocols. It prioritizes EVM formatting when multiple protocols are specified and
- * an address is valid for both. If EVM formatting fails but Starknet is supported,
- * it falls back to Starknet formatting. Throws an error if any address cannot be
- * formatted according to the specified protocols.
- *
- * @param addresses - Array of blockchain addresses to format
- * @param protocols - Array of protocol names to validate against. Defaults to ['evm'].
- *                   Valid protocols are 'evm' and 'starknet'.
- *
- * @returns Array of formatted addresses in the same order as input
- */
-export function getFormattedAddressesByProtocol(
-  addresses: string[],
-  protocols: Protocol[] = ['evm']
-): string[] {
-  validateProtocols(protocols);
-
-  return addresses.map((address) => {
-    if (protocols.includes('evm')) {
-      try {
-        return snapshot.utils.getFormattedAddress(address, 'evm');
-      } catch (e) {
-        // Continue to starknet if evm formatting fails and starknet is supported
-      }
+function categorizeAddressesByProtocol(
+  addresses: string[]
+): Record<Protocol, string[]> {
+  const results: Record<Protocol, string[]> = {
+    evm: [],
+    starknet: []
+  };
+
+  for (const address of addresses) {
+    const addressType = detectProtocol(address);
+    if (!addressType) {
+      throw new Error(`Invalid address format: ${address}`);
     }
 
-    if (protocols.includes('starknet')) {
-      try {
-        return snapshot.utils.getFormattedAddress(address, 'starknet');
-      } catch (e) {
-        // Address format not supported by any protocol
+    try {
+      const formattedAddress = snapshot.utils.getFormattedAddress(
+        address,
+        addressType
+      );
+      if (!results[addressType].includes(formattedAddress)) {
+        results[addressType].push(formattedAddress);
       }
+    } catch {
+      throw new Error(`Invalid ${addressType} address: ${address}`);
     }
+  }
 
-    throw new Error(
-      `Address "${address}" is not a valid ${protocols.join(' or ')} address`
-    );
-  });
+  return results;
+}
+
+function filterAddressesForStrategy(
+  addressesByProtocol: Record<Protocol, string[]>,
+  strategyName: string
+): string[] {
+  return _strategies[strategyName].supportedProtocols.flatMap(
+    (protocol: Protocol) => addressesByProtocol[protocol] || []
+  );
+}
+
+function validateStrategies(strategies: any[]): void {
+  const invalidStrategies = strategies
+    .filter((strategy) => !_strategies[strategy.name])
+    .map((strategy) => strategy.name);
+
+  if (invalidStrategies.length > 0) {
+    throw new Error(`Invalid strategies: ${invalidStrategies.join(', ')}`);
+  }
 }
 
 export const {
@@ -204,7 +212,6 @@ export default {
   sha256,
   getScoresDirect,
   customFetch,
-  getFormattedAddressesByProtocol,
   multicall,
   Multicaller,
   subgraphRequest,

test/integration/__snapshots__/utils.test.ts.snap

@@ -1,6 +1,6 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`utils getVp 1`] = `
+exports[`utils getVp() should calculate VP for EVM address on evm protocol 1`] = `
 {
   "vp": 21.55404462002206,
   "vp_by_strategy": [
@@ -12,3 +12,35 @@ exports[`utils getVp 1`] = `
   "vp_state": "final",
 }
 `;
+
+exports[`utils getVp() should calculate VP for EVM address on mixed protocol 1`] = `
+{
+  "vp": 10.998985610441185,
+  "vp_by_strategy": [
+    1,
+    9.998985610441185,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`utils getVp() should calculate VP for Starknet address on evm protocol 1`] = `
+{
+  "vp": 0,
+  "vp_by_strategy": [
+    0,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`utils getVp() should calculate VP for Starknet address on mixed protocol 1`] = `
+{
+  "vp": 1,
+  "vp_by_strategy": [
+    1,
+    0,
+  ],
+  "vp_state": "final",
+}
+`;

test/integration/fixtures/vp-fixtures.ts

@@ -0,0 +1,84 @@
+export const testConfig = {
+  network: '1',
+  snapshot: 15354134,
+  space: 'cvx.eth',
+  evmAddress: '0xeF8305E140ac520225DAf050e2f71d5fBcC543e7',
+  starknetAddress:
+    '0x07f71118e351c02f6EC7099C8CDf93AED66CEd8406E94631cC91637f7D7F203A'
+};
+
+export const strategies = {
+  withDelegation: [
+    {
+      name: 'erc20-balance-of',
+      params: {
+        symbol: 'CVX',
+        address: '0x72a19342e8F1838460eBFCCEf09F6585e32db86E',
+        decimals: 18
+      }
+    },
+    {
+      name: 'eth-balance',
+      network: '100',
+      params: {}
+    },
+    {
+      name: 'eth-balance',
+      network: '1',
+      params: {}
+    },
+    {
+      name: 'eth-balance',
+      network: '10',
+      params: {}
+    }
+  ],
+  mixed: [
+    {
+      name: 'whitelist',
+      params: {
+        addresses: [testConfig.evmAddress, testConfig.starknetAddress]
+      }
+    },
+    {
+      name: 'eth-balance',
+      network: '100',
+      params: {}
+    }
+  ],
+  evmOnly: [
+    {
+      name: 'eth-balance',
+      network: '100',
+      params: {}
+    }
+  ],
+  singleInvalid: [
+    {
+      name: 'whitelist-invalid',
+      params: {
+        addresses: [testConfig.evmAddress, testConfig.starknetAddress]
+      }
+    },
+    {
+      name: 'eth-balance',
+      network: '100',
+      params: {}
+    }
+  ],
+  multipleInvalid: [
+    {
+      name: 'strategy-one-invalid',
+      params: {}
+    },
+    {
+      name: 'strategy-two-invalid',
+      params: {}
+    },
+    {
+      name: 'eth-balance',
+      network: '100',
+      params: {}
+    }
+  ]
+};