diff --git a/whisper/whisperv6/libp2p_glue.go b/whisper/whisperv6/libp2p_glue.go index 4762430b86..2ed9cb1a9e 100644 --- a/whisper/whisperv6/libp2p_glue.go +++ b/whisper/whisperv6/libp2p_glue.go @@ -20,6 +20,7 @@ import ( "bytes" "encoding/binary" "fmt" + "math" "github.com/ethereum/go-ethereum/p2p" inet "github.com/libp2p/go-libp2p-net" @@ -69,13 +70,21 @@ func (stream *LibP2PStream) ReadMsg() (p2p.Msg, error) { // WriteMsg implements the MsgReadWriter interface to write messages // to lilbp2p streams. func (stream *LibP2PStream) WriteMsg(msg p2p.Msg) error { + // Refuse to write messages with an unsigned size greater than + // a signed 32-bit integer size. This is because len() returns + // an int, forcing a conversion at some locations in the code, + // and on some blatforms that might cause an issue. + if msg.Size > math.MaxInt32 { + return fmt.Errorf("Payload size must be a maximum of %d bytes", math.MaxInt32) + } + data := make([]byte, msg.Size+codeLength+payloadSizeLength) binary.LittleEndian.PutUint64(data[0:codeLength], msg.Code) binary.LittleEndian.PutUint32(data[codeLength:codeLength+payloadSizeLength], msg.Size) nbytes, err := msg.Payload.Read(data[codeLength+payloadSizeLength:]) - if (nbytes&0xFFFFFFFF) != nbytes || uint32(nbytes) != msg.Size { + if nbytes > math.MaxInt32 || uint32(nbytes) != msg.Size { return fmt.Errorf("Invalid size read in libp2p stream: read %d bytes, was expecting %d bytes", nbytes, msg.Size) } else if err != nil { return err diff --git a/whisper/whisperv6/libp2p_glue_test.go b/whisper/whisperv6/libp2p_glue_test.go index b1758f9304..b56f3a4baa 100644 --- a/whisper/whisperv6/libp2p_glue_test.go +++ b/whisper/whisperv6/libp2p_glue_test.go @@ -17,6 +17,7 @@ package whisperv6 import ( + "math" "encoding/binary" "context" "bytes" @@ -135,6 +136,8 @@ func TestSimpleDecode(t *testing.T) { hosts := createTestNetwork(ctx, t, 2) hosts[0].SetStreamHandler(testProtocolID, func (s inet.Stream) { + defer s.Close() + lps := LibP2PStream{ stream: s, } @@ -163,4 +166,112 @@ func TestSimpleDecode(t *testing.T) { } stream.Close() -} \ No newline at end of file +} + +func TestSimpleCodeDecode(t *testing.T) { + ctx := context.Background() + hosts := createTestNetwork(ctx, t, 2) + + code := rand.Uint64() + size := rand.Uint32() % 512 + payload := make([]byte, size) + n, err := rand.Read(payload) + if err != nil || uint32(n) != size { + t.Fatalf("Read %d random bytes instead of the expected %d, err: %v", n, size, err) + } + + hosts[0].SetStreamHandler(testProtocolID, func (s inet.Stream) { + defer s.Close() + + lps := LibP2PStream{ + stream: s, + } + + msg, err := lps.ReadMsg() + if err != nil { + t.Fatalf("Error decoding message: %s", err) + } + + if msg.Code != code { + t.Fatalf("Error decoding message code %d instead of %d", msg.Code, code) + } + if int(msg.Size) != len(payload) { + t.Fatalf("Error decoding message size %d instead of %d", msg.Size, len(payload)) + } + + readPayload := make([]byte, len(payload)) + sizeRead, err := msg.Payload.Read(readPayload) + if err != nil || sizeRead != len(payload) { + t.Fatalf("Error reading payload from source: %s (%d bytes read for %d expected)", err, sizeRead, len(payload)) + } else if !bytes.Equal(payload, readPayload) { + t.Fatal("Encoded payload differ from source") + } + }) + stream, err := hosts[1].NewStream(ctx, hosts[0].ID(), testProtocolID) + if err != nil { + t.Fatal(err) + } + + msg := p2p.Msg{ + Code: code, + Size: size, + Payload: bytes.NewReader(payload), + } + + lps := LibP2PStream{ + stream: stream, + } + + err = lps.WriteMsg(msg) + + if err != nil { + t.Fatalf("Error encoding a message to the stream: %s", err) + } + + stream.Close() +} + +func TestMaxWriteSize(t *testing.T) { + ctx := context.Background() + hosts := createTestNetwork(ctx, t, 2) + + code := rand.Uint64() + // This isn't the size that will be reported, but if I actually + // require 2GB or RAM the CI servers will fail. + size := 10 + payload := make([]byte, size) + n, err := rand.Read(payload) + if err != nil || n != size { + t.Fatalf("Read %d random bytes instead of the expected %d, err: %v", n, size, err) + } + + hosts[0].SetStreamHandler(testProtocolID, func (s inet.Stream) { + defer s.Close() + + dummy := []byte{0x0} + _, _ = s.Read(dummy) + }) + stream, err := hosts[1].NewStream(ctx, hosts[0].ID(), testProtocolID) + if err != nil { + t.Fatal(err) + } + + msg := p2p.Msg{ + Code: code, + Size: math.MaxInt32+1, + Payload: bytes.NewReader(payload), + } + + lps := LibP2PStream{ + stream: stream, + } + + err = lps.WriteMsg(msg) + + if err.Error() != "Payload size must be a maximum of 2147483647 bytes" { + t.Fatal("Should have returned an error with invalid payload size") + } + + stream.Close() +} +