1 files changed, 174 insertions, 0 deletions
diff --git a/p2p/rlpx.go b/p2p/rlpx.go
new file mode 100644
index 000000000..6b533e275
--- /dev/null
+++ b/p2p/rlpx.go
@@ -0,0 +1,174 @@
+package p2p
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/hmac"
+	"errors"
+	"hash"
+	"io"
+
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+var (
+	// this is used in place of actual frame header data.
+	// TODO: replace this when Msg contains the protocol type code.
+	zeroHeader = []byte{0xC2, 0x80, 0x80}
+
+	// sixteen zero bytes
+	zero16 = make([]byte, 16)
+
+	maxUint24 = ^uint32(0) >> 8
+)
+
+// rlpxFrameRW implements a simplified version of RLPx framing.
+// chunked messages are not supported and all headers are equal to
+// zeroHeader.
+//
+// rlpxFrameRW is not safe for concurrent use from multiple goroutines.
+type rlpxFrameRW struct {
+	conn io.ReadWriter
+	enc  cipher.Stream
+	dec  cipher.Stream
+
+	macCipher  cipher.Block
+	egressMAC  hash.Hash
+	ingressMAC hash.Hash
+}
+
+func newRlpxFrameRW(conn io.ReadWriter, s secrets) *rlpxFrameRW {
+	macc, err := aes.NewCipher(s.MAC)
+	if err != nil {
+		panic("invalid MAC secret: " + err.Error())
+	}
+	encc, err := aes.NewCipher(s.AES)
+	if err != nil {
+		panic("invalid AES secret: " + err.Error())
+	}
+	// we use an all-zeroes IV for AES because the key used
+	// for encryption is ephemeral.
+	iv := make([]byte, encc.BlockSize())
+	return &rlpxFrameRW{
+		conn:       conn,
+		enc:        cipher.NewCTR(encc, iv),
+		dec:        cipher.NewCTR(encc, iv),
+		macCipher:  macc,
+		egressMAC:  s.EgressMAC,
+		ingressMAC: s.IngressMAC,
+	}
+}
+
+func (rw *rlpxFrameRW) WriteMsg(msg Msg) error {
+	ptype, _ := rlp.EncodeToBytes(msg.Code)
+
+	// write header
+	headbuf := make([]byte, 32)
+	fsize := uint32(len(ptype)) + msg.Size
+	if fsize > maxUint24 {
+		return errors.New("message size overflows uint24")
+	}
+	putInt24(fsize, headbuf) // TODO: check overflow
+	copy(headbuf[3:], zeroHeader)
+	rw.enc.XORKeyStream(headbuf[:16], headbuf[:16]) // first half is now encrypted
+
+	// write header MAC
+	copy(headbuf[16:], updateMAC(rw.egressMAC, rw.macCipher, headbuf[:16]))
+	if _, err := rw.conn.Write(headbuf); err != nil {
+		return err
+	}
+
+	// write encrypted frame, updating the egress MAC hash with
+	// the data written to conn.
+	tee := cipher.StreamWriter{S: rw.enc, W: io.MultiWriter(rw.conn, rw.egressMAC)}
+	if _, err := tee.Write(ptype); err != nil {
+		return err
+	}
+	if _, err := io.Copy(tee, msg.Payload); err != nil {
+		return err
+	}
+	if padding := fsize % 16; padding > 0 {
+		if _, err := tee.Write(zero16[:16-padding]); err != nil {
+			return err
+		}
+	}
+
+	// write frame MAC. egress MAC hash is up to date because
+	// frame content was written to it as well.
+	fmacseed := rw.egressMAC.Sum(nil)
+	mac := updateMAC(rw.egressMAC, rw.macCipher, fmacseed)
+	_, err := rw.conn.Write(mac)
+	return err
+}
+
+func (rw *rlpxFrameRW) ReadMsg() (msg Msg, err error) {
+	// read the header
+	headbuf := make([]byte, 32)
+	if _, err := io.ReadFull(rw.conn, headbuf); err != nil {
+		return msg, err
+	}
+	// verify header mac
+	shouldMAC := updateMAC(rw.ingressMAC, rw.macCipher, headbuf[:16])
+	if !hmac.Equal(shouldMAC, headbuf[16:]) {
+		return msg, errors.New("bad header MAC")
+	}
+	rw.dec.XORKeyStream(headbuf[:16], headbuf[:16]) // first half is now decrypted
+	fsize := readInt24(headbuf)
+	// ignore protocol type for now
+
+	// read the frame content
+	var rsize = fsize // frame size rounded up to 16 byte boundary
+	if padding := fsize % 16; padding > 0 {
+		rsize += 16 - padding
+	}
+	framebuf := make([]byte, rsize)
+	if _, err := io.ReadFull(rw.conn, framebuf); err != nil {
+		return msg, err
+	}
+
+	// read and validate frame MAC. we can re-use headbuf for that.
+	rw.ingressMAC.Write(framebuf)
+	fmacseed := rw.ingressMAC.Sum(nil)
+	if _, err := io.ReadFull(rw.conn, headbuf[:16]); err != nil {
+		return msg, err
+	}
+	shouldMAC = updateMAC(rw.ingressMAC, rw.macCipher, fmacseed)
+	if !hmac.Equal(shouldMAC, headbuf[:16]) {
+		return msg, errors.New("bad frame MAC")
+	}
+
+	// decrypt frame content
+	rw.dec.XORKeyStream(framebuf, framebuf)
+
+	// decode message code
+	content := bytes.NewReader(framebuf[:fsize])
+	if err := rlp.Decode(content, &msg.Code); err != nil {
+		return msg, err
+	}
+	msg.Size = uint32(content.Len())
+	msg.Payload = content
+	return msg, nil
+}
+
+// updateMAC reseeds the given hash with encrypted seed.
+// it returns the first 16 bytes of the hash sum after seeding.
+func updateMAC(mac hash.Hash, block cipher.Block, seed []byte) []byte {
+	aesbuf := make([]byte, aes.BlockSize)
+	block.Encrypt(aesbuf, mac.Sum(nil))
+	for i := range aesbuf {
+		aesbuf[i] ^= seed[i]
+	}
+	mac.Write(aesbuf)
+	return mac.Sum(nil)[:16]
+}
+
+func readInt24(b []byte) uint32 {
+	return uint32(b[2]) | uint32(b[1])<<8 | uint32(b[0])<<16
+}
+
+func putInt24(v uint32, b []byte) {
+	b[0] = byte(v >> 16)
+	b[1] = byte(v >> 8)
+	b[2] = byte(v)
+}