From 1e037288711023e939582fbb249808fff1219708 Mon Sep 17 00:00:00 2001
From: Tudor Malene <tudor.malene@gmail.com>
Date: Fri, 13 Dec 2024 13:54:17 +0200
Subject: [PATCH] re-read the enclave key

---
 go/common/types.go         |  4 ++++
 go/enclave/rpc/vk_utils.go |  2 +-
 go/rpc/encrypted_client.go | 19 ++++++++++++++++++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/go/common/types.go b/go/common/types.go
index 53b05ed06f..f826d09aab 100644
--- a/go/common/types.go
+++ b/go/common/types.go
@@ -1,6 +1,7 @@
 package common
 
 import (
+	"errors"
 	"fmt"
 	"math/big"
 
@@ -77,6 +78,9 @@ type (
 	EnclaveID = common.Address
 )
 
+// FailedDecryptErr - when the TEN enclave fails to decrypt an RPC request
+var FailedDecryptErr = errors.New("failed to decrypt RPC payload. please use the correct enclave key")
+
 // EncryptedRPCRequest - an encrypted request with extra plaintext metadata
 type EncryptedRPCRequest struct {
 	Req  EncryptedRequest
diff --git a/go/enclave/rpc/vk_utils.go b/go/enclave/rpc/vk_utils.go
index 16a75e3701..8b60f6e46f 100644
--- a/go/enclave/rpc/vk_utils.go
+++ b/go/enclave/rpc/vk_utils.go
@@ -55,7 +55,7 @@ func HandleEncryptedRPC(ctx context.Context,
 	// 1. Decrypt request
 	plaintextRequest, err := encManager.DecryptBytes(encReq)
 	if err != nil {
-		return responses.AsPlaintextError(fmt.Errorf("could not decrypt params - %w", err)), nil
+		return responses.AsPlaintextError(common.FailedDecryptErr), nil
 	}
 
 	// 2. Unmarshall
diff --git a/go/rpc/encrypted_client.go b/go/rpc/encrypted_client.go
index d9ed541f95..f9eb0a6448 100644
--- a/go/rpc/encrypted_client.go
+++ b/go/rpc/encrypted_client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"reflect"
 	"sync/atomic"
@@ -71,7 +72,23 @@ func (c *EncRPCClient) CallContext(ctx context.Context, result interface{}, meth
 	}
 
 	if rpc.IsEncryptedMethod(method) {
-		return c.executeEncryptedCall(ctx, result, method, args...)
+		err := c.executeEncryptedCall(ctx, result, method, args...)
+		// this should only be triggered during testing
+		if err != nil && errors.Is(err, common.FailedDecryptErr) {
+			c.logger.Warn("Reconnecting to new backend. Reading the enclave key.")
+			newKey, err := ReadEnclaveKey(c.obscuroClient)
+			if err != nil {
+				return fmt.Errorf("could not refresh enclave key: %w", err)
+			}
+			enclPubECDSA, err := crypto.DecompressPubkey(newKey)
+			if err != nil {
+				return fmt.Errorf("failed to decompress key for RPC client: %w", err)
+			}
+			c.enclavePublicKey = ecies.ImportECDSAPublic(enclPubECDSA)
+			// retry with the updated key
+			return c.executeEncryptedCall(ctx, result, method, args...)
+		}
+		return err
 	}
 
 	// for non-sensitive methods or when viewing keys are disabled we just delegate directly to the geth RPC client