From 211dbb719711420358a3cdf718ce476c212adcf5 Mon Sep 17 00:00:00 2001
From: Jordan Krage <jmank88@gmail.com>
Date: Wed, 2 Nov 2022 09:29:33 -0500
Subject: [PATCH] rpc: handle wrong HTTP batch response length (#26064)

---
 rpc/client.go      |  1 +
 rpc/client_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 rpc/http.go        |  3 +++
 3 files changed, 52 insertions(+)

diff --git a/rpc/client.go b/rpc/client.go
index d3ce02977..fcd31319a 100644
--- a/rpc/client.go
+++ b/rpc/client.go
@@ -31,6 +31,7 @@ import (
 )
 
 var (
+	ErrBadResult                 = errors.New("bad result in JSON-RPC response")
 	ErrClientQuit                = errors.New("client is closed")
 	ErrNoResult                  = errors.New("no result in JSON-RPC response")
 	ErrSubscriptionQueueOverflow = errors.New("subscription queue overflow")
diff --git a/rpc/client_test.go b/rpc/client_test.go
index 04c847d0d..5ae549d86 100644
--- a/rpc/client_test.go
+++ b/rpc/client_test.go
@@ -19,6 +19,7 @@ package rpc
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/rand"
 	"net"
@@ -144,6 +145,53 @@ func TestClientBatchRequest(t *testing.T) {
 	}
 }
 
+func TestClientBatchRequest_len(t *testing.T) {
+	b, err := json.Marshal([]jsonrpcMessage{
+		{Version: "2.0", ID: json.RawMessage("1"), Method: "foo", Result: json.RawMessage(`"0x1"`)},
+		{Version: "2.0", ID: json.RawMessage("2"), Method: "bar", Result: json.RawMessage(`"0x2"`)},
+	})
+	if err != nil {
+		t.Fatal("failed to encode jsonrpc message:", err)
+	}
+	s := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		_, err := rw.Write(b)
+		if err != nil {
+			t.Error("failed to write response:", err)
+		}
+	}))
+	t.Cleanup(s.Close)
+
+	client, err := Dial(s.URL)
+	if err != nil {
+		t.Fatal("failed to dial test server:", err)
+	}
+	defer client.Close()
+
+	t.Run("too-few", func(t *testing.T) {
+		batch := []BatchElem{
+			{Method: "foo"},
+			{Method: "bar"},
+			{Method: "baz"},
+		}
+		ctx, cancelFn := context.WithTimeout(context.Background(), time.Second)
+		defer cancelFn()
+		if err := client.BatchCallContext(ctx, batch); !errors.Is(err, ErrBadResult) {
+			t.Errorf("expected %q but got: %v", ErrBadResult, err)
+		}
+	})
+
+	t.Run("too-many", func(t *testing.T) {
+		batch := []BatchElem{
+			{Method: "foo"},
+		}
+		ctx, cancelFn := context.WithTimeout(context.Background(), time.Second)
+		defer cancelFn()
+		if err := client.BatchCallContext(ctx, batch); !errors.Is(err, ErrBadResult) {
+			t.Errorf("expected %q but got: %v", ErrBadResult, err)
+		}
+	})
+}
+
 func TestClientNotify(t *testing.T) {
 	server := newTestServer()
 	defer server.Stop()
diff --git a/rpc/http.go b/rpc/http.go
index 858d80858..b82ea2707 100644
--- a/rpc/http.go
+++ b/rpc/http.go
@@ -173,6 +173,9 @@ func (c *Client) sendBatchHTTP(ctx context.Context, op *requestOp, msgs []*jsonr
 	if err := json.NewDecoder(respBody).Decode(&respmsgs); err != nil {
 		return err
 	}
+	if len(respmsgs) != len(msgs) {
+		return fmt.Errorf("batch has %d requests but response has %d: %w", len(msgs), len(respmsgs), ErrBadResult)
+	}
 	for i := 0; i < len(respmsgs); i++ {
 		op.resp <- &respmsgs[i]
 	}