From 759c3c15faf30f171802ea7b9eec1c10bff1445b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bla=C5=BE=20Hrastnik?= <blaz@mxxn.io>
Date: Fri, 15 Nov 2024 21:53:55 +0900
Subject: [PATCH] "robust" chainID handling

---
 deployment/common/view/nops.go              |  1 +
 deployment/environment.go                   | 11 +++++++--
 deployment/environment/memory/job_client.go | 26 ++++++++++++++++-----
 deployment/environment/memory/node.go       |  4 ++++
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/deployment/common/view/nops.go b/deployment/common/view/nops.go
index 12b8dad96f4..5861a225ab9 100644
--- a/deployment/common/view/nops.go
+++ b/deployment/common/view/nops.go
@@ -59,6 +59,7 @@ func GenerateNopsView(nodeIds []string, oc deployment.OffchainClient) (map[strin
 			IsEnabled:    nodeDetails.Node.IsEnabled,
 		}
 		for sel, ocrConfig := range node.SelToOCRConfig {
+			// TODO: support non-EVM here
 			chainid, err := chainsel.ChainIdFromSelector(sel)
 			if err != nil {
 				return nv, err
diff --git a/deployment/environment.go b/deployment/environment.go
index 766a92187e8..59e2693aefb 100644
--- a/deployment/environment.go
+++ b/deployment/environment.go
@@ -309,6 +309,7 @@ func NodeInfo(nodeIDs []string, oc NodeChainConfigsLister) (Nodes, error) {
 				break
 			}
 			var sel uint64
+			var ok bool
 			switch chainConfig.Chain.Type {
 			case nodev1.ChainType_CHAIN_TYPE_EVM:
 				evmChainID, err := strconv.Atoi(chainConfig.Chain.Id)
@@ -324,9 +325,15 @@ func NodeInfo(nodeIDs []string, oc NodeChainConfigsLister) (Nodes, error) {
 				if err != nil {
 					return nil, err
 				}
-				sel = chain_selectors.AptosChainIdToChainSelector()[uint64(aptosChainID)]
+				sel, ok = chain_selectors.AptosChainIdToChainSelector()[uint64(aptosChainID)]
+				if !ok {
+					return nil, fmt.Errorf("chain not found for chain %v", chainConfig.Chain.Id)
+				}
 			case nodev1.ChainType_CHAIN_TYPE_SOLANA:
-				sel = chain_selectors.SolanaChainIdToChainSelector()[chainConfig.Chain.Id]
+				sel, ok = chain_selectors.SolanaChainIdToChainSelector()[chainConfig.Chain.Id]
+				if !ok {
+					return nil, fmt.Errorf("chain not found for chain %v", chainConfig.Chain.Id)
+				}
 			default:
 				// we skip unknown types
 				continue
diff --git a/deployment/environment/memory/job_client.go b/deployment/environment/memory/job_client.go
index df1e3d5c5d5..ad42135cb1c 100644
--- a/deployment/environment/memory/job_client.go
+++ b/deployment/environment/memory/job_client.go
@@ -184,28 +184,42 @@ func (j JobClient) ListNodeChainConfigs(ctx context.Context, in *nodev1.ListNode
 		if err != nil {
 			return nil, err
 		}
-		chainID, err := chainsel.ChainIdFromSelector(selector)
-		if err != nil {
-			return nil, err
-		}
-
 		if family == chainsel.FamilyEVM {
 			// already handled above
 			continue
 		}
 
+		var chainID string
+
 		var ocrtype chaintype.ChainType
 		switch family {
 		case chainsel.FamilyEVM:
 			ocrtype = chaintype.EVM
+			cid, err := chainsel.ChainIdFromSelector(selector)
+			if err != nil {
+				return nil, err
+			}
+			chainID = strconv.Itoa(int(cid))
 		case chainsel.FamilySolana:
 			ocrtype = chaintype.Solana
+			cid, err := chainsel.SolanaChainIdFromSelector(selector)
+			if err != nil {
+				return nil, err
+			}
+			chainID = cid
 		case chainsel.FamilyStarknet:
 			ocrtype = chaintype.StarkNet
+			// TODO: support cid
 		case chainsel.FamilyCosmos:
 			ocrtype = chaintype.Cosmos
+			// TODO: support cid
 		case chainsel.FamilyAptos:
 			ocrtype = chaintype.Aptos
+			cid, err := chainsel.AptosChainIdFromSelector(selector)
+			if err != nil {
+				return nil, err
+			}
+			chainID = strconv.Itoa(int(cid))
 		default:
 			panic(fmt.Sprintf("Unsupported chain family %v", family))
 		}
@@ -238,7 +252,7 @@ func (j JobClient) ListNodeChainConfigs(ctx context.Context, in *nodev1.ListNode
 
 		chainConfigs = append(chainConfigs, &nodev1.ChainConfig{
 			Chain: &nodev1.Chain{
-				Id:   strconv.Itoa(int(chainID)),
+				Id:   chainID,
 				Type: ctype,
 			},
 			AccountAddress: "", // TODO: support AccountAddress
diff --git a/deployment/environment/memory/node.go b/deployment/environment/memory/node.go
index c2e4e457fbd..fd08d3cf17b 100644
--- a/deployment/environment/memory/node.go
+++ b/deployment/environment/memory/node.go
@@ -78,6 +78,10 @@ func NewNode(
 ) *Node {
 	evmchains := make(map[uint64]EVMChain)
 	for _, chain := range chains {
+		// we're only mapping evm chains here
+		if family, err := chainsel.GetSelectorFamily(chain.Selector); err != nil || family != chainsel.FamilyEVM {
+			continue
+		}
 		evmChainID, err := chainsel.ChainIdFromSelector(chain.Selector)
 		if err != nil {
 			t.Fatal(err)