1 files changed, 357 insertions, 0 deletions
diff --git a/p2p/message.go b/p2p/message.go
new file mode 100644
index 000000000..7adad4b09
--- /dev/null
+++ b/p2p/message.go
@@ -0,0 +1,357 @@
+package p2p
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"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
+// consumed during sending. It is not possible to create a Msg and
+// send it any number of times. If you want to reuse an encoded
+// structure, encode the payload into a byte array and create a
+// separate Msg with a bytes.Reader as Payload for each send.
+type Msg struct {
+	Code    uint64
+	Size    uint32 // size of the paylod
+	Payload io.Reader
+}
+
+// NewMsg creates an RLP-encoded message with the given code.
+func NewMsg(code uint64, params ...interface{}) Msg {
+	buf := new(bytes.Buffer)
+	for _, p := range params {
+		buf.Write(ethutil.Encode(p))
+	}
+	return Msg{Code: code, Size: uint32(buf.Len()), Payload: buf}
+}
+
+func encodePayload(params ...interface{}) []byte {
+	buf := new(bytes.Buffer)
+	for _, p := range params {
+		buf.Write(ethutil.Encode(p))
+	}
+	return buf.Bytes()
+}
+
+// Decode parse the RLP content of a message into
+// the given value, which must be a pointer.
+//
+// For the decoding rules, please see package rlp.
+func (msg Msg) Decode(val interface{}) error {
+	s := rlp.NewListStream(msg.Payload, uint64(msg.Size))
+	if err := s.Decode(val); err != nil {
+		return newPeerError(errInvalidMsg, "(code %#x) (size %d) %v", msg.Code, msg.Size, err)
+	}
+	return nil
+}
+
+func (msg Msg) String() string {
+	return fmt.Sprintf("msg #%v (%v bytes)", msg.Code, msg.Size)
+}
+
+// Discard reads any remaining payload data into a black hole.
+func (msg Msg) Discard() error {
+	_, err := io.Copy(ioutil.Discard, msg.Payload)
+	return err
+}
+
+type MsgReader interface {
+	ReadMsg() (Msg, error)
+}
+
+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 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
+}
+
+// EncodeMsg writes an RLP-encoded message with the given code and
+// data elements.
+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))
+	listhdr := makeListHeader(msg.Size + uint32(len(code)))
+	payloadLen := uint32(len(listhdr)) + uint32(len(code)) + msg.Size
+
+	start := make([]byte, 8)
+	copy(start, magicToken)
+	binary.BigEndian.PutUint32(start[4:], payloadLen)
+
+	for _, b := range [][]byte{start, listhdr, code} {
+		if _, err := w.Write(b); err != nil {
+			return err
+		}
+	}
+	_, err := io.CopyN(w, msg.Payload, int64(msg.Size))
+	return err
+}
+
+func makeListHeader(length uint32) []byte {
+	if length < 56 {
+		return []byte{byte(length + 0xc0)}
+	}
+	enc := big.NewInt(int64(length)).Bytes()
+	lenb := byte(len(enc)) + 0xf7
+	return append([]byte{lenb}, enc...)
+}
+
+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(rw.bufconn, start); err != nil {
+		return msg, err
+	}
+	if !bytes.HasPrefix(start, magicToken) {
+		return msg, fmt.Errorf("bad magic token %x", start[:4])
+	}
+	size := binary.BigEndian.Uint32(start[4:])
+
+	// decode start of RLP message to get the message code
+	posr := &postrack{rw.bufconn, 0}
+	s := rlp.NewStream(posr)
+	if _, err := s.List(); err != nil {
+		return msg, err
+	}
+	msg.Code, err = s.Uint()
+	if err != nil {
+		return msg, err
+	}
+	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.
+type postrack struct {
+	r rlp.ByteReader
+	p uint32
+}
+
+func (r *postrack) Read(buf []byte) (int, error) {
+	n, err := r.r.Read(buf)
+	r.p += uint32(n)
+	return n, err
+}
+
+func (r *postrack) ReadByte() (byte, error) {
+	b, err := r.r.ReadByte()
+	if err == nil {
+		r.p++
+	}
+	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.
+func MsgPipe() (*MsgPipeRW, *MsgPipeRW) {
+	var (
+		c1, c2  = make(chan Msg), make(chan Msg)
+		closing = make(chan struct{})
+		closed  = new(int32)
+		rw1     = &MsgPipeRW{c1, c2, closing, closed}
+		rw2     = &MsgPipeRW{c2, c1, closing, closed}
+	)
+	return rw1, rw2
+}
+
+// ErrPipeClosed is returned from pipe operations after the
+// pipe has been closed.
+var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe")
+
+// MsgPipeRW is an endpoint of a MsgReadWriter pipe.
+type MsgPipeRW struct {
+	w       chan<- Msg
+	r       <-chan Msg
+	closing chan struct{}
+	closed  *int32
+}
+
+// WriteMsg sends a messsage on the pipe.
+// It blocks until the receiver has consumed the message payload.
+func (p *MsgPipeRW) WriteMsg(msg Msg) error {
+	if atomic.LoadInt32(p.closed) == 0 {
+		consumed := make(chan struct{}, 1)
+		msg.Payload = &eofSignal{msg.Payload, msg.Size, consumed}
+		select {
+		case p.w <- msg:
+			if msg.Size > 0 {
+				// wait for payload read or discard
+				<-consumed
+			}
+			return nil
+		case <-p.closing:
+		}
+	}
+	return ErrPipeClosed
+}
+
+// ReadMsg returns a message sent on the other end of the pipe.
+func (p *MsgPipeRW) ReadMsg() (Msg, error) {
+	if atomic.LoadInt32(p.closed) == 0 {
+		select {
+		case msg := <-p.r:
+			return msg, nil
+		case <-p.closing:
+		}
+	}
+	return Msg{}, ErrPipeClosed
+}
+
+// Close unblocks any pending ReadMsg and WriteMsg calls on both ends
+// of the pipe. They will return ErrPipeClosed. Note that Close does
+// not interrupt any reads from a message payload.
+func (p *MsgPipeRW) Close() error {
+	if atomic.AddInt32(p.closed, 1) != 1 {
+		// someone else is already closing
+		atomic.StoreInt32(p.closed, 1) // avoid overflow
+		return nil
+	}
+	close(p.closing)
+	return nil
+}