p2p: better handling of RLP errors in wit

p2p/peer.go

@@ -129,10 +129,17 @@ type Peer struct {
 
 // NewPeer returns a peer for testing purposes.
 func NewPeer(id enode.ID, pubkey [64]byte, name string, caps []Cap, metricsEnabled bool) *Peer {
+	return NewPeerWithProtocols(id, pubkey, name, caps, nil, metricsEnabled)
+}
+
+// NewPeerWithProtocols returns a peer for testing purposes with the given
+// protocols registered in its running map. Caps and protocols must match
+// for a protocol to appear as running.
+func NewPeerWithProtocols(id enode.ID, pubkey [64]byte, name string, caps []Cap, protocols []Protocol, metricsEnabled bool) *Peer {
 	pipe, _ := net.Pipe()
 	node := enode.SignNull(new(enr.Record), id)
 	conn := &conn{fd: pipe, transport: nil, node: node, caps: caps, name: name}
-	peer := newPeer(log.Root(), conn, nil, pubkey, metricsEnabled)
+	peer := newPeer(log.Root(), conn, protocols, pubkey, metricsEnabled)
 	close(peer.closed) // ensures Disconnect doesn't block
 	return peer
 }

p2p/sentry/sentry_grpc_server.go

@@ -655,6 +655,7 @@ func runWitPeer(
 			var query wit.NewWitnessPacket
 			if err := rlp.DecodeBytes(b, &query); err != nil {
 				logger.Error("decoding NewWitnessMsg: %w, data: %x", err, b)
+				return p2p.NewPeerError(p2p.PeerErrorInvalidMessage, p2p.DiscSubprotocolError, err, "decoding NewWitnessMsg")
 			}
 
 			peerInfo.AddKnownWitness(query.Witness.Header().Hash())
@@ -674,6 +675,7 @@ func runWitPeer(
 			var query wit.NewWitnessHashesPacket
 			if err := rlp.DecodeBytes(b, &query); err != nil {
 				logger.Error("decoding NewWitnessHashesMsg: %w, data: %x", err, b)
+				return p2p.NewPeerError(p2p.PeerErrorInvalidMessage, p2p.DiscSubprotocolError, err, "decoding NewWitnessHashesMsg")
 			}
 
 			for _, hash := range query.Hashes {

p2p/sentry/sentry_grpc_server_test.go

@@ -31,6 +31,7 @@ import (
 	"github.com/erigontech/erigon/p2p/enode"
 	"github.com/erigontech/erigon/p2p/forkid"
 	"github.com/erigontech/erigon/p2p/protocols/eth"
+	"github.com/erigontech/erigon/p2p/protocols/wit"
 )
 
 // Handles RLP encoding/decoding for p2p.Msg
@@ -705,3 +706,110 @@ func TestSentryServerImpl_SetStatusInitPanic(t *testing.T) {
 		t.Fatalf("error expected")
 	}
 }
+
+// newTestPeerInfoWithEth creates a PeerInfo backed by a *p2p.Peer that has
+// the eth protocol in its running map (so WaitForEth won't fail) and marks
+// the eth handshake as already completed.
+func newTestPeerInfoWithEth(t *testing.T) (*PeerInfo, [64]byte) {
+	t.Helper()
+	var pubkey [64]byte
+	pubkey[0] = 0x01
+	id := enode.ID{}
+	copy(id[:], pubkey[:])
+
+	caps := []p2p.Cap{
+		{Name: eth.ProtocolName, Version: direct.ETH68},
+	}
+	protocols := []p2p.Protocol{
+		{Name: eth.ProtocolName, Version: direct.ETH68, Length: eth.ProtocolLengths[direct.ETH68]},
+	}
+	peer := p2p.NewPeerWithProtocols(id, pubkey, "test-peer", caps, protocols, false)
+
+	rw := NewRLPReadWriter()
+	t.Cleanup(rw.Close)
+
+	pi := NewPeerInfo(peer, rw)
+	// Mark eth handshake as done so WaitForEth returns immediately.
+	pi.SetEthProtocol(direct.ETH68)
+
+	return pi, pubkey
+}
+
+// TestRunWitPeer_MalformedNewWitnessMsg verifies that a malformed
+// NewWitnessMsg causes a PeerError (peer disconnect) rather than a
+// nil-pointer panic. This is the regression test for the DoS
+// vulnerability where a missing 'continue' after RLP decode failure
+// led to query.Witness.Header().Hash() panicking on a nil Witness.
+func TestRunWitPeer_MalformedNewWitnessMsg(t *testing.T) {
+	t.Parallel()
+
+	peerInfo, peerID := newTestPeerInfoWithEth(t)
+
+	rw := NewRLPReadWriter()
+	t.Cleanup(rw.Close)
+
+	logger := log.Root()
+
+	send := func(msgId sentryproto.MessageId, peerID [64]byte, b []byte) {}
+	hasSubscribers := func(msgId sentryproto.MessageId) bool { return true }
+	getWitnessRequest := func(hash common.Hash, peerID [64]byte) bool { return false }
+
+	// Feed a NewWitnessMsg with garbage RLP payload.
+	garbage := []byte{0xff, 0xfe, 0xfd}
+	rw.readCh <- p2p.Msg{
+		Code:    wit.NewWitnessMsg,
+		Size:    uint32(len(garbage)),
+		Payload: io.NopCloser(bytes.NewReader(garbage)),
+	}
+
+	errCh := make(chan *p2p.PeerError, 1)
+	go func() {
+		errCh <- runWitPeer(t.Context(), peerID, rw, peerInfo, send, hasSubscribers, getWitnessRequest, logger)
+	}()
+
+	select {
+	case peerErr := <-errCh:
+		require.NotNil(t, peerErr, "expected a PeerError for malformed message")
+		assert.Equal(t, p2p.PeerErrorInvalidMessage, peerErr.Code)
+	case <-time.After(5 * time.Second):
+		t.Fatal("runWitPeer did not return within timeout")
+	}
+}
+
+// TestRunWitPeer_MalformedNewWitnessHashesMsg verifies the same
+// protection for NewWitnessHashesMsg.
+func TestRunWitPeer_MalformedNewWitnessHashesMsg(t *testing.T) {
+	t.Parallel()
+
+	peerInfo, peerID := newTestPeerInfoWithEth(t)
+
+	rw := NewRLPReadWriter()
+	t.Cleanup(rw.Close)
+
+	logger := log.Root()
+
+	send := func(msgId sentryproto.MessageId, peerID [64]byte, b []byte) {}
+	hasSubscribers := func(msgId sentryproto.MessageId) bool { return true }
+	getWitnessRequest := func(hash common.Hash, peerID [64]byte) bool { return false }
+
+	// Feed a NewWitnessHashesMsg with garbage RLP payload.
+	garbage := []byte{0xff, 0xfe, 0xfd}
+	rw.readCh <- p2p.Msg{
+		Code:    wit.NewWitnessHashesMsg,
+		Size:    uint32(len(garbage)),
+		Payload: io.NopCloser(bytes.NewReader(garbage)),
+	}
+
+	errCh := make(chan *p2p.PeerError, 1)
+	go func() {
+		errCh <- runWitPeer(t.Context(), peerID, rw, peerInfo, send, hasSubscribers, getWitnessRequest, logger)
+	}()
+
+	select {
+	case peerErr := <-errCh:
+		require.NotNil(t, peerErr, "expected a PeerError for malformed message")
+		assert.Equal(t, p2p.PeerErrorInvalidMessage, peerErr.Code)
+	case <-time.After(5 * time.Second):
+		t.Fatal("runWitPeer did not return within timeout")
+	}
+}