p2p: add basic RLPx frame I/O

2 files changed, 252 insertions, 0 deletions

diff --git a/p2p/rlpx.go b/p2p/rlpx.go
new file mode 100644
index 000000000..9fd1aed1f
--- /dev/null
+++ b/p2p/rlpx.go
@@ -0,0 +1,129 @@
+package p2p
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/hmac"
+	"errors"
+	"hash"
+	"io"
+
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+var (
+	zeroHeader = []byte{0xC2, 0x80, 0x80}
+	zero16     = make([]byte, 16)
+)
+
+type rlpxFrameRW struct {
+	conn io.ReadWriter
+
+	macCipher  cipher.Block
+	egressMAC  hash.Hash
+	ingressMAC hash.Hash
+}
+
+func newRlpxFrameRW(conn io.ReadWriter, macSecret []byte, egressMAC, ingressMAC hash.Hash) *rlpxFrameRW {
+	cipher, err := aes.NewCipher(macSecret)
+	if err != nil {
+		panic("invalid macSecret: " + err.Error())
+	}
+	return &rlpxFrameRW{conn: conn, macCipher: cipher, egressMAC: egressMAC, ingressMAC: 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
+	putInt24(fsize, headbuf) // TODO: check overflow
+	copy(headbuf[3:], zeroHeader)
+	copy(headbuf[16:], updateHeaderMAC(rw.egressMAC, rw.macCipher, headbuf[:16]))
+	if _, err := rw.conn.Write(headbuf); err != nil {
+		return err
+	}
+
+	// write frame, updating the egress MAC while writing to conn.
+	tee := 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 packet-mac. egress MAC is up to date because
+	// frame content was written to it as well.
+	_, err := rw.conn.Write(rw.egressMAC.Sum(nil))
+	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
+	}
+	fsize := readInt24(headbuf)
+	// ignore protocol type for now
+	shouldMAC := updateHeaderMAC(rw.ingressMAC, rw.macCipher, headbuf[:16])
+	if !hmac.Equal(shouldMAC[:16], headbuf[16:]) {
+		return msg, errors.New("bad header MAC")
+	}
+
+	// read the frame content
+	framebuf := make([]byte, fsize)
+	if _, err := io.ReadFull(rw.conn, framebuf); err != nil {
+		return msg, err
+	}
+	rw.ingressMAC.Write(framebuf)
+	if padding := fsize % 16; padding > 0 {
+		if _, err := io.CopyN(rw.ingressMAC, rw.conn, int64(16-padding)); err != nil {
+			return msg, err
+		}
+	}
+	// read and validate frame MAC. we can re-use headbuf for that.
+	if _, err := io.ReadFull(rw.conn, headbuf); err != nil {
+		return msg, err
+	}
+	if !hmac.Equal(rw.ingressMAC.Sum(nil), headbuf) {
+		return msg, errors.New("bad frame MAC")
+	}
+
+	// decode message code
+	content := bytes.NewReader(framebuf)
+	if err := rlp.Decode(content, &msg.Code); err != nil {
+		return msg, err
+	}
+	msg.Size = uint32(content.Len())
+	msg.Payload = content
+	return msg, nil
+}
+
+func updateHeaderMAC(mac hash.Hash, block cipher.Block, header []byte) []byte {
+	aesbuf := make([]byte, aes.BlockSize)
+	block.Encrypt(aesbuf, mac.Sum(nil))
+	for i := range aesbuf {
+		aesbuf[i] ^= header[i]
+	}
+	mac.Write(aesbuf)
+	return mac.Sum(nil)
+}
+
+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)
+}
diff --git a/p2p/rlpx_test.go b/p2p/rlpx_test.go
new file mode 100644
index 000000000..380d9aba6
--- /dev/null
+++ b/p2p/rlpx_test.go
@@ -0,0 +1,123 @@
+package p2p
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/crypto/sha3"
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+func TestRlpxFrameFake(t *testing.T) {
+	buf := new(bytes.Buffer)
+	secret := crypto.Sha3()
+	hash := fakeHash([]byte{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
+	rw := newRlpxFrameRW(buf, secret, hash, hash)
+
+	golden := unhex(`
+000006C2808000000000000000000000
+01010101010101010101010101010101
+08C40102030400000000000000000000
+01010101010101010101010101010101
+01010101010101010101010101010101
+`)
+
+	// Check WriteMsg. This puts a message into the buffer.
+	if err := EncodeMsg(rw, 8, []interface{}{1, 2, 3, 4}); err != nil {
+		t.Fatalf("WriteMsg error: %v", err)
+	}
+	written := buf.Bytes()
+	if !bytes.Equal(written, golden) {
+		t.Fatalf("output mismatch:\n  got:  %x\n  want: %x", written, golden)
+	}
+
+	// Check ReadMsg. It reads the message encoded by WriteMsg, which
+	// is equivalent to the golden message above.
+	msg, err := rw.ReadMsg()
+	if err != nil {
+		t.Fatalf("ReadMsg error: %v", err)
+	}
+	if msg.Size != 5 {
+		t.Errorf("msg size mismatch: got %d, want %d", msg.Size, 5)
+	}
+	if msg.Code != 8 {
+		t.Errorf("msg code mismatch: got %d, want %d", msg.Code, 8)
+	}
+	payload, _ := ioutil.ReadAll(msg.Payload)
+	wantPayload := unhex("C401020304")
+	if !bytes.Equal(payload, wantPayload) {
+		t.Errorf("msg payload mismatch:\ngot  %x\nwant %x", payload, wantPayload)
+	}
+}
+
+type fakeHash []byte
+
+func (fakeHash) Write(p []byte) (int, error) { return len(p), nil }
+func (fakeHash) Reset()                      {}
+func (fakeHash) BlockSize() int              { return 0 }
+
+func (h fakeHash) Size() int           { return len(h) }
+func (h fakeHash) Sum(b []byte) []byte { return append(b, h...) }
+
+func unhex(str string) []byte {
+	b, err := hex.DecodeString(strings.Replace(str, "\n", "", -1))
+	if err != nil {
+		panic(fmt.Sprintf("invalid hex string: %q", str))
+	}
+	return b
+}
+
+func TestRlpxFrameRW(t *testing.T) {
+	var (
+		macSecret      = make([]byte, 16)
+		egressMACinit  = make([]byte, 32)
+		ingressMACinit = make([]byte, 32)
+	)
+	for _, s := range [][]byte{macSecret, egressMACinit, ingressMACinit} {
+		rand.Read(s)
+	}
+
+	conn := new(bytes.Buffer)
+
+	em1 := sha3.NewKeccak256()
+	em1.Write(egressMACinit)
+	im1 := sha3.NewKeccak256()
+	im1.Write(ingressMACinit)
+	rw1 := newRlpxFrameRW(conn, macSecret, em1, im1)
+
+	em2 := sha3.NewKeccak256()
+	em2.Write(ingressMACinit)
+	im2 := sha3.NewKeccak256()
+	im2.Write(egressMACinit)
+	rw2 := newRlpxFrameRW(conn, macSecret, em2, im2)
+
+	// send some messages
+	for i := 0; i < 10; i++ {
+		// write message into conn buffer
+		wmsg := []interface{}{"foo", "bar", strings.Repeat("test", i)}
+		err := EncodeMsg(rw1, uint64(i), wmsg)
+		if err != nil {
+			t.Fatalf("WriteMsg error (i=%d): %v", i, err)
+		}
+
+		// read message that rw1 just wrote
+		msg, err := rw2.ReadMsg()
+		if err != nil {
+			t.Fatalf("ReadMsg error (i=%d): %v", i, err)
+		}
+		if msg.Code != uint64(i) {
+			t.Fatalf("msg code mismatch: got %d, want %d", msg.Code, i)
+		}
+		payload, _ := ioutil.ReadAll(msg.Payload)
+		wantPayload, _ := rlp.EncodeToBytes(wmsg)
+		if !bytes.Equal(payload, wantPayload) {
+			t.Fatalf("msg payload mismatch:\ngot  %x\nwant %x", payload, wantPayload)
+		}
+	}
+}