From 60c94a2c2502390aa9e3a4f2a1bc4dedbef1b6fd Mon Sep 17 00:00:00 2001
From: Oleg Lomaka <oleg.lomaka@gmail.com>
Date: Mon, 29 Jan 2024 12:57:30 -0500
Subject: [PATCH] Add NetworkByChainID function to return network and
 blockchain for registered ChainID

---
 .github/workflows/lint.yaml |  4 +--
 .github/workflows/test.yaml |  6 ++--
 chain.go                    | 71 +++++++++++++++++++++++++++----------
 chain_test.go               | 16 +++++++++
 4 files changed, 74 insertions(+), 23 deletions(-)
 create mode 100644 chain_test.go

diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
index 520e55d..0226065 100644
--- a/.github/workflows/lint.yaml
+++ b/.github/workflows/lint.yaml
@@ -13,7 +13,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-go@v3
         with:
-          go-version: 1.19.2
+          go-version: 1.21.6
       - uses: golangci/golangci-lint-action@v3
         with:
-          version: v1.50.0
+          version: v1.55.2
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 3813d65..7c10f84 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -11,10 +11,10 @@ jobs:
     strategy:
       matrix:
         containers:
-          - 1.17.13-bullseye
           - 1.18.10-bullseye
-          - 1.19.5-bullseye
-          - 1.20.0-bullseye
+          - 1.19.13-bullseye
+          - 1.20.13-bullseye
+          - 1.21.6-bookworm
     runs-on: ubuntu-20.04
     container: golang:${{matrix.containers}}
     steps:
diff --git a/chain.go b/chain.go
index e3c1121..a2f84bf 100644
--- a/chain.go
+++ b/chain.go
@@ -1,7 +1,9 @@
 package core
 
 import (
+	"errors"
 	"fmt"
+	"sync"
 
 	"github.com/iden3/go-iden3-core/v2/w3c"
 )
@@ -9,15 +11,26 @@ import (
 // ChainID is alias for int32 that represents ChainID
 type ChainID int32
 
-// ChainIDs Object containing chain IDs for various blockchains and networks.
-var chainIDs = map[string]ChainID{
-	"eth:main":       1,
-	"eth:goerli":     5,
-	"eth:sepolia":    11155111,
-	"polygon:main":   137,
-	"polygon:mumbai": 80001,
-	"zkevm:main":     1101,
-	"zkevm:test":     1442,
+type chainIDKey struct {
+	blockchain Blockchain
+	networkID  NetworkID
+}
+
+var ErrChainIDNotRegistered = errors.New("chainID is not registered")
+
+var chainIDsLock sync.RWMutex
+
+// chainIDs Object containing chain IDs for various blockchains and networks.
+// It can be modified using RegisterChainID public function. So it is guarded
+// by chainIDsLock mutex.
+var chainIDs = map[chainIDKey]ChainID{
+	{Ethereum, Main}:    1,
+	{Ethereum, Goerli}:  5,
+	{Ethereum, Sepolia}: 11155111,
+	{Polygon, Main}:     137,
+	{Polygon, Mumbai}:   80001,
+	{ZkEVM, Main}:       1101,
+	{ZkEVM, Test}:       1442,
 }
 
 // ChainIDfromDID returns chain name from w3c.DID
@@ -38,17 +51,18 @@ func ChainIDfromDID(did w3c.DID) (ChainID, error) {
 		return 0, err
 	}
 
-	chainID, ok := chainIDs[fmt.Sprintf("%s:%s", blockchain, networkID)]
-	if !ok {
-		return 0, fmt.Errorf("chainID not found for %s:%s", blockchain, networkID)
-	}
-
-	return chainID, nil
+	return GetChainID(blockchain, networkID)
 }
 
 // RegisterChainID registers chainID for blockchain and network
 func RegisterChainID(blockchain Blockchain, network NetworkID, chainID int) error {
-	k := fmt.Sprintf("%s:%s", blockchain, network)
+	chainIDsLock.Lock()
+	defer chainIDsLock.Unlock()
+
+	k := chainIDKey{
+		blockchain: blockchain,
+		networkID:  network,
+	}
 	existingChainID, ok := chainIDs[k]
 	if ok && existingChainID == ChainID(chainID) {
 		return nil
@@ -67,10 +81,31 @@ func RegisterChainID(blockchain Blockchain, network NetworkID, chainID int) erro
 
 // GetChainID returns chainID for blockchain and network
 func GetChainID(blockchain Blockchain, network NetworkID) (ChainID, error) {
-	k := fmt.Sprintf("%s:%s", blockchain, network)
+	chainIDsLock.RLock()
+	defer chainIDsLock.RUnlock()
+
+	k := chainIDKey{
+		blockchain: blockchain,
+		networkID:  network,
+	}
 	if _, ok := chainIDs[k]; !ok {
-		return 0, fmt.Errorf("chainID not registered for %s:%s", blockchain, network)
+		return 0, fmt.Errorf("%w for %s:%s", ErrChainIDNotRegistered, blockchain,
+			network)
 	}
 
 	return chainIDs[k], nil
 }
+
+// NetworkByChainID returns blockchain and networkID for registered chain ID.
+// Or ErrChainIDNotRegistered error if chainID is not registered.
+func NetworkByChainID(chainID ChainID) (Blockchain, NetworkID, error) {
+	chainIDsLock.RLock()
+	defer chainIDsLock.RUnlock()
+
+	for k, v := range chainIDs {
+		if v == chainID {
+			return k.blockchain, k.networkID, nil
+		}
+	}
+	return NoChain, NoNetwork, ErrChainIDNotRegistered
+}
diff --git a/chain_test.go b/chain_test.go
new file mode 100644
index 0000000..625a5cc
--- /dev/null
+++ b/chain_test.go
@@ -0,0 +1,16 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNetworkByChainID(t *testing.T) {
+	chainID, err := GetChainID(Ethereum, Main)
+	require.NoError(t, err)
+	blockchain, networkID, err := NetworkByChainID(chainID)
+	require.NoError(t, err)
+	require.Equal(t, Ethereum, blockchain)
+	require.Equal(t, Main, networkID)
+}