@@ -31,6 +31,7 @@ import (
3131 "github.com/erigontech/erigon/p2p/enode"
3232 "github.com/erigontech/erigon/p2p/forkid"
3333 "github.com/erigontech/erigon/p2p/protocols/eth"
34+ "github.com/erigontech/erigon/p2p/protocols/wit"
3435)
3536
3637// Handles RLP encoding/decoding for p2p.Msg
@@ -705,3 +706,110 @@ func TestSentryServerImpl_SetStatusInitPanic(t *testing.T) {
705706 t .Fatalf ("error expected" )
706707 }
707708}
709+
710+ // newTestPeerInfoWithEth creates a PeerInfo backed by a *p2p.Peer that has
711+ // the eth protocol in its running map (so WaitForEth won't fail) and marks
712+ // the eth handshake as already completed.
713+ func newTestPeerInfoWithEth (t * testing.T ) (* PeerInfo , [64 ]byte ) {
714+ t .Helper ()
715+ var pubkey [64 ]byte
716+ pubkey [0 ] = 0x01
717+ id := enode.ID {}
718+ copy (id [:], pubkey [:])
719+
720+ caps := []p2p.Cap {
721+ {Name : eth .ProtocolName , Version : direct .ETH68 },
722+ }
723+ protocols := []p2p.Protocol {
724+ {Name : eth .ProtocolName , Version : direct .ETH68 , Length : eth .ProtocolLengths [direct .ETH68 ]},
725+ }
726+ peer := p2p .NewPeerWithProtocols (id , pubkey , "test-peer" , caps , protocols , false )
727+
728+ rw := NewRLPReadWriter ()
729+ t .Cleanup (rw .Close )
730+
731+ pi := NewPeerInfo (peer , rw )
732+ // Mark eth handshake as done so WaitForEth returns immediately.
733+ pi .SetEthProtocol (direct .ETH68 )
734+
735+ return pi , pubkey
736+ }
737+
738+ // TestRunWitPeer_MalformedNewWitnessMsg verifies that a malformed
739+ // NewWitnessMsg causes a PeerError (peer disconnect) rather than a
740+ // nil-pointer panic. This is the regression test for the DoS
741+ // vulnerability where a missing 'continue' after RLP decode failure
742+ // led to query.Witness.Header().Hash() panicking on a nil Witness.
743+ func TestRunWitPeer_MalformedNewWitnessMsg (t * testing.T ) {
744+ t .Parallel ()
745+
746+ peerInfo , peerID := newTestPeerInfoWithEth (t )
747+
748+ rw := NewRLPReadWriter ()
749+ t .Cleanup (rw .Close )
750+
751+ logger := log .Root ()
752+
753+ send := func (msgId sentryproto.MessageId , peerID [64 ]byte , b []byte ) {}
754+ hasSubscribers := func (msgId sentryproto.MessageId ) bool { return true }
755+ getWitnessRequest := func (hash common.Hash , peerID [64 ]byte ) bool { return false }
756+
757+ // Feed a NewWitnessMsg with garbage RLP payload.
758+ garbage := []byte {0xff , 0xfe , 0xfd }
759+ rw .readCh <- p2p.Msg {
760+ Code : wit .NewWitnessMsg ,
761+ Size : uint32 (len (garbage )),
762+ Payload : io .NopCloser (bytes .NewReader (garbage )),
763+ }
764+
765+ errCh := make (chan * p2p.PeerError , 1 )
766+ go func () {
767+ errCh <- runWitPeer (t .Context (), peerID , rw , peerInfo , send , hasSubscribers , getWitnessRequest , logger )
768+ }()
769+
770+ select {
771+ case peerErr := <- errCh :
772+ require .NotNil (t , peerErr , "expected a PeerError for malformed message" )
773+ assert .Equal (t , p2p .PeerErrorInvalidMessage , peerErr .Code )
774+ case <- time .After (5 * time .Second ):
775+ t .Fatal ("runWitPeer did not return within timeout" )
776+ }
777+ }
778+
779+ // TestRunWitPeer_MalformedNewWitnessHashesMsg verifies the same
780+ // protection for NewWitnessHashesMsg.
781+ func TestRunWitPeer_MalformedNewWitnessHashesMsg (t * testing.T ) {
782+ t .Parallel ()
783+
784+ peerInfo , peerID := newTestPeerInfoWithEth (t )
785+
786+ rw := NewRLPReadWriter ()
787+ t .Cleanup (rw .Close )
788+
789+ logger := log .Root ()
790+
791+ send := func (msgId sentryproto.MessageId , peerID [64 ]byte , b []byte ) {}
792+ hasSubscribers := func (msgId sentryproto.MessageId ) bool { return true }
793+ getWitnessRequest := func (hash common.Hash , peerID [64 ]byte ) bool { return false }
794+
795+ // Feed a NewWitnessHashesMsg with garbage RLP payload.
796+ garbage := []byte {0xff , 0xfe , 0xfd }
797+ rw .readCh <- p2p.Msg {
798+ Code : wit .NewWitnessHashesMsg ,
799+ Size : uint32 (len (garbage )),
800+ Payload : io .NopCloser (bytes .NewReader (garbage )),
801+ }
802+
803+ errCh := make (chan * p2p.PeerError , 1 )
804+ go func () {
805+ errCh <- runWitPeer (t .Context (), peerID , rw , peerInfo , send , hasSubscribers , getWitnessRequest , logger )
806+ }()
807+
808+ select {
809+ case peerErr := <- errCh :
810+ require .NotNil (t , peerErr , "expected a PeerError for malformed message" )
811+ assert .Equal (t , p2p .PeerErrorInvalidMessage , peerErr .Code )
812+ case <- time .After (5 * time .Second ):
813+ t .Fatal ("runWitPeer did not return within timeout" )
814+ }
815+ }
0 commit comments