diff --git a/whisper/whisperv6/libp2p_glue.go b/whisper/whisperv6/libp2p_glue.go index 44aea7fcfb..e3626efc8d 100644 --- a/whisper/whisperv6/libp2p_glue.go +++ b/whisper/whisperv6/libp2p_glue.go @@ -17,6 +17,7 @@ package whisperv6 import ( + "io" "github.com/ethereum/go-ethereum/log" "bytes" "context" @@ -101,7 +102,7 @@ func (stream *LibP2PStream) WriteMsg(msg p2p.Msg) error { nbytes, err := msg.Payload.Read(data[codeLength+payloadSizeLength:]) 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 { + } else if err != nil && err != io.EOF { return err } diff --git a/whisper/whisperv6/libp2p_glue_test.go b/whisper/whisperv6/libp2p_glue_test.go index ec46f3ed26..7e79b67d5b 100644 --- a/whisper/whisperv6/libp2p_glue_test.go +++ b/whisper/whisperv6/libp2p_glue_test.go @@ -17,6 +17,7 @@ package whisperv6 import ( + "io" "bytes" "context" "encoding/binary" @@ -36,6 +37,21 @@ const ( testProtocolID = "/whispertesting/6.1" ) +func newTestMsg(t *testing.T, size int) p2p.Msg { + code := rand.Uint64() + 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) + } + + return p2p.Msg{ + Code: code, + Size: math.MaxInt32 + 1, + Payload: bytes.NewReader(payload), + } +} + // Create a network with n mock hosts. Each host in the array is linked to // all hosts preceding it, and has dialed them. func createTestNetwork(ctx context.Context, t *testing.T, n int) []host.Host { @@ -236,15 +252,9 @@ 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) - } + msg := newTestMsg(t, 10) hosts[0].SetStreamHandler(testProtocolID, func(s inet.Stream) { defer s.Close() @@ -257,12 +267,6 @@ func TestMaxWriteSize(t *testing.T) { t.Fatal(err) } - msg := p2p.Msg{ - Code: code, - Size: math.MaxInt32 + 1, - Payload: bytes.NewReader(payload), - } - lps := LibP2PStream{ stream: stream, } @@ -311,3 +315,20 @@ func TestMaxReadSize(t *testing.T) { stream.Close() } + +// TestEndOfStreamDoesNotBlockWriteMsg checks that an EOF in the +// stream will not return an error. +func TestEndOfStreamDoesNotBlockWriteMsg(t *testing.T) { + r, w := io.Pipe() + + s := mocknet.NewStream(w, r) + + lps := &LibP2PStream{s} + + msg := newTestMsg(t, 5) + + err := lps.WriteMsg(msg) + if err == io.EOF { + t.Fatalf("EOF got reported as an error") + } +} \ No newline at end of file