package p2p
import (
"bytes"
"encoding/binary"
"fmt"
"io"
"io/ioutil"
"math/big"
"github.com/ethereum/go-ethereum/ethutil"
)
// 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()
}
// Data returns the decoded RLP payload items in a message.
func (msg Msg) Data() (*ethutil.Value, error) {
// TODO: avoid copying when we have a better RLP decoder
buf := new(bytes.Buffer)
var s []interface{}
if _, err := buf.ReadFrom(msg.Payload); err != nil {
return nil, err
}
for buf.Len() > 0 {
s = append(s, ethutil.DecodeWithReader(buf))
}
return ethutil.NewValue(s), nil
}
// 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 an existing message.
// The Payload reader of the message is consumed.
// Note that messages can be sent only once.
WriteMsg(Msg) error
// EncodeMsg writes an RLP-encoded message with the given
// code and data elements.
EncodeMsg(code uint64, data ...interface{}) error
}
// MsgReadWriter provides reading and writing of encoded messages.
type MsgReadWriter interface {
MsgReader
MsgWriter
}
// MsgLoop reads messages off the given reader and
// calls the handler function for each decoded message until
// it returns an error or the peer connection is closed.
//
// If a message is larger than the given maximum size,
// MsgLoop returns an appropriate error.
func MsgLoop(r MsgReader, maxsize uint32, f func(code uint64, data *ethutil.Value) error) error {
for {
msg, err := r.ReadMsg()
if err != nil {
return err
}
if msg.Size > maxsize {
return newPeerError(errInvalidMsg, "size %d exceeds maximum size of %d", msg.Size, maxsize)
}
value, err := msg.Data()
if err != nil {
return err
}
if err := f(msg.Code, value); err != nil {
return err
}
}
}
var magicToken = []byte{34, 64, 8, 145}
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...)
}
type byteReader interface {
io.Reader
io.ByteReader
}
// readMsg reads a message header from r.
func readMsg(r byteReader) (msg Msg, err error) {
// read magic and payload size
start := make([]byte, 8)
if _, err = io.ReadFull(r, start); err != nil {
return msg, newPeerError(errRead, "%v", err)
}
if !bytes.HasPrefix(start, magicToken) {
return msg, newPeerError(errMagicTokenMismatch, "got %x, want %x", start[:4], magicToken)
}
size := binary.BigEndian.Uint32(start[4:])
// decode start of RLP message to get the message code
_, hdrlen, err := readListHeader(r)
if err != nil {
return msg, err
}
code, codelen, err := readMsgCode(r)
if err != nil {
return msg, err
}
rlpsize := size - hdrlen - codelen
return Msg{
Code: code,
Size: rlpsize,
Payload: io.LimitReader(r, int64(rlpsize)),
}, nil
}
// readListHeader reads an RLP list header from r.
func readListHeader(r byteReader) (len uint64, hdrlen uint32, err error) {
b, err := r.ReadByte()
if err != nil {
return 0, 0, err
}
if b < 0xC0 {
return 0, 0, fmt.Errorf("expected list start byte >= 0xC0, got %x", b)
} else if b < 0xF7 {
len = uint64(b - 0xc0)
hdrlen = 1
} else {
lenlen := b - 0xF7
lenbuf := make([]byte, 8)
if _, err := io.ReadFull(r, lenbuf[8-lenlen:]); err != nil {
return 0, 0, err
}
len = binary.BigEndian.Uint64(lenbuf)
hdrlen = 1 + uint32(lenlen)
}
return len, hdrlen, nil
}
// readUint reads an RLP-encoded unsigned integer from r.
func readMsgCode(r byteReader) (code uint64, codelen uint32, err error) {
b, err := r.ReadByte()
if err != nil {
return 0, 0, err
}
if b < 0x80 {
return uint64(b), 1, nil
} else if b < 0x89 { // max length for uint64 is 8 bytes
codelen = uint32(b - 0x80)
if codelen == 0 {
return 0, 1, nil
}
buf := make([]byte, 8)
if _, err := io.ReadFull(r, buf[8-codelen:]); err != nil {
return 0, 0, err
}
return binary.BigEndian.Uint64(buf), codelen, nil
}
return 0, 0, fmt.Errorf("bad RLP type for message code: %x", b)
}
