diff --git a/p2p/crypto.go b/p2p/crypto.go index 242bad7977..e7ad29ec39 100644 --- a/p2p/crypto.go +++ b/p2p/crypto.go @@ -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, - } -} diff --git a/p2p/crypto_test.go b/p2p/crypto_test.go index 582eabe14b..2cb9c9bbe4 100644 --- a/p2p/crypto_test.go +++ b/p2p/crypto_test.go @@ -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() { diff --git a/p2p/message.go b/p2p/message.go index 09eefb7658..e4406adc39 100644 --- a/p2p/message.go +++ b/p2p/message.go @@ -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 diff --git a/p2p/message_test.go b/p2p/message_test.go index 5cde9abf5b..34b30065f1 100644 --- a/p2p/message_test.go +++ b/p2p/message_test.go @@ -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) } diff --git a/p2p/messenger.go b/p2p/messenger.go index 2999ebda9b..2bed6228b3 100644 --- a/p2p/messenger.go +++ b/p2p/messenger.go @@ -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 + in chan Msg + out chan Msg + errc chan error + unblock chan bool + 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)), + in: make(chan Msg), + out: make(chan Msg), + errc: make(chan error), + unblock: make(chan bool, 1), + 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") diff --git a/p2p/peer.go b/p2p/peer.go index 119e98e128..7704c764b8 100644 --- a/p2p/peer.go +++ b/p2p/peer.go @@ -1,6 +1,7 @@ package p2p import ( + "bufio" "bytes" "crypto/rand" "fmt" @@ -65,8 +66,8 @@ type Peer struct { // These fields maintain the running protocols. protocols []Protocol - runBaseProtocol bool // for testing - cryptoHandshake bool // for testing + runBaseProtocol 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 { - // 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) + if err = p.handleCryptoHandshake(); err != nil { + // from here on everything can be encrypted, authenticated + return DiscProtocolError, err // no graceful disconnect } 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,34 +294,62 @@ 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) { - // cryptoId is just created for the lifecycle of the handshake - // it is survived by an encrypted readwriter - var initiator bool - var sessionToken []byte - sessionToken = make([]byte, keyLen) - if _, err = rand.Read(sessionToken); err != nil { - return +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)}) } - if p.dialAddr != nil { // this should have its own method Outgoing() bool - initiator = true - } - // create crypto layer - // this could in principle run only once but maybe we want to allow - // identity switching - var crypto *cryptoId - if crypto, err = newCryptoId(p.ourID); err != nil { - return - } - // 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 + 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 + var sessionToken []byte + sessionToken = make([]byte, keyLen) + if _, err = rand.Read(sessionToken); err != nil { + return + } + if p.dialAddr != nil { // this should have its own method Outgoing() bool + initiator = true + } + // create crypto layer + // this could in principle run only once but maybe we want to allow + // identity switching + var crypto *cryptoId + if crypto, err = newCryptoId(p.ourID); err != nil { + return + } + // run on peer + // this bit handles the handshake and creates a secure communications channel with + 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) +} diff --git a/p2p/peer_error.go b/p2p/peer_error.go index 0eb7ec838d..d67c42a36d 100644 --- a/p2p/peer_error.go +++ b/p2p/peer_error.go @@ -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 { diff --git a/p2p/peer_test.go b/p2p/peer_test.go index affba1d7c7..7419488acc 100644 --- a/p2p/peer_test.go +++ b/p2p/peer_test.go @@ -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): diff --git a/p2p/server_test.go b/p2p/server_test.go index b7cc40ea22..1375344244 100644 --- a/p2p/server_test.go +++ b/p2p/server_test.go @@ -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 {