From 459542a9782f2c0c818c4c1be176f0018afdc2e3 Mon Sep 17 00:00:00 2001 From: Roman Khimov Date: Wed, 28 Aug 2019 15:43:56 +0300 Subject: [PATCH] pkg/network: convert to using binaryReader/Writer --- pkg/network/message.go | 37 ++++++--------- pkg/network/payload/address.go | 45 +++++++++--------- pkg/network/payload/getblocks.go | 23 ++++----- pkg/network/payload/headers.go | 13 +++-- pkg/network/payload/inventory.go | 31 +++++------- pkg/network/payload/merkleblock.go | 15 +++--- pkg/network/payload/version.go | 76 ++++++++++-------------------- 7 files changed, 98 insertions(+), 142 deletions(-) diff --git a/pkg/network/message.go b/pkg/network/message.go index fa39c0591..638289192 100644 --- a/pkg/network/message.go +++ b/pkg/network/message.go @@ -12,6 +12,7 @@ import ( "github.com/CityOfZion/neo-go/pkg/core/transaction" "github.com/CityOfZion/neo-go/pkg/crypto/hash" "github.com/CityOfZion/neo-go/pkg/network/payload" + "github.com/CityOfZion/neo-go/pkg/util" ) const ( @@ -147,17 +148,13 @@ func (m *Message) CommandType() CommandType { // Decode a Message from the given reader. func (m *Message) Decode(r io.Reader) error { - if err := binary.Read(r, binary.LittleEndian, &m.Magic); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &m.Command); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &m.Length); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &m.Checksum); err != nil { - return err + br := util.BinReader{R: r} + br.ReadLE(&m.Magic) + br.ReadLE(&m.Command) + br.ReadLE(&m.Length) + br.ReadLE(&m.Checksum) + if br.Err != nil { + return br.Err } // return if their is no payload. if m.Length == 0 { @@ -233,17 +230,13 @@ func (m *Message) decodePayload(r io.Reader) error { // Encode a Message to any given io.Writer. func (m *Message) Encode(w io.Writer) error { - if err := binary.Write(w, binary.LittleEndian, m.Magic); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, m.Command); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, m.Length); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, m.Checksum); err != nil { - return err + br := util.BinWriter{W: w} + br.WriteLE(m.Magic) + br.WriteLE(m.Command) + br.WriteLE(m.Length) + br.WriteLE(m.Checksum) + if br.Err != nil { + return br.Err } if m.Payload != nil { return m.Payload.EncodeBinary(w) diff --git a/pkg/network/payload/address.go b/pkg/network/payload/address.go index 044d652a8..723275338 100644 --- a/pkg/network/payload/address.go +++ b/pkg/network/payload/address.go @@ -1,7 +1,6 @@ package payload import ( - "encoding/binary" "io" "time" @@ -26,30 +25,22 @@ func NewAddressAndTime(e util.Endpoint, t time.Time) *AddressAndTime { // DecodeBinary implements the Payload interface. func (p *AddressAndTime) DecodeBinary(r io.Reader) error { - if err := binary.Read(r, binary.LittleEndian, &p.Timestamp); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.Services); err != nil { - return err - } - if err := binary.Read(r, binary.BigEndian, &p.Endpoint.IP); err != nil { - return err - } - return binary.Read(r, binary.BigEndian, &p.Endpoint.Port) + br := util.BinReader{R: r} + br.ReadLE(&p.Timestamp) + br.ReadLE(&p.Services) + br.ReadBE(&p.Endpoint.IP) + br.ReadBE(&p.Endpoint.Port) + return br.Err } // EncodeBinary implements the Payload interface. func (p *AddressAndTime) EncodeBinary(w io.Writer) error { - if err := binary.Write(w, binary.LittleEndian, p.Timestamp); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.Services); err != nil { - return err - } - if err := binary.Write(w, binary.BigEndian, p.Endpoint.IP); err != nil { - return err - } - return binary.Write(w, binary.BigEndian, p.Endpoint.Port) + bw := util.BinWriter{W: w} + bw.WriteLE(p.Timestamp) + bw.WriteLE(p.Services) + bw.WriteBE(p.Endpoint.IP) + bw.WriteBE(p.Endpoint.Port) + return bw.Err } // AddressList is a list with AddrAndTime. @@ -59,7 +50,11 @@ type AddressList struct { // DecodeBinary implements the Payload interface. func (p *AddressList) DecodeBinary(r io.Reader) error { - listLen := util.ReadVarUint(r) + br := util.BinReader{R: r} + listLen := br.ReadVarUint() + if br.Err != nil { + return br.Err + } p.Addrs = make([]*AddressAndTime, listLen) for i := 0; i < int(listLen); i++ { @@ -73,8 +68,10 @@ func (p *AddressList) DecodeBinary(r io.Reader) error { // EncodeBinary implements the Payload interface. func (p *AddressList) EncodeBinary(w io.Writer) error { - if err := util.WriteVarUint(w, uint64(len(p.Addrs))); err != nil { - return err + bw := util.BinWriter{W: w} + bw.WriteVarUint(uint64(len(p.Addrs))) + if bw.Err != nil { + return bw.Err } for _, addr := range p.Addrs { if err := addr.EncodeBinary(w); err != nil { diff --git a/pkg/network/payload/getblocks.go b/pkg/network/payload/getblocks.go index 6614f2388..0e0a3419e 100644 --- a/pkg/network/payload/getblocks.go +++ b/pkg/network/payload/getblocks.go @@ -1,7 +1,6 @@ package payload import ( - "encoding/binary" "io" "github.com/CityOfZion/neo-go/pkg/util" @@ -25,24 +24,22 @@ func NewGetBlocks(start []util.Uint256, stop util.Uint256) *GetBlocks { // DecodeBinary implements the payload interface. func (p *GetBlocks) DecodeBinary(r io.Reader) error { - lenStart := util.ReadVarUint(r) + br := util.BinReader{R: r} + lenStart := br.ReadVarUint() p.HashStart = make([]util.Uint256, lenStart) - if err := binary.Read(r, binary.LittleEndian, &p.HashStart); err != nil { - return err - } - return binary.Read(r, binary.LittleEndian, &p.HashStop) + br.ReadLE(&p.HashStart) + br.ReadLE(&p.HashStop) + return br.Err } // EncodeBinary implements the payload interface. func (p *GetBlocks) EncodeBinary(w io.Writer) error { - if err := util.WriteVarUint(w, uint64(len(p.HashStart))); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.HashStart); err != nil { - return err - } - return binary.Write(w, binary.LittleEndian, p.HashStop) + bw := util.BinWriter{W: w} + bw.WriteVarUint(uint64(len(p.HashStart))) + bw.WriteLE(p.HashStart) + bw.WriteLE(p.HashStop) + return bw.Err } // Size implements the payload interface. diff --git a/pkg/network/payload/headers.go b/pkg/network/payload/headers.go index 8699f7e9c..f4e63925a 100644 --- a/pkg/network/payload/headers.go +++ b/pkg/network/payload/headers.go @@ -14,7 +14,11 @@ type Headers struct { // DecodeBinary implements the Payload interface. func (p *Headers) DecodeBinary(r io.Reader) error { - lenHeaders := util.ReadVarUint(r) + br := util.BinReader{R: r} + lenHeaders := br.ReadVarUint() + if br.Err != nil { + return br.Err + } p.Hdrs = make([]*core.Header, lenHeaders) @@ -31,9 +35,12 @@ func (p *Headers) DecodeBinary(r io.Reader) error { // EncodeBinary implements the Payload interface. func (p *Headers) EncodeBinary(w io.Writer) error { - if err := util.WriteVarUint(w, uint64(len(p.Hdrs))); err != nil { - return err + bw := util.BinWriter{W: w} + bw.WriteVarUint(uint64(len(p.Hdrs))) + if bw.Err != nil { + return bw.Err } + for _, header := range p.Hdrs { if err := header.EncodeBinary(w); err != nil { return err diff --git a/pkg/network/payload/inventory.go b/pkg/network/payload/inventory.go index 9a9da7efd..7bed53eaa 100644 --- a/pkg/network/payload/inventory.go +++ b/pkg/network/payload/inventory.go @@ -1,7 +1,6 @@ package payload import ( - "encoding/binary" "io" "github.com/CityOfZion/neo-go/pkg/util" @@ -58,36 +57,28 @@ func NewInventory(typ InventoryType, hashes []util.Uint256) *Inventory { // DecodeBinary implements the Payload interface. func (p *Inventory) DecodeBinary(r io.Reader) error { - if err := binary.Read(r, binary.LittleEndian, &p.Type); err != nil { - return err - } + br := util.BinReader{R: r} + br.ReadLE(&p.Type) - listLen := util.ReadVarUint(r) + listLen := br.ReadVarUint() p.Hashes = make([]util.Uint256, listLen) for i := 0; i < int(listLen); i++ { - if err := binary.Read(r, binary.LittleEndian, &p.Hashes[i]); err != nil { - return err - } + br.ReadLE(&p.Hashes[i]) } - return nil + return br.Err } // EncodeBinary implements the Payload interface. func (p *Inventory) EncodeBinary(w io.Writer) error { - if err := binary.Write(w, binary.LittleEndian, p.Type); err != nil { - return err - } + bw := util.BinWriter{W: w} + bw.WriteLE(p.Type) listLen := len(p.Hashes) - if err := util.WriteVarUint(w, uint64(listLen)); err != nil { - return err - } - for i := 0; i < len(p.Hashes); i++ { - if err := binary.Write(w, binary.LittleEndian, p.Hashes[i]); err != nil { - return err - } + bw.WriteVarUint(uint64(listLen)) + for i := 0; i < listLen; i++ { + bw.WriteLE(p.Hashes[i]) } - return nil + return bw.Err } diff --git a/pkg/network/payload/merkleblock.go b/pkg/network/payload/merkleblock.go index bf46bbe81..f39b9789a 100644 --- a/pkg/network/payload/merkleblock.go +++ b/pkg/network/payload/merkleblock.go @@ -1,7 +1,6 @@ package payload import ( - "encoding/binary" "io" "github.com/CityOfZion/neo-go/pkg/core" @@ -20,18 +19,16 @@ func (m *MerkleBlock) DecodeBinary(r io.Reader) error { if err := m.BlockBase.DecodeBinary(r); err != nil { return err } + br := util.BinReader{R: r} - m.TxCount = int(util.ReadVarUint(r)) - n := util.ReadVarUint(r) + m.TxCount = int(br.ReadVarUint()) + n := br.ReadVarUint() m.Hashes = make([]util.Uint256, n) for i := 0; i < len(m.Hashes); i++ { - if err := binary.Read(r, binary.LittleEndian, &m.Hashes[i]); err != nil { - return err - } + br.ReadLE(&m.Hashes[i]) } - var err error - m.Flags, err = util.ReadVarBytes(r) - return err + m.Flags = br.ReadBytes() + return br.Err } func (m *MerkleBlock) EncodeBinary(w io.Writer) error { diff --git a/pkg/network/payload/version.go b/pkg/network/payload/version.go index f5be4b615..35f97a40f 100644 --- a/pkg/network/payload/version.go +++ b/pkg/network/payload/version.go @@ -1,9 +1,10 @@ package payload import ( - "encoding/binary" "io" "time" + + "github.com/CityOfZion/neo-go/pkg/util" ) const minVersionSize = 27 @@ -53,63 +54,36 @@ func NewVersion(id uint32, p uint16, ua string, h uint32, r bool) *Version { // DecodeBinary implements the Payload interface. func (p *Version) DecodeBinary(r io.Reader) error { - if err := binary.Read(r, binary.LittleEndian, &p.Version); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.Services); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.Timestamp); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.Port); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.Nonce); err != nil { - return err - } + br := util.BinReader{R: r} + br.ReadLE(&p.Version) + br.ReadLE(&p.Services) + br.ReadLE(&p.Timestamp) + br.ReadLE(&p.Port) + br.ReadLE(&p.Nonce) var lenUA uint8 - if err := binary.Read(r, binary.LittleEndian, &lenUA); err != nil { - return err - } + br.ReadLE(&lenUA) p.UserAgent = make([]byte, lenUA) - if err := binary.Read(r, binary.LittleEndian, &p.UserAgent); err != nil { - return err - } - if err := binary.Read(r, binary.LittleEndian, &p.StartHeight); err != nil { - return err - } - return binary.Read(r, binary.LittleEndian, &p.Relay) + br.ReadLE(&p.UserAgent) + br.ReadLE(&p.StartHeight) + br.ReadLE(&p.Relay) + return br.Err } // EncodeBinary implements the Payload interface. func (p *Version) EncodeBinary(w io.Writer) error { - if err := binary.Write(w, binary.LittleEndian, p.Version); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.Services); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.Timestamp); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.Port); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.Nonce); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, uint8(len(p.UserAgent))); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.UserAgent); err != nil { - return err - } - if err := binary.Write(w, binary.LittleEndian, p.StartHeight); err != nil { - return err - } - return binary.Write(w, binary.LittleEndian, p.Relay) + br := util.BinWriter{W: w} + br.WriteLE(p.Version) + br.WriteLE(p.Services) + br.WriteLE(p.Timestamp) + br.WriteLE(p.Port) + br.WriteLE(p.Nonce) + + br.WriteLE(uint8(len(p.UserAgent))) + br.WriteLE(p.UserAgent) + br.WriteLE(p.StartHeight) + br.WriteLE(&p.Relay) + return br.Err } // Size implements the payloader interface.