From cf2417c74c683aa89de7a4b560a9768828d830d6 Mon Sep 17 00:00:00 2001
From: rapha <raphael.thurnherr1990@gmail.com>
Date: Tue, 5 Mar 2024 16:57:51 +0100
Subject: [PATCH] feat: support multiple rest/rpc endpoints

---
 tools/kysor/cmd/chain/client.go |  9 ++--
 tools/kysor/cmd/config/kysor.go | 79 ++++++++++++++++++++++++++++-----
 tools/kysor/cmd/start.go        | 14 +++++-
 3 files changed, 84 insertions(+), 18 deletions(-)

diff --git a/tools/kysor/cmd/chain/client.go b/tools/kysor/cmd/chain/client.go
index 0379c6d3..61cee007 100644
--- a/tools/kysor/cmd/chain/client.go
+++ b/tools/kysor/cmd/chain/client.go
@@ -81,17 +81,18 @@ type KyveClient struct {
 }
 
 func NewKyveClient(cfg config.KysorConfig, valaccConfigs []config.ValaccountConfig) (*KyveClient, error) {
-	if cfg.RPC == "" {
-		return nil, fmt.Errorf("rpc address must not be empty")
+	rpc, err := cfg.GetWorkingRPC()
+	if err != nil {
+		return nil, err
 	}
 
-	httpClient, err := libclient.DefaultHTTPClient(cfg.RPC)
+	httpClient, err := libclient.DefaultHTTPClient(rpc)
 	if err != nil {
 		return nil, err
 	}
 
 	httpClient.Timeout = 10 * time.Second
-	rpcClient, err := rpchttp.NewWithClient(cfg.RPC, "/websocket", httpClient)
+	rpcClient, err := rpchttp.NewWithClient(rpc, "/websocket", httpClient)
 	if err != nil {
 		return nil, err
 	}
diff --git a/tools/kysor/cmd/config/kysor.go b/tools/kysor/cmd/config/kysor.go
index 98c540f0..25bcf7a6 100644
--- a/tools/kysor/cmd/config/kysor.go
+++ b/tools/kysor/cmd/config/kysor.go
@@ -2,10 +2,13 @@ package config
 
 import (
 	"fmt"
+	"net"
+	"net/http"
 	"os"
 	"path/filepath"
 	"regexp"
 	"strings"
+	"time"
 
 	"github.com/mitchellh/go-homedir"
 
@@ -54,6 +57,46 @@ func (c KysorConfig) GetChainPrettyName() string {
 	return c.ChainID
 }
 
+func (c KysorConfig) GetWorkingRPC() (string, error) {
+	var client = http.Client{
+		Transport: &http.Transport{
+			DialContext: (&net.Dialer{
+				Timeout: 5 * time.Second,
+			}).DialContext,
+		},
+	}
+
+	// Ping the rpc to check if it is working
+	// If it is not working, try the next one
+	for _, rpc := range strings.Split(c.RPC, ",") {
+		_, err := client.Get(rpc + "/status")
+		if err == nil {
+			return rpc, nil
+		}
+	}
+	return "", fmt.Errorf("no working rpc found")
+}
+
+func (c KysorConfig) GetWorkingREST() (string, error) {
+	var client = http.Client{
+		Transport: &http.Transport{
+			DialContext: (&net.Dialer{
+				Timeout: 5 * time.Second,
+			}).DialContext,
+		},
+	}
+
+	// Ping the rest to check if it is working
+	// If it is not working, try the next one
+	for _, rest := range strings.Split(c.REST, ",") {
+		_, err := client.Get(rest)
+		if err == nil {
+			return rest, nil
+		}
+	}
+	return "", fmt.Errorf("no working rest found")
+}
+
 func (c KysorConfig) Save(path string) error {
 	return save(c, path)
 }
@@ -167,21 +210,33 @@ func loadKysorConfig(cmd *cobra.Command, _ []string) error {
 		return fmt.Errorf("error unmarshalling config file: %s", err)
 	}
 
-	// Add port to the RPC and REST URLs if they are missing
-	if !regexp.MustCompile(`:\d+$`).MatchString(config.RPC) {
-		if strings.HasPrefix(config.RPC, "https://") {
-			config.RPC += ":443"
-		} else if strings.HasPrefix(config.RPC, "http://") {
-			config.RPC += ":80"
+	// Add port to the RPC URLs if they are missing
+	var rpcs []string
+	for _, rpc := range strings.Split(config.RPC, ",") {
+		if !regexp.MustCompile(`:\d+$`).MatchString(rpc) {
+			if strings.HasPrefix(rpc, "https://") {
+				rpc += ":443"
+			} else if strings.HasPrefix(rpc, "http://") {
+				rpc += ":80"
+			}
 		}
-	}
-	if !regexp.MustCompile(`:\d+$`).MatchString(config.REST) {
-		if strings.HasPrefix(config.REST, "https://") {
-			config.REST += ":443"
-		} else if strings.HasPrefix(config.REST, "http://") {
-			config.REST += ":80"
+		rpcs = append(rpcs, rpc)
+	}
+	config.RPC = strings.Join(rpcs, ",")
+
+	// Add port to the REST URLs if they are missing
+	var rests []string
+	for _, rest := range strings.Split(config.REST, ",") {
+		if !regexp.MustCompile(`:\d+$`).MatchString(rest) {
+			if strings.HasPrefix(rest, "https://") {
+				rest += ":443"
+			} else if strings.HasPrefix(rest, "http://") {
+				rest += ":80"
+			}
 		}
+		rests = append(rests, rest)
 	}
+	config.REST = strings.Join(rests, ",")
 
 	return nil
 }
diff --git a/tools/kysor/cmd/start.go b/tools/kysor/cmd/start.go
index cd8786b2..1d9fbc00 100644
--- a/tools/kysor/cmd/start.go
+++ b/tools/kysor/cmd/start.go
@@ -387,10 +387,20 @@ func startContainers(cli *client.Client, valConfig config.ValaccountConfig, pool
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
+	rpc, err := config.GetConfigX().GetWorkingRPC()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	rest, err := config.GetConfigX().GetWorkingREST()
+	if err != nil {
+		return nil, nil, err
+	}
+
 	env, err := docker.CreateProtocolEnv(docker.ProtocolEnv{
 		Valaccount:  valConfig.Valaccount,
-		RpcAddress:  config.GetConfigX().RPC,
-		RestAddress: config.GetConfigX().REST,
+		RpcAddress:  rpc,
+		RestAddress: rest,
 		Host:        runtimeName,
 		PoolId:      pool.Id,
 		Debug:       debug,