1 files changed, 131 insertions, 12 deletions
diff --git a/p2p/message.go b/p2p/message.go
index daf2bf05c..07916f7b3 100644
--- a/p2p/message.go
+++ b/p2p/message.go
@@ -1,6 +1,7 @@
 package p2p
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/binary"
 	"errors"
@@ -8,12 +9,37 @@ import (
 	"io"
 	"io/ioutil"
 	"math/big"
+	"net"
+	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/ethereum/go-ethereum/ethutil"
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+// parameters for frameRW
+const (
+	// maximum time allowed for reading a message header.
+	// this is effectively the amount of time a connection can be idle.
+	frameReadTimeout = 1 * time.Minute
+
+	// maximum time allowed for reading the payload data of a message.
+	// this is shorter than (and distinct from) frameReadTimeout because
+	// the connection is not considered idle while a message is transferred.
+	// this also limits the payload size of messages to how much the connection
+	// can transfer within the timeout.
+	payloadReadTimeout = 5 * time.Second
+
+	// maximum amount of time allowed for writing a complete message.
+	msgWriteTimeout = 5 * time.Second
+
+	// messages smaller than this many bytes will be read at
+	// once before passing them to a protocol. this increases
+	// concurrency in the processing.
+	wholePayloadSize = 64 * 1024
+)
+
 // Msg defines the structure of a p2p message.
 //
 // Note that a Msg can only be sent once since the Payload reader is
@@ -74,11 +100,14 @@ type MsgWriter interface {
 	// WriteMsg sends a message. It will block until the message's
 	// Payload has been consumed by the other end.
 	//
-	// Note that messages can be sent only once.
+	// Note that messages can be sent only once because their
+	// payload reader is drained.
 	WriteMsg(Msg) error
 }
 
 // MsgReadWriter provides reading and writing of encoded messages.
+// Implementations should ensure that ReadMsg and WriteMsg can be
+// called simultaneously from multiple goroutines.
 type MsgReadWriter interface {
 	MsgReader
 	MsgWriter
@@ -90,8 +119,45 @@ func EncodeMsg(w MsgWriter, code uint64, data ...interface{}) error {
 	return w.WriteMsg(NewMsg(code, data...))
 }
 
+// frameRW is a MsgReadWriter that reads and writes devp2p message frames.
+// As required by the interface, ReadMsg and WriteMsg can be called from
+// multiple goroutines.
+type frameRW struct {
+	net.Conn // make Conn methods available. be careful.
+	bufconn  *bufio.ReadWriter
+
+	// this channel is used to 'lend' bufconn to a caller of ReadMsg
+	// until the message payload has been consumed. the channel
+	// receives a value when EOF is reached on the payload, unblocking
+	// a pending call to ReadMsg.
+	rsync chan struct{}
+
+	// this mutex guards writes to bufconn.
+	writeMu sync.Mutex
+}
+
+func newFrameRW(conn net.Conn, timeout time.Duration) *frameRW {
+	rsync := make(chan struct{}, 1)
+	rsync <- struct{}{}
+	return &frameRW{
+		Conn:    conn,
+		bufconn: bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
+		rsync:   rsync,
+	}
+}
+
 var magicToken = []byte{34, 64, 8, 145}
 
+func (rw *frameRW) WriteMsg(msg Msg) error {
+	rw.writeMu.Lock()
+	defer rw.writeMu.Unlock()
+	rw.SetWriteDeadline(time.Now().Add(msgWriteTimeout))
+	if err := writeMsg(rw.bufconn, msg); err != nil {
+		return err
+	}
+	return rw.bufconn.Flush()
+}
+
 func writeMsg(w io.Writer, msg Msg) error {
 	// TODO: handle case when Size + len(code) + len(listhdr) overflows uint32
 	code := ethutil.Encode(uint32(msg.Code))
@@ -120,31 +186,51 @@ func makeListHeader(length uint32) []byte {
 	return append([]byte{lenb}, enc...)
 }
 
-// readMsg reads a message header from r.
-// It takes an rlp.ByteReader to ensure that the decoding doesn't buffer.
-func readMsg(r rlp.ByteReader) (msg Msg, err error) {
+func (rw *frameRW) ReadMsg() (msg Msg, err error) {
+	<-rw.rsync // wait until bufconn is ours
+
+	rw.SetReadDeadline(time.Now().Add(frameReadTimeout))
+
 	// read magic and payload size
 	start := make([]byte, 8)
-	if _, err = io.ReadFull(r, start); err != nil {
-		return msg, newPeerError(errRead, "%v", err)
+	if _, err = io.ReadFull(rw.bufconn, start); err != nil {
+		return msg, err
 	}
 	if !bytes.HasPrefix(start, magicToken) {
-		return msg, newPeerError(errMagicTokenMismatch, "got %x, want %x", start[:4], magicToken)
+		return msg, fmt.Errorf("bad magic token %x", start[:4], magicToken)
 	}
 	size := binary.BigEndian.Uint32(start[4:])
 
 	// decode start of RLP message to get the message code
-	posr := &postrack{r, 0}
+	posr := &postrack{rw.bufconn, 0}
 	s := rlp.NewStream(posr)
 	if _, err := s.List(); err != nil {
 		return msg, err
 	}
-	code, err := s.Uint()
+	msg.Code, err = s.Uint()
 	if err != nil {
 		return msg, err
 	}
-	payloadsize := size - posr.p
-	return Msg{code, payloadsize, io.LimitReader(r, int64(payloadsize))}, nil
+	msg.Size = size - posr.p
+
+	rw.SetReadDeadline(time.Now().Add(payloadReadTimeout))
+
+	if msg.Size <= wholePayloadSize {
+		// msg is small, read all of it and move on to the next message.
+		pbuf := make([]byte, msg.Size)
+		if _, err := io.ReadFull(rw.bufconn, pbuf); err != nil {
+			return msg, err
+		}
+		rw.rsync <- struct{}{} // bufconn is available again
+		msg.Payload = bytes.NewReader(pbuf)
+	} else {
+		// lend bufconn to the caller until it has
+		// consumed the payload. eofSignal will send a value
+		// on rw.rsync when EOF is reached.
+		pr := &eofSignal{rw.bufconn, msg.Size, rw.rsync}
+		msg.Payload = pr
+	}
+	return msg, nil
 }
 
 // postrack wraps an rlp.ByteReader with a position counter.
@@ -167,6 +253,39 @@ func (r *postrack) ReadByte() (byte, error) {
 	return b, err
 }
 
+// eofSignal wraps a reader with eof signaling. the eof channel is
+// closed when the wrapped reader returns an error or when count bytes
+// have been read.
+type eofSignal struct {
+	wrapped io.Reader
+	count   uint32 // number of bytes left
+	eof     chan<- struct{}
+}
+
+// note: when using eofSignal to detect whether a message payload
+// has been read, Read might not be called for zero sized messages.
+func (r *eofSignal) Read(buf []byte) (int, error) {
+	if r.count == 0 {
+		if r.eof != nil {
+			r.eof <- struct{}{}
+			r.eof = nil
+		}
+		return 0, io.EOF
+	}
+
+	max := len(buf)
+	if int(r.count) < len(buf) {
+		max = int(r.count)
+	}
+	n, err := r.wrapped.Read(buf[:max])
+	r.count -= uint32(n)
+	if (err != nil || r.count == 0) && r.eof != nil {
+		r.eof <- struct{}{} // tell Peer that msg has been consumed
+		r.eof = nil
+	}
+	return n, err
+}
+
 // MsgPipe creates a message pipe. Reads on one end are matched
 // with writes on the other. The pipe is full-duplex, both ends
 // implement MsgReadWriter.
@@ -198,7 +317,7 @@ type MsgPipeRW struct {
 func (p *MsgPipeRW) WriteMsg(msg Msg) error {
 	if atomic.LoadInt32(p.closed) == 0 {
 		consumed := make(chan struct{}, 1)
-		msg.Payload = &eofSignal{msg.Payload, int64(msg.Size), consumed}
+		msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed}
 		select {
 		case p.w <- msg:
 			if msg.Size > 0 {