From bfeb579bd6d1765d4a09f222b35ada8dfb1bce7a Mon Sep 17 00:00:00 2001
From: Seri Choi <choi.seri65@gmail.com>
Date: Wed, 12 Feb 2025 10:16:51 -0500
Subject: [PATCH] clean up

---
 pkg/contractManager/contractManager.go | 61 ++++++++++++++++++++++++++
 pkg/indexer/indexer.go                 | 14 +++---
 2 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/pkg/contractManager/contractManager.go b/pkg/contractManager/contractManager.go
index 9c007f12..0a770e1b 100644
--- a/pkg/contractManager/contractManager.go
+++ b/pkg/contractManager/contractManager.go
@@ -43,3 +43,64 @@ func (cm *ContractManager) GetContractWithProxy(
 
 	return contract, nil
 }
+
+func (cm *ContractManager) CreateProxyContract(
+	contractAddress string,
+	proxyContractAddress string,
+	blockNumber uint64,
+	reindexContract bool,
+) (*contractStore.ProxyContract, error) {
+	proxyContract, found, err := cm.ContractStore.FindOrCreateProxyContract(blockNumber, contractAddress, proxyContractAddress)
+	if err != nil {
+		cm.Logger.Sugar().Errorw("Failed to create proxy contract",
+			zap.Error(err),
+			zap.String("contractAddress", contractAddress),
+			zap.String("proxyContractAddress", proxyContractAddress),
+		)
+	} else {
+		if found {
+			cm.Logger.Sugar().Debugw("Found existing proxy contract",
+				zap.String("contractAddress", contractAddress),
+				zap.String("proxyContractAddress", proxyContractAddress),
+			)
+		} else {
+			cm.Logger.Sugar().Debugw("Created proxy contract",
+				zap.String("contractAddress", contractAddress),
+				zap.String("proxyContractAddress", proxyContractAddress),
+			)
+		}
+	}
+	// Check to see if the contract we're proxying to is already in the database
+	proxiedContract, err := cm.ContractStore.GetContractForAddress(proxyContractAddress)
+	if err != nil {
+		cm.Logger.Sugar().Errorw("Failed to get contract for address",
+			zap.Error(err),
+			zap.String("contractAddress", proxyContractAddress),
+		)
+	}
+	if proxiedContract != nil {
+		cm.Logger.Sugar().Debugw("Found proxied contract",
+			zap.String("contractAddress", proxyContractAddress),
+			zap.String("proxiedContractAddress", proxiedContract.ContractAddress),
+		)
+		if proxiedContract.ContractAbi == "" {
+			updatedContract := cm.FindAndSetContractAbi(proxyContractAddress)
+			if updatedContract != nil {
+				proxiedContract = updatedContract
+			}
+		}
+	} else {
+		_, err := cm.CreateContract(proxyContractAddress, "", reindexContract)
+		if err != nil {
+			cm.Logger.Sugar().Errorw("Failed to create contract",
+				zap.Error(err),
+				zap.String("contractAddress", proxyContractAddress),
+			)
+		} else {
+			cm.Logger.Sugar().Debugw("Created contract",
+				zap.String("contractAddress", proxyContractAddress),
+			)
+		}
+	}
+	return proxyContract, nil
+}
diff --git a/pkg/indexer/indexer.go b/pkg/indexer/indexer.go
index f2bd6ff8..e03c0891 100644
--- a/pkg/indexer/indexer.go
+++ b/pkg/indexer/indexer.go
@@ -162,7 +162,6 @@ func (idx *Indexer) IndexContractUpgrade(ctx context.Context, blockNumber uint64
 
 	// Check the arguments for the new address. EIP-1967 contracts include this as an argument.
 	// Otherwise, we'll check the storage slot
-	// TODO: Arguments is a string
 	for _, arg := range upgradedLog.Arguments {
 		if arg.Name == "implementation" && arg.Value != "" && arg.Value != nil {
 			newProxiedAddress = arg.Value.(common.Address).String()
@@ -178,26 +177,29 @@ func (idx *Indexer) IndexContractUpgrade(ctx context.Context, blockNumber uint64
 				zap.Uint64("block", blockNumber),
 				zap.String("upgradedLogAddress", upgradedLog.Address)
 			)
-		} else if len(storageValue) != 66 {
+			return err
+		if len(storageValue) != 66 {
 			idx.Logger.Sugar().Errorw("Invalid storage value",
 				zap.Uint64("block", blockNumber),
 				zap.String("storageValue", storageValue)
 			)
-		} else {
-			newProxiedAddress = storageValue[26:]
+			return err
 		}
+		
+		newProxiedAddress = storageValue[26:]
 	}
 
 	if newProxiedAddress == "" {
 		idx.Logger.Sugar().Debugw("No new proxied address found", zap.String("address", upgradedLog.Address))
-		return
+		return nil
 	}
 
 	_, err := idx.ContractManager.CreateProxyContract(upgradedLog.Address, newProxiedAddress, blockNumber, reindexContract)
 	if err != nil {
 		idx.Logger.Sugar().Errorw("Failed to create proxy contract", zap.Error(err))
-		return
+		return err
 	}
+
 	idx.Logger.Sugar().Infow("Upgraded proxy contract", zap.String("contractAddress", upgradedLog.Address), zap.String("proxyContractAddress", newProxiedAddress))
 	return nil
 }