From 530d361d215be10d65192e9acd9c6d36fc296788 Mon Sep 17 00:00:00 2001
From: Matt Ketmo <matthieu@moquet.net>
Date: Fri, 7 Jun 2024 11:08:09 +0200
Subject: [PATCH] feat: add option to disable websocket (#76)

---
 pkg/app/run.go  | 11 ++++++++++-
 pkg/rpc/node.go | 18 ++++++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/pkg/app/run.go b/pkg/app/run.go
index 7f4ad68..52fe8de 100644
--- a/pkg/app/run.go
+++ b/pkg/app/run.go
@@ -232,7 +232,16 @@ func createNodePool(ctx context.Context, nodes []string) (*rpc.Pool, error) {
 			return nil, fmt.Errorf("failed to create client: %w", err)
 		}
 
-		rpcNodes[i] = rpc.NewNode(client)
+		opts := []rpc.NodeOption{}
+
+		// Check is query string websocket is present in the endpoint
+		if u, err := url.Parse(endpoint); err == nil {
+			if u.Query().Get("__websocket") == "0" {
+				opts = append(opts, rpc.DisableWebsocket())
+			}
+		}
+
+		rpcNodes[i] = rpc.NewNode(client, opts...)
 
 		status, err := rpcNodes[i].Status(ctx)
 		if err != nil {
diff --git a/pkg/rpc/node.go b/pkg/rpc/node.go
index 22041d1..3b96093 100644
--- a/pkg/rpc/node.go
+++ b/pkg/rpc/node.go
@@ -30,12 +30,20 @@ type OnNodeStatus func(ctx context.Context, n *Node, status *ctypes.ResultStatus
 
 type NodeOption func(*Node)
 
+func DisableWebsocket() NodeOption {
+	return func(n *Node) {
+		n.disableWebsocket = true
+	}
+}
+
 type Node struct {
 	Client *http.HTTP
 
 	// Save endpoint url for redacted logging
 	endpoint *url.URL
 
+	disableWebsocket bool
+
 	onStart  []OnNodeStart
 	onStatus []OnNodeStatus
 	onEvent  map[string][]OnNodeEvent
@@ -143,8 +151,10 @@ func (n *Node) Start(ctx context.Context) error {
 	initTicker.Stop()
 
 	// Start the websocket process
-	if err := n.Client.Start(); err != nil {
-		return fmt.Errorf("failed to start client: %w", err)
+	if !n.disableWebsocket {
+		if err := n.Client.Start(); err != nil {
+			return fmt.Errorf("failed to start client: %w", err)
+		}
 	}
 
 	blocksEvents, err := n.Subscribe(ctx, EventNewBlock)
@@ -351,6 +361,10 @@ func (n *Node) Stop(ctx context.Context) error {
 }
 
 func (n *Node) Subscribe(ctx context.Context, eventType string) (<-chan ctypes.ResultEvent, error) {
+	if n.disableWebsocket {
+		return make(chan ctypes.ResultEvent), nil
+	}
+
 	if res, ok := n.subscriptions[eventType]; ok {
 		return res, nil
 	}