fix: custom ABI loading for proxy contracts

credmark · Jan 16, 2024 · 96f03d5 · 96f03d5
1 parent d9fb418
commit 96f03d5
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 20 deletions.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -3,6 +3,9 @@
 History
 -------
 
+0.8.139 [2024-01-16]
+* fix: custom ABI loading for proxy contracts
+
 0.8.138 [2023-12-28]
 * fix: event contract ledger for proxy contracts
 

diff --git a/credmark/cmf/types/contract.py b/credmark/cmf/types/contract.py
@@ -1,14 +1,14 @@
 # pylint: disable=line-too-long
 
-import json
 from typing import List, Optional, Sequence, Union, cast
 
 from web3 import Web3
 from web3.contract.contract import Contract as Web3Contract
 
 import credmark.cmf.model
 from credmark.cmf.engine.cache import ContractMetaCache
-from credmark.cmf.model.errors import ModelDataError, ModelEngineError, ModelRunError
+from credmark.cmf.model.errors import (ModelDataError, ModelEngineError,
+                                       ModelRunError)
 from credmark.dto import DTO, DTOField, IterableListGenericDTO, PrivateAttr
 
 from .abi import ABI
@@ -152,17 +152,17 @@ class ContractMetaData(DTO):
     _meta: ContractMetaData = PrivateAttr(
         default_factory=lambda: Contract.ContractMetaData())  # pylint: disable=unnecessary-lambda
     _instance: Web3Contract | None = PrivateAttr(default=None)
-    _proxy_implementation = PrivateAttr(default=None)
-    _loaded = PrivateAttr(default=False)
-    _ledger = PrivateAttr(default=None)
+    _loaded: bool = PrivateAttr(default=False)
+    _ledger: Optional[ContractLedger] = PrivateAttr(default=None)
+    _custom_abi: Optional[ABI] = PrivateAttr(default=None)
 
     class Config:
         arbitrary_types_allowed = True
         underscore_attrs_are_private = True
         schema_extra = {
             'examples': Account.Config.schema_extra['examples'] +
             [{'address': '0x1F98431c8aD98523631AE4a59f267346ea31F984',
-              'abi': '(Optional) contract abi JSON string'
+              'abi': '(Optional) contract abi JSON string or list'
               }]
         }
 
@@ -176,10 +176,10 @@ def __init__(self, *args, **data):
             if isinstance(meta, type(self._meta)):
                 self._meta = meta
 
-        if isinstance(data.get('abi'), str):
-            self._meta.abi = json.loads(data['abi'])
+        if isinstance(data.get('abi'), (str, list)):
+            self._meta.abi = ABI(data['abi'])
+            self._custom_abi = self._meta.abi
         self._instance = None
-        self._proxy_implementation = None
         self._ledger = None
 
     def _load(self):
@@ -223,23 +223,26 @@ def _load(self):
             self._meta.abi = ABI(res.get('abi'))
             self._meta.is_transparent_proxy = res.get('proxy', 0) == "1"
 
+            implementation_address = None
             if self._meta.contract_name in ['BeaconProxy', 'BridgeToken'] and self._meta.is_transparent_proxy:
                 # TODO: Special case for BeaconProxy, proxy address may not up to date
-                proxy_address = res.get('implementation')
-                self._meta.proxy_implementation = Contract(
-                    address=proxy_address)
+                implementation_address = res.get('implementation')
             else:
                 slot_proxy_address = get_slot_proxy_address(
                     context, self.address, self._meta.contract_name, self._meta.abi)
 
                 if slot_proxy_address is not None and not slot_proxy_address.is_null():
                     # TODO: as we only store the latest implementation in DB but not for the history.
-                    self._meta.proxy_implementation = Contract(
-                        address=slot_proxy_address)
+                    implementation_address = slot_proxy_address
                 elif self._meta.is_transparent_proxy:
-                    proxy_address = res.get('implementation')
-                    self._meta.proxy_implementation = Contract(
-                        address=proxy_address)
+                    implementation_address = res.get('implementation')
+
+            if implementation_address is not None:
+                self._meta.proxy_implementation = Contract(
+                    address=implementation_address,
+                    # If a custom ABI is provided, we pass it on to
+                    # the implementation contract to use as fallback
+                    abi=self._custom_abi)
             self._loaded = True
         else:
             if self._meta.abi is None:

diff --git a/credmark/cmf/types/token_erc20.py b/credmark/cmf/types/token_erc20.py
@@ -179,15 +179,15 @@ def __init__(self, *args, **data):
             raise ModelDataError(
                 f'NULL address ({Address.null()}) is not a valid Token Address')
 
+        if 'abi' not in data:
+            data['abi'] = ERC20_BASE_ABI
+
         super().__init__(**data)
 
     def _load(self):
         if self._loaded:
             return
 
-        if self._meta.abi is None:
-            self._meta.abi = ABI(ERC20_BASE_ABI)
-
         super()._load()
 
     def as_erc20(self, set_loaded=False, use_alt=False):