forked from cerc-io/plugeth
7964f30dcb
With RLPx frames, the message code is contained in the frame and is no longer part of the encoded data. EncodeMsg, Msg.Decode have been updated to match. Code that decodes RLP directly from Msg.Payload will need to change.
125 lines
3.3 KiB
Go
125 lines
3.3 KiB
Go
package p2p
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/rand"
|
|
"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)
|
|
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, secrets{
|
|
AES: crypto.Sha3(),
|
|
MAC: crypto.Sha3(),
|
|
IngressMAC: hash,
|
|
EgressMAC: hash,
|
|
})
|
|
|
|
golden := unhex(`
|
|
00828ddae471818bb0bfa6b551d1cb42
|
|
01010101010101010101010101010101
|
|
ba628a4ba590cb43f7848f41c4382885
|
|
01010101010101010101010101010101
|
|
`)
|
|
|
|
// Check WriteMsg. This puts a message into the buffer.
|
|
if err := EncodeMsg(rw, 8, 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 TestRlpxFrameRW(t *testing.T) {
|
|
var (
|
|
aesSecret = make([]byte, 16)
|
|
macSecret = make([]byte, 16)
|
|
egressMACinit = make([]byte, 32)
|
|
ingressMACinit = make([]byte, 32)
|
|
)
|
|
for _, s := range [][]byte{aesSecret, macSecret, egressMACinit, ingressMACinit} {
|
|
rand.Read(s)
|
|
}
|
|
conn := new(bytes.Buffer)
|
|
|
|
s1 := secrets{
|
|
AES: aesSecret,
|
|
MAC: macSecret,
|
|
EgressMAC: sha3.NewKeccak256(),
|
|
IngressMAC: sha3.NewKeccak256(),
|
|
}
|
|
s1.EgressMAC.Write(egressMACinit)
|
|
s1.IngressMAC.Write(ingressMACinit)
|
|
rw1 := newRlpxFrameRW(conn, s1)
|
|
|
|
s2 := secrets{
|
|
AES: aesSecret,
|
|
MAC: macSecret,
|
|
EgressMAC: sha3.NewKeccak256(),
|
|
IngressMAC: sha3.NewKeccak256(),
|
|
}
|
|
s2.EgressMAC.Write(ingressMACinit)
|
|
s2.IngressMAC.Write(egressMACinit)
|
|
rw2 := newRlpxFrameRW(conn, s2)
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
}
|