From cf8994bb4df1de83dde5820614b3e8e7f8adc4b6 Mon Sep 17 00:00:00 2001
From: trunghai95 <bthai.1995@gmail.com>
Date: Mon, 22 Jul 2024 16:35:04 +0800
Subject: [PATCH] Always use L1 token address for bridge balance

---
 server/service_xlayer.go         | 1 +
 utils/messagebridge/processor.go | 8 +++-----
 utils/messagebridge/usdclxly.go  | 2 ++
 utils/messagebridge/wsteth.go    | 2 ++
 4 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/server/service_xlayer.go b/server/service_xlayer.go
index 30f4c170..af2f8bb2 100644
--- a/server/service_xlayer.go
+++ b/server/service_xlayer.go
@@ -579,6 +579,7 @@ func (s *bridgeService) GetWstEthTokenNotWithdrawn(ctx context.Context, req *pb.
 			Msg:  "internal: wstETH processor is not inited",
 		}, nil
 	}
+	// Always use L1 token address for storage
 	tokenAddr := processor.GetTokenAddressList()[0]
 	valueL1, errL1 := s.storage.GetBridgeBalance(ctx, tokenAddr, utils.GetMainNetworkId(), false, nil)
 	valueL2, errL2 := s.storage.GetBridgeBalance(ctx, tokenAddr, utils.GetRollupNetworkId(), false, nil)
diff --git a/utils/messagebridge/processor.go b/utils/messagebridge/processor.go
index 237cdd67..be758b49 100644
--- a/utils/messagebridge/processor.go
+++ b/utils/messagebridge/processor.go
@@ -25,6 +25,8 @@ var (
 // Current supported tokens: USDC, wstETH
 type Processor struct {
 	contractToTokenMapping map[common.Address]common.Address
+	contractAddressList    []common.Address
+	tokenAddressList       []common.Address
 	// DecodeMetadata decodes the metadata of the message bridge, returns the actual destination address and bridged amount
 	DecodeMetadataFn func(metadata []byte) (common.Address, *big.Int)
 }
@@ -40,11 +42,7 @@ func (u *Processor) GetContractAddressList() []common.Address {
 
 // GetTokenAddressList returns the list of original token addresses
 func (u *Processor) GetTokenAddressList() []common.Address {
-	result := make([]common.Address, 0)
-	for _, addr := range u.contractToTokenMapping {
-		result = append(result, addr)
-	}
-	return result
+	return u.tokenAddressList
 }
 
 // CheckContractAddress returns true if the input address is in the contract address list of this bridge
diff --git a/utils/messagebridge/usdclxly.go b/utils/messagebridge/usdclxly.go
index 92c0ebf5..bd3834db 100644
--- a/utils/messagebridge/usdclxly.go
+++ b/utils/messagebridge/usdclxly.go
@@ -25,6 +25,8 @@ func InitUSDCLxLyProcessor(usdcContractAddresses, usdcTokenAddresses []common.Ad
 	if len(contractToTokenMapping) > 0 {
 		processorMap[USDC] = &Processor{
 			contractToTokenMapping: contractToTokenMapping,
+			contractAddressList:    usdcContractAddresses,
+			tokenAddressList:       usdcTokenAddresses,
 			DecodeMetadataFn: func(metadata []byte) (common.Address, *big.Int) {
 				// Metadata structure:
 				// - Destination address: 32 bytes
diff --git a/utils/messagebridge/wsteth.go b/utils/messagebridge/wsteth.go
index 7d8f27af..2a27478b 100644
--- a/utils/messagebridge/wsteth.go
+++ b/utils/messagebridge/wsteth.go
@@ -25,6 +25,8 @@ func InitWstETHProcessor(wstETHContractAddresses, wstETHTokenAddresses []common.
 	if len(contractToTokenMapping) > 0 {
 		processorMap[WstETH] = &Processor{
 			contractToTokenMapping: contractToTokenMapping,
+			contractAddressList:    wstETHContractAddresses,
+			tokenAddressList:       wstETHTokenAddresses,
 			DecodeMetadataFn: func(metadata []byte) (common.Address, *big.Int) {
 				// Metadata structure:
 				// - Destination address: 32 bytes