fix keystore chaintypes

smartcontractkit · Dec 23, 2024 · 490a247 · 490a247
1 parent 3c241fd
commit 490a247
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 8 deletions.
diff --git a/core/services/keystore/chaintype/chaintype.go b/core/services/keystore/chaintype/chaintype.go
@@ -50,6 +50,8 @@ func NewChainType(typ uint8) (ChainType, error) {
 		return StarkNet, nil
 	case 5:
 		return Aptos, nil
+	case 6:
+		return Tron, nil
 	default:
 		return "", fmt.Errorf("unexpected chaintype.ChainType: %#v", typ)
 	}
@@ -67,6 +69,8 @@ func (c ChainType) Type() (uint8, error) {
 		return 4, nil
 	case Aptos:
 		return 5, nil
+	case Tron:
+		return 6, nil
 	default:
 		return 0, fmt.Errorf("unexpected chaintype.ChainType: %#v", c)
 	}

diff --git a/core/services/keystore/keys/ocr2key/export.go b/core/services/keystore/keys/ocr2key/export.go
@@ -48,6 +48,8 @@ func FromEncryptedJSON(keyJSON []byte, password string) (KeyBundle, error) {
 				kb = newKeyBundle(new(starkkey.OCR2Key))
 			case chaintype.Aptos:
 				kb = newKeyBundle(new(aptosKeyring))
+			case chaintype.Tron:
+				kb = newKeyBundle(new(evmKeyring))
 			default:
 				return nil, chaintype.NewErrInvalidChainType(export.ChainType)
 			}

diff --git a/core/services/keystore/keys/ocr2key/export_test.go b/core/services/keystore/keys/ocr2key/export_test.go
@@ -19,6 +19,7 @@ func TestExport(t *testing.T) {
 		{chain: chaintype.Solana},
 		{chain: chaintype.StarkNet},
 		{chain: chaintype.Aptos},
+		{chain: chaintype.Tron},
 	}
 	for _, tc := range tt {
 		tc := tc

diff --git a/core/services/keystore/keys/ocr2key/key_bundle.go b/core/services/keystore/keys/ocr2key/key_bundle.go
@@ -59,6 +59,8 @@ func New(chainType chaintype.ChainType) (KeyBundle, error) {
 		return newKeyBundleRand(chaintype.StarkNet, starkkey.NewOCR2Key)
 	case chaintype.Aptos:
 		return newKeyBundleRand(chaintype.Aptos, newAptosKeyring)
+	case chaintype.Tron:
+		return newKeyBundleRand(chaintype.Tron, newEVMKeyring)
 	}
 	return nil, chaintype.NewErrInvalidChainType(chainType)
 }
@@ -76,6 +78,8 @@ func MustNewInsecure(reader io.Reader, chainType chaintype.ChainType) KeyBundle
 		return mustNewKeyBundleInsecure(chaintype.StarkNet, starkkey.NewOCR2Key, reader)
 	case chaintype.Aptos:
 		return mustNewKeyBundleInsecure(chaintype.Aptos, newAptosKeyring, reader)
+	case chaintype.Tron:
+		return mustNewKeyBundleInsecure(chaintype.Tron, newEVMKeyring, reader)
 	}
 	panic(chaintype.NewErrInvalidChainType(chainType))
 }
@@ -126,6 +130,8 @@ func (raw Raw) Key() (kb KeyBundle) {
 		kb = newKeyBundle(new(starkkey.OCR2Key))
 	case chaintype.Aptos:
 		kb = newKeyBundle(new(aptosKeyring))
+	case chaintype.Tron:
+		kb = newKeyBundle(new(evmKeyring))
 	default:
 		return nil
 	}

diff --git a/core/services/keystore/keys/tronkey/key.go b/core/services/keystore/keys/tronkey/key.go
@@ -46,14 +46,7 @@ type Key struct {
 }
 
 func New() (Key, error) {
-	privKeyECDSA, err := ecdsa.GenerateKey(crypto.S256(), rand.Reader)
-	if err != nil {
-		return Key{}, err
-	}
-	return Key{
-		privKey: privKeyECDSA,
-		pubKey:  &privKeyECDSA.PublicKey,
-	}, nil
+	return newFrom(rand.Reader)
 }
 
 // MustNewInsecure return Key if no error