core-chain/p2p/message.go

package p2p

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"math/big"

	"github.com/ethereum/go-ethereum/ethutil"
)

type MsgCode uint64

// 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    MsgCode
	Size    uint32 // size of the paylod
	Payload io.Reader
}

// NewMsg creates an RLP-encoded message with the given code.
func NewMsg(code MsgCode, 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
}

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.
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(ReadError, "%v", err)
	}
	if !bytes.HasPrefix(start, magicToken) {
		return msg, NewPeerError(MagicTokenMismatch, "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 MsgCode, codelen uint32, err error) {
	b, err := r.ReadByte()
	if err != nil {
		return 0, 0, err
	}
	if b < 0x80 {
		return MsgCode(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 MsgCode(binary.BigEndian.Uint64(buf)), codelen, nil
	}
	return 0, 0, fmt.Errorf("bad RLP type for message code: %x", b)
}