diff options
Diffstat (limited to 'p2p')
| -rw-r--r-- | p2p/rlpx.go | 129 | ||||
| -rw-r--r-- | p2p/rlpx_test.go | 123 |
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) + } + } +} |