diff --git a/nodebalancer/cmd/nodebalancer/balances.go b/nodebalancer/cmd/nodebalancer/balances.go index d50823b1..1161343f 100644 --- a/nodebalancer/cmd/nodebalancer/balances.go +++ b/nodebalancer/cmd/nodebalancer/balances.go @@ -124,13 +124,6 @@ func getBalances(ctx context.Context, address string) (BalancesResponse, error) return } - // Get token list for this blockchain - tokens := getTokenList(blockchain) - if len(tokens) == 0 { - resultChan <- chainResult{blockchain: blockchain, balances: make(ChainBalances)} - return - } - // Connect to client client, err := ethclient.Dial(node.Endpoint.String()) if err != nil { @@ -139,14 +132,34 @@ func getBalances(ctx context.Context, address string) (BalancesResponse, error) } defer client.Close() - // Try Multicall3 first - chainBalances, err := getBalancesMulticall(chainCtx, client, tokens, checksumAddress, blockchain) + // Initialize chain balances + chainBalances := make(ChainBalances) + + // Get native token balance first + nativeBalance, err := client.BalanceAt(chainCtx, common.HexToAddress(checksumAddress), nil) if err != nil { - // Fallback to individual calls - chainBalances, err = getBalancesFallback(chainCtx, client, tokens, checksumAddress) + log.Printf("Failed to get native balance for %s: %v", blockchain, err) + } else { + nativeSymbol := getNativeTokenSymbol(blockchain) + chainBalances[nativeSymbol] = nativeBalance.String() + } + + // Get token list for this blockchain + tokens := getTokenList(blockchain) + if len(tokens) > 0 { + // Try Multicall3 first for ERC20 tokens + tokenBalances, err := getBalancesMulticall(chainCtx, client, tokens, checksumAddress, blockchain) if err != nil { - resultChan <- chainResult{blockchain: blockchain, err: err} - return + // Fallback to individual calls + tokenBalances, err = getBalancesFallback(chainCtx, client, tokens, checksumAddress) + if err != nil { + resultChan <- chainResult{blockchain: blockchain, err: err} + return + } + } + // Merge token balances into chain balances + for token, balance := range tokenBalances { + chainBalances[token] = balance } } @@ -236,3 +249,10 @@ func getTokenList(blockchain string) []string { } return nil } + +func getNativeTokenSymbol(blockchain string) string { + if chain, ok := contractsConfig[blockchain]; ok { + return chain.NativeToken + } + return "" +} diff --git a/nodebalancer/cmd/nodebalancer/configs.go b/nodebalancer/cmd/nodebalancer/configs.go index a3d5400b..0e2ec86a 100644 --- a/nodebalancer/cmd/nodebalancer/configs.go +++ b/nodebalancer/cmd/nodebalancer/configs.go @@ -25,6 +25,8 @@ var ( bugoutClient *bugout.BugoutClient + contractsConfig ContractsConfig + // Bugout client // TODO(kompotkot): Find out why it cuts out the port BUGOUT_BROOD_URL = "https://auth.bugout.dev" @@ -171,14 +173,13 @@ type NodeConfig struct { type TokenConfig map[string]string type ChainConfig struct { - Multicall3 string `json:"multicall3"` - Tokens TokenConfig `json:"tokens"` + Multicall3 string `json:"multicall3"` + Tokens TokenConfig `json:"tokens"` + NativeToken string `json:"native_token"` } type ContractsConfig map[string]ChainConfig -var contractsConfig ContractsConfig - func LoadConfig(configPath string) error { rawBytes, err := ioutil.ReadFile(configPath) if err != nil { @@ -195,10 +196,6 @@ func LoadConfig(configPath string) error { func LoadContractsConfig(configPath string) error { - if configPath == "" { - return nil - } - // Read config file data, err := os.ReadFile(filepath.Join(configPath)) if err != nil { diff --git a/nodebalancer/cmd/nodebalancer/routes.go b/nodebalancer/cmd/nodebalancer/routes.go index 92fe540b..cece7e40 100644 --- a/nodebalancer/cmd/nodebalancer/routes.go +++ b/nodebalancer/cmd/nodebalancer/routes.go @@ -130,6 +130,11 @@ func balancesRoute(w http.ResponseWriter, r *http.Request) { // Get address from query params address := r.URL.Query().Get("address") + if address == "" { + http.Error(w, "Address is required", http.StatusBadRequest) + return + } + // Check cache first cacheKey := fmt.Sprintf("balances:%s", address) if cachedData, found := balancesCache.Get(cacheKey); found {