From 89739f7bf6dfc67319be44c0d4a466ae8a31f89f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Zdyba=C5=82?= <tomek@zdybal.lap.pl>
Date: Tue, 18 Feb 2025 14:58:57 +0100
Subject: [PATCH] feat: generate JWT token on each call

---
 execution.go        | 21 ++++++++++++---------
 integration_test.go |  7 ++++++-
 mocks_test.go       |  3 +--
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/execution.go b/execution.go
index 8620a79..8ef1642 100644
--- a/execution.go
+++ b/execution.go
@@ -54,14 +54,18 @@ func NewEngineAPIExecutionClient(
 		return nil, err
 	}
 
-	authToken, err := getAuthToken(jwtSecret)
+	secret, err := decodeSecret(jwtSecret)
 	if err != nil {
-		ethClient.Close()
 		return nil, err
 	}
 
 	engineClient, err := rpc.DialOptions(context.Background(), engineURL,
 		rpc.WithHTTPAuth(func(h http.Header) error {
+			authToken, err := getAuthToken(secret)
+			if err != nil {
+				return err
+			}
+
 			if authToken != "" {
 				h.Set("Authorization", "Bearer "+authToken)
 			}
@@ -278,23 +282,22 @@ func (c *EngineAPIExecutionClient) derivePrevRandao(blockHeight uint64) common.H
 	return common.BigToHash(big.NewInt(int64(blockHeight))) //nolint:gosec // disable G115
 }
 
-// Add this function to execution.go
-func getAuthToken(jwtSecret string) (string, error) {
-	if jwtSecret == "" {
-		return "", nil
-	}
+func decodeSecret(jwtSecret string) ([]byte, error) {
 	secret, err := hex.DecodeString(strings.TrimPrefix(jwtSecret, "0x"))
 	if err != nil {
-		return "", fmt.Errorf("failed to decode JWT secret: %w", err)
+		return nil, fmt.Errorf("failed to decode JWT secret: %w", err)
 	}
+	return secret, nil
+}
 
+func getAuthToken(jwtSecret []byte) (string, error) {
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
 		"exp": time.Now().Add(time.Hour * 1).Unix(), // Expires in 1 hour
 		"iat": time.Now().Unix(),
 	})
 
 	// Sign the token with the decoded secret
-	authToken, err := token.SignedString(secret)
+	authToken, err := token.SignedString(jwtSecret)
 	if err != nil {
 		return "", fmt.Errorf("failed to sign JWT token: %w", err)
 	}
diff --git a/integration_test.go b/integration_test.go
index d3b0a27..82c8872 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -166,7 +166,12 @@ func waitForRethContainer(t *testing.T, jwtSecret string) error {
 					}
 					req.Header.Set("Content-Type", "application/json")
 
-					authToken, err := getAuthToken(jwtSecret)
+					secret, err := decodeSecret(jwtSecret)
+					if err != nil {
+						return err
+					}
+
+					authToken, err := getAuthToken(secret)
 					if err != nil {
 						return err
 					}
diff --git a/mocks_test.go b/mocks_test.go
index cf6a200..928e1b4 100644
--- a/mocks_test.go
+++ b/mocks_test.go
@@ -194,8 +194,7 @@ func NewMockEthAPI(t *testing.T) *MockEthAPI {
 
 		_ = json.NewEncoder(w).Encode(map[string]interface{}{
 			"jsonrpc": "2.0",
-			"id":      req["id"],
-			"result":  resp,
+			"id":      req["id"], "result": resp,
 		})
 	}))