From e8ab47edeae85c4cddf0e30a75f5c98345f24035 Mon Sep 17 00:00:00 2001
From: pk910 <philipp@pk910.de>
Date: Wed, 21 Aug 2024 22:45:58 +0200
Subject: [PATCH] restore functionality of `disableSSZRequests` killswitch

---
 clients/consensus/client.go        | 11 ++++++-----
 clients/consensus/rpc/beaconapi.go | 30 +++++++++++++++++-------------
 services/chainservice.go           |  7 ++++---
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/clients/consensus/client.go b/clients/consensus/client.go
index eb83fc7d..1b19aede 100644
--- a/clients/consensus/client.go
+++ b/clients/consensus/client.go
@@ -14,10 +14,11 @@ import (
 )
 
 type ClientConfig struct {
-	URL       string
-	Name      string
-	Headers   map[string]string
-	SshConfig *sshtunnel.SshConfig
+	URL        string
+	Name       string
+	Headers    map[string]string
+	SshConfig  *sshtunnel.SshConfig
+	DisableSSZ bool
 }
 
 type Client struct {
@@ -56,7 +57,7 @@ type Client struct {
 func (pool *Pool) newPoolClient(clientIdx uint16, endpoint *ClientConfig) (*Client, error) {
 	logger := pool.logger.WithField("client", endpoint.Name)
 
-	rpcClient, err := rpc.NewBeaconClient(endpoint.Name, endpoint.URL, endpoint.Headers, endpoint.SshConfig, logger)
+	rpcClient, err := rpc.NewBeaconClient(endpoint.Name, endpoint.URL, endpoint.Headers, endpoint.SshConfig, endpoint.DisableSSZ, logger)
 	if err != nil {
 		return nil, err
 	}
diff --git a/clients/consensus/rpc/beaconapi.go b/clients/consensus/rpc/beaconapi.go
index ad4eb931..cd49b3e4 100644
--- a/clients/consensus/rpc/beaconapi.go
+++ b/clients/consensus/rpc/beaconapi.go
@@ -28,21 +28,23 @@ import (
 )
 
 type BeaconClient struct {
-	name      string
-	endpoint  string
-	headers   map[string]string
-	sshtunnel *sshtunnel.SSHTunnel
-	clientSvc eth2client.Service
-	logger    logrus.FieldLogger
+	name       string
+	endpoint   string
+	headers    map[string]string
+	sshtunnel  *sshtunnel.SSHTunnel
+	disableSSZ bool
+	clientSvc  eth2client.Service
+	logger     logrus.FieldLogger
 }
 
 // NewBeaconClient is used to create a new beacon client
-func NewBeaconClient(name, endpoint string, headers map[string]string, sshcfg *sshtunnel.SshConfig, logger logrus.FieldLogger) (*BeaconClient, error) {
+func NewBeaconClient(name, endpoint string, headers map[string]string, sshcfg *sshtunnel.SshConfig, disableSSZ bool, logger logrus.FieldLogger) (*BeaconClient, error) {
 	client := &BeaconClient{
-		name:     name,
-		endpoint: endpoint,
-		headers:  headers,
-		logger:   logger,
+		name:       name,
+		endpoint:   endpoint,
+		headers:    headers,
+		disableSSZ: disableSSZ,
+		logger:     logger,
 	}
 
 	if sshcfg != nil {
@@ -104,8 +106,6 @@ func (bc *BeaconClient) Initialize(ctx context.Context) error {
 		http.WithAddress(bc.endpoint),
 		http.WithTimeout(10 * time.Minute),
 		http.WithLogLevel(zerolog.Disabled),
-		// TODO (when upstream PR is merged)
-		// http.WithConnectionCheck(false),
 		http.WithCustomSpecSupport(true),
 	}
 
@@ -114,6 +114,10 @@ func (bc *BeaconClient) Initialize(ctx context.Context) error {
 		cliParams = append(cliParams, http.WithExtraHeaders(bc.headers))
 	}
 
+	if bc.disableSSZ {
+		cliParams = append(cliParams, http.WithEnforceJSON(true))
+	}
+
 	clientSvc, err := http.New(ctx, cliParams...)
 	if err != nil {
 		return err
diff --git a/services/chainservice.go b/services/chainservice.go
index bf181043..b7c56be2 100644
--- a/services/chainservice.go
+++ b/services/chainservice.go
@@ -66,9 +66,10 @@ func (cs *ChainService) StartService() error {
 	// add consensus clients
 	for index, endpoint := range utils.Config.BeaconApi.Endpoints {
 		endpointConfig := &consensus.ClientConfig{
-			URL:     endpoint.Url,
-			Name:    endpoint.Name,
-			Headers: endpoint.Headers,
+			URL:        endpoint.Url,
+			Name:       endpoint.Name,
+			Headers:    endpoint.Headers,
+			DisableSSZ: utils.Config.KillSwitch.DisableSSZRequests,
 		}
 
 		if endpoint.Ssh != nil {