p2p: add basic RLPx frame I/O
This commit is contained in:
Felix Lange
parent
15f491e500
commit
936dd0f3bc
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue