fix: only process addresses from relevant protocols

Author: wa0x6e

src/utils.ts

@@ -122,10 +122,7 @@ function validateProtocols(protocols: Protocol[]): void {
  * 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.
+ * protocols.
  *
  * @param addresses - Array of blockchain addresses to format
  * @param protocols - Array of protocol names to validate against. Defaults to ['evm'].
@@ -140,20 +137,14 @@ export function getFormattedAddressesByProtocol(
   validateProtocols(protocols);
 
   return addresses.map((address) => {
-    if (protocols.includes('evm')) {
+    if (protocols.includes('evm') && address.length === 42) {
       try {
         return snapshot.utils.getFormattedAddress(address, 'evm');
-      } catch (e) {
-        // Continue to starknet if evm formatting fails and starknet is supported
-      }
-    }
-
-    if (protocols.includes('starknet')) {
+      } catch (e) {}
+    } else if (protocols.includes('starknet') && address.length === 66) {
       try {
         return snapshot.utils.getFormattedAddress(address, 'starknet');
-      } catch (e) {
-        // Address format not supported by any protocol
-      }
+      } catch (e) {}
     }
 
     throw new Error(
@@ -162,6 +153,46 @@ export function getFormattedAddressesByProtocol(
   });
 }
 
+/**
+ * Filters addresses by the specified protocols.
+ *
+ * This function takes a list of addresses and filters them based on the provided
+ * protocols. It checks if each address matches the expected format for EVM or
+ * Starknet protocols
+ *
+ * This function DOES NOT format/validate the addresses, and may return invalid addresses
+ *
+ * @param addresses - Array of addresses to filter
+ * @param protocols - Array of protocol names to validate against. Defaults to ['evm'].
+ *                   Valid protocols are 'evm' and 'starknet'.
+ *
+ * @returns Array of addresses that match the specified protocols
+ */
+export function getAddressesByProtocol(
+  addresses: string[],
+  protocols: Protocol[] = ['evm']
+): string[] {
+  validateProtocols(protocols);
+
+  return addresses.filter((address) => {
+    if (protocols.includes('evm') && /^0x[a-fA-F0-9]{40}$/.test(address)) {
+      return true;
+    }
+    if (protocols.includes('starknet') && /^0x[a-fA-F0-9]{64}$/.test(address)) {
+      return true;
+    }
+    return false;
+  });
+}
+
+export function validateAddress(address: string): void {
+  try {
+    getFormattedAddressesByProtocol([address], ['evm', 'starknet']);
+  } catch {
+    throw new Error('invalid address');
+  }
+}
+
 export const {
   multicall,
   Multicaller,
@@ -181,6 +212,7 @@ export default {
   getScoresDirect,
   customFetch,
   getFormattedAddressesByProtocol,
+  getAddressesByProtocol,
   multicall,
   Multicaller,
   subgraphRequest,

src/utils/vp.ts

@@ -1,16 +1,17 @@
 import { formatBytes32String } from '@ethersproject/strings';
-import { getAddress } from '@ethersproject/address';
+import { getAddress, isAddress } from '@ethersproject/address';
 import subgraphs from '@snapshot-labs/snapshot.js/src/delegationSubgraphs.json';
 import {
+  getAddressesByProtocol,
   getFormattedAddressesByProtocol,
   getProvider,
   getSnapshots,
   Multicaller,
-  subgraphRequest
+  subgraphRequest,
+  validateAddress
 } from '../utils';
 import _strategies from '../strategies';
 import { Score, Snapshot, VotingPower } from '../types';
-import { DEFAULT_SUPPORTED_PROTOCOLS } from '../constants';
 
 const DELEGATION_CONTRACT = '0x469788fE6E9E9681C6ebF3bF78e7Fd26Fc015446';
 const EMPTY_ADDRESS = '0x0000000000000000000000000000000000000000';
@@ -30,6 +31,12 @@ export async function getVp(
   space: string,
   delegation?: boolean
 ): Promise<VotingPower> {
+  if (strategies.some((s) => !_strategies[s.name])) {
+    throw new Error('invalid strategies');
+  }
+  // Ensure that the address is either a valid EVM address or Starknet address
+  validateAddress(address);
+
   const networks = [...new Set(strategies.map((s) => s.network || network))];
   const snapshots = await getSnapshots(
     network,
@@ -38,8 +45,10 @@ export async function getVp(
     networks
   );
 
+  // TODO: Delegation support for Starknet
+  const delegationSupported = delegation && isAddress(address);
   const delegations = {};
-  if (delegation) {
+  if (delegationSupported) {
     const ds = await Promise.all(
       networks.map((n) => getDelegations(address, n, snapshots[n], space))
     );
@@ -49,19 +58,20 @@ export async function getVp(
   const p: Score[] = strategies.map((strategy) => {
     const n = strategy.network || network;
     let addresses = [address];
+    const supportedProtocols = _strategies[strategy.name].supportedProtocols;
 
-    if (delegation) {
+    if (delegationSupported && supportedProtocols.includes('evm')) {
       addresses = delegations[n].in;
       if (!delegations[n].out) addresses.push(address);
       addresses = [...new Set(addresses)];
       if (addresses.length === 0) return {};
     }
 
-    addresses = getFormattedAddressesByProtocol(
-      addresses,
-      _strategies[strategy.name].supportedProtocols ??
-        DEFAULT_SUPPORTED_PROTOCOLS
-    );
+    addresses = getAddressesByProtocol(addresses, supportedProtocols);
+    addresses = getFormattedAddressesByProtocol(addresses, supportedProtocols);
+
+    if (addresses.length === 0) return {};
+
     return _strategies[strategy.name].strategy(
       space,
       n,
@@ -76,18 +86,17 @@ export async function getVp(
   const vpByStrategy = scores.map((score, i) => {
     const n = strategies[i].network || network;
     let addresses = [address];
+    const supportedProtocols =
+      _strategies[strategies[i].name].supportedProtocols;
 
-    if (delegation) {
+    if (delegationSupported && supportedProtocols.includes('evm')) {
       addresses = delegations[n].in;
       if (!delegations[n].out) addresses.push(address);
       addresses = [...new Set(addresses)];
     }
 
-    addresses = getFormattedAddressesByProtocol(
-      addresses,
-      _strategies[strategies[i].name].supportedProtocols ??
-        DEFAULT_SUPPORTED_PROTOCOLS
-    );
+    addresses = getAddressesByProtocol(addresses, supportedProtocols);
+    addresses = getFormattedAddressesByProtocol(addresses, supportedProtocols);
     return addresses.reduce((a, b) => a + (score[b] || 0), 0);
   });
   const vp = vpByStrategy.reduce((a, b) => a + b, 0);

test/__snapshots__/vp.test.ts.snap

@@ -0,0 +1,80 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP for EVM address on mixed protocol 1`] = `
+{
+  "vp": 10.998985610441185,
+  "vp_by_strategy": [
+    1,
+    9.998985610441185,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP for Starknet address on evm protocol 1`] = `
+{
+  "vp": 0,
+  "vp_by_strategy": [
+    0,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP for Starknet address on evm protocol with delegation enabled 1`] = `
+{
+  "vp": 0,
+  "vp_by_strategy": [
+    0,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP for Starknet address on mixed protocol 1`] = `
+{
+  "vp": 1,
+  "vp_by_strategy": [
+    1,
+    0,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP for Starknet address on mixed protocol with delegation enabled 1`] = `
+{
+  "vp": 1,
+  "vp_by_strategy": [
+    1,
+    0,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP with delegation enabled 1`] = `
+{
+  "vp": 10.077130335431244,
+  "vp_by_strategy": [
+    0,
+    10.028706352441185,
+    0,
+    0.04842398299005922,
+  ],
+  "vp_state": "final",
+}
+`;
+
+exports[`VP Calculation Integration Tests getVp() should calculate VP without delegation 1`] = `
+{
+  "vp": 21.55404462002206,
+  "vp_by_strategy": [
+    0,
+    9.998985610441185,
+    11.506635026590818,
+    0.04842398299005922,
+  ],
+  "vp_state": "final",
+}
+`;