Bound batch decode by item limit (PLT-820)

rpc/batch_limit_test.go

@@ -0,0 +1,150 @@
+package rpc
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+// makeScalarBatch builds a compact batch array of n integer elements. Each element is a
+// scalar rather than a JSON-RPC object, which is the shape that makes a small frame
+// expand into one jsonrpcMessage per element.
+func makeScalarBatch(n int) string {
+	var b strings.Builder
+	b.WriteByte('[')
+	for i := 0; i < n; i++ {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		b.WriteByte('1')
+	}
+	b.WriteByte(']')
+	return b.String()
+}
+
+// makeCallBatch builds a batch of n test_echo calls, the first of which carries id 1.
+func makeCallBatch(n int) string {
+	elems := make([]string, n)
+	for i := range elems {
+		elems[i] = fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"test_echo","params":["x",99]}`, i+1)
+	}
+	return "[" + strings.Join(elems, ",") + "]"
+}
+
+func TestParseMessageBatchItemLimit(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		raw       string
+		itemLimit int
+		wantLen   int
+		wantBatch bool
+	}{
+		{name: "under limit", raw: makeScalarBatch(3), itemLimit: 5, wantLen: 3, wantBatch: true},
+		{name: "at limit", raw: makeScalarBatch(5), itemLimit: 5, wantLen: 5, wantBatch: true},
+		{name: "one over limit", raw: makeScalarBatch(6), itemLimit: 5, wantLen: 6, wantBatch: true},
+		{name: "far over limit", raw: makeScalarBatch(100000), itemLimit: 5, wantLen: 6, wantBatch: true},
+		{name: "no limit", raw: makeScalarBatch(1000), itemLimit: 0, wantLen: 1000, wantBatch: true},
+		{name: "empty batch", raw: "[]", itemLimit: 5, wantLen: 0, wantBatch: true},
+		{name: "single message", raw: `{"jsonrpc":"2.0","id":1,"method":"test_echo"}`, itemLimit: 5, wantLen: 1, wantBatch: false},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			msgs, batch := parseMessage(json.RawMessage(test.raw), test.itemLimit)
+			if len(msgs) != test.wantLen {
+				t.Fatalf("decoded %d messages, want %d", len(msgs), test.wantLen)
+			}
+			if batch != test.wantBatch {
+				t.Fatalf("batch = %v, want %v", batch, test.wantBatch)
+			}
+		})
+	}
+}
+
+// TestParseMessageBatchAllocationsBounded is the regression guard for the report: a
+// compact array of scalars must not allocate a jsonrpcMessage per element before the
+// item limit is consulted.
+// It cannot run in parallel: testing.AllocsPerRun panics in a parallel test.
+func TestParseMessageBatchAllocationsBounded(t *testing.T) {
+	const (
+		elements  = 200000
+		itemLimit = 100
+	)
+	raw := json.RawMessage(makeScalarBatch(elements))
+
+	allocs := testing.AllocsPerRun(2, func() {
+		parseMessage(raw, itemLimit)
+	})
+	// Bound generously: the point is the order of magnitude, not the exact count. Before
+	// the limit was pushed into parseMessage this exceeded `elements` allocations.
+	if maxAllocs := float64(10 * itemLimit); allocs > maxAllocs {
+		t.Fatalf("parseMessage made %.0f allocations for a %d-element batch limited to %d, want at most %.0f",
+			allocs, elements, itemLimit, maxAllocs)
+	}
+}
+
+// TestWSOversizeBatchRejectedAndConnectionSurvives checks that the truncated decode still
+// produces the protocol-level rejection, and that the connection remains usable after it.
+func TestWSOversizeBatchRejectedAndConnectionSurvives(t *testing.T) {
+	t.Parallel()
+
+	srv := newTestServer()
+	srv.SetBatchLimits(4, 100000)
+	_, wsURL := startWSTestServer(t, srv)
+
+	conn := dialWS(t, wsURL)
+	defer conn.Close()
+
+	writeWSJSON(t, conn, makeScalarBatch(50000))
+	var resp []jsonrpcMessage
+	readWSJSON(t, conn, &resp)
+	if len(resp) != 1 {
+		t.Fatalf("got %d responses, want 1", len(resp))
+	}
+	if resp[0].Error == nil || resp[0].Error.Message != errMsgBatchTooLarge {
+		t.Fatalf("wrong response to oversize batch: %+v", resp[0])
+	}
+
+	// The connection must still serve the next request.
+	writeWSJSON(t, conn, `{"jsonrpc":"2.0","id":7,"method":"test_echo","params":["x",99]}`)
+	var next jsonrpcMessage
+	readWSJSON(t, conn, &next)
+	if next.Error != nil {
+		t.Fatalf("request after oversize batch failed: %v", next.Error)
+	}
+}
+
+// TestHTTPOversizeBatchRejected covers the single-request path, which builds its handler
+// separately from ServeCodec and so wires the codec up on its own.
+func TestHTTPOversizeBatchRejected(t *testing.T) {
+	t.Parallel()
+
+	srv := newTestServer()
+	defer srv.Stop()
+	srv.SetBatchLimits(4, 100000)
+	httpsrv := httptest.NewServer(srv)
+	defer httpsrv.Close()
+
+	resp, err := http.Post(httpsrv.URL, "application/json", strings.NewReader(makeCallBatch(50000)))
+	if err != nil {
+		t.Fatalf("post batch: %v", err)
+	}
+	defer resp.Body.Close()
+
+	var msgs []jsonrpcMessage
+	if err := json.NewDecoder(resp.Body).Decode(&msgs); err != nil {
+		t.Fatalf("decode response: %v", err)
+	}
+	if len(msgs) != 1 {
+		t.Fatalf("got %d responses, want 1", len(msgs))
+	}
+	if msgs[0].Error == nil || msgs[0].Error.Message != errMsgBatchTooLarge {
+		t.Fatalf("wrong response to oversize batch: %+v", msgs[0])
+	}
+}

rpc/client.go

@@ -126,19 +126,22 @@ func (c *Client) newClientConn(conn ServerCodec) *clientConn {
 	ctx = context.WithValue(ctx, clientContextKey{}, c)
 	ctx = context.WithValue(ctx, peerInfoContextKey{}, conn.peerInfo())
 	handler := newHandler(ctx, conn, c.idgen, c.services, c.batchItemLimit, c.batchResponseMaxSize, c.wsConcurrentBudget, c.readLimit, c.admissionEventHook, c.wsAdmissionTimeout)
-	attachBudgetHandler(conn, handler)
+	attachHandler(conn, handler)
 	return &clientConn{conn, handler}
 }
 
-// budgetHandlerSetter is implemented by codecs that need the handler wired in for
-// byte-budget admission (jsonCodec and, via embedding, websocketCodec).
-type budgetHandlerSetter interface {
-	setBudgetHandler(h *handler)
+// handlerSetter is implemented by codecs that need the handler wired in to enforce its
+// read limits (jsonCodec and, via embedding, websocketCodec).
+type handlerSetter interface {
+	setHandler(h *handler)
 }
 
-func attachBudgetHandler(codec ServerCodec, h *handler) {
-	if c, ok := codec.(budgetHandlerSetter); ok {
-		c.setBudgetHandler(h)
+// attachHandler wires h into codec so reads on it observe the handler's limits. A codec
+// that does not implement handlerSetter reads unlimited, and the handler's own checks
+// remain the backstop.
+func attachHandler(codec ServerCodec, h *handler) {
+	if c, ok := codec.(handlerSetter); ok {
+		c.setHandler(h)
 	}
 }
 

rpc/handler.go

@@ -393,7 +393,9 @@ func (h *handler) respondWithBatchTooLarge(cp *callProc, batch []*jsonrpcMessage
 	resp := errorMessage(&invalidRequestError{errMsgBatchTooLarge})
 	// Find the first call and add its "id" field to the error.
 	// This is the best we can do, given that the protocol doesn't have a way
-	// of reporting an error for the entire batch.
+	// of reporting an error for the entire batch. The batch is only decoded up to
+	// the item limit, so a batch whose every decoded element is a notification is
+	// answered with a null id even if a later element was a call.
 	for _, msg := range batch {
 		if msg.isCall() {
 			resp.ID = msg.ID

rpc/json.go

@@ -185,7 +185,7 @@ type jsonCodec struct {
 	encMu   sync.Mutex       // guards the encoder
 	encode  encodeFunc       // encoder to allow multiple transports
 	conn    deadlineCloser
-	handler *handler // set by the read loop for byte-budget admission
+	handler *handler // set by the read loop; source of this codec's read limits
 }
 
 type encodeFunc = func(v interface{}, isErrorResponse bool) error
@@ -221,11 +221,20 @@ func NewCodec(conn Conn) ServerCodec {
 	return NewFuncCodec(conn, encode, dec.Decode)
 }
 
-// setBudgetHandler wires in the handler used for byte-budget admission.
-func (c *jsonCodec) setBudgetHandler(h *handler) {
+// setHandler wires in the handler whose limits govern reads on this codec.
+func (c *jsonCodec) setHandler(h *handler) {
 	c.handler = h
 }
 
+// batchItemLimit returns the number of batch elements a read may decode, or 0 when no
+// limit applies.
+func (c *jsonCodec) batchItemLimit() int {
+	if c.handler == nil {
+		return 0
+	}
+	return c.handler.batchRequestLimit
+}
+
 func (c *jsonCodec) peerInfo() PeerInfo {
 	// This returns "ipc" because all other built-in transports have a separate codec type.
 	return PeerInfo{Transport: "ipc", RemoteAddr: c.remote}
@@ -252,7 +261,7 @@ func (c *jsonCodec) readBatch() (messages []*jsonrpcMessage, batch bool, rawLen
 		}
 		return nil, false, 0, err
 	}
-	messages, batch = parseMessage(rawmsg)
+	messages, batch = parseMessage(rawmsg, c.batchItemLimit())
 	for i, msg := range messages {
 		if msg == nil {
 			// Message is JSON 'null'. Replace with zero value so it
@@ -290,8 +299,9 @@ func (c *jsonCodec) closed() <-chan interface{} {
 // parseMessage parses raw bytes as a (batch of) JSON-RPC message(s). There are no error
 // checks in this function because the raw message has already been syntax-checked when it
 // is called. Any non-JSON-RPC messages in the input return the zero value of
-// jsonrpcMessage.
-func parseMessage(raw json.RawMessage) ([]*jsonrpcMessage, bool) {
+// jsonrpcMessage. A batch is decoded to at most itemLimit+1 messages; an itemLimit of 0
+// means unlimited.
+func parseMessage(raw json.RawMessage, itemLimit int) ([]*jsonrpcMessage, bool) {
 	if !isBatch(raw) {
 		msgs := []*jsonrpcMessage{{}}
 		json.Unmarshal(raw, &msgs[0])
@@ -301,6 +311,13 @@ func parseMessage(raw json.RawMessage) ([]*jsonrpcMessage, bool) {
 	dec.Token() // skip '['
 	var msgs []*jsonrpcMessage
 	for dec.More() {
+		// Stop one element past the limit rather than at it. That surplus element is
+		// what handleBatch's own count check reads to reject the batch, and decoding
+		// past it would allocate a jsonrpcMessage per element of an array the server
+		// has already decided not to serve.
+		if itemLimit > 0 && len(msgs) > itemLimit {
+			break
+		}
 		msgs = append(msgs, new(jsonrpcMessage))
 		dec.Decode(&msgs[len(msgs)-1])
 	}

rpc/server.go

@@ -223,6 +223,7 @@ func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {
 
 	h := newHandler(ctx, codec, s.idgen, &s.services, s.batchItemLimit, s.batchResponseLimit, nil, s.readLimit, nil, s.wsAdmissionTimeout)
 	h.allowSubscribe = false
+	attachHandler(codec, h)
 	defer h.close(io.EOF, nil)
 
 	reqs, batch, _, err := codec.readBatch()

rpc/websocket.go

@@ -362,7 +362,7 @@ func (wc *websocketCodec) readBatch() ([]*jsonrpcMessage, bool, int64, error) {
 		wc.fireOversizeFrameHook(err)
 		return nil, false, 0, err
 	}
-	messages, batch := parseMessage(rawmsg)
+	messages, batch := parseMessage(rawmsg, wc.batchItemLimit())
 	for i, msg := range messages {
 		if msg == nil {
 			messages[i] = new(jsonrpcMessage)

rpc/ws_admission_test.go

@@ -546,13 +546,6 @@ func TestWSBudgetWaitTimeoutOnActiveBurst(t *testing.T) {
 	}
 }
 
-// passthroughCodec embeds ServerCodec without implementing budgetHandlerSetter, so
-// attachBudgetHandler's type assertion fails for it — like any decorator that forwards
-// reads but not setBudgetHandler.
-type passthroughCodec struct {
-	ServerCodec
-}
-
 // TestBudgetCommitWithoutHandlerWiringDoesNotPanic guards against a regression where
 // commitFrameBudget released budget that was never acquired for an unwired codec,
 // panicking the semaphore on the first request.
@@ -569,6 +562,12 @@ func TestBudgetCommitWithoutHandlerWiringDoesNotPanic(t *testing.T) {
 	srv.SetWSConcurrentRequestBytes(budget)
 
 	p1, p2 := net.Pipe()
+
+	// passthroughCodec embeds ServerCodec without implementing handlerSetter, so
+	// attachHandler's type assertion fails for it.
+	type passthroughCodec struct {
+		ServerCodec
+	}
 	wrapped := &passthroughCodec{ServerCodec: NewCodec(p1)}
 	go srv.ServeCodec(wrapped, 0)
 	t.Cleanup(func() { p2.Close(); p1.Close(); srv.Stop() })