forks/go-ethereum
1
0
Fork 1
mirror of https://github.com/ethereum/go-ethereum.git synced 2026-07-22 04:36:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Refactor message read/write

Browse source 
- take readwriter out of crypto
- cryptoId.NewSession now gets a reader and writer simply and wraps them in a encrypted MsgReadWriter which it returns if handshake is successful
- refactor reader/writers MsgReadWriter and Messenger
- message got NewMsgFromRLP method
- remove package-wide readMsg and writeMsg (they need to be scoped under the encryption scheme)
- writeProtoMsg back to peer
- peer has new field CryptoType
- add errAuthentication and errEncryption to peer Errors
- testPeer has to wait till handshake done (rw is set up) <-peer.cryptoReady
This commit is contained in:
zelig 2015-01-23 11:44:45 +00:00
parent ba8318db6a
commit cd235b24c5
9 changed files with 190 additions and 157 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
p2p/crypto.go
View file

@ -5,7 +5,6 @@ import (
"crypto/rand"
"fmt"
"io"
"net"
"github.com/ethereum/go-ethereum/crypto"
"github.com/ethereum/go-ethereum/crypto/secp256k1"
@ -26,13 +25,6 @@ var (
rHSLen int = 210 // size of the final ECIES payload sent as receiver's handshake
)
// secretRW implements a message read writer with encryption and authentication
// it is initialised by cryptoId.Run() after a successful crypto handshake
// aesSecret, macSecret, egressMac, ingress
type secretRW struct {
aesSecret, macSecret, egressMac, ingressMac []byte
}
/*
cryptoId implements the crypto layer for the p2p networking
It is initialised on the node's own identity (which has access to the node's private key) and run separately on a peer connection to set up a secure session after a crypto handshake
@ -81,7 +73,7 @@ func (self hexkey) String() string {
}
/*
Run(connection, remotePublicKey, sessionToken) is called when the peer connection starts to set up a secure session by performing a crypto handshake.
NewSession is called when the peer connection starts to set up a secure session by performing a crypto handshake.
connection is (a buffered) network connection.
@ -94,7 +86,7 @@ Run(connection, remotePublicKey, sessionToken) is called when the peer connectio
It returns a secretRW which implements the MsgReadWriter interface.
*/
func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rwF func(net.Conn) MsgChanReadWriter, err error) {
func (self *cryptoId) NewSession(r io.Reader, w io.Writer, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rw MsgReadWriter, err error) {
var auth, initNonce, recNonce []byte
var read int
var randomPrivKey *ecdsa.PrivateKey
@ -112,12 +104,12 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
randomPublicKeyS, _ := ExportPublicKey(&randomPrivKey.PublicKey)
clogger.Debugf("initiator-random-public-key: %v", hexkey(randomPublicKeyS))
if _, err = conn.Write(auth); err != nil {
if _, err = w.Write(auth); err != nil {
return
}
clogger.Debugf("initiator handshake (sent to %v):\n%v", hexkey(remotePubKeyS), hexkey(auth))
var response []byte = make([]byte, rHSLen)
if read, err = conn.Read(response); err != nil || read == 0 {
if read, err = r.Read(response); err != nil || read == 0 {
return
}
if read != rHSLen {
@ -136,7 +128,7 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
} else {
auth = make([]byte, iHSLen)
clogger.Debugf("waiting for initiator handshake (from %v)", hexkey(remotePubKeyS))
if read, err = conn.Read(auth); err != nil {
if read, err = r.Read(auth); err != nil {
return
}
if read != iHSLen {
@ -153,12 +145,12 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
}
clogger.Debugf("receiver-nonce: %v", hexkey(recNonce))
clogger.Debugf("receiver-random-priv-key: %v", hexkey(crypto.FromECDSA(randomPrivKey)))
if _, err = conn.Write(response); err != nil {
if _, err = w.Write(response); err != nil {
return
}
clogger.Debugf("receiver handshake (sent to %v):\n%v", hexkey(remotePubKeyS), hexkey(response))
}
return self.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
return self.newSession(r, w, initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
}
/*
@ -365,7 +357,7 @@ func (self *cryptoId) completeHandshake(auth []byte) (respNonce []byte, remoteRa
/*
newSession is called after the handshake is completed. The arguments are values negotiated in the handshake and the return value is a new session : a new session Token to be remembered for the next time we connect with this peer. And a MsgReadWriter that implements an encrypted and authenticated connection with key material obtained from the crypto handshake key exchange
*/
func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecdsa.PrivateKey, remoteRandomPubKey *ecdsa.PublicKey) (sessionToken []byte, rwF func(net.Conn) MsgChanReadWriter, err error) {
func (self *cryptoId) newSession(r io.Reader, w io.Writer, initNonce, respNonce, auth []byte, privKey *ecdsa.PrivateKey, remoteRandomPubKey *ecdsa.PublicKey) (sessionToken []byte, rw MsgReadWriter, err error) {
// 3) Now we can trust ecdhe-random-pubk to derive new keys
//ecdhe-shared-secret = ecdh.agree(ecdhe-random, remote-ecdhe-random-pubk)
var dhSharedSecret []byte
@ -395,9 +387,7 @@ func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecd
clogger.Debugf("egress-mac: %v", hexkey(egressMac))
clogger.Debugf("ingress-mac: %v", hexkey(ingressMac))
rwF = func(conn net.Conn) MsgChanReadWriter {
return NewSecureMessenger(conn, aesSecret, macSecret, egressMac, ingressMac)
}
rw, err = NewCryptoMsgRW(r, w, aesSecret, macSecret, egressMac, ingressMac)
return
}
@ -410,17 +400,3 @@ func Xor(one, other []byte) (xor []byte) {
}
return
}
type SecureMessenger struct {
Messenger
aesSecret, macSecret, egressMac, ingressMac []byte
}
func NewSecureMessenger(conn net.Conn, aesSecret, macSecret, egressMac, ingressMac []byte) *SecureMessenger {
return &SecureMessenger{
aesSecret: aesSecret,
macSecret: macSecret,
egressMac: egressMac,
ingressMac: ingressMac,
}
}

20
p2p/crypto_test.go
View file

@ -5,6 +5,7 @@ import (
// "crypto/ecdsa"
// "crypto/elliptic"
// "crypto/rand"
"bufio"
"fmt"
"net"
"testing"
@ -105,13 +106,15 @@ func TestCryptoHandshake(t *testing.T) {
t.Errorf("%v", err)
}
conn0, conn1 := net.Pipe()
// now both parties should have the same session parameters
initSessionToken, initSecretRWF, err := initiator.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
initSessionToken, initRW, err := initiator.newSession(bufio.NewReader(conn0), conn0, initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
if err != nil {
t.Errorf("%v", err)
}
recSessionToken, recSecretRWF, err := receiver.newSession(remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
recSessionToken, recRW, err := receiver.newSession(bufio.NewReader(conn1), conn1, remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
if err != nil {
t.Errorf("%v", err)
}
@ -130,14 +133,13 @@ func TestCryptoHandshake(t *testing.T) {
t.Errorf("session tokens do not match")
}
conn0, conn1 := net.Pipe()
initSecretRW, ok := initSecretRWF(conn0).(*SecureMessenger)
initSecretRW, ok := initRW.(*CryptoMsgRW)
if !ok {
t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
t.Errorf("incorrect return type. expected *CryptoMsgRW, got %T", initRW)
}
recSecretRW, ok := recSecretRWF(conn1).(*SecureMessenger)
recSecretRW, ok := recRW.(*CryptoMsgRW)
if !ok {
t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
t.Errorf("incorrect return type. expected *CryptoMsgRW, got %T", recRW)
}
// aesSecret, macSecret, egressMac, ingressMac
@ -184,8 +186,8 @@ func TestPeersHandshake(t *testing.T) {
initiator.pubkeyHook = func(*peerAddr) error { return nil }
receiver.pubkeyHook = func(*peerAddr) error { return nil }
initiator.cryptoHandshake = true
receiver.cryptoHandshake = true
initiator.CryptoType = EthCrypto
receiver.CryptoType = EthCrypto
errc0 := make(chan error, 1)
errc1 := make(chan error, 1)
go func() {

15
p2p/message.go
View file

@ -65,6 +65,21 @@ func (msg Msg) Discard() error {
return err
}
func NewMsgFromRLP(size uint32, r rlp.ByteReader) (msg Msg, err error) {
// decode start of RLP message to get the message code
posr := &postrack{r, 0}
s := rlp.NewStream(posr)
if _, err := s.List(); err != nil {
return msg, err
}
code, err := s.Uint()
if err != nil {
return msg, err
}
payloadsize := size - posr.p
return Msg{code, payloadsize, io.LimitReader(r, int64(payloadsize))}, nil
}
/*
MsgReadWriter is an interface for reading and writing messages
It is aware of message structure and knows how to encode/decode

9
p2p/message_test.go
View file

@ -29,12 +29,13 @@ func TestNewMsg(t *testing.T) {
func TestEncodeDecodeMsg(t *testing.T) {
msg := NewMsg(3, 1, "000")
buf := new(bytes.Buffer)
if err := writeMsg(buf, msg); err != nil {
rw, _ := NewMsgRW(buf, buf)
if err := rw.WriteMsg(msg); err != nil {
t.Fatalf("encodeMsg error: %v", err)
}
// t.Logf("encoded: %x", buf.Bytes())
decmsg, err := readMsg(buf)
decmsg, err := rw.ReadMsg()
if err != nil {
t.Fatalf("readMsg error: %v", err)
}
@ -62,7 +63,9 @@ func TestEncodeDecodeMsg(t *testing.T) {
func TestDecodeRealMsg(t *testing.T) {
data := ethutil.Hex2Bytes("2240089100000080f87e8002b5457468657265756d282b2b292f5065657220536572766572204f6e652f76302e372e382f52656c656173652f4c696e75782f672b2bc082765fb84086dd80b7aefd6a6d2e3b93f4f300a86bfb6ef7bdc97cb03f793db6bb")
msg, err := readMsg(bytes.NewReader(data))
buf := bytes.NewReader(data)
rw, _ := NewMsgRW(buf, nil)
msg, err := rw.ReadMsg()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}

99
p2p/messenger.go
View file

@ -1,13 +1,10 @@
package p2p
import (
"bufio"
"bytes"
"encoding/binary"
"fmt"
"io"
"math/big"
"net"
"time"
"github.com/ethereum/go-ethereum/ethutil"
@ -27,15 +24,12 @@ const (
var magicToken = []byte{34, 64, 8, 145}
/*
writeMsg, readMsg (makeListHeader) should all be internal to the default legacy
MsgReadWriter implementation
A MsgChanReadWriter implementation will typically sit on a multiplexed peer connection and runs a single read and a write loop without need to use locking.
A MsgReadWriter implementation will typically sit on a multiplexed peer connection and runs a single read and a write loop without need to use locking.
It passes on incoming messages to its channel picked up by the interface methods Read
It passes on incoming messages to its channel ReadC()
the peer runs a dispatch loop that figures out which protocol to forward the message to.
Because framing and header structure will change there will be hardly any overlap with the new code so I do not abstract readers any further.
The channel for outcgoing messages (WriteC) is simply shared between the individual MsgReadWriter instances for each protocol
*/
type MsgChanReadWriter interface {
@ -47,26 +41,24 @@ type MsgChanReadWriter interface {
}
type Messenger struct {
msgReadTimeout time.Duration
msgWriteTimeout time.Duration
in chan Msg
out chan Msg
errc chan error
unblock chan bool
conn net.Conn
bufconn *bufio.ReadWriter
rw MsgReadWriter
}
func NewMessenger(conn net.Conn) *Messenger {
/*
Messenger is a simple implementation of a read and write loop using a MsgReadWriter to encode/decode individual messages
This MsgReadWriter can implement parsing from/to any kind of packet structure and employ encryption and authentication
*/
func NewMessenger(rw MsgReadWriter) *Messenger {
self := &Messenger{
msgReadTimeout: msgReadTimeout,
msgWriteTimeout: msgWriteTimeout,
in: make(chan Msg),
out: make(chan Msg),
errc: make(chan error),
unblock: make(chan bool, 1),
conn: conn,
bufconn: bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
rw: rw,
}
go self.readLoop()
go self.writeLoop()
@ -90,14 +82,14 @@ func (self *Messenger) ErrorC() chan error {
return self.errc
}
// ReadNextC <- true must be called before the next read is attempted
func (self *Messenger) ReadNextC() chan bool {
return self.unblock
}
func (self *Messenger) readLoop() {
for _ = range self.unblock {
self.conn.SetReadDeadline(time.Now().Add(self.msgReadTimeout))
if msg, err := readMsg(self.bufconn); err != nil {
if msg, err := self.rw.ReadMsg(); err != nil {
self.errc <- err
} else {
self.in <- msg
@ -108,15 +100,35 @@ func (self *Messenger) readLoop() {
func (self *Messenger) writeLoop() {
for msg := range self.out {
self.conn.SetWriteDeadline(time.Now().Add(self.msgWriteTimeout))
if err := writeMsg(self.bufconn, msg); err != nil {
if err := self.rw.WriteMsg(msg); err != nil {
self.errc <- newPeerError(errWrite, "%v", err)
}
self.bufconn.Flush()
}
}
func writeMsg(w io.Writer, msg Msg) error {
/*
MsgReadWriter is an interface for reading and writing messages
It is aware of message structure and knows how to encode/decode
MsgRW is a simple encoder implementing MsgReadWriter
It complies with the legacy devp2p packet structure and no encryption or authentication
*/
type MsgRW struct {
r rlp.ByteReader // this is implemented by bufio.ReadWriter
// r io.Reader
w io.Writer
}
func NewMsgRW(r rlp.ByteReader, w io.Writer) (*MsgRW, error) {
return &MsgRW{
r: r,
w: w,
}, nil
}
func (self *MsgRW) WriteMsg(msg Msg) error {
// TODO: handle case when Size + len(code) + len(listhdr) overflows uint32
code := ethutil.Encode(uint32(msg.Code))
listhdr := makeListHeader(msg.Size + uint32(len(code)))
@ -127,11 +139,11 @@ func writeMsg(w io.Writer, msg Msg) error {
binary.BigEndian.PutUint32(start[4:], payloadLen)
for _, b := range [][]byte{start, listhdr, code} {
if _, err := w.Write(b); err != nil {
if _, err := self.w.Write(b); err != nil {
return err
}
}
_, err := io.CopyN(w, msg.Payload, int64(msg.Size))
_, err := io.CopyN(self.w, msg.Payload, int64(msg.Size))
return err
}
@ -146,29 +158,19 @@ func makeListHeader(length uint32) []byte {
// readMsg reads a message header from r.
// It takes an rlp.ByteReader to ensure that the decoding doesn't buffer.
func readMsg(r rlp.ByteReader) (msg Msg, err error) {
func (self *MsgRW) ReadMsg() (msg Msg, err error) {
// read magic and payload size
start := make([]byte, 8)
if _, err = io.ReadFull(r, start); err != nil {
if _, err = io.ReadFull(self.r, start); err != nil {
return msg, newPeerError(errRead, "%v", err)
}
if !bytes.HasPrefix(start, magicToken) {
return msg, newPeerError(errMagicTokenMismatch, "got %x, want %x", start[:4], magicToken)
}
size := binary.BigEndian.Uint32(start[4:])
// decode start of RLP message to get the message code
posr := &postrack{r, 0}
s := rlp.NewStream(posr)
if _, err := s.List(); err != nil {
return msg, err
}
code, err := s.Uint()
if err != nil {
return msg, err
}
payloadsize := size - posr.p
return Msg{code, payloadsize, io.LimitReader(r, int64(payloadsize))}, nil
return NewMsgFromRLP(size, self.r)
}
// postrack wraps an rlp.ByteReader with a position counter.
@ -191,20 +193,6 @@ func (r *postrack) ReadByte() (byte, error) {
return b, err
}
// this duplicates functionality of proto.WriteMsg
// if we need this for broadcasting via a server interface then
// simply call the appropriate Write function of the protocol RW
// writeProtoMsg sends the given message on behalf of the given named protocol.
func (p *Peer) writeProtoMsg(protoName string, msg Msg) error {
p.runlock.RLock()
proto, ok := p.running[protoName]
p.runlock.RUnlock()
if !ok {
return fmt.Errorf("protocol %s not handled by peer", protoName)
}
return proto.WriteMsg(msg)
}
// proto will embed the same writer channel as given to the readwriter
// in the legacy code it knows about the code offset
// no need to go through peer for writing , so do not need to embed peer as field
@ -214,6 +202,7 @@ type proto struct {
maxcode, offset uint64
}
// WriteMsg proto implements MsgWriter interface
func (rw *proto) WriteMsg(msg Msg) error {
if msg.Code >= rw.maxcode {
return newPeerError(errInvalidMsgCode, "not handled")

64
p2p/peer.go
View file

@ -1,6 +1,7 @@
package p2p
import (
"bufio"
"bytes"
"crypto/rand"
"fmt"
@ -66,7 +67,7 @@ type Peer struct {
// These fields maintain the running protocols.
protocols []Protocol
runBaseProtocol bool // for testing
cryptoHandshake bool // for testing
CryptoType CryptoType //
cryptoReady chan struct{}
runlock sync.RWMutex // protects running
@ -209,17 +210,12 @@ func (p *Peer) loop() (reason DiscReason, err error) {
defer close(p.closed)
defer p.conn.Close()
if p.cryptoHandshake {
if crwF, err := p.handleCryptoHandshake(); err != nil {
if err = p.handleCryptoHandshake(); err != nil {
// from here on everything can be encrypted, authenticated
return DiscProtocolError, err // no graceful disconnect
} else {
p.crw = crwF(p.conn)
}
} else {
p.crw = NewMessenger(p.conn)
}
close(p.cryptoReady)
defer p.crw.Close()
// read loop
protoDone := make(chan struct{}, 1)
@ -271,7 +267,6 @@ loop:
// tell the remote end to disconnect
p.writeProtoMsg("", NewMsg(discMsg, reason))
// io.Copy(ioutil.Discard, p.conn)//??
p.crw.Close()
<-time.After(disconnectGracePeriod)
return reason, err
}
@ -299,9 +294,34 @@ func (p *Peer) dispatch(msg Msg, protoDone chan struct{}) (wait bool, err error)
return wait, nil
}
type readLoop func(chan<- Msg, chan<- error, <-chan bool)
type CryptoType byte
func (p *Peer) handleCryptoHandshake() (crw func(net.Conn) MsgChanReadWriter, err error) {
const (
NoCrypto CryptoType = iota
EthCrypto
)
var cryptoType = map[CryptoType]string{
NoCrypto: "no encryption",
EthCrypto: "AES256 CTR HMAC SHA256",
}
func (self CryptoType) String() (s string) {
s = cryptoType[self]
if len(s) == 0 {
s = string([]byte{byte(self)})
}
return
}
func (p *Peer) handleCryptoHandshake() (err error) {
var crw MsgReadWriter
switch p.CryptoType {
case NoCrypto:
if crw, err = NewMsgRW(bufio.NewReader(p.conn), p.conn); err != nil {
return
}
case EthCrypto:
// cryptoId is just created for the lifecycle of the handshake
// it is survived by an encrypted readwriter
var initiator bool
@ -322,11 +342,14 @@ func (p *Peer) handleCryptoHandshake() (crw func(net.Conn) MsgChanReadWriter, er
}
// run on peer
// this bit handles the handshake and creates a secure communications channel with
// var rw *secretRW
if sessionToken, crw, err = crypto.NewSession(p.conn, p.PublicKey(), sessionToken, initiator); err != nil {
p.Debugf("unable to setup secure session: %v", err)
return
if sessionToken, crw, err = crypto.NewSession(bufio.NewReader(p.conn), p.conn, p.PublicKey(), sessionToken, initiator); err != nil {
p.Errorf("unable to setup secure session: %v", err)
}
default:
err = fmt.Errorf("unrecognised crypto type %v", p.CryptoType)
p.Errorf("%v", err)
}
p.crw = NewMessenger(crw)
return
}
@ -406,3 +429,14 @@ func (p *Peer) closeProtocols() {
p.runlock.RUnlock()
p.protoWG.Wait()
}
// writeProtoMsg sends the given message on behalf of the given named protocol.
func (p *Peer) writeProtoMsg(protoName string, msg Msg) error {
p.runlock.RLock()
proto, ok := p.running[protoName]
p.runlock.RUnlock()
if !ok {
return fmt.Errorf("protocol %s not handled by peer", protoName)
}
return proto.WriteMsg(msg)
}

4
p2p/peer_error.go
View file

@ -19,6 +19,8 @@ const (
errPingTimeout
errInvalidNetworkId
errInvalidProtocolVersion
errAuthentication
errEncryption
)
var errorToString = map[int]string{
@ -36,6 +38,8 @@ var errorToString = map[int]string{
errPingTimeout: "Ping timeout",
errInvalidNetworkId: "Invalid network id",
errInvalidProtocolVersion: "Invalid protocol version",
errAuthentication: "Authentication error",
errEncryption: "Encryption error",
}
type peerError struct {

23
p2p/peer_test.go
View file

@ -33,12 +33,12 @@ func testPeer(protos []Protocol) (net.Conn, *Peer, <-chan error) {
peer := newPeer(conn1, protos, nil)
peer.ourID = &peerId{}
peer.pubkeyHook = func(*peerAddr) error { return nil }
peer.crw = NewMessenger(conn1)
errc := make(chan error, 1)
go func() {
_, err := peer.loop()
errc <- err
}()
<-peer.cryptoReady
return conn2, peer, errc
}
@ -72,9 +72,11 @@ func TestPeerProtoReadMsg(t *testing.T) {
net, peer, errc := testPeer([]Protocol{proto})
defer net.Close()
peer.startSubprotocols([]Cap{proto.cap()})
writeMsg(net, NewMsg(18, 1, "000"))
rw, _ := NewMsgRW(nil, net)
rw.WriteMsg(NewMsg(18, 1, "000"))
select {
case <-done:
case err := <-errc:
@ -108,9 +110,11 @@ func TestPeerProtoReadLargeMsg(t *testing.T) {
net, peer, errc := testPeer([]Protocol{proto})
defer net.Close()
peer.startSubprotocols([]Cap{proto.cap()})
writeMsg(net, NewMsg(18, make([]byte, msgsize)))
rw, _ := NewMsgRW(nil, net)
peer.startSubprotocols([]Cap{proto.cap()})
rw.WriteMsg(NewMsg(18, make([]byte, msgsize)))
select {
case <-done:
case err := <-errc:
@ -138,10 +142,11 @@ func TestPeerProtoEncodeMsg(t *testing.T) {
}
net, peer, _ := testPeer([]Protocol{proto})
defer net.Close()
peer.startSubprotocols([]Cap{proto.cap()})
bufr := bufio.NewReader(net)
msg, err := readMsg(bufr)
rw, _ := NewMsgRW(bufr, nil)
peer.startSubprotocols([]Cap{proto.cap()})
msg, err := rw.ReadMsg()
if err != nil {
t.Errorf("read error: %v", err)
}
@ -178,7 +183,8 @@ func TestPeerWrite(t *testing.T) {
read := make(chan struct{})
go func() {
bufr := bufio.NewReader(net)
msg, err := readMsg(bufr)
rw, _ := NewMsgRW(bufr, nil)
msg, err := rw.ReadMsg()
if err != nil {
t.Errorf("read error: %v", err)
} else if msg.Code != 16 {
@ -212,8 +218,9 @@ func TestPeerActivity(t *testing.T) {
sub := peer.activity.Subscribe(time.Time{})
defer sub.Unsubscribe()
rw, _ := NewMsgRW(nil, net)
for i := 0; i < 6; i++ {
writeMsg(net, NewMsg(16))
rw.WriteMsg(NewMsg(16))
select {
case <-sub.Chan():
case <-time.After(inactivityTimeout / 2):

11
p2p/server_test.go
View file

@ -1,6 +1,7 @@
package p2p
import (
"bufio"
"bytes"
"io"
"net"
@ -121,7 +122,8 @@ func TestServerBroadcast(t *testing.T) {
var connected sync.WaitGroup
srv := startTestServer(t, func(srv *Server, c net.Conn, dialAddr *peerAddr) *Peer {
peer := newPeer(c, []Protocol{discard}, dialAddr)
peer.crw = NewMessenger(peer.conn)
rw, _ := NewMsgRW(bufio.NewReader(c), c)
peer.crw = NewMessenger(rw)
peer.startSubprotocols([]Cap{discard.cap()})
connected.Done()
return peer
@ -146,9 +148,10 @@ func TestServerBroadcast(t *testing.T) {
// broadcast one message
srv.Broadcast("discard", 0, "foo")
goldbuf := new(bytes.Buffer)
writeMsg(goldbuf, NewMsg(16, "foo"))
golden := goldbuf.Bytes()
buf := new(bytes.Buffer)
rw, _ := NewMsgRW(buf, buf)
rw.WriteMsg(NewMsg(16, "foo"))
golden := buf.Bytes()
// check that the message has been written everywhere
for i, conn := range conns {